docs/distributed.md

Distributed story (0.2)

The diffusion split (the standard paradigm, 0.2)

Diffusion training has TWO sanctioned paths, by design:

1. Native amoe.diffusion.train — single-GPU + the same DDP posture as the LM trainer below (rank-sharded cache sampling, find_unused_parameters=False, rank-0 saves). This is the framework-level path for the certified single-card recipes; correct-by-construction, multi-GPU smoke deferred.
2. The diffusion-pipe fork (https://github.com/AbstractEyes/diffusion-pipe) — the PRODUCTION multi-GPU trainer: DeepSpeed pipeline engine with the aleph relays baked into the model integrations (attach at declared dtype, freeze-by-lr-0 trunk, plain-Adam branch, adapter-only saves in the amoe anchor format). Proven end-to-end on the Anima 2B DiT (r2 exp004). Multiband w_bands plumbing is working at pipeline_stages=1; stages>1 is documented-deferred.

No FSDP on either path, same reasons as below.

Implemented

• Device-following attach: each adapter/dispatch is placed on its wrapped block's device at attach time. This is the entire device_map="auto" compatibility story for inference — accelerate shards the trunk, the adapters follow, nothing else changes.
• DDP-aware train/align: when torch.distributed is initialized (launch with torchrun), the trainer rank-shards row sampling (per-rank generator offset), wraps in DDP with find_unused_parameters=False (all trainable params fire every step in the default recipes), and logs/saves on rank 0 only. Gradient sync is trivial at these sizes (6.3M train / ~8k align params). Verified on one GPU only — the DDP branch is correct by construction but multi-GPU smoke is deferred.
• Gradient checkpointing uses use_reentrant=False exclusively (the DDP-compatible mode).

Documented non-goals (0.1)

• FSDP: the layer-replacement wrap conflicts with auto-wrap policies; home/key_proj buffers need sharded-state-dict care; and sharding megabyte-scale adapters buys nothing. Only relevant when the trunk must shard. Sketch for later: attach after FSDP wrapping at block granularity with a custom wrap policy that treats BlockWithAdapter as a unit.
• Tensor/pipeline parallel, DeepSpeed: the dispatch reads the full residual per block; TP would require the address/adapter to be sharded consistently with the trunk's TP plan — out of scope.
• Edge case encoded: AlignConfig.train_new_anchor=True plus runtime masks during training creates unused parameters; the aligner flips find_unused_parameters=True in that configuration.