214 lines
7.2 KiB
Markdown
214 lines
7.2 KiB
Markdown
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 —
|