chore: bump Megatron-Bridge to latest main (7110a96) (#2223)

yuki-97 · web-flow · commit cb34225ca312 · 2026-04-10T10:57:08.000+08:00
Signed-off-by: Yuki Huang &lt;yukih@nvidia.com&gt;
diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge
@@ -1 +1 @@
-Subproject commit a2bb70b91b827bd6b085a77442c7cf60cfdb59fe
+Subproject commit 7110a964272a5c74dcb6b680b691087e190c220c
diff --git a/3rdparty/Megatron-Bridge-workspace/setup.py b/3rdparty/Megatron-Bridge-workspace/setup.py
@@ -25,6 +25,7 @@
 bridge_src_dir = "Megatron-Bridge/src/megatron/bridge"
 bridge_package_name = "megatron.bridge"
 
+# Default dependencies from pyproject.toml
 CACHED_DEPENDENCIES = [
     "transformers>=5.0.0,<=5.3.0",
     "peft>=0.18.1",
@@ -50,7 +51,7 @@
     # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13 base container
     "transformer-engine[pytorch,core_cu12]",
     "mamba-ssm",
-    "nvidia-resiliency-ext~=0.5.0",
+    "nvidia-resiliency-ext",
     "causal-conv1d",
     "flash-linear-attention",
     "timm",
diff --git a/3rdparty/Megatron-LM-workspace/Megatron-LM b/3rdparty/Megatron-LM-workspace/Megatron-LM
@@ -1 +1 @@
-Subproject commit 9e2810417315a7ee93b41d4e234454abd3c16af5
+Subproject commit 17a67b9a97fb11a75933fd7f76ad76e1ac98a53d
diff --git a/3rdparty/Megatron-LM-workspace/setup.py b/3rdparty/Megatron-LM-workspace/setup.py
@@ -50,7 +50,8 @@
     "nvidia-modelopt[torch]; sys_platform != 'darwin'",
     # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13 base container
     "transformer-engine[pytorch,core_cu12]",
-    "nvidia-resiliency-ext @ git+https://github.com/NVIDIA/nvidia-resiliency-ext.git@v0.5.0",
+    # VCS dependency - must match pyproject.toml [tool.uv.sources]
+    "nvidia-resiliency-ext @ git+https://github.com/NVIDIA/nvidia-resiliency-ext.git@63154570cea17f8805a7fd15cc3b8cc2919ba575",
     "tqdm",
     "einops~=0.8",
     "tensorstore~=0.1,!=0.1.46,!=0.1.72",
@@ -65,10 +66,10 @@
     "flashinfer-python~=0.5.0",
     "wget",
     "onnxscript",
-    # VCS dependency - must match pyproject.toml [tool.uv.sources]
-    "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1.0",
-    "datasets",
     "fastapi~=0.50",
+    "datasets",
+    # VCS dependency - must match pyproject.toml [tool.uv.sources]
+    "emerging_optimizers @ git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.2.0",
     "hypercorn",
     "quart",
     "openai[aiohttp]",
diff --git a/examples/configs/distillation_math.yaml b/examples/configs/distillation_math.yaml
@@ -113,6 +113,7 @@ policy: &POLICY_BASE
         moe_enable_deepep: false
         moe_token_dispatcher_type: "alltoall"
         moe_shared_expert_overlap: false
+        gradient_accumulation_fusion: false
         
         optimizer:
             optimizer: "adam"
diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml
@@ -147,6 +147,7 @@ policy:
     moe_enable_deepep: false
     moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
+    gradient_accumulation_fusion: false
 
     peft:
       enabled: false
diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml
@@ -155,6 +155,7 @@ policy:
     moe_enable_deepep: false
     moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
+    gradient_accumulation_fusion: false
 
     peft:
       enabled: false
diff --git a/examples/configs/recipes/llm/grpo-dapomath17k-dsv3-megatron.yaml b/examples/configs/recipes/llm/grpo-dapomath17k-dsv3-megatron.yaml
@@ -30,7 +30,6 @@ policy:
     sequence_parallel: true
     moe_permute_fusion: true
     apply_rope_fusion: false
-    gradient_accumulation_fusion: false
     moe_enable_deepep: true
     moe_token_dispatcher_type: flex
     optimizer:
diff --git a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-fp8-e2e.yaml b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron-fp8-e2e.yaml
@@ -26,7 +26,6 @@ policy:
     num_layers_in_first_pipeline_stage: 7
     num_layers_in_last_pipeline_stage: 6
     apply_rope_fusion: false
-    gradient_accumulation_fusion: false
     fp8_cfg:
       enabled: true
     optimizer:
diff --git a/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron.yaml b/examples/configs/recipes/llm/grpo-moonlight-16ba3b-4n8g-megatron.yaml
@@ -27,7 +27,6 @@ policy:
     num_layers_in_first_pipeline_stage: 7
     num_layers_in_last_pipeline_stage: 6
     apply_rope_fusion: false
-    gradient_accumulation_fusion: false
     optimizer:
       lr: 1.0e-06
     scheduler:
diff --git a/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml b/examples/configs/recipes/llm/performance/dapo-deepseek-v3-64n8g.yaml
@@ -53,7 +53,6 @@ policy:
     sequence_parallel: true
     moe_permute_fusion: true
     apply_rope_fusion: false
-    gradient_accumulation_fusion: false
     # MTP — disabled
     mtp_num_layers: 0
     optimizer:
diff --git a/examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml b/examples/configs/recipes/llm/performance/grpo-deepseek-v3-32n8g.yaml
@@ -27,7 +27,6 @@ policy:
     num_layers_in_first_pipeline_stage: 3
     num_layers_in_last_pipeline_stage: 2
     apply_rope_fusion: false
-    gradient_accumulation_fusion: false
     moe_permute_fusion: true
     defer_fp32_logits: true
     # MTP — disabled
diff --git a/examples/configs/rm.yaml b/examples/configs/rm.yaml
@@ -87,6 +87,7 @@ policy:
     num_layers_in_first_pipeline_stage: null
     num_layers_in_last_pipeline_stage: null
     sequence_parallel: false
+    gradient_accumulation_fusion: false
 
     optimizer:
       optimizer: "adam"
diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml
@@ -128,6 +128,7 @@ policy:
     moe_enable_deepep: false
     moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
+    gradient_accumulation_fusion: false
 
     peft:
       enabled: false
diff --git a/examples/nemo_gym/grpo_nanov3.yaml b/examples/nemo_gym/grpo_nanov3.yaml
@@ -107,6 +107,7 @@ policy:
     track_moe_metrics: True
     moe_per_layer_logging: True
     moe_shared_expert_overlap: false
+    gradient_accumulation_fusion: false
 
     optimizer:
       optimizer: "adam"
diff --git a/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml b/examples/nemo_gym/grpo_workplace_assistant_nemotron_nano_v2_9b.yaml
@@ -121,6 +121,7 @@ policy:
     moe_enable_deepep: false
     moe_token_dispatcher_type: "alltoall"
     moe_shared_expert_overlap: false
+    gradient_accumulation_fusion: false
 
     optimizer:
       optimizer: "adam"
diff --git a/nemo_rl/models/megatron/community_import.py b/nemo_rl/models/megatron/community_import.py
@@ -75,12 +75,9 @@ def import_model_from_hf_name(
         ]
         model_provider.pipeline_dtype = megatron_config["pipeline_dtype"]
         model_provider.sequence_parallel = megatron_config["sequence_parallel"]
-        if (
-            gradient_accumulation_fusion := megatron_config.get(
-                "gradient_accumulation_fusion"
-            )
-        ) is not None:
-            model_provider.gradient_accumulation_fusion = gradient_accumulation_fusion
+        model_provider.gradient_accumulation_fusion = megatron_config[
+            "gradient_accumulation_fusion"
+        ]
     model_provider.finalize()
 
     from megatron.core import parallel_state
diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py
@@ -147,13 +147,14 @@ def destroy_parallel_state():
         old_call_idx = getattr(
             megatron_async_utils._async_calls_queue, "call_idx", None
         )
-        num_unfinalized = (
-            megatron_async_utils._async_calls_queue.get_num_unfinalized_calls()
-        )
-        if num_unfinalized > 0:
-            print(
-                f"[WARNING] Resetting Megatron async calls queue with {num_unfinalized} unfinalized calls"
+        if megatron_async_utils._async_calls_queue is not None:
+            num_unfinalized = (
+                megatron_async_utils._async_calls_queue.get_num_unfinalized_calls()
             )
+            if num_unfinalized > 0:
+                print(
+                    f"[WARNING] Resetting Megatron async calls queue with {num_unfinalized} unfinalized calls"
+                )
         try:
             megatron_async_utils._async_calls_queue.close()
         except:
diff --git a/pyproject.toml b/pyproject.toml
@@ -50,11 +50,10 @@ dependencies = [
   "swanlab",
   "pyzmq",
   "decord2",
-
-  "nccl4py",                      # for non-colocated refit
-  "cuda-bindings",                # for non-colocated refit
-  "pybase64",                     # for sglang refit
-  "nvidia-cudnn-cu12==9.19.0.56", # for transformer-engine no build isolation
+  "nccl4py",                                                                                                          # for non-colocated refit
+  "cuda-bindings",                                                                                                    # for non-colocated refit
+  "pybase64",                                                                                                         # for sglang refit
+  "nvidia-cudnn-cu12==9.19.0.56",                                                                                     # for transformer-engine no build isolation
   # nvidia-resiliency-ext removed: no Python 3.13 wheels available (v0.5.0 only has cp310-cp312)
 ]
 
@@ -99,14 +98,14 @@ mcore = [
   # This dependency also needs to be compatible with the spec in Megatron-Bridge/pyproject.toml.
   # It is specified here since we don't directly use Megatron-Bridge/pyproject.toml, but a proxy setup.py+pyproject.toml combo
   # outside to allow "optionally" installing the megatron path. It's simpler to deal with transformer-engine here in the NeMo RL pyproject.toml
-  "transformer-engine[pytorch]==2.12.0",
+  "transformer-engine[pytorch,core_cu12] @ git+https://github.com/NVIDIA/TransformerEngine.git@71bbefbf153418f943640df0f7373625dc93fa46",
   "megatron-core",
   "megatron-bridge",
   # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular)
   # https://github.com/NVIDIA/TransformerEngine/blob/v2.3/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L108
   # https://github.com/facebookresearch/xformers/blob/8354497deb2c04c67fbb2e2ad911e86530da0e90/xformers/ops/fmha/flash.py#L76
   "flash-attn==2.8.1",
-  "emerging-optimizers==0.1.0",
+  "emerging-optimizers==0.2.0",
   "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480",
 ]
 nemo_gym = ["nemo_gym"]
@@ -185,6 +184,7 @@ nv-grouped-gemm = { git = "https://github.com/fanshiqing/grouped_gemm", tag = "v
 # From JustinTong0323/sglang branch update-transformers-v5 (sgl-project/sglang#17784)
 sglang = { git = "https://github.com/JustinTong0323/sglang.git", rev = "70aa688742dd2b75bf9e8e980249303f39295b0d", subdirectory = "python" }
 sgl-kernel = { git = "https://github.com/JustinTong0323/sglang.git", rev = "70aa688742dd2b75bf9e8e980249303f39295b0d", subdirectory = "sgl-kernel" }
+emerging-optimizers = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git", rev = "v0.2.0" }
 
 [tool.uv.workspace]
 members = [
@@ -236,7 +236,7 @@ link-mode = "copy"
 # The timm override is needed because current automodel pins to 1.0.16. This can be removed once we move ToT automodel
 # The nvidia-modelopt override is needed because mcore is still on 0.33
 override-dependencies = [
-  "transformer-engine[pytorch]==2.12.0",
+  "transformer-engine[pytorch,core_cu12] @ git+https://github.com/NVIDIA/TransformerEngine.git@71bbefbf153418f943640df0f7373625dc93fa46",
   "nvidia-cudnn-cu12==9.19.0.56",
   "opencv-python-headless>=4.11.0",
   "timm<=1.0.22",
@@ -367,13 +367,22 @@ name = "deep_gemm"
 version = "v2.0.0+7b6b556"
 requires-dist = ["torch", "packaging", "ninja"]
 
+[[tool.uv.dependency-metadata]]
+name = "transformer-engine"
+version = "2.14.0+71bbefb"
+requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"]
+
+[[tool.uv.dependency-metadata]]
+name = "transformer-engine-torch"
+version = "2.14.0+71bbefb"
+requires-dist = ["torch", "transformer-engine"]
+
 [[tool.uv.dependency-metadata]]
 name = "nv-grouped-gemm"
 # This version has to match the version in the commit/rev/tag used
 version = "v1.1.4.post7"
 requires-dist = ["setuptools", "wheel", "torch", "numpy"]
 
-
 [[tool.uv.dependency-metadata]]
 name = "sgl-kernel"
 # This version has to match the version in the commit/rev/tag used
@@ -483,7 +492,7 @@ requires-dist = [
   # TODO(https://github.com/NVIDIA-NeMo/RL/issues/2111): upgrade to core_cu13 when we move to CUDA 13
   "transformer-engine[pytorch,core_cu12]",
   "mamba-ssm",
-  "nvidia-resiliency-ext~=0.5.0",
+  "nvidia-resiliency-ext",
   "causal-conv1d",
   "flash-linear-attention",
   "timm",
@@ -493,6 +502,12 @@ requires-dist = [
   "torch>=2.6.0",
 ]
 
+# remove numpy upper-bound to resolve conflict of nvidia-resiliency-ext==0.6.0 and onnx>=1.21.0rc4
+[[tool.uv.dependency-metadata]]
+name = "logsage"
+version = "0.1.5"
+requires-dist = []
+
 [tool.black]
 line-length = 120
 include = '\.pyi?$'
diff --git a/research/template_project/configs/grpo_math_1B.yaml b/research/template_project/configs/grpo_math_1B.yaml
@@ -96,6 +96,7 @@ policy:
     #gives ~20% training perf speedup with sequence packing
     apply_rope_fusion: True
     defer_fp32_logits: null
+    gradient_accumulation_fusion: false
 
     optimizer:
       optimizer: "adam"
diff --git a/tests/functional/test_converter_roundtrip.py b/tests/functional/test_converter_roundtrip.py
@@ -265,7 +265,22 @@ def create_megatron_checkpoint(model_name: str, temp_dir: str) -> str:
 
     with temporary_distributed_context(backend="gloo"):
         megatron_checkpoint_path = os.path.join(temp_dir, "megatron_checkpoint")
-        import_model_from_hf_name(model_name, megatron_checkpoint_path)
+        megatron_config = {
+            "tensor_model_parallel_size": 1,
+            "pipeline_model_parallel_size": 1,
+            "context_parallel_size": 1,
+            "expert_tensor_parallel_size": 1,
+            "expert_model_parallel_size": 1,
+            "num_layers_in_first_pipeline_stage": None,
+            "num_layers_in_last_pipeline_stage": None,
+            "pipeline_dtype": "float32",
+            "sequence_parallel": False,
+            "gradient_accumulation_fusion": False,
+        }
+        # need to set `gradient_accumulation_fusion=False` to avoid RuntimeError
+        import_model_from_hf_name(
+            model_name, megatron_checkpoint_path, megatron_config=megatron_config
+        )
 
     if not check_file_exists(megatron_checkpoint_path):
         raise FileNotFoundError(
diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py
@@ -245,6 +245,7 @@ def get_basic_megatron_test_config(
             "apply_rope_fusion": True,
             "bias_activation_fusion": True,
             "moe_per_layer_logging": False,
+            "gradient_accumulation_fusion": False,
             "train_iters": 100,  # Required for Megatron training
             "optimizer": {
                 "optimizer": "adam",
diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py
@@ -144,6 +144,7 @@ def create_megatron_test_config(
             "defer_fp32_logits": defer_fp32_logits,
             "use_linear_ce_fusion_loss": False,
             "linear_ce_fusion_chunk_size": 256,
+            "gradient_accumulation_fusion": False,
             "train_iters": 100,  # Required for Megatron training
             "optimizer": {
                 "optimizer": "adam",
diff --git a/tools/refit_verifier.py b/tools/refit_verifier.py
@@ -221,6 +221,7 @@ def setup_configs(args, tokenizer):
             "moe_per_layer_logging": False,
             "freeze_moe_router": False,
             "apply_rope_fusion": False,
+            "gradient_accumulation_fusion": False,
             "optimizer": {
                 "optimizer": "adam",
                 "lr": 5.0e-6,
diff --git a/uv.lock b/uv.lock
