From ec63ae5ab7ce19fdedaada581a33872fa0fd7d63 Mon Sep 17 00:00:00 2001 From: Taylor Wallgren <5424153+twallgren@users.noreply.github.com> Date: Thu, 17 Sep 2026 13:57:44 -0600 Subject: [PATCH] Support uneven tensor-parallel sharding in distributed Linear layers Sharding a dimension across the group required it to divide evenly, so a model whose attention heads or MLP size do not divide by the number of devices could not be sharded at all. Split a dimension as evenly as possible instead, and let a caller pass explicit per-rank sizes when it needs the split to follow a structure the layer does not know about, such as whole grouped-query KV head groups. A quantized module's packed weight and its grouped scales and biases are split from one list of sizes in unpacked elements so that they stay at matching boundaries, and those sizes are multiples of the group size so a paired all-to-sharded and sharded-to-all layer split their common dimension the same way. --- python/mlx/nn/layers/distributed.py | 284 ++++++++++++++++++++++---- python/tests/mlx_distributed_tests.py | 258 +++++++++++++++++++++++ 2 files changed, 502 insertions(+), 40 deletions(-) diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index 16979b5097..a5b89cb3c0 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/python/mlx/nn/layers/distributed.py @@ -37,20 +37,137 @@ def _split(weight, segments, axis): return mx.split(weight, indices, axis=axis) +def _rank_sizes(dim, N, block=1): + """Split ``dim`` as evenly as possible across ``N`` ranks in multiples of + ``block``, giving the remainder to the first ranks.""" + base, extra = divmod(dim // block, N) + return [(base + (r < extra)) * block for r in range(N)] + + +def _quantized_output_sizes(dim, N, group_size): + """Split the output rows of a quantized layer. + + Rows can be split anywhere, but splitting at multiples of ``group_size`` + when possible matches the default split of a paired sharded-to-all layer. + """ + if dim % group_size == 0 and dim // group_size >= N: + return _rank_sizes(dim, N, group_size) + return _rank_sizes(dim, N) + + +def _resolve_sizes(dim, N, sizes, name, block=1): + """Validate the size of each rank's shard of ``dim``, splitting it evenly + in multiples of ``block`` if ``sizes`` is not given.""" + if sizes is None: + sizes = _rank_sizes(dim, N, block) + if len(sizes) != N or sum(sizes) != dim: + raise ValueError(f"Expected {N} sizes that sum to {dim} but got {sizes}.") + if min(sizes) <= 0: + raise ValueError(f"Cannot shard the {name} of size {dim} across {N} devices.") + if any(s % block for s in sizes): + raise ValueError(f"The sizes {sizes} must be multiples of {block}.") + return list(sizes) + + +def _layer_sizes(dim, N, segments, sizes, default, name): + """Return the per-rank sizes of a layer and the sizes to shard its + parameters with. Layers with more than one segment are split evenly.""" + if segments == 1: + sizes = default if sizes is None else sizes + return sizes, sizes + if sizes is not None: + raise ValueError("Explicit sizes are only supported with segments=1.") + if dim % N != 0: + raise ValueError(f"Cannot shard the {name} of size {dim} across {N} devices.") + return [dim // N] * N, None + + +def _split_sizes(weight, sizes, axis): + """Split ``weight`` along ``axis`` into parts proportional to ``sizes``. + + ``sizes`` may sum to more than the length of ``axis``, as is the case for + packed quantized weights and their scales, but every boundary must land on + an integer index. + """ + dim = weight.shape[axis] + total = sum(sizes) + indices = [] + boundary = 0 + for s in sizes[:-1]: + boundary += s + index, remainder = divmod(boundary * dim, total) + if remainder != 0: + raise ValueError( + f"Cannot split an axis of size {dim} according to sizes {sizes}." + ) + indices.append(index) + return mx.split(weight, indices, axis=axis) + + +def _quantized_sizes(parameters, sharding_predicate, N, quantized_paths): + """Return the per-rank sizes of each quantized module's sharded axis. + + A quantized weight is packed along its last axis and its scales and biases + are grouped along the same axis, so all three are split from one list of + sizes in unpacked elements to keep them at matching boundaries. The sizes + are multiples of ``group_size`` so that a paired all-to-sharded and + sharded-to-all layer split their common dimension the same way. + """ + rank_sizes = {} + for path, weight in tree_flatten(parameters): + module, _, name = path.rpartition(".") + if name != "weight" or module not in quantized_paths: + continue + shard_spec = sharding_predicate(path, weight) + if shard_spec is None: + continue + if isinstance(shard_spec, tuple): + axis, segments = shard_spec + else: + axis, segments = shard_spec, 1 + if segments != 1: + continue + group_size, bits = quantized_paths[module] + if axis % weight.ndim == weight.ndim - 1: + dim = (weight.shape[axis] * 32) // bits + sizes = _rank_sizes(dim, N, group_size) + else: + dim = weight.shape[axis] + sizes = _quantized_output_sizes(dim, N, group_size) + if min(sizes) <= 0: + raise ValueError( + f"Cannot shard the quantized {module or 'module'} of size " + f"{dim} across {N} devices." + ) + rank_sizes[module] = sizes + return rank_sizes + + def _shard( parameters: dict, sharding_predicate: Callable, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, + quantized_paths: Optional[dict] = None, ): """Returns a new parameter tree with the weights sharded according to the sharding_predicate. The sharding predicate should return the sharding axis and optionally also - the segments that comprise the weight. + the segments that comprise the weight. If ``sizes`` is provided, each rank + gets the corresponding proportion of the sharded axis instead of an equal + part. ``sizes`` requires a single segment. ``quantized_paths`` maps the + path of a quantized module to its ``(group_size, bits)`` so that its + packed weight and grouped metadata stay at matching boundaries. """ group = group or mx.distributed.init() N = group.size() r = group.rank() + quantized_sizes = ( + _quantized_sizes(parameters, sharding_predicate, N, quantized_paths) + if sizes is None and quantized_paths + else {} + ) def _shard_fn(path, weight): if not isinstance(weight, mx.array): @@ -71,6 +188,21 @@ def _shard_fn(path, weight): "The sharding function should return int or tuple[int, list]" ) + rank_sizes = sizes + if rank_sizes is None: + rank_sizes = quantized_sizes.get(path.rpartition(".")[0]) + if rank_sizes is None and segments == 1: + # Split as evenly as possible rather than requiring the axis to + # divide by the number of devices. + rank_sizes = _rank_sizes(weight.shape[axis], N) + if min(rank_sizes) <= 0: + raise ValueError( + f"Cannot shard {path!r} of size {weight.shape[axis]} " + f"across {N} devices." + ) + if rank_sizes is not None: + return mx.contiguous(_split_sizes(weight, rank_sizes, axis)[r]) + return mx.contiguous( mx.concatenate( [_split(part, N, axis)[r] for part in _split(weight, segments, axis)], @@ -152,7 +284,18 @@ def shard_inplace( if sharding == "all-to-sharded" else _sharded_to_all(segments) ) - module.update(_shard(module.parameters(), sharding, group)) + # Detect quantized modules, including third party ones, so that their + # packed weights and grouped metadata are split at the same boundaries. + quantized_paths = { + path: (child.group_size, child.bits) + for path, child in module.named_modules() + if isinstance(getattr(child, "group_size", None), int) + and isinstance(getattr(child, "bits", None), int) + and "scales" in child + } + module.update( + _shard(module.parameters(), sharding, group, quantized_paths=quantized_paths) + ) def shard_linear( @@ -161,6 +304,7 @@ def shard_linear( *, segments: Union[int, list] = 1, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): """Create a new linear layer that has its parameters sharded and also performs distributed communication either in the forward or backward @@ -177,6 +321,10 @@ def shard_linear( segments (int or list): The segments to use. Default: ``1``. group (mlx.core.distributed.Group): The distributed group to shard across. If not set, the global group will be used. Default: ``None``. + sizes (list, optional): The size of each rank's shard of the sharded + dimension. If not set, the dimension is split as evenly as + possible with the remainder going to the first ranks. Uneven + splits require ``segments=1``. Default: ``None``. """ _check_sharding(sharding) fns = { @@ -186,7 +334,7 @@ def shard_linear( ("sharded-to-all", False): QuantizedShardedToAllLinear.from_quantized_linear, } return fns[sharding, isinstance(module, Linear)]( - module, segments=segments, group=group + module, segments=segments, group=group, sizes=sizes ) @@ -204,6 +352,9 @@ class AllToShardedLinear(Module): group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. + sizes (list, optional): The size of each rank's shard of the output + features. If not set, they are split as evenly as possible. + Default: ``None``. """ def __init__( @@ -212,6 +363,7 @@ def __init__( output_dims: int, bias: bool = True, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -220,27 +372,25 @@ def __init__( self.group = group or mx.distributed.init() N = self.group.size() - if (output_dims % N) != 0: - raise ValueError( - f"Cannot shard the output of size {output_dims} across {N} devices." - ) + sizes = _resolve_sizes(output_dims, N, sizes, "output") + local_output_dims = sizes[self.group.rank()] + self._output_dims = output_dims self.weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims // N, input_dims), + shape=(local_output_dims, input_dims), ) if bias: self.bias = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims // N,), + shape=(local_output_dims,), ) def _extra_repr(self) -> str: - out_dims, in_dims = self.weight.shape - N = self.group.size() - out_dims *= N + in_dims = self.weight.shape[1] + out_dims = self._output_dims return f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}" def __call__(self, x: mx.array) -> mx.array: @@ -261,12 +411,21 @@ def from_linear( *, segments: Union[int, list] = 1, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): group = group or mx.distributed.init() + N = group.size() output_dims, input_dims = linear_layer.weight.shape + sizes, shard_sizes = _layer_sizes( + output_dims, N, segments, sizes, _rank_sizes(output_dims, N), "output" + ) - sl = cls(input_dims, output_dims, hasattr(linear_layer, "bias"), group) - sl.update(_shard(linear_layer.parameters(), _all_to_sharded(segments), group)) + sl = cls(input_dims, output_dims, hasattr(linear_layer, "bias"), group, sizes) + sl.update( + _shard( + linear_layer.parameters(), _all_to_sharded(segments), group, shard_sizes + ) + ) return sl @@ -288,6 +447,9 @@ class ShardedToAllLinear(Module): group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. + sizes (list, optional): The size of each rank's shard of the input + features. If not set, they are split as evenly as possible. + Default: ``None``. """ def __init__( @@ -296,6 +458,7 @@ def __init__( output_dims: int, bias: bool = True, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -304,15 +467,14 @@ def __init__( self.group = group or mx.distributed.init() N = self.group.size() - if (input_dims % N) != 0: - raise ValueError( - f"The input of size {input_dims} cannot be sharded across {N} devices." - ) + sizes = _resolve_sizes(input_dims, N, sizes, "input") + local_input_dims = sizes[self.group.rank()] + self._input_dims = input_dims self.weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims, input_dims // N), + shape=(output_dims, local_input_dims), ) if bias: self.bias = mx.random.uniform( @@ -322,9 +484,8 @@ def __init__( ) def _extra_repr(self) -> str: - N = self.group.size() - out_dims, in_dims = self.weight.shape - in_dims *= N + out_dims = self.weight.shape[0] + in_dims = self._input_dims return f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}" def __call__(self, x: mx.array) -> mx.array: @@ -344,12 +505,21 @@ def from_linear( *, segments: Union[int, list] = 1, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): group = group or mx.distributed.init() + N = group.size() output_dims, input_dims = linear_layer.weight.shape + sizes, shard_sizes = _layer_sizes( + input_dims, N, segments, sizes, _rank_sizes(input_dims, N), "input" + ) - sl = cls(input_dims, output_dims, hasattr(linear_layer, "bias"), group) - sl.update(_shard(linear_layer.parameters(), _sharded_to_all(segments), group)) + sl = cls(input_dims, output_dims, hasattr(linear_layer, "bias"), group, sizes) + sl.update( + _shard( + linear_layer.parameters(), _sharded_to_all(segments), group, shard_sizes + ) + ) return sl @@ -376,6 +546,9 @@ class QuantizedAllToShardedLinear(Module): group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. + sizes (list, optional): The size of each rank's shard of the output + features. If not set, they are split as evenly as possible, + preferring multiples of ``group_size``. Default: ``None``. """ def __init__( @@ -387,6 +560,7 @@ def __init__( bits: int = 4, mode: str = "affine", group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -400,15 +574,16 @@ def __init__( self.group = group or mx.distributed.init() N = self.group.size() - if (output_dims % N) != 0: - raise ValueError( - f"Cannot shard the output of size {output_dims} across {N} devices." - ) + if sizes is None: + sizes = _quantized_output_sizes(output_dims, N, group_size) + sizes = _resolve_sizes(output_dims, N, sizes, "output") + local_output_dims = sizes[self.group.rank()] + self._output_dims = output_dims weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims // N, input_dims), + shape=(local_output_dims, input_dims), ) self.weight, self.scales, *biases = mx.quantize( weight, group_size, bits, mode=mode @@ -417,7 +592,7 @@ def __init__( # And bias if needed if bias: - self.bias = mx.zeros((output_dims // N,)) + self.bias = mx.zeros((local_output_dims,)) # Freeze this model's parameters self.freeze() @@ -431,7 +606,7 @@ def unfreeze(self, *args, **kwargs): def _extra_repr(self) -> str: out_dims, in_dims = self.weight.shape in_dims = (in_dims * 32) // self.bits - out_dims *= self.group.size() + out_dims = self._output_dims return ( f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, " f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}" @@ -462,25 +637,38 @@ def from_quantized_linear( *, segments: Union[int, list] = 1, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): group = group or mx.distributed.init() + N = group.size() output_dims, input_dims = quantized_linear_layer.weight.shape input_dims = (input_dims * 32) // quantized_linear_layer.bits + group_size = quantized_linear_layer.group_size + sizes, shard_sizes = _layer_sizes( + output_dims, + N, + segments, + sizes, + _quantized_output_sizes(output_dims, N, group_size), + "output", + ) sl = cls( input_dims, output_dims, hasattr(quantized_linear_layer, "bias"), - group_size=quantized_linear_layer.group_size, + group_size=group_size, bits=quantized_linear_layer.bits, mode=getattr(quantized_linear_layer, "mode", "affine"), group=group, + sizes=sizes, ) sl.update( _shard( quantized_linear_layer.parameters(), _all_to_sharded(segments), group, + shard_sizes, ) ) @@ -511,6 +699,9 @@ class QuantizedShardedToAllLinear(Module): group (mx.distributed.Group, optional): The sharding will happen across this group. If not set then the global group is used. Default is ``None``. + sizes (list, optional): The size of each rank's shard of the input + features. Each size must be a multiple of ``group_size``. If not + set, they are split as evenly as possible. Default: ``None``. """ def __init__( @@ -522,6 +713,7 @@ def __init__( bits: int = 4, mode: str = "affine", group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -535,15 +727,14 @@ def __init__( self.group = group or mx.distributed.init() N = self.group.size() - if (input_dims % N) != 0: - raise ValueError( - f"The input of size {input_dims} cannot be sharded across {N} devices." - ) + sizes = _resolve_sizes(input_dims, N, sizes, "input", group_size) + local_input_dims = sizes[self.group.rank()] + self._input_dims = input_dims weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims, input_dims // N), + shape=(output_dims, local_input_dims), ) self.weight, self.scales, *biases = mx.quantize( weight, group_size, bits, mode=mode @@ -564,8 +755,8 @@ def unfreeze(self, *args, **kwargs): self.freeze(recurse=False) def _extra_repr(self) -> str: - out_dims, in_dims = self.weight.shape - in_dims = (in_dims * 32) // self.bits * self.group.size() + out_dims = self.weight.shape[0] + in_dims = self._input_dims return ( f"input_dims={in_dims}, output_dims={out_dims}, bias={'bias' in self}, " f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}" @@ -594,25 +785,38 @@ def from_quantized_linear( *, segments: Union[int, list] = 1, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): group = group or mx.distributed.init() + N = group.size() output_dims, input_dims = quantized_linear_layer.weight.shape input_dims = (input_dims * 32) // quantized_linear_layer.bits + group_size = quantized_linear_layer.group_size + sizes, shard_sizes = _layer_sizes( + input_dims, + N, + segments, + sizes, + _rank_sizes(input_dims, N, group_size), + "input", + ) sl = cls( input_dims, output_dims, hasattr(quantized_linear_layer, "bias"), - group_size=quantized_linear_layer.group_size, + group_size=group_size, bits=quantized_linear_layer.bits, mode=getattr(quantized_linear_layer, "mode", "affine"), group=group, + sizes=sizes, ) sl.update( _shard( quantized_linear_layer.parameters(), _sharded_to_all(segments), group, + shard_sizes, ) ) diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index 7d6e56aaac..0252d39506 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -9,6 +9,21 @@ from mlx.nn.utils import average_gradients, clip_grad_norm_sharded +class _FakeGroup: + """A group that only reports its size and rank so that every rank's shard + can be checked in a single process.""" + + def __init__(self, size, rank): + self._size = size + self._rank = rank + + def size(self): + return self._size + + def rank(self): + return self._rank + + class MLXDistributedCommonTestCase(mlx_tests.MLXTestCase): def test_average_gradients(self): original_all_sum = mx.distributed.all_sum @@ -265,6 +280,249 @@ def dummy_loss(model, x, y): ) ) + def test_shard_linear_uneven(self): + mx.random.seed(0xF0F0F0F0) + world = mx.distributed.init() + N = world.size() + + def part(sizes): + start = sum(sizes[: world.rank()]) + return slice(None), slice(start, start + sizes[world.rank()]) + + # Explicit sizes and the default remainder split + explicit = [16 * (i + 1) for i in range(N)] + default = [17] + [16] * (N - 1) + for sizes, kwargs in ((explicit, {"sizes": explicit}), (default, {})): + dims = sum(sizes) + x = mx.random.normal((4, dims)) + lin = nn.Linear(dims, dims, bias=True) + slin1 = shard_linear(lin, "all-to-sharded", **kwargs) + slin2 = shard_linear(lin, "sharded-to-all", **kwargs) + y = lin(x) + self.assertTrue( + mx.allclose(y[part(sizes)], slin1(x), atol=self.atol, rtol=self.rtol) + ) + self.assertTrue( + mx.allclose(y, slin2(x[part(sizes)]), atol=self.atol, rtol=self.rtol) + ) + + # QuantizedMatmul is not supported on CUDA + if not mx.cuda.is_available(): + # Explicit sizes and the default split keep quantization groups intact + explicit = [32 * (i + 1) for i in range(N)] + default = [96] + [64] * (N - 1) + for sizes, kwargs in ((explicit, {"sizes": explicit}), (default, {})): + dims = sum(sizes) + x = mx.random.normal((4, dims)) + qlin = nn.Linear(dims, dims).to_quantized(group_size=32, bits=4) + slin1 = shard_linear(qlin, "all-to-sharded", **kwargs) + slin2 = shard_linear(qlin, "sharded-to-all", **kwargs) + y = qlin(x) + # Uneven splits change the per-rank matmul sizes, so the + # quantized accumulation order (and thus rounding) differs + # from the unsharded reference more than an even split does. + self.assertTrue( + mx.allclose(y[part(sizes)], slin1(x), atol=1e-5, rtol=1e-3) + ) + self.assertTrue( + mx.allclose(y, slin2(x[part(sizes)]), atol=1e-5, rtol=1e-3) + ) + + # Check the backward pass + def dummy_loss(model, x, y): + return (model(x) * y).sum() + + sizes = [16 * (i + 1) for i in range(N)] + dims = sum(sizes) + mod = nn.Sequential(nn.Linear(dims, dims), nn.Linear(dims, dims)) + smod = nn.Sequential( + shard_linear(mod.layers[0], "all-to-sharded", sizes=sizes), + shard_linear(mod.layers[1], "sharded-to-all", sizes=sizes), + ) + x = mx.random.normal((4, dims)) + y = mx.random.normal((4, dims)) + l1, g1 = nn.value_and_grad(mod, dummy_loss)(mod, x, y) + l2, g2 = nn.value_and_grad(smod, dummy_loss)(smod, x, y) + mx.eval(l1, g1, l2, g2) + + # Uneven splits change the summation order of the distributed + # gradient reduction, so the tolerance is looser than an even split + # needs. + rows = part(sizes)[1] + self.assertTrue(mx.allclose(l1, l2, atol=1e-5, rtol=1e-3)) + for key in ("weight", "bias"): + self.assertTrue( + mx.allclose( + g1["layers"][0][key][rows], + g2["layers"][0][key], + atol=1e-5, + rtol=1e-3, + ) + ) + self.assertTrue( + mx.allclose( + g1["layers"][1]["weight"][:, rows], + g2["layers"][1]["weight"], + atol=1e-5, + rtol=1e-3, + ) + ) + + def test_shard_linear_uneven_shards(self): + lin = nn.Linear(11, 12) + + # Explicit sizes + sizes = [5, 4, 3] + shards = [ + shard_linear(lin, "all-to-sharded", sizes=sizes, group=_FakeGroup(3, r)) + for r in range(3) + ] + self.assertEqual([s.weight.shape for s in shards], [(5, 11), (4, 11), (3, 11)]) + self.assertTrue( + mx.array_equal(mx.concatenate([s.weight for s in shards]), lin.weight) + ) + self.assertTrue( + mx.array_equal(mx.concatenate([s.bias for s in shards]), lin.bias) + ) + self.assertIn("output_dims=12", repr(shards[0])) + + # The default split gives the remainder to the first ranks + shards = [ + shard_linear(lin, "sharded-to-all", group=_FakeGroup(3, r)) + for r in range(3) + ] + self.assertEqual([s.weight.shape for s in shards], [(12, 4), (12, 4), (12, 3)]) + self.assertTrue( + mx.array_equal( + mx.concatenate([s.weight for s in shards], axis=1), lin.weight + ) + ) + self.assertIn("input_dims=11", repr(shards[0])) + + def test_shard_linear_uneven_quantized_shards(self): + # Split the packed weights, scales and biases at the same boundaries + sizes = [96, 64, 32] + for bits in (2, 3, 4, 5, 6, 8): + qlin = nn.Linear(192, 16).to_quantized(group_size=32, bits=bits) + shards = [ + shard_linear( + qlin, "sharded-to-all", sizes=sizes, group=_FakeGroup(3, r) + ) + for r in range(3) + ] + for key in ("weight", "scales", "biases"): + parts = [s[key] for s in shards] + self.assertEqual( + [p.shape[1] for p in parts], + [qlin[key].shape[1] * s // 192 for s in sizes], + ) + self.assertTrue( + mx.array_equal(mx.concatenate(parts, axis=1), qlin[key]) + ) + self.assertIn("input_dims=192", repr(shards[0])) + + # Output rows can be split anywhere + qlin = nn.Linear(64, 10).to_quantized(group_size=32, bits=4) + shards = [ + shard_linear(qlin, "all-to-sharded", sizes=[7, 3], group=_FakeGroup(2, r)) + for r in range(2) + ] + for key in ("weight", "scales", "biases", "bias"): + self.assertTrue( + mx.array_equal(mx.concatenate([s[key] for s in shards]), qlin[key]) + ) + self.assertIn("output_dims=10", repr(shards[0])) + + # Paired default splits match when the output rows allow it + qin = nn.Linear(64, 160).to_quantized(group_size=32, bits=4) + qout = nn.Linear(160, 64).to_quantized(group_size=32, bits=4) + for r, rows in enumerate((64, 64, 32)): + s1 = shard_linear(qin, "all-to-sharded", group=_FakeGroup(3, r)) + s2 = shard_linear(qout, "sharded-to-all", group=_FakeGroup(3, r)) + self.assertEqual(s1.weight.shape[0], rows) + self.assertEqual(s2.scales.shape[1] * s2.group_size, rows) + + def test_shard_linear_uneven_errors(self): + lin = nn.Linear(10, 12) + group = _FakeGroup(3, 0) + + # The sizes must have one entry per rank and sum to the sharded dimension + with self.assertRaises(ValueError): + shard_linear(lin, "all-to-sharded", sizes=[5, 4, 2], group=group) + with self.assertRaises(ValueError): + shard_linear(lin, "all-to-sharded", sizes=[6, 6], group=group) + + # Every rank needs a non-empty shard + with self.assertRaises(ValueError): + shard_linear(lin, "all-to-sharded", sizes=[12, 0, 0], group=group) + with self.assertRaises(ValueError): + shard_linear(lin, "sharded-to-all", sizes=[11, -1, 0], group=group) + with self.assertRaises(ValueError): + shard_linear(nn.Linear(10, 2), "all-to-sharded", group=group) + + # Uneven splits require a single segment + with self.assertRaises(ValueError): + shard_linear( + lin, "all-to-sharded", segments=2, sizes=[4, 4, 4], group=group + ) + with self.assertRaises(ValueError): + shard_linear(nn.Linear(10, 14), "all-to-sharded", segments=2, group=group) + + # Quantized input shards must be multiples of the group size + qlin = nn.Linear(96, 10).to_quantized(group_size=32, bits=4) + with self.assertRaisesRegex(ValueError, "multiples of 32"): + shard_linear(qlin, "sharded-to-all", sizes=[48, 48], group=_FakeGroup(2, 0)) + with self.assertRaises(ValueError): + shard_linear(qlin, "sharded-to-all", group=_FakeGroup(4, 0)) + with self.assertRaises(ValueError): + shard_linear(qlin, "all-to-sharded", group=_FakeGroup(11, 0)) + + def test_shard_inplace_uneven(self): + # A dimension that does not divide by the number of devices is split + # as evenly as possible instead of raising. + lin = nn.Linear(12, 11) + shards = [] + for r in range(3): + m = nn.Linear(12, 11) + m.update(lin.parameters()) + shard_inplace(m, "all-to-sharded", group=_FakeGroup(3, r)) + shards.append(m.weight) + self.assertEqual([s.shape[0] for s in shards], [4, 4, 3]) + self.assertTrue(mx.array_equal(mx.concatenate(shards), lin.weight)) + + # QuantizedMatmul is not supported on CUDA + if not mx.cuda.is_available(): + # A quantized module's packed weight and its grouped scales and + # biases must be split at the same boundaries, so the sizes are + # multiples of the group size. + qlin = nn.Linear(1408, 64).to_quantized(group_size=64, bits=4) + parts = {"weight": [], "scales": [], "biases": []} + for r in range(3): + m = nn.Linear(1408, 64).to_quantized(group_size=64, bits=4) + m.update(qlin.parameters()) + shard_inplace(m, "sharded-to-all", group=_FakeGroup(3, r)) + for key in parts: + parts[key].append(m[key]) + # 1408 is 22 groups of 64, split as 8, 7 and 7 groups. + self.assertEqual([p.shape[-1] for p in parts["weight"]], [64, 56, 56]) + self.assertEqual([p.shape[-1] for p in parts["scales"]], [8, 7, 7]) + for key, whole in parts.items(): + self.assertTrue( + mx.array_equal(mx.concatenate(whole, axis=-1), qlin[key]) + ) + + # A paired all-to-sharded layer splits the same dimension the same + # way, so the two layers agree on each rank's share. + qout = nn.Linear(64, 1408).to_quantized(group_size=64, bits=4) + rows = [] + for r in range(3): + m = nn.Linear(64, 1408).to_quantized(group_size=64, bits=4) + m.update(qout.parameters()) + shard_inplace(m, "all-to-sharded", group=_FakeGroup(3, r)) + rows.append(m.weight.shape[0]) + self.assertEqual(rows, [512, 448, 448]) + self.assertEqual(rows, [p.shape[-1] * 64 for p in parts["scales"]]) + def test_shard_predicate(self): mx.random.seed(0xF0F0F0F0)