Nemotron masked-diffusion training currently divides each microbatch's
weighted loss sum by its own random mask count, then lets Trainer
average those local means. With identical realized masks, batch size 1 /
gradient accumulation 2 therefore computes `mean(S_g / N_g)` instead of
`sum(S_g) / sum(N_g)`. This changes token weights even when both answers
have the same length.
[NVIDIA's original
forward](16c67f0560/modeling_nemotron_labs_diffusion.py)
returns `(sum_loss, num_mask_tokens)` for external normalization by the
global token count. This PR implements that normalization across the
accumulation window.
### Changes
- Separate mask/noise preparation from the forward pass, and retain each
sampled mask for reuse. Count actual masked tokens across the window,
including the existing EOS-padding rules. Counting non-ignored labels
would count eligible targets rather than the sampled subset.
- Divide each loss sum by the shared count and tell Trainer that
accumulation normalization is already included. Forward and backward
still run one microbatch at a time; only input/noise tensors are
retained across the window.
- When `average_tokens_across_devices=True`, gather counts across ranks
and compensate for DDP's gradient averaging. Respect rank-local
normalization when the option is disabled.
- Enable this normalization for the Nemotron profile. Evaluation and
direct calls without a window count retain their single-batch
denominator; other profiles retain their normalization. Zero-mask
windows remain finite.
The noise distribution and inverse-probability weighting are preserved.
Sampling happens earlier in the training loop, so the comparison fixes
the actual noise draws rather than assuming that changing batch shape
with the same RNG seed produces identical masks.
### Validation
The independent oracle executes NVIDIA's unmodified model class and
original forward/loss using a tiny random configuration: two layers,
hidden size 32, vocabulary 97. No model outputs or gradients are
stubbed. The full 8B checkpoint was not loaded, and these checks make no
claim about downstream model quality or throughput.
Example on CUDA: two equal-length answers, `micro_batch_size=1`,
`gradient_accumulation_steps=2`, and realized mask counts **7 and 10**.
The reference processes both examples together with those same masks.
| Check | Before | After | NVIDIA reference |
| --- | ---: | ---: | ---: |
| Loss | 9.78913879 | 9.90060806 | 9.90060997 |
| Maximum absolute parameter-gradient error | 0.18355381 | 2.38e-7 | — |
| Maximum parameter error after one SGD update | 1.84e-3 | 3.73e-9 | — |
- **8 CPU and 8 CUDA cases passed** after the fix: three noise seeds, a
partial accumulation window, one/all empty microbatches, a frozen output
head, and a single-microbatch control. Losses, parameter gradients, and
SGD updates are compared to the original NVIDIA implementation.
- Complete CPU and CUDA `Trainer.train()` runs passed with window sizes
**2, 2, 1**, including evaluation parity. After three SGD updates,
maximum parameter error was **7.45e-9** on both runs.
- **4 real two-process CPU/Gloo DDP cases passed**: unequal per-rank
mask counts (18 and 22), one empty rank, all ranks empty, and explicitly
requested rank-local normalization.
- Existing `tests/diffusion/test_diffusion_masked.py`: **40 passed, 4
skipped**. Pre-commit hooks and `git diff --check` passed.
CUDA validation uses the real Unsloth import and Trainer, including the
normal Trainer patch installed by FastModel, on PyTorch 2.12.1+cu130 /
Transformers 5.5.0 / TRL 0.24.0. `UNSLOTH_FORCE_FLOAT32=1` keeps the
full training loop and reference in FP32. CPU validation uses PyTorch
2.13.0 / Transformers 5.14.1 / TRL 1.9.2, isolating Unsloth package
startup while retaining the real production diffusion methods and
TRL/Accelerate execution. Distributed validation covers CPU/Gloo, not
multi-GPU, FSDP, or DeepSpeed.
This PR contains only the two production files. The additional
validation scripts were run separately and are not added to the
repository. Tested commit: `dbd56922`; base: `b0e74e74`. The affected
upstream files were rechecked at `f0466637` and are unchanged from that
base.
---------
Co-authored-by: Daniel Han <23090290+danielhanchen@users.noreply.github.com>
7ec86bbc71