Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
akoumpa wants to merge 1 commit into
base: jingyu-ml/dmd2-qwen-image-trainer
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 48 additions & 28 deletions nemo_automodel/recipes/diffusion/step_distillation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)

(
Expand All @@ -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:
Expand Down Expand Up @@ -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 (
Expand Down
2 changes: 2 additions & 0 deletions tests/unit_tests/recipes/test_diffusion_step_distillation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down Expand Up @@ -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

Expand Down
Loading