From a0cd67839edf563b0901d2cb1cb659a961bdf385 Mon Sep 17 00:00:00 2001 From: Alexandros Koumparoulis Date: Tue, 25 Aug 2026 14:59:49 -0700 Subject: [PATCH] fix(diffusion): rank-symmetric discriminator all-reduce and deferred FSDP2 grad sync in DMD2 Two independent distributed issues in the DMD2 objective. 1. `_synchronize_discriminator_gradients` skipped parameters whose `.grad` was `None`, which makes the *number* of collectives rank-dependent. NCCL matches collectives by call order, so a single rank taking a different branch hangs the whole job with no diagnostic -- no timeout message and no rank identification. The replicated discriminator uses every parameter on every forward today, so there is no reachable trigger from the shipped recipes, but the guard buys nothing: materializing a zero gradient costs one allocation and keeps every rank issuing the same collectives in the same order. 2. The microbatch loop never deferred FSDP2's gradient reduce-scatter to the final microbatch. It 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, i.e. a full reduce-scatter set per microbatch instead of only 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. The test fixtures gain `defer_fsdp_grad_sync`, which the real `TrainDiffusionRecipe` sets but the partial stand-ins did not. Signed-off-by: Alexandros Koumparoulis --- .../recipes/diffusion/step_distillation.py | 76 ++++++++++++------- .../test_diffusion_step_distillation.py | 2 + 2 files changed, 50 insertions(+), 28 deletions(-) diff --git a/nemo_automodel/recipes/diffusion/step_distillation.py b/nemo_automodel/recipes/diffusion/step_distillation.py index f7a2173ab0..724d8ac336 100644 --- a/nemo_automodel/recipes/diffusion/step_distillation.py +++ b/nemo_automodel/recipes/diffusion/step_distillation.py @@ -26,6 +26,7 @@ from torch import nn from nemo_automodel.components.checkpoint.stateful_wrappers import ModelState, OptimizerState +from nemo_automodel.components.distributed.utils import get_sync_ctx from nemo_automodel.components.training.utils import ( clip_grad_norm, prepare_after_first_microbatch, @@ -339,7 +340,8 @@ def train_batch_group( num_microbatches = len(batch_group) for index, micro_batch in enumerate(batch_group): - if index == num_microbatches - 1: + is_final_microbatch = index == num_microbatches - 1 + if is_final_microbatch: prepare_for_final_backward([active_model], pp_enabled=False) ( @@ -354,30 +356,42 @@ def train_batch_group( "encoder_hidden_states": text_embeddings, "encoder_hidden_states_mask": text_mask, } - if student_phase: - losses = self.dmd_pipeline.compute_student_loss( - latents, - noise, - negative_encoder_hidden_states=negative_embeddings, - negative_encoder_hidden_states_mask=negative_mask, - **kwargs, - ) - else: - losses = self.dmd_pipeline.compute_fake_score_loss(latents, noise, **kwargs) - - total = losses["total"] - if recipe.check_loss and not torch.isfinite(total).all(): - raise FloatingPointError(f"Non-finite DMD2 {phase} loss at step {global_step}.") - (total / num_microbatches).backward() - reported_total = total.detach() - - if not student_phase and self.discriminator_optimizer is not None: - discriminator_losses = self.dmd_pipeline.compute_discriminator_loss(latents, noise, **kwargs) - discriminator_total = discriminator_losses["total"] - if recipe.check_loss and not torch.isfinite(discriminator_total).all(): - raise FloatingPointError(f"Non-finite DMD2 discriminator loss at step {global_step}.") - (discriminator_total / num_microbatches).backward() - reported_total = reported_total + discriminator_total.detach() + # Defer FSDP2's gradient reduce-scatter to the last microbatch. The + # prepare_for_* helpers above are gated on a method only the MoE mixin + # defines, so both are no-ops for a diffusers transformer -- without this + # context `fsdp.defer_fsdp_grad_sync` has no effect here and every + # microbatch pays a full reduce-scatter set. Mirrors what the non-DMD2 + # branch of this recipe already does. + sync_context = get_sync_ctx( + active_model, + is_final_microbatch, + defer_fsdp_grad_sync=recipe.defer_fsdp_grad_sync, + ) + with sync_context: + if student_phase: + losses = self.dmd_pipeline.compute_student_loss( + latents, + noise, + negative_encoder_hidden_states=negative_embeddings, + negative_encoder_hidden_states_mask=negative_mask, + **kwargs, + ) + else: + losses = self.dmd_pipeline.compute_fake_score_loss(latents, noise, **kwargs) + + total = losses["total"] + if recipe.check_loss and not torch.isfinite(total).all(): + raise FloatingPointError(f"Non-finite DMD2 {phase} loss at step {global_step}.") + (total / num_microbatches).backward() + reported_total = total.detach() + + if not student_phase and self.discriminator_optimizer is not None: + discriminator_losses = self.dmd_pipeline.compute_discriminator_loss(latents, noise, **kwargs) + discriminator_total = discriminator_losses["total"] + if recipe.check_loss and not torch.isfinite(discriminator_total).all(): + raise FloatingPointError(f"Non-finite DMD2 discriminator loss at step {global_step}.") + (discriminator_total / num_microbatches).backward() + reported_total = reported_total + discriminator_total.detach() total_losses.append(reported_total) if index == 0: @@ -577,9 +591,15 @@ def _synchronize_discriminator_gradients(self, recipe: TrainDiffusionRecipe) -> return dp_size = recipe._get_dp_group_size() for parameter in self.discriminator.parameters(): - if parameter.grad is not None: - dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM, group=recipe._get_dp_group()) - parameter.grad.div_(dp_size) + if parameter.grad is None: + # Skipping would make the number of collectives rank-dependent, and one + # rank taking a different branch deadlocks the job with no diagnostic -- + # no timeout message, no rank identification. Materializing a zero costs + # one allocation and keeps every rank issuing the same collectives in the + # same order. + parameter.grad = torch.zeros_like(parameter) + dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM, group=recipe._get_dp_group()) + parameter.grad.div_(dp_size) def _require_ready(self) -> None: if ( diff --git a/tests/unit_tests/recipes/test_diffusion_step_distillation.py b/tests/unit_tests/recipes/test_diffusion_step_distillation.py index 033d080697..93591bc896 100644 --- a/tests/unit_tests/recipes/test_diffusion_step_distillation.py +++ b/tests/unit_tests/recipes/test_diffusion_step_distillation.py @@ -143,6 +143,7 @@ def _discriminator_loss( clip_grad_max_norm=100.0, grad_clip_foreach=False, check_loss=True, + defer_fsdp_grad_sync=True, ) micro_batch = { "image_latents": torch.zeros(1, 1, 1, 1), @@ -212,6 +213,7 @@ def _checkpoint_recipe(checkpoint_dir: Path, base_weight: float) -> TrainDiffusi recipe.clip_grad_max_norm = 100.0 recipe.grad_clip_foreach = False recipe.check_loss = True + recipe.defer_fsdp_grad_sync = True recipe.cpu_offload = False recipe.peft_config = None