From a9e9e5aa389f54dd6e702835cdc77cc8be1481b7 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 25 Jun 2026 14:54:33 +0200 Subject: [PATCH 1/2] fsdp --- python/mlx/nn/utils.py | 217 +++++++++++++++++------------------------ 1 file changed, 88 insertions(+), 129 deletions(-) diff --git a/python/mlx/nn/utils.py b/python/mlx/nn/utils.py index b53e9efe21..586bee8f88 100644 --- a/python/mlx/nn/utils.py +++ b/python/mlx/nn/utils.py @@ -7,6 +7,7 @@ from ..utils import tree_flatten, tree_map, tree_reduce, tree_unflatten from .layers.base import Module +from .layers.distributed import _shard def value_and_grad(model: Module, fn: Callable): @@ -173,7 +174,7 @@ def average_gradients( return tree_unflatten(new_flat_grads) -def _clip_grads_fsdp(grads_slice, max_norm, group=None): +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) @@ -183,139 +184,97 @@ def _clip_grads_fsdp(grads_slice, max_norm, group=None): 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, -): - """Perform a distributed optimizer step by sharding gradients and optimizer states across ranks. +def _make_gather_fn(group, full_shapes, shard_sizes, cast_dtype): + S = group.size() + indices = reduce(lambda acc, w: acc + [acc[-1] + w], shard_sizes, [0]) + split_indices = indices[1:-1] + shard_shapes = [(shape[0] // S,) + tuple(shape[1:]) for shape in full_shapes] - 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. + def _maybe_cast(x, dtype): + if dtype is None or x.dtype == dtype: + return x + return x.astype(dtype) - This is similar to PyTorch's FSDP with `reshard_after_forward=False`. + @mx.custom_function + def gather(shards): + big_shard = mx.concatenate( + [_maybe_cast(s.reshape(1, -1), cast_dtype) for s in shards], axis=1 + ) + big_full = mx.distributed.all_gather(big_shard, group=group) + parts = mx.split(big_full, split_indices, axis=1) + return [p.reshape(shape) for p, shape in zip(parts, full_shapes)] + + @gather.vjp + def gather_vjp(shards, cotangents, _): + big_cot_full = mx.concatenate([c.reshape(S, -1) for c in cotangents], axis=1) + big_cot_shard = mx.distributed.sum_scatter(big_cot_full, group=group) / S + parts = mx.split(big_cot_shard, split_indices, axis=1) + return [p.reshape(shape) for p, shape in zip(parts, shard_shapes)] + + 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): + def __init__(self, module, group, cast_dtype): + 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"FSDP: parameter {path} is a 0-D scalar and cannot be sharded." + ) + if a.shape[0] % N != 0: + raise ValueError( + f"FSDP: parameter {path} has shape {a.shape}; axis 0 must " + f"be divisible by the FSDP group size {N}." + ) - 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``. + self._paths = [k for k, _ in flat] + full_shapes = [a.shape for _, a in flat] + shard_sizes = [a.size // N for _, a in flat] - 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) - """ - fsdp_group = fsdp_group or mx.distributed.init() - N = fsdp_group.size() * (dp_group.size() if dp_group is not None else 1) + module.update(_shard(shard_params, lambda p, w: 0, group)) - 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] + self.module = module + self._gather_fn = _make_gather_fn(group, full_shapes, shard_sizes, cast_dtype) - # 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 - ) + 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) - # optimizer step - updated_param_slices = optimizer.apply_gradients(grad_slices, param_slices) + def __call__(self, *args, **kwargs): + return self._gathered_call(self.module, *args, **kwargs) - # 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 + 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, + cast_dtype: Optional[mx.Dtype] = None, +) -> Module: + group = group or mx.distributed.init() + if group.size() == 1: + return module + if isinstance(module, FullyShardedModule): + return module + + wrapped = FullyShardedModule(module, group, cast_dtype) + return wrapped if wrapped._paths else module From 63ca041bccf8bd79bd8c59fa5e573f3591badc04 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Sun, 12 Jul 2026 19:54:14 +0200 Subject: [PATCH 2/2] refactoring and better docs --- docs/src/python/nn.rst | 1 - docs/src/python/nn/distributed.rst | 2 + python/mlx/nn/__init__.py | 1 - python/mlx/nn/layers/__init__.py | 2 + python/mlx/nn/layers/distributed.py | 148 +++++++++++++- python/mlx/nn/utils.py | 132 +++--------- python/tests/mlx_distributed_tests.py | 40 +++- python/tests/nccl_test_distributed.py | 279 ++++++-------------------- 8 files changed, 277 insertions(+), 328 deletions(-) diff --git a/docs/src/python/nn.rst b/docs/src/python/nn.rst index 0ea3a56fbb..00f7f95456 100644 --- a/docs/src/python/nn.rst +++ b/docs/src/python/nn.rst @@ -175,7 +175,6 @@ In detail: value_and_grad quantize average_gradients - fsdp_apply_gradients .. toctree:: diff --git a/docs/src/python/nn/distributed.rst b/docs/src/python/nn/distributed.rst index 07dd8e2308..23998b71a3 100644 --- a/docs/src/python/nn/distributed.rst +++ b/docs/src/python/nn/distributed.rst @@ -15,6 +15,7 @@ create sharded layers from existing :class:`Modules `. shard_linear shard_inplace + fully_shard Layers ^^^^^^ @@ -28,3 +29,4 @@ Layers ShardedToAllLinear QuantizedAllToShardedLinear QuantizedShardedToAllLinear + FullyShardedModule diff --git a/python/mlx/nn/__init__.py b/python/mlx/nn/__init__.py index 6c3bd0aa24..5b51f8a01d 100644 --- a/python/mlx/nn/__init__.py +++ b/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, ) diff --git a/python/mlx/nn/layers/__init__.py b/python/mlx/nn/layers/__init__.py index c2fba58347..25b96c0a7f 100644 --- a/python/mlx/nn/layers/__init__.py +++ b/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 diff --git a/python/mlx/nn/layers/distributed.py b/python/mlx/nn/layers/distributed.py index c4799bd073..69804756c9 100644 --- a/python/mlx/nn/layers/distributed.py +++ b/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 diff --git a/python/mlx/nn/utils.py b/python/mlx/nn/utils.py index 586bee8f88..a60468db64 100644 --- a/python/mlx/nn/utils.py +++ b/python/mlx/nn/utils.py @@ -7,7 +7,6 @@ from ..utils import tree_flatten, tree_map, tree_reduce, tree_unflatten from .layers.base import Module -from .layers.distributed import _shard def value_and_grad(model: Module, fn: Callable): @@ -174,107 +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 _make_gather_fn(group, full_shapes, shard_sizes, cast_dtype): - S = group.size() - indices = reduce(lambda acc, w: acc + [acc[-1] + w], shard_sizes, [0]) - split_indices = indices[1:-1] - shard_shapes = [(shape[0] // S,) + 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): - big_shard = mx.concatenate( - [_maybe_cast(s.reshape(1, -1), cast_dtype) for s in shards], axis=1 - ) - big_full = mx.distributed.all_gather(big_shard, group=group) - parts = mx.split(big_full, split_indices, axis=1) - return [p.reshape(shape) for p, shape in zip(parts, full_shapes)] - - @gather.vjp - def gather_vjp(shards, cotangents, _): - big_cot_full = mx.concatenate([c.reshape(S, -1) for c in cotangents], axis=1) - big_cot_shard = mx.distributed.sum_scatter(big_cot_full, group=group) / S - parts = mx.split(big_cot_shard, split_indices, axis=1) - return [p.reshape(shape) for p, shape in zip(parts, shard_shapes)] - - 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): - def __init__(self, module, group, cast_dtype): - 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"FSDP: parameter {path} is a 0-D scalar and cannot be sharded." - ) - if a.shape[0] % N != 0: - raise ValueError( - f"FSDP: parameter {path} has shape {a.shape}; axis 0 must " - f"be divisible by the FSDP group size {N}." - ) - - self._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, cast_dtype) - - 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 clip_grad_norm_sharded( + gradients: Any, + max_norm: float, + group: Optional[mx.distributed.Group] = None, +): + """Clip the global norm of gradients that are sharded across a group. + 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`. -def fully_shard( - module: Module, - group: Optional["mx.distributed.Group"] = None, - cast_dtype: Optional[mx.Dtype] = None, -) -> Module: - group = group or mx.distributed.init() - if group.size() == 1: - return module - if isinstance(module, FullyShardedModule): - return module + Args: + 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``. - wrapped = FullyShardedModule(module, group, cast_dtype) - return wrapped if wrapped._paths else module + Returns: + (Any, mlx.core.array): The possibly rescaled local shard of the + gradients and the global gradient norm. + """ + 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 diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index 00d3182bc0..77bdea6a1f 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -1,10 +1,12 @@ # Copyright © 2025 Apple Inc. +import math + 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.utils import average_gradients +from mlx.nn.utils import average_gradients, clip_grad_norm_sharded class MLXDistributedCommonTestCase(mlx_tests.MLXTestCase): @@ -322,3 +324,39 @@ def test_all_gather(self): y = mx.distributed.all_gather(x) self.assertEqual(y.shape, (world.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) + + def test_clip_grad_norm_sharded(self): + world = mx.distributed.init() + N = world.size() + + value = 3.0 + grads_slice = {"a": mx.ones((4, 3)) * value, "b": mx.ones((5,)) * value} + local_numel = 4 * 3 + 5 + expected_norm = math.sqrt(N * local_numel) * value + + clipped, grad_norm = clip_grad_norm_sharded( + grads_slice, max_norm=1e9, group=world + ) + mx.eval(clipped, grad_norm) + self.assertTrue( + mx.allclose( + grad_norm, mx.array(expected_norm), atol=self.atol, rtol=self.rtol + ) + ) + for k in grads_slice: + self.assertTrue( + mx.allclose(clipped[k], grads_slice[k], atol=self.atol, rtol=self.rtol) + ) + + max_norm = 1.0 + clipped, grad_norm = clip_grad_norm_sharded( + grads_slice, max_norm=max_norm, group=world + ) + mx.eval(clipped, grad_norm) + scale = max_norm / (expected_norm + 1e-6) + for k in grads_slice: + self.assertTrue( + mx.allclose( + clipped[k], grads_slice[k] * scale, atol=self.atol, rtol=self.rtol + ) + ) diff --git a/python/tests/nccl_test_distributed.py b/python/tests/nccl_test_distributed.py index 7a6bcce04b..db9bbb2e62 100644 --- a/python/tests/nccl_test_distributed.py +++ b/python/tests/nccl_test_distributed.py @@ -1,10 +1,11 @@ # Copyright © 2024 Apple Inc. import mlx.core as mx -import mlx.optimizers as optim +import mlx.nn as nn import mlx_distributed_tests import mlx_tests -from mlx.nn.utils import average_gradients, fsdp_apply_gradients +from mlx.nn.utils import average_gradients +from mlx.utils import tree_flatten, tree_map class TestNCCLDistributed(mlx_distributed_tests.MLXDistributedCommonTestCase): @@ -116,230 +117,66 @@ def test_all_gather_split(self): self.assertEqual(y.shape, (sub.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) - def test_fsdp_apply_gradients(self): - world = mx.distributed.init() - N = world.size() - - params = { - "w1": mx.ones((N * 10, 8)), - "w2": mx.ones((N * 20,)), - } - grads = { - "w1": mx.ones((N * 10, 8)) * 0.1, - "w2": mx.ones((N * 20,)) * 0.1, - } - - optimizer = optim.SGD(learning_rate=0.1) - updated_params_fsdp = fsdp_apply_gradients(grads, params, optimizer) - mx.eval(updated_params_fsdp) - - self.assertEqual(updated_params_fsdp["w1"].shape, (N * 10, 8)) - self.assertEqual(updated_params_fsdp["w2"].shape, (N * 20,)) - - self.assertTrue( - mx.allclose( - updated_params_fsdp["w1"], mx.ones((N * 10, 8)) * 0.99, atol=1e-6 - ) - ) - self.assertTrue( - mx.allclose(updated_params_fsdp["w2"], mx.ones((N * 20,)) * 0.99, atol=1e-6) - ) - - grads = { - "w1": mx.ones((N * 10, 8)) * 10.0, - "w2": mx.ones((N * 20,)) * 10.0, - } - - new_params_clipped, grad_norm = fsdp_apply_gradients( - grads, params, optimizer, max_norm=1.0 - ) - mx.eval(new_params_clipped, grad_norm) - - self.assertIsNotNone(grad_norm) - expected_norm = mx.sqrt((N * 10 * 8 + N * 20) * 100.0) - self.assertTrue(mx.allclose(grad_norm, expected_norm, atol=1e-4, rtol=1e-4)) - self.assertEqual(new_params_clipped["w1"].shape, (N * 10, 8)) - self.assertEqual(new_params_clipped["w2"].shape, (N * 20,)) - - scale = 1.0 / expected_norm - expected_update = 1.0 - 0.1 * 10.0 * scale - self.assertTrue( - mx.allclose( - new_params_clipped["w1"], - mx.ones((N * 10, 8)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - self.assertTrue( - mx.allclose( - new_params_clipped["w2"], - mx.ones((N * 20,)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - params = {"w": mx.ones((N * 4,))} - grads = {"w": mx.ones((N * 4,)) * 0.5} - - optimizer_fsdp = optim.SGD(learning_rate=0.1) - updated_params_fsdp = fsdp_apply_gradients(grads, params, optimizer_fsdp) - - optimizer_ddp = optim.SGD(learning_rate=0.1) - avg_grads = average_gradients(grads) - updated_params_ddp = optimizer_ddp.apply_gradients(avg_grads, params) - mx.eval(updated_params_ddp, updated_params_fsdp) - - self.assertTrue( - mx.allclose( - updated_params_fsdp["w"], updated_params_ddp["w"], atol=1e-6, rtol=1e-4 - ), - ) + def test_fully_shard_grads(self): + dtypes = [ + (mx.float32, 1e-6, 1e-6), + (mx.bfloat16, 1e-3, 1e-3), + ] - def test_fsdp_ddp_apply_gradients(self): - world = mx.distributed.init() - N = world.size() - S = 4 - fsdp_group = world.split(world.rank() // S) - dp_group = world.split(world.rank() % S) - - self.assertEqual(fsdp_group.size(), S) - self.assertEqual(dp_group.size(), N // S) - - params = { - "w1": mx.ones((S * 10, 8)), - "w2": mx.ones((S * 20,)), - } - grads = { - "w1": mx.ones((S * 10, 8)) * 0.1, - "w2": mx.ones((S * 20,)) * 0.1, - } - - optimizer = optim.SGD(learning_rate=0.1) - updated = fsdp_apply_gradients( - grads, - params, - optimizer, - fsdp_group=fsdp_group, - dp_group=dp_group, - ) - mx.eval(updated) - - self.assertEqual(updated["w1"].shape, (S * 10, 8)) - self.assertEqual(updated["w2"].shape, (S * 20,)) - - self.assertTrue( - mx.allclose(updated["w1"], mx.ones((S * 10, 8)) * 0.99, atol=1e-6) - ) - self.assertTrue( - mx.allclose(updated["w2"], mx.ones((S * 20,)) * 0.99, atol=1e-6) - ) - - grads_big = { - "w1": mx.ones((S * 10, 8)) * 10.0, - "w2": mx.ones((S * 20,)) * 10.0, - } - - optimizer2 = optim.SGD(learning_rate=0.1) - clipped, grad_norm = fsdp_apply_gradients( - grads_big, - params, - optimizer2, - fsdp_group=fsdp_group, - dp_group=dp_group, - max_norm=1.0, - ) - mx.eval(clipped, grad_norm) - - self.assertIsNotNone(grad_norm) - expected_norm = mx.sqrt((S * 10 * 8 + S * 20) * 100.0) - self.assertTrue(mx.allclose(grad_norm, expected_norm, atol=1e-4, rtol=1e-4)) - self.assertEqual(clipped["w1"].shape, (S * 10, 8)) - self.assertEqual(clipped["w2"].shape, (S * 20,)) - - scale = 1.0 / expected_norm - expected_update = 1.0 - 0.1 * 10.0 * scale - self.assertTrue( - mx.allclose( - clipped["w1"], - mx.ones((S * 10, 8)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - self.assertTrue( - mx.allclose( - clipped["w2"], - mx.ones((S * 20,)) * expected_update, - atol=1e-4, - rtol=1e-4, - ) - ) - - params_eq = {"w": mx.ones((S * 4,))} - grads_eq = {"w": mx.ones((S * 4,)) * 0.5} - - optimizer_hybrid = optim.SGD(learning_rate=0.1) - updated_hybrid = fsdp_apply_gradients( - grads_eq, - params_eq, - optimizer_hybrid, - fsdp_group=fsdp_group, - dp_group=dp_group, - ) - - optimizer_ddp = optim.SGD(learning_rate=0.1) - avg_grads = average_gradients(grads_eq) - updated_ddp = optimizer_ddp.apply_gradients(avg_grads, params_eq) - mx.eval(updated_hybrid, updated_ddp) - - self.assertTrue( - mx.allclose(updated_hybrid["w"], updated_ddp["w"], atol=1e-6, rtol=1e-4), - ) - - def test_fsdp_peak_memory(self): world = mx.distributed.init() N = world.size() - mx.random.seed(42) - params = { - "w1": mx.random.normal((N * 1024, 1024)), - "w2": mx.random.normal((N * 2048, 512)), - } - grads = { - "w1": mx.random.normal((N * 1024, 1024)), - "w2": mx.random.normal((N * 2048, 512)), - } - mx.eval(params, grads) - optimizer_ddp = optim.Adam(learning_rate=0.01) - optimizer_fsdp = optim.Adam(learning_rate=0.01) - - def pseudo_step_ddp(grads, params, optimizer): - grads = average_gradients(grads) - grads, grad_norm = optim.clip_grad_norm(grads, max_norm=1.0) - params = optimizer.apply_gradients(grads, params) - return grad_norm, params - - def pseudo_step_fsdp(grads, params, optimizer): - params, grad_norm = fsdp_apply_gradients( - grads, params, optimizer, max_norm=1.0 + rank = world.rank() + dims = 8 * N + part = slice(rank * dims // N, (rank + 1) * dims // N) + + class MLP(nn.Module): + def __init__(self, dims): + super().__init__() + self.l1 = nn.Linear(dims, dims) + self.l2 = nn.Linear(dims, dims) + + def __call__(self, x): + return self.l2(nn.relu(self.l1(x))) + + def loss_fn(model, x, y): + logits = model(x).astype(mx.float32) + return ((logits - y) ** 2).mean() + + for dtype, atol, rtol in dtypes: + + mx.random.seed(0xF0F0F0F0) + + kx, ky = mx.random.split(mx.random.key(rank)) + x = mx.random.normal((4, dims), dtype=dtype, key=kx) + y = mx.random.normal((4, dims), key=ky) + + # DDP reference: replicated params, gradients averaged across ranks + model = MLP(dims) + params = model.trainable_parameters() + + # Cast parameters to dtype for forward + model.update(tree_map(lambda p: p.astype(dtype), params)) + loss, grads = mx.value_and_grad(loss_fn)(model, x, y) + grads = average_gradients(grads, group=world) + loss = mx.distributed.all_sum(loss, group=world) / N + mx.eval(loss, grads) + + # Shard the model + model_sharded = MLP(dims) + model_sharded.update(params) + model_sharded = nn.fully_shard(model_sharded, compute_dtype=dtype) + loss_sharded, grads_sharded = mx.value_and_grad(loss_fn)( + model_sharded, x, y ) - return grad_norm, params - - mx.reset_peak_memory() - - for i in range(10): - grad_norm, params = pseudo_step_ddp(grads, params, optimizer_ddp) - mx.eval(grad_norm, params) - - ddp_peak_memory = mx.get_peak_memory() - mx.reset_peak_memory() - - for i in range(10): - grad_norm, params = pseudo_step_fsdp(grads, params, optimizer_fsdp) - mx.eval(grad_norm, params) - fsdp_peak_memory = mx.get_peak_memory() - self.assertTrue(fsdp_peak_memory < ddp_peak_memory) + loss_sharded = mx.distributed.all_sum(loss_sharded, group=world) / N + mx.eval(loss_sharded, grads_sharded) + grads_ref = dict(tree_flatten(grads)) + self.assertTrue(mx.allclose(loss, loss_sharded, atol=1e-4, rtol=1e-4)) + for k, gs in tree_flatten(grads_sharded["module"]): + self.assertTrue( + mx.allclose(gs, grads_ref[k][part], atol=atol, rtol=rtol) + ) if __name__ == "__main__":