fix: cache a copy of the GPT dataset loss mask

[ZhiyuLi-Nvidia]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

The bug

GPTDataset.__getitem__ caches the loss mask, then masks in place:

if not cacheable or not cached:                   # FILL — first sample only
    loss_mask = _get_ltor_masks_and_position_ids(...)
    if cacheable:
        self.cached_loss_mask = loss_mask         # 303: stores the tensor
else:                                             # HIT — every later sample
    loss_mask = self.cached_loss_mask.clone()     # 308: stores a copy

loss_mask[labels == self._pad_token_id] = 0.0     # 312: in-place

Line 303 stores the tensor itself rather than a copy, so line 312 writes through that reference
into the cache. Only the first sample served can do this — FILL runs once per dataset
(masks_and_position_ids_are_cached is set True there and never reset) and every later sample
masks a clone.

·  loss counted    ×  masked

Case B — the masked sample is served anywhere but first       NO EFFECT
──────────────────────────────────────────────────────────────────────
              self.cached_loss_mask       per-call tensor        returned

 1st  s0      ┌───────────────┐
      FILL    │ A ··········· │ ← the only write that ever happens     A
              └───────────────┘

 2nd  s_k     ┌───────────────┐
      HIT     │ A ··········· │ ─clone→ B ···········                  B
              └───────────────┘  read    312 masks B: ······×××××
                A not touched            the zeros land in B only

 3rd  s_j     ┌───────────────┐
      HIT     │ A ··········· │ ─clone→ C ···········                  C
              └───────────────┘         clones A — NOT B


Case A — the masked sample is served FIRST                        BUG
──────────────────────────────────────────────────────────────────────
 1st  s_k     ┌───────────────┐
      FILL    │ A ··········· │ ←── 303 stored the tensor itself, so the
              └───────────────┘     slot and the tensor masked at 312
                     ↑              are THE SAME OBJECT
                     └── 312 writes here: A ······×××××            A

 2nd  s0      ┌───────────────┐
      HIT     │ A ······××××× │ ─clone→ B ······×××××                  B  WRONG
              └───────────────┘         a dirty base

 3rd  s_j     ┌───────────────┐
      HIT     │ A ······××××× │ ─clone→ C ······×××××                  C  WRONG
              └───────────────┘                ^^^^^
                                        s_k's zeros, on every sample

Step 3 does not reuse "the cache from step 2" — step 2 never wrote one. Line 308
always clones the slot, which has held A since step 1.

Impact

The cached mask is masked in place, so the first sample's padding poisons every batch after
it. Line 303 stores the tensor itself, then line 312 writes into that same object — so whatever
the first sample masks (real padding, or the pad id among its tokens) is written straight into the
cache and silently inherited by every following batch. Later samples are harmless by comparison:
they clone the cache first, so their own padding lands on a throwaway copy.

Two ways a sample gets masked positions:

• Real padding, when a sample is shorter than sequence_length. Training never hits this
  (drop_last_partial_sequence is hardcoded True at line 248), but validation can, with
  drop_last_partial_validation_sequence=False.
• The pad id as ordinary data. _pad_token_id is tokenizer.pad verbatim
  (megatron_dataset.py:75), so any occurrence in the token stream is masked, on any sample.
  MockGPTLowLevelDataset generates (arange(length - 1) + 1) % vocab_size — a ramp over the whole
  vocabulary — so under --mock-data this fires on essentially every sample.

It breaks bitwise checkpoint resume, silently. Which sample runs FILL is a property of the
process: a fresh run fills from its first training sample, a resumed run from whatever sample its
step maps to. The two then serve different masks for the same sample, so gradients diverge from the
first step after the resume even with checkpoint state, tokens, labels and position_ids all
bitwise identical. Nothing errors — the loss stays plausible because it is normalised by the mask sum.

MockGPTDataset and GPTFIMDataset inherit this __getitem__. SFTDataset defines its own and is
unaffected.

The fix

Cache a clone at 303, mirroring the clone already made at 308. One line — the first
sample's in-place write then lands on a throwaway and the slot stays clean.

Test

test_mask_cache_does_not_leak_padding builds two dataset instances, serves the
masked sample first in one, and asserts both return the same loss_mask for
index 0. It first asserts that sample really is masked, so it fails loudly rather
than passing vacuously if the fixture stops producing one. It reproduces the defect
via trigger 1, which is the one constructible from an in-tree tokenizer; both
triggers exercise the same aliasing.

Fails before the change (AssertionError: padding from the first served sample leaked into a later sample's loss_mask), passes after.

cached_attention_mask and cached_position_ids are stored by reference too.
Neither is mutated today so neither is buggy; left unchanged to keep the diff
minimal.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.