Repository navigation
Expand file tree
/
Copy pathshard_and_stack.py
More file actions
149 lines (123 loc) · 5.33 KB
/
Copy pathshard_and_stack.py
File metadata and controls
149 lines (123 loc) · 5.33 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
# ===----------------------------------------------------------------------=== #
# Copyright (c) 2026, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #
"""Op implementation for shard_and_stack."""
from __future__ import annotations
from collections.abc import Sequence
from ..dim import StaticDim
from ..type import Device, DeviceRef, TensorType
from ..value import TensorValue, TensorValueLike
from .custom import custom
from .validation import (
assert_on_host,
assert_same_device,
assert_same_dtype,
assert_same_shape,
assert_valid_axis,
)
def shard_and_stack(
inputs: Sequence[TensorValueLike],
devices: Sequence[Device | DeviceRef],
axis: int = 0,
) -> list[TensorValue]:
"""Shards a list of input tensors along a specified axis, producing multiple outputs.
This operation takes multiple input tensors, splits each along the specified axis
into ``len(devices)`` chunks, and returns one output tensor per device. Each output
contains the chunks at the corresponding index stacked from all inputs along
a new dimension 0.
This is useful for distributing model weights across multiple devices in
tensor parallel configurations.
For example, with 2 inputs A and B, axis=0, and 2 devices:
- Input A shape [10, D], Input B shape [10, D]
- Output 0: stack([A[0:5], B[0:5]]) -> shape [2, 5, D] on devices[0]
- Output 1: stack([A[5:10], B[5:10]]) -> shape [2, 5, D] on devices[1]
With axis=1 and 2 devices:
- Input A shape [D, 10], Input B shape [D, 10]
- Output 0: stack([A[:, 0:5], B[:, 0:5]]) -> shape [2, D, 5] on devices[0]
- Output 1: stack([A[:, 5:10], B[:, 5:10]]) -> shape [2, D, 5] on devices[1]
Args:
inputs: A list of symbolic tensors to shard. All tensors must have
the same shape, dtype, and device.
devices: Target devices for each output tensor. The number of devices
determines the number of splits. Each output tensor
will be placed on the corresponding device. This enables direct
host-to-device transfer without intermediate CPU storage.
axis: The axis along which to split each input tensor. Defaults to 0.
Supports negative indexing (for example, -1 for last axis).
Returns:
A list of ``len(devices)`` tensors, each with shape
[num_inputs, D0, ..., Daxis//len(devices), ..., Dn-1] where the input
shape is [D0, ..., Daxis, ..., Dn-1]. Output i contains the stacked
chunks at position i from all input tensors, placed on devices[i].
Raises:
ValueError: If inputs list is empty, if devices list is empty, if input
tensors don't have matching shapes, if the dimension size at the
axis is not evenly divisible by len(devices), or if axis is out of
bounds.
"""
if not inputs:
raise ValueError("Must provide at least one input tensor.")
if not devices:
raise ValueError("Must provide at least one device.")
num_splits = len(devices)
# Convert inputs to TensorValue
input_tensors = [TensorValue(x) for x in inputs]
# Validate all inputs have the same shape
first_shape = input_tensors[0].shape
first_dtype = input_tensors[0].dtype
input_rank = input_tensors[0].rank
assert_same_device(*input_tensors)
assert_same_dtype(*input_tensors)
assert_same_shape(*input_tensors)
# Validate axis in bounds (before normalization)
assert_valid_axis(input_tensors[0], axis)
# Validate inputs are on the host device
assert_on_host(*input_tensors)
# Normalize negative axis
if axis < 0:
axis = input_rank + axis
# Validate divisibility (best-effort: only check if dimension is static)
axis_dim = first_shape[axis]
if isinstance(axis_dim, StaticDim):
if int(axis_dim) % num_splits != 0:
raise ValueError(
f"Dimension size along axis {axis} ({axis_dim}) must be evenly "
f"divisible by number of devices ({num_splits})"
)
# Calculate output shape
# Each output has new dimension 0 for stacking, then original shape with
# axis dimension divided
chunk_size = axis_dim // num_splits
# Build output shape: [num_inputs, D0, ..., Daxis//num_splits, ..., Dn-1]
output_shape = [len(input_tensors), *first_shape]
# Position axis+1 because we added new dimension 0
output_shape[axis + 1] = chunk_size
# Create one output type per device
output_types = [
TensorType(
dtype=first_dtype,
shape=output_shape,
device=device,
)
for device in devices
]
# Call the custom kernel operation
results = custom(
"mo.shard_and_stack",
input_tensors[0].device,
input_tensors,
out_types=output_types,
parameters={
"axis": axis,
},
)
return [r.tensor for r in results]