fix: token_dtype_code assigns uint16 to vocab sizes needing uint64

[andrewwhitecdw]

This PR addresses the following issue in src/megatron/bridge/training/setup.py: token_dtype_code assigns uint16 to vocab sizes needing uint64.

Changes

• src/megatron/bridge/training/setup.py: token_dtype_code assigns uint16 to vocab sizes needing uint64.

Details

--- a/src/megatron/bridge/training/setup.py
+++ b/src/megatron/bridge/training/setup.py
@@ -1,6 +1,4 @@
-    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
+    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:
+            cfg.dataset.token_dtype_code = _get_token_dtype_code(vocab_size)

Tests

• tests/unit_tests/bridge/training/test_setup.py

--- a/tests/unit_tests/bridge/training/test_setup.py
+++ b/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
+    assert _get_token_dtype_code(65537) == 8

Contributor guidelines

Per this repo's CONTRIBUTING.md:

• All commits are signed off (Signed-off-by trailer, DCO).

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

[kamran-nvidia]

/ok to test f376614

[kamran-nvidia]

@andrewwhitecdw Please address the CI failures