[ckpt] fix: Prevent int32 overflow for high train sample counts.

[BlueCrescent]

What does this PR do ?

Previously, when creating the TrainState for checkpointing int32 was used as type for encoding the number of consumed samples and related values. For longer trainings, torch will detect this as potential overflow and crash with RuntimeError: value cannot be converted to type int32 without overflow. We saw this crash while checkpointing step 700,000 while training with sequence length 8192 and batch size of 3072, so at 2,150,400,000 (> 2^31 - 1 = 2,147,483,648) samples.

Changelog

• Switched from int32 to int64 for all values using this type in TrainState.

Summary by CodeRabbit

Bug Fixes

• Improved integer precision in training state serialization. Counter values are now stored as 64-bit integers instead of 32-bit, preventing potential overflow issues during extended training runs. This ensures reliable checkpoint persistence and stability when handling large training step and sample count values.

[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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: e8c35d23-d746-4271-b7ee-f73d9a8e5323

📥 Commits

Reviewing files that changed from the base of the PR and between 7ea8a45 and e357135.

📒 Files selected for processing (1)

• src/megatron/bridge/training/state.py

---

📝 Walkthrough

Walkthrough

Updated the TrainState.state_dict() method in the training state module to serialize integer training counters (step, consumed_train_samples, skipped_train_samples, consumed_valid_samples) as torch.int64 tensors instead of torch.int32 tensors.

Changes

Cohort / File(s)	Summary
Training State Serialization src/megatron/bridge/training/state.py	Changed dtype of integer training counter tensors from torch.int32 to torch.int64 in the state_dict() method (4 lines modified).

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 4

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: preventing int32 overflow for high train sample counts by switching to int64 serialization.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Test Results For Major Changes	✅ Passed	PR contains minor changes (8 lines in one method) and comprehensive unit tests verifying int64 serialization with values exceeding int32 limits.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[BlueCrescent]

Hmm, with this change I was able to write out the iteration 700,000 checkpoint. But when warmstarting from it I saw cuda OOM twice and iteration 700004: Unexpected result nan (message='found NaN in local forward loss calculation') once.
Need to investigate further.