Merge pull request #764 from srivatsankrishnan/m-bridge

Megatron Bridge in CloudAI

Author: srivatsankrishnan

conf/experimental/megatron_bridge/test/megatron_bridge_qwen_30b.toml

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name = "megatron_bridge_qwen_30b"
+description = "Megatron-Bridge run via CloudAI SlurmSystem for Qwen3 30B A3B"
+test_template_name = "MegatronBridge"
+
+extra_container_mounts = []
+
+[cmd_args]
+gpu_type = "gb200"
+container_image = "nvcr.io#nvidia/nemo:25.11.01"
+
+model_name = "qwen3"
+model_size = "30b_a3b"
+gpus_per_node = 4
+num_gpus = 8
+domain = "llm"
+task = "pretrain"
+compute_dtype = "fp8_mx"
+
+hf_token = "REPLACE_ME_WITH_HF_TOKEN"

conf/experimental/megatron_bridge/test_scenario/megatron_bridge_qwen_30b.toml

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name = "megatron_bridge_qwen_30b"
+
+[[Tests]]
+id = "megatron_bridge_qwen_30b"
+test_name = "megatron_bridge_qwen_30b"
+num_nodes = "2"

src/cloudai/registration.py

@@ -91,6 +91,11 @@ def register_all():
         JaxToolboxSlurmCommandGenStrategy,
         NemotronTestDefinition,
     )
+    from cloudai.workloads.megatron_bridge import (
+        MegatronBridgeReportGenerationStrategy,
+        MegatronBridgeSlurmCommandGenStrategy,
+        MegatronBridgeTestDefinition,
+    )
     from cloudai.workloads.megatron_run import (
         CheckpointTimingReportGenerationStrategy,
         MegatronRunSlurmCommandGenStrategy,
@@ -185,6 +190,9 @@ def register_all():
     Registry().add_command_gen_strategy(SlurmSystem, MegatronRunTestDefinition, MegatronRunSlurmCommandGenStrategy)
     Registry().add_command_gen_strategy(SlurmSystem, NCCLTestDefinition, NcclTestSlurmCommandGenStrategy)
     Registry().add_command_gen_strategy(SlurmSystem, DDLBTestDefinition, DDLBTestSlurmCommandGenStrategy)
+    Registry().add_command_gen_strategy(
+        SlurmSystem, MegatronBridgeTestDefinition, MegatronBridgeSlurmCommandGenStrategy
+    )
 
     Registry().add_command_gen_strategy(SlurmSystem, NeMoLauncherTestDefinition, NeMoLauncherSlurmCommandGenStrategy)
     Registry().add_command_gen_strategy(SlurmSystem, NeMoRunTestDefinition, NeMoRunSlurmCommandGenStrategy)
@@ -234,6 +242,7 @@ def register_all():
     Registry().add_test_definition("JaxToolboxNemotron", NemotronTestDefinition)
     Registry().add_test_definition("SlurmContainer", SlurmContainerTestDefinition)
     Registry().add_test_definition("MegatronRun", MegatronRunTestDefinition)
+    Registry().add_test_definition("MegatronBridge", MegatronBridgeTestDefinition)
     Registry().add_test_definition("TritonInference", TritonInferenceTestDefinition)
     Registry().add_test_definition("NIXLBench", NIXLBenchTestDefinition)
     Registry().add_test_definition("AIDynamo", AIDynamoTestDefinition)
@@ -250,6 +259,7 @@ def register_all():
     Registry().add_report(GPTTestDefinition, JaxToolboxReportGenerationStrategy)
     Registry().add_report(GrokTestDefinition, JaxToolboxReportGenerationStrategy)
     Registry().add_report(MegatronRunTestDefinition, CheckpointTimingReportGenerationStrategy)
+    Registry().add_report(MegatronBridgeTestDefinition, MegatronBridgeReportGenerationStrategy)
     Registry().add_report(NCCLTestDefinition, NcclTestPerformanceReportGenerationStrategy)
     Registry().add_report(NeMoLauncherTestDefinition, NeMoLauncherReportGenerationStrategy)
     Registry().add_report(NeMoRunTestDefinition, NeMoRunReportGenerationStrategy)

src/cloudai/workloads/megatron_bridge/__init__.py

@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
+# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .megatron_bridge import MegatronBridgeCmdArgs, MegatronBridgeTestDefinition
+from .report_generation_strategy import MegatronBridgeReportGenerationStrategy
+from .slurm_command_gen_strategy import MegatronBridgeSlurmCommandGenStrategy
+
+__all__ = [
+    "MegatronBridgeCmdArgs",
+    "MegatronBridgeReportGenerationStrategy",
+    "MegatronBridgeSlurmCommandGenStrategy",
+    "MegatronBridgeTestDefinition",
+]