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

[olehtika]

What is wrong

DiffusionPretrainTrainer.forward_step's validation loss function returns

return loss_sum, {"loss": (loss_sum.detach(), sample_count.detach())}

Megatron's forward_step takes the first element as the tensor to backpropagate and rescales it in place before storing the reported dict:

output_tensor, loss_reduced = outputs
output_tensor *= cp_group_size
output_tensor /= num_microbatches
forward_data_store.append(loss_reduced)

loss_sum.detach() shares storage with loss_sum, so the reported numerator is rescaled along with the backpropagated tensor. What reaches the caller is the true validation loss divided by the number of microbatches (and, with context parallelism, multiplied by cp_group_size).

Why it matters

At one microbatch per rank per step the divisor is 1, so nothing is visible — which is why every shape measured until now looked correct. At two microbatches the reported validation loss halves, and under mlperf_mode that halved value is what the convergence gate is compared against. A run then reports convergence at roughly half the samples it actually needed, and reports it as a pass. Any Flux MLPerf configuration that reaches its global batch through gradient accumulation is affected.

Measured on Flux 12B (MI355X), matched global batch 512, same stack, same node:

micro batch	accumulation	step 100	step 200	step 300
64	1	1.266634	0.860864	0.799235
32 (before)	2	0.630950	0.436635	— early-stopped at the gate
32 (after)	2	1.261901	0.871321	0.797608

Before the fix the ratio between the two configurations is 2.007 at step 100 and 1.972 at step 200. After it they agree to 0.7% at step 100 and 0.01% at step 200. The same check on two nodes at global batch 1024: micro batch 32 now reports 0.767331 at step 256 against micro batch 64's 0.766881, where before the fix it reported 0.383666 and tripped the 0.586 gate at its first evaluation.

Evaluation coverage was never at fault — both configurations report covered 29696 samples — only the reported value.

The fix

Clone rather than alias, which is what the training path a few lines below already does for reduced_train_loss.

Test plan

• New unit test test_validation_report_survives_megatrons_in_place_rescale reproduces Megatron's in-place rescale and asserts the reported pair does not move. It fails with assert 8.0 == 16.0 without the fix.
• tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py passes in full (8 passed).
• End-to-end on Flux 12B MXFP6, single node and two nodes, as tabulated above.

Made with Cursor

[Copilot]

🟢 Approval recommended

The fix directly addresses a confirmed aliasing bug and is covered by a focused regression test that fails pre-fix and passes post-fix.

Pull request overview

Fixes incorrect validation-loss reporting in Megatron diffusion training by preventing Megatron’s in-place rescale of the backprop tensor from mutating the reported validation numerator (which previously shared storage via detach()).

Changes:

• Clone the detached validation-loss sum before returning it in the reporting dict so Megatron’s in-place rescale cannot affect the reported metric.
• Add a unit test that reproduces Megatron’s in-place rescale behavior and asserts the reported (loss_sum, sample_count) pair remains stable.

File summaries

File	Description
primus/backends/megatron/diffusion_trainer.py	Prevents validation metric aliasing by returning a detached clone for the reported loss numerator.
tests/unit_tests/backends/megatron/diffusion/training/test_diffusion_trainer.py	Adds a regression test ensuring reported validation loss is invariant to Megatron’s in-place rescale of the backprop tensor.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟢 Approval recommended

The fix addresses the confirmed tensor-aliasing root cause and is covered by a targeted regression unit test that fails without the change.

Review details

• Files reviewed: 2/2 changed files
• Comments generated: 0 new
• Review effort level: Lite