Skip to content

[WIP] [CUDA] fsdp #3768

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
nastya236 wants to merge 1 commit into
base: main
Choose a base branch
from
+88 −129
Draft

[WIP] [CUDA] fsdp #3768

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
217 changes: 88 additions & 129 deletions python/mlx/nn/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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)
Expand All @@ -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
Loading