RustyTorch++ — Phase 3 Plan (Distributed Training & Multi‑GPU Execution)

Objective (4–5 months): Enable scalable training across multiple GPUs and nodes with topology‑aware communication, hybrid parallelism (data/tensor/pipeline/sequence), elastic recovery, and sharded states (FSDP/ZeRO‑style). Maintain determinism where declared and ship CI gates, benchmarks, and observability tailored to distributed runs.

Reference paths

RUSTYTORCH_ROOT: /home/osobh/projects/rustytorch

RUSTG_COMPILER: /home/osobh/projects/rust/rustg (RTX 5090, sm_120)

STRATOSWARM_ROOT: /home/osobh/projects/stratoswarm (deployment in later phases)

Darwin apps remain deferred. Phase 3 focuses on RustyTorch++ core scaling.

1. Scope & Non‑Goals

In scope

Communication backend: NCCL (CUDA) / RCCL (ROCm) integration and process group API

Topology discovery & hints (PCIe, NVLink, NVSwitch, InfiniBand)

Parallelism: data parallel (DP), tensor parallel (TP), pipeline parallel (PP), sequence parallel (SP); auto‑hybrid partitioner

Sharded training: FSDP/ZeRO‑style (gradients, optimizer, and parameters)

Elastic engine: membership changes (join/leave), pre‑emption‑safe checkpoints (WAL), retry/backoff

Overlap: comm/compute scheduling; bucketization; gradient accumulation

Determinism surface for distributed runs (documented limits)

Benchmarks & CI at 1, 2, 4, 8 GPUs (single‑node) and 2–4 nodes (where available)

Out of scope (Phase 4+)

Advanced auto‑kernel synthesis & compiler autotuning (beyond basic sweeps)

Inference serving/scheduling stack (separate track)

2. Success Criteria (Exit / Phase‑Gate)

Functional DP/TP/PP/SP with a simple configuration API; hybrids validated on reference models.

Sharded training reduces per‑GPU memory by ≥40% versus DP baseline at equal global batch.

Throughput scaling: ≥0.8× efficiency from 1→8 GPUs (single node); ≥0.7× to multi‑node on IB/NVLink.

Elasticity: mid‑run failure of one rank recovers automatically; pre‑emption resumes from latest durable checkpoint.

Determinism: declared deterministic mode yields repeatable results within tolerance across identical hardware and seeds; non‑deterministic surfaces are documented.

Observability: per‑rank and aggregated timelines/metrics; bottleneck attribution (compute vs. network) with actionable hints.

3. Architecture Work

A. Process Group & Collectives

Process group abstraction (rtx-dist) with pluggable backends (NCCL/RCCL; CPU fallback)

Collectives: allreduce, reduce‑scatter, allgather, broadcast; async handles with stream selection

Fault detection (timeout/health) and retry/backoff policies

B. Topology & Placement

Discover PCIe/NVLink/NVSwitch/IB; build bandwidth/latency matrix

Rank mapping strategies (ring, mesh, tree) selected by topology and model parallel mode

Affinity & NUMA pinning policies

C. Parallelism Implementations

DP with gradient bucketing & overlap

TP (tensor‑slicing of large layers) with fused all‑reduce/all‑gather epilogues

PP with micro‑batching & 1F1B schedule; activation checkpointing integration

SP for attention/sequence‑length scaling

Auto‑hybrid planner chooses DP/TP/PP mix from model graph & hardware profile

D. Sharded Training (FSDP/ZeRO‑style)

Partition parameters/gradients/optimizer states

State prefetch & offload hooks (CPU/NVMe) with backpressure

Checkpoint format with shard metadata; streaming save/restore

E. Elasticity & Reliability

Rendezvous service; rank reconfiguration without restart

WAL checkpoints + periodic durable snapshots

Deterministic resume: resume graph/seed/capture hashes; verify parity

4. APIs & Configuration

dist.init(backend, devices, env); returns process group

parallel.configure({dp, tp, pp, sp, micro_batch, bucket_bytes})

fsdp.configure({shard_strategy, offload={'cpu'|'nvme'|none}, auto_wrap=true})

elastic.configure({rendezvous, preempt_policy, checkpoint_interval})

Minimal user code to switch from single‑GPU to distributed (≤ 10 LOC)

5. Benchmarks & Baselines

Models

LLM‑block (attention + MLP) and GPT‑style stack; ViT block for CV parity

Scales

Single node: 1, 2, 4, 8 GPUs (RTX 5090)

Multi‑node (if available): 2–4 nodes with IB

Metrics

Tokens/sec (or images/sec), step time, p50/p95/p99, allreduce time %, overlap %, network BW utilization, peak memory, frag%

6. CI Gates (Phase 3)

Scaling Gate: efficiency thresholds (≥0.8× 1→8 GPUs single node)

Memory Gate: FSDP/ZeRO memory reduction ≥40% vs DP

Reliability Gate: injected rank failure recovers; determinism resume parity

Security Gate: dependency pinning; SBOM/sign; secrets hygiene

7. Observability & Telemetry

Per‑rank Nsight/rocprof traces; aggregated timeline

Prometheus: step_time_ms, tokens_per_s, allreduce_pct, overlap_pct, net_bw_gbps, alloc_frag_pct, oom_count

Bottleneck classifier (compute‑bound vs. network‑bound) with suggestions

8. Deliverables

rtx-dist crate (process group, collectives, rendezvous)

Parallelism modules (DP/TP/PP/SP) + auto‑hybrid planner

rtx-fsdp (or integrated) sharding with offload hooks & checkpoint format

Elastic runtime with WAL checkpoints and pre‑emption support

Benchmarks & dashboards for scaling curves; CI gates

Documentation: docs/distributed.md, docs/fsdp.md, examples

9. RACI — Phase 3 Agent Ownership

Conductor (Orchestrator) — Accountable: cross‑agent plan, gate readiness

Distributed Strategy Agent — Responsible: topology discovery, planner, overlap scheduling

Rust Engineer — Responsible: runtime hooks, stream/collective integration, checkpoint I/O

Kernel Synthesizer Agent — Consulted: fused epilogues for collective‑adjacent ops

Memory Planner Agent — Responsible: sharding integration; offload/backpressure policies

Performance Engineer — Accountable: scaling benches, overlap tuning, efficiency gates

Security Engineer & Auditor — Accountable: dependency/secret policies; SBOM/sign

DevOps Engineer — Responsible: multi‑GPU runners; optional multi‑node staging; artifact storage

Governance & Provenance Agent — Responsible: checkpoint provenance; reproducible resumes

Agent Organizer — Accountable: routing, tool allowlists, freeze switch on gate failures

10. Risks & Mitigations

Allreduce bottlenecks → bucket tuning; overlap; topology‑aware rings/trees

Elasticity edge cases → chaos tests; forced rank drops; partial shard recovery

Determinism drift → documented nondet surfaces; seed discipline; stable reductions; resume checks

Network variance → pin NIC/drivers; IB tuning; retry/backoff

Memory pressure → shard/offload thresholds; activation checkpointing; streaming checkpoints

11. Timeline (suggested)

Weeks 1–2: Process group, collectives, rendezvous prototype; single‑node DP

Weeks 3–4: TP/PP/SP implementations; overlap scheduling; hybrid planner

Weeks 5–6: FSDP/ZeRO sharding & checkpoint format; offload hooks

Weeks 7–8: Elastic runtime; chaos tests; multi‑GPU scaling benches; CI gates; docs; phase‑gate review

12. Phase‑4 Handover Seeds

Compiler autotuning & auto‑kernel synthesis plan (hardware profile DB, caches)

Inference scheduler design seeds (continuous batching, KV paging)

Governance extension for distributed reproducibility across clusters

— End of Phase 3 Plan —