diff --git a/megatron/core/resharding/copy_services/base.py b/megatron/core/resharding/copy_services/base.py index 3ea58e8d141..00dc884767c 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,45 @@ 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" + ) + # 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" + ) + 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..b22ea91c851 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,16 +19,18 @@ 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]] = [] + # 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( @@ -74,43 +57,23 @@ 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] + # 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_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(), 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] if dst_tensor.device != src_tensor.device: dst_tensor.copy_(src_tensor.to(dst_tensor.device)) else: @@ -128,13 +91,15 @@ 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. + # Wait only on default-stream staging copies; _copy_stream keeps running. if cpu_send_bufs: - torch.cuda.synchronize() + torch.cuda.current_stream().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) @@ -150,18 +115,17 @@ def run(self): 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() + # 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: + dst_tensor.copy_(recv.tensor, non_blocking=True) + else: + dst_tensor.copy_(recv.tensor) + + # 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: logger.info("GlooCopyService: batched communication completed") diff --git a/megatron/core/resharding/copy_services/nccl_copy_service.py b/megatron/core/resharding/copy_services/nccl_copy_service.py index 77df4381136..a1b63334bbf 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: diff --git a/megatron/core/resharding/copy_services/nvshmem_copy_service.py b/megatron/core/resharding/copy_services/nvshmem_copy_service.py index f0ade45518d..4e82d60de11 100644 --- a/megatron/core/resharding/copy_services/nvshmem_copy_service.py +++ b/megatron/core/resharding/copy_services/nvshmem_copy_service.py @@ -24,9 +24,8 @@ 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 @@ -39,6 +38,11 @@ def __init__(self, group=None): 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,7 +54,7 @@ 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() @@ -80,7 +84,7 @@ 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() @@ -141,7 +145,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 +154,13 @@ def run(self): # - 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") 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 e1b75bb0a70..8b06ef1cd14 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,67 @@ def execute_reshard_plan( transform's prepare_send / prepare_recv / finalize_recv methods instead of the default slice-and-copy logic. """ - # 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)} + # 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 {} # 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] = {} + # 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] = {} - 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] + 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: + 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. + 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(): + src_view = src_view.contiguous() + 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 + # 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,43 +147,51 @@ 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 + + # 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(): + 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(): + 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 (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( + _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() @@ -152,48 +202,32 @@ def _get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: # For quantized destination params (fp8_param=true on receiver), # accumulate ALL BF16 slices per-param before calling quantize_() once. # This avoids corrupting MXFP8 per-block scales from partial-slice updates. - # # Since refit overwrites every slice of each param, we allocate a fresh # BF16 buffer (torch.empty) instead of dequantizing the existing MXFP8 # weights — this avoids a full-model-sized dequantize+clone. - # - # 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 + 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' — 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)) + 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() + + # 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) - # Free the BF16 accumulation buffers eagerly. pending_quantized.clear() # Ensure all writeback copies are visible to subsequent CUDA ops (e.g. CUDA @@ -206,6 +240,8 @@ def _get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: # Release transient BF16 recv/accumulation buffers back to the CUDA driver. # Without this the caching allocator retains the peak allocation, which can # be significant for MXFP8 destinations (full model weight size in BF16). - torch.cuda.empty_cache() + # Skip the (expensive) empty_cache walk when no MXFP8 staging happened. + if had_mxfp8_staging: + torch.cuda.empty_cache() logger.info("Reshard complete") diff --git a/megatron/core/resharding/planner.py b/megatron/core/resharding/planner.py index 242450bd835..f0e38004d0e 100644 --- a/megatron/core/resharding/planner.py +++ b/megatron/core/resharding/planner.py @@ -73,142 +73,110 @@ def _build_descriptors_for_param( return descriptors -def _plan_multi_dim_lcm( +def _emit_lcm_block_ops( + *, param_name: str, - src_metadata: ParameterMetadata, - dst_metadata: ParameterMetadata, - descriptors: list[ShardingDescriptor], - my_global_rank: int, -) -> list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]: + 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. """ - TP-only planner using LCM tiling to support strides on source/destination. - - Requires exactly one TP descriptor - - Supports arbitrary integer strides (contiguous micro-tiles) - """ - 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})" - ) - # 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: + 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}: TP dim{dim} full_len {full_len} not divisible by LCM {L} " + f"{param_name}: {block_label} length {full_block_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)): + 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 - # 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_global_rank = src_dim_ranks[s_idx % src_world] 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_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))) - _sort_ops_by_dst_offset(ops, dim) - return ops - -def _plan_block_interleaved( +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, ...]]]: - """ - 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. This function treats each - block independently: it gathers (or scatters) each block across TP ranks - before moving to the next block. - - ``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``. + descriptor: ShardingDescriptor, + src_shape: tuple[int, ...], + dst_shape: tuple[int, ...], +) -> list[tuple[int, int, int, int, int, str]]: + """Compute the per-block layout for a TP transfer. + + Returns a list of ``(src_offset, dst_offset, full_block_len, src_stride, + dst_stride, label)`` tuples that the LCM micro-tiler iterates. + + - Plain TP (no ``partition_sizes``): single block covering the full + partition dim with the descriptor's strides. + - Block-interleaved TP (``partition_sizes`` present, e.g. Mamba ``in_proj``): + one block per packed component, each independently sharded with stride=1. """ - 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) - - # 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") + 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}")] - # 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] else: 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 + 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( @@ -216,53 +184,62 @@ 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 - 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 - - 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 +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. - # 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 + 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_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_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_block_offset += src_blk_sz - dst_block_offset += dst_blk_sz + blocks = _tp_block_layout(param_name, src_metadata, dst_metadata, d, src_shape, dst_shape) - _sort_ops_by_dst_offset(ops, dim) + ops: list[tuple[int, tuple[slice, ...], tuple[slice, ...]]] = [] + for src_off, dst_off, full_len, src_stride, dst_stride, label in blocks: + _emit_lcm_block_ops( + param_name=param_name, + src_shape=src_shape, + dst_shape=dst_shape, + dim=d.dim, + src_world=src_world, + dst_world=dst_world, + src_stride=src_stride, + dst_stride=dst_stride, + full_block_len=full_len, + dst_local_rank=dst_local_rank, + src_dim_ranks=d.src_dim_ranks, + src_block_offset=src_off, + dst_block_offset=dst_off, + block_label=label, + ops=ops, + ) + _sort_ops_by_dst_offset(ops, d.dim) return ops @@ -294,23 +271,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( @@ -321,20 +286,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, @@ -406,12 +362,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 +376,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..36ba914d33b 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,17 +239,14 @@ 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 @@ -350,12 +348,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 +360,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 +373,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 +383,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 +418,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 +447,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..94c5767c217 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,116 @@ 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. - - 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 +) -> list[ParameterMetadata]: + """In non-collocated mode with matching EP size, restrict candidates to the + source rank holding the same global experts as ``dst_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. """ - 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 + ) + # EP resharding (sizes differ) — skip filter; keep all source candidates. + 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: + # Sizes match but no local rank match — configuration bug. + 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. + + 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 + # 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 + # 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) + # 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. + 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 + + return _round_robin_dp(candidates, dst_rank) 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_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 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) 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