From e137aa162b77aa1112a0211549f8df770fac9472 Mon Sep 17 00:00:00 2001 From: William Dykas Date: Tue, 12 May 2026 12:53:28 -0700 Subject: [PATCH 1/4] fri --- .../core/resharding/copy_services/base.py | 62 ++++ .../copy_services/gloo_copy_service.py | 77 +---- .../copy_services/nccl_copy_service.py | 62 +--- .../copy_services/nvshmem_copy_service.py | 32 +- megatron/core/resharding/execution.py | 279 +++++++++--------- megatron/core/resharding/planner.py | 229 +++++++------- megatron/core/resharding/refit.py | 178 ++++++----- megatron/core/resharding/transforms.py | 10 +- megatron/core/resharding/utils.py | 241 +++++++-------- .../unit_tests/resharding/test_mxfp8_refit.py | 21 ++ 10 files changed, 565 insertions(+), 626 deletions(-) diff --git a/megatron/core/resharding/copy_services/base.py b/megatron/core/resharding/copy_services/base.py index 3ea58e8d141..de47705fd15 100644 --- a/megatron/core/resharding/copy_services/base.py +++ b/megatron/core/resharding/copy_services/base.py @@ -2,9 +2,29 @@ from __future__ import annotations from abc import ABC, abstractmethod +from dataclasses import dataclass from typing import Optional import torch +import torch.distributed as dist + + +@dataclass +class SendOp: + """Single send operation pending in a CopyService queue.""" + + task_id: int | None + tensor: torch.Tensor + dest_rank: int + + +@dataclass +class RecvOp: + """Single receive operation pending in a CopyService queue.""" + + task_id: int | None + tensor: torch.Tensor + src_rank: int class CopyService(ABC): @@ -17,6 +37,13 @@ class CopyService(ABC): remote transfers simply ignore it. """ + def __init__(self, group=None): + self.group = group + # group.rank()/size() supports cross-cluster ProcessGroups where members + # have independent default PGs. + self.rank = group.rank() if group is not None else dist.get_rank() + self.world_size = group.size() if group is not None else dist.get_world_size() + @abstractmethod def submit_send(self, src_tensor: torch.Tensor, dest_rank: int, task_id: Optional[int] = None): """Register a tensor send from the current rank to ``dest_rank``.""" @@ -31,3 +58,38 @@ def submit_recv(self, dest_tensor: torch.Tensor, src_rank: int, task_id: Optiona def run(self): """Execute all previously submitted send/recv operations as a single batch.""" ... + + def close(self) -> None: + """Release backend-owned resources. Default no-op; NVSHMEM overrides.""" + + +def match_local_ops_by_task_id( + local_sends: list, local_recvs: list, backend_name: str, rank: int +) -> list[tuple]: + """Pair same-rank send/recv ops by task_id, raising on any mismatch. + + Returns a list of ``(send_op, recv_op)`` tuples for the caller to apply + backend-specific local-copy logic. Either op type may be a backend-local + wrapper as long as it exposes ``.task_id``. + """ + sends_by_id = {op.task_id: op for op in local_sends} + recvs_by_id = {op.task_id: op for op in local_recvs} + if None in sends_by_id or None in recvs_by_id: + raise RuntimeError( + f"{backend_name}: local (same-rank) transfer requires a task_id " + "to match sends with recvs" + ) + if len(sends_by_id) != len(local_sends) or len(recvs_by_id) != len(local_recvs): + raise RuntimeError( + f"{backend_name}: unmatched local ops on rank {rank}: " + f"{len(local_sends)} local sends vs {len(local_recvs)} local recvs" + ) + pairs = [] + for task_id, recv_op in recvs_by_id.items(): + send_op = sends_by_id.get(task_id) + if send_op is None: + raise RuntimeError( + f"{backend_name}: missing local send for task_id={task_id} on rank {rank}" + ) + pairs.append((send_op, recv_op)) + return pairs diff --git a/megatron/core/resharding/copy_services/gloo_copy_service.py b/megatron/core/resharding/copy_services/gloo_copy_service.py index 0857f8dfb52..5846a501e77 100644 --- a/megatron/core/resharding/copy_services/gloo_copy_service.py +++ b/megatron/core/resharding/copy_services/gloo_copy_service.py @@ -2,35 +2,16 @@ from __future__ import annotations import logging -from dataclasses import dataclass from typing import List, Optional, Tuple import torch import torch.distributed as dist -from .base import CopyService +from .base import CopyService, RecvOp, SendOp, match_local_ops_by_task_id logger = logging.getLogger(__name__) -@dataclass -class SendOp: - """Simple container describing a single send operation.""" - - task_id: int | None - tensor: torch.Tensor - dest_rank: int - - -@dataclass -class RecvOp: - """Simple container describing a single receive operation.""" - - task_id: int | None - tensor: torch.Tensor - src_rank: int - - class GlooCopyService(CopyService): """ CopyService implementation that routes refit traffic over a CPU/Gloo @@ -38,17 +19,15 @@ class GlooCopyService(CopyService): """ def __init__(self, group=None): + super().__init__(group=group) if group is not None: self.gloo_pg = group - self.rank = group.rank() - self.world_size = group.size() else: - self.rank = dist.get_rank() - self.world_size = dist.get_world_size() self.gloo_pg = dist.new_group(backend="gloo") self.send_ops: List[SendOp] = [] + # Each recv op is paired with its GPU destination tensor; the SendOp/RecvOp + # itself carries a pinned-CPU staging buffer for Gloo's CPU PG. self.recv_ops: List[Tuple[RecvOp, torch.Tensor]] = [] - self._copy_stream = torch.cuda.Stream() if self.rank == 0: logger.info( f"GlooCopyService initialized on rank {self.rank} with {self.world_size} ranks" @@ -58,7 +37,7 @@ def submit_send(self, src_tensor: torch.Tensor, dest_rank: int, task_id: Optiona self.send_ops.append(SendOp(task_id=task_id, tensor=src_tensor, dest_rank=dest_rank)) def submit_recv(self, dest_tensor: torch.Tensor, src_rank: int, task_id: Optional[int] = None): - # Allocate a pinned CPU buffer for faster CPU↔GPU transfer. + # Pinned CPU staging buffer for the Gloo recv; copied back to dest_tensor in run(). cpu_buffer = torch.empty( dest_tensor.shape, dtype=dest_tensor.dtype, device="cpu", pin_memory=True ) @@ -74,43 +53,21 @@ def run(self): f"{len(self.send_ops)} sends + {len(self.recv_ops)} recvs = {total_ops} ops" ) - p2p_ops: List[dist.P2POp] = [] - - # Short-circuit self transfers into local device copies. local_sends = [op for op in self.send_ops if op.dest_rank == self.rank] remote_sends = [op for op in self.send_ops if op.dest_rank != self.rank] local_recvs = [(recv, dst) for (recv, dst) in self.recv_ops if recv.src_rank == self.rank] remote_recvs = [(recv, dst) for (recv, dst) in self.recv_ops if recv.src_rank != self.rank] if local_sends or local_recvs: - local_sends_by_id = {op.task_id: op for op in local_sends} - if None in local_sends_by_id: - raise RuntimeError( - "GlooCopyService: local (same-rank) transfer requires a task_id " - "to match sends with recvs" - ) - local_recvs_by_id = {recv.task_id: (recv, dst) for (recv, dst) in local_recvs} - if None in local_recvs_by_id: - raise RuntimeError( - "GlooCopyService: local (same-rank) transfer requires a task_id " - "to match sends with recvs" - ) - if len(local_sends_by_id) != len(local_sends) or len(local_recvs_by_id) != len( - local_recvs - ): - raise RuntimeError( - f"GlooCopyService: unmatched local ops on rank {self.rank}: " - f"{len(local_sends)} local sends vs {len(local_recvs)} local recvs" - ) - for task_id, (recv_op, dst_tensor) in local_recvs_by_id.items(): - send_op = local_sends_by_id.get(task_id) - if send_op is None: - raise RuntimeError( - f"GlooCopyService: missing local send for task_id={task_id} " - f"on rank {self.rank}" - ) - with torch.no_grad(): + local_recv_objs = [recv for recv, _ in local_recvs] + dst_by_task_id = {recv.task_id: dst for recv, dst in local_recvs} + pairs = match_local_ops_by_task_id( + local_sends, local_recv_objs, "GlooCopyService", self.rank + ) + with torch.no_grad(): + for send_op, recv_op in pairs: src_tensor = send_op.tensor + dst_tensor = dst_by_task_id[recv_op.task_id] if dst_tensor.device != src_tensor.device: dst_tensor.copy_(src_tensor.to(dst_tensor.device)) else: @@ -128,13 +85,14 @@ def run(self): ) cpu_tensor.copy_(op.tensor.detach(), non_blocking=True) cpu_send_bufs.append(cpu_tensor) - # Single sync after all GPU→CPU copies are issued. if cpu_send_bufs: torch.cuda.synchronize() for op in remote_sends: + # Drop the GPU reference now that staging is complete. op.tensor = None + p2p_ops: List[dist.P2POp] = [] for cpu_tensor, op in zip(cpu_send_bufs, remote_sends): p2p_ops.append( dist.P2POp(dist.isend, cpu_tensor, group=self.gloo_pg, group_peer=op.dest_rank) @@ -149,17 +107,12 @@ def run(self): for req in reqs: req.wait() - # Copy received CPU buffers back into the original destination tensors. - # Use non_blocking with pinned memory for overlap. for recv, dst_tensor in remote_recvs: if dst_tensor.is_cuda: dst_tensor.copy_(recv.tensor, non_blocking=True) else: dst_tensor.copy_(recv.tensor) - if self._copy_stream is not None: - torch.cuda.current_stream().wait_stream(self._copy_stream) - # Ensure all async CPU→GPU copies are complete. torch.cuda.synchronize() diff --git a/megatron/core/resharding/copy_services/nccl_copy_service.py b/megatron/core/resharding/copy_services/nccl_copy_service.py index 77df4381136..9fc25e401cf 100644 --- a/megatron/core/resharding/copy_services/nccl_copy_service.py +++ b/megatron/core/resharding/copy_services/nccl_copy_service.py @@ -2,35 +2,16 @@ from __future__ import annotations import logging -from dataclasses import dataclass from typing import List, Optional import torch import torch.distributed as dist -from .base import CopyService +from .base import CopyService, RecvOp, SendOp, match_local_ops_by_task_id logger = logging.getLogger(__name__) -@dataclass -class SendOp: - """Simple container describing a single send operation.""" - - task_id: int | None - tensor: torch.Tensor - dest_rank: int - - -@dataclass -class RecvOp: - """Simple container describing a single receive operation.""" - - task_id: int | None - tensor: torch.Tensor - src_rank: int - - class NCCLCopyService(CopyService): """ Thin wrapper around torch.distributed batch_isend_irecv to submit and execute @@ -38,10 +19,7 @@ class NCCLCopyService(CopyService): """ def __init__(self, group=None): - self.group = group - # Use group.rank()/size() to support cross-cluster ProcessGroups - self.rank = group.rank() if group is not None else dist.get_rank() - self.world_size = group.size() if group is not None else dist.get_world_size() + super().__init__(group=group) self.send_ops: List[SendOp] = [] self.recv_ops: List[RecvOp] = [] # Dedicated stream for local (same-rank) copies to avoid unnecessary @@ -72,35 +50,12 @@ def run(self): remote_recvs = [op for op in self.recv_ops if op.src_rank != self.rank] if local_sends or local_recvs: - local_sends_by_id = {op.task_id: op for op in local_sends} - if None in local_sends_by_id: - raise RuntimeError( - "NCCLCopyService: local (same-rank) transfer requires a task_id " - "to match sends with recvs" - ) - local_recvs_by_id = {op.task_id: op for op in local_recvs} - if None in local_recvs_by_id: - raise RuntimeError( - "NCCLCopyService: local (same-rank) transfer requires a task_id " - "to match sends with recvs" - ) - if len(local_sends_by_id) != len(local_sends) or len(local_recvs_by_id) != len( - local_recvs - ): - raise RuntimeError( - f"NCCLCopyService: unmatched local ops on rank {self.rank}: " - f"{len(local_sends)} local sends vs {len(local_recvs)} local recvs" - ) - for task_id, recv_op in local_recvs_by_id.items(): - send_op = local_sends_by_id.get(task_id) - if send_op is None: - raise RuntimeError( - f"NCCLCopyService: missing local send for task_id={task_id} " - f"on rank {self.rank}" - ) - with torch.no_grad(): - with torch.cuda.stream(self._copy_stream): - recv_op.tensor.copy_(send_op.tensor) + pairs = match_local_ops_by_task_id( + local_sends, local_recvs, "NCCLCopyService", self.rank + ) + with torch.no_grad(), torch.cuda.stream(self._copy_stream): + for send_op, recv_op in pairs: + recv_op.tensor.copy_(send_op.tensor) p2p_ops = [] for op in remote_sends: @@ -113,7 +68,6 @@ def run(self): for req in reqs: req.wait() - # Make sure the copy stream is finished torch.cuda.current_stream().wait_stream(self._copy_stream) if self.rank == 0: diff --git a/megatron/core/resharding/copy_services/nvshmem_copy_service.py b/megatron/core/resharding/copy_services/nvshmem_copy_service.py index f0ade45518d..9260dfd374c 100644 --- a/megatron/core/resharding/copy_services/nvshmem_copy_service.py +++ b/megatron/core/resharding/copy_services/nvshmem_copy_service.py @@ -24,21 +24,24 @@ class NVSHMEMCopyService(CopyService): def __init__(self, group=None): if not dist.is_initialized(): raise RuntimeError("torch.distributed must be initialized before NVSHMEMCopyService()") + super().__init__(group=group) - self._group = group - self.rank = group.rank() if group is not None else dist.get_rank() self._remote = RemoteCopyService(group=group) - # Lazily initialized on first use to avoid side effects at import time self._initialized = False - # NOTE: keep the original typed tensors here (not uint8 views) so local copies - # preserve shape/strides semantics and avoid byte-offset pitfalls. + # Keep original typed tensors (not uint8 views) so local copies preserve + # shape/strides semantics and avoid byte-offset pitfalls. self._local_send_ops: Dict[int, torch.Tensor] = {} self._local_recv_ops: Dict[int, torch.Tensor] = {} self._local_copy_stream = torch.cuda.Stream() logger.info("NVSHMEMCopyService constructed") + def close(self) -> None: + if self._initialized: + self._remote.finalize() + self._initialized = False + def _ensure_initialized(self): if not self._initialized: self._remote.init(log_level="INFO") @@ -50,14 +53,13 @@ def _ensure_initialized(self): def submit_send(self, src_tensor: torch.Tensor, dest_rank: int, task_id: Optional[int] = None): if task_id is None: raise RuntimeError( - "NVSHMEMCopyService requires a task_id for every transfer; " "got task_id=None" + "NVSHMEMCopyService requires a task_id for every transfer; got task_id=None" ) self._ensure_initialized() if not src_tensor.is_contiguous(): src_tensor = src_tensor.contiguous() - # Local transfers: keep them out of RemoteCopyService entirely. if dest_rank == self.rank: self._local_send_ops[task_id] = src_tensor return @@ -80,14 +82,13 @@ def submit_send(self, src_tensor: torch.Tensor, dest_rank: int, task_id: Optiona def submit_recv(self, dest_tensor: torch.Tensor, src_rank: int, task_id: Optional[int] = None): if task_id is None: raise RuntimeError( - "NVSHMEMCopyService requires a task_id for every transfer; " "got task_id=None" + "NVSHMEMCopyService requires a task_id for every transfer; got task_id=None" ) self._ensure_initialized() if not dest_tensor.is_contiguous(): dest_tensor = dest_tensor.contiguous() - # Local transfers: keep them out of RemoteCopyService entirely. if src_rank == self.rank: self._local_recv_ops[task_id] = dest_tensor return @@ -117,7 +118,7 @@ def run(self): """ self._ensure_initialized() - # 1) Run same-rank copies (match by task_id), like NCCL backend. + # Local copies match by task_id (the NVSHMEM RemoteCopyService never sees them). if self._local_send_ops or self._local_recv_ops: missing_sends = set(self._local_recv_ops.keys()) - set(self._local_send_ops.keys()) missing_recvs = set(self._local_send_ops.keys()) - set(self._local_recv_ops.keys()) @@ -145,12 +146,11 @@ def run(self): self._local_send_ops.clear() self._local_recv_ops.clear() - # 2) Execute remote schedule (if any remote sends/recvs were registered). - # NOTE: ALL ranks must call schedule() and run() because they contain collective - # operations that require all ranks to participate: - # - schedule() has dist.all_gather_object() (torch distributed collective) - # - run() has nvshmem.core.barrier_all() (nvshmem collective) - # This is critical for non-collocated refit where some ranks may have no work. + # ALL ranks must call schedule() and run() because they contain collectives + # that require all ranks to participate: + # - schedule() uses dist.all_gather_object() + # - run() uses nvshmem.core.barrier_all() + # Critical for non-collocated refit where some ranks may have no work. logger.info("NVSHMEMCopyService: building NVSHMEM schedule and executing") self._remote.schedule() self._remote.run() diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py index e1b75bb0a70..c548c49aad1 100644 --- a/megatron/core/resharding/execution.py +++ b/megatron/core/resharding/execution.py @@ -2,25 +2,57 @@ from __future__ import annotations import logging +from dataclasses import dataclass from typing import Optional import torch import torch.distributed as dist +from megatron.core.fp8_utils import is_mxfp8tensor + from .copy_services.base import CopyService from .transforms import ReshardTransform, _ensure_sendable -from .utils import ReshardPlan, named_refit_tensors +from .utils import ReshardPlan, get_refit_tensor_dict logger = logging.getLogger(__name__) -def _is_mxfp8_tensor(param): - """Check if param is a TE MXFP8Tensor (fp8_param=true).""" - return ( - hasattr(param, 'quantize_') - and hasattr(param, 'dequantize') - and hasattr(param, '_rowwise_data') - ) +@dataclass +class _Writeback: + """Tagged-union for what to do with a received tensor after service.run(). + + 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 + accumulation when the dest is quantized). ``transform`` hands the + received buffers to a ``ReshardTransform.finalize_recv`` call. + """ + + kind: str # 'direct' | 'copy' | 'transform' + recv_buffer: Optional[torch.Tensor] = None + dst_param: Optional[torch.Tensor] = None + dst_slice: Optional[tuple] = None + param_name: Optional[str] = None + 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. + + 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. + """ + 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, []) + pending[param_id] = entry + return entry[1], entry[2] def execute_reshard_plan( @@ -47,57 +79,40 @@ def execute_reshard_plan( transform's prepare_send / prepare_recv / finalize_recv methods instead of the default slice-and-copy logic. """ + # Refit tensors (parameters + persistent buffers) are cached on each module + # so the named_modules() walk happens once per model, not per refit. + 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 {} + + # Dequantized BF16 views of MXFP8 source params are reused across multiple + # send ops for the same param. + sendable_cache: dict[str, torch.Tensor] = {} + + def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: + if param_name not in sendable_cache: + sendable_cache[param_name] = _ensure_sendable(param) + return sendable_cache[param_name] - # Extract parameters and persistent buffers from models if present. - # Persistent buffers carry training state (e.g. MoE router expert_bias) - # and must be refit alongside parameters. - src_params = {} - dst_params = {} - if src_module is not None: - src_params = {name: p for name, p in named_refit_tensors(src_module)} - if dst_module is not None: - dst_params = {name: p for name, p in named_refit_tensors(dst_module)} - - # Cache dequantized BF16 views of MXFP8 source params so that multiple - # send ops for the same param reuse one dequant instead of repeating it. - _sendable_cache: dict[str, torch.Tensor] = {} - - def _get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: - if param_name not in _sendable_cache: - _sendable_cache[param_name] = _ensure_sendable(param) - return _sendable_cache[param_name] - - # Submit sends (only if we have source model) for op in plan.send_ops: + src_param = src_params.get(op.param_name) + if src_param is None: + continue if transform is not None and transform.should_transform(op.param_name): - src_param = src_params.get(op.param_name) - if src_param is not None: - tensors = transform.prepare_send(op.param_name, op.my_slice, src_param) - for t in tensors: - service.submit_send(t.contiguous(), op.peer_rank, task_id=op.task_id) + tensors = transform.prepare_send(op.param_name, op.my_slice, src_param) + for t in tensors: + service.submit_send(t.contiguous(), op.peer_rank, task_id=op.task_id) else: - src_param = src_params.get(op.param_name) - if src_param is not None: - sendable = _get_sendable(op.param_name, src_param) - src_view = sendable[op.my_slice] - # Only copy if the slice is non-contiguous. - if not src_view.is_contiguous(): - src_view = src_view.contiguous() - service.submit_send(src_view, op.peer_rank, task_id=op.task_id) - - # Free the dequant cache — slices have been submitted and the service - # holds its own references to the contiguous buffers it needs. - _sendable_cache.clear() - - # Submit recvs (only if we have destination model) - # Writebacks: each entry is one of: - # ('direct',) — recv'd in-place, no writeback - # ('default', recv_buffer, dst_param, dst_slice) — copy recv_buffer → dst_param - # ('transform', param_name, dst_slice, [recv_bufs]) — transform.finalize_recv - recv_writebacks: list = [] - - # Pre-allocate BF16 accumulation buffers for TE MXFP8 destination params so - # we can recv directly into views instead of allocating per-slice buffers. + sendable = get_sendable(op.param_name, src_param) + src_view = sendable[op.my_slice] + if not src_view.is_contiguous(): + src_view = src_view.contiguous() + service.submit_send(src_view, op.peer_rank, task_id=op.task_id) + + sendable_cache.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]] = {} for op in plan.recv_ops: @@ -105,107 +120,85 @@ def _get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: recv_bufs = transform.prepare_recv(op.param_name, op.my_slice) for buf in recv_bufs: service.submit_recv(buf, op.peer_rank, task_id=op.task_id) - recv_writebacks.append(('transform', op.param_name, op.my_slice, recv_bufs)) - else: - dst_param = dst_params.get(op.param_name) - if dst_param is not None: - # Try to recv directly into the destination parameter slice to - # avoid allocating a separate buffer + a writeback copy. This - # is safe when 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] - if dst_slice_view.is_contiguous() and not _is_mxfp8_tensor(dst_param): - # Recv directly into destination — no writeback needed. - service.submit_recv(dst_slice_view, op.peer_rank, task_id=op.task_id) - recv_writebacks.append(('direct',)) - elif _is_mxfp8_tensor(dst_param): - # TE MXFP8: recv directly into pre-allocated accumulation - # buffer to avoid per-slice BF16 allocations. - param_id = id(dst_param) - if param_id not in pending_quantized: - full_bf16 = torch.empty( - dst_param.shape, dtype=torch.bfloat16, device=dst_param.device - ) - pending_quantized[param_id] = (dst_param, full_bf16, []) - accum_view = pending_quantized[param_id][1][op.my_slice] - if accum_view.is_contiguous(): - service.submit_recv(accum_view, op.peer_rank, task_id=op.task_id) - recv_writebacks.append(('direct',)) - else: - recv_buffer = torch.empty_like(dst_slice_view.contiguous()) - service.submit_recv(recv_buffer, op.peer_rank, task_id=op.task_id) - recv_writebacks.append(('default', recv_buffer, dst_param, op.my_slice)) - else: - recv_buffer = torch.empty_like(dst_slice_view.contiguous()) - service.submit_recv(recv_buffer, op.peer_rank, task_id=op.task_id) - recv_writebacks.append(('default', recv_buffer, dst_param, op.my_slice)) - - # Execute + writebacks.append( + _Writeback( + kind='transform', + param_name=op.param_name, + dst_slice=op.my_slice, + recv_bufs=recv_bufs, + ) + ) + continue + + dst_param = dst_params.get(op.param_name) + if dst_param is None: + continue + + dst_slice_view = dst_param.data[op.my_slice] + dst_is_mxfp8 = is_mxfp8tensor(dst_param) + + if not dst_is_mxfp8 and dst_slice_view.is_contiguous(): + # Plain tensor: recv straight into the destination slice. + service.submit_recv(dst_slice_view, op.peer_rank, task_id=op.task_id) + writebacks.append(_Writeback(kind='direct')) + continue + + if dst_is_mxfp8: + full_bf16, _slices = _get_mxfp8_accumulator(pending_quantized, dst_param) + accum_view = full_bf16[op.my_slice] + if accum_view.is_contiguous(): + # Recv straight into the BF16 accumulator slice. + service.submit_recv(accum_view, op.peer_rank, task_id=op.task_id) + writebacks.append(_Writeback(kind='direct')) + continue + + # Fallback: stage into a temporary BF16 buffer. + recv_buffer = torch.empty_like(dst_slice_view.contiguous()) + service.submit_recv(recv_buffer, op.peer_rank, task_id=op.task_id) + writebacks.append( + _Writeback( + kind='copy', recv_buffer=recv_buffer, dst_param=dst_param, dst_slice=op.my_slice + ) + ) + logger.info(f"Executing {len(plan.send_ops)} sends + {len(plan.recv_ops)} recvs") service.run() torch.cuda.synchronize() dist.barrier(group=group) - # Write back received buffers into their destination parameter slices. - # - # 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. - # - # pending_quantized was pre-populated during recv submission so that - # contiguous MXFP8 slices recv'd directly into the accumulation buffer. - # Non-contiguous fallback slices are copied into it here. - for i in range(len(recv_writebacks)): - wb = recv_writebacks[i] - recv_writebacks[i] = None # Eagerly drop reference to free recv buffers + # Writebacks: ``direct`` already landed in place; ``transform`` hands off to + # the transform; ``copy`` copies the staging buffer into the destination + # slice (deferring MXFP8 accumulation to one quantize_() per param). + for i in range(len(writebacks)): + wb = writebacks[i] + writebacks[i] = None # Drop reference eagerly so recv buffers can free. with torch.no_grad(): - if wb[0] == 'direct': - # Already written in-place during recv — nothing to do. - pass - elif wb[0] == 'transform': - _, param_name, dst_slice, recv_bufs = wb - transform.finalize_recv(param_name, dst_slice, recv_bufs) + if wb.kind == 'direct': + continue + if wb.kind == 'transform': + transform.finalize_recv(wb.param_name, wb.dst_slice, wb.recv_bufs) + continue + # 'copy' + 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)) + full_bf16[wb.dst_slice].copy_(wb.recv_buffer) else: - _, recv_buffer, dst_param, dst_slice = wb - if _is_mxfp8_tensor(dst_param): - # Non-contiguous fallback: copy into pre-allocated accum buffer. - param_id = id(dst_param) - if param_id not in pending_quantized: - # Allocate empty BF16 buffer — no need to dequantize - # existing weights since all slices will be overwritten. - full_bf16 = torch.empty( - dst_param.shape, dtype=torch.bfloat16, device=dst_param.device - ) - pending_quantized[param_id] = (dst_param, full_bf16, []) - pending_quantized[param_id][2].append((dst_slice, recv_buffer)) - pending_quantized[param_id][1][dst_slice].copy_(recv_buffer) - else: - dst_param.data[dst_slice].copy_(recv_buffer) - # Free writeback list — recv_buffers are no longer needed after copy. - recv_writebacks.clear() - - # Finalize deferred quantized param updates - for param_id, (dst_param, full_bf16, slices) in pending_quantized.items(): + wb.dst_param.data[wb.dst_slice].copy_(wb.recv_buffer) + writebacks.clear() + + for _param_id, (dst_param, full_bf16, _slices) in pending_quantized.items(): with torch.no_grad(): dst_param.quantize_(full_bf16) - # Free the BF16 accumulation buffers eagerly. pending_quantized.clear() - # Ensure all writeback copies are visible to subsequent CUDA ops (e.g. CUDA - # graph warmup). The synchronize() above fires *before* the writeback loop, - # so without this second sync the .copy_() kernels are still async when - # execute_reshard_plan returns — creating a race with callers that immediately - # inspect or capture (via CUDA graphs) the destination parameters. + # Second sync: the writeback loop's .copy_() kernels are still async when + # execute_reshard_plan returns; without this CUDA-graph capture or callers + # that read params immediately race against the writes. torch.cuda.synchronize() - # 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). + # Release transient BF16 staging/accumulation buffers back to the CUDA + # driver. Significant for MXFP8 destinations (full model BF16 footprint). torch.cuda.empty_cache() logger.info("Reshard complete") diff --git a/megatron/core/resharding/planner.py b/megatron/core/resharding/planner.py index 242450bd835..b3008ed61da 100644 --- a/megatron/core/resharding/planner.py +++ b/megatron/core/resharding/planner.py @@ -73,6 +73,60 @@ def _build_descriptors_for_param( 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 _plan_multi_dim_lcm( param_name: str, src_metadata: ParameterMetadata, @@ -80,11 +134,7 @@ def _plan_multi_dim_lcm( descriptors: list[ShardingDescriptor], my_global_rank: int, ) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: - """ - TP-only planner using LCM tiling to support strides on source/destination. - - Requires exactly one TP descriptor - - Supports arbitrary integer strides (contiguous micro-tiles) - """ + """TP-only planner using LCM tiling with arbitrary integer src/dst strides.""" if not descriptors: return [] if len(descriptors) != 1: @@ -110,44 +160,25 @@ def _plan_multi_dim_lcm( f"(src_world={src_world}, src_local={src_local}, " f"dst_world={dst_world}, dst_local={dst_local})" ) - # LCM tiling with strides - Ns = src_world * max(1, d.src_stride) - Nd = dst_world * max(1, d.dst_stride) - full_len = dst_local * dst_world - g = math.gcd(Ns, Nd) - L = (Ns // g) * Nd - if full_len % L != 0: - raise RuntimeError( - f"{param_name}: TP dim{dim} full_len {full_len} not divisible by LCM {L} " - f"(Ns={Ns}, Nd={Nd})" - ) - unit = full_len // L # micro-tile length - cps = L // Ns # micro-tiles per source segment - cpd = L // Nd # micro-tiles per destination segment - seg_src = cps * unit # contiguous length per source segment - seg_dst = cpd * unit # contiguous length per destination segment - dst_local_rank = _get_rank_in_group(my_global_rank, d.dst_dim_ranks) - ops: list[tuple[int, tuple[slice, ...], tuple[slice, ...]]] = [] - # Sweep destination segments owned by this rank (handle destination stride) - for k in range(max(1, d.dst_stride)): - g_dst_seg = dst_local_rank + k * dst_world - # Within this segment, enumerate the cpd micro-tiles - for off in range(cpd): - g_micro = g_dst_seg * cpd + off - s_idx = g_micro // cps - in_seg = g_micro % cps - src_owner_in_dim = s_idx % src_world - src_global_rank = d.src_dim_ranks[src_owner_in_dim] - src_local_seg_idx = s_idx // src_world - src_start = src_local_seg_idx * seg_src + in_seg * unit - dst_start = k * seg_dst + off * unit - # Build full N-D slices - 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))) + ops: list[tuple[int, tuple[slice, ...], tuple[slice, ...]]] = [] + _emit_lcm_block_ops( + param_name=param_name, + src_shape=src_shape, + dst_shape=dst_shape, + dim=dim, + src_world=src_world, + dst_world=dst_world, + src_stride=d.src_stride, + dst_stride=d.dst_stride, + full_block_len=dst_local * dst_world, + dst_local_rank=_get_rank_in_group(my_global_rank, d.dst_dim_ranks), + src_dim_ranks=d.src_dim_ranks, + src_block_offset=0, + dst_block_offset=0, + block_label=f"TP dim{dim}", + ops=ops, + ) _sort_ops_by_dst_offset(ops, dim) return ops @@ -164,9 +195,8 @@ def _plan_block_interleaved( When a parameter packs multiple independently-sharded components of *different* sizes (e.g. Mamba in_proj packs z, x, B, C, dt), a simple - contiguous concat produces the wrong layout. This function treats each - block independently: it gathers (or scatters) each block across TP ranks - before moving to the next block. + contiguous concat produces the wrong layout. Each block is gathered + (or scattered) across TP ranks independently before moving to the next. ``partition_sizes`` lists the per-TP-rank block sizes along the partition dim. Block *i* occupies ``[sum(sizes[:i]), sum(sizes[:i+1]))`` in the @@ -187,14 +217,12 @@ def _plan_block_interleaved( dst_world = len(d.dst_dim_ranks) dst_local_rank = _get_rank_in_group(my_global_rank, d.dst_dim_ranks) - # Use partition_sizes from whichever side has it (prefer src) src_sizes = src_metadata.partition_sizes dst_sizes = dst_metadata.partition_sizes if src_sizes is None and dst_sizes is None: raise RuntimeError(f"{param_name}: _plan_block_interleaved called without partition_sizes") - # Derive the full (un-sharded) block sizes if src_sizes is not None: num_blocks = len(src_sizes) full_sizes = [s * src_world for s in src_sizes] @@ -202,13 +230,11 @@ def _plan_block_interleaved( num_blocks = len(dst_sizes) full_sizes = [s * dst_world for s in dst_sizes] - # Compute per-rank block sizes for both sides 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] - # Validate conservation for i in range(num_blocks): if src_sizes[i] * src_world != dst_sizes[i] * dst_world: raise RuntimeError( @@ -218,49 +244,28 @@ def _plan_block_interleaved( ) ops: list[tuple[int, tuple[slice, ...], tuple[slice, ...]]] = [] - - # For each block, compute the transfer ops independently - src_block_offset = 0 # cumulative offset in source local tensor - dst_block_offset = 0 # cumulative offset in destination local tensor - + src_block_offset = 0 + dst_block_offset = 0 for blk in range(num_blocks): - src_blk_sz = src_sizes[blk] # per-src-rank size of this block - dst_blk_sz = dst_sizes[blk] # per-dst-rank size of this block - full_blk_sz = full_sizes[blk] - - # Within this block, use simple LCM tiling (stride=1) - Ns = src_world - Nd = dst_world - g = math.gcd(Ns, Nd) - L = (Ns // g) * Nd - if full_blk_sz % L != 0: - raise RuntimeError( - f"{param_name}: block {blk} full_size {full_blk_sz} not divisible by LCM {L}" - ) - unit = full_blk_sz // L - cps = L // Ns - cpd = L // Nd - - # This dst rank's segment within the block - g_dst_seg = dst_local_rank - for off in range(cpd): - g_micro = g_dst_seg * cpd + off - s_idx = g_micro // cps - in_seg = g_micro % cps - src_owner_in_dim = s_idx % src_world - src_global_rank = d.src_dim_ranks[src_owner_in_dim] - src_local_seg_idx = s_idx // src_world - src_start = src_block_offset + src_local_seg_idx * (cps * unit) + in_seg * unit - dst_start = dst_block_offset + 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))) - - src_block_offset += src_blk_sz - dst_block_offset += dst_blk_sz + _emit_lcm_block_ops( + param_name=param_name, + src_shape=src_shape, + dst_shape=dst_shape, + dim=dim, + src_world=src_world, + dst_world=dst_world, + src_stride=1, + dst_stride=1, + full_block_len=full_sizes[blk], + dst_local_rank=dst_local_rank, + src_dim_ranks=d.src_dim_ranks, + src_block_offset=src_block_offset, + dst_block_offset=dst_block_offset, + block_label=f"block {blk}", + ops=ops, + ) + src_block_offset += src_sizes[blk] + dst_block_offset += dst_sizes[blk] _sort_ops_by_dst_offset(ops, dim) return ops @@ -294,23 +299,11 @@ def _finalize_dp_transfers( full_slice = tuple(slice(None) for _ in range(len(dst_shape))) return [(my_global_rank, full_slice, full_slice)] - # Different DP groups - use round-robin based on destination global rank for - # better load balancing across source ranks. This ensures that destination - # ranks are distributed across source ranks even when they have the same - # position within their respective DP groups. - # - # In non-collocated mode, src_dp_ranks might include ranks that don't - # have the source model (e.g., idle ranks or destination ranks). Filter to only - # include the rank that provided this metadata (src_metadata.owner_rank). - # src_metadata was selected by select_src_metadata_balanced, so owner_rank is the - # actual source rank for this parameter. - actual_src_rank = src_metadata.owner_rank - src_global_rank = src_dp_ranks[my_global_rank % len(src_dp_ranks)] - # Override with the actual source rank if the selected rank doesn't have the parameter - if src_global_rank != actual_src_rank: - src_global_rank = actual_src_rank + # 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_global_rank, full_slice, full_slice)] + return [(src_metadata.owner_rank, full_slice, full_slice)] def _determine_source_ranks_for_dst_param( @@ -406,12 +399,11 @@ def _extract_metadata(module, rank_offset): my_src_metadata = _extract_metadata(src_module, src_rank_offset) my_dst_metadata = _extract_metadata(dst_module, dst_rank_offset) - # Gather metadata to rank 0 only (not all ranks) to save CPU memory. - # Other ranks don't need the full metadata — they only need their own plan. - all_src_metadata_by_rank = [None] * world_size if my_global_rank == 0 else None - all_dst_metadata_by_rank = [None] * world_size if my_global_rank == 0 else None - dist.gather_object(my_src_metadata, all_src_metadata_by_rank, group_dst=0, group=group) - dist.gather_object(my_dst_metadata, all_dst_metadata_by_rank, group_dst=0, group=group) + # Gather (src, dst) tuples in one collective so we pay one pickle round-trip + # instead of two. Only rank 0 needs the full picture; other ranks just need + # their own plan from the later scatter. + gathered_pairs = [None] * world_size if my_global_rank == 0 else None + dist.gather_object((my_src_metadata, my_dst_metadata), gathered_pairs, group_dst=0, group=group) # Free local metadata — no longer needed after gather. del my_src_metadata, my_dst_metadata @@ -421,14 +413,13 @@ def _extract_metadata(module, rank_offset): src_param_metadata: dict[str, list[ParameterMetadata]] = {} if my_global_rank == 0: - for rank_id, rank_metadata_list in enumerate(all_dst_metadata_by_rank): - dst_param_metadata_by_rank[rank_id] = {m.resolved_name: m for m in rank_metadata_list} - for rank_metadata_list in all_src_metadata_by_rank: - for metadata in rank_metadata_list: + for rank_id, (src_meta_list, dst_meta_list) in enumerate(gathered_pairs): + dst_param_metadata_by_rank[rank_id] = {m.resolved_name: m for m in dst_meta_list} + for metadata in src_meta_list: src_param_metadata.setdefault(metadata.resolved_name, []).append(metadata) - # Free the raw gathered lists — data is now in the indexed dicts. - del all_src_metadata_by_rank, all_dst_metadata_by_rank + # Free the raw gathered list — data is now in the indexed dicts. + del gathered_pairs # Build the plan on global rank 0 and broadcast to all ranks if my_global_rank == 0: diff --git a/megatron/core/resharding/refit.py b/megatron/core/resharding/refit.py index 2cb6ba4479f..8574ef9b69e 100644 --- a/megatron/core/resharding/refit.py +++ b/megatron/core/resharding/refit.py @@ -27,7 +27,7 @@ from .copy_services.nccl_copy_service import NCCLCopyService from .copy_services.nvshmem_copy_service import NVSHMEMCopyService from .transforms import MXFP8ReshardTransform, ReshardTransform -from .utils import named_persistent_buffers +from .utils import invalidate_refit_tensor_cache, named_persistent_buffers # Supported refit backend names RefitBackendName = Literal["nccl", "gloo", "nvshmem"] @@ -44,55 +44,56 @@ class _PlanCacheKey: src_config: Optional[Tuple[int, int, int, int, int]] dst_config: Optional[Tuple[int, int, int, int, int]] num_experts: Optional[int] + # Rank offsets distinguish non-collocated configurations that would otherwise + # share the same (rank, sizes, num_experts) tuple but route to different + # global ranks. + src_rank_offset: int = 0 + dst_rank_offset: 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. + """Extract (TP, PP, EP, DP, expt_tp) sizes from a model core, memoized on the core. - Returns: - Tuple of (TP, PP, EP, DP, expt_tp) sizes, or None if core is None. - - TP: Tensor parallelism - - PP: Pipeline parallelism - - EP: Expert parallelism - - DP: Data parallelism - - expt_tp: Expert tensor parallelism + 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 + the hot path (each refit looks the key up 2-3x). """ if core is None: return None + cached = getattr(core, '_refit_config_tuple', None) + if cached is not None: + return cached pg = core.pg_collection - return ( - len(torch.distributed.get_process_group_ranks(pg.tp)) if pg.tp else 1, - len(torch.distributed.get_process_group_ranks(pg.pp)) if pg.pp else 1, - len(torch.distributed.get_process_group_ranks(pg.ep)) if pg.ep else 1, - len(torch.distributed.get_process_group_ranks(pg.dp)) if pg.dp else 1, - ( - len(torch.distributed.get_process_group_ranks(pg.expt_tp)) - if hasattr(pg, 'expt_tp') and pg.expt_tp - else 1 - ), + 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, ) + core._refit_config_tuple = result + return result def _build_plan_cache_key( - src_core, tgt_core, num_experts: Optional[int], group=None + src_core, + tgt_core, + num_experts: Optional[int], + group=None, + src_rank_offset: int = 0, + dst_rank_offset: int = 0, ) -> _PlanCacheKey: - """Build cache key for reshard plan. - - Args: - src_core: Source model core (or None for non-collocated destination/idle ranks) - tgt_core: Target model core (or None for non-collocated source/idle ranks) - num_experts: Number of MoE experts (or None for non-MoE models) - group: Optional process group for rank query - - Returns: - Cache key that uniquely identifies this reshard configuration for this rank - """ - # Use group.rank() to support cross-cluster ProcessGroups + """Build cache key for reshard plan.""" + # group.rank() supports cross-cluster ProcessGroups. rank = group.rank() if group is not None else torch.distributed.get_rank() - src_config = _get_config_tuple(src_core) - dst_config = _get_config_tuple(tgt_core) return _PlanCacheKey( - rank=rank, src_config=src_config, dst_config=dst_config, num_experts=num_experts + rank=rank, + src_config=_get_config_tuple(src_core), + dst_config=_get_config_tuple(tgt_core), + num_experts=num_experts, + src_rank_offset=src_rank_offset, + dst_rank_offset=dst_rank_offset, ) @@ -131,19 +132,12 @@ def clear_service_cache(): """Clear the cached refit services. Call this if you need to invalidate the cache, for example when - reinitializing distributed state. - - This properly finalizes services to free GPU buffers - before clearing the cache. + reinitializing distributed state. Services are ``close()``-d first so + backends owning GPU buffers (NVSHMEM) release them cleanly. """ global _service_cache - - # Finalize services to free resources for NVSHMEM backend - # NCCL/Gloo services have no cleanup needed - for backend_name, service in _service_cache.items(): - if hasattr(service, '_remote') and hasattr(service._remote, 'finalize'): - service._remote.finalize() - + for service in _service_cache.values(): + service.close() _service_cache.clear() @@ -204,7 +198,14 @@ def _build_or_get_plan(src_core, tgt_core, num_experts, group, src_rank_offset, yet cached, because build_centralized_reshard_plan uses collective communication. """ global _plan_cache - cache_key = _build_plan_cache_key(src_core, tgt_core, num_experts, group=group) + cache_key = _build_plan_cache_key( + src_core, + tgt_core, + num_experts, + group=group, + src_rank_offset=src_rank_offset, + dst_rank_offset=dst_rank_offset, + ) if cache_key not in _plan_cache: _plan_cache[cache_key] = build_centralized_reshard_plan( src_core, @@ -238,34 +239,28 @@ def _setup_mxfp8_transform_on_plan(plan, target_model) -> None: 2. Quantizes the target model's decoder weights to FlashInfer MXFP8Tensor (creating persistent buffers whose addresses are later captured by CUDA graphs). - 3. Builds an ``MXFP8ReshardTransform`` and attaches it to the plan as - ``plan.transform``. + 3. Builds an ``MXFP8ReshardTransform`` and attaches it to ``plan.transform``. - If the model doesn't need MXFP8, ``plan.transform`` is set to None. - Subsequent calls are no-ops if the plan already has a transform attribute. + Idempotent: skips re-setup if ``plan.transform`` is already populated. """ - if hasattr(plan, 'transform'): - return # Already set up + if plan.transform is not None: + return if not _needs_mxfp8_conversion(target_model): - plan.transform = None return lm = target_model[0] if isinstance(target_model, (list, tuple)) else target_model core = unwrap_model(lm) decoder = core.decoder if hasattr(core, 'decoder') else core - # 1. Compute which parameters are eligible for MXFP8 conversion. - # Must be done while params are still visible as nn.Parameter (BF16). + # Eligible params must be computed while still visible as nn.Parameter (BF16). convertible: set[str] = set() for name, param in decoder.named_parameters(): if _should_quantize_param(param): convertible.add(f"decoder.{name}") - # 2. Quantize decoder weights → persistent MXFP8Tensor buffers. persistent_buffers = quantize_params_to_mxfp8(decoder) - # 3. Build the transform and attach it to the plan. plan.transform = MXFP8ReshardTransform( convertible_params=convertible, persistent_buffers=persistent_buffers, @@ -350,12 +345,10 @@ def swap_model_weights( """ if isinstance(refit_method, str): service = get_or_create_service(refit_method, group=group) - elif hasattr(refit_method, 'submit_send') and hasattr(refit_method, 'run'): + elif isinstance(refit_method, CopyService): service = refit_method else: - raise TypeError( - "refit_method must be a str backend name or a CopyService-compatible instance" - ) + raise TypeError("refit_method must be a str backend name or a CopyService instance") # Auto-resolve MXFP8 transform from the cached plan when no # explicit transform was provided. @@ -364,7 +357,7 @@ def swap_model_weights( plan = _build_or_get_plan( src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset ) - transform = getattr(plan, 'transform', None) + transform = plan.transform reshard_model_weights( src_model, @@ -377,7 +370,7 @@ def swap_model_weights( ) -def _harmonize_buffer_dtypes(src_core, tgt_core, group=None): +def _harmonize_buffer_dtypes(plan, src_core, tgt_core, group=None): """Bring destination persistent-buffer dtypes into agreement with source. Some buffers (notably the MoE router ``expert_bias``) are upcast to fp32 @@ -387,40 +380,34 @@ def _harmonize_buffer_dtypes(src_core, tgt_core, group=None): sending fp32 bytes into a bf16 receive buffer corrupts the data — so dst's buffer must match src's dtype before the transfer. - Works for both collocated and non-collocated transfers: every rank reports - its source-side persistent-buffer dtypes via a single - ``all_gather_object`` on ``group``. Destination-side ranks then look up - each of their own buffers in the gathered map and replace the tensor with - one in src's dtype. Source-only and idle ranks contribute empty dicts and - skip the apply step, but still participate in the collective so it is - well-formed across every rank. - - Buffer matching is by raw module path (e.g. ``decoder.layers.0.…``); the - planner's PP-aware ``resolved_name`` is intentionally not used here because - we only need the dtype, which is uniform for a given buffer kind across - layers in practice. + The canonical dtype map is collected once via ``all_gather_object`` and + cached on the plan. Subsequent refits reuse the cached map and only do + the per-buffer dtype check / replacement (no collective). """ - # Build local map of source-side persistent buffer dtypes. - local_src_dtypes: dict[str, torch.dtype] = {} - if src_core is not None: - for full_name, _sub, _buf_name, buf in named_persistent_buffers(src_core): - local_src_dtypes[full_name] = buf.dtype - - world_size = group.size() if group is not None else torch.distributed.get_world_size() - gathered: list = [None] * world_size - torch.distributed.all_gather_object(gathered, local_src_dtypes, group=group) - - canonical: dict[str, torch.dtype] = {} - for d in gathered: - if not d: - continue - for name, dtype in d.items(): - # Replicated buffers agree across ranks; first writer wins. - canonical.setdefault(name, dtype) + if plan.buffer_dtypes is None: + local_src_dtypes: dict[str, torch.dtype] = {} + if src_core is not None: + for full_name, _sub, _buf_name, buf in named_persistent_buffers(src_core): + local_src_dtypes[full_name] = buf.dtype + + world_size = group.size() if group is not None else torch.distributed.get_world_size() + gathered: list = [None] * world_size + torch.distributed.all_gather_object(gathered, local_src_dtypes, group=group) + + canonical: dict[str, torch.dtype] = {} + for d in gathered: + if not d: + continue + for name, dtype in d.items(): + # Replicated buffers agree across ranks; first writer wins. + canonical.setdefault(name, dtype) + plan.buffer_dtypes = canonical if tgt_core is None: return + canonical = plan.buffer_dtypes + invalidated = False for full_name, sub, buf_name, dst_buf in named_persistent_buffers(tgt_core): expected = canonical.get(full_name) if expected is not None and dst_buf.dtype != expected: @@ -428,6 +415,9 @@ def _harmonize_buffer_dtypes(src_core, tgt_core, group=None): # recvs write the right number of bytes and the in-model lookup # (``self.expert_bias``) sees the new storage. sub._buffers[buf_name] = dst_buf.to(expected) + invalidated = True + if invalidated: + invalidate_refit_tensor_cache(tgt_core) def reshard_model_weights( @@ -454,10 +444,10 @@ def reshard_model_weights( transform: Optional ReshardTransform for custom format conversion. """ src_core, tgt_core, num_experts = _unwrap_model_cores(src_model, target_model) - _harmonize_buffer_dtypes(src_core, tgt_core, group=group) plan = _build_or_get_plan( src_core, tgt_core, num_experts, group, src_rank_offset, dst_rank_offset ) + _harmonize_buffer_dtypes(plan, src_core, tgt_core, group=group) execute_reshard_plan( plan, src_core, tgt_core, service=service, group=group, transform=transform ) diff --git a/megatron/core/resharding/transforms.py b/megatron/core/resharding/transforms.py index 213bef32c6e..de69a8d8c94 100644 --- a/megatron/core/resharding/transforms.py +++ b/megatron/core/resharding/transforms.py @@ -11,6 +11,7 @@ import torch +from megatron.core.fp8_utils import dequantize_fp8_tensor, is_mxfp8tensor from megatron.core.inference.quantization.mxfp8_tensor import MXFP8Tensor @@ -104,13 +105,8 @@ def _ensure_sendable(param: torch.Tensor) -> torch.Tensor: dequantized to their original precision (usually BF16). Standard parameters are returned via ``.data`` (unwrapped from autograd). """ - try: - from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor as _TEMXFP8 - - if isinstance(param, _TEMXFP8): - return param.dequantize() - except ImportError: - pass + 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 1c748f5ff98..443555681d3 100644 --- a/megatron/core/resharding/utils.py +++ b/megatron/core/resharding/utils.py @@ -1,12 +1,16 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. from __future__ import annotations +import re from dataclasses import dataclass -from typing import Mapping, Optional +from typing import TYPE_CHECKING, Mapping, Optional import torch import torch.distributed as dist +if TYPE_CHECKING: + from .transforms import ReshardTransform + # ----------------------------------------------------------------------------- # Dataclasses used by the planner # ----------------------------------------------------------------------------- @@ -99,6 +103,11 @@ class ReshardPlan: send_ops: list[TransferOp] recv_ops: list[TransferOp] + transform: Optional["ReshardTransform"] = None + # Cache of canonical persistent-buffer dtypes keyed by raw module path. + # Populated by _harmonize_buffer_dtypes on first call; reused thereafter to + # skip the all_gather_object + named_modules() walks on the hot path. + buffer_dtypes: Optional[dict[str, torch.dtype]] = None def __str__(self): return f"ReshardPlan(sends={len(self.send_ops)}, recvs={len(self.recv_ops)})" @@ -119,17 +128,15 @@ def _get_rank_in_group(global_rank: int, group_ranks: list[int]) -> int: ) +_EXPERT_PARAM_RE = re.compile(r"^(weight|bias)(\d+)$") + + def _detect_expert_index_from_param_name(param_name: str) -> Optional[int]: """Extract expert index from parameter name for TEGroupedMLP per-expert tensors.""" for part in param_name.split('.'): - if ( - part.startswith('weight') - and len(part) > len('weight') - and part[len('weight') :].isdigit() - ): - return int(part[len('weight') :]) - if part.startswith('bias') and len(part) > len('bias') and part[len('bias') :].isdigit(): - return int(part[len('bias') :]) + m = _EXPERT_PARAM_RE.match(part) + if m is not None: + return int(m.group(2)) return None @@ -167,15 +174,10 @@ def assign_ep_resolved_name_inplace( meta.global_expert_index = global_idx # Replace trailing integer in "weightK"/"biasK" with global_idx - parts = base.split('.') new_parts = [] - for p in parts: - if p.startswith('weight') and len(p) > len('weight') and p[len('weight') :].isdigit(): - new_parts.append('weight' + str(global_idx)) - elif p.startswith('bias') and len(p) > len('bias') and p[len('bias') :].isdigit(): - new_parts.append('bias' + str(global_idx)) - else: - new_parts.append(p) + for p in base.split('.'): + m = _EXPERT_PARAM_RE.match(p) + new_parts.append(f"{m.group(1)}{global_idx}" if m else p) meta.resolved_name = '.'.join(new_parts) @@ -229,6 +231,31 @@ def named_refit_tensors(module: torch.nn.Module): yield full_name, buf +_REFIT_TENSOR_CACHE_ATTR = "_refit_tensor_cache" + + +def get_refit_tensor_dict(module: torch.nn.Module) -> dict[str, torch.Tensor]: + """Return the cached ``{name: tensor}`` dict for ``module``, building it if needed. + + Walking ``named_modules()`` is hundreds of ms for multi-B-parameter models, + and the parameter/persistent-buffer set is stable across refits — so we + cache the dict on the module itself. ``invalidate_refit_tensor_cache`` + must be called whenever a persistent buffer is replaced (e.g. by + ``_harmonize_buffer_dtypes``) so the cache picks up the new tensor. + """ + cached = getattr(module, _REFIT_TENSOR_CACHE_ATTR, None) + if cached is None: + cached = dict(named_refit_tensors(module)) + setattr(module, _REFIT_TENSOR_CACHE_ATTR, cached) + return cached + + +def invalidate_refit_tensor_cache(module: torch.nn.Module) -> None: + """Drop the cached refit tensor dict so the next call rebuilds it.""" + if hasattr(module, _REFIT_TENSOR_CACHE_ATTR): + delattr(module, _REFIT_TENSOR_CACHE_ATTR) + + def _build_layer_module_prefix_map(module: torch.nn.Module) -> dict[str, str]: """Build a mapping local_module_prefix -> global_module_prefix for PP layer modules. @@ -412,134 +439,86 @@ def _offset_ranks(ranks: list[int]) -> list[int]: return meta -def select_src_metadata_balanced( +def _filter_by_ep_local_rank( src_meta_list: list[ParameterMetadata], dst_metadata: ParameterMetadata, dst_rank: int -) -> ParameterMetadata: - """Choose a representative source `ParameterMetadata` for a destination rank. +) -> list[ParameterMetadata]: + """In non-collocated mode with matching EP size, restrict candidates to the + source rank holding the same global experts as ``dst_rank``. - The selected metadata provides topology information (TP/EP/DP group ranks) that the - LCM transfer planner uses to compute actual source ranks and slices. This function - doesn't perform transfers itself - it just picks which source configuration to use - as reference for planning. - - Two scenarios for EP-sharded parameters: - 1. Non-collocated mode (same EP size, different rank numbering): - - Filter by matching EP local rank to pair ranks with same expert position - - Example: src ranks [0-63] and dst ranks [64-127] both with EP=8 - - Dst EP local 0 should use src EP local 0 as reference (same experts) - - 2. Resharding mode (different EP sizes): - - Skip EP local rank filtering (sizes don't correspond) - - Example: EP=8→EP=16 means dst EP local 8 has no matching src EP local - - Expert matching handled by resolved_name; LCM handles TP dimension changes - - Finally, balances across data-parallel (DP) groups to distribute load: - - Groups src_meta_list by DP group - - Selects source DP group via round-robin: dst_rank % num_src_dp_groups - - Ensures even distribution of transfer load across source DP replicas + When EP sizes differ (resharding), expert matching is handled via + ``resolved_name`` and no filter is applied here. """ - if not src_meta_list: - raise ValueError("src_meta_list must be non-empty") - - # ============================================================================ - # EXPERT PARALLELISM (EP) LOCAL RANK FILTERING - # ============================================================================ - # Purpose: In non-collocated mode with same EP size, ensure destination ranks - # use source metadata from ranks with the same EP local position (same experts). - # - # Why size check matters: - # - Same size (EP=8→EP=8): Local ranks 0-7 exist in both src and dst - # → Filter ensures dst EP local 0 uses src EP local 0 (same global experts) - # - Different size (EP=8→EP=16): Local ranks 0-15 in dst, only 0-7 in src - # → Dst EP local 8 has no corresponding src EP local rank - # → Skip filter; expert reassignment handled by resolved_name matching - # - # Expert routing: When EP size changes, each expert parameter is matched via - # resolved_name (which includes global expert index). The LCM/TP planner - # handles any TP dimension changes, and DP round-robin distributes load. - # ============================================================================ dst_ep_group = dst_metadata.expert_parallel_group_ranks - if dst_ep_group is not None: - dst_ep_local = dst_ep_group.index(dst_rank) - # Check if EP sizes match between source and destination - src_ep_size = ( - len(src_meta_list[0].expert_parallel_group_ranks) - if src_meta_list[0].expert_parallel_group_ranks - else None + if dst_ep_group is None: + return src_meta_list + + dst_ep_local = dst_ep_group.index(dst_rank) + src_ep_size = ( + len(src_meta_list[0].expert_parallel_group_ranks) + if src_meta_list[0].expert_parallel_group_ranks + else None + ) + if src_ep_size != len(dst_ep_group): + return src_meta_list + + matching = [ + m + for m in src_meta_list + if m.expert_parallel_group_ranks + and m.expert_parallel_group_ranks.index(m.owner_rank) == dst_ep_local + ] + if not matching: + available = [ + ( + m.owner_rank, + ( + m.expert_parallel_group_ranks.index(m.owner_rank) + if m.expert_parallel_group_ranks + else None + ), + ) + for m in src_meta_list + ] + raise ValueError( + f"No source metadata with EP local rank {dst_ep_local}" + f" found for dst rank {dst_rank}. Available: {available}" ) - dst_ep_size = len(dst_ep_group) - - # Only filter by EP local rank when sizes match (non-collocated, not resharding) - if src_ep_size == dst_ep_size: - matching_ep = [ - m - for m in src_meta_list - if m.expert_parallel_group_ranks - and m.expert_parallel_group_ranks.index(m.owner_rank) == dst_ep_local - ] - if not matching_ep: - # This indicates a configuration bug: sizes match but no local rank match - def _ep_local(m): - return ( - m.expert_parallel_group_ranks.index(m.owner_rank) - if m.expert_parallel_group_ranks - else None - ) - - available = [(m.owner_rank, _ep_local(m)) for m in src_meta_list] - raise ValueError( - f"No source metadata with EP local rank {dst_ep_local}" - f" found for dst rank {dst_rank}. Available: {available}" - ) - src_meta_list = matching_ep - # else: EP resharding mode (sizes differ) - skip filter, keep all source candidates - - # ============================================================================ - # LOCAL COPY OPTIMIZATION (COLLOCATED MODE) - # ============================================================================ - # In collocated mode, prefer local copies when available. If dst_rank appears - # in the source metadata list (after TP/EP filtering), use it directly to - # avoid unnecessary data transfers. - # - # A local copy is essentially free - # (tensor.copy_() on same GPU), while any remote transfer incurs significant - # overhead even within the same node. - # ============================================================================ - local_meta = [m for m in src_meta_list if m.owner_rank == dst_rank] - if local_meta: - # Found local metadata - use it for a free local copy - return local_meta[0] - - # ============================================================================ - # DATA PARALLELISM (DP) LOAD BALANCING - # ============================================================================ - # After TP/EP filtering (if applicable), balance transfer load across source - # data-parallel replicas. Each DP group holds a complete copy of the model, - # so we can read from any DP group - choosing via round-robin spreads load. - # - # Load distribution: dst_rank % num_src_dp_groups ensures even distribution - # even when destination has different DP configuration than source. - # ============================================================================ + return matching + + +def _round_robin_dp(src_meta_list: list[ParameterMetadata], dst_rank: int) -> ParameterMetadata: + """Round-robin across source DP groups so transfer load spreads evenly.""" grouped_by_dp: dict[tuple[int, ...], list[ParameterMetadata]] = {} for meta in src_meta_list: dp_group = tuple(meta.data_parallel_group_ranks or []) grouped_by_dp.setdefault(dp_group, []).append(meta) - # Fast path: only one DP group present; no balancing necessary if len(grouped_by_dp) == 1: return src_meta_list[0] - # Round-robin selection across source DP groups based on destination global rank - # This ensures even distribution: if we have 4 src DP groups and 128 dst ranks, - # each src DP group will be selected by 32 dst ranks (128 / 4 = 32) sorted_dp_groups = sorted(grouped_by_dp.keys()) chosen_group = sorted_dp_groups[dst_rank % len(sorted_dp_groups)] - - # Within the chosen DP group, distribute across available metadata entries - # to balance load across all TP groups in the DP replica. - # Example: With 4 TP groups in a DP group, dst_ranks will cycle through all 4 - # instead of always using the first one, better distributing transfer load. group_metadata = grouped_by_dp[chosen_group] within_group_idx = (dst_rank // len(sorted_dp_groups)) % len(group_metadata) - selected = group_metadata[within_group_idx] - return selected + return group_metadata[within_group_idx] + + +def select_src_metadata_balanced( + src_meta_list: list[ParameterMetadata], dst_metadata: ParameterMetadata, dst_rank: int +) -> 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. + """ + if not src_meta_list: + raise ValueError("src_meta_list must be non-empty") + + candidates = _filter_by_ep_local_rank(src_meta_list, dst_metadata, dst_rank) + + for meta in candidates: + if meta.owner_rank == dst_rank: + return meta + + return _round_robin_dp(candidates, dst_rank) diff --git a/tests/unit_tests/resharding/test_mxfp8_refit.py b/tests/unit_tests/resharding/test_mxfp8_refit.py index 815d4eeedac..23f02a88ac1 100644 --- a/tests/unit_tests/resharding/test_mxfp8_refit.py +++ b/tests/unit_tests/resharding/test_mxfp8_refit.py @@ -128,6 +128,23 @@ def test_finalize_recv_1d_scale_wrong_element_count(self): # =========================================================================== +def _pre_quantize_linear(model: torch.nn.Module) -> None: + """Replace every Linear's BF16 weight with an ``nn.Parameter`` wrapping a + Transformer-Engine MXFP8 tensor. ``quantize_params_to_mxfp8`` accepts + inputs whose ``.data`` is a TEMXFP8Tensor; it does not accept plain BF16 + ``nn.Parameter`` (production callers wrap weights via TE's ``fp8_param`` + machinery before calling this function). + """ + import transformer_engine_torch as tex + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + for submodule in model.modules(): + if isinstance(submodule, torch.nn.Linear): + te_mxfp8 = quantizer(submodule.weight.data) + submodule.weight = torch.nn.Parameter(te_mxfp8, requires_grad=False) + + class TestQuantizeParamsToMXFP8: """Tests for persistent buffer quantization (quantization/utils.py). @@ -140,6 +157,7 @@ def test_basic_quantization_replaces_param(self): from megatron.core.inference.quantization.utils import quantize_params_to_mxfp8 model = torch.nn.Linear(128, 64, bias=False).to(dtype=torch.bfloat16, device="cuda") + _pre_quantize_linear(model) buffers = quantize_params_to_mxfp8(model) assert "weight" in buffers @@ -152,11 +170,13 @@ def test_persistent_buffer_reuse_preserves_addresses(self): from megatron.core.inference.quantization.utils import quantize_params_to_mxfp8 model = torch.nn.Linear(128, 64, bias=False).to(dtype=torch.bfloat16, device="cuda") + _pre_quantize_linear(model) buffers = quantize_params_to_mxfp8(model) data_ptr = buffers["weight"].data.data_ptr() scale_ptr = buffers["weight"].scale.data_ptr() model2 = torch.nn.Linear(128, 64, bias=False).to(dtype=torch.bfloat16, device="cuda") + _pre_quantize_linear(model2) quantize_params_to_mxfp8(model2, persistent_buffers=buffers) assert buffers["weight"].data.data_ptr() == data_ptr @@ -170,6 +190,7 @@ def test_nested_module_fqn(self): model = torch.nn.Sequential( torch.nn.Linear(128, 64, bias=False), torch.nn.Linear(64, 32, bias=False) ).to(dtype=torch.bfloat16, device="cuda") + _pre_quantize_linear(model) buffers = quantize_params_to_mxfp8(model) assert "0.weight" in buffers and "1.weight" in buffers From 0b0f0890762e5227394c9d40fc61fb658b5b2fcd Mon Sep 17 00:00:00 2001 From: William Dykas Date: Tue, 12 May 2026 15:19:35 -0700 Subject: [PATCH 2/4] test --- .../copy_services/gloo_copy_service.py | 42 ++++++++++++------- .../copy_services/nvshmem_copy_service.py | 6 ++- megatron/core/resharding/execution.py | 35 ++++++++++++++-- 3 files changed, 63 insertions(+), 20 deletions(-) diff --git a/megatron/core/resharding/copy_services/gloo_copy_service.py b/megatron/core/resharding/copy_services/gloo_copy_service.py index 5846a501e77..89afb075687 100644 --- a/megatron/core/resharding/copy_services/gloo_copy_service.py +++ b/megatron/core/resharding/copy_services/gloo_copy_service.py @@ -28,6 +28,10 @@ def __init__(self, group=None): # Each recv op is paired with its GPU destination tensor; the SendOp/RecvOp # itself carries a pinned-CPU staging buffer for Gloo's CPU PG. self.recv_ops: List[Tuple[RecvOp, torch.Tensor]] = [] + # Dedicated stream for GPU-side work (local same-rank copies and the + # final CPU->GPU writebacks) so they overlap with the GPU->CPU staging + # copies issued on the default stream during ``run()``. + self._copy_stream = torch.cuda.Stream() if self.rank == 0: logger.info( f"GlooCopyService initialized on rank {self.rank} with {self.world_size} ranks" @@ -58,13 +62,15 @@ def run(self): local_recvs = [(recv, dst) for (recv, dst) in self.recv_ops if recv.src_rank == self.rank] remote_recvs = [(recv, dst) for (recv, dst) in self.recv_ops if recv.src_rank != self.rank] + # Local copies run on a dedicated stream so they overlap with the + # GPU->CPU staging copies issued on the default stream below. if local_sends or local_recvs: local_recv_objs = [recv for recv, _ in local_recvs] dst_by_task_id = {recv.task_id: dst for recv, dst in local_recvs} pairs = match_local_ops_by_task_id( local_sends, local_recv_objs, "GlooCopyService", self.rank ) - with torch.no_grad(): + with torch.no_grad(), torch.cuda.stream(self._copy_stream): for send_op, recv_op in pairs: src_tensor = send_op.tensor dst_tensor = dst_by_task_id[recv_op.task_id] @@ -73,11 +79,11 @@ def run(self): else: dst_tensor.copy_(src_tensor) - # Build Gloo P2P ops over CPU tensors. For sends we stage all - # GPU→CPU copies with non_blocking, sync once, then build P2P ops. - # Use group_peer (not peer) to pass ranks directly in group space, - # avoiding the global-to-group rank conversion in P2POp which doesn't - # work for cross-world ProcessGroups. + # Stage all remote sends GPU->CPU on the default stream with non_blocking + # copies, then wait on a single event so other streams (e.g. _copy_stream) + # aren't blocked. Use group_peer (not peer) to pass ranks directly in + # group space, avoiding the global-to-group rank conversion in P2POp + # which doesn't work for cross-world ProcessGroups. cpu_send_bufs: List[torch.Tensor] = [] for op in remote_sends: cpu_tensor = torch.empty( @@ -86,7 +92,9 @@ def run(self): cpu_tensor.copy_(op.tensor.detach(), non_blocking=True) cpu_send_bufs.append(cpu_tensor) if cpu_send_bufs: - torch.cuda.synchronize() + stage_event = torch.cuda.Event() + stage_event.record() # default stream + stage_event.synchronize() for op in remote_sends: # Drop the GPU reference now that staging is complete. @@ -107,14 +115,18 @@ def run(self): for req in reqs: req.wait() - for recv, dst_tensor in remote_recvs: - if dst_tensor.is_cuda: - dst_tensor.copy_(recv.tensor, non_blocking=True) - else: - dst_tensor.copy_(recv.tensor) - - # Ensure all async CPU→GPU copies are complete. - torch.cuda.synchronize() + # CPU->GPU writebacks run on _copy_stream so they pipeline with each + # other and don't block the default stream. + with torch.cuda.stream(self._copy_stream): + for recv, dst_tensor in remote_recvs: + if dst_tensor.is_cuda: + dst_tensor.copy_(recv.tensor, non_blocking=True) + else: + dst_tensor.copy_(recv.tensor) + + # Join _copy_stream into the default stream so subsequent default-stream + # work sees the local copies and the writebacks. + torch.cuda.current_stream().wait_stream(self._copy_stream) if self.rank == 0: logger.info("GlooCopyService: batched communication completed") diff --git a/megatron/core/resharding/copy_services/nvshmem_copy_service.py b/megatron/core/resharding/copy_services/nvshmem_copy_service.py index 9260dfd374c..a8cf14b8a6c 100644 --- a/megatron/core/resharding/copy_services/nvshmem_copy_service.py +++ b/megatron/core/resharding/copy_services/nvshmem_copy_service.py @@ -142,7 +142,6 @@ def run(self): ) dst.copy_(src, non_blocking=True) - torch.cuda.current_stream().wait_stream(self._local_copy_stream) self._local_send_ops.clear() self._local_recv_ops.clear() @@ -151,8 +150,13 @@ def run(self): # - schedule() uses dist.all_gather_object() # - run() uses nvshmem.core.barrier_all() # Critical for non-collocated refit where some ranks may have no work. + # Local copies on `_local_copy_stream` run concurrently with this remote + # NVSHMEM pipeline because the join below happens after `run()`. logger.info("NVSHMEMCopyService: building NVSHMEM schedule and executing") self._remote.schedule() self._remote.run() self._remote.clear_requests() + + # Join local-copy stream after remote pipeline so they overlapped. + torch.cuda.current_stream().wait_stream(self._local_copy_stream) logger.info("NVSHMEMCopyService: NVSHMEM transfers complete") diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py index c548c49aad1..2c13d418a15 100644 --- a/megatron/core/resharding/execution.py +++ b/megatron/core/resharding/execution.py @@ -85,8 +85,28 @@ def execute_reshard_plan( dst_params = get_refit_tensor_dict(dst_module) if dst_module is not None else {} # Dequantized BF16 views of MXFP8 source params are reused across multiple - # send ops for the same param. + # send ops for the same param. 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 default stream). sendable_cache: dict[str, torch.Tensor] = {} + sendable_events: dict[str, torch.cuda.Event] = {} + + mxfp8_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 mxfp8_param_names: + prefetch_stream = torch.cuda.Stream() + with torch.cuda.stream(prefetch_stream): + for param_name in mxfp8_param_names: + sendable_cache[param_name] = _ensure_sendable(src_params[param_name]) + ev = torch.cuda.Event() + ev.record() + sendable_events[param_name] = ev def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: if param_name not in sendable_cache: @@ -102,6 +122,9 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: for t in tensors: service.submit_send(t.contiguous(), op.peer_rank, task_id=op.task_id) else: + ev = sendable_events.get(op.param_name) + if ev is not None: + torch.cuda.current_stream().wait_event(ev) sendable = get_sendable(op.param_name, src_param) src_view = sendable[op.my_slice] if not src_view.is_contiguous(): @@ -109,6 +132,7 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: service.submit_send(src_view, op.peer_rank, task_id=op.task_id) sendable_cache.clear() + sendable_events.clear() writebacks: list[_Writeback] = [] # Maps id(dst_param) -> (dst_param, full_bf16, slices) for MXFP8 dests that @@ -187,6 +211,7 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: wb.dst_param.data[wb.dst_slice].copy_(wb.recv_buffer) writebacks.clear() + had_mxfp8_staging = bool(pending_quantized) for _param_id, (dst_param, full_bf16, _slices) in pending_quantized.items(): with torch.no_grad(): dst_param.quantize_(full_bf16) @@ -197,8 +222,10 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: # that read params immediately race against the writes. torch.cuda.synchronize() - # Release transient BF16 staging/accumulation buffers back to the CUDA - # driver. Significant for MXFP8 destinations (full model BF16 footprint). - torch.cuda.empty_cache() + # MXFP8 destinations allocate a full-model-sized BF16 staging buffer that + # can dwarf the rest of the working set. Reclaim it back to the driver + # when present; skip the (expensive) empty_cache walk otherwise. + if had_mxfp8_staging: + torch.cuda.empty_cache() logger.info("Reshard complete") From 76b0025815fd6dd4eb2efda249233f808cf6c83e Mon Sep 17 00:00:00 2001 From: William Dykas Date: Wed, 13 May 2026 05:58:56 -0700 Subject: [PATCH 3/4] revert --- .../copy_services/gloo_copy_service.py | 25 ++- .../copy_services/nccl_copy_service.py | 1 + .../copy_services/nvshmem_copy_service.py | 20 +- megatron/core/resharding/execution.py | 54 ++++-- megatron/core/resharding/planner.py | 183 +++++++----------- megatron/core/resharding/refit.py | 5 +- megatron/core/resharding/utils.py | 25 ++- tests/unit_tests/resharding/test_planner.py | 52 ++--- 8 files changed, 183 insertions(+), 182 deletions(-) diff --git a/megatron/core/resharding/copy_services/gloo_copy_service.py b/megatron/core/resharding/copy_services/gloo_copy_service.py index 89afb075687..b22ea91c851 100644 --- a/megatron/core/resharding/copy_services/gloo_copy_service.py +++ b/megatron/core/resharding/copy_services/gloo_copy_service.py @@ -41,7 +41,7 @@ def submit_send(self, src_tensor: torch.Tensor, dest_rank: int, task_id: Optiona self.send_ops.append(SendOp(task_id=task_id, tensor=src_tensor, dest_rank=dest_rank)) def submit_recv(self, dest_tensor: torch.Tensor, src_rank: int, task_id: Optional[int] = None): - # Pinned CPU staging buffer for the Gloo recv; copied back to dest_tensor in run(). + # Allocate a pinned CPU buffer for faster CPU↔GPU transfer. cpu_buffer = torch.empty( dest_tensor.shape, dtype=dest_tensor.dtype, device="cpu", pin_memory=True ) @@ -79,11 +79,11 @@ def run(self): else: dst_tensor.copy_(src_tensor) - # Stage all remote sends GPU->CPU on the default stream with non_blocking - # copies, then wait on a single event so other streams (e.g. _copy_stream) - # aren't blocked. Use group_peer (not peer) to pass ranks directly in - # group space, avoiding the global-to-group rank conversion in P2POp - # which doesn't work for cross-world ProcessGroups. + # Build Gloo P2P ops over CPU tensors. For sends we stage all + # GPU→CPU copies with non_blocking, sync once, then build P2P ops. + # Use group_peer (not peer) to pass ranks directly in group space, + # avoiding the global-to-group rank conversion in P2POp which doesn't + # work for cross-world ProcessGroups. cpu_send_bufs: List[torch.Tensor] = [] for op in remote_sends: cpu_tensor = torch.empty( @@ -91,10 +91,9 @@ def run(self): ) cpu_tensor.copy_(op.tensor.detach(), non_blocking=True) cpu_send_bufs.append(cpu_tensor) + # Wait only on default-stream staging copies; _copy_stream keeps running. if cpu_send_bufs: - stage_event = torch.cuda.Event() - stage_event.record() # default stream - stage_event.synchronize() + torch.cuda.current_stream().synchronize() for op in remote_sends: # Drop the GPU reference now that staging is complete. @@ -115,8 +114,9 @@ def run(self): for req in reqs: req.wait() - # CPU->GPU writebacks run on _copy_stream so they pipeline with each - # other and don't block the default stream. + # Copy received CPU buffers back into the original destination tensors. + # Use non_blocking with pinned memory for overlap. Routed through + # _copy_stream so subsequent default-stream work isn't blocked. with torch.cuda.stream(self._copy_stream): for recv, dst_tensor in remote_recvs: if dst_tensor.is_cuda: @@ -124,8 +124,7 @@ def run(self): else: dst_tensor.copy_(recv.tensor) - # Join _copy_stream into the default stream so subsequent default-stream - # work sees the local copies and the writebacks. + # Ensure all async CPU→GPU copies are complete and local copies have landed. torch.cuda.current_stream().wait_stream(self._copy_stream) if self.rank == 0: diff --git a/megatron/core/resharding/copy_services/nccl_copy_service.py b/megatron/core/resharding/copy_services/nccl_copy_service.py index 9fc25e401cf..a1b63334bbf 100644 --- a/megatron/core/resharding/copy_services/nccl_copy_service.py +++ b/megatron/core/resharding/copy_services/nccl_copy_service.py @@ -68,6 +68,7 @@ def run(self): for req in reqs: req.wait() + # Make sure the copy stream is finished torch.cuda.current_stream().wait_stream(self._copy_stream) if self.rank == 0: diff --git a/megatron/core/resharding/copy_services/nvshmem_copy_service.py b/megatron/core/resharding/copy_services/nvshmem_copy_service.py index a8cf14b8a6c..4e82d60de11 100644 --- a/megatron/core/resharding/copy_services/nvshmem_copy_service.py +++ b/megatron/core/resharding/copy_services/nvshmem_copy_service.py @@ -27,10 +27,11 @@ def __init__(self, group=None): super().__init__(group=group) self._remote = RemoteCopyService(group=group) + # Lazily initialized on first use to avoid side effects at import time self._initialized = False - # Keep original typed tensors (not uint8 views) so local copies preserve - # shape/strides semantics and avoid byte-offset pitfalls. + # NOTE: keep the original typed tensors here (not uint8 views) so local copies + # preserve shape/strides semantics and avoid byte-offset pitfalls. self._local_send_ops: Dict[int, torch.Tensor] = {} self._local_recv_ops: Dict[int, torch.Tensor] = {} self._local_copy_stream = torch.cuda.Stream() @@ -60,6 +61,7 @@ def submit_send(self, src_tensor: torch.Tensor, dest_rank: int, task_id: Optiona if not src_tensor.is_contiguous(): src_tensor = src_tensor.contiguous() + # Local transfers: keep them out of RemoteCopyService entirely. if dest_rank == self.rank: self._local_send_ops[task_id] = src_tensor return @@ -89,6 +91,7 @@ def submit_recv(self, dest_tensor: torch.Tensor, src_rank: int, task_id: Optiona if not dest_tensor.is_contiguous(): dest_tensor = dest_tensor.contiguous() + # Local transfers: keep them out of RemoteCopyService entirely. if src_rank == self.rank: self._local_recv_ops[task_id] = dest_tensor return @@ -118,7 +121,7 @@ def run(self): """ self._ensure_initialized() - # Local copies match by task_id (the NVSHMEM RemoteCopyService never sees them). + # 1) Run same-rank copies (match by task_id), like NCCL backend. if self._local_send_ops or self._local_recv_ops: missing_sends = set(self._local_recv_ops.keys()) - set(self._local_send_ops.keys()) missing_recvs = set(self._local_send_ops.keys()) - set(self._local_recv_ops.keys()) @@ -145,11 +148,12 @@ def run(self): self._local_send_ops.clear() self._local_recv_ops.clear() - # ALL ranks must call schedule() and run() because they contain collectives - # that require all ranks to participate: - # - schedule() uses dist.all_gather_object() - # - run() uses nvshmem.core.barrier_all() - # Critical for non-collocated refit where some ranks may have no work. + # 2) Execute remote schedule (if any remote sends/recvs were registered). + # NOTE: ALL ranks must call schedule() and run() because they contain collective + # operations that require all ranks to participate: + # - schedule() has dist.all_gather_object() (torch distributed collective) + # - run() has nvshmem.core.barrier_all() (nvshmem collective) + # This is critical for non-collocated refit where some ranks may have no work. # Local copies on `_local_copy_stream` run concurrently with this remote # NVSHMEM pipeline because the join below happens after `run()`. logger.info("NVSHMEMCopyService: building NVSHMEM schedule and executing") diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py index 2c13d418a15..8b06ef1cd14 100644 --- a/megatron/core/resharding/execution.py +++ b/megatron/core/resharding/execution.py @@ -79,15 +79,18 @@ def execute_reshard_plan( transform's prepare_send / prepare_recv / finalize_recv methods instead of the default slice-and-copy logic. """ - # Refit tensors (parameters + persistent buffers) are cached on each module - # so the named_modules() walk happens once per model, not per refit. + # Extract parameters and persistent buffers from models if present. + # Persistent buffers carry training state (e.g. MoE router expert_bias) + # and must be refit alongside parameters. Cached on each module so the + # named_modules() walk happens once per model, not per refit. 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 {} - # Dequantized BF16 views of MXFP8 source params are reused across multiple - # send ops for the same param. 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 default stream). + # Cache dequantized BF16 views of MXFP8 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] = {} @@ -158,25 +161,29 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: if dst_param is None: continue + # Try to recv directly into the destination parameter slice to avoid + # allocating a separate buffer + a writeback copy. This is safe when + # 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) if not dst_is_mxfp8 and dst_slice_view.is_contiguous(): - # Plain tensor: recv straight into the destination slice. 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 + # buffer to avoid per-slice BF16 allocations. full_bf16, _slices = _get_mxfp8_accumulator(pending_quantized, dst_param) accum_view = full_bf16[op.my_slice] if accum_view.is_contiguous(): - # Recv straight into the BF16 accumulator slice. service.submit_recv(accum_view, op.peer_rank, task_id=op.task_id) writebacks.append(_Writeback(kind='direct')) continue - # Fallback: stage into a temporary BF16 buffer. + # Fallback: stage into a temporary BF16 buffer (non-contiguous slice). recv_buffer = torch.empty_like(dst_slice_view.contiguous()) service.submit_recv(recv_buffer, op.peer_rank, task_id=op.task_id) writebacks.append( @@ -190,9 +197,14 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: torch.cuda.synchronize() dist.barrier(group=group) - # Writebacks: ``direct`` already landed in place; ``transform`` hands off to - # the transform; ``copy`` copies the staging buffer into the destination - # slice (deferring MXFP8 accumulation to one quantize_() per param). + # Write back received buffers into their destination parameter slices. + # + # 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. for i in range(len(writebacks)): wb = writebacks[i] writebacks[i] = None # Drop reference eagerly so recv buffers can free. @@ -202,7 +214,7 @@ 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' + # '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)) @@ -211,20 +223,24 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: 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(): dst_param.quantize_(full_bf16) pending_quantized.clear() - # Second sync: the writeback loop's .copy_() kernels are still async when - # execute_reshard_plan returns; without this CUDA-graph capture or callers - # that read params immediately race against the writes. + # Ensure all writeback copies are visible to subsequent CUDA ops (e.g. CUDA + # graph warmup). The synchronize() above fires *before* the writeback loop, + # so without this second sync the .copy_() kernels are still async when + # execute_reshard_plan returns — creating a race with callers that immediately + # inspect or capture (via CUDA graphs) the destination parameters. torch.cuda.synchronize() - # MXFP8 destinations allocate a full-model-sized BF16 staging buffer that - # can dwarf the rest of the working set. Reclaim it back to the driver - # when present; skip the (expensive) empty_cache walk otherwise. + # 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: torch.cuda.empty_cache() diff --git a/megatron/core/resharding/planner.py b/megatron/core/resharding/planner.py index b3008ed61da..e740ca6608f 100644 --- a/megatron/core/resharding/planner.py +++ b/megatron/core/resharding/planner.py @@ -127,101 +127,41 @@ def _emit_lcm_block_ops( ops.append((src_global_rank, tuple(src_slice), tuple(dst_slice))) -def _plan_multi_dim_lcm( +def _tp_block_layout( param_name: str, src_metadata: ParameterMetadata, dst_metadata: ParameterMetadata, - descriptors: list[ShardingDescriptor], - my_global_rank: int, -) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: - """TP-only planner using LCM tiling with arbitrary integer src/dst strides.""" - if not descriptors: - return [] - if len(descriptors) != 1: - raise NotImplementedError( - f"{param_name}: _plan_multi_dim_lcm supports TP-only (one descriptor)" - ) - if descriptors[0].name != "tp": - raise NotImplementedError(f"{param_name}: _plan_multi_dim_lcm expects TP 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) - dim = d.dim - src_world = len(d.src_dim_ranks) - dst_world = len(d.dst_dim_ranks) - 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})" - ) - - ops: list[tuple[int, tuple[slice, ...], tuple[slice, ...]]] = [] - _emit_lcm_block_ops( - param_name=param_name, - src_shape=src_shape, - dst_shape=dst_shape, - dim=dim, - src_world=src_world, - dst_world=dst_world, - src_stride=d.src_stride, - dst_stride=d.dst_stride, - full_block_len=dst_local * dst_world, - dst_local_rank=_get_rank_in_group(my_global_rank, d.dst_dim_ranks), - src_dim_ranks=d.src_dim_ranks, - src_block_offset=0, - dst_block_offset=0, - block_label=f"TP dim{dim}", - ops=ops, - ) - _sort_ops_by_dst_offset(ops, dim) - return ops + 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. -def _plan_block_interleaved( - param_name: str, - src_metadata: ParameterMetadata, - dst_metadata: ParameterMetadata, - descriptors: list[ShardingDescriptor], - my_global_rank: int, -) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: - """ - Block-interleaved TP planner for parameters with ``partition_sizes``. - - When a parameter packs multiple independently-sharded components of - *different* sizes (e.g. Mamba in_proj packs z, x, B, C, dt), a simple - contiguous concat produces the wrong layout. Each block is gathered - (or scattered) across TP ranks independently before moving to the next. - - ``partition_sizes`` lists the per-TP-rank block sizes along the partition - dim. Block *i* occupies ``[sum(sizes[:i]), sum(sizes[:i+1]))`` in the - local tensor on every TP rank. In the *full* (TP-gathered) tensor, block - *i* occupies ``[sum(full_sizes[:i]), sum(full_sizes[:i+1]))`` where - ``full_sizes[i] = sizes[i] * src_tp_world``. + - 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. """ - if not descriptors or descriptors[0].name != "tp": - return [] - d = descriptors[0] - if my_global_rank not in d.dst_dim_ranks: - return [] - + d = descriptor dim = d.dim - 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) - src_sizes = src_metadata.partition_sizes dst_sizes = dst_metadata.partition_sizes if src_sizes is None and dst_sizes is None: - raise RuntimeError(f"{param_name}: _plan_block_interleaved called without partition_sizes") + 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) @@ -229,12 +169,14 @@ def _plan_block_interleaved( 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( @@ -242,32 +184,64 @@ def _plan_block_interleaved( 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, ...]]] = [] - src_block_offset = 0 - dst_block_offset = 0 - for blk in range(num_blocks): + 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=dim, + dim=d.dim, src_world=src_world, dst_world=dst_world, - src_stride=1, - dst_stride=1, - full_block_len=full_sizes[blk], + 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_block_offset, - dst_block_offset=dst_block_offset, - block_label=f"block {blk}", + src_block_offset=src_off, + dst_block_offset=dst_off, + block_label=label, ops=ops, ) - src_block_offset += src_sizes[blk] - dst_block_offset += dst_sizes[blk] - - _sort_ops_by_dst_offset(ops, dim) + _sort_ops_by_dst_offset(ops, d.dim) return ops @@ -314,20 +288,11 @@ def _determine_source_ranks_for_dst_param( ) -> 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 + # 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: - # Use block-interleaved planner when partition_sizes is present - # (e.g. Mamba in_proj packs components of different sizes) - if src_metadata.partition_sizes is not None or dst_metadata.partition_sizes is not None: - return _plan_block_interleaved( - param_name=param_name, - src_metadata=src_metadata, - dst_metadata=dst_metadata, - descriptors=descriptors, - my_global_rank=my_global_rank, - ) - return _plan_multi_dim_lcm( + return _plan_tp( param_name=param_name, src_metadata=src_metadata, dst_metadata=dst_metadata, diff --git a/megatron/core/resharding/refit.py b/megatron/core/resharding/refit.py index 8574ef9b69e..36ba914d33b 100644 --- a/megatron/core/resharding/refit.py +++ b/megatron/core/resharding/refit.py @@ -253,14 +253,17 @@ def _setup_mxfp8_transform_on_plan(plan, target_model) -> None: core = unwrap_model(lm) decoder = core.decoder if hasattr(core, 'decoder') else core - # Eligible params must be computed while still visible as nn.Parameter (BF16). + # 1. Compute which parameters are eligible for MXFP8 conversion. + # Must be done while params are still visible as nn.Parameter (BF16). convertible: set[str] = set() for name, param in decoder.named_parameters(): if _should_quantize_param(param): convertible.add(f"decoder.{name}") + # 2. Quantize decoder weights → persistent MXFP8Tensor buffers. persistent_buffers = quantize_params_to_mxfp8(decoder) + # 3. Build the transform and attach it to the plan. plan.transform = MXFP8ReshardTransform( convertible_params=convertible, persistent_buffers=persistent_buffers, diff --git a/megatron/core/resharding/utils.py b/megatron/core/resharding/utils.py index 443555681d3..b2725a0fd43 100644 --- a/megatron/core/resharding/utils.py +++ b/megatron/core/resharding/utils.py @@ -447,6 +447,13 @@ def _filter_by_ep_local_rank( When EP sizes differ (resharding), expert matching is handled via ``resolved_name`` and no filter is applied here. + + Why size check matters: + - Same size (EP=8→EP=8): local ranks 0-7 exist in both src and dst → + 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. """ dst_ep_group = dst_metadata.expert_parallel_group_ranks if dst_ep_group is None: @@ -458,6 +465,7 @@ def _filter_by_ep_local_rank( if src_meta_list[0].expert_parallel_group_ranks else None ) + # EP resharding (sizes differ) — skip filter; keep all source candidates. if src_ep_size != len(dst_ep_group): return src_meta_list @@ -468,6 +476,7 @@ def _filter_by_ep_local_rank( and m.expert_parallel_group_ranks.index(m.owner_rank) == dst_ep_local ] if not matching: + # Sizes match but no local rank match — configuration bug. available = [ ( m.owner_rank, @@ -487,12 +496,22 @@ def _filter_by_ep_local_rank( def _round_robin_dp(src_meta_list: list[ParameterMetadata], dst_rank: int) -> ParameterMetadata: - """Round-robin across source DP groups so transfer load spreads evenly.""" + """Round-robin across source DP groups so transfer load spreads evenly. + + Each DP group holds a complete copy of the model, so we can read from any + DP group; choosing via ``dst_rank % num_src_dp_groups`` ensures even + distribution even when destination has different DP configuration. E.g. + with 4 src DP groups and 128 dst ranks, each src DP group is selected by + 32 dst ranks (128/4=32). Within the chosen DP group we further cycle + across available metadata entries to balance load across TP groups in the + DP replica. + """ grouped_by_dp: dict[tuple[int, ...], list[ParameterMetadata]] = {} for meta in src_meta_list: dp_group = tuple(meta.data_parallel_group_ranks or []) grouped_by_dp.setdefault(dp_group, []).append(meta) + # Fast path: only one DP group present; no balancing necessary. if len(grouped_by_dp) == 1: return src_meta_list[0] @@ -511,12 +530,16 @@ def select_src_metadata_balanced( 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. + A local copy is essentially free (``tensor.copy_()`` on same GPU), while + any remote transfer incurs significant overhead even within the same node. """ if not src_meta_list: raise ValueError("src_meta_list must be non-empty") candidates = _filter_by_ep_local_rank(src_meta_list, dst_metadata, dst_rank) + # Local copy optimization (collocated mode): if dst_rank owns a source + # replica after EP filtering, use it directly to skip the network entirely. for meta in candidates: if meta.owner_rank == dst_rank: return meta diff --git a/tests/unit_tests/resharding/test_planner.py b/tests/unit_tests/resharding/test_planner.py index 5ad087c925a..55c4019dc09 100644 --- a/tests/unit_tests/resharding/test_planner.py +++ b/tests/unit_tests/resharding/test_planner.py @@ -2,9 +2,9 @@ """Unit tests for the resharding planner functions. -These test the TP planners (_plan_multi_dim_lcm, _plan_block_interleaved), -DP fallback (_finalize_dp_transfers), and descriptor building in isolation -without requiring distributed init or GPU. +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. """ import math @@ -14,8 +14,7 @@ from megatron.core.resharding.planner import ( _build_descriptors_for_param, _finalize_dp_transfers, - _plan_block_interleaved, - _plan_multi_dim_lcm, + _plan_tp, ) from megatron.core.resharding.utils import ParameterMetadata, ShardingDescriptor @@ -84,7 +83,7 @@ def _verify_full_coverage(ops, dim, expected_full_len): # =========================================================================== -# _plan_multi_dim_lcm +# _plan_tp # =========================================================================== @@ -99,7 +98,7 @@ def test_tp2_to_tp1(self): dst = _meta(shape=(64, 128), is_tp=False, tp_ranks=[0]) desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0]) - ops = _plan_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=0) + 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} @@ -113,13 +112,13 @@ def test_tp1_to_tp2(self): desc = _tp_descriptor(dim=1, src_ranks=[0], dst_ranks=[0, 1]) # Rank 0 receives first half - ops_r0 = _plan_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=0) + 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) # Rank 1 receives second half - ops_r1 = _plan_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=1) + 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 @@ -130,7 +129,7 @@ def test_tp2_to_tp4(self): desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0, 1, 2, 3]) for rank in range(4): - ops = _plan_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=rank) + 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) @@ -140,7 +139,7 @@ def test_same_tp_size(self): 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_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=0) + ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) assert len(ops) == 1 assert ops[0][0] == 0 # from self @@ -150,14 +149,14 @@ def test_rank_not_in_dst(self): 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_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=0) + 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_multi_dim_lcm("weight", src, dst, [], my_global_rank=0) + ops = _plan_tp("weight", src, dst, [], my_global_rank=0) assert ops == [] def test_size_mismatch_raises(self): @@ -167,7 +166,7 @@ def test_size_mismatch_raises(self): desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0]) with pytest.raises(RuntimeError, match="size mismatch"): - _plan_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=0) + _plan_tp("weight", src, dst, [desc], my_global_rank=0) def test_dim0_partition(self): """TP on dimension 0 (row-parallel).""" @@ -175,7 +174,7 @@ def test_dim0_partition(self): 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_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=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) @@ -187,7 +186,7 @@ def test_conservation_all_ranks(self): all_ops = [] for rank in range(2): - ops = _plan_multi_dim_lcm("weight", src, dst, [desc], my_global_rank=rank) + 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 @@ -196,7 +195,7 @@ def test_conservation_all_ranks(self): # =========================================================================== -# _plan_block_interleaved +# _plan_tp # =========================================================================== @@ -213,7 +212,7 @@ def test_tp2_to_tp1_two_blocks(self): 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_block_interleaved("weight", src, dst, [desc], my_global_rank=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) @@ -225,22 +224,13 @@ def test_tp1_to_tp2_two_blocks(self): ) desc = _tp_descriptor(dim=1, src_ranks=[0], dst_ranks=[0, 1]) - ops_r0 = _plan_block_interleaved("weight", src, dst, [desc], my_global_rank=0) - ops_r1 = _plan_block_interleaved("weight", src, dst, [desc], my_global_rank=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_no_partition_sizes_raises(self): - """Both src and dst missing partition_sizes should raise.""" - src = _meta(shape=(64, 48), is_tp=True, partition_dim=1, tp_ranks=[0, 1]) - dst = _meta(shape=(64, 96), is_tp=False, tp_ranks=[0]) - desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0]) - - with pytest.raises(RuntimeError, match="partition_sizes"): - _plan_block_interleaved("weight", src, dst, [desc], my_global_rank=0) - def test_rank_not_in_dst(self): """Rank not in destination returns empty.""" src = _meta( @@ -249,7 +239,7 @@ def test_rank_not_in_dst(self): dst = _meta(shape=(64, 96), tp_ranks=[2]) desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[2]) - ops = _plan_block_interleaved("weight", src, dst, [desc], my_global_rank=0) + ops = _plan_tp("weight", src, dst, [desc], my_global_rank=0) assert ops == [] def test_three_blocks_tp2_to_tp4(self): @@ -268,7 +258,7 @@ def test_three_blocks_tp2_to_tp4(self): desc = _tp_descriptor(dim=1, src_ranks=[0, 1], dst_ranks=[0, 1, 2, 3]) for rank in range(4): - ops = _plan_block_interleaved("weight", src, dst, [desc], my_global_rank=rank) + 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) From 4aa90b99ed0fdfd5617a669181fcda99fe97bcee Mon Sep 17 00:00:00 2001 From: William Dykas Date: Wed, 13 May 2026 06:33:17 -0700 Subject: [PATCH 4/4] tests --- .../core/resharding/copy_services/base.py | 9 +- megatron/core/resharding/planner.py | 4 +- megatron/core/resharding/utils.py | 7 + .../resharding/test_copy_services.py | 166 ++++++++++ .../unit_tests/resharding/test_refit_cache.py | 311 ++++++++++++++++++ 5 files changed, 493 insertions(+), 4 deletions(-) create mode 100644 tests/unit_tests/resharding/test_copy_services.py create mode 100644 tests/unit_tests/resharding/test_refit_cache.py diff --git a/megatron/core/resharding/copy_services/base.py b/megatron/core/resharding/copy_services/base.py index de47705fd15..00dc884767c 100644 --- a/megatron/core/resharding/copy_services/base.py +++ b/megatron/core/resharding/copy_services/base.py @@ -79,7 +79,14 @@ def match_local_ops_by_task_id( f"{backend_name}: local (same-rank) transfer requires a task_id " "to match sends with recvs" ) - if len(sends_by_id) != len(local_sends) or len(recvs_by_id) != len(local_recvs): + # Count mismatch catches both imbalanced send/recv lists (which would + # otherwise silently drop the longer side) and duplicate task_ids (which + # collapse to fewer dict entries than list entries). + if ( + len(local_sends) != len(local_recvs) + or len(sends_by_id) != len(local_sends) + or len(recvs_by_id) != len(local_recvs) + ): raise RuntimeError( f"{backend_name}: unmatched local ops on rank {rank}: " f"{len(local_sends)} local sends vs {len(local_recvs)} local recvs" diff --git a/megatron/core/resharding/planner.py b/megatron/core/resharding/planner.py index e740ca6608f..f0e38004d0e 100644 --- a/megatron/core/resharding/planner.py +++ b/megatron/core/resharding/planner.py @@ -218,9 +218,7 @@ def _plan_tp( 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 - ) + 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: diff --git a/megatron/core/resharding/utils.py b/megatron/core/resharding/utils.py index b2725a0fd43..94c5767c217 100644 --- a/megatron/core/resharding/utils.py +++ b/megatron/core/resharding/utils.py @@ -515,8 +515,15 @@ def _round_robin_dp(src_meta_list: list[ParameterMetadata], dst_rank: int) -> Pa if len(grouped_by_dp) == 1: return src_meta_list[0] + # Round-robin selection across source DP groups based on destination global rank. + # This ensures even distribution: if we have 4 src DP groups and 128 dst ranks, + # each src DP group will be selected by 32 dst ranks (128 / 4 = 32). sorted_dp_groups = sorted(grouped_by_dp.keys()) chosen_group = sorted_dp_groups[dst_rank % len(sorted_dp_groups)] + # Within the chosen DP group, distribute across available metadata entries + # to balance load across all TP groups in the DP replica. + # Example: With 4 TP groups in a DP group, dst_ranks will cycle through all 4 + # instead of always using the first one, better distributing transfer load. group_metadata = grouped_by_dp[chosen_group] within_group_idx = (dst_rank // len(sorted_dp_groups)) % len(group_metadata) return group_metadata[within_group_idx] diff --git a/tests/unit_tests/resharding/test_copy_services.py b/tests/unit_tests/resharding/test_copy_services.py new file mode 100644 index 00000000000..0fe9a40bf60 --- /dev/null +++ b/tests/unit_tests/resharding/test_copy_services.py @@ -0,0 +1,166 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for shared copy_services utilities. + +Covers: +- ``match_local_ops_by_task_id``: pairing semantics and every error branch. +- ``CopyService.close()`` default no-op contract. +""" + +import pytest +import torch + +from megatron.core.resharding.copy_services.base import ( + CopyService, + RecvOp, + SendOp, + match_local_ops_by_task_id, +) + + +def _t(): + return torch.zeros(4) + + +class TestMatchLocalOps: + """match_local_ops_by_task_id pairs by task_id and rejects malformed inputs.""" + + def test_single_pair(self): + sends = [SendOp(task_id=1, tensor=_t(), dest_rank=0)] + recvs = [RecvOp(task_id=1, tensor=_t(), src_rank=0)] + pairs = match_local_ops_by_task_id(sends, recvs, "Test", rank=0) + assert len(pairs) == 1 + send_op, recv_op = pairs[0] + assert send_op is sends[0] + assert recv_op is recvs[0] + + def test_pairs_match_across_order(self): + """Order of sends vs recvs doesn't matter; pairing is by task_id.""" + sends = [ + SendOp(task_id=1, tensor=_t(), dest_rank=0), + SendOp(task_id=2, tensor=_t(), dest_rank=0), + ] + recvs = [ + RecvOp(task_id=2, tensor=_t(), src_rank=0), + RecvOp(task_id=1, tensor=_t(), src_rank=0), + ] + pairs = match_local_ops_by_task_id(sends, recvs, "Test", rank=0) + pair_ids = {(s.task_id, r.task_id) for s, r in pairs} + assert pair_ids == {(1, 1), (2, 2)} + + def test_empty_lists(self): + pairs = match_local_ops_by_task_id([], [], "Test", rank=0) + assert pairs == [] + + def test_none_send_task_id_raises(self): + sends = [SendOp(task_id=None, tensor=_t(), dest_rank=0)] + recvs = [RecvOp(task_id=1, tensor=_t(), src_rank=0)] + with pytest.raises(RuntimeError, match="requires a task_id"): + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=0) + + def test_none_recv_task_id_raises(self): + sends = [SendOp(task_id=1, tensor=_t(), dest_rank=0)] + recvs = [RecvOp(task_id=None, tensor=_t(), src_rank=0)] + with pytest.raises(RuntimeError, match="requires a task_id"): + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=0) + + def test_more_sends_than_recvs_raises(self): + sends = [ + SendOp(task_id=1, tensor=_t(), dest_rank=0), + SendOp(task_id=2, tensor=_t(), dest_rank=0), + ] + recvs = [RecvOp(task_id=1, tensor=_t(), src_rank=0)] + with pytest.raises(RuntimeError, match="unmatched local ops"): + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=7) + + def test_more_recvs_than_sends_raises(self): + sends = [SendOp(task_id=1, tensor=_t(), dest_rank=0)] + recvs = [ + RecvOp(task_id=1, tensor=_t(), src_rank=0), + RecvOp(task_id=2, tensor=_t(), src_rank=0), + ] + with pytest.raises(RuntimeError, match="unmatched local ops"): + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=7) + + def test_duplicate_send_task_ids_raises(self): + """Duplicate task_ids in sends collapse the dict — triggers count-mismatch raise.""" + sends = [ + SendOp(task_id=1, tensor=_t(), dest_rank=0), + SendOp(task_id=1, tensor=_t(), dest_rank=0), # duplicate + ] + recvs = [ + RecvOp(task_id=1, tensor=_t(), src_rank=0), + RecvOp(task_id=2, tensor=_t(), src_rank=0), + ] + with pytest.raises(RuntimeError, match="unmatched local ops"): + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=7) + + def test_duplicate_recv_task_ids_raises(self): + sends = [ + SendOp(task_id=1, tensor=_t(), dest_rank=0), + SendOp(task_id=2, tensor=_t(), dest_rank=0), + ] + recvs = [ + RecvOp(task_id=1, tensor=_t(), src_rank=0), + RecvOp(task_id=1, tensor=_t(), src_rank=0), # duplicate + ] + with pytest.raises(RuntimeError, match="unmatched local ops"): + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=7) + + def test_mismatched_task_ids_raises(self): + """Equal counts but task_ids don't overlap — should raise missing-send.""" + sends = [SendOp(task_id=1, tensor=_t(), dest_rank=0)] + recvs = [RecvOp(task_id=99, tensor=_t(), src_rank=0)] + with pytest.raises(RuntimeError, match="missing local send"): + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=0) + + def test_error_message_includes_backend_and_rank(self): + sends = [SendOp(task_id=None, tensor=_t(), dest_rank=0)] + recvs = [RecvOp(task_id=1, tensor=_t(), src_rank=0)] + with pytest.raises(RuntimeError) as exc_info: + match_local_ops_by_task_id(sends, recvs, "TestBackend", rank=0) + assert "TestBackend" in str(exc_info.value) + + +class TestCopyServiceClose: + """CopyService.close() default is a no-op; subclasses may override.""" + + def test_default_close_is_noop_and_returns_none(self): + class _Stub(CopyService): + def __init__(self): # bypass dist requirement + pass + + def submit_send(self, *args, **kwargs): + pass + + def submit_recv(self, *args, **kwargs): + pass + + def run(self): + pass + + svc = _Stub() + result = svc.close() + assert result is None + + def test_subclass_can_override_close(self): + class _ClosingService(CopyService): + def __init__(self): + self.closed = False + + def submit_send(self, *args, **kwargs): + pass + + def submit_recv(self, *args, **kwargs): + pass + + def run(self): + pass + + def close(self): + self.closed = True + + svc = _ClosingService() + assert svc.closed is False + svc.close() + assert svc.closed is True diff --git a/tests/unit_tests/resharding/test_refit_cache.py b/tests/unit_tests/resharding/test_refit_cache.py new file mode 100644 index 00000000000..b7c7fba2c8a --- /dev/null +++ b/tests/unit_tests/resharding/test_refit_cache.py @@ -0,0 +1,311 @@ +# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for the refit/reshard caches. + +Covers: +- ``_PlanCacheKey`` separation across configurations that route to different + global ranks (the rank-offset bug — two non-collocated configs with identical + parallel sizes used to silently share a plan). +- ``get_refit_tensor_dict`` / ``invalidate_refit_tensor_cache`` (module-level + named_refit_tensors cache + invalidation when ``_harmonize_buffer_dtypes`` + replaces a buffer). +""" + +import torch +import torch.nn as nn + +from megatron.core.resharding.refit import _PlanCacheKey +from megatron.core.resharding.utils import get_refit_tensor_dict, invalidate_refit_tensor_cache + + +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 + ) + assert k1 == k2 + assert hash(k1) == hash(k2) + + 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), + 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), + num_experts=None, + src_rank_offset=8, + dst_rank_offset=12, + ) + assert k1 != k2 + assert hash(k1) != hash(k2) + + 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), + 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), + num_experts=None, + src_rank_offset=0, + dst_rank_offset=8, + ) + assert k1 != k2 + + 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 + ) + k2 = _PlanCacheKey( + rank=3, + src_config=(2, 1, 1, 4, 1), + dst_config=(2, 1, 1, 4, 1), + num_experts=None, + src_rank_offset=0, + dst_rank_offset=0, + ) + assert k1 == k2 + + def test_num_experts_distinguishes(self): + k1 = _PlanCacheKey(rank=0, src_config=None, dst_config=None, num_experts=8) + k2 = _PlanCacheKey(rank=0, src_config=None, dst_config=None, num_experts=16) + assert k1 != k2 + + +class TestPlanCacheKeyNonCollocated: + """Non-collocated ranks set src_config or dst_config to None. + + Cache key must distinguish the three rank classes (source-only, dest-only, + idle) so they don't share plans across roles. + """ + + 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) + 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 + ) + assert idle != with_src + assert idle != with_dst + assert with_src != with_dst + + 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) + # 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, + num_experts=None, + src_rank_offset=0, + dst_rank_offset=4, + ) + layout_b = _PlanCacheKey( + rank=0, + src_config=sizes, + dst_config=sizes, + num_experts=None, + src_rank_offset=0, + dst_rank_offset=8, + ) + assert layout_a != layout_b + + +class TestNeedsMxfp8Conversion: + """_needs_mxfp8_conversion gracefully handles non-target ranks (model=None).""" + + def test_none_returns_false(self): + """Source-only and idle ranks pass target_model=None to _setup_mxfp8_...""" + from megatron.core.resharding.refit import _needs_mxfp8_conversion + + assert _needs_mxfp8_conversion(None) is False + + def test_mxfp8_model_returns_true(self): + from megatron.core.resharding.refit import _needs_mxfp8_conversion + + class _Cfg: + transformer_impl = "inference_optimized" + fp8_recipe = "mxfp8" + + class _Model: + config = _Cfg() + + assert _needs_mxfp8_conversion(_Model()) is True + + def test_non_inference_optimized_returns_false(self): + from megatron.core.resharding.refit import _needs_mxfp8_conversion + + class _Cfg: + transformer_impl = "transformer_engine" + fp8_recipe = "mxfp8" + + class _Model: + config = _Cfg() + + assert _needs_mxfp8_conversion(_Model()) is False + + def test_non_mxfp8_recipe_returns_false(self): + from megatron.core.resharding.refit import _needs_mxfp8_conversion + + class _Cfg: + transformer_impl = "inference_optimized" + fp8_recipe = "delayed" + + class _Model: + config = _Cfg() + + assert _needs_mxfp8_conversion(_Model()) is False + + def test_list_wrapped_model(self): + """The function unwraps a single-element list/tuple.""" + from megatron.core.resharding.refit import _needs_mxfp8_conversion + + class _Cfg: + transformer_impl = "inference_optimized" + fp8_recipe = "mxfp8" + + class _Model: + config = _Cfg() + + assert _needs_mxfp8_conversion([_Model()]) is True + + +class TestSetupMxfp8TransformOnPlan: + """_setup_mxfp8_transform_on_plan is a no-op on non-target ranks and idempotent.""" + + def test_target_none_leaves_transform_unset(self): + """Source-only / idle ranks should leave plan.transform at None.""" + from megatron.core.resharding.refit import _setup_mxfp8_transform_on_plan + from megatron.core.resharding.utils import ReshardPlan + + plan = ReshardPlan(send_ops=[], recv_ops=[]) + _setup_mxfp8_transform_on_plan(plan, None) + assert plan.transform is None + + def test_non_mxfp8_target_leaves_transform_unset(self): + from megatron.core.resharding.refit import _setup_mxfp8_transform_on_plan + from megatron.core.resharding.utils import ReshardPlan + + class _Cfg: + transformer_impl = "transformer_engine" + fp8_recipe = None + + class _Model: + config = _Cfg() + + plan = ReshardPlan(send_ops=[], recv_ops=[]) + _setup_mxfp8_transform_on_plan(plan, _Model()) + assert plan.transform is None + + def test_already_populated_skips_rebuild(self): + """Idempotent: if plan.transform is already set, do not re-quantize.""" + from megatron.core.resharding.refit import _setup_mxfp8_transform_on_plan + from megatron.core.resharding.transforms import ReshardTransform + from megatron.core.resharding.utils import ReshardPlan + + sentinel = ReshardTransform() + plan = ReshardPlan(send_ops=[], recv_ops=[], transform=sentinel) + + # Even with an MXFP8 model, the existing transform should not be replaced. + class _Cfg: + transformer_impl = "inference_optimized" + fp8_recipe = "mxfp8" + + class _Model: + config = _Cfg() + + _setup_mxfp8_transform_on_plan(plan, _Model()) + assert plan.transform is sentinel + + +class TestRefitTensorCache: + """get_refit_tensor_dict caches the param/buffer dict on the module.""" + + def test_returns_same_dict_on_repeat(self): + model = nn.Linear(4, 4, bias=False) + d1 = get_refit_tensor_dict(model) + d2 = get_refit_tensor_dict(model) + assert d1 is d2 + + def test_contains_parameters(self): + model = nn.Linear(4, 4) + d = get_refit_tensor_dict(model) + assert "weight" in d and "bias" in d + + def test_contains_persistent_buffers(self): + model = nn.Module() + model.register_buffer("running_mean", torch.zeros(4)) + d = get_refit_tensor_dict(model) + assert "running_mean" in d + + def test_excludes_non_persistent_buffers(self): + model = nn.Module() + model.register_buffer("tmp", torch.zeros(4), persistent=False) + d = get_refit_tensor_dict(model) + assert "tmp" not in d + + def test_invalidate_drops_cache(self): + model = nn.Linear(4, 4, bias=False) + d1 = get_refit_tensor_dict(model) + invalidate_refit_tensor_cache(model) + d2 = get_refit_tensor_dict(model) + assert d1 is not d2 + + def test_invalidate_picks_up_replaced_buffer(self): + """Mirrors _harmonize_buffer_dtypes: replace _buffers entry, invalidate, re-read.""" + model = nn.Module() + model.register_buffer("buf", torch.zeros(4, dtype=torch.bfloat16)) + d1 = get_refit_tensor_dict(model) + old_buf = d1["buf"] + + model._buffers["buf"] = old_buf.to(torch.float32) + invalidate_refit_tensor_cache(model) + + d2 = get_refit_tensor_dict(model) + assert d2["buf"].dtype == torch.float32 + assert d2["buf"] is not old_buf + + def test_invalidate_when_no_cache_is_safe(self): + """Calling invalidate before any get_refit_tensor_dict call should not raise.""" + model = nn.Linear(4, 4, bias=False) + invalidate_refit_tensor_cache(model) # no-op + + def test_cache_is_per_module(self): + m1 = nn.Linear(4, 4, bias=False) + m2 = nn.Linear(4, 4, bias=False) + d1 = get_refit_tensor_dict(m1) + d2 = get_refit_tensor_dict(m2) + assert d1 is not d2