`core.resharding.planner`#

Module Contents#

Functions#

`_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.
`_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

API#

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

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._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, ) → 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.

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, ) → 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.

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, ) → core.resharding.utils.ReshardPlan#: Deprecated compatibility wrapper for :func:build_local_reshard_plan.

core.resharding.planner#

Module Contents#

Functions#

Data#

API#

`core.resharding.planner`#