Why GTP can't be turned on with the setting this PR reads

Explainer for NVIDIA-NeMo/RL #4512. All permalinks pin the PR head 339e2020; Megatron-Bridge at ec835530, Megatron-LM at b4d72b79.

Before: The driver had no idea of GTP. Its world-size check used PP × CP × TP only.
This PR: Reads megatron_cfg.model_overrides.gtp_weight_remat_size to size GTP, and adds it to the world-size check and the data layout.
Status: Open, not fixed in this PR. That setting is rejected on every worker, so the new GTP code never runs. The review asks to read Bridge's real setting instead.

Three pages on this PR, one question each. This one: how GTP gets turned on.

Background. NeMo-RL builds each Megatron model from a Bridge model provider. model_overrides copies any field the provider accepts in its constructor. Bridge's finalize() computes the GTP size itself. Dotted words have a short explanation: hover or tap them.

How GTP is turned on. You set one number, tensor_parallel_num_weight_shards: how many pieces each weight is cut into across GPUs. Bridge divides it by TP to get the GTP size, 12 // 4 = 3, and hands that to Megatron, which works out DP = 24 GPUs // (TP 4 × GTP 3) = 2 (initialize.py:822-826). Nobody sets DP. gtp_weight_remat_size is the computed result, so Bridge declares it field(init=False, default=1), meaning it cannot be passed in (transformer_config.py:173-176, :204-218). Megatron-LM calls the shard count "the user-facing knob" (model_parallel_config.py:56-67). NeMo-RL's override merge only accepts fields that can be passed in (setup.py:1248), but the PR's driver helper reads only gtp_weight_remat_size (lm_policy.py:78-88). The PR's test calls that helper on a plain dict, so it never meets Bridge (test_gtp_config.py).

Three configs: the two you can write today, and the proposed one

Which DP: 2 or 6?

Nobody sets DP, in any of A, B or C. NeMo-RL has no DP setting. Both sides work it out from the GPU count. But the word "DP" names two different numbers here:

Copies of the model: 2. GPUs ÷ (TP × GTP) = 24 ÷ 12. Megatron calls this the replicate degree.
Data lanes: 6. GPU groups that each train on their own batch: 2 copies × 3 GTP peers. Megatron's get_data_parallel_world_size() returns this one by default (with_gtp_remat=True, parallel_state.py:2198-2207).

Everything that splits data or counts tokens for the loss must use 6. Megatron sums each gradient over all 6 lanes and divides by the token count of all 6 (finalize_model_grads.py:546-556). NeMo-RL counts the tokens over Megatron's default group, which is those 6 lanes (megatron_policy_worker.py:2096-2100). So the loss is right. Counting over only the 2 copies would make each gradient about 3× too big.

24 GPUs, TP 4	A	B (today)	PR's code, once its key works	C (proposed)
You set DP?	no	no	no	no
Driver `data_parallel` axis (lm_policy.py:500)	—	6	2, plus a separate GTP axis of 3	6
Megatron: copies / data lanes	— (crashes)	2 / 6	2 / 6	2 / 6
Loss token count summed over	—	6 lanes ✓	6 lanes ✓	6 lanes ✓
Calls other than SFT v2 training	—	all 24 GPUs ✓	only 8 GPUs get the call → hang	all 24 GPUs ✓

The loss row is ✓ everywhere because workers ask Megatron for the lane count; they never use the driver's number. The driver's number decides who gets each call and how a batch is split, e.g. gbs // dp_size in the older train path (megatron_policy_worker.py:1181). There, 6 is right and the PR's 2 is wrong. That column is why C goes with page 1's fix: keep the driver at 6 lanes.

What would fix it (C): proposed, not in the PR

Where the fix lands: RL only, three places. A typed tensor_parallel_num_weight_shards field in MegatronConfig; always set on the provider in _apply_parallelism_config, like TP and CP; and read by the driver helper for the world-size check only. Setting it every run matters because a provider loaded from a checkpoint can carry a shard count saved at a different TP. No Bridge change: Bridge already derives GTP from the shard count.

Nothing below is implemented. The expert-side setting (expert_tensor_parallel_num_weight_shards) gets the same treatment.

@@ lm_policy.py L78-L88 · _megatron_gtp_weight_remat_size
 def _megatron_gtp_weight_remat_size(config: PolicyConfig) -> int:
-    """Return the dense GTP data-shard size, or one when GTP is disabled."""
-    overrides = (config.get("megatron_cfg") or {}).get("model_overrides") or {}
-    size = overrides.get("gtp_weight_remat_size", 1)
-    if size is None:
-        return 1
-    if isinstance(size, bool) or not isinstance(size, int) or size < 1:
-        raise ValueError(
-            "model_overrides.gtp_weight_remat_size must be a positive integer"
-        )
-    return size
+    """Return the dense GTP size Bridge derives from tensor_parallel_num_weight_shards."""
+    megatron_cfg = config["megatron_cfg"]
+    shards = megatron_cfg["tensor_parallel_num_weight_shards"]
+    if shards is None:
+        return 1
+    # Same rule as Bridge's TransformerConfig._finalize_gtp_weight_shards.
+    return shards // megatron_cfg["tensor_model_parallel_size"]

B and C run the same model on the same GPUs. The only difference is step 1: the driver now knows GTP is 3, so its world-size check asks for at least TP × GTP = 12 GPUs and stops an 8-GPU run before any worker starts. This must ship together with page 1's fix: once the driver sees GTP 3, the PR's code would turn on the fifth layout axis, and that axis hangs every call except SFT v2 training. The test should run the real _merge_model_overrides on a Bridge provider and compare its finalized GTP size with this helper.

🔴 review comment, lm_policy.py:78

So what. As written, no config turns GTP on through this PR's code: the key it reads crashes every worker, and the key that works is invisible to it. GTP data still lands on the right GPUs with the working key, because the layout NeMo-RL already had handles it (page 1). Reading the real setting fixes the world-size check; page 1's fix removes the code that would hang once this is fixed.