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)