diff --git a/megatron/core/resharding/README.md b/megatron/core/resharding/README.md index 134e9c9225c..4d0a29605d7 100644 --- a/megatron/core/resharding/README.md +++ b/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,15 +89,16 @@ 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** @@ -104,6 +106,27 @@ through the network stack. 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` | diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py index ce45f124b58..02632a80fc9 100644 --- a/megatron/core/resharding/execution.py +++ b/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,7 +120,7 @@ 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 @@ -94,18 +128,18 @@ def execute_reshard_plan( 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") diff --git a/megatron/core/resharding/planner.py b/megatron/core/resharding/planner.py index f89f39a4a0c..22a641ddb49 100644 --- a/megatron/core/resharding/planner.py +++ b/megatron/core/resharding/planner.py @@ -2,19 +2,17 @@ from __future__ import annotations import logging -import math import warnings import torch import torch.distributed as dist +from .shard_planner import plan_sharded_transfer from .utils import ( ParameterMetadata, ReshardPlan, - ShardingDescriptor, TransferOp, _build_layer_module_prefix_map, - _get_rank_in_group, extract_param_metadata, named_refit_tensors, select_src_metadata_balanced, @@ -23,285 +21,6 @@ logger = logging.getLogger(__name__) -def _sort_ops_by_dst_offset(ops, dim): - """Sort transfer ops by destination offset on the sharded dimension.""" - ops.sort(key=lambda op: op[2][dim].start if isinstance(op[2][dim], slice) else 0) - - -def _build_descriptors_for_param( - src_metadata: ParameterMetadata, dst_metadata: ParameterMetadata -) -> list[ShardingDescriptor]: - """Construct sharding descriptors (currently TP) for this parameter based on actual layout. - Guard TP descriptor with size conservation so we don't mis-classify replicated tensors. - """ - descriptors: list[ShardingDescriptor] = [] - - # TP descriptor: allow when either side participates in TP - if src_metadata.is_tp or dst_metadata.is_tp: - # Prefer destination partition_dim, else source - tp_dim = dst_metadata.partition_dim if dst_metadata.is_tp else src_metadata.partition_dim - src_tp_ranks = src_metadata.tensor_parallel_group_ranks - dst_tp_ranks = dst_metadata.tensor_parallel_group_ranks - if src_tp_ranks is None or dst_tp_ranks is None: - # Not enough context to build TP descriptor - return descriptors - src_stride = src_metadata.partition_stride if src_metadata.is_tp else 1 - dst_stride = dst_metadata.partition_stride if dst_metadata.is_tp else 1 - - # Size conservation check on partition dim - src_world = len(src_tp_ranks) - dst_world = len(dst_tp_ranks) - src_local = src_metadata.shape[tp_dim] - dst_local = dst_metadata.shape[tp_dim] - if src_world * src_local != dst_world * dst_local: - raise RuntimeError( - f"Cannot build TP descriptor for {dst_metadata.name} dim{tp_dim}: " - f"src_world*src_local={src_world}*{src_local} != {dst_world}*{dst_local}. " - "This usually means the param is marked TP but is effectively replicated on that " - "dim or partition_dim/metadata is inconsistent between source and destination." - ) - - descriptors.append( - ShardingDescriptor( - name="tp", - dim=tp_dim, - src_stride=src_stride, - dst_stride=dst_stride, - src_dim_ranks=src_tp_ranks, - dst_dim_ranks=dst_tp_ranks, - ) - ) - return descriptors - - -def _emit_lcm_block_ops( - *, - param_name: str, - src_shape: tuple[int, ...], - dst_shape: tuple[int, ...], - dim: int, - src_world: int, - dst_world: int, - src_stride: int, - dst_stride: int, - full_block_len: int, - dst_local_rank: int, - src_dim_ranks: list[int], - src_block_offset: int, - dst_block_offset: int, - block_label: str, - ops: list, -) -> None: - """Emit (src_rank, src_slice, dst_slice) ops for one LCM-tiled block. - - Used both by the single-block stride-aware TP planner and by the - per-block loop of the block-interleaved planner. - """ - Ns = src_world * max(1, src_stride) - Nd = dst_world * max(1, dst_stride) - L = math.lcm(Ns, Nd) - if full_block_len % L != 0: - raise RuntimeError( - f"{param_name}: {block_label} length {full_block_len} not divisible by LCM {L} " - f"(Ns={Ns}, Nd={Nd})" - ) - unit = full_block_len // L - cps = L // Ns - cpd = L // Nd - seg_src = cps * unit - seg_dst = cpd * unit - - for k in range(max(1, dst_stride)): - g_dst_seg = dst_local_rank + k * dst_world - for off in range(cpd): - g_micro = g_dst_seg * cpd + off - s_idx = g_micro // cps - in_seg = g_micro % cps - src_global_rank = src_dim_ranks[s_idx % src_world] - src_local_seg_idx = s_idx // src_world - src_start = src_block_offset + src_local_seg_idx * seg_src + in_seg * unit - dst_start = dst_block_offset + k * seg_dst + off * unit - src_slice = [slice(None)] * len(src_shape) - dst_slice = [slice(None)] * len(dst_shape) - src_slice[dim] = slice(src_start, src_start + unit) - dst_slice[dim] = slice(dst_start, dst_start + unit) - ops.append((src_global_rank, tuple(src_slice), tuple(dst_slice))) - - -def _tp_block_layout( - param_name: str, - src_metadata: ParameterMetadata, - dst_metadata: ParameterMetadata, - descriptor: ShardingDescriptor, - src_shape: tuple[int, ...], - dst_shape: tuple[int, ...], -) -> list[tuple[int, int, int, int, int, str]]: - """Compute the per-block layout for a TP transfer. - - Returns a list of ``(src_offset, dst_offset, full_block_len, src_stride, - dst_stride, label)`` tuples that the LCM micro-tiler iterates. - - - Plain TP (no ``partition_sizes``): single block covering the full - partition dim with the descriptor's strides. - - Block-interleaved TP (``partition_sizes`` present, e.g. Mamba ``in_proj``): - one block per packed component, each independently sharded with stride=1. - """ - d = descriptor - dim = d.dim - src_world = len(d.src_dim_ranks) - dst_world = len(d.dst_dim_ranks) - src_sizes = src_metadata.partition_sizes - dst_sizes = dst_metadata.partition_sizes - - if src_sizes is None and dst_sizes is None: - src_local = src_shape[dim] - dst_local = dst_shape[dim] - if src_world * src_local != dst_world * dst_local: - raise RuntimeError( - f"{param_name}: size mismatch on TP dim{dim} " - f"(src_world={src_world}, src_local={src_local}, " - f"dst_world={dst_world}, dst_local={dst_local})" - ) - return [(0, 0, dst_local * dst_world, d.src_stride, d.dst_stride, f"TP dim{dim}")] - - if src_sizes is not None: - num_blocks = len(src_sizes) - full_sizes = [s * src_world for s in src_sizes] - else: - num_blocks = len(dst_sizes) - full_sizes = [s * dst_world for s in dst_sizes] - if src_sizes is None: - src_sizes = [f // src_world for f in full_sizes] - if dst_sizes is None: - dst_sizes = [f // dst_world for f in full_sizes] - - blocks: list[tuple[int, int, int, int, int, str]] = [] - src_off = 0 - dst_off = 0 - for i in range(num_blocks): - if src_sizes[i] * src_world != dst_sizes[i] * dst_world: - raise RuntimeError( - f"{param_name}: block {i} size mismatch: " - f"src_sizes[{i}]={src_sizes[i]}*{src_world} != " - f"dst_sizes[{i}]={dst_sizes[i]}*{dst_world}" - ) - blocks.append((src_off, dst_off, full_sizes[i], 1, 1, f"block {i}")) - src_off += src_sizes[i] - dst_off += dst_sizes[i] - return blocks - - -def _plan_tp( - param_name: str, - src_metadata: ParameterMetadata, - dst_metadata: ParameterMetadata, - descriptors: list[ShardingDescriptor], - my_global_rank: int, -) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: - """Plan TP transfers via LCM tiling, supporting both plain and block-interleaved TP. - - The block layout is derived once by ``_tp_block_layout`` — the inner - LCM micro-tile math (``_emit_lcm_block_ops``) is identical for both cases, - so the single-block plain-TP path is just a special case of the - multi-block partitioned path. - """ - if not descriptors: - return [] - if len(descriptors) != 1 or descriptors[0].name != "tp": - raise NotImplementedError(f"{param_name}: _plan_tp supports TP-only (one descriptor)") - d = descriptors[0] - if my_global_rank not in d.dst_dim_ranks: - return [] - - src_shape = tuple(src_metadata.shape) - dst_shape = tuple(dst_metadata.shape) - src_world = len(d.src_dim_ranks) - dst_world = len(d.dst_dim_ranks) - dst_local_rank = _get_rank_in_group(my_global_rank, d.dst_dim_ranks) - - blocks = _tp_block_layout(param_name, src_metadata, dst_metadata, d, src_shape, dst_shape) - - ops: list[tuple[int, tuple[slice, ...], tuple[slice, ...]]] = [] - for src_off, dst_off, full_len, src_stride, dst_stride, label in blocks: - _emit_lcm_block_ops( - param_name=param_name, - src_shape=src_shape, - dst_shape=dst_shape, - dim=d.dim, - src_world=src_world, - dst_world=dst_world, - src_stride=src_stride, - dst_stride=dst_stride, - full_block_len=full_len, - dst_local_rank=dst_local_rank, - src_dim_ranks=d.src_dim_ranks, - src_block_offset=src_off, - dst_block_offset=dst_off, - block_label=label, - ops=ops, - ) - _sort_ops_by_dst_offset(ops, d.dim) - return ops - - -def _finalize_dp_transfers( - param_name: str, - src_metadata: ParameterMetadata, - dst_metadata: ParameterMetadata, - my_global_rank: int, -) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: - """Return receiver-side transfer for a parameter that is not TP-sharded. - - This is reached when we cannot build a TP sharding descriptor for the parameter - (i.e., it is effectively replicated with respect to sharding). We use this when the - destination and source mode have no TP or the parameter is replicted on all ranks - such as layernorm. If the source and destination DP groups match, we return a local - full-tensor copy; otherwise we pick a source rank from the source DP group in a - deterministic round-robin manner based on the receiver's global rank for better load - distribution. - """ - dst_dp_ranks = dst_metadata.data_parallel_group_ranks - src_dp_ranks = src_metadata.data_parallel_group_ranks - if my_global_rank not in dst_dp_ranks: - return [] - - dst_shape = dst_metadata.shape - - # Same DP layout - local copy (only if this rank has the source parameter) - if src_dp_ranks == dst_dp_ranks and my_global_rank in src_dp_ranks: - full_slice = tuple(slice(None) for _ in range(len(dst_shape))) - return [(my_global_rank, full_slice, full_slice)] - - # Use the owner of the metadata picked by select_src_metadata_balanced. - # That selection already handles DP round-robin and non-collocated cases - # (where some src DP ranks don't actually own the source model). - full_slice = tuple(slice(None) for _ in range(len(dst_shape))) - return [(src_metadata.owner_rank, full_slice, full_slice)] - - -def _determine_source_ranks_for_dst_param( - param_name: str, - src_metadata: ParameterMetadata, - dst_metadata: ParameterMetadata, - my_global_rank: int, -) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: - """Route to dimension-specific planner based on parameter sharding type.""" - - # Regular TP/DP planning with EP-resolved metadata. _plan_tp handles both - # plain TP and block-interleaved TP (partition_sizes-driven) layouts. - descriptors = _build_descriptors_for_param(src_metadata=src_metadata, dst_metadata=dst_metadata) - if descriptors: - return _plan_tp( - param_name=param_name, - src_metadata=src_metadata, - dst_metadata=dst_metadata, - descriptors=descriptors, - my_global_rank=my_global_rank, - ) - # DP / replicated fallback - return _finalize_dp_transfers(param_name, src_metadata, dst_metadata, my_global_rank) - - def _iter_global_transfer_ops( dst_param_metadata_by_rank: dict[int, dict[str, ParameterMetadata]], src_param_metadata: dict[str, list[ParameterMetadata]], @@ -347,8 +66,8 @@ def _iter_global_transfer_ops( ) # Choose a representative source metadata with DP round-robin balancing. src_metadata = select_src_metadata_balanced(src_meta_list, dst_metadata, dst_rank) - sources = _determine_source_ranks_for_dst_param( - resolved_name, src_metadata, dst_metadata, dst_rank + sources = plan_sharded_transfer( + resolved_name, src_meta_list, src_metadata, dst_metadata ) for src_rank, src_slice, dst_slice in sources: task_id = next_task_id diff --git a/megatron/core/resharding/refit.py b/megatron/core/resharding/refit.py index 5798b48c29a..722980b66f1 100644 --- a/megatron/core/resharding/refit.py +++ b/megatron/core/resharding/refit.py @@ -9,7 +9,7 @@ """ from dataclasses import dataclass -from typing import Any, Literal, Optional, Tuple, Union +from typing import Any, Literal, NamedTuple, Optional, Union import torch @@ -34,6 +34,18 @@ RefitBackendName = Literal["nccl", "gloo", "nvshmem", "nixl"] +class _ParallelConfig(NamedTuple): + """Parallel group sizes that determine a refit plan.""" + + tp_size: int + pp_size: int + ep_size: int + dp_size: int + expert_tp_size: int + gtp_remat_size: int + expert_gtp_remat_size: int + + @dataclass(frozen=True) class _PlanCacheKey: """ @@ -41,9 +53,8 @@ class _PlanCacheKey: """ rank: int - # Parallelism configuration: (TP, PP, EP, DP, expt_tp) or None for non-collocated ranks - src_config: Optional[Tuple[int, int, int, int, int]] - dst_config: Optional[Tuple[int, int, int, int, int]] + src_config: Optional[_ParallelConfig] + dst_config: Optional[_ParallelConfig] num_experts: Optional[int] # Adding inference nodes leaves the configs and offsets unchanged, so without # world_size the stale pre-growth plan would be reused. @@ -58,8 +69,8 @@ class _PlanCacheKey: pool_index: int = 0 -def _get_config_tuple(core) -> Optional[Tuple[int, int, int, int, int]]: - """Extract (TP, PP, EP, DP, expt_tp) sizes from a model core, memoized on the core. +def _get_parallel_config(core) -> Optional[_ParallelConfig]: + """Extract TP/PP/EP/DP/expert-TP/GTP-remat sizes, memoized on the core. Process-group sizes don't change after init, so the result is cached on the core object itself to avoid repeated ``get_process_group_ranks`` calls on @@ -67,19 +78,23 @@ def _get_config_tuple(core) -> Optional[Tuple[int, int, int, int, int]]: """ if core is None: return None - cached = getattr(core, '_refit_config_tuple', None) + cached = getattr(core, '_refit_parallel_config', None) if cached is not None: return cached pg = core.pg_collection expt_tp = getattr(pg, 'expt_tp', None) - result = ( - pg.tp.size() if pg.tp else 1, - pg.pp.size() if pg.pp else 1, - pg.ep.size() if pg.ep else 1, - pg.dp.size() if pg.dp else 1, - expt_tp.size() if expt_tp else 1, + gtp_remat = getattr(pg, 'gtp_remat', None) + expt_gtp_remat = getattr(pg, 'expt_gtp_remat', None) + result = _ParallelConfig( + tp_size=pg.tp.size() if pg.tp else 1, + pp_size=pg.pp.size() if pg.pp else 1, + ep_size=pg.ep.size() if pg.ep else 1, + dp_size=pg.dp.size() if pg.dp else 1, + expert_tp_size=expt_tp.size() if expt_tp else 1, + gtp_remat_size=gtp_remat.size() if gtp_remat else 1, + expert_gtp_remat_size=expt_gtp_remat.size() if expt_gtp_remat else 1, ) - core._refit_config_tuple = result + core._refit_parallel_config = result return result @@ -99,8 +114,8 @@ def _build_plan_cache_key( world_size = group.size() if group is not None else torch.distributed.get_world_size() return _PlanCacheKey( rank=rank, - src_config=_get_config_tuple(src_core), - dst_config=_get_config_tuple(tgt_core), + src_config=_get_parallel_config(src_core), + dst_config=_get_parallel_config(tgt_core), num_experts=num_experts, world_size=world_size, src_rank_offset=src_rank_offset, @@ -192,7 +207,7 @@ def _unwrap_model_cores(src_model, target_model): raise RuntimeError("Source model missing pg_collection required for reshard") # Fill missing DP group on the source using Megatron's parallel state if not provided if getattr(src_core.pg_collection, "dp", None) is None: - src_core.pg_collection.dp = parallel_state.get_data_parallel_group() + src_core.pg_collection.dp = parallel_state.get_data_parallel_group(with_gtp_remat=False) if target_model is not None: tgt_lm = target_model[0] if isinstance(target_model, (list, tuple)) else target_model diff --git a/megatron/core/resharding/shard_planner.py b/megatron/core/resharding/shard_planner.py new file mode 100644 index 00000000000..d8eb13b95e2 --- /dev/null +++ b/megatron/core/resharding/shard_planner.py @@ -0,0 +1,255 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Transfer planning in logical weight coordinates.""" + +import math +from itertools import product +from typing import NamedTuple + +from .utils import ParameterMetadata, _get_rank_in_group + + +class _Segment(NamedTuple): + """A contiguous local interval and its position in the logical global weight.""" + + local_start: int + global_start: int + length: int + + +class _GTPShardLayout(NamedTuple): + """Location of one stored dimension in its unpadded TP-local layout. + + ``tp_local_size`` is the size after GTP shards are joined and padding is + removed, but before TP shards are joined. For TP=2, GTP=2, stored dim-0 + size 3, and two padding rows (P): + + GTP: TP 0 [a b c] + [d P P] -> [a b c d] + TP 1 [e f g] + [h P P] -> [e f g h] + TP: [a b c d] + [e f g h] -> [a b c d e f g h] + + Each stored shard has size 3. Joining GTP gives ``tp_local_size = + 3 * 2 - 2 = 4``; joining TP then gives 8 global rows. + """ + + tp_local_start: int + tp_local_stop: int + tp_local_size: int + + +def _gtp_shard_layout(metadata: ParameterMetadata, dim: int) -> _GTPShardLayout: + """Describe this rank's GTP shard within the TP-local dimension.""" + stored_size = metadata.shape[dim] + if not metadata.is_gtp or dim != 0: + return _GTPShardLayout(0, stored_size, stored_size) + + group = metadata.gtp_remat_group_ranks + if not group: + raise RuntimeError(f"{metadata.name}: missing GTP rematerialization group") + + tp_local_size = stored_size * len(group) - metadata.gtp_pad_length + if tp_local_size <= 0: + raise RuntimeError( + f"{metadata.name}: invalid GTP padding ({metadata.gtp_pad_length}) " + f"for dim 0 size {stored_size} and group size {len(group)}" + ) + + gtp_rank = _get_rank_in_group(metadata.owner_rank, group) + tp_local_start = gtp_rank * stored_size + tp_local_stop = min(tp_local_start + stored_size, tp_local_size) + return _GTPShardLayout(tp_local_start, tp_local_stop, tp_local_size) + + +def _tp_segments(metadata: ParameterMetadata, dim: int, tp_local_size: int) -> list[_Segment]: + """Map a TP-local dimension to the logical global weight.""" + if not metadata.is_tp or metadata.partition_dim != dim: + return [_Segment(0, 0, tp_local_size)] + + group = metadata.tensor_parallel_group_ranks + if not group: + raise RuntimeError(f"{metadata.name}: missing tensor-parallel group") + + tp_rank = _get_rank_in_group(metadata.owner_rank, group) + tp_size = len(group) + + if metadata.partition_sizes is not None: + if sum(metadata.partition_sizes) != tp_local_size: + raise RuntimeError( + f"{metadata.name}: partition_sizes sum to {sum(metadata.partition_sizes)}, " + f"expected TP-local size {tp_local_size}" + ) + + segments = [] + local_offset = 0 + global_offset = 0 + for block_size in metadata.partition_sizes: + segments.append( + _Segment(local_offset, global_offset + tp_rank * block_size, block_size) + ) + local_offset += block_size + global_offset += block_size * tp_size + return segments + + stride = max(1, metadata.partition_stride) + if tp_local_size % stride: + raise RuntimeError( + f"{metadata.name}: TP-local size {tp_local_size} is not divisible by " + f"partition_stride={stride}" + ) + + segment_size = tp_local_size // stride + global_block_size = segment_size * tp_size + return [ + _Segment( + stride_idx * segment_size, + stride_idx * global_block_size + tp_rank * segment_size, + segment_size, + ) + for stride_idx in range(stride) + ] + + +def _local_segments(metadata: ParameterMetadata, dim: int) -> list[_Segment]: + """Map local storage to the unpadded logical global weight. + + GTP always takes a contiguous dim-0 slice of the TP-local layout. Intersecting + that slice with the TP segments naturally handles column, row, strided, and + packed tensor-parallel layouts. + """ + gtp_layout = _gtp_shard_layout(metadata, dim) + segments = [] + for tp_segment in _tp_segments(metadata, dim, gtp_layout.tp_local_size): + tp_segment_stop = tp_segment.local_start + tp_segment.length + tp_local_start = max(gtp_layout.tp_local_start, tp_segment.local_start) + tp_local_stop = min(gtp_layout.tp_local_stop, tp_segment_stop) + if tp_local_start < tp_local_stop: + segments.append( + _Segment( + tp_local_start - gtp_layout.tp_local_start, + tp_segment.global_start + tp_local_start - tp_segment.local_start, + tp_local_stop - tp_local_start, + ) + ) + return segments + + +def _global_shape(metadata: ParameterMetadata) -> tuple[int, ...]: + """Return the unpadded shape after materializing TP and GTP.""" + shape = [] + for dim in range(len(metadata.shape)): + tp_local_size = _gtp_shard_layout(metadata, dim).tp_local_size + if metadata.is_tp and metadata.partition_dim == dim: + group = metadata.tensor_parallel_group_ranks + if not group: + raise RuntimeError(f"{metadata.name}: missing tensor-parallel group") + tp_local_size *= len(group) + shape.append(tp_local_size) + return tuple(shape) + + +def _source_shards( + all_src_metadata: list[ParameterMetadata], selected: ParameterMetadata +) -> list[ParameterMetadata]: + """Find the TP x GTP shard grid containing the selected source replica. + + Walking both groups transitively finds the grid without assuming a + particular global rank layout. + """ + by_rank = {metadata.owner_rank: metadata for metadata in all_src_metadata} + pending = [selected.owner_rank] + ranks = {selected.owner_rank} + + while pending: + metadata = by_rank[pending.pop()] + groups = [] + if metadata.is_tp: + groups.append(metadata.tensor_parallel_group_ranks) + if metadata.is_gtp: + groups.append(metadata.gtp_remat_group_ranks) + for group in groups: + for rank in group or (): + if rank in by_rank and rank not in ranks: + ranks.add(rank) + pending.append(rank) + + return [by_rank[rank] for rank in sorted(ranks)] + + +def _intersect_segments( + src_segments: list[_Segment], dst_segments: list[_Segment] +) -> list[tuple[slice, slice]]: + """Intersect two segment lists in logical global coordinates.""" + overlaps = [] + for src in src_segments: + for dst in dst_segments: + start = max(src.global_start, dst.global_start) + stop = min(src.global_start + src.length, dst.global_start + dst.length) + if start < stop: + overlaps.append( + ( + slice( + src.local_start + start - src.global_start, + src.local_start + stop - src.global_start, + ), + slice( + dst.local_start + start - dst.global_start, + dst.local_start + stop - dst.global_start, + ), + ) + ) + return overlaps + + +def _rectangles_overlap(left: tuple[slice, ...], right: tuple[slice, ...]) -> bool: + """Return whether two slice rectangles cover any common element.""" + return all( + max(left_part.start, right_part.start) < min(left_part.stop, right_part.stop) + for left_part, right_part in zip(left, right) + ) + + +def plan_sharded_transfer( + param_name: str, + all_src_metadata: list[ParameterMetadata], + selected_src: ParameterMetadata, + dst_metadata: ParameterMetadata, +) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: + """Plan a transfer by intersecting source and destination logical shards.""" + src_shards = _source_shards(all_src_metadata, selected_src) + dst_shape = _global_shape(dst_metadata) + dst_segments = [_local_segments(dst_metadata, dim) for dim in range(len(dst_metadata.shape))] + + expected_elements = math.prod( + sum(segment.length for segment in segments) for segments in dst_segments + ) + + transferred_elements = 0 + ops = [] + for src in src_shards: + src_shape = _global_shape(src) + if src_shape != dst_shape: + raise RuntimeError( + f"{param_name}: logical shape mismatch: source rank {src.owner_rank} " + f"has {src_shape}, destination rank {dst_metadata.owner_rank} has {dst_shape}" + ) + + overlaps_by_dim = [ + _intersect_segments(_local_segments(src, dim), dst_segments[dim]) + for dim in range(len(dst_metadata.shape)) + ] + for rectangle in product(*overlaps_by_dim): + src_slice = tuple(slices[0] for slices in rectangle) + dst_slice = tuple(slices[1] for slices in rectangle) + if any(_rectangles_overlap(dst_slice, op[2]) for op in ops): + raise RuntimeError(f"{param_name}: overlapping destination coverage") + transferred_elements += math.prod(part.stop - part.start for part in dst_slice) + ops.append((src.owner_rank, src_slice, dst_slice)) + + if transferred_elements != expected_elements: + raise RuntimeError( + f"{param_name}: covered {transferred_elements} of {expected_elements} " + "destination elements " + f"from source ranks {[metadata.owner_rank for metadata in src_shards]}" + ) + + return sorted(ops, key=lambda op: tuple(part.start for part in op[2])) diff --git a/megatron/core/resharding/transforms.py b/megatron/core/resharding/transforms.py index de69a8d8c94..9a839ed49bc 100644 --- a/megatron/core/resharding/transforms.py +++ b/megatron/core/resharding/transforms.py @@ -11,8 +11,9 @@ import torch -from megatron.core.fp8_utils import dequantize_fp8_tensor, is_mxfp8tensor +from megatron.core.fp8_utils import dequantize_fp8_tensor, is_float8tensor, is_mxfp8tensor from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor +from megatron.core.tensor_parallel import gtp_api class ReshardTransform: @@ -105,6 +106,12 @@ def _ensure_sendable(param: torch.Tensor) -> torch.Tensor: dequantized to their original precision (usually BF16). Standard parameters are returned via ``.data`` (unwrapped from autograd). """ + if gtp_api.HAVE_GTP and gtp_api.is_gtp_param(param) and is_float8tensor(param): + # Native quantized GTP parameters use a dynamic GTP_ + # subclass. TransformerEngine dispatches dequantization on the exact + # base class, so use GTP's temporary reclassification helper. Despite + # its name, the helper handles both native FP8 and NVFP4 parameters. + return gtp_api.dequantize_gtp_native_fp8(param) if is_mxfp8tensor(param): return dequantize_fp8_tensor(param) return param.data diff --git a/megatron/core/resharding/utils.py b/megatron/core/resharding/utils.py index 94c5767c217..156dd885736 100644 --- a/megatron/core/resharding/utils.py +++ b/megatron/core/resharding/utils.py @@ -53,6 +53,14 @@ class ParameterMetadata: # interleaves these blocks rather than doing a simple contiguous concat. partition_sizes: list[int] | None = None + # GTP always shards dim 0 after any TP-local layout has been formed. + is_gtp: bool = False + # Ordered global ranks that own contiguous dim-0 shards. The list position + # determines each owner's GTP rank and therefore its shard offset. + gtp_remat_group_ranks: list[int] | None = None + # Alignment-only rows at the tail of the TP-local layout. + gtp_pad_length: int = 0 + # EP sharding info (fused/grouped MoE) is_ep: bool = False num_experts: Optional[int] = None @@ -87,7 +95,7 @@ class ParameterMetadata: @dataclass class ShardingDescriptor: - """Descriptor for a sharded dimension for a parameter.""" + """Legacy sharding descriptor kept for import compatibility.""" name: str # "tp" | "ep" | custom label dim: int @@ -334,6 +342,12 @@ def extract_param_metadata( if partition_sizes is not None: partition_sizes = list(partition_sizes) + # GTP parameters carry their actual dense/expert rematerialization group + # directly. Prefer that over selecting a group by parameter name or model + # type: it is authoritative for both GTP and expert GTP. + is_gtp = bool(getattr(param, 'is_gtp_weight_remat', False)) + gtp_pad_length = int(getattr(param, 'pad_length', 0)) if is_gtp else 0 + # EP detection: Megatron convention - expert params are not allreduced is_ep = not bool(getattr(param, 'allreduce', True)) @@ -347,6 +361,7 @@ def extract_param_metadata( ) tensor_parallel_group_ranks: list[int] | None = None + gtp_remat_group_ranks: list[int] | None = None expert_parallel_group_ranks: list[int] | None = None data_parallel_group_ranks: list[int] | None = None pipeline_parallel_group_ranks: list[int] | None = None @@ -368,6 +383,12 @@ def _offset_ranks(ranks: list[int]) -> list[int]: result = [r + rank_offset for r in ranks] if rank_offset else ranks return _dedup_ranks(result) + if is_gtp: + gtp_group = getattr(param, 'group', None) + if gtp_group is None: + raise ValueError(f"GTP parameter {param_name!r} is missing its rematerialization group") + gtp_remat_group_ranks = _offset_ranks(dist.get_process_group_ranks(gtp_group)) + if is_ep or is_expert_param: if is_ep: expert_parallel_group_ranks = _offset_ranks( @@ -424,6 +445,9 @@ def _offset_ranks(ranks: list[int]) -> list[int]: partition_dim=partition_dim, partition_stride=partition_stride, partition_sizes=partition_sizes, + is_gtp=is_gtp, + gtp_remat_group_ranks=gtp_remat_group_ranks, + gtp_pad_length=gtp_pad_length, is_ep=is_ep, num_experts=num_experts, owner_rank=owner_rank, @@ -453,7 +477,7 @@ def _filter_by_ep_local_rank( filter ensures dst EP local 0 uses src EP local 0 (same global experts). - Different size (EP=8→EP=16): dst EP local 8 has no corresponding src EP local → skip filter; expert reassignment is handled by resolved_name - matching, and the LCM/TP planner handles any TP dimension changes. + matching, and the shard planner handles any TP dimension changes. """ dst_ep_group = dst_metadata.expert_parallel_group_ranks if dst_ep_group is None: @@ -534,9 +558,9 @@ def select_src_metadata_balanced( ) -> ParameterMetadata: """Choose a representative source `ParameterMetadata` for a destination rank. - The selected metadata supplies topology (TP/EP/DP group ranks) to the LCM - planner. Selection prefers a local copy when ``dst_rank`` itself owns a - source replica, then round-robins across source DP groups to balance load. + The selected metadata identifies one complete source replica. Selection + prefers a local copy when ``dst_rank`` itself owns a source replica, then + round-robins across source DP groups to balance load. A local copy is essentially free (``tensor.copy_()`` on same GPU), while any remote transfer incurs significant overhead even within the same node. """ diff --git a/tests/unit_tests/resharding/test_planner.py b/tests/unit_tests/resharding/test_planner.py index fb8fcd6e9f9..7151b0f7cf3 100644 --- a/tests/unit_tests/resharding/test_planner.py +++ b/tests/unit_tests/resharding/test_planner.py @@ -1,29 +1,17 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -"""Unit tests for the resharding planner functions. - -These test the TP planner (_plan_tp, covering both plain and block-interleaved -layouts), DP fallback (_finalize_dp_transfers), and descriptor building in -isolation without requiring distributed init or GPU. -""" +"""Unit tests for logical-coordinate reshard planning.""" import math +from itertools import product import pytest +import torch import megatron.core.resharding.planner as planner -from megatron.core.resharding.planner import ( - _build_descriptors_for_param, - _finalize_dp_transfers, - _plan_tp, - build_plan_from_rosters, - index_metadata_rosters, -) -from megatron.core.resharding.utils import ParameterMetadata, ShardingDescriptor - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- +from megatron.core.resharding.planner import build_plan_from_rosters, index_metadata_rosters +from megatron.core.resharding.shard_planner import plan_sharded_transfer +from megatron.core.resharding.utils import ParameterMetadata def _meta( @@ -37,11 +25,12 @@ def _meta( tp_ranks=None, dp_ranks=None, ep_ranks=None, + is_gtp=False, + gtp_ranks=None, + gtp_pad_length=0, resolved_name=None, ): - """Create a ParameterMetadata for testing.""" - import torch - + """Create parameter metadata for testing.""" return ParameterMetadata( name=name, shape=shape, @@ -51,6 +40,9 @@ def _meta( partition_dim=partition_dim, partition_stride=partition_stride, partition_sizes=partition_sizes, + is_gtp=is_gtp, + gtp_remat_group_ranks=gtp_ranks, + gtp_pad_length=gtp_pad_length, owner_rank=owner_rank, tensor_parallel_group_ranks=tp_ranks, data_parallel_group_ranks=dp_ranks, @@ -59,384 +51,412 @@ def _meta( ) -def _tp_descriptor(dim, src_ranks, dst_ranks, src_stride=1, dst_stride=1): - """Create a TP ShardingDescriptor for testing.""" - return ShardingDescriptor( - name="tp", - dim=dim, - src_stride=src_stride, - dst_stride=dst_stride, - src_dim_ranks=src_ranks, - dst_dim_ranks=dst_ranks, - ) +def _tp_metadata(shape, partition_dim, tp_ranks, **kwargs): + """Create metadata for every rank in one TP group.""" + return [ + _meta( + shape=shape, + is_tp=len(tp_ranks) > 1, + partition_dim=partition_dim, + owner_rank=rank, + tp_ranks=tp_ranks, + dp_ranks=[rank], + **kwargs, + ) + for rank in tp_ranks + ] -def _verify_full_coverage(ops, dim, expected_full_len): - """Verify that transfer ops cover every element of the destination tensor exactly once.""" +def _verify_nd_coverage(ops, expected_shape): + """Verify that destination slices cover each logical local element once.""" covered = set() for _, _, dst_slice in ops: - s = dst_slice[dim] - for i in range(s.start, s.stop): - assert i not in covered, f"Duplicate coverage at dst offset {i}" - covered.add(i) - assert covered == set(range(expected_full_len)), ( - f"Expected coverage [0, {expected_full_len}), got gaps: " - f"{set(range(expected_full_len)) - covered}" - ) - - -# =========================================================================== -# _plan_tp -# =========================================================================== + ranges = [range(part.start, part.stop) for part in dst_slice] + for index in product(*ranges): + assert index not in covered, f"Duplicate coverage at destination index {index}" + covered.add(index) + assert len(covered) == math.prod(expected_shape) + + +def _tp_global_index(local_index, tp_rank, tp_size, local_size, stride): + """Map one plain/strided TP-local index to its logical global index.""" + segment_size = local_size // stride + segment, offset = divmod(local_index, segment_size) + return segment * segment_size * tp_size + tp_rank * segment_size + offset + + +class TestLogicalShardPlanner: + """One algorithm covers replicated, TP, packed, strided, and GTP layouts.""" + + @staticmethod + def _tp2_gtp2_metadata(shape, partition_dim, **kwargs): + topology = { + 0: ([0, 1], [0, 2]), + 1: ([0, 1], [1, 3]), + 2: ([2, 3], [0, 2]), + 3: ([2, 3], [1, 3]), + } + return [ + _meta( + shape=shape, + is_tp=True, + partition_dim=partition_dim, + owner_rank=rank, + tp_ranks=tp_ranks, + dp_ranks=[rank], + is_gtp=True, + gtp_ranks=gtp_ranks, + **kwargs, + ) + for rank, (tp_ranks, gtp_ranks) in topology.items() + ] + def test_replicated_transfer(self): + src = [_meta(owner_rank=2, dp_ranks=[2, 3])] + dst = _meta(owner_rank=4, dp_ranks=[4, 5]) -class TestPlanMultiDimLcm: - """Tests for the LCM-based TP planner.""" + ops = plan_sharded_transfer("weight", src, src[0], dst) - def test_tp2_to_tp1(self): - """TP2 → TP1: destination rank 0 should receive from both source ranks.""" - # Source: TP2, each rank has shape (64, 64) on dim 1 - # Destination: TP1, rank 0 has shape (64, 128) on dim 1 - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 128), is_tp=False, tp_ranks=[0]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0]) + assert ops == [(2, (slice(0, 64), slice(0, 128)), (slice(0, 64), slice(0, 128)))] - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - assert len(ops) == 2 - # Should receive from rank 0 and rank 1 - src_ranks = {op[0] for op in ops} - assert src_ranks == {0, 1} - _verify_full_coverage(ops, dim=1, expected_full_len=128) + def test_tp2_to_unsharded(self): + src = _tp_metadata(shape=(64, 64), partition_dim=1, tp_ranks=[0, 1]) + dst = _meta(shape=(64, 128), owner_rank=2, tp_ranks=[2], dp_ranks=[2]) - def test_tp1_to_tp2(self): - """TP1 → TP2: each destination rank receives half.""" - src = _meta(shape=(64, 128), is_tp=False, tp_ranks=[0]) - dst = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - desc = _tp_descriptor(dim=1, src_ranks=[0], dst_ranks=[0, 1]) + ops = plan_sharded_transfer("weight", src, src[0], dst) - # Rank 0 receives first half - ops_r0 = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - assert len(ops_r0) == 1 - assert ops_r0[0][0] == 0 # from rank 0 - _verify_full_coverage(ops_r0, dim=1, expected_full_len=64) + assert ops == [ + (0, (slice(0, 64), slice(0, 64)), (slice(0, 64), slice(0, 64))), + (1, (slice(0, 64), slice(0, 64)), (slice(0, 64), slice(64, 128))), + ] - # Rank 1 receives second half - ops_r1 = _plan_tp("weight", src, dst, [desc], my_global_rank=1) - assert len(ops_r1) == 1 - assert ops_r1[0][0] == 0 # from rank 0 + def test_unsharded_to_tp2(self): + src = [_meta(shape=(64, 128), owner_rank=0, tp_ranks=[0], dp_ranks=[0])] + dst_group = [1, 2] + + for tp_rank, rank in enumerate(dst_group): + dst = _meta( + shape=(64, 64), + is_tp=True, + partition_dim=1, + owner_rank=rank, + tp_ranks=dst_group, + dp_ranks=[rank], + ) + ops = plan_sharded_transfer("weight", src, src[0], dst) + assert ops == [ + ( + 0, + (slice(0, 64), slice(tp_rank * 64, (tp_rank + 1) * 64)), + (slice(0, 64), slice(0, 64)), + ) + ] def test_tp2_to_tp4(self): - """TP2 → TP4: each dst rank receives from 1 or more src ranks.""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 32), is_tp=True, partition_dim=1, tp_ranks=[0, 1, 2, 3]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0, 1, 2, 3]) - - for rank in range(4): - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=rank) - assert len(ops) >= 1 - _verify_full_coverage(ops, dim=1, expected_full_len=32) - - def test_same_tp_size(self): - """TP2 → TP2: each dst rank receives from its corresponding src rank.""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0, 1]) - - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - assert len(ops) == 1 - assert ops[0][0] == 0 # from self - - def test_rank_not_in_dst(self): - """Rank not in destination group returns empty ops.""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 128), is_tp=False, tp_ranks=[2]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[2]) - - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - assert ops == [] - - def test_empty_descriptors(self): - """No descriptors returns empty ops.""" - src = _meta() - dst = _meta() - ops = _plan_tp("weight", src, dst, [], my_global_rank=0) - assert ops == [] - - def test_size_mismatch_raises(self): - """Mismatched TP dimensions should raise.""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 100), is_tp=False, tp_ranks=[0]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0]) - - with pytest.raises(RuntimeError, match="size mismatch"): - _plan_tp("weight", src, dst, [desc], my_global_rank=0) - - def test_dim0_partition(self): - """TP on dimension 0 (row-parallel).""" - src = _meta(shape=(32, 128), is_tp=True, partition_dim=0, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 128), is_tp=False, tp_ranks=[0]) - desc = _tp_descriptor(dim=0, src_ranks=[0, 1], dst_ranks=[0]) - - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - assert len(ops) == 2 - _verify_full_coverage(ops, dim=0, expected_full_len=64) - - def test_conservation_all_ranks(self): - """TP4 → TP2: verify all source elements are accounted for across all dst ranks.""" - src = _meta(shape=(64, 32), is_tp=True, partition_dim=1, tp_ranks=[0, 1, 2, 3]) - dst = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1, 2, 3], dst_ranks=[0, 1]) - - all_ops = [] - for rank in range(2): - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=rank) - all_ops.extend(ops) - - # Total transferred elements should equal full tensor size on dim 1 - total = sum(op[2][1].stop - op[2][1].start for op in all_ops) - assert total == 128 # 64 per rank * 2 ranks - - -# =========================================================================== -# _plan_tp -# =========================================================================== - - -class TestPlanBlockInterleaved: - """Tests for the block-interleaved TP planner (Mamba in_proj style).""" - - def test_tp2_to_tp1_two_blocks(self): - """TP2 → TP1 with two blocks of different sizes.""" - # Block sizes per rank: [32, 16] → full sizes: [64, 32] - # Source local dim = 32+16 = 48, Dest local dim = 64+32 = 96 - src = _meta( - shape=(64, 48), is_tp=True, partition_dim=1, partition_sizes=[32, 16], tp_ranks=[0, 1] - ) - dst = _meta(shape=(64, 96), is_tp=False, partition_sizes=None, tp_ranks=[0]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0]) - - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - assert len(ops) > 0 - _verify_full_coverage(ops, dim=1, expected_full_len=96) - - def test_tp1_to_tp2_two_blocks(self): - """TP1 → TP2 with two blocks.""" - src = _meta(shape=(64, 96), is_tp=False, partition_sizes=None, tp_ranks=[0]) - dst = _meta( - shape=(64, 48), is_tp=True, partition_dim=1, partition_sizes=[32, 16], tp_ranks=[0, 1] - ) - desc = _tp_descriptor(dim=1, src_ranks=[0], dst_ranks=[0, 1]) - - ops_r0 = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - ops_r1 = _plan_tp("weight", src, dst, [desc], my_global_rank=1) - assert len(ops_r0) > 0 - assert len(ops_r1) > 0 - _verify_full_coverage(ops_r0, dim=1, expected_full_len=48) - _verify_full_coverage(ops_r1, dim=1, expected_full_len=48) - - def test_rank_not_in_dst(self): - """Rank not in destination returns empty.""" - src = _meta( - shape=(64, 48), is_tp=True, partition_dim=1, partition_sizes=[32, 16], tp_ranks=[0, 1] + src = _tp_metadata(shape=(64, 64), partition_dim=1, tp_ranks=[0, 1]) + dst_group = [2, 3, 4, 5] + + for rank in dst_group: + dst = _meta( + shape=(64, 32), + is_tp=True, + partition_dim=1, + owner_rank=rank, + tp_ranks=dst_group, + dp_ranks=[rank], + ) + ops = plan_sharded_transfer("weight", src, src[0], dst) + _verify_nd_coverage(ops, (64, 32)) + + @pytest.mark.parametrize( + ("partition_dim", "src_size", "dst_size", "src_stride", "dst_stride"), + [(0, 2, 1, 1, 1), (1, 4, 2, 1, 1), (0, 3, 4, 1, 1), (1, 2, 3, 2, 3), (0, 4, 2, 3, 2)], + ) + def test_tp_layouts_preserve_logical_indices( + self, partition_dim, src_size, dst_size, src_stride, dst_stride + ): + global_size = 144 + src_group = list(range(src_size)) + dst_group = list(range(100, 100 + dst_size)) + src_shape = [6, 6] + dst_shape = [6, 6] + src_shape[partition_dim] = global_size // src_size + dst_shape[partition_dim] = global_size // dst_size + src = _tp_metadata( + shape=tuple(src_shape), + partition_dim=partition_dim, + partition_stride=src_stride, + tp_ranks=src_group, ) - dst = _meta(shape=(64, 96), tp_ranks=[2]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[2]) - - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) - assert ops == [] - def test_three_blocks_tp2_to_tp4(self): - """TP2 → TP4 with three blocks (simulates Mamba z,x,B,C,dt packing).""" - # Per-rank sizes: [16, 8, 4] → full: [32, 16, 8] - src = _meta( - shape=(64, 28), is_tp=True, partition_dim=1, partition_sizes=[16, 8, 4], tp_ranks=[0, 1] + for dst_tp_rank, dst_rank in enumerate(dst_group): + dst = _meta( + shape=tuple(dst_shape), + is_tp=dst_size > 1, + partition_dim=partition_dim, + partition_stride=dst_stride, + owner_rank=dst_rank, + tp_ranks=dst_group, + dp_ranks=[dst_rank], + ) + ops = plan_sharded_transfer("weight", src, src[0], dst) + _verify_nd_coverage(ops, tuple(dst_shape)) + for src_rank, src_slice, dst_slice in ops: + for src_index, dst_index in zip( + range(src_slice[partition_dim].start, src_slice[partition_dim].stop), + range(dst_slice[partition_dim].start, dst_slice[partition_dim].stop), + ): + src_global = _tp_global_index( + src_index, src_rank, src_size, src_shape[partition_dim], src_stride + ) + dst_global = _tp_global_index( + dst_index, dst_tp_rank, dst_size, dst_shape[partition_dim], dst_stride + ) + assert src_global == dst_global + + def test_strided_tp(self): + src = _tp_metadata(shape=(64, 32), partition_dim=1, partition_stride=2, tp_ranks=[0, 1]) + dst = _meta(shape=(64, 64), owner_rank=2, tp_ranks=[2], dp_ranks=[2]) + + ops = plan_sharded_transfer("weight", src, src[0], dst) + + assert [rank for rank, _, _ in ops] == [0, 1, 0, 1] + _verify_nd_coverage(ops, (64, 64)) + + def test_packed_tp2_to_tp4(self): + src = _tp_metadata( + shape=(64, 28), partition_dim=1, partition_sizes=[16, 8, 4], tp_ranks=[0, 1] ) + dst_group = [2, 3, 4, 5] + + for rank in dst_group: + dst = _meta( + shape=(64, 14), + is_tp=True, + partition_dim=1, + partition_sizes=[8, 4, 2], + owner_rank=rank, + tp_ranks=dst_group, + dp_ranks=[rank], + ) + ops = plan_sharded_transfer("weight", src, src[0], dst) + _verify_nd_coverage(ops, (64, 14)) + + def test_logical_shape_mismatch_raises(self): + src = _tp_metadata(shape=(64, 64), partition_dim=1, tp_ranks=[0, 1]) + dst = _meta(shape=(64, 100), owner_rank=2, tp_ranks=[2], dp_ranks=[2]) + + with pytest.raises(RuntimeError, match="logical shape mismatch"): + plan_sharded_transfer("weight", src, src[0], dst) + + def test_overlapping_source_metadata_raises(self): + src = [ + _meta( + shape=(4, 3), + owner_rank=rank, + tp_ranks=[rank], + dp_ranks=[rank], + is_gtp=True, + gtp_ranks=group, + ) + for rank, group in ((0, [0, 1]), (1, [1, 0])) + ] + dst = _meta(shape=(8, 3), owner_rank=2, tp_ranks=[2], dp_ranks=[2]) + + with pytest.raises(RuntimeError, match="overlapping destination coverage"): + plan_sharded_transfer("weight", src, src[0], dst) + + def test_missing_tp_group_raises(self): + src = [_meta(shape=(64, 64), is_tp=True, partition_dim=1, owner_rank=0)] + dst = _meta(shape=(64, 128), owner_rank=1) + + with pytest.raises(RuntimeError, match="missing tensor-parallel group"): + plan_sharded_transfer("weight", src, src[0], dst) + + def test_gtp2_to_unsharded(self): + src = [ + _meta( + shape=(4, 3), + owner_rank=rank, + tp_ranks=[rank], + dp_ranks=[rank], + is_gtp=True, + gtp_ranks=[0, 1], + ) + for rank in (0, 1) + ] + dst = _meta(shape=(8, 3), owner_rank=2, tp_ranks=[2], dp_ranks=[2]) + + ops = plan_sharded_transfer("weight", src, src[0], dst) + + assert [rank for rank, _, _ in ops] == [0, 1] + _verify_nd_coverage(ops, (8, 3)) + + def test_uses_only_selected_source_replica(self): + src = [] + for group in ([0, 1], [2, 3]): + src.extend( + _meta( + shape=(4, 3), + owner_rank=rank, + tp_ranks=[rank], + dp_ranks=[rank % 2, rank % 2 + 2], + is_gtp=True, + gtp_ranks=group, + ) + for rank in group + ) + dst = _meta(shape=(8, 3), owner_rank=4, tp_ranks=[4], dp_ranks=[4]) + + ops = plan_sharded_transfer("weight", src, src[2], dst) + + assert [rank for rank, _, _ in ops] == [2, 3] + _verify_nd_coverage(ops, (8, 3)) + + def test_unsharded_to_padded_gtp2(self): + src = [_meta(shape=(6, 2), owner_rank=0, tp_ranks=[0], dp_ranks=[0])] dst = _meta( - shape=(64, 14), - is_tp=True, - partition_dim=1, - partition_sizes=[8, 4, 2], - tp_ranks=[0, 1, 2, 3], + shape=(4, 2), + owner_rank=2, + tp_ranks=[2], + dp_ranks=[2], + is_gtp=True, + gtp_ranks=[1, 2], + gtp_pad_length=2, ) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0, 1, 2, 3]) - - for rank in range(4): - ops = _plan_tp("weight", src, dst, [desc], my_global_rank=rank) - assert len(ops) > 0 - _verify_full_coverage(ops, dim=1, expected_full_len=14) - - -# =========================================================================== -# _finalize_dp_transfers -# =========================================================================== - - -class TestFinalizeDpTransfers: - """Tests for the DP/replicated parameter fallback planner.""" - - def test_same_dp_local_copy(self): - """Same DP group → local copy from self.""" - src = _meta(owner_rank=0, dp_ranks=[0, 1]) - dst = _meta(owner_rank=0, dp_ranks=[0, 1]) - ops = _finalize_dp_transfers("weight", src, dst, my_global_rank=0) - assert len(ops) == 1 - assert ops[0][0] == 0 # source is self - # Full tensor copy - assert ops[0][1] == (slice(None), slice(None)) - assert ops[0][2] == (slice(None), slice(None)) + ops = plan_sharded_transfer("weight", src, src[0], dst) - def test_different_dp_uses_owner_rank(self): - """Different DP groups → uses src_metadata.owner_rank.""" - src = _meta(owner_rank=2, dp_ranks=[2, 3]) - dst = _meta(owner_rank=0, dp_ranks=[0, 1]) + assert ops == [(0, (slice(4, 6), slice(0, 2)), (slice(0, 2), slice(0, 2)))] - ops = _finalize_dp_transfers("weight", src, dst, my_global_rank=0) - assert len(ops) == 1 - assert ops[0][0] == 2 # src owner rank + def test_column_tp2_gtp2_to_unsharded(self): + src = self._tp2_gtp2_metadata(shape=(2, 3), partition_dim=0) + dst = _meta(shape=(8, 3), owner_rank=4, tp_ranks=[4], dp_ranks=[4]) - def test_rank_not_in_dst_dp(self): - """Rank not in destination DP group returns empty.""" - src = _meta(owner_rank=0, dp_ranks=[0, 1]) - dst = _meta(owner_rank=2, dp_ranks=[2, 3]) + ops = plan_sharded_transfer("weight", src, src[0], dst) - ops = _finalize_dp_transfers("weight", src, dst, my_global_rank=0) - assert ops == [] + assert [rank for rank, _, _ in ops] == [0, 2, 1, 3] + _verify_nd_coverage(ops, (8, 3)) - def test_non_collocated_dp(self): - """Non-collocated: src and dst have completely disjoint ranks.""" - src = _meta(owner_rank=0, dp_ranks=[0, 1, 2, 3]) - dst = _meta(owner_rank=4, dp_ranks=[4, 5, 6, 7]) + def test_row_tp2_gtp2_to_unsharded(self): + src = self._tp2_gtp2_metadata(shape=(4, 4), partition_dim=1) + dst = _meta(shape=(8, 8), owner_rank=4, tp_ranks=[4], dp_ranks=[4]) - ops = _finalize_dp_transfers("weight", src, dst, my_global_rank=4) - assert len(ops) == 1 - assert ops[0][0] == 0 # from src owner + ops = plan_sharded_transfer("weight", src, src[0], dst) + assert len(ops) == 4 + assert {rank for rank, _, _ in ops} == {0, 1, 2, 3} + _verify_nd_coverage(ops, (8, 8)) -# =========================================================================== -# _build_descriptors_for_param -# =========================================================================== + def test_strided_tp_gtp_to_unsharded(self): + src = self._tp2_gtp2_metadata(shape=(2, 3), partition_dim=0, partition_stride=2) + dst = _meta(shape=(8, 3), owner_rank=4, tp_ranks=[4], dp_ranks=[4]) + ops = plan_sharded_transfer("weight", src, src[0], dst) -class TestBuildDescriptors: - """Tests for TP descriptor construction.""" - - def test_tp_both_sides(self): - """Both src and dst are TP → produces TP descriptor.""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 128), is_tp=True, partition_dim=1, tp_ranks=[0]) - - descs = _build_descriptors_for_param(src, dst) - assert len(descs) == 1 - assert descs[0].name == "tp" - assert descs[0].dim == 1 - - def test_tp_one_side_only(self): - """Only src is TP → still produces TP descriptor.""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 128), is_tp=False, tp_ranks=[0]) - - descs = _build_descriptors_for_param(src, dst) - assert len(descs) == 1 - - def test_neither_tp(self): - """Neither side is TP → no descriptors.""" - src = _meta(shape=(64, 128), is_tp=False, tp_ranks=None) - dst = _meta(shape=(64, 128), is_tp=False, tp_ranks=None) - - descs = _build_descriptors_for_param(src, dst) - assert descs == [] - - def test_size_conservation_failure(self): - """Mismatched global sizes should raise.""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - # Global = 128, but dst claims global = 100 (1 rank * 100) - dst = _meta(shape=(64, 100), is_tp=True, partition_dim=1, tp_ranks=[0]) - - with pytest.raises(RuntimeError, match="Cannot build TP descriptor"): - _build_descriptors_for_param(src, dst) - - def test_missing_tp_ranks(self): - """Missing TP group ranks → no descriptors (not enough context).""" - src = _meta(shape=(64, 64), is_tp=True, partition_dim=1, tp_ranks=None) - dst = _meta(shape=(64, 128), is_tp=False, tp_ranks=[0]) + assert [rank for rank, _, _ in ops] == [0, 1, 2, 3] + assert [dst_slice[0] for _, _, dst_slice in ops] == [ + slice(0, 2), + slice(2, 4), + slice(4, 6), + slice(6, 8), + ] + _verify_nd_coverage(ops, (8, 3)) - descs = _build_descriptors_for_param(src, dst) - assert descs == [] + def test_packed_tp_gtp_layout_with_padding(self): + src = self._tp2_gtp2_metadata( + shape=(2, 2), partition_dim=0, partition_sizes=[2, 1], gtp_pad_length=1 + ) + dst = _meta( + shape=(6, 2), + is_tp=True, + partition_dim=0, + partition_sizes=[4, 2], + owner_rank=4, + tp_ranks=[4], + dp_ranks=[4], + ) + ops = plan_sharded_transfer("weight", src, src[0], dst) -# =========================================================================== -# build_plan_from_rosters (local, deterministic planning + node-add stability) -# =========================================================================== + assert [rank for rank, _, _ in ops] == [0, 1, 2, 3] + _verify_nd_coverage(ops, (6, 2)) def _plan_edges(plans): - """Collect (task_id, src_rank, dst_rank) transfers from a {rank: ReshardPlan}. - - Reads them once from every send op and once from every recv op; the two sets - must be equal for the plan to be consistent (a matching, same-task_id recv for - every send). - """ - sends = {(op.task_id, r, op.peer_rank) for r, p in plans.items() for op in p.send_ops} - recvs = {(op.task_id, op.peer_rank, r) for r, p in plans.items() for op in p.recv_ops} + """Collect matching send and receive edges from per-rank plans.""" + sends = { + (op.task_id, rank, op.peer_rank) for rank, plan in plans.items() for op in plan.send_ops + } + recvs = { + (op.task_id, op.peer_rank, rank) for rank, plan in plans.items() for op in plan.recv_ops + } return sends, recvs def _build_all(gathered_pairs): - """Build every rank's plan from a rank-ordered list of (src_meta, dst_meta).""" + """Build every rank's plan from a rank-ordered metadata roster.""" dst_by_rank, src_by_name = index_metadata_rosters(gathered_pairs) return {rank: build_plan_from_rosters(dst_by_rank, src_by_name, rank) for rank in dst_by_rank} def _recv_sig(plan): - """Identity of a plan's recv ops (task_id + slices), for stability comparisons.""" + """Return the stable identity of a plan's receive operations.""" return [(op.task_id, op.peer_rank, op.my_slice, op.peer_slice) for op in plan.recv_ops] class TestBuildPlanFromRosters: - """Local plan building replayed independently per rank.""" + """The global deterministic schedule uses the same shard algorithm.""" + + def test_gtp_schedule_matches_across_ranks(self): + src = [ + _meta( + shape=(4, 3), + owner_rank=rank, + tp_ranks=[rank], + dp_ranks=[rank], + is_gtp=True, + gtp_ranks=[0, 1], + ) + for rank in (0, 1) + ] + dst = _meta(shape=(8, 3), owner_rank=2, tp_ranks=[2], dp_ranks=[2]) + plans = _build_all([([src[0]], []), ([src[1]], []), ([], [dst])]) - def test_task_ids_match_across_ranks(self): - """Sender and receiver, planned independently, agree on task_id per transfer. + sends, recvs = _plan_edges(plans) + assert sends == recvs + assert {(src_rank, dst_rank) for _, src_rank, dst_rank in sends} == {(0, 2), (1, 2)} + assert [op.my_slice[0] for op in plans[2].recv_ops] == [slice(0, 4), slice(4, 8)] - rank 0 sources a replicated weight; ranks 1 and 2 each receive a full copy. - """ + def test_task_ids_match_across_ranks(self): gathered = [ - ([_meta(owner_rank=0, tp_ranks=[0], dp_ranks=[0])], []), # rank 0: source - ([], [_meta(owner_rank=1, tp_ranks=[1], dp_ranks=[1])]), # rank 1: dest - ([], [_meta(owner_rank=2, tp_ranks=[2], dp_ranks=[2])]), # rank 2: dest + ([_meta(owner_rank=0, tp_ranks=[0], dp_ranks=[0])], []), + ([], [_meta(owner_rank=1, tp_ranks=[1], dp_ranks=[1])]), + ([], [_meta(owner_rank=2, tp_ranks=[2], dp_ranks=[2])]), ] plans = _build_all(gathered) sends, recvs = _plan_edges(plans) - # Every send has a matching recv with the same task_id, and vice versa. assert sends == recvs - # Two transfers: 0->1 and 0->2, with distinct task_ids. assert len(sends) == 2 - assert {(s, d) for _, s, d in sends} == {(0, 1), (0, 2)} - assert len({tid for tid, _, _ in sends}) == 2 + assert {(src, dst) for _, src, dst in sends} == {(0, 1), (0, 2)} + assert len({task_id for task_id, _, _ in sends}) == 2 def test_node_add_keeps_existing_task_ids_stable(self): - """Appending a rank rebuilds locally without renumbering existing transfers.""" base = [ ([_meta(owner_rank=0, tp_ranks=[0], dp_ranks=[0])], []), ([], [_meta(owner_rank=1, tp_ranks=[1], dp_ranks=[1])]), ([], [_meta(owner_rank=2, tp_ranks=[2], dp_ranks=[2])]), ] before = _build_all(base) + after = _build_all(base + [([], [_meta(owner_rank=3, tp_ranks=[3], dp_ranks=[3])])]) - # A new destination rank 3 joins; everyone replays over the grown roster. - grown = base + [([], [_meta(owner_rank=3, tp_ranks=[3], dp_ranks=[3])])] - after = _build_all(grown) - - sends_after, recvs_after = _plan_edges(after) - assert sends_after == recvs_after - # Existing receivers keep the exact same recv ops (task_id + slices). + sends, recvs = _plan_edges(after) + assert sends == recvs for rank in (1, 2): assert _recv_sig(before[rank]) == _recv_sig(after[rank]) - # The new rank added exactly one transfer with a fresh task_id. - assert len(sends_after) == 3 - assert {(s, d) for _, s, d in sends_after} == {(0, 1), (0, 2), (0, 3)} + assert len(sends) == 3 + assert {(src, dst) for _, src, dst in sends} == {(0, 1), (0, 2), (0, 3)} def test_centralized_planner_compatibility_wrapper(monkeypatch): diff --git a/tests/unit_tests/resharding/test_refit_cache.py b/tests/unit_tests/resharding/test_refit_cache.py index b7c7fba2c8a..0ad9e1885d5 100644 --- a/tests/unit_tests/resharding/test_refit_cache.py +++ b/tests/unit_tests/resharding/test_refit_cache.py @@ -14,20 +14,28 @@ import torch import torch.nn as nn -from megatron.core.resharding.refit import _PlanCacheKey +from megatron.core.resharding.refit import _get_parallel_config, _ParallelConfig, _PlanCacheKey from megatron.core.resharding.utils import get_refit_tensor_dict, invalidate_refit_tensor_cache +def _config(tp=1, pp=1, ep=1, dp=1, expert_tp=1, gtp_remat=1, expert_gtp_remat=1): + return _ParallelConfig( + tp_size=tp, + pp_size=pp, + ep_size=ep, + dp_size=dp, + expert_tp_size=expert_tp, + gtp_remat_size=gtp_remat, + expert_gtp_remat_size=expert_gtp_remat, + ) + + class TestPlanCacheKey: """Plan cache must distinguish configs that route to different global ranks.""" def test_equality_with_same_inputs(self): - k1 = _PlanCacheKey( - rank=0, src_config=(1, 1, 1, 1, 1), dst_config=(1, 1, 1, 1, 1), num_experts=None - ) - k2 = _PlanCacheKey( - rank=0, src_config=(1, 1, 1, 1, 1), dst_config=(1, 1, 1, 1, 1), num_experts=None - ) + k1 = _PlanCacheKey(rank=0, src_config=_config(), dst_config=_config(), num_experts=None) + k2 = _PlanCacheKey(rank=0, src_config=_config(), dst_config=_config(), num_experts=None) assert k1 == k2 assert hash(k1) == hash(k2) @@ -35,16 +43,16 @@ def test_different_src_rank_offset_distinguishes(self): """Same sizes + rank, different src_rank_offset → different cache key.""" k1 = _PlanCacheKey( rank=0, - src_config=(2, 1, 1, 2, 1), - dst_config=(2, 1, 1, 2, 1), + src_config=_config(tp=2, dp=2), + dst_config=_config(tp=2, dp=2), num_experts=None, src_rank_offset=0, dst_rank_offset=4, ) k2 = _PlanCacheKey( rank=0, - src_config=(2, 1, 1, 2, 1), - dst_config=(2, 1, 1, 2, 1), + src_config=_config(tp=2, dp=2), + dst_config=_config(tp=2, dp=2), num_experts=None, src_rank_offset=8, dst_rank_offset=12, @@ -55,16 +63,16 @@ def test_different_src_rank_offset_distinguishes(self): def test_different_dst_rank_offset_distinguishes(self): k1 = _PlanCacheKey( rank=0, - src_config=(2, 1, 1, 2, 1), - dst_config=(2, 1, 1, 2, 1), + src_config=_config(tp=2, dp=2), + dst_config=_config(tp=2, dp=2), num_experts=None, src_rank_offset=0, dst_rank_offset=4, ) k2 = _PlanCacheKey( rank=0, - src_config=(2, 1, 1, 2, 1), - dst_config=(2, 1, 1, 2, 1), + src_config=_config(tp=2, dp=2), + dst_config=_config(tp=2, dp=2), num_experts=None, src_rank_offset=0, dst_rank_offset=8, @@ -74,12 +82,12 @@ def test_different_dst_rank_offset_distinguishes(self): def test_default_offsets_match_collocated(self): """Collocated callers (no offsets specified) reuse the same plan.""" k1 = _PlanCacheKey( - rank=3, src_config=(2, 1, 1, 4, 1), dst_config=(2, 1, 1, 4, 1), num_experts=None + rank=3, src_config=_config(tp=2, dp=4), dst_config=_config(tp=2, dp=4), num_experts=None ) k2 = _PlanCacheKey( rank=3, - src_config=(2, 1, 1, 4, 1), - dst_config=(2, 1, 1, 4, 1), + src_config=_config(tp=2, dp=4), + dst_config=_config(tp=2, dp=4), num_experts=None, src_rank_offset=0, dst_rank_offset=0, @@ -91,6 +99,43 @@ def test_num_experts_distinguishes(self): k2 = _PlanCacheKey(rank=0, src_config=None, dst_config=None, num_experts=16) assert k1 != k2 + def test_gtp_remat_sizes_distinguish(self): + base = _config(tp=2, dp=2) + plain = _PlanCacheKey(rank=0, src_config=base, dst_config=base, num_experts=None) + + for config in (_config(tp=2, dp=2, gtp_remat=4), _config(tp=2, dp=2, expert_gtp_remat=2)): + assert plain != _PlanCacheKey( + rank=0, src_config=config, dst_config=config, num_experts=None + ) + + +def test_parallel_config_includes_gtp_remat_sizes(): + class Group: + def __init__(self, size): + self._size = size + + def size(self): + return self._size + + class Core: + pg_collection = type( + "PG", + (), + { + "tp": Group(2), + "pp": Group(3), + "ep": Group(4), + "dp": Group(5), + "expt_tp": Group(6), + "gtp_remat": Group(7), + "expt_gtp_remat": Group(8), + }, + )() + + assert _get_parallel_config(Core()) == _config( + tp=2, pp=3, ep=4, dp=5, expert_tp=6, gtp_remat=7, expert_gtp_remat=8 + ) + class TestPlanCacheKeyNonCollocated: """Non-collocated ranks set src_config or dst_config to None. @@ -102,20 +147,16 @@ class TestPlanCacheKeyNonCollocated: def test_source_only_vs_dest_only_distinguish(self): """Source-only (dst_config=None) and dest-only (src_config=None) on the same global rank must produce different plans.""" - sizes = (2, 1, 1, 2, 1) - src_only = _PlanCacheKey(rank=0, src_config=sizes, dst_config=None, num_experts=None) - dst_only = _PlanCacheKey(rank=0, src_config=None, dst_config=sizes, num_experts=None) + config = _config(tp=2, dp=2) + src_only = _PlanCacheKey(rank=0, src_config=config, dst_config=None, num_experts=None) + dst_only = _PlanCacheKey(rank=0, src_config=None, dst_config=config, num_experts=None) assert src_only != dst_only def test_idle_rank_distinguishes_from_active(self): """Idle rank (both configs None) is distinct from a rank with either model.""" idle = _PlanCacheKey(rank=5, src_config=None, dst_config=None, num_experts=None) - with_src = _PlanCacheKey( - rank=5, src_config=(1, 1, 1, 1, 1), dst_config=None, num_experts=None - ) - with_dst = _PlanCacheKey( - rank=5, src_config=None, dst_config=(1, 1, 1, 1, 1), num_experts=None - ) + with_src = _PlanCacheKey(rank=5, src_config=_config(), dst_config=None, num_experts=None) + with_dst = _PlanCacheKey(rank=5, src_config=None, dst_config=_config(), num_experts=None) assert idle != with_src assert idle != with_dst assert with_src != with_dst @@ -123,20 +164,20 @@ def test_idle_rank_distinguishes_from_active(self): def test_non_collocated_offset_combinations(self): """src_rank_offset and dst_rank_offset together distinguish non-collocated layouts that share parallel sizes.""" - sizes = (2, 1, 1, 2, 1) + config = _config(tp=2, dp=2) # Two non-collocated layouts: world=[src 0-3, dst 4-7] vs [src 0-3, dst 8-11]. layout_a = _PlanCacheKey( rank=0, - src_config=sizes, - dst_config=sizes, + src_config=config, + dst_config=config, num_experts=None, src_rank_offset=0, dst_rank_offset=4, ) layout_b = _PlanCacheKey( rank=0, - src_config=sizes, - dst_config=sizes, + src_config=config, + dst_config=config, num_experts=None, src_rank_offset=0, dst_rank_offset=8,