[training] fix: use uint64 instead of int64 for TrainState counters

Unsigned 64-bit doubles the positive range to ~1.8e19 and better
reflects the semantics of non-negative counters (step, consumed/skipped
samples). Backward compatible via .item().

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/training/state.py

@@ -70,10 +70,10 @@ def state_dict(self) -> dict[str, torch.Tensor]:
             their corresponding tensor representations.
         """
         return {
-            "step": torch.tensor(self.step, dtype=torch.int64),
-            "consumed_train_samples": torch.tensor(self.consumed_train_samples, dtype=torch.int64),
-            "skipped_train_samples": torch.tensor(self.skipped_train_samples, dtype=torch.int64),
-            "consumed_valid_samples": torch.tensor(self.consumed_valid_samples, dtype=torch.int64),
+            "step": torch.tensor(self.step, dtype=torch.uint64),
+            "consumed_train_samples": torch.tensor(self.consumed_train_samples, dtype=torch.uint64),
+            "skipped_train_samples": torch.tensor(self.skipped_train_samples, dtype=torch.uint64),
+            "consumed_valid_samples": torch.tensor(self.consumed_valid_samples, dtype=torch.uint64),
             "floating_point_operations_so_far": torch.tensor(
                 self.floating_point_operations_so_far, dtype=torch.float64
             ),

tests/unit_tests/training/test_state.py

@@ -86,11 +86,11 @@ def test_state_dict_structure_and_types(self):
         }
         assert set(state_dict.keys()) == expected_keys
 
-        # Check tensor types (int64 to avoid overflow on long training runs)
-        assert state_dict["step"].dtype == torch.int64
-        assert state_dict["consumed_train_samples"].dtype == torch.int64
-        assert state_dict["skipped_train_samples"].dtype == torch.int64
-        assert state_dict["consumed_valid_samples"].dtype == torch.int64
+        # Check tensor types (uint64 to avoid overflow on long training runs)
+        assert state_dict["step"].dtype == torch.uint64
+        assert state_dict["consumed_train_samples"].dtype == torch.uint64
+        assert state_dict["skipped_train_samples"].dtype == torch.uint64
+        assert state_dict["consumed_valid_samples"].dtype == torch.uint64
         assert state_dict["floating_point_operations_so_far"].dtype == torch.float64
         assert state_dict["do_train"].dtype == torch.bool
         assert state_dict["do_valid"].dtype == torch.bool