`core.resharding.planner`#

Module Contents#

Classes#

_NativeParameterPart

Functions#

`_find_source_metadata`	Find source metadata, including the tied-output embedding alias.
`_sort_ops_by_dst_offset`	Sort transfer ops by destination offset on the sharded dimension.
`_build_descriptors_for_param`	Construct sharding descriptors (currently TP) for this parameter based on actual layout. Guard TP descriptor with size conservation so we don’t mis-classify replicated tensors.
`_emit_lcm_block_ops`	Emit (src_rank, src_slice, dst_slice) ops for one LCM-tiled block.
`_tp_block_layout`	Compute the per-block layout for a TP transfer.
`_plan_tp`	Plan TP transfers via LCM tiling, supporting both plain and block-interleaved TP.
`_finalize_dp_transfers`	Return receiver-side transfer for a parameter that is not TP-sharded.
`_determine_source_ranks_for_dst_param`	Route to dimension-specific planner based on parameter sharding type.
`_iter_global_transfer_ops`	Yield the whole reshard schedule in a deterministic order.
`_build_execution_batch_ids`	Assign complete logical parameters to deterministic memory-bounded batches.
`_tensor_mesh`	Return the TP mesh that owns a local parameter shard or replica.
`_global_shape_for_mesh`
`_validate_native_mesh_metadata`	Return why a parameter mesh cannot use whole-tensor native resharding.
`_native_parameter_parts`	Split a packed TP parameter into regular tensors for native resharding.
`_build_tensor_reshard_specs`	Build native mesh specs without affecting the generic plan.
`_extract_module_metadata`	Metadata for a module’s params and persistent buffers, or [] if None.
`index_metadata_rosters`	Turn a rank-ordered list of `(src_meta, dst_meta)` (index == rank) into the two rosters the plan builder consumes: dst params keyed by rank, and src params keyed by resolved_name. The list may come from the all-gather, or be reassembled in rank order as nodes are added, before calling build_plan_from_rosters.
`build_plan_from_rosters`	Replay the deterministic global schedule and keep only this rank’s ops.
`build_local_reshard_plan`	Build this rank’s reshard plan locally: all-gather the parameter metadata, replay the global schedule (see _iter_global_transfer_ops), and keep only the ops where this rank is the sender or receiver. No rank-0 bottleneck and no scatter, since sender and receiver derive matching task_ids from the same metadata.
`build_centralized_reshard_plan`	Deprecated compatibility wrapper for :func:`build_local_reshard_plan`.

Data#

`logger`
`DEFAULT_MAX_PARAMS_PER_BATCH`

API#

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

core.resharding.planner.DEFAULT_MAX_PARAMS_PER_BATCH#: 32

class core.resharding.planner._NativeParameterPart#

global_shape: tuple[int, ...]#: None

src_local_shape: tuple[int, ...]#: None

dst_local_shape: tuple[int, ...]#: None

src_slice: tuple[slice, ...] | None#: None

dst_slice: tuple[slice, ...] | None#: None

core.resharding.planner._find_source_metadata( src_param_metadata: dict[str, list[core.resharding.utils.ParameterMetadata]], resolved_name: str, ) → list[core.resharding.utils.ParameterMetadata] | None#: Find source metadata, including the tied-output embedding alias.

core.resharding.planner._sort_ops_by_dst_offset(ops, dim)#: Sort transfer ops by destination offset on the sharded dimension.

core.resharding.planner._build_descriptors_for_param( src_metadata: core.resharding.utils.ParameterMetadata, dst_metadata: core.resharding.utils.ParameterMetadata, ) → list[core.resharding.utils.ShardingDescriptor]#: Construct sharding descriptors (currently TP) for this parameter based on actual layout. Guard TP descriptor with size conservation so we don’t mis-classify replicated tensors.

core.resharding.planner._emit_lcm_block_ops( *, param_name: str, src_shape: tuple[int, ...], dst_shape: tuple[int, ...], dim: int, src_world: int, dst_world: int, src_stride: int, dst_stride: int, full_block_len: int, dst_local_rank: int, src_dim_ranks: list[int], src_block_offset: int, dst_block_offset: int, block_label: str, ops: list, ) → None#

Emit (src_rank, src_slice, dst_slice) ops for one LCM-tiled block.

Used both by the single-block stride-aware TP planner and by the per-block loop of the block-interleaved planner.

core.resharding.planner._tp_block_layout( param_name: str, src_metadata: core.resharding.utils.ParameterMetadata, dst_metadata: core.resharding.utils.ParameterMetadata, descriptor: core.resharding.utils.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.

core.resharding.planner._plan_tp( param_name: str, src_metadata: core.resharding.utils.ParameterMetadata, dst_metadata: core.resharding.utils.ParameterMetadata, descriptors: list[core.resharding.utils.ShardingDescriptor], my_global_rank: int, ) → list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]#

Plan TP transfers via LCM tiling, supporting both plain and block-interleaved TP.

The block layout is derived once by _tp_block_layout — the inner LCM micro-tile math (_emit_lcm_block_ops) is identical for both cases, so the single-block plain-TP path is just a special case of the multi-block partitioned path.

core.resharding.planner._finalize_dp_transfers( param_name: str, src_metadata: core.resharding.utils.ParameterMetadata, dst_metadata: core.resharding.utils.ParameterMetadata, my_global_rank: int, ) → list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]#

Return receiver-side transfer for a parameter that is not TP-sharded.

This is reached when we cannot build a TP sharding descriptor for the parameter (i.e., it is effectively replicated with respect to sharding). We use this when the destination and source mode have no TP or the parameter is replicted on all ranks such as layernorm. If the source and destination DP groups match, we return a local full-tensor copy; otherwise we pick a source rank from the source DP group in a deterministic round-robin manner based on the receiver’s global rank for better load distribution.

core.resharding.planner._determine_source_ranks_for_dst_param( param_name: str, src_metadata: core.resharding.utils.ParameterMetadata, dst_metadata: core.resharding.utils.ParameterMetadata, my_global_rank: int, ) → list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]#: Route to dimension-specific planner based on parameter sharding type.

core.resharding.planner._iter_global_transfer_ops( dst_param_metadata_by_rank: dict[int, dict[str, core.resharding.utils.ParameterMetadata]], src_param_metadata: dict[str, list[core.resharding.utils.ParameterMetadata]], )#

Yield the whole reshard schedule in a deterministic order.

The iteration order (dst rank ascending, then that rank’s dst params in gathered order, then per-source sub-ops) depends only on the rosters, so replaying this on any rank produces the same sequence and assigns the same task_id to the same transfer. That’s what lets each rank build its own send/recv ops while sender and receiver still agree on task_id.

Ranks are taken from the roster keys rather than range(world_size), so a sparse or growing rank set (nodes added later) rebuilds identically.

Yields (task_id, dst_rank, src_rank, src_slice, dst_slice, src_metadata, dst_metadata). PP is handled implicitly: each rank contributes metadata only for the params it owns, and any source holding the same resolved_name can serve as sender (with DP balancing).

core.resharding.planner._build_execution_batch_ids( dst_param_metadata_by_rank: dict[int, dict[str, core.resharding.utils.ParameterMetadata]], src_param_metadata: dict[str, list[core.resharding.utils.ParameterMetadata]], max_batch_bytes: int | None = None, ) → tuple[dict[str, int], int]#

Assign complete logical parameters to deterministic memory-bounded batches.

The planner is replayed from the same global metadata roster on every rank, so these IDs let the generic executor call CopyService.run() in lockstep without another collective. Source and destination bytes are accumulated per rank; starting a new batch when any rank would cross the soft limit bounds both sender-side dequantization and receiver-side staging. All replicas and shards of one resolved parameter stay in one batch. None disables the byte limit; batches are then bounded only by the per-batch parameter cap (DEFAULT_MAX_PARAMS_PER_BATCH / MEGATRON_REFIT_MAX_PARAMS_PER_BATCH), which keeps every rank’s NCCL P2P group small enough to stay in a single kernel plan.

core.resharding.planner._tensor_mesh( metadata: core.resharding.utils.ParameterMetadata, ) → tuple[int, ...]#: Return the TP mesh that owns a local parameter shard or replica.

core.resharding.planner._global_shape_for_mesh( metadata: core.resharding.utils.ParameterMetadata, mesh: tuple[int, ...], ) → tuple[int, ...]#

core.resharding.planner._validate_native_mesh_metadata( resolved_name: str, side: str, mesh: tuple[int, ...], metadata_by_rank: dict[int, core.resharding.utils.ParameterMetadata], ) → str | None#: Return why a parameter mesh cannot use whole-tensor native resharding.

core.resharding.planner._native_parameter_parts( resolved_name: str, src_metadata: core.resharding.utils.ParameterMetadata, dst_metadata: core.resharding.utils.ParameterMetadata, src_mesh: tuple[int, ...], dst_mesh: tuple[int, ...], ) → tuple[list[core.resharding.planner._NativeParameterPart] | None, str | None]#

Split a packed TP parameter into regular tensors for native resharding.

Mamba-style projections concatenate several components that are each independently TP-sharded. The concatenation itself is not a regular sharded tensor, but each component is. Plain parameters remain a single full-tensor part.

core.resharding.planner._build_tensor_reshard_specs( dst_param_metadata_by_rank: dict[int, dict[str, core.resharding.utils.ParameterMetadata]], src_param_metadata: dict[str, list[core.resharding.utils.ParameterMetadata]], my_global_rank: int, ) → tuple[list[core.resharding.utils.TensorReshardSpec] | None, str | None]#: Build native mesh specs without affecting the generic plan.

core.resharding.planner._extract_module_metadata( module, owner_rank, num_experts, rank_offset, rank_list_cache, ) → list[core.resharding.utils.ParameterMetadata]#

Metadata for a module’s params and persistent buffers, or [] if None.

Persistent buffers travel too so training state (e.g. MoE router expert_bias) refits with the weights.

core.resharding.planner.index_metadata_rosters(gathered_pairs: list)#: Turn a rank-ordered list of (src_meta, dst_meta) (index == rank) into the two rosters the plan builder consumes: dst params keyed by rank, and src params keyed by resolved_name. The list may come from the all-gather, or be reassembled in rank order as nodes are added, before calling build_plan_from_rosters.

core.resharding.planner.build_plan_from_rosters( dst_param_metadata_by_rank: dict[int, dict[str, core.resharding.utils.ParameterMetadata]], src_param_metadata: dict[str, list[core.resharding.utils.ParameterMetadata]], my_global_rank: int, execution_batch_bytes: int | None = None, ) → core.resharding.utils.ReshardPlan#

Replay the deterministic global schedule and keep only this rank’s ops.

Pure and collective-free, so it can be tested or reused with preassembled rosters without touching the process group. Live membership orchestration is intentionally outside this module. execution_batch_bytes is an optional soft per-rank limit; a single complete parameter may exceed it. None keeps the model-wide submission behavior.

core.resharding.planner.build_local_reshard_plan( src_module: torch.nn.Module, dst_module: torch.nn.Module, num_experts: int | None = None, group=None, src_rank_offset: int = 0, dst_rank_offset: int = 0, execution_batch_bytes: int | None = None, ) → core.resharding.utils.ReshardPlan#

Build this rank’s reshard plan locally: all-gather the parameter metadata, replay the global schedule (see _iter_global_transfer_ops), and keep only the ops where this rank is the sender or receiver. No rank-0 bottleneck and no scatter, since sender and receiver derive matching task_ids from the same metadata.

The metadata gather (the one collective) and the plan build are split into index_metadata_rosters + build_plan_from_rosters, so the deterministic build can also be tested against preassembled rosters without a process group.

src_module/dst_module may be None for non-collocated ranks (destination-only, source-only, or idle). Each rank contributes metadata only for the models it owns, including its parallel-group membership. execution_batch_bytes is an optional soft per-rank limit for transient generic-executor staging; None keeps the model-wide submission behavior.

core.resharding.planner.build_centralized_reshard_plan( src_module: torch.nn.Module, dst_module: torch.nn.Module, num_experts: int | None = None, group=None, src_rank_offset: int = 0, dst_rank_offset: int = 0, execution_batch_bytes: int | None = None, ) → core.resharding.utils.ReshardPlan#: Deprecated compatibility wrapper for :func:build_local_reshard_plan.

core.resharding.planner#