core.resharding.shard_planner

Transfer planning in logical weight coordinates.

Module Contents

Classes

_Segment	A contiguous local interval and its position in the logical global weight.
_GTPShardLayout	Location of one stored dimension in its unpadded TP-local layout.

Functions

_gtp_shard_layout	Describe this rank’s GTP shard within the TP-local dimension.
_tp_segments	Map a TP-local dimension to the logical global weight.
_local_segments	Map local storage to the unpadded logical global weight.
_global_shape	Return the unpadded shape after materializing TP and GTP.
_source_shards	Find the TP x GTP shard grid containing the selected source replica.
_intersect_segments	Intersect two segment lists in logical global coordinates.
_rectangles_overlap	Return whether two slice rectangles cover any common element.
plan_sharded_transfer	Plan a transfer by intersecting source and destination logical shards.

API

class core.resharding.shard_planner._Segment

Bases: typing.NamedTuple

A contiguous local interval and its position in the logical global weight.

local_start: int

None

global_start: int

None

length: int

None

class core.resharding.shard_planner._GTPShardLayout

Bases: typing.NamedTuple

Location of one stored dimension in its unpadded TP-local layout.

tp_local_size is the size after GTP shards are joined and padding is removed, but before TP shards are joined. For TP=2, GTP=2, stored dim-0 size 3, and two padding rows (P):

GTP: TP 0 [a b c] + [d P P] -> [a b c d]
     TP 1 [e f g] + [h P P] -> [e f g h]
TP:       [a b c d] + [e f g h] -> [a b c d e f g h]

Each stored shard has size 3. Joining GTP gives tp_local_size = 3 * 2 - 2 = 4; joining TP then gives 8 global rows.

tp_local_start: int

None

tp_local_stop: int

None

tp_local_size: int

None

core.resharding.shard_planner._gtp_shard_layout(

metadata: core.resharding.utils.ParameterMetadata,

dim: int,

) → core.resharding.shard_planner._GTPShardLayout

Describe this rank’s GTP shard within the TP-local dimension.

core.resharding.shard_planner._tp_segments(

metadata: core.resharding.utils.ParameterMetadata,

dim: int,

tp_local_size: int,

) → list[core.resharding.shard_planner._Segment]

Map a TP-local dimension to the logical global weight.

core.resharding.shard_planner._local_segments(

metadata: core.resharding.utils.ParameterMetadata,

dim: int,

) → list[core.resharding.shard_planner._Segment]

Map local storage to the unpadded logical global weight.

GTP always takes a contiguous dim-0 slice of the TP-local layout. Intersecting that slice with the TP segments naturally handles column, row, strided, and packed tensor-parallel layouts.

core.resharding.shard_planner._global_shape(

metadata: core.resharding.utils.ParameterMetadata,

) → tuple[int, ...]

Return the unpadded shape after materializing TP and GTP.

core.resharding.shard_planner._source_shards(

all_src_metadata: list[core.resharding.utils.ParameterMetadata],

selected: core.resharding.utils.ParameterMetadata,

) → list[core.resharding.utils.ParameterMetadata]

Find the TP x GTP shard grid containing the selected source replica.

Walking both groups transitively finds the grid without assuming a particular global rank layout.

core.resharding.shard_planner._intersect_segments(

src_segments: list[core.resharding.shard_planner._Segment],

dst_segments: list[core.resharding.shard_planner._Segment],

) → list[tuple[slice, slice]]

Intersect two segment lists in logical global coordinates.

core.resharding.shard_planner._rectangles_overlap(

left: tuple[slice, ...],

right: tuple[slice, ...],

) → bool

Return whether two slice rectangles cover any common element.

core.resharding.shard_planner.plan_sharded_transfer(

param_name: str,

all_src_metadata: list[core.resharding.utils.ParameterMetadata],

selected_src: core.resharding.utils.ParameterMetadata,

dst_metadata: core.resharding.utils.ParameterMetadata,

) → list[tuple[int, tuple[slice, ...], tuple[slice, ...]]]

Plan a transfer by intersecting source and destination logical shards.