fix(diffusion): rank-symmetric DMD2 all-reduce and deferred FSDP2 grad sync

[akoumpa]

Two independent distributed issues in the DMD2 objective.

Stacked on #3298 — step_distillation.py does not exist on main, so this targets
that branch rather than main.

1. Rank-dependent collective count in the discriminator all-reduce

_synchronize_discriminator_gradients skipped parameters whose .grad was None:

for parameter in self.discriminator.parameters():
    if parameter.grad is not None:          # <- makes the collective COUNT rank-dependent
        dist.all_reduce(parameter.grad, ...)

NCCL matches collectives by call order, so one rank taking a different branch hangs
the whole job with no diagnostic — no timeout message, no rank identification.

The replicated discriminator uses every parameter on every forward today, so there is
no reachable trigger from the shipped recipes — this is a latent hazard rather than
a live bug. But the guard buys nothing: materializing a zero gradient costs one
allocation and keeps every rank issuing the same collectives in the same order.

Verified: the collective count is now constant at 4 whether or not a gradient is
None; before the change it was 3 vs 4 across ranks.

2. FSDP2 gradient sync was never deferred to the final microbatch

The microbatch loop relies on prepare_for_grad_accumulation /
prepare_for_final_backward, which are gated on a method only the MoE mixin defines —
so both are no-ops for a diffusers transformer, and fsdp.defer_fsdp_grad_sync had
no effect on this path.

Measured on a sharded Wan transformer: 20 reduce_scatter_tensor calls for 4
microbatches where 5 are needed — a full reduce-scatter set per microbatch instead of
only on the last.

This is a performance issue, not a correctness one (reduce-scatter of a sum equals the
sum of reduce-scatters), and it is invisible at the shipped configs, whose
global_batch_size // (local_batch_size * dp_size) is 1. It becomes real as soon as
global_batch_size is raised.

The non-DMD2 branch of this same recipe already handles it via get_sync_ctx; the DMD2
branch now does too.

Testing

• tests/unit_tests/recipes/test_diffusion_step_distillation.py: 42 passed
• Wider sweep (recipes, datasets/diffusion, _diffusers, flow_matching): 1721
  passed; the 5 test_base_recipe.py failures are pre-existing and reproduce on a
  clean main checkout
• ruff check / ruff format --check clean

The test fixtures gain defer_fsdp_grad_sync, which the real TrainDiffusionRecipe
sets but the partial stand-ins did not — added there rather than defensively
getattr-ing it in production, since that would mask the same class of bug later.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.