Enable GTP for heterogeneous MIMO training by yashaswikarnati · Pull Request #6249 · NVIDIA/Megatron-LM

yashaswikarnati · 2026-08-04T16:35:41Z

Summary

extend heterogeneous MIMO grids with dense GTP and expert GTP axes and thread their process groups through data distribution, bridge communication, RNG, gradient synchronization, optimizer ownership, and checkpointing
pass ProcessGroupCollection through Mamba projections, shared experts, language embeddings, and the output head
distinguish full data-distribution groups from dense/expert replica groups in distributed checkpoint metadata and support fresh MIMO hybrid checkpoint restore
parameterize the 20-layer Nemotron mock launcher for multi-node GTP/EGTP execution

Why

The heterogeneous MIMO path builds disjoint module grids instead of using MPU-global groups. On top of #6234, it still lacked the GTP/EGTP axes and process-group plumbing required by the dense, expert, bridge, optimizer, and distributed-checkpoint paths.

Depends on #6234.

Validation

24-rank TP2/GTP8/EP4/EGTP2, 128-expert, 20-layer training completed 100 iterations with finite loss and gradient norm, skipped=0, and nan=0
saved a distributed checkpoint at iteration 100; a fresh 24-rank process loaded the model and distributed Adam state and completed a finite iteration 101
focused unit tests: 15 passed; distributed topology/bridge: 2 passed; TE GTP DCP: 16 passed, 2 skipped
Ruff, isort, Python compilation, shell syntax, and diff checks passed
runtime validation used Transformer Engine 2.19.dev0 or newer, as required by [fix] Cleanup usage of pg_collection in gtp #6234

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

- Thread pg_collection from TransformerLayer through Attention/MLP into the TE and core linear layers. - Assert in finalize_model_grads when the config enables the axis but the group is missing, instead of treating it as "GTP inactive" and training on wrong gradients. - Add test_gtp_custom_pgs.py: two TransformerBlocks, identical degrees (TP=1, CP=1, GTP_remat=2 over world=4), identical weights and input, one from the MPU groups and one from a custom collection with permuted gtp_remat rank membership. Output and gradients must match; verified it fails without the plumbing. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

hasattr() is always True on ProcessGroupCollection: __getattr__ returns None for declared-but-unset fields. resolve_gtp_remat_group() gated its MPU fallback on hasattr, so the fallback was unreachable. * A collection omitting gtp_remat resolved to None -> weights silently built UNSHARDED under GTP. Check `attr in vars(...)` instead. * Add test_pg_collection_without_gtp_remat_falls_back_to_mpu. The existing test always set gtp_remat explicitly, so it could not catch this. Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

Signed-off-by: ykarnati <ykarnati@nvidia.com>

copy-pr-bot · 2026-08-04T16:35:45Z

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.

yashaswikarnati · 2026-08-16T17:01:21Z

 | `test_gtp_fp8_param_gather.py` | Native-FP8 GTP_remat (§1.3): fp8-vs-BF16 loss parity (TP1/TP2, MoE), post-save-spike guard. |
-
-The fp32-accumulation primitive itself is covered outside this suite, by `tests/unit_tests/distributed/test_reduce_scatter_with_fp32_accumulation.py`, which does not require GTP_remat.
+| `test_gtp_custom_pgs.py` | `pg_collection` plumbing: a custom `gtp_remat` group (permuted ranks, same size) must give the same fwd/bwd results as the MPU groups — catches modules reading `parallel_state` instead of the collection passed to them. |


yashaswikarnati · 2026-08-16T17:02:40Z


 set -euo pipefail

+export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}


why did we add this back ? not needed ?

yashaswikarnati · 2026-08-16T17:02:57Z

 #!/bin/bash

-# Run an eight-rank heterogeneous mock training loop with Nemotron6-MoE VLM 20L.
+# Run heterogeneous mock training with the Nemotron6-MoE VLM 20L recipe.


no need to change this comment, avoid unnecessary edits

yashaswikarnati · 2026-08-16T17:04:09Z

+EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS=${EXPERT_TENSOR_PARALLEL_NUM_WEIGHT_SHARDS:-${LLM_EXPT_TP}}
+
+if ((TENSOR_PARALLEL_NUM_WEIGHT_SHARDS % LLM_TP != 0)); then
+  echo "TENSOR_PARALLEL_NUM_WEIGHT_SHARDS must be divisible by LLM_TP" >&2


avoid these verbose checks in the sbatch script, also dont introduce new env variables if not necessary

yashaswikarnati · 2026-08-16T17:04:30Z

 TORCHRUN_ARGS=(
-  --standalone
-  --nproc-per-node 8
+  --nproc-per-node "${NPROC_PER_NODE}"


the diff should be minimal dont change unnecessary things

yashaswikarnati · 2026-08-16T17:04:48Z

-  --llm-tp 2 \
-  --llm-cp 1 \
-  --llm-pp 1 \
+  --encoder-tp "${ENCODER_TP}" \


we dont need env variables for this

yashaswikarnati · 2026-08-16T17:07:23Z

-    # so it yields llm_dp. The physical world incl. encoder ranks is restored below.
+    # Stock validation owns the derived training arguments. Give it the language module's
+    # parallel degrees and logical world so its DP/GTP accounting matches the explicit MIMO grid.
+    args.tensor_model_parallel_size = args.llm_tp


why did we newly introduce this args.tensor_model_parallel_size = args.llm_tp
args.pipeline_model_parallel_size = args.llm_pp
args.context_parallel_size = args.llm_cp
args.expert_model_parallel_size = args.llm_ep
args.expert_tensor_parallel_size = args.llm_expt_tp or 1

yashaswikarnati · 2026-08-16T17:10:23Z

        embd_group = parallel_state.get_embedding_group(check_initialized=False)
        pos_emb_group = parallel_state.get_position_embedding_group(check_initialized=False)
        dp_cp_group = parallel_state.get_data_parallel_group(with_context_parallel=True)
+        gtp_remat_group = parallel_state.get_gtp_weight_remat_group(check_initialized=False)


is these changes already on main ?

yashaswikarnati · 2026-08-16T17:13:52Z

    """Communication role for ranks in bridge communication.

-    SENDER: Leader tp-cp rank within each DP replica of source grid.
+    SENDER: Leader tp-cp rank within each data lane of source grid.


this change should already be in main ? for bridge comm? why we need these changes again here ? seems only cosmetic ? #6263

yashaswikarnati · 2026-08-16T17:15:09Z

    pg_collection.dp = parallel_state.get_data_parallel_group(
        with_context_parallel=False, partial_data_parallel=False
    )
+    # gtp_remat axis: consumers read these with getattr and silently skip the gtp_remat


this change already seems to be on main

yashaswikarnati · 2026-08-16T17:15:26Z

                if is_pp_last_stage(p2p_communicator.pp_group):
                    recv_next = False
-                (input_tensor, output_tensor_grad) = (
+                input_tensor, output_tensor_grad = (


cosmetic changes and already on main not necessary

yashaswikarnati · 2026-08-16T17:17:00Z

    tp_group,
    dp_cp_group,
    intra_dp_cp_group=None,
+    intra_expt_dp_group=None,


why do we need changes to this file ? seems not relevant ?

yashaswikarnati · 2026-08-16T17:17:21Z

        tp_comm_buffer_name: Optional[str] = None,
        tp_group: Optional[torch.distributed.ProcessGroup] = None,
        name: str | None = None,
+        pg_collection: Optional[ProcessGroupCollection] = None,


inference layers changes is out of scope

yashaswikarnati · 2026-08-16T17:18:00Z

            is_expert=False,
            tp_comm_buffer_name='proj',
            tp_group=self.pg_collection.tp,
+            pg_collection=self.pg_collection,


are these changes not on main already ?

fanshiqing and others added 5 commits August 3, 2026 20:02

Refine usage of pg_collection in gtp

4e87571

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

fix lint format

a6fffab

Signed-off-by: Shiqing Fan <shiqingf@nvidia.com>

Enable GTP for heterogeneous MIMO training

c43272e

Signed-off-by: ykarnati <ykarnati@nvidia.com>

nvidia-megatron-lm-release-bot Bot added the community-request label Aug 4, 2026

yashaswikarnati mentioned this pull request Aug 4, 2026

Integrate GTP with heterogeneous MIMO training #6268

Closed

chtruong814 removed the community-request label Aug 5, 2026

yashaswikarnati commented Aug 16, 2026

View reviewed changes

lhb8125 mentioned this pull request Aug 27, 2026

Megatron-LM Monthly Activity Report — August 2026 (through August 23, 2026) #6921

Open

yashaswikarnati closed this Aug 27, 2026


		set -euo pipefail

		export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}

Uh oh!

Conversation

yashaswikarnati commented Aug 4, 2026

Summary

Why

Validation

Uh oh!

copy-pr-bot Bot commented Aug 4, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants