zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-22 21:48:08 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

rebase

Browse Source
This commit is contained in:
CircleCI Docs
2025-01-09 21:56:20 +00:00
parent 04b749a588
commit d8d647015b
2642 changed files with 137687 additions and 70861 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
docs/build/html/_sources/usage/distributed.rst vendored
View File

@@ -141,12 +141,13 @@ everything else remaining the same.
from mlx.utils import tree_map
def all_reduce_grads(grads):
N = mx.distributed.init()
N = mx.distributed.init().size()
if N == 1:
return grads
return tree_map(
lambda x: mx.distributed.all_sum(x) / N,
grads)
lambda x: mx.distributed.all_sum(x) / N,
grads
)
def step(model, x, y):
loss, grads = loss_grad_fn(model, x, y)