fix: token_dtype_code assigns uint16 to vocab sizes needing uint64

src/megatron/bridge/training/setup.py

@@ -335,9 +335,7 @@ def setup(
     if hasattr(cfg.dataset, "token_dtype_code") and cfg.dataset.token_dtype_code is None:
         vocab_size = getattr(tokenizer, "vocab_size", None)
         if vocab_size is not None:
-            import numpy
-
-            cfg.dataset.token_dtype_code = 4 if vocab_size > numpy.iinfo(numpy.uint16).max + 1 else 8
+            cfg.dataset.token_dtype_code = _get_token_dtype_code(vocab_size)
 
     if cfg.train.num_epochs is not None:
         if should_fire(callback_manager, "on_data_init_start"):
@@ -780,6 +778,16 @@ def _validate_and_set_vocab_size(
         return model_vocab_size, False
 
 
+def _get_token_dtype_code(vocab_size: int) -> int:
+    """Return the numpy dtype code for token IDs that can hold `vocab_size` tokens.
+
+    uint16 (code 4) is used when all token ids (0..vocab_size-1) fit in uint16,
+    otherwise uint64 (code 8) is used.
+    """
+    import numpy as np
+
+    return 4 if vocab_size <= np.iinfo(np.uint16).max + 1 else 8
+
 def maybe_log_and_save_config(cfg: ConfigContainer) -> None:
     """Save configuration to disk and log non-default values on rank 0.
 

tests/unit_tests/bridge/training/test_setup.py

@@ -0,0 +1,22 @@
+import numpy as np
+import pytest
+
+from megatron.bridge.training.setup import _get_token_dtype_code
+
+
+def test_get_token_dtype_code_uint16():
+    """Vocab sizes up to 65536 fit in uint16 (numpy code 4)."""
+    assert _get_token_dtype_code(1) == np.dtype(np.uint16).num
+    assert _get_token_dtype_code(65535) == np.dtype(np.uint16).num
+    assert _get_token_dtype_code(65536) == np.dtype(np.uint16).num
+
+
+def test_get_token_dtype_code_uint64():
+    """Vocab sizes larger than 65536 require uint64 (numpy code 8)."""
+    assert _get_token_dtype_code(65537) == np.dtype(np.uint64).num
+    assert _get_token_dtype_code(100000) == np.dtype(np.uint64).num
+
+
+def test_get_token_dtype_code_boundary():
+    """The boundary between uint16 and uint64 is exactly 65536."""
+    assert _get_token_dtype_code(65536) == 4