[WIP] [CUDA] fsdp

docs/src/python/nn.rst

@@ -175,7 +175,6 @@ In detail:
    value_and_grad
    quantize
    average_gradients
-   fsdp_apply_gradients
 
 .. toctree::
 

docs/src/python/nn/distributed.rst

@@ -15,6 +15,7 @@ create sharded layers from existing :class:`Modules <mlx.nn.Module>`.
 
    shard_linear
    shard_inplace
+   fully_shard
 
 Layers
 ^^^^^^
@@ -28,3 +29,4 @@ Layers
    ShardedToAllLinear
    QuantizedAllToShardedLinear
    QuantizedShardedToAllLinear
+   FullyShardedModule

python/mlx/nn/__init__.py

@@ -4,6 +4,5 @@
 from mlx.nn.layers import *
 from mlx.nn.utils import (
     average_gradients,
-    fsdp_apply_gradients,
     value_and_grad,
 )

python/mlx/nn/layers/__init__.py

@@ -64,9 +64,11 @@
 )
 from mlx.nn.layers.distributed import (
     AllToShardedLinear,
+    FullyShardedModule,
     QuantizedAllToShardedLinear,
     QuantizedShardedToAllLinear,
     ShardedToAllLinear,
+    fully_shard,
 )
 from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d
 from mlx.nn.layers.embedding import Embedding

python/mlx/nn/layers/distributed.py

@@ -1,14 +1,14 @@
 # Copyright © 2024 Apple Inc.
 
 import math
-from functools import lru_cache
+from functools import lru_cache, reduce
 from typing import Callable, Optional, Union
 
 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_map_with_path
+from mlx.utils import tree_flatten, tree_map_with_path, tree_unflatten
 
 
 @lru_cache
@@ -617,3 +617,147 @@ def from_quantized_linear(
         )
 
         return sl
+
+
+def _make_gather_fn(group, full_shapes, shard_sizes, compute_dtype):
+    N = group.size()
+    indices = reduce(lambda acc, w: acc + [acc[-1] + w], shard_sizes, [0])
+    split_indices = indices[1:-1]
+    shard_shapes = [(shape[0] // N,) + tuple(shape[1:]) for shape in full_shapes]
+
+    def _maybe_cast(x, dtype):
+        if dtype is None or x.dtype == dtype:
+            return x
+        return x.astype(dtype)
+
+    @mx.custom_function
+    def gather(shards):
+        shard = mx.concatenate(
+            [_maybe_cast(s.reshape(1, -1), compute_dtype) for s in shards], axis=1
+        )
+        full = mx.distributed.all_gather(shard, group=group)
+        parts = mx.split(full, split_indices, axis=1)
+        return [p.reshape(shape) for p, shape in zip(parts, full_shapes)]
+
+    @gather.vjp
+    def gather_vjp(shards, cotangents, _):
+        local_full = mx.concatenate([c.reshape(N, -1) for c in cotangents], axis=1)
+        local_shard = mx.distributed.sum_scatter(local_full, group=group) / N
+        parts = mx.split(local_shard, split_indices, axis=1)
+        return [
+            _maybe_cast(p.reshape(shape), s.dtype)
+            for p, shape, s in zip(parts, shard_shapes, shards)
+        ]
+
+    return gather
+
+
+def _maybe_shard(m, k, v):
+    if isinstance(v, FullyShardedModule):
+        return False
+    return Module.valid_parameter_filter(m, k, v)
+
+
+class FullyShardedModule(Module):
+    """Wrap a module so each member of the group holds only a shard of its
+    parameters.
+
+    The full parameters are gathered for the forward pass and the gradients
+    are reduce-scattered in the backward pass, so during training
+    each member of the group stores and updates only its own shard.
+
+    Every parameter is sharded along axis 0, so each parameter's size along
+    that axis must be divisible by the size of ``group``.
+
+    Use :func:`fully_shard` to wrap a module.
+
+    Args:
+        module (mlx.nn.Module): The module whose parameters will be sharded.
+        group (mlx.core.distributed.Group, optional): The group to shard
+            across. If not set, the global group is used. Default: ``None``.
+        compute_dtype (mlx.core.Dtype, optional): If set, the gathered
+            parameters are cast to this dtype for the forward pass.
+            Default: ``None``.
+    """
+
+    def __init__(
+        self,
+        module: Module,
+        group: Optional[mx.distributed.Group] = None,
+        compute_dtype: Optional[mx.Dtype] = None,
+    ):
+        super().__init__()
+        group = group or mx.distributed.init()
+        N = group.size()
+
+        shard_params = module.filter_and_map(_maybe_shard)
+        flat = tree_flatten(shard_params)
+        for path, a in flat:
+            if a.ndim == 0:
+                raise ValueError(
+                    f"Cannot shard parameter '{path}' because it is a scalar."
+                )
+            if a.shape[0] % N != 0:
+                raise ValueError(
+                    f"Cannot shard parameter '{path}' with shape {a.shape} "
+                    f"across {N} devices: axis 0 must be divisible by {N}."
+                )
+
+        super(Module, self).__setattr__("_paths", [k for k, _ in flat])
+        full_shapes = [a.shape for _, a in flat]
+        shard_sizes = [a.size // N for _, a in flat]
+
+        module.update(_shard(shard_params, lambda p, w: 0, group))
+
+        self.module = module
+        self._gather_fn = _make_gather_fn(
+            group, full_shapes, shard_sizes, compute_dtype
+        )
+
+    def _extra_repr(self) -> str:
+        return f"num_sharded_params={len(self._paths)}"
+
+    def _gathered_call(self, fn, *args, **kwargs):
+        shard_tree = self.module.filter_and_map(_maybe_shard)
+        shards = [a for _, a in tree_flatten(shard_tree)]
+        fulls = self._gather_fn(shards)
+        self.module.update(tree_unflatten(list(zip(self._paths, fulls))))
+        try:
+            return fn(*args, **kwargs)
+        finally:
+            self.module.update(shard_tree)
+
+    def __call__(self, *args, **kwargs):
+        return self._gathered_call(self.module, *args, **kwargs)
+
+    def as_linear(self, *args, **kwargs):
+        return self._gathered_call(self.module.as_linear, *args, **kwargs)
+
+
+def fully_shard(
+    module: Module,
+    *,
+    group: Optional[mx.distributed.Group] = None,
+    compute_dtype: Optional[mx.Dtype] = None,
+) -> Module:
+    """Wrap ``module`` in a :class:`FullyShardedModule`.
+
+    Args:
+        module (mlx.nn.Module): The module to wrap.
+        group (mlx.core.distributed.Group, optional): The group to shard
+            across. If not set, the global group is used. Default: ``None``.
+        compute_dtype (mlx.core.Dtype, optional): If set, the gathered
+            parameters are cast to this dtype for the forward pass.
+            Default: ``None``.
+
+    Returns:
+        The wrapped :class:`FullyShardedModule`, or ``module`` unchanged.
+    """
+    group = group or mx.distributed.init()
+    if group.size() == 1:
+        return module
+    if isinstance(module, FullyShardedModule):
+        return module
+
+    wrapped = FullyShardedModule(module, group=group, compute_dtype=compute_dtype)
+    return wrapped if wrapped._paths else module

python/mlx/nn/utils.py

@@ -173,149 +173,36 @@ def average_gradients(
         return tree_unflatten(new_flat_grads)
 
 
-def _clip_grads_fsdp(grads_slice, max_norm, group=None):
-    local_norm_sq = tree_reduce(lambda acc, g: acc + g.square().sum(), grads_slice, 0.0)
-    global_norm_sq = mx.distributed.all_sum(local_norm_sq, group=group)
-    grad_norm = mx.sqrt(global_norm_sq)
-    normalizer = mx.minimum(max_norm / (grad_norm + 1e-6), 1.0)
-    grads_slice = tree_map(lambda g: g * normalizer, grads_slice)
-
-    return grads_slice, grad_norm
-
-
-def fsdp_apply_gradients(
-    gradients,
-    parameters,
-    optimizer,
-    fsdp_group=None,
-    dp_group=None,
-    communication_size=32 * 1024**2,
-    communication_stream=None,
-    max_norm=None,
+def clip_grad_norm_sharded(
+    gradients: Any,
+    max_norm: float,
+    group: Optional[mx.distributed.Group] = None,
 ):
-    """Perform a distributed optimizer step by sharding gradients and optimizer states across ranks.
-
-    This helper function performs the following steps:
-    1. Reduce-scatter the gradients across ranks so each rank gets a shard of the averaged gradients.
-    2. Optionally clip the sharded gradients by global norm.
-    3. Apply the optimizer update on the local parameter slice using the sharded gradients.
-    4. All-gather the updated parameter slices from all ranks to reconstruct the full parameters tree.
+    """Clip the global norm of gradients that are sharded across a group.
 
-    This is similar to PyTorch's FSDP with `reshard_after_forward=False`.
+    This is the sharded equivalent of
+    :func:`mlx.optimizers.clip_grad_norm`. Each member of the group holds only
+    a shard of the gradients, so the global norm is computed by summing the
+    local squared norms across the group before rescaling. It is useful for
+    clipping the gradients of a module wrapped with :func:`mlx.nn.fully_shard`.
 
     Args:
-        gradients (Any): The Python tree containing the full gradients (it should
-            have the same structure as ``parameters``). Each gradient's first
-            dimension must be divisible by ``fsdp_group.size()``.
-        parameters (Any): The Python tree containing the full parameters (it should
-            have the same structure across processes). Each parameter's first
-            dimension must be divisible by ``fsdp_group.size()``.
-        optimizer: Optimizer with an ``apply_gradients`` method.
-        fsdp_group (Optional[mlx.core.distributed.Group]): The group of processes
-            for FSDP sharding. If ``None``, the global group is used.
-        dp_group (Optional[mlx.core.distributed.Group]): The group of processes
-            for data-parallel gradient averaging. Required when ``fsdp_group`` is
-            smaller than the world (e.g. FSDP intra-node, DDP inter-node).
-            Default: ``None``.
-        communication_size (int): Group arrays until their size in bytes exceeds
-            this number. Perform one communication step per group of arrays. If
-            less or equal to 0 array grouping is disabled. Default: ``32MiB``.
-        communication_stream (Optional[mlx.core.Stream]): The stream to use
-            for the communication. If unspecified the default communication
-            stream is used which can vary by back-end. Default: ``None``.
-        max_norm (Optional[float]): If provided, clip gradients to this
-            maximum global norm before applying the optimizer update.
-            Default: ``None``.
+        gradients (Any): A Python tree containing the local shard of the
+            gradient arrays.
+        max_norm (float): The maximum allowed global norm of the gradients.
+        group (Optional[mlx.core.distributed.Group]): The group across which
+            the gradients are sharded. If set to ``None`` the global group is
+            used. Default: ``None``.
 
     Returns:
-        If ``max_norm`` is ``None``, returns the updated full-parameter tree.
-        Otherwise returns ``(parameters, grad_norm)``, where ``grad_norm`` is
-        the global gradient norm before clipping.
-
-    Example:
-
-        >>> optimizer = optim.SGD(learning_rate=0.01)
-        >>> # Without gradient clipping
-        >>> updated_params = fsdp_apply_gradients(grads, params, optimizer)
-        >>> model.update(updated_params)
-        >>>
-        >>> # With gradient clipping
-        >>> updated_params, grad_norm = fsdp_apply_gradients(
-        ...     grads, params, optimizer, max_norm=1.0
-        ... )
-        >>> model.update(updated_params)
+        (Any, mlx.core.array): The possibly rescaled local shard of the
+        gradients and the global gradient norm.
     """
-    fsdp_group = fsdp_group or mx.distributed.init()
-    N = fsdp_group.size() * (dp_group.size() if dp_group is not None else 1)
-
-    if N == 1:
-        if max_norm is not None:
-            gradients, grad_norm = _clip_grads_fsdp(gradients, max_norm)
-            return optimizer.apply_gradients(gradients, parameters), grad_norm
-        return optimizer.apply_gradients(gradients, parameters)
-
-    flat_grads = tree_flatten(gradients)
-    flat_params = tree_flatten(parameters)
-
-    keys, shapes, sizes, dtypes = _extract_info(flat_grads)
-    itemsize = dtypes[0].size
-
-    groups = _group_by_size(keys, sizes, itemsize, communication_size)
-
-    S = fsdp_group.size()
-    fsdp_rank = fsdp_group.rank()
-    # reduce-scatter gradients, shard parameters
-    grad_slices = {}
-    param_slices = {}
-    for group_idx, arr_group in enumerate(groups):
-        big_grad = mx.concatenate(
-            [flat_grads[i][1].reshape(S, -1) for i in arr_group], axis=1
-        )
-        grad_slices[group_idx] = (
-            mx.distributed.sum_scatter(
-                big_grad, group=fsdp_group, stream=communication_stream
-            )
-            / N
-        )
-        if dp_group is not None:
-            grad_slices[group_idx] = mx.distributed.all_sum(
-                grad_slices[group_idx], group=dp_group, stream=communication_stream
-            )
-        big_param = mx.concatenate(
-            [flat_params[i][1].reshape(S, -1) for i in arr_group], axis=1
-        )
-        param_slices[group_idx] = big_param[fsdp_rank]
-
-    # clip gradients if needed
-    grad_norm = None
-    if max_norm is not None:
-        grad_slices, grad_norm = _clip_grads_fsdp(
-            grad_slices, max_norm, group=fsdp_group
-        )
-
-    # optimizer step
-    updated_param_slices = optimizer.apply_gradients(grad_slices, param_slices)
-
-    # all-gather and reconstruct
-    new_flat = []
-    for group_idx, arr_group in enumerate(groups):
-        big_gathered = mx.distributed.all_gather(
-            updated_param_slices[group_idx],
-            group=fsdp_group,
-            stream=communication_stream,
-        )
-        split_sizes = [sizes[i] // S for i in arr_group]
-        split_indices = []
-        acc = 0
-        for s in split_sizes:
-            acc += s
-            split_indices.append(acc)
-
-        parts = mx.split(big_gathered, split_indices[:-1], axis=1)
-        for idx_in_group, i in enumerate(arr_group):
-            new_flat.append((keys[i], parts[idx_in_group].reshape(shapes[i])))
-
-    result = tree_unflatten(new_flat)
-    if max_norm is not None:
-        return result, grad_norm
-    return result
+    local_norm_squared = tree_reduce(
+        lambda acc, g: acc + g.square().sum(), gradients, 0.0
+    )
+    global_norm_squared = mx.distributed.all_sum(local_norm_squared, group=group)
+    grad_norm = mx.sqrt(global_norm_squared)
+    normalizer = mx.minimum(max_norm / (grad_norm + 1e-6), 1.0)
+    clipped_gradients = tree_map(lambda g: g * normalizer, gradients)
+    return clipped_gradients, grad_norm