diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index 16979b5097..072dd20440 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/python/mlx/nn/layers/distributed.py @@ -7,7 +7,6 @@ import mlx.core as mx from mlx.nn.layers.base import Module from mlx.nn.layers.linear import Linear -from mlx.nn.layers.quantized import QuantizedLinear from mlx.utils import tree_flatten, tree_map_with_path, tree_unflatten @@ -37,43 +36,302 @@ 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. + + Give remainder blocks to the first ranks. All sizes are multiples of + ``block``. + """ + if dim % block != 0: + raise ValueError( + f"Cannot shard dimension of size {dim} across {N} devices: " + f"the size is not a multiple of the quantization group size " + f"({block}), so it cannot be split into quantize-able chunks." + ) + + n_blocks = dim // block + base_blocks = n_blocks // N + extra_blocks = n_blocks - base_blocks * N + return [(base_blocks + (1 if r < extra_blocks else 0)) * block for r in range(N)] + + +def _quantized_output_sizes(dim, N, group_size): + """Use quantization-group boundaries when a paired input shard can too.""" + block = group_size if dim % group_size == 0 and dim // group_size >= N else 1 + return _rank_sizes(dim, N, block=block) + + +def _resolve_sizes(dim, n_segments, N, sizes, block=1): + """Resolve total and per-segment rank sizes for equal-sized segments. + + Return ``(cls_sizes, shard_sizes)``, where ``cls_sizes`` contains each + rank's total and ``shard_sizes`` contains its share of one segment. + """ + if dim % n_segments != 0: + raise ValueError( + f"Cannot split dimension of size {dim} into {n_segments} equal " + f"segments." + ) + seg_dim = dim // n_segments + if sizes is None: + shard_sizes = _rank_sizes(seg_dim, N, block=block) + else: + if len(sizes) != N: + raise ValueError( + f"Explicit per-rank sizes {sizes} has {len(sizes)} entries, " + f"but the group has {N} ranks -- these must match." + ) + if sum(sizes) != dim: + raise ValueError(f"Explicit per-rank sizes {sizes} do not sum to {dim}.") + if any(s % n_segments != 0 for s in sizes): + raise ValueError( + f"Explicit per-rank sizes {sizes} must each be evenly " + f"divisible by the segment count {n_segments} -- each " + f"segment needs an equal, exact share of every rank's " + f"allocation." + ) + shard_sizes = [s // n_segments for s in sizes] + if block != 1 and any(s % block != 0 for s in shard_sizes): + raise ValueError( + f"Explicit per-rank sizes {sizes}, divided evenly across " + f"{n_segments} segments, give each segment's share as " + f"{shard_sizes} -- these must each be a multiple of the " + f"quantization group_size ({block})." + ) + cls_sizes = [s * n_segments for s in shard_sizes] + return cls_sizes, shard_sizes + + +def _check_no_zero_shares(sizes): + """Reject zero-width quantized shards on every rank.""" + if any(s == 0 for s in sizes): + raise ValueError( + f"Resolved per-rank sizes {sizes} give rank(s) " + f"{[r for r, s in enumerate(sizes) if s == 0]} a zero-width " + f"shard -- for a quantized layer this can silently produce " + f"wrong (not just zero) results from mx.quantized_matmul, not " + f"merely a wasted rank. Use fewer ranks, fewer segments, or a " + f"smaller quantization block size." + ) + + +def _n_segments(segments, reason="of an explicit `sizes` split"): + if isinstance(segments, int) and not isinstance(segments, bool): + return segments + raise ValueError( + f"segments={segments!r}: only a plain int segment count is " + f"supported here (a list of fractional or index-based segment " + f"boundaries isn't, since unequal segments have no single " + f"well-defined per-segment share {reason})." + ) + + +def _segment_sizes_from_spec(dim, segments): + """Return the segment sizes produced by ``_split`` for ``dim``.""" + if isinstance(segments, int) or isinstance(segments[0], int): + indices_or_sections = segments + else: + indices_or_sections = [int(s * dim) for s in segments] + return [int(part.size) for part in mx.split(mx.arange(dim), indices_or_sections)] + + +def _split_uneven(weight, N, axis, sizes=None): + """Split ``weight`` along ``axis`` into ``N`` pieces. + + ``sizes`` may use a logical dimension whose length differs from the array + axis, as with packed quantized weights and grouped scales. In that case, + scale each boundary to the array axis and require an exact integer index. + """ + dim = weight.shape[axis] + if sizes is None: + local_sizes = _rank_sizes(dim, N) + else: + if len(sizes) != N: + raise ValueError( + f"sizes {sizes} has {len(sizes)} entries, but N={N} ranks " + f"were requested -- these must match, or ranks beyond " + f"len(sizes) would silently get no shard at all." + ) + total = sum(sizes) + if dim == total: + local_sizes = list(sizes) + elif total != 0 and all( + (sum(sizes[: i + 1]) * dim) % total == 0 for i in range(len(sizes) - 1) + ): + # Scale logical boundaries instead of sizes to support fractional + # packing ratios such as 3-bit weights. + ratio_num, ratio_den = dim, total + local_sizes = [] + acc_real = 0 + acc_local = 0 + for s in sizes[:-1]: + acc_real += s + nxt = (acc_real * ratio_num) // ratio_den + local_sizes.append(nxt - acc_local) + acc_local = nxt + local_sizes.append(dim - acc_local) + else: + raise ValueError( + f"Explicit per-rank sizes {sizes} (summing to {total}) are " + f"not compatible with this array's axis size {dim}: at " + f"least one cumulative boundary, scaled by the ratio " + f"{dim}/{total}, doesn't land on a whole number." + ) + indices = [] + acc = 0 + for s in local_sizes[:-1]: + acc += s + indices.append(acc) + return mx.split(weight, indices, axis=axis) + + +def _normalize_axis(axis, ndim): + if ( + not isinstance(axis, int) + or isinstance(axis, bool) + or axis < -ndim + or axis >= ndim + ): + raise ValueError(f"Invalid sharding axis {axis} for an array with {ndim} axes.") + return axis % ndim + + def _shard( parameters: dict, sharding_predicate: Callable, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, + quantized_paths=None, ): - """Returns a new parameter tree with the weights sharded according to the - sharding_predicate. + """Return parameters sharded according to ``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. ``sizes`` sets logical per-rank + sizes. ``quantized_paths`` keeps packed weights and grouped metadata at + matching logical boundaries. """ group = group or mx.distributed.init() N = group.size() r = group.rank() + def _group_of(path): + return path.rsplit(".", 1)[0] if "." in path else "" + + def _leaf_name(path): + return path.rsplit(".", 1)[-1] + + quantized_paths = {} if quantized_paths is None else dict(quantized_paths) + predicate_results = {} + sibling_sizes = {} + if sizes is None and quantized_paths: + for path, leaf in tree_flatten(parameters): + if ( + not isinstance(leaf, mx.array) + or _leaf_name(path) != "weight" + or _group_of(path) not in quantized_paths + ): + continue + shard_spec = sharding_predicate(path, leaf) + predicate_results[path] = shard_spec + if shard_spec is None: + continue + axis = shard_spec if isinstance(shard_spec, int) else shard_spec[0] + axis = _normalize_axis(axis, leaf.ndim) + segments = 1 if isinstance(shard_spec, int) else shard_spec[1] + group_size, bits = quantized_paths[_group_of(path)] + # Quantization packs and groups only the last weight axis. + is_packed_axis = axis == leaf.ndim - 1 + weight_seg_sizes = [ + part.shape[axis] for part in _split(leaf, segments, axis) + ] + if is_packed_axis: + real_seg_sizes = [sz * 32 // bits for sz in weight_seg_sizes] + weight_seg_rank_sizes = [ + _rank_sizes(sz, N, block=group_size) for sz in real_seg_sizes + ] + else: + weight_seg_rank_sizes = [ + _quantized_output_sizes(sz, N, group_size) + for sz in weight_seg_sizes + ] + # Quantized matmul receives the concatenated segments. + total_rank_sizes = [ + sum(seg[rk] for seg in weight_seg_rank_sizes) for rk in range(N) + ] + _check_no_zero_shares(total_rank_sizes) + sibling_sizes[_group_of(path)] = ( + axis, + weight_seg_sizes, + weight_seg_rank_sizes, + ) + def _shard_fn(path, weight): if not isinstance(weight, mx.array): return weight - s = sharding_predicate(path, weight) - if s is None: + if path in predicate_results: + shard_spec = predicate_results[path] + else: + shard_spec = sharding_predicate(path, weight) + if shard_spec is None: return weight axis = None segments = 1 - if isinstance(s, int): - axis = s - elif isinstance(s, tuple): - axis, segments = s + if isinstance(shard_spec, int): + axis = shard_spec + elif isinstance(shard_spec, tuple): + axis, segments = shard_spec else: raise ValueError( "The sharding function should return int or tuple[int, list]" ) + axis = _normalize_axis(axis, weight.ndim) + + quantized_ref = None + if sizes is None and _group_of(path) in quantized_paths: + quantized_ref = sibling_sizes.get(_group_of(path)) + if quantized_ref is not None and quantized_ref[0] != axis: + quantized_ref = None + if quantized_ref is None and _leaf_name(path) in ("scales", "biases"): + raise ValueError( + f"Cannot shard quantized parameter {path!r}: its " + f"'weight' sibling was not sharded consistently on " + f"axis {axis} (the sharding predicate returned `None` " + f"for it, or a different axis) -- there is no safe " + f"reference split to keep this leaf's boundaries " + f"consistent with weight's." + ) + + segment_sizes = None + if sizes and isinstance(sizes[0], (list, tuple)): + segment_sizes = sizes + logical_sizes = [sum(s) for s in segment_sizes] + parts = _split_uneven(weight, len(logical_sizes), axis, sizes=logical_sizes) + elif quantized_ref is not None: + # Reuse the weight boundaries for its packed/grouped metadata. + parts = _split_uneven( + weight, len(quantized_ref[1]), axis, sizes=quantized_ref[1] + ) + else: + parts = _split(weight, segments, axis) + per_segment_sizes = quantized_ref[2] if quantized_ref is not None else None return mx.contiguous( mx.concatenate( - [_split(part, N, axis)[r] for part in _split(weight, segments, axis)], + [ + _split_uneven( + part, + N, + axis, + ( + per_segment_sizes[i] + if per_segment_sizes + else segment_sizes[i] if segment_sizes else sizes + ), + )[r] + for i, part in enumerate(parts) + ], axis=axis, ) ) @@ -152,7 +410,22 @@ def shard_inplace( if sharding == "all-to-sharded" else _sharded_to_all(segments) ) - module.update(_shard(module.parameters(), sharding, group)) + # Detect compatible third-party quantized layers by their parameters. + 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 +434,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 @@ -174,9 +448,18 @@ def shard_linear( module (mlx.nn.Module): The linear layer to be sharded. sharding (str): One of "all-to-sharded" and "sharded-to-all" that defines the type of sharding to perform. - segments (int or list): The segments to use. Default: ``1``. + segments (int or list): The segments to split independently before + sharding. Explicit ``sizes`` require an integer segment count. + 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): Explicit sizes for each rank. The sizes must + sum to the sharded dimension. Quantized input sizes must be + multiples of ``group_size``. Default: ``None``. + + .. note:: + Paired quantized layers with different ``group_size`` values should + use explicit matching ``sizes``. """ _check_sharding(sharding) fns = { @@ -186,7 +469,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 +487,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): Explicit per-rank sizes (one entry per rank in + ``group``, summing to the full sharded dimension) to use instead + of the automatic remainder-aware split. Default: ``None``. """ def __init__( @@ -212,6 +498,7 @@ def __init__( output_dims: int, bias: bool = True, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -219,28 +506,45 @@ def __init__( scale = math.sqrt(1.0 / input_dims) self.group = group or mx.distributed.init() N = self.group.size() + r = self.group.rank() - if (output_dims % N) != 0: - raise ValueError( - f"Cannot shard the output of size {output_dims} across {N} devices." - ) + # Each rank gets a possibly uneven slice of the output features. + if sizes is not None: + if len(sizes) != N: + raise ValueError( + f"Explicit per-rank sizes {sizes} has {len(sizes)} " + f"entries, but the group has {N} ranks -- these must " + f"match." + ) + if any(s < 0 for s in sizes): + raise ValueError( + f"Explicit per-rank sizes {sizes} must all be " f"non-negative." + ) + if sum(sizes) != output_dims: + raise ValueError( + f"Explicit per-rank sizes {sizes} do not sum to " + f"output_dims={output_dims}." + ) + my_output_dims = sizes[r] + else: + my_output_dims = _rank_sizes(output_dims, N)[r] + self._total_output_dims = output_dims self.weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims // N, input_dims), + shape=(my_output_dims, input_dims), ) if bias: self.bias = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims // N,), + shape=(my_output_dims,), ) def _extra_repr(self) -> str: out_dims, in_dims = self.weight.shape - N = self.group.size() - out_dims *= N + out_dims = self._total_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 +565,37 @@ 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 + if sizes is None and not isinstance(segments, int): + # Split unequal segments independently across ranks. + seg_sizes = _segment_sizes_from_spec(output_dims, segments) + cls_sizes = [sum(_rank_sizes(s, N)[r] for s in seg_sizes) for r in range(N)] + shard_sizes = None + else: + n_segments = _n_segments(segments) + if n_segments == 1: + cls_sizes = sizes if sizes is not None else _rank_sizes(output_dims, N) + shard_sizes = cls_sizes + else: + cls_sizes, shard_sizes = _resolve_sizes( + output_dims, n_segments, N, sizes + ) - 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, cls_sizes + ) + sl.update( + _shard( + linear_layer.parameters(), + _all_to_sharded(segments), + group, + sizes=shard_sizes, + ) + ) return sl @@ -288,6 +617,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): Explicit per-rank sizes (one entry per rank in + ``group``, summing to the full sharded dimension) to use instead + of the automatic remainder-aware split. Default: ``None``. """ def __init__( @@ -296,6 +628,7 @@ def __init__( output_dims: int, bias: bool = True, group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -303,16 +636,34 @@ def __init__( scale = math.sqrt(1.0 / input_dims) self.group = group or mx.distributed.init() N = self.group.size() + r = self.group.rank() - if (input_dims % N) != 0: - raise ValueError( - f"The input of size {input_dims} cannot be sharded across {N} devices." - ) + # Each rank gets a possibly uneven slice of the input features. + if sizes is not None: + if len(sizes) != N: + raise ValueError( + f"Explicit per-rank sizes {sizes} has {len(sizes)} " + f"entries, but the group has {N} ranks -- these must " + f"match." + ) + if any(s < 0 for s in sizes): + raise ValueError( + f"Explicit per-rank sizes {sizes} must all be " f"non-negative." + ) + if sum(sizes) != input_dims: + raise ValueError( + f"Explicit per-rank sizes {sizes} do not sum to " + f"input_dims={input_dims}." + ) + my_input_dims = sizes[r] + else: + my_input_dims = _rank_sizes(input_dims, N)[r] + self._total_input_dims = input_dims self.weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims, input_dims // N), + shape=(output_dims, my_input_dims), ) if bias: self.bias = mx.random.uniform( @@ -322,9 +673,8 @@ def __init__( ) def _extra_repr(self) -> str: - N = self.group.size() out_dims, in_dims = self.weight.shape - in_dims *= N + in_dims = self._total_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 +694,37 @@ 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 + if sizes is None and not isinstance(segments, int): + # Split unequal segments independently across ranks. + seg_sizes = _segment_sizes_from_spec(input_dims, segments) + cls_sizes = [sum(_rank_sizes(s, N)[r] for s in seg_sizes) for r in range(N)] + shard_sizes = None + else: + n_segments = _n_segments(segments) + if n_segments == 1: + cls_sizes = sizes if sizes is not None else _rank_sizes(input_dims, N) + shard_sizes = cls_sizes + else: + cls_sizes, shard_sizes = _resolve_sizes( + input_dims, n_segments, N, sizes + ) - 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, cls_sizes + ) + sl.update( + _shard( + linear_layer.parameters(), + _sharded_to_all(segments), + group, + sizes=shard_sizes, + ) + ) return sl @@ -376,6 +751,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): Explicit per-rank sizes (one entry per rank in + ``group``, summing to the full sharded dimension) to use instead + of the automatic remainder-aware split. Default: ``None``. """ def __init__( @@ -387,6 +765,7 @@ def __init__( bits: int = 4, mode: str = "affine", group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -399,16 +778,36 @@ def __init__( scale = math.sqrt(1.0 / input_dims) self.group = group or mx.distributed.init() N = self.group.size() + r = self.group.rank() - if (output_dims % N) != 0: - raise ValueError( - f"Cannot shard the output of size {output_dims} across {N} devices." - ) + # Prefer group boundaries to match a paired input-sharded layer. + if sizes is not None: + if len(sizes) != N: + raise ValueError( + f"Explicit per-rank sizes {sizes} has {len(sizes)} " + f"entries, but the group has {N} ranks -- these must " + f"match." + ) + if any(s < 0 for s in sizes): + raise ValueError( + f"Explicit per-rank sizes {sizes} must all be " f"non-negative." + ) + if sum(sizes) != output_dims: + raise ValueError( + f"Explicit per-rank sizes {sizes} do not sum to " + f"output_dims={output_dims}." + ) + my_output_dims = sizes[r] + else: + sizes = _quantized_output_sizes(output_dims, N, group_size) + my_output_dims = sizes[r] + _check_no_zero_shares(sizes) + self._total_output_dims = output_dims weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims // N, input_dims), + shape=(my_output_dims, input_dims), ) self.weight, self.scales, *biases = mx.quantize( weight, group_size, bits, mode=mode @@ -417,7 +816,7 @@ def __init__( # And bias if needed if bias: - self.bias = mx.zeros((output_dims // N,)) + self.bias = mx.zeros((my_output_dims,)) # Freeze this model's parameters self.freeze() @@ -431,7 +830,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._total_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,10 +861,46 @@ 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 + # Use the same sizes for allocation and parameter slicing. + if sizes is None and not isinstance(segments, int): + seg_sizes = _segment_sizes_from_spec(output_dims, segments) + shard_sizes = [ + _quantized_output_sizes(s, N, quantized_linear_layer.group_size) + for s in seg_sizes + ] + cls_sizes = [sum(s[r] for s in shard_sizes) for r in range(N)] + else: + n_segments = _n_segments(segments) + if n_segments == 1: + cls_sizes = ( + sizes + if sizes is not None + else _quantized_output_sizes( + output_dims, N, quantized_linear_layer.group_size + ) + ) + shard_sizes = cls_sizes + else: + # Output rows need not use group boundaries when there are + # fewer groups than ranks. + seg_dim = output_dims // n_segments + group_size = quantized_linear_layer.group_size + block = ( + group_size + if sizes is None + and seg_dim % group_size == 0 + and seg_dim // group_size >= N + else 1 + ) + cls_sizes, shard_sizes = _resolve_sizes( + output_dims, n_segments, N, sizes, block=block + ) sl = cls( input_dims, @@ -475,12 +910,14 @@ def from_quantized_linear( bits=quantized_linear_layer.bits, mode=getattr(quantized_linear_layer, "mode", "affine"), group=group, + sizes=cls_sizes, ) sl.update( _shard( quantized_linear_layer.parameters(), _all_to_sharded(segments), group, + sizes=shard_sizes, ) ) @@ -511,6 +948,8 @@ 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): Explicit per-rank sizes. Each size must be a + multiple of ``group_size``. Default: ``None``. """ def __init__( @@ -522,6 +961,7 @@ def __init__( bits: int = 4, mode: str = "affine", group: Optional[mx.distributed.Group] = None, + sizes: Optional[list] = None, ): super().__init__() @@ -534,16 +974,41 @@ def __init__( scale = math.sqrt(1.0 / input_dims) self.group = group or mx.distributed.init() N = self.group.size() + r = self.group.rank() - if (input_dims % N) != 0: - raise ValueError( - f"The input of size {input_dims} cannot be sharded across {N} devices." - ) + # Keep each input shard aligned to quantization groups. + if sizes is not None: + if len(sizes) != N: + raise ValueError( + f"Explicit per-rank sizes {sizes} has {len(sizes)} " + f"entries, but the group has {N} ranks -- these must " + f"match." + ) + if any(s < 0 for s in sizes): + raise ValueError( + f"Explicit per-rank sizes {sizes} must all be " f"non-negative." + ) + if sum(sizes) != input_dims: + raise ValueError( + f"Explicit per-rank sizes {sizes} do not sum to " + f"input_dims={input_dims}." + ) + if any(s % group_size != 0 for s in sizes): + raise ValueError( + f"Explicit per-rank sizes {sizes} must each be a " + f"multiple of group_size={group_size}." + ) + my_input_dims = sizes[r] + else: + sizes = _rank_sizes(input_dims, N, block=group_size) + my_input_dims = sizes[r] + _check_no_zero_shares(sizes) + self._total_input_dims = input_dims weight = mx.random.uniform( low=-scale, high=scale, - shape=(output_dims, input_dims // N), + shape=(output_dims, my_input_dims), ) self.weight, self.scales, *biases = mx.quantize( weight, group_size, bits, mode=mode @@ -565,7 +1030,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 * self.group.size() + in_dims = self._total_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,10 +1059,31 @@ 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 + # Use the same logical sizes for packed weights and grouped metadata. + group_size = quantized_linear_layer.group_size + if sizes is None and not isinstance(segments, int): + seg_sizes = _segment_sizes_from_spec(input_dims, segments) + shard_sizes = [_rank_sizes(s, N, block=group_size) for s in seg_sizes] + cls_sizes = [sum(s[r] for s in shard_sizes) for r in range(N)] + else: + n_segments = _n_segments(segments) + if n_segments == 1: + cls_sizes = ( + sizes + if sizes is not None + else _rank_sizes(input_dims, N, block=group_size) + ) + shard_sizes = cls_sizes + else: + cls_sizes, shard_sizes = _resolve_sizes( + input_dims, n_segments, N, sizes, block=group_size + ) sl = cls( input_dims, @@ -607,12 +1093,14 @@ def from_quantized_linear( bits=quantized_linear_layer.bits, mode=getattr(quantized_linear_layer, "mode", "affine"), group=group, + sizes=cls_sizes, ) sl.update( _shard( quantized_linear_layer.parameters(), _sharded_to_all(segments), group, + sizes=shard_sizes, ) ) diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index cbb9663046..ed34df1132 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -5,7 +5,12 @@ import mlx.core as mx import mlx.nn as nn import mlx_tests -from mlx.nn.layers.distributed import shard_inplace, shard_linear +from mlx.nn.layers.distributed import ( + _rank_sizes, + _split_uneven, + shard_inplace, + shard_linear, +) from mlx.nn.utils import average_gradients, clip_grad_norm_sharded @@ -118,12 +123,12 @@ def test_shard_linear(self): # Prepare inputs world = mx.distributed.init() + # Match shard_linear's remainder distribution. + _sizes1024 = _rank_sizes(1024, world.size()) + _start1024 = sum(_sizes1024[: world.rank()]) part = ( slice(None), - slice( - world.rank() * 1024 // world.size(), - (world.rank() + 1) * 1024 // world.size(), - ), + slice(_start1024, _start1024 + _sizes1024[world.rank()]), ) x = mx.random.normal((4, 1024)) @@ -139,16 +144,24 @@ def test_shard_linear(self): # And their quant versions (QuantizedMatmul is not supported on CUDA) if not mx.cuda.is_available(): + # Quantized shards use group-aligned boundaries. + def _quant_part(group_size): + sizes = _rank_sizes(1024, world.size(), block=group_size) + start = sum(sizes[: world.rank()]) + return (slice(None), slice(start, start + sizes[world.rank()])) + + qpart = _quant_part(64) qlin = lin.to_quantized() slin1 = shard_linear(qlin, "all-to-sharded") slin2 = shard_linear(qlin, "sharded-to-all") y = qlin(x) y1 = slin1(x) - y2 = slin2(x[part]) + y2 = slin2(x[qpart]) self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol)) - self.assertTrue(mx.allclose(y[part], y1)) + self.assertTrue(mx.allclose(y[qpart], y1)) # Test non-affine quantization modes (mxfp8) + qpart_mxfp8 = _quant_part(32) qlin_mxfp8 = lin.to_quantized(group_size=32, bits=8, mode="mxfp8") self.assertEqual(qlin_mxfp8.mode, "mxfp8") @@ -165,9 +178,9 @@ def test_shard_linear(self): y = qlin_mxfp8(x) y1 = slin1_mxfp8(x) - y2 = slin2_mxfp8(x[part]) + y2 = slin2_mxfp8(x[qpart_mxfp8]) self.assertTrue(mx.allclose(y, y2, atol=self.atol, rtol=self.rtol)) - self.assertTrue(mx.allclose(y[part], y1)) + self.assertTrue(mx.allclose(y[qpart_mxfp8], y1)) # Check the backward works as expected def dummy_loss(model, x, y): @@ -196,9 +209,10 @@ def dummy_loss(model, x, y): l2, g2 = grad2(smod, x, y) mx.eval(l1, g1, l2, g2) - part = slice( - world.rank() * 128 // world.size(), (world.rank() + 1) * 128 // world.size() - ) + # Match shard_linear's remainder distribution. + _sizes128 = _rank_sizes(128, world.size()) + _start128 = sum(_sizes128[: world.rank()]) + part = slice(_start128, _start128 + _sizes128[world.rank()]) self.assertTrue(mx.allclose(l1, l2)) self.assertTrue( mx.allclose( @@ -265,6 +279,867 @@ def dummy_loss(model, x, y): ) ) + def test_rank_sizes(self): + # Evenly-divisible cases reduce to the plain even split. + self.assertEqual(_rank_sizes(12, 3), [4, 4, 4]) + self.assertEqual(_rank_sizes(256, 4, block=64), [64, 64, 64, 64]) + # Never split a block between ranks. + self.assertEqual(_rank_sizes(128, 4, block=64), [64, 64, 0, 0]) + + # Give remainder units to the first ranks. + self.assertEqual(_rank_sizes(10, 3), [4, 3, 3]) + self.assertEqual(_rank_sizes(11, 3), [4, 4, 3]) + self.assertEqual(_rank_sizes(1, 3), [1, 0, 0]) + + # Block-aware splits preserve whole groups. + self.assertEqual(_rank_sizes(320, 3, block=64), [128, 128, 64]) + for N in (2, 3, 4, 5, 7): + sizes = _rank_sizes(320, N, block=64) + self.assertEqual(sum(sizes), 320) + self.assertTrue(all(s % 64 == 0 for s in sizes)) + self.assertLessEqual(max(sizes) - min(sizes), 64) + + # A dim that isn't a multiple of block can't be split at all. + with self.assertRaises(ValueError): + _rank_sizes(10, 3, block=64) + + def test_split_uneven(self): + mx.random.seed(0) + w = mx.random.normal((10, 8)) + + # sizes=None falls back to the remainder-aware split of this + # array's own axis length. + parts = _split_uneven(w, 3, axis=0) + self.assertEqual([p.shape[0] for p in parts], [4, 3, 3]) + self.assertTrue(mx.array_equal(mx.concatenate(parts, axis=0), w)) + + # Explicit sizes that already match this array's axis length are + # used directly. + parts = _split_uneven(w, 3, axis=0, sizes=[5, 3, 2]) + self.assertEqual([p.shape[0] for p in parts], [5, 3, 2]) + self.assertTrue(mx.array_equal(mx.concatenate(parts, axis=0), w)) + + # Scale logical sizes to a packed axis. + packed = mx.random.normal((10, 4)) # axis 1 packed 2x relative to 8 + parts = _split_uneven(packed, 3, axis=1, sizes=[4, 2, 2]) + self.assertEqual([p.shape[1] for p in parts], [2, 1, 1]) + self.assertTrue(mx.array_equal(mx.concatenate(parts, axis=1), packed)) + + # Incompatible sizes (neither divides the other) raise. + with self.assertRaises(ValueError): + _split_uneven(w, 3, axis=0, sizes=[3, 3, 3]) # sums to 9, not 10 + + # Require one size per rank. + with self.assertRaises(ValueError): + _split_uneven(w, 3, axis=0, sizes=[5, 3, 2, 0]) # 4 sizes, N=3 + with self.assertRaises(ValueError): + _split_uneven(w, 2, axis=0, sizes=[5, 3, 2]) # 3 sizes, N=2 + + def test_shard_linear_uneven(self): + # Inspect every shard without distributed communication. + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + mx.random.seed(0xF0F0F0F0) + + lin = nn.Linear(10, 12, bias=True) + x = mx.random.normal((4, 10)) + y = lin(x) + + # Check an explicit uneven output split. + sizes = [5, 4, 3] + N = len(sizes) + shards = [ + shard_linear(lin, "all-to-sharded", sizes=sizes, group=_FakeGroup(N, r)) + for r in range(N) + ] + starts = [0, 5, 9] + for r, (s, start) in enumerate(zip(shards, starts)): + self.assertEqual(s.weight.shape, (sizes[r], 10)) + self.assertEqual(s.bias.shape, (sizes[r],)) + self.assertTrue( + mx.allclose(s(x), y[:, start : start + sizes[r]], atol=1e-6, rtol=1e-4) + ) + self.assertEqual(s._extra_repr().count("output_dims=12"), 1) + # The shards reconstruct the original weight. + self.assertTrue( + mx.array_equal( + mx.concatenate([s.weight for s in shards], axis=0), lin.weight + ) + ) + + # Check the same split along the input dimension. + lin2 = nn.Linear(12, 10, bias=True) + shards2 = [ + shard_linear(lin2, "sharded-to-all", sizes=sizes, group=_FakeGroup(N, r)) + for r in range(N) + ] + for r, s in enumerate(shards2): + self.assertEqual(s.weight.shape, (10, sizes[r])) + self.assertEqual(s._extra_repr().count("input_dims=12"), 1) + self.assertTrue( + mx.array_equal( + mx.concatenate([s.weight for s in shards2], axis=1), lin2.weight + ) + ) + + # Automatic sizes also distribute a remainder instead of requiring + # the dimension to divide evenly. + auto = nn.Linear(10, 11, bias=True) + auto_shards = [ + shard_linear(auto, "all-to-sharded", group=_FakeGroup(3, r)) + for r in range(3) + ] + self.assertEqual([s.weight.shape[0] for s in auto_shards], [4, 4, 3]) + self.assertTrue( + mx.array_equal( + mx.concatenate([s.weight for s in auto_shards], axis=0), auto.weight + ) + ) + + auto_in = nn.Linear(11, 10, bias=True) + auto_in_shards = [ + shard_linear(auto_in, "sharded-to-all", group=_FakeGroup(3, r)) + for r in range(3) + ] + self.assertEqual([s.weight.shape[1] for s in auto_in_shards], [4, 4, 3]) + self.assertTrue( + mx.array_equal( + mx.concatenate([s.weight for s in auto_in_shards], axis=1), + auto_in.weight, + ) + ) + + # Plain segmented coverage must run on CUDA too. + segmented = nn.Linear(10, 14, bias=True) + segmented_shards = [ + shard_linear( + segmented, "all-to-sharded", segments=[9], group=_FakeGroup(2, r) + ) + for r in range(2) + ] + self.assertEqual([s.weight.shape[0] for s in segmented_shards], [8, 6]) + + # Explicit sizes must sum to the sharded dimension. + with self.assertRaises(ValueError): + shard_linear( + lin, "all-to-sharded", sizes=[5, 4, 2], group=_FakeGroup(3, 0) + ) # sums to 11, not 12 + + # Explicit sizes must contain one entry per rank. + with self.assertRaises(ValueError): + shard_linear(lin, "all-to-sharded", sizes=sizes, group=_FakeGroup(2, 0)) + + # QuantizedMatmul is not supported on CUDA. + if not mx.cuda.is_available(): + # Output splits need not align to quantization groups. + xq = mx.random.normal((4, 64)) + qlin = nn.Linear(64, 96, bias=True).to_quantized(group_size=32, bits=4) + yq = qlin(xq) + + # Default output shards use the same block-aware split as a paired + # quantized sharded-to-all layer. + qout = nn.Linear(64, 320, bias=True).to_quantized(group_size=64, bits=4) + qin = nn.Linear(320, 64, bias=True).to_quantized(group_size=64, bits=4) + qout_shards = [ + shard_linear(qout, "all-to-sharded", group=_FakeGroup(3, r)) + for r in range(3) + ] + qin_shards = [ + shard_linear(qin, "sharded-to-all", group=_FakeGroup(3, r)) + for r in range(3) + ] + self.assertEqual([s.weight.shape[0] for s in qout_shards], [128, 128, 64]) + self.assertEqual([s._total_input_dims for s in qin_shards], [320, 320, 320]) + self.assertEqual( + [s.weight.shape[1] * 32 // s.bits for s in qin_shards], + [128, 128, 64], + ) + q_sizes = [64, 32] # multiples of group_size=32, sums to 96 + for r, start in zip(range(2), [0, 64]): + s = shard_linear( + qlin, "all-to-sharded", sizes=q_sizes, group=_FakeGroup(2, r) + ) + self.assertEqual(s.weight.shape[0], q_sizes[r]) + self.assertTrue( + mx.allclose( + s(xq), yq[:, start : start + q_sizes[r]], atol=1e-3, rtol=1e-2 + ) + ) + # Not a multiple of group_size=32, but still valid: splitting output + # rows doesn't touch a quantization group. + for r, start in zip(range(2), [0, 48]): + s = shard_linear( + qlin, "all-to-sharded", sizes=[48, 48], group=_FakeGroup(2, r) + ) + self.assertEqual(s.weight.shape[0], 48) + self.assertTrue( + mx.allclose(s(xq), yq[:, start : start + 48], atol=1e-3, rtol=1e-2) + ) + + # Input splits preserve groups for every supported bit width. + xq2 = mx.random.normal((4, 96)) + sq_sizes = [64, 32] + for bits in (2, 3, 4, 5, 6, 8): + qlinb = nn.Linear(96, 64, bias=True).to_quantized( + group_size=32, bits=bits + ) + shardsb = [ + shard_linear( + qlinb, "sharded-to-all", sizes=sq_sizes, group=_FakeGroup(2, r) + ) + for r in range(2) + ] + for r in range(2): + self.assertEqual( + shardsb[r].weight.shape, (64, sq_sizes[r] * bits // 32) + ) + self.assertEqual(shardsb[r].scales.shape, (64, sq_sizes[r] // 32)) + self.assertEqual(shardsb[0]._extra_repr().count("input_dims=96"), 1) + # Check each local matmul without the distributed all-sum. + starts = [0, 64] + for r, start in zip(range(2), starts): + partial = mx.quantized_matmul( + xq2[:, start : start + sq_sizes[r]], + shardsb[r].weight, + scales=shardsb[r].scales, + biases=shardsb[r].get("biases"), + transpose=True, + group_size=shardsb[r].group_size, + bits=shardsb[r].bits, + mode=shardsb[r].mode, + ) + ref_x = mx.zeros((4, 96)) + ref_x[:, start : start + sq_sizes[r]] = xq2[ + :, start : start + sq_sizes[r] + ] + self.assertTrue( + mx.allclose( + partial, + qlinb(ref_x) - qlinb.bias, + atol=1e-2, + rtol=5e-2, + ) + ) + if bits == 4: + qlin2 = qlinb # reused by the rejection test below + + # Check grouped boundaries across three fused segments. + qlin3 = nn.Linear(192, 8, bias=True).to_quantized(group_size=32, bits=4) + seg_qshards = [ + shard_linear( + qlin3, "sharded-to-all", segments=3, group=_FakeGroup(2, r) + ) + for r in range(2) + ] + seg_cols = [ + list(range(0, 32)) + list(range(64, 96)) + list(range(128, 160)), + list(range(32, 64)) + list(range(96, 128)) + list(range(160, 192)), + ] + xq3 = mx.random.normal((4, 192)) + total = mx.zeros((4, 8)) + for r in range(2): + s = seg_qshards[r] + self.assertEqual(len(seg_cols[r]), 96) # 32 per segment * 3 segments + self.assertEqual(s.weight.shape, (8, 96 * 4 // 32)) + partial = mx.quantized_matmul( + xq3[:, seg_cols[r]], + s.weight, + scales=s.scales, + biases=s.get("biases"), + transpose=True, + group_size=s.group_size, + bits=s.bits, + mode=s.mode, + ) + total = total + partial + self.assertTrue( + mx.allclose(total + qlin3.bias, qlin3(xq3), atol=1e-2, rtol=5e-2) + ) + + # List boundaries remain valid when every segment can be divided + # at quantization-group boundaries. + qlin4 = nn.Linear(128, 8, bias=True).to_quantized(group_size=32, bits=4) + xq4 = mx.random.normal((4, 128)) + segment_cols = [ + list(range(0, 32)) + list(range(64, 96)), + list(range(32, 64)) + list(range(96, 128)), + ] + for segments in ([0.5], [64]): + list_qshards = [ + shard_linear( + qlin4, + "sharded-to-all", + segments=segments, + group=_FakeGroup(2, r), + ) + for r in range(2) + ] + self.assertEqual( + [s.weight.shape[1] * 32 // s.bits for s in list_qshards], + [64, 64], + ) + self.assertEqual([s.scales.shape[1] for s in list_qshards], [2, 2]) + total = mx.zeros((4, 8)) + for shard, cols in zip(list_qshards, segment_cols): + total += mx.quantized_matmul( + xq4[:, cols], + shard.weight, + scales=shard.scales, + biases=shard.get("biases"), + transpose=True, + group_size=shard.group_size, + bits=shard.bits, + mode=shard.mode, + ) + self.assertTrue( + mx.allclose(total + qlin4.bias, qlin4(xq4), atol=1e-2, rtol=5e-2) + ) + with self.assertRaises(ValueError): + # 48 is not a multiple of group_size=32 + shard_linear( + qlin2, "sharded-to-all", sizes=[48, 48], group=_FakeGroup(2, 0) + ) + + # Divide each explicit size evenly among segments. + fused = nn.Linear(10, 12, bias=True) + with self.assertRaises(ValueError): + shard_linear( + fused, + "all-to-sharded", + segments=3, + sizes=[5, 4, 3], + group=_FakeGroup(3, 0), + ) + + # Split each fused segment independently. + seg_shards = [ + shard_linear( + fused, "all-to-sharded", segments=3, group=_FakeGroup(2, r) + ) + for r in range(2) + ] + expected_rank0 = mx.concatenate( + [fused.weight[0:2], fused.weight[4:6], fused.weight[8:10]], axis=0 + ) + expected_rank1 = mx.concatenate( + [fused.weight[2:4], fused.weight[6:8], fused.weight[10:12]], axis=0 + ) + self.assertEqual(seg_shards[0].weight.shape, (6, 10)) + self.assertEqual(seg_shards[1].weight.shape, (6, 10)) + self.assertTrue(mx.array_equal(seg_shards[0].weight, expected_rank0)) + self.assertTrue(mx.array_equal(seg_shards[1].weight, expected_rank1)) + + # Apply explicit sizes to each fused segment. + seg_shards2 = [ + shard_linear( + fused, + "all-to-sharded", + segments=3, + sizes=[9, 3], + group=_FakeGroup(2, r), + ) + for r in range(2) + ] + expected_rank0_2 = mx.concatenate( + [fused.weight[0:3], fused.weight[4:7], fused.weight[8:11]], axis=0 + ) + expected_rank1_2 = mx.concatenate( + [fused.weight[3:4], fused.weight[7:8], fused.weight[11:12]], axis=0 + ) + self.assertEqual(seg_shards2[0].weight.shape, (9, 10)) + self.assertEqual(seg_shards2[1].weight.shape, (3, 10)) + self.assertTrue(mx.array_equal(seg_shards2[0].weight, expected_rank0_2)) + self.assertTrue(mx.array_equal(seg_shards2[1].weight, expected_rank1_2)) + + # Unequal segments do not support explicit sizes. + with self.assertRaises(ValueError): + shard_linear( + fused, + "all-to-sharded", + segments=[6], + sizes=[5, 4, 3], + group=_FakeGroup(3, 0), + ) + + # Split unequal plain segments independently. + asym = nn.Linear(10, 14, bias=True) + seg_list_shards = [ + shard_linear( + asym, "all-to-sharded", segments=[9], group=_FakeGroup(2, r) + ) + for r in range(2) + ] + expected_list_rank0 = mx.concatenate( + [asym.weight[0:5], asym.weight[9:12]], axis=0 + ) + expected_list_rank1 = mx.concatenate( + [asym.weight[5:9], asym.weight[12:14]], axis=0 + ) + self.assertEqual(seg_list_shards[0].weight.shape, (8, 10)) + self.assertEqual(seg_list_shards[1].weight.shape, (6, 10)) + self.assertTrue( + mx.array_equal(seg_list_shards[0].weight, expected_list_rank0) + ) + self.assertTrue( + mx.array_equal(seg_list_shards[1].weight, expected_list_rank1) + ) + + # Quantized output segments also split independently. + seg_list_qshards = [ + shard_linear( + qlin, "all-to-sharded", segments=[48], group=_FakeGroup(2, r) + ) + for r in range(2) + ] + self.assertEqual(seg_list_qshards[0].weight.shape[0], 48) + self.assertEqual(seg_list_qshards[1].weight.shape[0], 48) + self.assertTrue( + mx.allclose( + seg_list_qshards[0](xq), + mx.concatenate([yq[:, 0:24], yq[:, 48:72]], axis=1), + atol=1e-3, + rtol=1e-2, + ) + ) + self.assertTrue( + mx.allclose( + seg_list_qshards[1](xq), + mx.concatenate([yq[:, 24:48], yq[:, 72:96]], axis=1), + atol=1e-3, + rtol=1e-2, + ) + ) + + # A list-valued split is rejected when a segment cannot be divided + # at quantization-group boundaries. + with self.assertRaises(ValueError): + shard_linear( + qlin2, "sharded-to-all", segments=[48], group=_FakeGroup(2, 0) + ) + + # Reject zero-width quantized shards on every rank. + tiny_qlin = nn.Linear(192, 64, bias=True).to_quantized( + group_size=64, bits=4 + ) + for r in (0, 1): + with self.assertRaises(ValueError): + shard_linear( + tiny_qlin, "sharded-to-all", segments=3, group=_FakeGroup(2, r) + ) + + def test_shard_inplace_custom_parameters(self): + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + class CustomModule(nn.Module): + def __init__(self): + super().__init__() + self.weight = mx.zeros((10,)) + self.scales = mx.zeros((6,)) + + # Names used by quantized layers do not couple unrelated arrays in a + # generic module. + expected = [(4, 2), (3, 2), (3, 2)] + for r, shapes in enumerate(expected): + module = CustomModule() + shard_inplace(module, lambda path, value: 0, group=_FakeGroup(3, r)) + self.assertEqual((module.weight.size, module.scales.size), shapes) + + module = CustomModule() + with self.assertRaises(ValueError): + shard_inplace(module, lambda path, value: 1, group=_FakeGroup(2, 0)) + + def test_shard_inplace_third_party_quantized_module(self): + # Detect third-party quantized modules by their parameters. + if not mx.cuda.is_available(): + + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + class ThirdPartyQuantized(nn.Module): + def __init__(self, input_dims, output_dims, group_size, bits): + super().__init__() + self.group_size = group_size + self.bits = bits + lin = nn.Linear(input_dims, output_dims, bias=True).to_quantized( + group_size=group_size, bits=bits + ) + self.weight = lin.weight + self.scales = lin.scales + if lin.get("biases") is not None: + self.biases = lin.biases + self.bias = lin.bias + + # Check an uneven grouped input split. + ref = ThirdPartyQuantized(96, 64, group_size=32, bits=4) + x = mx.random.normal((4, 96)) + y_full = ( + mx.quantized_matmul( + x, + ref.weight, + scales=ref.scales, + biases=ref.get("biases"), + transpose=True, + group_size=32, + bits=4, + ) + + ref.bias + ) + total = None + cols = [range(0, 64), range(64, 96)] + for r in (0, 1): + m = ThirdPartyQuantized(96, 64, group_size=32, bits=4) + m.update(ref.parameters()) + shard_inplace(m, "sharded-to-all", group=_FakeGroup(2, r)) + self.assertEqual(m.weight.shape[1] * 32 // 4, [64, 32][r]) + partial = mx.quantized_matmul( + x[:, list(cols[r])], + m.weight, + scales=m.scales, + biases=m.get("biases"), + transpose=True, + group_size=32, + bits=4, + ) + total = partial if total is None else total + partial + self.assertTrue(mx.allclose(total + ref.bias, y_full, atol=1e-2, rtol=5e-2)) + + def test_shard_inplace_predicate_called_once(self): + if mx.cuda.is_available(): + return + + class _FakeGroup: + def size(self): + return 2 + + def rank(self): + return 0 + + module = nn.Linear(128, 64).to_quantized(group_size=32, bits=4) + calls = {} + + def sharding(path, value): + calls[path] = calls.get(path, 0) + 1 + return None if path.endswith("bias") else -1 + + shard_inplace(module, sharding, group=_FakeGroup()) + self.assertTrue(calls) + self.assertTrue(all(count == 1 for count in calls.values())) + + def test_shard_inplace_quantized_consistency(self): + # QuantizedMatmul is not supported on CUDA. + if not mx.cuda.is_available(): + # Keep packed weights and grouped metadata aligned. + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + # Check an uneven split of 96 inputs into 64 and 32. + qlin = nn.Linear(96, 64, bias=True).to_quantized(group_size=32, bits=4) + x = mx.random.normal((4, 96)) + y_full = qlin(x) + total = None + cols = [range(0, 64), range(64, 96)] + for r in (0, 1): + m = nn.Linear(96, 64, bias=True).to_quantized(group_size=32, bits=4) + m.update(qlin.parameters()) + shard_inplace(m, "sharded-to-all", group=_FakeGroup(2, r)) + partial = mx.quantized_matmul( + x[:, list(cols[r])], + m.weight, + scales=m.scales, + biases=m.get("biases"), + transpose=True, + group_size=32, + bits=4, + ) + total = partial if total is None else total + partial + self.assertTrue( + mx.allclose(total + qlin.bias, y_full, atol=1e-2, rtol=5e-2) + ) + + # Handle quantized parameters in a nested module. + class Wrap(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(96, 64, bias=True).to_quantized( + group_size=32, bits=4 + ) + + w = Wrap() + for r in (0, 1): + wm = Wrap() + wm.update(w.parameters()) + shard_inplace(wm, "sharded-to-all", group=_FakeGroup(2, r)) + self.assertEqual(wm.linear.weight.shape[1], [8, 4][r]) + + # Preserve the evenly divisible case. + qlin2 = nn.Linear(128, 64, bias=True).to_quantized(group_size=32, bits=4) + x2 = mx.random.normal((4, 128)) + y_full2 = qlin2(x2) + total2 = None + for r in (0, 1): + m2 = nn.Linear(128, 64, bias=True).to_quantized(group_size=32, bits=4) + m2.update(qlin2.parameters()) + shard_inplace(m2, "sharded-to-all", group=_FakeGroup(2, r)) + xr = x2[:, r * 64 : (r + 1) * 64] + partial = mx.quantized_matmul( + xr, + m2.weight, + scales=m2.scales, + biases=m2.get("biases"), + transpose=True, + group_size=32, + bits=4, + ) + total2 = partial if total2 is None else total2 + partial + self.assertTrue( + mx.allclose(total2 + qlin2.bias, y_full2, atol=1e-2, rtol=5e-2) + ) + + # Metadata cannot be sharded without its weight. + qlin3 = nn.Linear(96, 64).to_quantized(group_size=32, bits=4) + + def skip_weight(path, value): + return None if path.endswith("weight") else 0 + + m3 = nn.Linear(96, 64).to_quantized(group_size=32, bits=4) + m3.update(qlin3.parameters()) + with self.assertRaises(ValueError): + shard_inplace(m3, skip_weight, group=_FakeGroup(2, 0)) + + def test_shard_inplace_quantized_list_segments(self): + # QuantizedMatmul is not supported on CUDA. + if not mx.cuda.is_available(): + + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + # Reject segment boundaries that differ after packing and grouping. + qlin = nn.Linear(1024, 8, bias=True).to_quantized(group_size=64, bits=4) + for r in (0, 1): + m = nn.Linear(1024, 8, bias=True).to_quantized(group_size=64, bits=4) + m.update(qlin.parameters()) + with self.assertRaises(ValueError): + shard_inplace( + m, "sharded-to-all", segments=[0.3], group=_FakeGroup(2, r) + ) + + # Accept boundaries that align in every representation. + qlin2 = nn.Linear(256, 8, bias=True).to_quantized(group_size=32, bits=4) + x = mx.random.normal((4, 256)) + y_full = qlin2(x) + # Each rank gets 64 columns from both segments. + cols = [ + list(range(0, 64)) + list(range(128, 192)), + list(range(64, 128)) + list(range(192, 256)), + ] + total = None + for r in (0, 1): + m2 = nn.Linear(256, 8, bias=True).to_quantized(group_size=32, bits=4) + m2.update(qlin2.parameters()) + shard_inplace( + m2, "sharded-to-all", segments=[0.5], group=_FakeGroup(2, r) + ) + partial = mx.quantized_matmul( + x[:, cols[r]], + m2.weight, + scales=m2.scales, + biases=m2.get("biases"), + transpose=True, + group_size=32, + bits=4, + ) + total = partial if total is None else total + partial + self.assertTrue( + mx.allclose(total + qlin2.bias, y_full, atol=1e-2, rtol=5e-2) + ) + + def test_shard_linear_quantized_fused_segments_output(self): + # Fall back to plain output splits when groups cannot cover all ranks. + output_dims, n_segments, N, group_size = 256, 2, 4, 64 + qlin = nn.Linear(64, output_dims, bias=True).to_quantized( + group_size=group_size, bits=4 + ) + + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + shards = [ + shard_linear( + qlin, "all-to-sharded", segments=n_segments, group=_FakeGroup(N, r) + ) + for r in range(N) + ] + self.assertEqual([s.weight.shape[0] for s in shards], [64, 64, 64, 64]) + + def test_shard_inplace_quantized_zero_share_rejected(self): + # shard_inplace must also reject zero-width quantized shards. + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + # 2 output rows across 4 ranks: ranks 2 and 3 get a zero-width share. + qlin = nn.Linear(128, 2, bias=True).to_quantized(group_size=32, bits=4) + for r in range(4): + m = nn.Linear(128, 2, bias=True).to_quantized(group_size=32, bits=4) + m.update(qlin.parameters()) + with self.assertRaises(ValueError): + shard_inplace(m, "all-to-sharded", group=_FakeGroup(4, r)) + + def test_shard_inplace_matches_shard_linear_quantized(self): + # Both APIs must resolve the same quantized shard sizes. + if not mx.cuda.is_available(): + + class _FakeGroup: + def __init__(self, n, r): + self._n, self._r = n, r + + def size(self): + return self._n + + def rank(self): + return self._r + + def _shard_linear_sizes(module, sharding, N): + try: + return [ + ( + shard_linear( + module, sharding, group=_FakeGroup(N, r) + ).weight.shape[0 if sharding == "all-to-sharded" else 1] + ) + for r in range(N) + ], None + except ValueError as e: + return None, str(e) + + def _shard_inplace_sizes(reference, module_factory, sharding, N): + sizes = [] + for r in range(N): + m = module_factory() + m.update(reference.parameters()) + try: + shard_inplace(m, sharding, group=_FakeGroup(N, r)) + except ValueError as e: + return None, str(e) + sizes.append( + m.weight.shape[0 if sharding == "all-to-sharded" else 1] + ) + return sizes, None + + for group_size in (32, 64): + for dim_mult in (3, 5, 12, 17): + dim = group_size * dim_mult + for N in (2, 3, 5, 8): + qlin_out = nn.Linear(group_size, dim, bias=True).to_quantized( + group_size=group_size, bits=4 + ) + qlin_in = nn.Linear(dim, group_size, bias=True).to_quantized( + group_size=group_size, bits=4 + ) + + sl_sizes, sl_err = _shard_linear_sizes( + qlin_out, "all-to-sharded", N + ) + si_sizes, si_err = _shard_inplace_sizes( + qlin_out, + lambda: nn.Linear(group_size, dim, bias=True).to_quantized( + group_size=group_size, bits=4 + ), + "all-to-sharded", + N, + ) + # The APIs must agree on acceptance and shard sizes. + self.assertEqual( + sl_err is None, + si_err is None, + f"all-to-sharded raise-disagreement: dim={dim}, " + f"group_size={group_size}, N={N}: " + f"shard_linear={sl_err!r} shard_inplace={si_err!r}", + ) + if sl_err is None: + self.assertEqual( + sl_sizes, + si_sizes, + f"all-to-sharded disagreement: dim={dim}, " + f"group_size={group_size}, N={N}", + ) + + sl_in_sizes, sl_in_err = _shard_linear_sizes( + qlin_in, "sharded-to-all", N + ) + si_in_sizes, si_in_err = _shard_inplace_sizes( + qlin_in, + lambda: nn.Linear(dim, group_size, bias=True).to_quantized( + group_size=group_size, bits=4 + ), + "sharded-to-all", + N, + ) + self.assertEqual( + sl_in_err is None, + si_in_err is None, + f"sharded-to-all raise-disagreement: dim={dim}, " + f"group_size={group_size}, N={N}: " + f"shard_linear={sl_in_err!r} shard_inplace={si_in_err!r}", + ) + if sl_in_err is None: + self.assertEqual( + sl_in_sizes, + si_in_sizes, + f"sharded-to-all disagreement: dim={dim}, " + f"group_size={group_size}, N={N}", + ) + def test_shard_predicate(self): mx.random.seed(0xF0F0F0F0)