# Training Frameworks

> Distributed Training Frameworks & Parallelism

- Skill: `sanjeevrg89/training-frameworks` (Agent Skill, multi-file: 5 files)
- Install (CLI): `npx skillmds@latest add sanjeevrg89/training-frameworks`
- Raw SKILL.md: https://api.skillmd.com/api/skills/sanjeevrg89/training-frameworks/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: Coding & Dev Tools
- Author: sanjeevrg89 (https://skillmd.com/u/sanjeevrg89)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/sanjeevrg89/training-frameworks

---


# Distributed Training Frameworks & Parallelism

Apply the judgment of an engineer who has trained frontier-scale dense and MoE models on thousands of
accelerators for years: pick the **fewest parallelism axes that make the model fit**, map the chattiest
comms onto the fastest interconnect, keep every device doing math (high MFU), and survive failures
without losing the run.

## How to use this skill

1. **Read `training-frameworks-guide.md`** in this directory — the full reference (parallelism
   strategies, framework selection, comms, checkpointing/fault tolerance, decision guide,
   anti-patterns, troubleshooting). Apply it to the task.
2. For concrete artifacts to imitate — an FSDP2 `torchrun` launch, a Megatron/DeepSpeed config sketch,
   and a `torchrun`→JobSet/Kueue mapping — read **`examples.md`**.
3. Match the surrounding repo's framework and conventions; apply the correctness rules
   (topology-aware mesh shape, comm overlap, sharded checkpoints, resumable data loader) regardless.
   Never fabricate a flag/config key — when unsure, verify against the pinned version's docs.

## Essentials (full detail in `training-frameworks-guide.md`)

- **Static state for Adam mixed precision ≈ 16–18 bytes/param** (bf16 params+grads + fp32 master +
  fp32 m,v). That plus activations is your memory budget. If it fits on one GPU → **DDP**; else shard.
- **Parallelism is memory-vs-comms.** DDP replicates everything (limit: model must fit one device).
  **ZeRO/FSDP** shard optimizer(1)→grads(2)→params(3=FSDP `FULL_SHARD`). **TP** splits inside a layer
  (Megatron). **PP** splits across layers (bubble ∝ `(p-1)/m`). **CP/SP** split the sequence. **EP**
  splits MoE experts (all-to-all). Compose into a **3D/ND device mesh**.
- **Topology is the design constraint:** TP/EP must stay inside NVLink/one node (chattiest collectives);
  PP/DP/FSDP can cross nodes; gang-schedule the whole job into one fabric domain.
- **FSDP2** (`fully_shard`, DTensor-based) is the current PyTorch path; use a **HYBRID/2-D mesh**
  multi-node so the param all-gather stays on NVLink. Wrap per transformer block, not whole-model.
- **Always-on memory tools:** selective activation checkpointing, gradient accumulation (skip the
  all-reduce on non-final micro-batches), bf16 (fp8 on the big GEMMs with careful scaling), fused/
  distributed optimizer (ZeRO-1).
- **Comm overlap is MFU.** Overlap grad all-reduce / FSDP all-gather+reduce-scatter / TP comms with
  compute; lost overlap, a CPU-bound data loader, or `m≈p` pipelines are the top MFU killers.
- **Collectives:** all-reduce (DDP/TP), reduce-scatter (FSDP grads), all-gather (FSDP params),
  all-to-all (MoE), P2P (PP). NCCL/RCCL on GPU; XLA collectives on TPU.
- **Checkpoint sharded + async** (PyTorch DCP / Megatron dist / DeepSpeed universal) — never gather to
  rank 0. Use elastic `torchrun`/torchft/Pathways + a **resumable data loader** to recover from node
  loss. Optimize **MFU**, then **goodput**.
- **Framework fit:** TorchTitan/FSDP2 for PyTorch-from-scratch; Megatron-Core/NeMo for frontier
  dense/MoE on NVIDIA; DeepSpeed for ZeRO-3 + offload; HF Accelerate/Trainer/TRL for fine-tune;
  JAX MaxText/Levanter/Paxml+Pathways on TPU; Ray Train to orchestrate. Don't add an axis you don't need.

## Related skills

- `[[ml-frameworks]]` — PyTorch/JAX/XLA compute internals, DTensor, `torch.compile`, GPU/TPU (sibling).
- `[[aiml-on-kubernetes]]` — umbrella for training/inference/RL on K8s & GKE.
- `[[jobset-leaderworkerset]]` — the multi-host gang primitive a `torchrun` rendezvous targets.
- `[[kueue-advanced]]` — gang scheduling, quota, and Topology-Aware Scheduling for training jobs.
- `[[slurm-hpc-on-kubernetes]]` — Slurm/MPI/RDMA, MPIJob, Volcano for HPC-style launches.
- `[[gke-master]]` — TPU/GPU node pools, networking, and placement on GKE.
- `[[serving-frameworks]]` — inference/serving (vLLM, SGLang, TensorRT-LLM), the other side of the stack.