Gtp refit support

megatron/core/resharding/README.md

@@ -1,7 +1,7 @@
 # Resharding (Refit)
 
 Transfer model weights between different parallelism configurations
-(TP, PP, EP, DP) with optional format conversion (e.g. BF16 to MXFP8).
+(TP, GTP_remat, PP, EP, DP) with optional format conversion (e.g. BF16 to MXFP8).
 Used primarily in RL loops to move weights from a training model to an
 inference model that may use a different parallelism layout.
 
@@ -12,6 +12,7 @@ refit.py            High-level API: swap_model_weights, caching, MXFP8 auto-dete
     |
 planner.py          Local plan builder (every rank all-gathers metadata, replays
                     the same deterministic schedule, keeps only its own ops)
+shard_planner.py    Logical-coordinate planner for replicated, TP, and GTP layouts
     |
 execution.py        Submits send/recv ops to a CopyService, handles writebacks
     |
@@ -88,22 +89,44 @@ through the network stack.
 
 ## How the Reshard Plan Works
 
-1. Each rank extracts parameter metadata (shape, sharding, TP/EP/PP groups).
+1. Each rank extracts parameter metadata (shape, sharding, TP/GTP_remat/EP/PP groups).
 2. Metadata is all-gathered so **every** rank has the full picture
    (`dist.all_gather_object()`) — no rank-0 bottleneck, no scatter.
 3. Every rank independently replays the **same deterministic schedule**
    (`_iter_global_transfer_ops`):
    - Iterate destination ranks, then each rank's destination params in gathered
      order; for each destination param, find the matching source param(s) by name.
-   - Route to a dimension-specific planner (LCM tiling for standard TP,
-     block-interleaved for partitioned params like Mamba `in_proj`).
+   - Map source and destination storage into logical global-weight coordinates,
+     then intersect those regions. The same algorithm handles replicated, TP,
+     packed/strided TP, GTP, and combined TP x GTP layouts.
    - Assign a monotonic `task_id` per sub-op.  Because the iteration order and
      counter are a pure function of the gathered metadata, the send op computed
      on the sender and the recv op computed on the receiver get the **same**
      `task_id` without any central authority.
 4. Each rank keeps only the ops where it is the sender or receiver.
 5. The plan is cached so repeated refits skip steps 1-4.
 
+## Generalized Tensor Parallelism
+
+See the [GTP user guide](../../../docs/api-guide/core/generalized_tensor_parallel.md)
+for the current GTP design and configuration. GTP_remat stores a contiguous
+dim-0 slice of the already TP-local weight. The shard
+planner represents every local weight as intervals in the logical, unpadded
+global weight and intersects source intervals with destination intervals.
+GTP_remat therefore uses the same algorithm as ordinary TP. This single rule covers:
+
+- column parallel weights, where TP and GTP_remat both cut dim 0;
+- row parallel weights, where TP cuts dim 1 and GTP_remat cuts dim 0;
+- strided TP weights such as gated MLP projections;
+- packed TP layouts such as Mamba `in_proj`; and
+- transitions to or from a non-GTP_remat model, or between different GTP_remat sizes.
+
+Alignment-only GTP_remat padding is excluded from transfers. Quantized destinations
+are assembled once in a zeroed BF16 buffer and requantized once, so padding
+cannot affect the scale of logical weight values. Native TransformerEngine
+MXFP8/NVFP4 GTP_remat parameters are temporarily presented as their base TE class
+while dequantizing or updating them.
+
 The deterministic schedule stays stable when a larger roster is supplied: existing
 transfers keep their `task_id`s and newly appended destination ranks receive new
 ones. Live process-group membership changes and their orchestration remain future
@@ -134,7 +157,7 @@ across refits.
 | Cache | Key | Contents | Why |
 |-------|-----|----------|-----|
 | `_service_cache` | Backend name | `CopyService` instance | Avoid re-creating CUDA streams / NVSHMEM buffers |
-| `_plan_cache` | (rank, src_config, dst_config, num_experts) | `ReshardPlan` + attached transform | Avoid collective plan rebuild on repeated refits |
+| `_plan_cache` | (rank, src_config, dst_config, num_experts) | `ReshardPlan` + attached transform | Avoid collective plan rebuild on repeated refits; configs include dense/expert GTP_remat sizes |
 
 Call `clear_all_caches()` before destroying distributed process groups
 to avoid stale references.  This also finalizes NVSHMEM resources.
@@ -151,13 +174,16 @@ attribute with the following groups:
 | `pp` | If PP > 1 | Pipeline stage / layer index remapping |
 | `ep` | If MoE | Expert parallelism routing |
 | `expt_tp` | If expert TP | Expert-specific tensor parallelism |
+| `gtp_remat` | If dense GTP_remat | Dense weight-rematerialization shards |
+| `expt_gtp_remat` | If expert GTP_remat | Expert weight-rematerialization shards |
 
 ## File Reference
 
 | File | Role |
 |------|------|
 | `refit.py` | Public API, caching, MXFP8 auto-detection |
-| `planner.py` | Local deterministic plan builder (metadata, LCM/block-interleaved planners) |
+| `planner.py` | Local deterministic schedule builder |
+| `shard_planner.py` | Replicated/TP/GTP local-to-global mapping and intersection planner |
 | `execution.py` | Plan executor (send/recv submission, writeback, format conversion) |
 | `transforms.py` | `ReshardTransform` base class, `MXFP8ReshardTransform` |
 | `utils.py` | `TransferOp`, `ReshardPlan`, `ParameterMetadata`, `ShardingDescriptor` |

megatron/core/resharding/execution.py

@@ -2,13 +2,15 @@
 from __future__ import annotations
 
 import logging
+from contextlib import nullcontext
 from dataclasses import dataclass
 from typing import Optional
 
 import torch
 import torch.distributed as dist
 
-from megatron.core.fp8_utils import is_mxfp8tensor
+from megatron.core.fp8_utils import is_float8tensor, is_mxfp8tensor
+from megatron.core.tensor_parallel import gtp_api
 
 from .copy_services.base import CopyService
 from .transforms import ReshardTransform, _ensure_sendable
@@ -24,7 +26,7 @@ class _Writeback:
     Exactly one of the three kinds applies; the other fields are unused for
     that kind.  ``direct`` means the data landed in its final destination
     during recv and there's nothing to copy.  ``copy`` copies a staging
-    ``recv_buffer`` into a slice of ``dst_param`` (deferring to MXFP8
+    ``recv_buffer`` into a slice of ``dst_param`` (deferring to quantized
     accumulation when the dest is quantized).  ``transform`` hands the
     received buffers to a ``ReshardTransform.finalize_recv`` call.
     """
@@ -37,22 +39,54 @@ class _Writeback:
     recv_bufs: Optional[list[torch.Tensor]] = None
 
 
-def _get_mxfp8_accumulator(
-    pending: dict[int, tuple], dst_param: torch.Tensor
-) -> tuple[torch.Tensor, list]:
-    """Get or lazily allocate the BF16 accumulation buffer for an MXFP8 dest param.
+def _requires_bf16_staging(param: torch.Tensor) -> bool:
+    """Return whether refit must materialize this parameter in BF16.
+
+    Quantized source storage is dequantized before slicing for transfer. On the
+    destination, updating quantized storage slice-by-slice is unsafe because
+    scale blocks can cross slice boundaries, so refit assembles the complete
+    local BF16 weight and quantizes it once after all receives finish.
+    """
+    is_gtp = gtp_api.HAVE_GTP and gtp_api.is_gtp_param(param)
+    # Ordinary params retain the existing MXFP8 path. GTP also accepts TE's
+    # generic quantized tensor type, which includes native NVFP4 weights.
+    return is_mxfp8tensor(param) or (is_gtp and is_float8tensor(param))
+
+
+def _get_quantized_accumulator(pending: dict[int, tuple], dst_param: torch.Tensor) -> torch.Tensor:
+    """Get or lazily allocate the BF16 accumulation buffer for a quantized destination.
 
     All slices for the same dst_param land in this buffer; ``quantize_`` is
-    called once after all slices have been written.  Allocates empty (not
-    dequantized) because every slice will be overwritten.
+    called once after all slices have been written.
     """
     param_id = id(dst_param)
     entry = pending.get(param_id)
     if entry is None:
-        full_bf16 = torch.empty(dst_param.shape, dtype=torch.bfloat16, device=dst_param.device)
-        entry = (dst_param, full_bf16, [])
+        has_gtp_padding = bool(
+            gtp_api.HAVE_GTP
+            and gtp_api.is_gtp_param(dst_param)
+            and getattr(dst_param, "pad_length", 0)
+        )
+        # GTP padding receives no data. Zero it so uninitialized values cannot
+        # distort a quantization block shared with logical weight values.
+        allocate = torch.zeros if has_gtp_padding else torch.empty
+        full_bf16 = allocate(dst_param.shape, dtype=torch.bfloat16, device=dst_param.device)
+        entry = (dst_param, full_bf16)
         pending[param_id] = entry
-    return entry[1], entry[2]
+    return entry[1]
+
+
+def _native_gtp_load_context(module: torch.nn.Module | None, pending: dict[int, tuple]):
+    """Return the special update context required by native quantized GTP params."""
+    if not gtp_api.HAVE_GTP:
+        return nullcontext()
+
+    if module is None or not any(
+        gtp_api.is_gtp_param(param) for param, _buffer in pending.values()
+    ):
+        return nullcontext()
+
+    return gtp_api.gtp_native_fp8_load_context(module)
 
 
 def execute_reshard_plan(
@@ -86,26 +120,26 @@ def execute_reshard_plan(
     src_params = get_refit_tensor_dict(src_module) if src_module is not None else {}
     dst_params = get_refit_tensor_dict(dst_module) if dst_module is not None else {}
 
-    # Cache dequantized BF16 views of MXFP8 source params so that multiple
+    # Cache dequantized BF16 views of quantized source params so that multiple
     # send ops for the same param reuse one dequant instead of repeating it.
     # Issue all dequants on a side stream and record per-param events so each
     # send op only waits on its own dequant (later dequants can overlap with
     # earlier sends' slicing on the default stream).
     sendable_cache: dict[str, torch.Tensor] = {}
     sendable_events: dict[str, torch.cuda.Event] = {}
 
-    mxfp8_param_names: set[str] = set()
+    quantized_param_names: set[str] = set()
     for op in plan.send_ops:
         if transform is not None and transform.should_transform(op.param_name):
             continue
         src_param = src_params.get(op.param_name)
-        if src_param is not None and is_mxfp8tensor(src_param):
-            mxfp8_param_names.add(op.param_name)
+        if src_param is not None and _requires_bf16_staging(src_param):
+            quantized_param_names.add(op.param_name)
 
-    if mxfp8_param_names:
+    if quantized_param_names:
         prefetch_stream = torch.cuda.Stream()
         with torch.cuda.stream(prefetch_stream):
-            for param_name in mxfp8_param_names:
+            for param_name in quantized_param_names:
                 sendable_cache[param_name] = _ensure_sendable(src_params[param_name])
                 ev = torch.cuda.Event()
                 ev.record()
@@ -138,9 +172,9 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor:
     sendable_events.clear()
 
     writebacks: list[_Writeback] = []
-    # Maps id(dst_param) -> (dst_param, full_bf16, slices) for MXFP8 dests that
-    # need deferred quantize_() after all slices are written.
-    pending_quantized: dict[int, tuple[torch.nn.Parameter, torch.Tensor, list]] = {}
+    # Quantized destinations are assembled in BF16 and quantized once all
+    # logical slices have arrived.
+    pending_quantized: dict[int, tuple[torch.Tensor, torch.Tensor]] = {}
 
     for op in plan.recv_ops:
         if transform is not None and transform.should_transform(op.param_name):
@@ -166,17 +200,17 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor:
         # the slice view is already contiguous AND the parameter is a plain
         # tensor (not quantized — quantized params need deferred accumulation).
         dst_slice_view = dst_param.data[op.my_slice]
-        dst_is_mxfp8 = is_mxfp8tensor(dst_param)
+        dst_requires_staging = _requires_bf16_staging(dst_param)
 
-        if not dst_is_mxfp8 and dst_slice_view.is_contiguous():
+        if not dst_requires_staging and dst_slice_view.is_contiguous():
             service.submit_recv(dst_slice_view, op.peer_rank, task_id=op.task_id)
             writebacks.append(_Writeback(kind='direct'))
             continue
 
-        if dst_is_mxfp8:
-            # TE MXFP8: recv directly into pre-allocated BF16 accumulation
+        if dst_requires_staging:
+            # Quantized parameter: recv directly into pre-allocated BF16 accumulation
             # buffer to avoid per-slice BF16 allocations.
-            full_bf16, _slices = _get_mxfp8_accumulator(pending_quantized, dst_param)
+            full_bf16 = _get_quantized_accumulator(pending_quantized, dst_param)
             accum_view = full_bf16[op.my_slice]
             if accum_view.is_contiguous():
                 service.submit_recv(accum_view, op.peer_rank, task_id=op.task_id)
@@ -202,10 +236,9 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor:
     #
     # For quantized destination params (fp8_param=true on receiver),
     # accumulate ALL BF16 slices per-param before calling quantize_() once.
-    # This avoids corrupting MXFP8 per-block scales from partial-slice updates.
-    # Since refit overwrites every slice of each param, we allocate a fresh
-    # BF16 buffer (torch.empty) instead of dequantizing the existing MXFP8
-    # weights — this avoids a full-model-sized dequantize+clone.
+    # This avoids corrupting block scales through partial updates. A zeroed
+    # buffer also gives padded GTP rows neutral values without dequantizing the
+    # old weight.
     for i in range(len(writebacks)):
         wb = writebacks[i]
         writebacks[i] = None  # Drop reference eagerly so recv buffers can free.
@@ -215,19 +248,18 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor:
             if wb.kind == 'transform':
                 transform.finalize_recv(wb.param_name, wb.dst_slice, wb.recv_bufs)
                 continue
-            # 'copy' — direct buffer copy, with deferred MXFP8 accumulation if needed.
-            if is_mxfp8tensor(wb.dst_param):
-                full_bf16, slices = _get_mxfp8_accumulator(pending_quantized, wb.dst_param)
-                slices.append((wb.dst_slice, wb.recv_buffer))
+            # 'copy' — direct buffer copy, with deferred quantization if needed.
+            if _requires_bf16_staging(wb.dst_param):
+                full_bf16 = _get_quantized_accumulator(pending_quantized, wb.dst_param)
                 full_bf16[wb.dst_slice].copy_(wb.recv_buffer)
             else:
                 wb.dst_param.data[wb.dst_slice].copy_(wb.recv_buffer)
     writebacks.clear()
 
     # Finalize deferred quantized param updates.
-    had_mxfp8_staging = bool(pending_quantized)
-    for _param_id, (dst_param, full_bf16, _slices) in pending_quantized.items():
-        with torch.no_grad():
+    had_quantized_staging = bool(pending_quantized)
+    with _native_gtp_load_context(dst_module, pending_quantized), torch.no_grad():
+        for dst_param, full_bf16 in pending_quantized.values():
             dst_param.quantize_(full_bf16)
     pending_quantized.clear()
 
@@ -240,9 +272,9 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor:
 
     # Release transient BF16 recv/accumulation buffers back to the CUDA driver.
     # Without this the caching allocator retains the peak allocation, which can
-    # be significant for MXFP8 destinations (full model weight size in BF16).
-    # Skip the (expensive) empty_cache walk when no MXFP8 staging happened.
-    if had_mxfp8_staging:
+    # be significant for quantized destinations (full model weight size in BF16).
+    # Skip the (expensive) empty_cache walk when no staging happened.
+    if had_quantized_staging:
         torch.cuda.empty_cache()
 
     logger.info("Reshard complete")