From c64a481295437f004b34fe5859e0b5eed39de936 Mon Sep 17 00:00:00 2001 From: Ossi Lehtikangas Date: Wed, 2 Sep 2026 04:27:54 +0000 Subject: [PATCH 1/2] fix(flux): clone the validation loss out of Megatron's rescale path Megatron's forward_step takes the first element of what a loss function returns as the tensor to backpropagate and rescales it IN PLACE -- `output_tensor *= cp_group_size`, then `output_tensor /= num_microbatches` -- and only then stores the reported dict. The diffusion validation path put `loss_sum.detach()` in that dict, and a detached tensor shares storage with the one it came from, so the reported loss was rescaled along with it. What reached the caller was therefore the true validation loss divided by the number of microbatches. At one microbatch per rank per step the divisor is 1 and nothing shows, which is why every shape measured so far looked right. At two it halves, and under mlperf_mode the halved value is what the convergence gate is compared against -- so a run reports convergence at roughly half the samples it actually needed, and reports it as a pass. Measured on Flux 12B at matched global batch 512: micro batch 32 reported 0.630950 at step 100 where micro batch 64 reported 1.266634, a ratio of 2.007. With this fix the two agree to 0.7%, and to 0.01% by step 200. Context parallelism has the mirror problem, inflating the reported loss by cp_group_size. The training path a few lines below already clones for this reason. --- primus/backends/megatron/diffusion_trainer.py | 12 +++- .../training/test_diffusion_trainer.py | 55 +++++++++++++++++++ 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/primus/backends/megatron/diffusion_trainer.py b/primus/backends/megatron/diffusion_trainer.py index dd91f3285..e0c7521a7 100644 --- a/primus/backends/megatron/diffusion_trainer.py +++ b/primus/backends/megatron/diffusion_trainer.py @@ -278,7 +278,17 @@ def val_loss_func(output_tensor, non_loss_data=False): sample_count = torch.tensor( loss_per_sample.numel(), dtype=loss_sum.dtype, device=loss_sum.device ) - return loss_sum, {"loss": (loss_sum.detach(), sample_count.detach())} + # CLONE, NOT JUST DETACH. Megatron's forward_step treats the first + # element of this pair as the tensor to backpropagate and rescales it + # IN PLACE before storing the dict below -- `output_tensor *= + # cp_group_size`, then `output_tensor /= num_microbatches`. A detached + # view shares that storage, so the reported loss gets rescaled with it + # and what reaches the caller is the true loss divided by the number of + # microbatches. That is invisible at one microbatch per rank per step + # and halves the reported validation loss at two, which under + # mlperf_mode trips the convergence gate at roughly half the samples it + # should. The training path below clones for the same reason. + return loss_sum, {"loss": (loss_sum.detach().clone(), sample_count.detach())} return noise_pred, val_loss_func diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py index 487af4193..f7bda24a3 100644 --- a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py @@ -341,3 +341,58 @@ def test_forward_step_loss_func_with_non_loss_data(self, monkeypatch: pytest.Mon # Should return output_tensor directly assert result is output + + def test_validation_report_survives_megatrons_in_place_rescale( + self, monkeypatch: pytest.MonkeyPatch + ): + """The reported validation loss must not move when Megatron rescales the + tensor it backpropagates. + + megatron.core.pipeline_parallel.schedules.forward_step takes the first + element of the pair the loss function returns and rescales it in place -- + ``output_tensor *= cp_group_size``, then ``output_tensor /= + num_microbatches`` -- and only then stores the reported dict. While that + dict held a detached *view* of the same tensor, the report was rescaled + along with it, so the reported validation loss was the true loss divided + by the microbatch count: right at one microbatch per rank per step, half + the true value at two. Under mlperf_mode that halves the number the + convergence gate is compared against, so a run claimed to converge at + about half the samples it really needed. + """ + import torch + + trainer = _build_diffusion_trainer(monkeypatch) + trainer._scheduler = Mock() + trainer.runtime_state = Mock() + trainer.runtime_state.update_metrics = Mock() + + # Four samples, so a summed loss and a per-sample mean cannot be confused. + noise_pred = torch.full((4, 1), 3.0) + clean_latents = torch.zeros(4, 1) + noise = torch.ones(4, 1) + + monkeypatch.setattr( + "primus.backends.megatron.training.diffusion.forward_step.flux_forward_step_func", + lambda *args, **kwargs: (noise_pred, clean_latents, noise, None, {}, True), + ) + + model = Mock() + model.training = False + _, val_loss_func = trainer.forward_step(Mock(), model) + + loss_sum, reported = val_loss_func(noise_pred) + + # target = noise - clean_latents = 1, so each element contributes + # (3 - 1) ** 2 = 4, and the sum over four samples is 16. + assert reported["loss"][0].item() == pytest.approx(16.0) + assert reported["loss"][1].item() == pytest.approx(4.0) + + # Exactly what forward_step does to the tensor it backpropagates. + loss_sum *= 1 # cp_group_size, 1 without context parallelism + loss_sum /= 2 # num_microbatches, 2 at micro batch 32 and GBS 1024 on 16 ranks + + # The rescale has to have happened, or this asserts nothing. + assert loss_sum.item() == pytest.approx(8.0) + + assert reported["loss"][0].item() == pytest.approx(16.0) + assert reported["loss"][1].item() == pytest.approx(4.0) From 0e8f67de6ee5f0c0a692e993100bc0f71f73cbb3 Mon Sep 17 00:00:00 2001 From: Ossi Lehtikangas Date: Wed, 2 Sep 2026 06:00:11 +0000 Subject: [PATCH 2/2] chore: satisfy black line-length 110 on the new regression test signature --- .../megatron/diffusion/training/test_diffusion_trainer.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py index f7bda24a3..fb54284d7 100644 --- a/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py +++ b/tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py @@ -342,9 +342,7 @@ def test_forward_step_loss_func_with_non_loss_data(self, monkeypatch: pytest.Mon # Should return output_tensor directly assert result is output - def test_validation_report_survives_megatrons_in_place_rescale( - self, monkeypatch: pytest.MonkeyPatch - ): + def test_validation_report_survives_megatrons_in_place_rescale(self, monkeypatch: pytest.MonkeyPatch): """The reported validation loss must not move when Megatron rescales the tensor it backpropagates.