Avoid redundant gradient synchronization during accumulation

[ReichtumQian]

Summary

Flame currently synchronizes replicated gradients after every microbatch during gradient accumulation, resulting in redundant communication.

Replica synchronization is now deferred until the final microbatch. Composable DDP disables gradient synchronization on non-final microbatches, while HSDP disables only the replica all-reduce and keeps reduce-scatter enabled. Pure FSDP behavior is unchanged.

Problem

The training loop calls backward() once per microbatch. With composable DDP, each backward currently triggers gradient synchronization, even though the optimizer is updated only after all accumulated microbatches have completed.

As a result, K gradient accumulation steps perform K replica synchronizations instead of synchronizing once after the accumulated gradient is ready. This introduces redundant communication and can substantially reduce training throughput.

For HSDP, the same redundant replica all-reduce can be avoided while still performing reduce-scatter after each microbatch, preserving sharded gradient storage.

Benchmark

Benchmarked with Flame's official transformer_340M.json configuration:

• 4 × NVIDIA RTX 4090
• 373,867,520 parameters
• sequence length: 2,048
• per-device batch size: 1
• BF16 composable DDP
• torch.compile disabled
• median per-device throughput, excluding the first two warmup steps

Gradient accumulation	Before	After	Improvement
1	1,301.5 tokens/s	1,313.5 tokens/s	+0.9%
4	1,253.5 tokens/s	4,057.5 tokens/s	+223.7% (3.24×)

Gradient accumulation 1 is effectively unchanged, while accumulation 4 improves throughput by 3.24×. The checked loss matched between the two versions, and peak memory remained approximately 9.76 GiB per GPU.