fix(flux): clone the validation loss out of Megatron's rescale path

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
 

tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py

@@ -341,3 +341,56 @@ 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)