fix(ckpt): support GTP checkpoint groups and HF conversion (#6116)

Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

docs/training/checkpointing.md

@@ -191,6 +191,59 @@ checkpoint = CheckpointConfig(
 
 **Important**: When using Megatron FSDP (`use_megatron_fsdp=True`), you must set `ckpt_format="fsdp_dtensor"`. Other formats are not compatible with FSDP's sharded parameter layout. See {doc}`megatron-fsdp` for complete FSDP configuration details.
 
+### Generalized Tensor Parallelism (GTP)
+
+Use `torch_dist` checkpoints for BF16 GTP models with TransformerEngine 2.19 or newer.
+Checkpoint metadata and fully parallel save/load include the GTP rematerialization
+axis. RNG checkpoint entries preserve distinct streams on dense and expert GTP ranks.
+
+HF import splits each mapped TP-local tensor into its stored GTP row shard. HF export
+gathers those rows and removes alignment padding before applying the existing TP/EP/PP
+mappings. Use `AutoBridge.export_hf_weights()` or `save_hf_weights()` for this export;
+the local HF parameter-view API cannot describe the extra GTP shard axis. Quantized
+HF export and LoRA adapter export/merge also gather before applying their mappings.
+Raw FP8 export tasks and local native MXFP8 parameter export reject GTP before
+yielding weights, because their storage views cannot preserve this layout.
+
+When loading a native GTP checkpoint with `load_megatron_model()`, pass `mp_overrides`
+with the saved TP size and weight shard count for each axis that uses GTP:
+`tensor_model_parallel_size` / `tensor_parallel_num_weight_shards` for dense weights,
+and `expert_tensor_parallel_size` / `expert_tensor_parallel_num_weight_shards` for
+expert weights. Axes without GTP can still change TP independently. Changing a GTP
+weight-sharding topology is
+rejected because the current Megatron-Core SwiGLU checkpoint layout can reorder
+gate/up rows. Training resume and finetuning enforce the same restriction, including
+loading a non-GTP native checkpoint into a GTP model. To change layouts, export HF
+weights from the original topology first, then import the HF weights into the new
+topology.
+
+Energon shards samples across DP and dense GTP ranks while CP peers consume the
+same samples. Dataloader checkpoints use the same group for file ownership and
+restore; old duplicated GTP states with a different stream count are rejected.
+With GTP, do not enable both `dist_ckpt_optim_fully_reshardable` and
+`distrib_optim_fully_reshardable_mem_efficient`: Core does not create the Gloo
+groups needed by that combination.
+
+The focused two-GPU regression covers exact BF16 HF round-trips, model and named
+RNG-stream resume, and rejection of unsafe topology changes, with fully parallel
+checkpoint I/O enabled and disabled:
+
+```bash
+uv run python -m torch.distributed.run --standalone --nproc_per_node=2 -m pytest \
+  tests/functional_tests/test_groups/converter/test_gtp_checkpoint_conversion.py \
+  tests/functional_tests/test_groups/converter/test_gtp_native_fp8_conversion.py \
+  tests/functional_tests/test_groups/data/energon/test_gtp_checkpoint_state.py -q
+```
+
+The H100 L1 launcher `L1_Launch_gtp_checkpoint_conversion.sh` runs this matrix in CI.
+GTP cases require TransformerEngine 2.19 or newer and skip on older versions; the
+ordinary TP baseline cases still run.
+
+The additional tests check native FP8 HF import against an independent TE cast and
+exact Energon sample continuation after restore. Native MXFP8 requires Blackwell
+and skips on H100; this does not validate native MXFP8 training or Muon optimizer
+state. Quantized and LoRA GTP mapping regressions use controlled unit fixtures.
+
 ## Performance Optimizations
 
 | Parameter | Type | Default | Description |

src/megatron/bridge/data/energon/base_energon_datamodule.py

@@ -23,6 +23,8 @@
 from megatron.energon.flavors.webdataset.default_generic_webdataset import DefaultGenericWebdatasetFactory
 from megatron.energon.typed_converter import JsonParser
 
+from megatron.bridge.training.gtp import _get_dataloader_process_group
+
 
 logger = logging.getLogger(__name__)
 
@@ -254,16 +256,11 @@ def __init__(
         self.kwargs = kwargs
 
     def _build_worker_config(self, num_workers: int, split: str = "train") -> WorkerConfig:
-        """Build a WorkerConfig using pg_collection, falling back to default_worker_config.
-
-        NOTE: We intentionally use the pure DP rank (pg_collection.dp)
-        rather than the combined DP-CP rank. With Megatron's rank ordering
-        (default "tp-cp-ep-dp-pp"), all CP ranks within the same DP replica
-        already share the same pure DP rank. This ensures that CP ranks
-        processing different sequence portions of the same batch receive
-        identical data from the dataloader.
-        Using dp_cp would be INCORRECT here — it would assign each CP rank
-        a unique rank, causing them to read different data shards.
+        """Build workers for distinct microbatches across data and GTP ranks.
+
+        Context-parallel peers consume the same sample before splitting its
+        sequence, so the worker group excludes CP. GTP peers consume different
+        microbatches and must have distinct worker ranks, including on resume.
         """
         if self.pg_collection is None or self.pg_collection.dp is None:
             logger.info(
@@ -272,9 +269,9 @@ def _build_worker_config(self, num_workers: int, split: str = "train") -> Worker
             )
             return WorkerConfig.default_worker_config(num_workers)
 
-        rank = self.pg_collection.dp.rank()
-        world_size = self.pg_collection.dp.size()
-        data_parallel_group = self.pg_collection.dp
+        data_parallel_group = _get_dataloader_process_group(self.pg_collection)
+        rank = data_parallel_group.rank()
+        world_size = data_parallel_group.size()
         cp_rank = self.pg_collection.cp.rank() if self.pg_collection.cp is not None else 0
         cp_size = self.pg_collection.cp.size() if self.pg_collection.cp is not None else 1
 

src/megatron/bridge/models/conversion/auto_bridge.py

@@ -2038,11 +2038,15 @@ def get_model(
                 ]
 
             model_config.finalize()
+            from megatron.bridge.training.gtp import classify_gtp_remat_chains, configure_gtp_remat
+
+            configure_gtp_remat(model_config)
             builder = model_config.get_builder_cls()(model_config)
             if pg_collection is None:
                 pg_collection = self._get_or_initialize_pg_collection(transformer_config)
             kwargs.setdefault("data_parallel_random_init", False)
             models = builder.build_distributed_models(pg_collection=pg_collection, **kwargs)
+            classify_gtp_remat_chains(models, model_config)
             for model in models:
                 model.model_config = model_config
             succeeded = True
@@ -2083,6 +2087,8 @@ def _get_or_initialize_pg_collection(
                 context_parallel_size=transformer_config.context_parallel_size or 1,
                 expert_model_parallel_size=transformer_config.expert_model_parallel_size or 1,
                 expert_tensor_parallel_size=transformer_config.expert_tensor_parallel_size,
+                gtp_remat_size=transformer_config.gtp_weight_remat_size,
+                expert_gtp_remat_size=transformer_config.expert_gtp_weight_remat_size,
             )
             if torch.cuda.is_available():
                 from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed

src/megatron/bridge/models/conversion/gtp.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Translate between stored GTP shards and the TP-local tensors used by mappings."""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+
+import torch
+
+
+def _is_gtp_param(param: torch.Tensor | None) -> bool:
+    return getattr(param, "is_gtp_weight_remat", False) is True
+
+
+def _get_mapping_shape(param: torch.Tensor) -> torch.Size:
+    """Return the logical TP-local shape, excluding GTP padding and sharding."""
+    if _is_gtp_param(param):
+        shape = list(param.shape)
+        shape[0] = shape[0] * param.group.size() - param.pad_length
+        return torch.Size(shape)
+    return param.shape
+
+
+@contextmanager
+def _gtp_weight_load_context(param: torch.Tensor, module: torch.nn.Module) -> Iterator[None]:
+    """Let TE update native FP8 storage while preserving the GTP parameter subclass."""
+    if _is_gtp_param(param) and getattr(param, "_gtp_native_fp8", False):
+        from megatron.core.tensor_parallel.gtp_api import gtp_native_fp8_load_context
+
+        with gtp_native_fp8_load_context(module):
+            yield
+    else:
+        yield
+
+
+def _slice_gtp_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
+    """Take the stored GTP row shard after the mapping has applied TP splitting."""
+    if not _is_gtp_param(param):
+        return weight
+    expected_shape = _get_mapping_shape(param)
+    if weight.shape != expected_shape:
+        raise ValueError(f"Expected TP-local GTP weight shape {expected_shape}, got {weight.shape}")
+    if param.pad_length:
+        weight = torch.cat((weight, weight.new_zeros(param.pad_length, *weight.shape[1:])), dim=0)
+    return weight.narrow(0, param.group.rank() * param.shape[0], param.shape[0]).contiguous()
+
+
+def _gather_gtp_weight(param: torch.Tensor | None) -> torch.Tensor | None:
+    """Gather only the GTP axis, leaving TP/EP/PP handling to the mapping."""
+    if not _is_gtp_param(param):
+        return param
+    if getattr(param, "_gtp_native_fp8", False):
+        from megatron.core.tensor_parallel.gtp_api import dequantize_gtp_native_fp8
+
+        local_weight = dequantize_gtp_native_fp8(param)
+    else:
+        # GTP's detach preserves the subclass but drops its runtime attributes.
+        # Collectives should operate on a plain tensor without triggering prefetch.
+        local_weight = param.as_subclass(torch.Tensor).detach()
+    local_weight = local_weight.contiguous()
+    padded_shape = (param.shape[0] * param.group.size(), *param.shape[1:])
+    weight = torch.empty(padded_shape, dtype=local_weight.dtype, device=local_weight.device)
+    torch.distributed.all_gather_into_tensor(weight, local_weight, group=param.group)
+    return weight[: _get_mapping_shape(param)[0]]

src/megatron/bridge/models/conversion/model_bridge.py

@@ -31,6 +31,7 @@
     Generic,
     Iterable,
     List,
+    Literal,
     Mapping,
     NamedTuple,
     Optional,
@@ -59,6 +60,13 @@
     detect_fp8_export_layout,
     get_fp8_export_tensors,
 )
+from megatron.bridge.models.conversion.gtp import (
+    _gather_gtp_weight,
+    _get_mapping_shape,
+    _gtp_weight_load_context,
+    _is_gtp_param,
+    _slice_gtp_weight,
+)
 from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
 from megatron.bridge.models.conversion.param_mapping import (
     LocalHFParam,
@@ -1529,7 +1537,8 @@ def load_weights_hf_to_megatron(
                     continue
 
                 # Check shape compatibility before copying
-                if converted_weights.shape != task.param_weight.shape:
+                expected_shape = _get_mapping_shape(task.param_weight)
+                if converted_weights.shape != expected_shape:
                     # Check whitelist
                     is_whitelisted = False
                     if allowed_mismatched_params:
@@ -1548,11 +1557,12 @@ def load_weights_hf_to_megatron(
 
                     raise ValueError(
                         f"Shape mismatch for megatron param {task.mapping.megatron_param}:\n"
-                        f"  Expected shape: {task.param_weight.shape}\n"
+                        f"  Expected shape: {expected_shape}\n"
                         f"  Got shape: {converted_weights.shape}\n"
                         f"  Bridge type: {type(task.mapping).__name__}\n"
                         f"  HF mapping: {task.mapping.hf_param}"
                     )
+                converted_weights = _slice_gtp_weight(converted_weights, task.param_weight)
                 if capture_unquantized_state_dict:
                     vp_stage = task.vp_stage if task.vp_stage is not None else 0
                     chunk_key = f"model{vp_stage}"
@@ -1562,7 +1572,7 @@ def load_weights_hf_to_megatron(
                 # Float8BlockwiseQTensor) that is a leaf with requires_grad=True.
                 # In-place updates under grad mode will raise:
                 # "a leaf Variable that requires grad is being used in an in-place operation."
-                with torch.no_grad():
+                with torch.no_grad(), _gtp_weight_load_context(task.param_weight, task.megatron_module):
                     task.param_weight.copy_(converted_weights)
         self.finalize_hf_import(megatron_model)
         if use_megatron_fsdp:
@@ -1681,6 +1691,7 @@ def stream_weights_hf_to_megatron(
             hf_weights = self.maybe_modify_loaded_hf_weight(task.mapping.hf_param, hf_state_dict)
             converted_weights = self._convert_loaded_hf_weight(task, hf_weights)
             if converted_weights is not None:
+                converted_weights = _slice_gtp_weight(converted_weights, task.param_weight)
                 # Assert that vp_stage is not None for HF->Megatron tasks
                 yield MegatronWeightTuple(task.param_name, converted_weights, task.vp_stage)
 
@@ -1800,7 +1811,7 @@ def stream_weights_megatron_to_hf(
 
                 megatron_weights = uneven_dtensor_to_full_tensor(task.param_weight)
             else:
-                megatron_weights = task.param_weight
+                megatron_weights = _gather_gtp_weight(task.param_weight)
             megatron_module = task.megatron_module
             if self._should_skip_mtp_duplicate_embedding_export(task, megatron_model):
                 megatron_weights = None
@@ -2345,7 +2356,7 @@ def _detect_fp8_params(
         Returns:
             Dictionary mapping global parameter names to FP8 export layouts.
         """
-        local_fp8_layouts: Dict[str, FP8ExportLayout] = {}
+        local_fp8_layouts: Dict[str, FP8ExportLayout | Literal[-1]] = {}
         global_name_set = set(sorted_global_param_names_all_pp_ranks)
 
         for vp_stage, model in enumerate(megatron_model):
@@ -2368,17 +2379,24 @@ def _detect_fp8_params(
                     fp8_scale_inv_attr=fp8_scale_inv_attr,
                 )
                 if layout is not None:
-                    local_fp8_layouts[global_name] = layout
+                    # Reject on every PP rank before creating raw storage views,
+                    # which do not retain the parameter's GTP sharding metadata.
+                    local_fp8_layouts[global_name] = -1 if _is_gtp_param(local_weights) else layout
 
         # Gather across PP ranks to ensure consistent insertion decisions
-        fp8_layouts_list: list[Dict[str, FP8ExportLayout]] = [None] * get_pg_size(pp_group)
+        fp8_layouts_list: list[Dict[str, FP8ExportLayout | Literal[-1]] | None] = [None] * get_pg_size(pp_group)
         torch.distributed.all_gather_object(fp8_layouts_list, local_fp8_layouts, group=pp_group)
         global_fp8_layouts: Dict[str, FP8ExportLayout] = {}
         for layouts in fp8_layouts_list:
             if not layouts:
                 continue
             # Validate only after every PP rank has completed the collective.
             for global_name, layout in layouts.items():
+                if layout == -1:
+                    raise ValueError(
+                        f"{global_name}: raw FP8 export does not support GTP sharding; "
+                        "use gathered dequantized export or gathered quantized export instead"
+                    )
                 try:
                     layout.validate()
                 except ValueError as error:

src/megatron/bridge/models/conversion/param_mapping.py

@@ -32,6 +32,7 @@
 )
 from torch.distributed._tensor import DTensor
 
+from megatron.bridge.models.conversion.gtp import _get_mapping_shape, _is_gtp_param
 from megatron.bridge.models.conversion.utils import (
     get_module_and_param_from_name,
     is_modelopt_dynamic_module,
@@ -360,6 +361,11 @@ def local_hf_params(
             )
         if isinstance(weight, DTensor) or _module_uses_fsdp(megatron_module):
             raise ValueError(f"{global_param_name}: local HF parameter iteration does not support DTensor/FSDP")
+        if _is_gtp_param(weight):
+            raise ValueError(
+                f"{global_param_name}: local HF parameter views cannot represent GTP sharding; "
+                "use the gathered HF weight export instead"
+            )
 
         specs = self.local_hf_param_specs(global_param_name)
         if not specs:
@@ -1297,7 +1303,7 @@ def hf_to_megatron(
             actual_dim0_size = hf_weights.shape[0]
             # DTensor.shape is already the global shape across TP ranks, while
             # a regular Megatron parameter stores only its local TP shard.
-            expect_dim0_size = target_param.shape[0]
+            expect_dim0_size = _get_mapping_shape(target_param)[0]
             if not isinstance(target_param, DTensor):
                 expect_dim0_size *= self.tp_size
             if actual_dim0_size != expect_dim0_size:
@@ -1321,7 +1327,7 @@ def hf_to_megatron(
         if isinstance(target_param, DTensor):
             output_shape = target_param.orig_param.shape
         else:
-            output_shape = target_param.shape
+            output_shape = _get_mapping_shape(target_param)
         # Scatter to all ranks. Each rank gets its sharded shape from its module.
         return self.scatter_to_tp_ranks(
             splits,
@@ -1485,13 +1491,14 @@ def hf_to_megatron(
 
         # HF fused expert weights (e.g. down_proj) may be stored in [in, out]
         # layout while Megatron expects [out, in]. Detect via the unsharded dim:
-        # for RowParallel, dim 0 is never split, so hf_weights.shape[0] must
-        # equal target_param.shape[0].
+        # RowParallel TP does not split dim 0. GTP may split it further in storage,
+        # so compare against the logical TP-local shape.
+        target_shape = _get_mapping_shape(target_param)
         if (
             hf_weights is not None
             and hf_weights.ndim == 2
-            and hf_weights.shape[0] != target_param.shape[0]
-            and hf_weights.shape[1] == target_param.shape[0]
+            and hf_weights.shape[0] != target_shape[0]
+            and hf_weights.shape[1] == target_shape[0]
         ):
             hf_weights = hf_weights.t().contiguous()
 
@@ -1522,7 +1529,7 @@ def hf_to_megatron(
         if isinstance(target_param, DTensor):
             output_shape = target_param.orig_param.shape
         else:
-            output_shape = target_param.shape
+            output_shape = _get_mapping_shape(target_param)
         # Scatter to all ranks. Each rank gets its sharded shape from its module.
         return self.scatter_to_tp_ranks(
             splits,
@@ -3176,7 +3183,7 @@ def hf_to_megatron(
         if isinstance(target_param, DTensor):
             output_shape = target_param.orig_param.shape
         else:
-            output_shape = target_param.shape
+            output_shape = _get_mapping_shape(target_param)
         # Scatter the concatenated shards to each rank
         return self.scatter_to_tp_ranks(
             splits,
@@ -3581,7 +3588,7 @@ def hf_to_megatron(self, hf_weights: torch.Tensor, megatron_module: nn.Module) -
 
         normalized_param = self._normalize_expert_param_name(self.megatron_param)
         _, target_param = get_module_and_param_from_name(megatron_module, normalized_param)
-        target_shape = target_param.shape
+        target_shape = _get_mapping_shape(target_param)
 
         if target_shape[0] % 2 != 0:
             raise ValueError(f"Expected even fused dim for {self.megatron_param}, got {target_shape}.")