`core.resharding.execution`#

Module Contents#

Classes#

_Writeback

Tagged-union for what to do with a received tensor after service.run().

Functions#

`refresh_module_caches`	Refresh parameter-derived caches in the destination module(s).
`_requires_bf16_staging`	Return whether refit must materialize this parameter in BF16.
`_get_quantized_accumulator`	Get or lazily allocate the BF16 accumulation buffer for a quantized destination.
`_native_gtp_load_context`	Return the special update context required by native quantized GTP params.
`_validate_execution_batches`	Validate locally checkable invariants of a batched reshard plan.
`_get_execution_batches`	Return the plan’s validated batch grouping, building it only once.
`_execute_batch`	Submit, execute, and finalize one memory-bounded operation batch.
`execute_reshard_plan`	Execute a reshard plan (built locally on each rank). A communication service must be provided to abstract transport. Expected service API: submit_send(tensor, dest_rank, task_id), submit_recv(tensor, src_rank, task_id), run().

Data#

logger

API#

core.resharding.execution.logger#: ‘getLogger(…)’

core.resharding.execution.refresh_module_caches( dst_module: torch.nn.Module | list[torch.nn.Module] | tuple[torch.nn.Module, ...] | None, ) → None#

Refresh parameter-derived caches in the destination module(s).

Lists and tuples let external refit callers, including Megatron Bridge, pass virtual-pipeline model chunks directly. None is a no-op for send-only ranks. The dst_module parameter name remains stable for callers that use keyword arguments.

class core.resharding.execution._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 quantized accumulation when the dest is quantized). transform hands the received buffers to a ReshardTransform.finalize_recv call.

kind: str#: None

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

core.resharding.execution._requires_bf16_staging(param: torch.Tensor) → bool#

Return whether refit must materialize this parameter in BF16.

Quantized source storage is dequantized before slicing for transfer. On the destination, updating quantized storage slice-by-slice is unsafe because scale blocks can cross slice boundaries, so refit assembles the complete local BF16 weight and quantizes it once after all receives finish.

core.resharding.execution._get_quantized_accumulator( pending: dict[int, tuple], dst_param: torch.Tensor, ) → torch.Tensor#

Get or lazily allocate the BF16 accumulation buffer for a quantized destination.

All slices for the same dst_param land in this buffer; quantize_ is called once after all slices have been written.

core.resharding.execution._native_gtp_load_context( module: torch.nn.Module | None, pending: dict[int, tuple], )#: Return the special update context required by native quantized GTP params.

core.resharding.execution._validate_execution_batches( plan: core.resharding.utils.ReshardPlan, ) → None#: Validate locally checkable invariants of a batched reshard plan.

core.resharding.execution._get_execution_batches( plan: core.resharding.utils.ReshardPlan, ) → tuple[tuple[int, list[core.resharding.utils.TransferOp], list[core.resharding.utils.TransferOp]], ...]#: Return the plan’s validated batch grouping, building it only once.

core.resharding.execution._execute_batch( send_ops: list[core.resharding.utils.TransferOp], recv_ops: list[core.resharding.utils.TransferOp], src_params: dict[str, torch.Tensor], dst_params: dict[str, torch.Tensor], service: core.resharding.copy_services.base.CopyService, dst_module: torch.nn.Module | None, transform: Optional[core.resharding.transforms.ReshardTransform], prefetch_stream: Optional[torch.cuda.Stream], ) → bool#: Submit, execute, and finalize one memory-bounded operation batch.

core.resharding.execution.execute_reshard_plan( plan: core.resharding.utils.ReshardPlan, src_module: torch.nn.Module, dst_module: torch.nn.Module, service: core.resharding.copy_services.base.CopyService, group=None, transform: Optional[core.resharding.transforms.ReshardTransform] = None, ) → None#

Execute a reshard plan (built locally on each rank). A communication service must be provided to abstract transport. Expected service API: submit_send(tensor, dest_rank, task_id), submit_recv(tensor, src_rank, task_id), run().

Supports None for src_module and/or dst_module to allow ranks in non-collocated mode:

src_module=None: Rank only receives data (destination-only)
dst_module=None: Rank only sends data (source-only)
Both provided: Rank participates in both send and recv (collocated mode)

When transform is provided, parameters for which transform.should_transform(param_name) returns True use the transform’s prepare_send / prepare_recv / finalize_recv methods instead of the default slice-and-copy logic.

core.resharding.execution#