distributed-training-backends

Distributed Training Backends

Use this sub-skill when an agent must choose or configure an Accelerate backend beyond a plain single-process training loop: DeepSpeed ZeRO, FSDP/FSDP2, Megatron-LM, torch native tensor/context/sequence parallelism, TPU/XLA, Gaudi/HPU, FP8 or low precision, bitsandbytes/torchao quantization, torch.compile, Local SGD, or DDP communication hooks.

Routing

• For accelerate config, accelerate launch, config-file locations, or generic CLI flag syntax, use ../configuration-and-cli/ first, then return here for backend-specific keys and compatibility.
• For where to instantiate Accelerator, how to call prepare, backward, accumulate, or unwrap models in a training loop, use ../training-loop-integration/.
• For save_state, load_state, FSDP weight merging, DeepSpeed checkpoint folders, or tracker artifacts, use ../checkpointing-and-tracking/.
• For backend choice, plugin objects, optional dependencies, hardware constraints, or config-vs-plugin precedence, stay in this sub-skill.

Fast Backend Triage

• Choose DeepSpeed for ZeRO optimizer/parameter partitioning, CPU/NVMe offload, DeepSpeed-specific optimizers/schedulers, or existing DeepSpeed JSON; see references/deepspeed-fsdp.md.
• Choose FSDP/FSDP2 for PyTorch-native sharding, transformer_based_wrap, state-dict control, CPU-RAM-efficient Transformers loading, and FSDP2 composition with native parallelism; see references/deepspeed-fsdp.md.
• Choose torch native parallelism when using parallelism_config for data replication/sharding, tensor parallelism, context parallelism, or DeepSpeed sequence parallelism; see references/parallelism-and-precision.md.
• Choose Megatron-LM only when the training stack is Megatron-aware and needs tensor/pipeline/sequence parallelism plus Megatron dummy optimizer/scheduler integration.
• Choose TPU/XLA only when torch_xla and TPU runtime are available; CPU-only inspection can validate config shape, not TPU execution.
• Choose FP8/quantization/compile/communication options as add-ons after the distributed backend is selected; verify optional dependency and hardware support before promising speedups.

Required References

• references/deepspeed-fsdp.md — DeepSpeed, FSDP/FSDP2 selection, launch/config implications, and plugin usage.
• references/parallelism-and-precision.md — torch native parallelism, Megatron-LM, TPU/XLA, Gaudi/HPU, FP8, low precision, quantization, compilation, Local SGD, and DDP hooks.
• references/api-reference.md — primary Accelerate plugin/config classes, common constructor fields, environment/config key mapping, and minimal snippets.
• references/troubleshooting.md — dependency, hardware, ZeRO/offload, FSDP wrapping, FP8 mismatch, multi-node, and precedence diagnostics.

Bundled Helper

Use scripts/validate_backend_config.py to statically validate JSON or YAML snippets for backend key shape and common contradictions:

python sub-skills/distributed-training-backends/scripts/validate_backend_config.py path/to/config.yaml
python sub-skills/distributed-training-backends/scripts/validate_backend_config.py --format deepspeed deepspeed.json
python sub-skills/distributed-training-backends/scripts/validate_backend_config.py --check-imports config.yaml

The helper does not launch distributed jobs, import Accelerate, initialize process groups, or require GPUs/TPUs. It reports missing optional packages only as hints when --check-imports is passed.

Safety Notes

• Label DeepSpeed, FSDP2, Megatron-LM, TPU/XLA, FP8, bitsandbytes, torchao, transformer-engine, ms-amp, and Gaudi as optional dependency or hardware paths.
• Do not promise that config parsing proves runtime correctness; distributed execution still depends on process count, network/rendezvous setup, device availability, backend package versions, and model compatibility.
• Prefer explicit plugin objects in Python when code must be reproducible, and prefer config files when users rely on accelerate config/launch workflows. If both are present, check precedence and mismatch rules in references/troubleshooting.md.