Where each GTP GPU gets its data: the layout NeMo-RL already had does it, and the PR's new axis breaks other calls

Explainer for NVIDIA-NeMo/RL #4512 (SFT processor config and GTP data placement). All permalinks pin the PR head 339e2020; Megatron-LM at b4d72b79.

Before
The driver laid GPUs out in 4 named axes (PP, DP, CP, TP) and did not count GTP in its world-size check.
This PR
Adds a fifth axis for GTP, plus code in the Energon loader planner and the SFT v2 train dispatch that sends one batch per (DP, GTP) pair.
Status
Open, not fixed in this PR. The fifth axis is not needed, and every call except SFT v2 training would skip two thirds of the GPUs. The review asks to keep the 4-D layout.

Three pages on this PR, one question each. This one: which GPUs get work.

Background. GTP splits each layer's weights across G more GPUs. Unlike tensor parallelism, every GPU in a GTP group trains on its own batch. The driver gives every call a layout: a grid of GPU numbers with named axes. A call names the axes it splits data over and the axes that get a copy. Any other axis only gets the call at index 0. Dotted words have a short explanation: hover or tap them.

Megatron numbers GPUs TP fastest, then CP, then GTP, then DP, then PP (parallel_state.py:602-618), and its data-parallel rank already counts GTP GPUs as separate lanes by default (parallel_state.py:2225-2231). So the old 4-D layout's data_parallel index for GPU (d, g) is d·G + g, which is exactly the lane the SFT loader already asks Megatron for (sft_worker.py:82) and the number the PR now computes by hand (topology.py:89). The PR's fifth axis gives the same lanes for SFT v2 training, but every other call still splits only data_parallel.

Both layouts number the same 6 lanes the same way

A lane here is a group of GPUs that trains on its own batch. With DP = 2 and GTP = 3 there are 6 lanes, and each lane is 4 GPUs (the TP group). Layout A gives each lane one number directly. Layout B gives it two coordinates, and the PR combines them into one number. The figure lines the numbering up per lane.

Same 24 GPUs (DP = 2, GTP = 3, TP = 4). Each column = the 4 TP GPUs that share one (d, g). 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 Megatron's coordinates of those 4 GPUs (rank = t + 4·(g + 3·d)) parallel_state.py:602-618 d=0, g=0 d=0, g=1 d=0, g=2 d=1, g=0 d=1, g=1 d=1, g=2 A · old 4-D layout: index on the data_parallel axis (= rank // 4) lm_policy.py:341-354 (base) 0 1 2 3 4 5 B · PR's 5-D layout: (data_parallel, gtp_weight_remat) coordinates lm_policy.py:356-370 (0, 0) (0, 1) (0, 2) (1, 0) (1, 1) (1, 2) … which the PR turns into one lane number: dp·3 + gtp topology.py:89 0 1 2 3 4 5 Megatron's own get_data_parallel_rank() on those GPUs (the SFT loader asks this) sft_worker.py:82 0 1 2 3 4 5 lane number: which batch these 4 GPUs train on. All three green rows agree column by column, so batch k goes to the same 4 GPUs in A, in B, and in Megatron's own numbering.

Why A already matches: Megatron numbers GPUs TP fastest, then GTP, then DP, so GPU r = t + 4·(g + 3·d). Reshaping 24 GPUs as (pp=1, -1, cp=1, tp=4) leaves r // 4 = 3·d + g on the middle axis, which is the same lane number B computes by hand. So splitting the producer's batches along A's data_parallel axis already sends batch k to lane k.

How SFT v2 feeds each lane: the loader never needs (dp, gtp).
  1. One loader per lane. SFT v2 builds a loader on every GPU with TP = CP = PP = 0: GPUs 0, 4, 8, 12, 16, 20, the same 6 GPUs in either layout (sft_v2.py:268-272).
  2. Each loader asks Megatron which lane it is. It calls get_data_parallel_rank() and get_data_parallel_world_size(), which count GTP by default, so the answers are 0..5 of 6. Energon reads shard k of 6 from that (sft_worker.py:82-83, unchanged by the PR).
  3. The driver collects the 6 batches in lane order. It already checks they come back as exactly 0..5 (sft_v2.py:300-305).
  4. The driver sends batch k to lane k. Before the PR: split on data_parallel (6 entries) (tq_policy.py:801-803, base). The PR: the same list reshaped into [[0,1,2],[3,4,5]] and split on (dp, gtp) (tq_policy.py:763-770). Same GPUs, same batches.
The only place G has to be known is Megatron itself (the PR's setup calls, which stay) and the driver's world-size check (page 2). The PR's (dp, gtp) code builds the pair from one number and turns it straight back into that number (topology.py:88-89).

Same 24 GPUs, two layouts, two kinds of call

24 GPUs: DP = 2, GTP = 3, TP = 4 (CP = PP = 1). Megatron numbers them TP fastest, then CP, GTP, DP, PP. 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 d0 g0 d0 g1 d0 g2 d1 g0 d1 g1 d1 g2 order: parallel_state.py:602-618 A · before the PR: 4-D layout reshape(pp, -1, cp, tp) lm_policy.py:341-354 (base) data_parallel has 6 entries; GPU (d, g) is at index d·3 + g, the same as Megatron's own data-parallel rank dp idx 0 dp idx 1 dp idx 2 dp idx 3 dp idx 4 dp idx 5 placed SFT v2 train get_logprobs, train, … ✓ both calls reach all 24 GPUs; 6 distinct batches, one per lane B · this PR: 5-D layout reshape(pp, dp, gtp, cp, tp) lm_policy.py:356-370 new axis gtp_weight_remat (size 3); data_parallel shrinks to 2 placed SFT v2 train (shards data_parallel + gtp) tq_policy.py:754-770 get_logprobs, train, … (still shard data_parallel only) lm_policy.py:764 ✗ only the 8 GPUs with g = 0 are called; the other 16 never start, so the GTP weight gather waits worker_groups.py:1188-1196 (an axis that is neither sharded nor replicated only runs index 0) gets the call, own batch never called GPU (rank number) The hang in B is read from the dispatch code, not run on a GPU. Today no working config builds layout B (page 2), so this is a trap, not a live break.

In layout B the other calls see gtp_weight_remat as an axis they neither split nor copy, so only index 0 runs (worker_groups.py:1188-1196). The PR's only guard is on SFT v2 training (tq_policy.py:813-818).

What would fix it: proposed, not in the PR

Where the fix lands: RL only. Keep the 4-D layout in Policy.__init__; delete GTP_WEIGHT_REMAT_AXIS (named_sharding.py), the GTP loop in topology.py, and the GTP nesting and guard in tq_policy.py. Keep the GTP setup calls in setup.py. No Megatron change: Megatron's default data-parallel group already includes GTP.

Nothing below is implemented. With GTP = 1 this is exactly the code before the PR.

@@ lm_policy.py L356-L370 · Policy.__init__
-        # MCore rank order is TP, CP, GTP, DP, PP, with TP varying fastest.
-        shape = [pp_size, -1, cp_size, tp_size]
-        axes = [
-            "pipeline_parallel",
-            "data_parallel",
-            "context_parallel",
-            "tensor_parallel",
-        ]
-        if gtp_size > 1:
-            shape.insert(2, gtp_size)
-            axes.insert(2, GTP_WEIGHT_REMAT_AXIS)
-        self.sharding_annotations = NamedSharding(
-            layout=np.arange(cluster.world_size()).reshape(shape),
-            names=axes,
-        )
+        # MCore's rank order is TP, CP, GTP, DP, PP (TP fastest), and its
+        # data-parallel rank includes GTP peers (with_gtp_remat=True by default).
+        # So GTP folds into data_parallel: index = dp * gtp_size + gtp.
+        self.sharding_annotations = NamedSharding(
+            layout=np.arange(cluster.world_size()).reshape(
+                pp_size,  # PP
+                -1,  # DP x GTP
+                cp_size,  # CP
+                tp_size,  # TP
+            ),
+            names=["pipeline_parallel", "data_parallel", "context_parallel", "tensor_parallel"],
+        )

One unit test should pin the folding: build the 4-D layout and check that GPU (d, g) lands at data_parallel index d·G + g.

So what does the author do instead of the GTP axis? Nothing new. Nobody ever writes DP, not 2 and not 6. The only GTP input is one setting, the shard count (page 2). Only the two places that care how weights are split read it. Everything that moves data uses the flat lane number 0..5, which both sides already compute on their own. Same example: 24 GPUs, TP 4, shard count 12.

PieceNeeds G?After the fix
Config (YAML)inputTP 4 and shard count 12. No DP setting exists.
Megatron setupyesBridge: GTP = 12 // 4 = 3 (transformer_config.py), passed to Megatron (initialize.py), which works out DP = 24 // 12 = 2. Already works; no change.
Driver world-size checkyesGTP = shard count // TP, from the same setting (page 2). The only place in NeMo-RL that needs G.
Driver layoutnoreshape(PP, -1, CP, TP): the -1 becomes 24 / 4 = 6 lanes.
Loader plannoBack to the code before the PR: one loader per lane 0..5 (topology.py, base).
Energon loadernoAlready asks Megatron: lane 0..5 out of 6 (sft_worker.py). Unchanged.
DispatchnoSplits on data_parallel only, like every other call (tq_policy.py, base).

Driver and Megatron count the same 6 lanes without talking to each other. The driver gets 24 / (PP·CP·TP) = 6. Megatron's data-parallel size includes GTP by default, so it gets DP × GTP = 2 × 3 = 6. The figure at the top shows they also number them the same way.

🔴 review comment, lm_policy.py:356-370

So what. The PR's per-(DP, GTP) code only does what the old layout already did, and only for SFT v2 training. Every other call (logprobs, the older train path, scoring) would skip the GTP GPUs and wait forever. Keeping the 4-D layout makes all calls correct and deletes about 60 lines. Today the fifth axis is never built, because the setting it reads cannot be set: see page 2.