feat(megatron): support GTP weight rematerialization for training + refit

examples/configs/grpo_math_1B_megatron.yaml

@@ -116,7 +116,14 @@ policy:
     # Use ["moe"] for MoE models to recompute expert activations only.
     recompute_modules: null
     tensor_model_parallel_size: 1
+    # Generalized Tensor Parallelism (GTP): total shards per weight across the TP
+    # and GTP axes. Must be a multiple of tensor_model_parallel_size; the quotient
+    # is the GTP degree, an extra dim-0 weight sharding taken out of the DP axis
+    # and all-gathered on demand. null (or equal to TP) leaves GTP off.
+    tensor_parallel_num_weight_shards: null
     expert_tensor_parallel_size: 1
+    # Expert-side counterpart, relative to expert_tensor_parallel_size.
+    expert_tensor_parallel_num_weight_shards: null
     expert_model_parallel_size: 1
     pipeline_model_parallel_size: 1
     num_layers_in_first_pipeline_stage: null

nemo_rl/models/generation/megatron/config.py

@@ -117,6 +117,17 @@ def dedicated_inference_megatron_cfg(
     """
     inference_mcfg = merged_inference_megatron_cfg(policy_config)
     inference_mcfg["context_parallel_size"] = 1
+    # Inference never uses GTP either: refit reassembles the training model's
+    # extra dim-0 weight shards into whole inference weights, and rematerializing
+    # them on every forward would be pure overhead. Pin the inference weight-shard
+    # counts to its own TP degrees so a GTP-sharded training model always takes
+    # the dedicated-inference-model reshard path.
+    inference_mcfg["tensor_parallel_num_weight_shards"] = inference_mcfg[
+        "tensor_model_parallel_size"
+    ]
+    inference_mcfg["expert_tensor_parallel_num_weight_shards"] = inference_mcfg[
+        "expert_tensor_parallel_size"
+    ]
 
     train_mcfg = cast(dict[str, Any], policy_config["megatron_cfg"])
     layout_keys = (
@@ -125,8 +136,24 @@ def dedicated_inference_megatron_cfg(
         "expert_model_parallel_size",
         "expert_tensor_parallel_size",
         "context_parallel_size",
+        "tensor_parallel_num_weight_shards",
+        "expert_tensor_parallel_num_weight_shards",
     )
-    layout_differs = any(inference_mcfg[k] != train_mcfg[k] for k in layout_keys)
+    # A null/absent weight-shard count means "no GTP", i.e. exactly the TP
+    # degree. Normalize before comparing so `null` and an explicit TP-equal
+    # value are not mistaken for a layout change.
+    normalized_train = {
+        **train_mcfg,
+        "tensor_parallel_num_weight_shards": (
+            train_mcfg.get("tensor_parallel_num_weight_shards")
+            or train_mcfg["tensor_model_parallel_size"]
+        ),
+        "expert_tensor_parallel_num_weight_shards": (
+            train_mcfg.get("expert_tensor_parallel_num_weight_shards")
+            or train_mcfg["expert_tensor_parallel_size"]
+        ),
+    }
+    layout_differs = any(inference_mcfg[k] != normalized_train[k] for k in layout_keys)
     impl_differs = inference_mcfg.get("transformer_impl") != train_mcfg.get(
         "transformer_impl"
     )

nemo_rl/models/megatron/setup.py

@@ -13,13 +13,16 @@
 # limitations under the License.
 
 import copy
+import functools
 import hashlib
+import inspect
 import json
 import os
 import threading
 import time
 import warnings
 from collections.abc import Mapping
+from contextlib import contextmanager
 from dataclasses import fields, is_dataclass, replace
 from typing import Any, Callable, Optional, TypeVar
 
@@ -62,6 +65,9 @@
 from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
 from megatron.core import parallel_state
 from megatron.core.inference.shards import build_inference_pg_collection
+from megatron.core.model_parallel_config import (
+    resolve_tensor_parallel_weight_shards,
+)
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer import MegatronModule
 from megatron.core.transformer.enums import AttnBackend, InferenceCudaGraphScope
@@ -684,6 +690,9 @@ def setup_model_config(
     # Apply MoE settings
     _apply_moe_config(model_cfg, config)
 
+    # Apply GTP settings (after MoE, which finalizes expert_tensor_parallel_size)
+    _apply_gtp_config(model_cfg, config)
+
     # Apply MTP settings
     _apply_mtp_config(model_cfg, config)
 
@@ -894,6 +903,125 @@ def _apply_parallelism_config(model_cfg: Any, config: PolicyConfig) -> None:
         )
 
 
+def _apply_gtp_config(model_cfg: Any, config: PolicyConfig) -> None:
+    """Apply Generalized Tensor Parallelism (GTP) weight-shard configuration.
+
+    ``tensor_parallel_num_weight_shards`` is MCore's user-facing knob: the total
+    number of shards a weight is split into across the TP and GTP axes. GTP is
+    the quotient ``tensor_parallel_num_weight_shards // tensor_model_parallel_size``
+    and shards each weight further along dim 0, rematerializing it on demand.
+
+    The provider dataclass was already constructed by the time NeMo-RL applies
+    its settings, so ``ModelParallelConfig.__post_init__`` has already run and
+    will not re-derive the internal ``gtp_weight_remat_size``. Derive it here
+    with MCore's own reconciliation helper so the two stay consistent.
+    """
+    num_weight_shards = config["megatron_cfg"].get(
+        "tensor_parallel_num_weight_shards", None
+    )
+    expert_num_weight_shards = config["megatron_cfg"].get(
+        "expert_tensor_parallel_num_weight_shards", None
+    )
+    if num_weight_shards is None and expert_num_weight_shards is None:
+        return
+
+    (
+        model_cfg.tensor_parallel_num_weight_shards,
+        model_cfg.gtp_weight_remat_size,
+    ) = resolve_tensor_parallel_weight_shards(
+        model_cfg.tensor_model_parallel_size,
+        num_weight_shards,
+        model_cfg.gtp_weight_remat_size,
+    )
+    (
+        model_cfg.expert_tensor_parallel_num_weight_shards,
+        model_cfg.expert_gtp_weight_remat_size,
+    ) = resolve_tensor_parallel_weight_shards(
+        model_cfg.expert_tensor_parallel_size,
+        expert_num_weight_shards,
+        model_cfg.expert_gtp_weight_remat_size,
+        shards_field="expert_tensor_parallel_num_weight_shards",
+        tp_field="expert_tensor_parallel_size",
+    )
+
+
+@contextmanager
+def _gtp_process_groups(model_cfg: Any):
+    """Make Megatron-Bridge create the GTP rematerialization process groups.
+
+    Bridge's ``initialize_megatron`` calls
+    ``parallel_state.initialize_model_parallel`` without ``gtp_remat_size`` /
+    ``expert_gtp_remat_size``, so those groups are never created and
+    ``gtp_weight_remat_size`` on the model config silently has no effect (MCore
+    layers only wrap their weights for GTP when the group exists and has size
+    > 1). Bridge resolves the function off the ``parallel_state`` module at call
+    time, so temporarily wrapping the module attribute is enough to forward the
+    sizes. No-op unless GTP is actually requested.
+
+    TODO: drop this once Megatron-Bridge forwards the GTP sizes itself.
+    """
+    gtp_remat_size = getattr(model_cfg, "gtp_weight_remat_size", 1)
+    expert_gtp_remat_size = getattr(model_cfg, "expert_gtp_weight_remat_size", 1)
+    if gtp_remat_size <= 1 and expert_gtp_remat_size <= 1:
+        yield
+        return
+
+    original = parallel_state.initialize_model_parallel
+    accepted = inspect.signature(original).parameters
+    missing = [
+        name
+        for name in ("gtp_remat_size", "expert_gtp_remat_size")
+        if name not in accepted
+    ]
+    if missing:
+        raise RuntimeError(
+            "GTP was requested via megatron_cfg.tensor_parallel_num_weight_shards "
+            f"(gtp={gtp_remat_size}, expert_gtp={expert_gtp_remat_size}), but the "
+            f"pinned Megatron-LM's initialize_model_parallel does not accept {missing}. "
+            "Bump Megatron-LM to a revision with GTP support."
+        )
+
+    # Accepting the kwargs is not the same as being able to use GTP: MCore's GTP
+    # core requires a newer TransformerEngine than it requires overall, and when
+    # that check fails it degrades to import stubs rather than raising. Nothing
+    # downstream re-checks -- MCore's own layer constructors gate GTP on the
+    # process-group size alone -- so without this the run dies ~6 frames deep in
+    # generalized_tensor_parallelism.py on `isinstance() arg 2 must be a type`,
+    # which says nothing about TE. Fail here with the actual requirement instead.
+    # Local import: gtp_api does not exist on Megatron-LM revisions predating GTP,
+    # and setup.py must keep importing against those.
+    from megatron.core.tensor_parallel import gtp_api
+
+    if not gtp_api.HAVE_GTP:
+        raise RuntimeError(
+            "GTP was requested via megatron_cfg.tensor_parallel_num_weight_shards "
+            f"(gtp={gtp_remat_size}, expert_gtp={expert_gtp_remat_size}), but the "
+            "installed Megatron-LM reports GTP unavailable (HAVE_GTP=False). This is "
+            "almost always TransformerEngine being too old -- MCore's GTP core "
+            "requires a newer TE than MCore itself does. Check "
+            "megatron.core.tensor_parallel.generalized_tensor_parallelism for the "
+            "minimum version, or leave tensor_parallel_num_weight_shards unset/null "
+            "to run without GTP."
+        )
+
+    @functools.wraps(original)
+    def initialize_model_parallel_with_gtp(*args: Any, **kwargs: Any) -> Any:
+        kwargs.setdefault("gtp_remat_size", gtp_remat_size)
+        kwargs.setdefault("expert_gtp_remat_size", expert_gtp_remat_size)
+        return original(*args, **kwargs)
+
+    parallel_state.initialize_model_parallel = initialize_model_parallel_with_gtp
+    print(
+        f"[gtp] enabling GTP weight rematerialization (gtp_remat_size={gtp_remat_size} "
+        f"expert_gtp_remat_size={expert_gtp_remat_size})",
+        flush=True,
+    )
+    try:
+        yield
+    finally:
+        parallel_state.initialize_model_parallel = original
+
+
 def _apply_moe_config(model_cfg: Any, config: PolicyConfig) -> None:
     """Apply Mixture of Experts configuration."""
     model_cfg.expert_tensor_parallel_size = config["megatron_cfg"][
@@ -1602,6 +1730,25 @@ def build_inference_model(
         inference_provider.sequence_parallel
         and inference_provider.tensor_model_parallel_size > 1
     )
+    # GTP shards weights to save training memory and rematerializes them on every
+    # forward/backward — pure overhead for inference. Refit is what maps the
+    # training model's GTP shards onto these whole weights, so force the layout
+    # off here rather than inheriting it from the training provider snapshot.
+    #
+    # Forcing it off on the provider is necessary but not sufficient: megatron's
+    # layers resolve their GTP axis from the process-group collection, not from
+    # the provider, and `build_inference_pg_collection` must declare the axis
+    # explicitly off for that resolution to not fall back to the *training* MPU
+    # globals. That is NVIDIA/Megatron-LM#6940; until it merges, running with
+    # `tensor_parallel_num_weight_shards` < TP needs a megatron-core carrying it.
+    inference_provider.tensor_parallel_num_weight_shards = (
+        inference_provider.tensor_model_parallel_size
+    )
+    inference_provider.gtp_weight_remat_size = 1
+    inference_provider.expert_tensor_parallel_num_weight_shards = (
+        inference_provider.expert_tensor_parallel_size
+    )
+    inference_provider.expert_gtp_weight_remat_size = 1
     # Inference never trains: disable recompute.
     inference_provider.recompute_granularity = None
     inference_provider.recompute_method = None
@@ -1670,11 +1817,12 @@ def setup_model_and_optimizer(
     state.initialize_async_checkpoint_worker()
 
     megatron_cfg.dist.external_gpu_device_mapping = True
-    initialize_megatron(
-        cfg=megatron_cfg,
-        get_embedding_ranks=get_embedding_ranks,
-        get_position_embedding_ranks=get_position_embedding_ranks,
-    )
+    with _gtp_process_groups(megatron_cfg.model):
+        initialize_megatron(
+            cfg=megatron_cfg,
+            get_embedding_ranks=get_embedding_ranks,
+            get_position_embedding_ranks=get_position_embedding_ranks,
+        )
 
     if megatron_cfg.ft and megatron_cfg.ft.enable_ft_package:
         fault_tolerance.setup(megatron_cfg, state)

nemo_rl/models/policy/__init__.py

@@ -348,6 +348,13 @@ class MegatronConfig(TypedDict):
     # when None. Use ["moe"] to recompute only expert activations (production-proven config).
     recompute_modules: NotRequired[list[str] | None]
     tensor_model_parallel_size: int
+    # Generalized Tensor Parallelism (GTP). Total number of shards each weight is
+    # split into across the tensor-parallel and GTP axes; must be a multiple of
+    # tensor_model_parallel_size. The quotient is MCore's gtp_weight_remat_size:
+    # an extra dim-0 sharding of every weight, carved out of the data-parallel
+    # axis and all-gathered on demand in the forward and backward passes. Absent
+    # or null leaves GTP off (equivalent to tensor_model_parallel_size).
+    tensor_parallel_num_weight_shards: NotRequired[int | None]
     pipeline_model_parallel_size: int
     num_layers_in_first_pipeline_stage: int | None
     num_layers_in_last_pipeline_stage: int | None
@@ -365,6 +372,9 @@ class MegatronConfig(TypedDict):
     sequence_parallel: bool
     freeze_moe_router: bool
     expert_tensor_parallel_size: int
+    # Expert-side counterpart of tensor_parallel_num_weight_shards; must be a
+    # multiple of expert_tensor_parallel_size. Independent of the dense knob.
+    expert_tensor_parallel_num_weight_shards: NotRequired[int | None]
     expert_model_parallel_size: int
     # If True, defer the casting of logits to float32 until the backward pass.
     # If you are using logprob_chunk_size, you must set this to True.

tests/functional/L1_Functional_Tests_Megatron_4.sh

@@ -48,11 +48,31 @@ megatron_generation_supported() {
     return 0
 }
 
+# GTP needs a newer TransformerEngine than MCore itself does (MCore's GTP core
+# declares its own TE floor and degrades to import stubs below it). NeMo-RL
+# currently pins TE release_v2.15, which is under that floor, so the GTP test
+# cannot pass yet -- registering it unguarded would just make L1 permanently red.
+# Ask MCore directly rather than hardcoding a version here, so this re-enables
+# itself the moment the TE pin is bumped.
+# TODO: remove this guard once NeMo-RL's TE pin satisfies MCore's GTP floor.
+megatron_gtp_supported() {
+    if ! uv run --no-sync python -c \
+        'import sys; from megatron.core.tensor_parallel.gtp_api import HAVE_GTP; sys.exit(0 if HAVE_GTP else 1)' \
+        &> /dev/null; then
+        echo "WARNING: Skipping GTP test; Megatron-LM reports HAVE_GTP=False (TransformerEngine too old)"
+        return 1
+    fi
+    return 0
+}
+
 if megatron_generation_supported; then
     run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_topology.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_non_colocated.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard.sh
+    if megatron_gtp_supported; then
+        run_test  uv run --no-sync bash ./tests/functional/grpo_megatron_generation_gtp.sh
+    fi
     run_test fast uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_async_grpo.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_colocated_reshard_async_grpo.sh
     run_test      uv run --no-sync bash ./tests/functional/grpo_megatron_generation_async_gym.sh

tests/functional/grpo_megatron_generation_gtp.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
+PROJECT_ROOT=$(realpath $SCRIPT_DIR/../..)
+# Mark the current repo as safe, since wandb fetches metadata about the repo
+git config --global --add safe.directory $PROJECT_ROOT
+
+set -eou pipefail
+
+EXP_NAME=$(basename $0 .sh)
+EXP_DIR=$SCRIPT_DIR/$EXP_NAME
+LOG_DIR=$EXP_DIR/logs
+JSON_METRICS=$EXP_DIR/metrics.json
+RUN_LOG=$EXP_DIR/run.log
+export PYTHONPATH=${PROJECT_ROOT}:${PYTHONPATH:-}
+
+rm -rf $EXP_DIR $LOG_DIR
+mkdir -p $EXP_DIR $LOG_DIR
+
+# GTP refit: training runs TP1 x GTP2 (tensor_parallel_num_weight_shards=2 over
+# TP=1), so every training weight is split along dim 0 across the 2 ranks and
+# rematerialized on demand. Inference runs TP1 with whole weights, so refit must
+# reassemble each weight from its GTP shards -- including the alignment padding
+# that carries no logical data. Getting that wrong yields subtly wrong inference
+# weights, which shows up as generation/training logprob disagreement, hence the
+# token_mult_prob_error gate.
+#
+# Requires a Megatron-LM with GTP support in megatron/core/resharding
+# (NVIDIA/Megatron-LM#6133). Older revisions reject the plan outright.
+cd $PROJECT_ROOT
+uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJECT_ROOT/nemo_rl \
+    $PROJECT_ROOT/examples/run_grpo.py \
+    --config $PROJECT_ROOT/examples/configs/grpo_math_1B_megatron.yaml \
+    policy.model_name=Qwen/Qwen2.5-0.5B \
+    grpo.num_prompts_per_step=2 \
+    grpo.num_generations_per_prompt=4 \
+    policy.train_global_batch_size=4 \
+    policy.logprob_batch_size=4 \
+    policy.train_micro_batch_size=1 \
+    policy.megatron_cfg.tensor_model_parallel_size=1 \
+    policy.megatron_cfg.tensor_parallel_num_weight_shards=2 \
+    policy.generation.backend=megatron \
+    policy.generation.mcore_generation_config.refit_backend=nccl \
+    cluster.gpus_per_node=2 \
+    grpo.max_num_steps=2 \
+    logger.tensorboard_enabled=true \
+    logger.log_dir=$LOG_DIR \
+    logger.wandb_enabled=false \
+    logger.monitor_gpus=true \
+    checkpointing.enabled=false \
+    $@ \
+    2>&1 | tee $RUN_LOG
+
+uv run tests/json_dump_tb_logs.py $LOG_DIR --output_path $JSON_METRICS
+
+uv run tests/check_metrics.py $JSON_METRICS \
+    'max(data["train/token_mult_prob_error"]) < 1.05'
+
+# Guard against a vacuous pass. GTP is only active when MCore actually creates
+# the gtp_remat process group; if that silently degrades to size 1 the run is an
+# ordinary TP1 job that would pass every metric above without testing anything.
+if ! grep -q "\[gtp\] enabling GTP weight rematerialization (gtp_remat_size=2" $RUN_LOG; then
+    echo "FAIL: GTP marker not found; the training model was not GTP-sharded"
+    exit 1
+fi
+
+# GTP-sharded training weights must be reassembled into a dedicated inference
+# model. Without this, generation would run on the training model directly and
+# the refit path under test would never execute.
+if ! grep -q "\[colocated-reshard\] building dedicated inference model" $RUN_LOG; then
+    echo "FAIL: colocated-reshard marker not found; dedicated inference model was never built"
+    exit 1
+fi