`dockyard_rl.distributed.named_sharding`¶

Example¶

layout = [[[0, 1, 2, 3], [4, 5, 6, 7]]]
names  = ["dp", "pp", "tp"]
# DP=1, PP=2, TP=4
sharding = NamedSharding(layout, names)
sharding.shape                      # {"dp": 1, "pp": 2, "tp": 4}
sharding.get_ranks(dp=0, pp=1)      # NamedSharding over ranks [4,5,6,7]
sharding.get_ranks_by_coord(pp=1)   # [4, 5, 6, 7]
sharding.get_worker_coords(5)       # {"dp": 0, "pp": 1, "tp": 1}

N-dimensional rank layout with named axes for sharding, replication, and data routing across TP / PP / DP / EP parallelism dimensions.

Module Contents¶

Classes¶

NamedSharding

N-dimensional arrangement of ranks with named axes.

Data¶

REPLICATED_AXES

API¶

dockyard_rl.distributed.named_sharding.REPLICATED_AXES: tuple[str, ...]¶: (‘tensor_parallel’, ‘context_parallel’, ‘pipeline_parallel’)

class dockyard_rl.distributed.named_sharding.NamedSharding(layout: Sequence[Any] | numpy.ndarray, names: list[str])¶

N-dimensional arrangement of ranks with named axes.

Facilitates data sharding, replication, and collection based on named parallelism axes (e.g. “dp”, “pp”, “tp”, “ep”).

Initialization

Initialise NamedSharding.

Args: layout: Nested sequence (list-of-lists) or ndarray representing the N-D rank layout. All leaf values must be integer rank IDs and must be unique across the layout. names: Axis names ordered from outermost to innermost dimension.

Raises: ValueError: Non-integer values, duplicate ranks, or shape/name length mismatch.

property shape: dict[str, int]¶: Shape of the rank layout as a name → size mapping.

property names: list[str]¶: Axis names (copy).

property ndim: int¶: Number of dimensions.

property size: int¶: Total number of ranks.

property layout: numpy.ndarray¶: Underlying NumPy rank array (copy).

get_worker_coords(worker_id: int) → dict[str, int]¶

Return the axis coordinates of a given rank ID.

Args: worker_id: Integer rank ID to look up.

Returns: Dict mapping each axis name to its coordinate index for this rank.

Raises: ValueError: If worker_id is not present in the layout.

static is_axis_zero(coords: dict[str, int], axes: Sequence[str]) → bool¶

Return True when coords is 0 on every axis in axes.

Replica-leader check shared with TQWorkerMixin._local_coords on the worker side; driver-side callers can pair it with get_worker_coords. Axes absent from coords are treated as 0.

get_ranks_by_coord(**coords: int) → list[int]¶

Return all ranks matching the specified axis coordinates.

Unspecified axes match all coordinates along that axis.

Args: **coords: Axis-name → coordinate pairs to filter by.

Returns: Sorted list of matching rank IDs (empty list if none match).

Raises: ValueError: If an unknown axis name is supplied.

get_ranks(**kwargs: int) → Union[dockyard_rl.distributed.named_sharding.NamedSharding, int]¶

Slice the layout along named axes, returning a sub-sharding.

If all axes are specified, returns the single integer rank ID.

Args: **kwargs: Axis-name → index pairs for the axes to slice.

Returns: A new NamedSharding over the remaining axes, or an int when all axes are fully specified.

Raises: ValueError / IndexError: Unknown axis name or out-of-bounds index.

get_axis_index(name: str) → int¶: Return the numerical index of a named axis.

get_axis_size(name: str) → int¶: Return the size of a named axis.

dockyard_rl.distributed.named_sharding¶