feat: Add dataset compile helper

pyproject.toml

@@ -116,7 +116,7 @@ override-dependencies = [
 ]
 
 [tool.uv.sources]
-megatron-core = { path = "3rdparty/Megatron-LM/" }
+megatron-core = { path = "3rdparty/Megatron-LM/", editable = true }
 nvidia-modelopt = { git = "https://github.com/NVIDIA/TensorRT-Model-Optimizer.git", rev = "0a4f0a8b933121f7af080261a0a5a7717f2c5d49" }
 nvidia-resiliency-ext = { git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git", rev = "v0.4.1" } # Requires a source install to compile cupti for cuda13
 

src/megatron/bridge/training/initialize.py

@@ -14,13 +14,15 @@
 
 import datetime
 import os
+import time
 import warnings
 from typing import Callable, Optional
 
 import torch
 import torch.distributed
 import torch.nn.functional as F
 from megatron.core import parallel_state, tensor_parallel
+from megatron.core.datasets.utils import compile_helpers
 from megatron.core.fusions.fused_bias_dropout import bias_dropout_add_fused_train
 from megatron.core.fusions.fused_bias_gelu import bias_gelu
 from megatron.core.fusions.fused_bias_swiglu import bias_swiglu
@@ -115,7 +117,7 @@ def initialize_megatron(
     init_rerun_state(rerun_state_machine_config)
 
     # torch.distributed initialization
-    return torch_dist_init(
+    result = torch_dist_init(
         model_config=model_config,
         dist_config=dist_config,
         rng_config=rng_config,
@@ -128,6 +130,22 @@ def initialize_megatron(
         use_inprocess_restart=use_inprocess_restart,
     )
 
+    # Compile dataset helpers after distributed initialization
+    if torch.distributed.is_initialized():
+        if get_rank_safe() == 0:
+            start_time = time.time()
+            print("> compiling dataset index builder ...")
+            compile_helpers()
+            print(
+                ">>> done with dataset index builder. Compilation time: {:.3f} seconds".format(
+                    time.time() - start_time
+                ),
+                flush=True,
+            )
+        torch.distributed.barrier()
+
+    return result
+
 
 def torch_dist_init(
     model_config: GPTModelProvider | T5ModelProvider,