[StarrickLiu]

Summary

DedicatedParamGroup.replicate_grads_enabled = False is meant to mirror FSDP2's set_requires_all_reduce(False): run the Stage-1 shard reduce on every micro-batch, hold the shard-only partial in _partial_reduce_by_param, and fold it into the next Stage-2 replicate reduce. On main that path never worked:

• Unreachable stash branch. It was guarded by has_replicate and not self.replicate_grads_enabled, but has_replicate already requires the gate to be on. Gated micro-batches fell through to the shard-only branch, and wait_for_reduce handed the shard-only partial to the optimizer as a finished gradient, missing the replicate peers' contributions.
• Double counting. Even if the branch had been reached, it kept _pending_reduce, so the partial would also have been unpacked into _reduced_grad.
• Stream ordering. The fold ran on the current stream instead of on reduce_stream after the Stage-1 collectives.
• Sharded AdamW. Sharded AdamW partials were not stashed at all.

Changes

• fsdp2/group.py: when the replicate group exists but the gate is off, fold Stage-1 outputs into _partial_reduce_by_param on reduce_stream. This applies on the shard-owner rank for owner-managed params and on every shard rank for sharded AdamW. _pending_reduce is dropped (after record_stream) so nothing is unpacked, and _sharded_adamw_grad is cleared for the gated micro-batch. The Stage-2 flush adds the flat partial and pins it on replicate_stream.
• optim/muon.py: zero_grad clears partials that were never flushed because the gate was never re-enabled.
• New tests/distributed/test_hsdp_replicate_gate_accum.py (2×2 HSDP, Muon, sharded AdamW and FSDP2 params): micro-batch 1 runs with the replicate gate off, micro-batch 2 runs normally. The test compares the gradients the optimizer consumes against no_sync accumulation over 3 steps, and checks that no partial survives a step.

Validation (4× H100, torch 2.10 + cu128)

	worst relative grad diff vs no_sync
main	1.08 (Muon and sharded AdamW wrong every step)
this PR	1.4e-7

test_hsdp_correctness, test_grad_accum, test_hsdp_async_correctness, test_muon_step and test_checkpoint still pass. CPU unit tests are unchanged.

Note for #19

#19 reformats fsdp2/group.py and moves sharded AdamW grads into _pending_reduce, so whichever PR lands second needs a rebase. After #19, the gated branch must keep sharded AdamW entries from _pending_reduce on every shard rank (not only owner_shard) and should not collect them twice through sharded_partial_pending.

🤖 Generated with Claude Code