Do refit or checkpoint saving need a GTP axis on the driver? No: both learn GTP from Megatron on each GPU

Explainer for NVIDIA-NeMo/RL #4512. NeMo-RL permalinks pin the PR head e3248a60; Megatron-Bridge at ec835530, Megatron-LM at b4d72b79 (the PR's pins). Read from the code; nothing here was run on a GPU.

Before
The driver lays GPUs out in 4 named axes (PP, DP, CP, TP). Refit and checkpoint save already go to every GPU.
This PR
Adds a fifth driver axis for GTP. It does not touch the refit or checkpoint code.
Status
Driver axis not needed. Neither refit nor checkpointing picks GPUs by a GTP axis. Two gaps are open, and neither is fixed in this PR: GTP refit exports the wrong weights, and a GTP checkpoint save will likely fail on one RNG entry. Both fixes belong in Megatron-Bridge, on the GPU side.

Three pages on this PR, one question each. This one: whether refit or checkpoint saving need the PR's GTP axis.

Background. GTP works like FSDP inside each TP slice. Each GPU keeps 1/G of its TP slice of a weight. Before a layer runs, the G GPUs in a GTP group all-gather the weight (gtp.py:1565), and after backward they reduce-scatter the gradients (gtp.py:2271). Each GPU in the group trains on its own batch. Refit is the copy of trained weights into vLLM. The driver decides which GPUs run a call. Megatron on each GPU knows its own TP and GTP rank. Dotted words have a short explanation: hover or tap them.

The example below uses TP = 2, G = 3 and one model copy, so 6 GPUs. Megatron numbers them TP fastest, then GTP: GPU r = t + 2·g (parallel_state.py:602). One weight W has 24 rows. TP splits the rows in half, and GTP splits each half in 3, so each GPU holds 4 rows (gtp.py:619). The driver sends all three calls that matter here to every GPU, without naming an axis.

Driver: picks which GPUs run a call. It names no axis for these three calls: lm_policy.py:1273 refit (IPC) lm_policy.py:1361 refit (broadcast) lm_policy.py:1473 save_checkpoint → all 6 GPUs run_all_workers_single_data GTP peer g = 0GTP peer g = 1GTP peer g = 2 TP slice t = 0(rows 0–11) TP slice t = 1(rows 12–23) r0 (t0, g0)holds rows 0–3 r2 (t0, g1)holds rows 4–7 r4 (t0, g2)holds rows 8–11 r1 (t1, g0)holds rows 12–15 r3 (t1, g1)holds rows 16–19 r5 (t1, g2)holds rows 20–23 Each GPU's Megatron knows its own t and g. GTP groups: {r0,r2,r4} and {r1,r3,r5}. TP groups: {r0,r1}, {r2,r3}, {r4,r5}. Refit: every GPU runs Bridge's export today, r0: TP gather {r0,r1} → rows 0–3 + 12–15 = [8,k], not [24,k] ✗ fix, r0: GTP gather {r0,r2,r4} → rows 0–11 then TP gather {r0,r1} → rows 0–23 ✓ Order matters: TP first, then GTP, gives rows 0–3, 12–15, 4–7, 16–19, … (all rows, wrong order). param_mapping.py:1331 TP-only gather model_bridge.py:1779 where the gather goes Checkpoint: every GPU writes its own pieces weights: piece = t·3 + g, 6 pieces, each once ✓ mcore utils.py:1130 offset rule loader state: 3 lanes, from Megatron's DP rank ✓ sft_worker.py:82 lane = d·3 + g rng_state at (pp0, tp0, dp0): r0, r2, r4 all write it as the main copy ✗ bridge checkpointing.py:579 uses dp_cp validation.py:458 rejects duplicate keys correct today broken today; the fix is in Megatron-Bridge on each GPU, not in the driver

Refit: vLLM needs no GTP support; Bridge must gather GTP first (fix proposed, not in the PR)

Every GPU runs Bridge's export_hf_weights (megatron_policy_worker.py:3372). Bridge hands the local weight to the mapping as-is (model_bridge.py:2459), and the mapping gathers over TP only. Bridge's only GTP code is training setup (training/gtp.py). vLLM never sees trainer slices: it receives full Hugging Face tensors and splits them for its own TP (vllm_backend.py:1615), so vLLM needs no GTP support. Once each GPU gathers its GTP group, every GTP peer holds the full TP slice and the rest of refit works unchanged.

Where the fix lands: Megatron-Bridge export, plus the same gather in NeMo-RL's NCCL-reshard refit path (megatron_policy_worker.py:3463, which ships local slices the same way). No driver change: refit already runs on every GPU, and it never picks senders by axis. The broadcast path sends from GPU 0, which is hard-coded (:3735).

Nothing below is implemented. It is a sketch and has not been run. Native-FP8 GTP weights also need dequantize_gtp_native_fp8 first.

@@ megatron/bridge/models/conversion/model_bridge.py L1779-L1787 · stream_weights_megatron_to_hf task loop
         for task in self._with_progress_tracking(megatron_to_hf_tasks, "Converting to HuggingFace", show_progress):
             if isinstance(task.param_weight, DTensor):
                 ...
                 megatron_weights = uneven_dtensor_to_full_tensor(task.param_weight)
+            elif is_gtp_param(task.param_weight):
+                # GTP keeps 1/G of this GPU's TP slice; rebuild the TP slice
+                # before the mapping does its TP gather. All GTP peers must call this.
+                w = task.param_weight
+                full = w.new_empty((w.shape[0] * w.group.size(), *w.shape[1:]))
+                torch.distributed.all_gather_into_tensor(full, w.get_padded_shard().contiguous(), group=w.group)
+                megatron_weights = w._strip_padding(full)
             else:
                 megatron_weights = task.param_weight

is_gtp_param, .group, get_padded_shard and _strip_padding are Megatron-LM's own (gtp.py:863, :1255, :1375). There is no public "unshard" call; the forward pass gathers into a reused buffer, so export should not borrow it. Refit is not reachable in this PR's scope: SFT v2 never refits (sft_v2.py:579). Until Bridge has the fix, the review asks for a setup-time error when GTP is combined with refit.

🔴 review comment C3 — setup.py:2612

Checkpoints: Megatron already writes GTP pieces; one RNG entry collides (fix proposed, not in the PR)

The driver sends save_checkpoint to every GPU (figure), and each GPU loads at start-up (sft_v2.py:579). Megatron describes each GTP weight piece by its place in the full weight, t·G + g of TP·G pieces, and marks model copies as replicas (utils.py:1130, gtp.py:2854). The optimizer reuses the same replica marks (distrib_optimizer.py:1804). The only state saved per data lane is the SFT loader. Its lane number comes from Megatron on the GPU (sft_worker.py:82). The owners are picked by skipping only TP, CP and PP copies (sft_v2.py:385), which gives r0, r2 and r4 under both the 4-axis and 5-axis layouts.

The gap is the RNG state. Bridge keys it by the data-parallel group without GTP (process_groups_config.py:274), so all 3 GTP peers write the same key as the main copy. Bridge already has a helper that fixes exactly this for weights (checkpointing.py:4085), but the RNG code does not use it. No test saves with GTP > 1: the PR has none, and Bridge's GTP functional test turns saving off (test_pretrain.py:202).

Where the fix lands: Megatron-Bridge, one function. get_rng_state should use _checkpoint_dp_cp_group. Stopgap in NeMo-RL: set save_rng=False when GTP > 1. NeMo-RL already loads with load_rng=False (setup.py:1985). No driver change.

Nothing below is implemented. It has not been run.

@@ megatron/bridge/training/checkpointing.py L578-L601 · get_rng_state
     rng_state_list = None
-    if torch.distributed.is_initialized() and pg_collection.dp_cp.size() > 1 and data_parallel_random_init:
-        rng_state_list = [None for i in range(pg_collection.dp_cp.size())]
-        torch.distributed.all_gather_object(rng_state_list, rng_state, group=pg_collection.dp_cp)
+    # Include GTP peers: they share (pp, tp, dp) and would all claim one key.
+    dp_cp = _checkpoint_dp_cp_group(pg_collection)
+    if torch.distributed.is_initialized() and dp_cp.size() > 1 and data_parallel_random_init:
+        rng_state_list = [None for i in range(dp_cp.size())]
+        torch.distributed.all_gather_object(rng_state_list, rng_state, group=dp_cp)
     ...
         rng_state_list = ShardedObject(
             key,
             rng_state_list,
-            (pp_size, tp_size, pg_collection.dp_cp.size()),
-            (pp_rank, tp_rank, pg_collection.dp_cp.rank()),
+            (pp_size, tp_size, dp_cp.size()),
+            (pp_rank, tp_rank, dp_cp.rank()),
             replica_id=0,
         )

In the example, the key's last coordinate becomes 0, 1 and 2 for r0, r2 and r4, so the 6 GPUs write 6 different RNG entries instead of 2 keys claimed 3 times each.

The review asks for a test before any fix: run the SFT v2 save/resume test (a baseline run, an interrupted run and a resumed run, compared step by step) with GTP turned on.

🔴 review comment C8 — setup.py:2624

So what. The worry is reasonable, but neither refit nor checkpointing chooses GPUs by a GTP axis. Every GTP detail they need lives in Megatron on each GPU, so the driver can keep its 4-axis layout. The check did find two real gaps, both in Megatron-Bridge. Refit is out of SFT v2's reach and is blocked by C3. The RNG save is in SFT v2's path: from the code, a GTP run's checkpoint save looks likely to fail, and no CI test saves with GTP on. C8 asks to turn GTP on in the SFT v2 save/resume test. Decide the layout before anyone saves: the loader's placement check includes the layout's axes (topology.py:114), so a checkpoint saved under one layout will not resume under the other.