From a9e9e5aa389f54dd6e702835cdc77cc8be1481b7 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 25 Jun 2026 14:54:33 +0200 Subject: [PATCH] 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