Skip to content

add runnable data parallel example#3905

Open
tylergibbs1 wants to merge 1 commit into
ml-explore:mainfrom
tylergibbs1:agent/add-data-parallel-example
Open

add runnable data parallel example#3905
tylergibbs1 wants to merge 1 commit into
ml-explore:mainfrom
tylergibbs1:agent/add-data-parallel-example

Conversation

@tylergibbs1

Copy link
Copy Markdown

proposed changes

#2930 asks for standalone distributed examples. #2973 added the explanation, but examples/python still has no data parallel script users can run end to end.

this adds a linear regression example where each rank creates its own synthetic data shard, starts from the same parameters, and averages gradients with nn.average_gradients. the final loss is evaluated on every rank before rank zero prints. without that step, lazy evaluation can leave one process waiting on a collective after its peer exits.

run it with:

mlx.launch -n 2 python examples/python/data_parallel.py

on my m4 mac, the one rank run finished with an l2 distance of 0.00971. the two rank run finished at 0.00464.

checks

  • pre-commit run --all-files
  • python examples/python/data_parallel.py
  • mlx.launch -n 2 python examples/python/data_parallel.py

checklist

  • i have read the contributing document
  • i have run pre-commit run --all-files
  • i have added automated tests. this pr adds a runnable example and the checks above exercise both paths
  • i have updated the necessary documentation. no api or existing docs changed

the distributed guide explains gradient averaging with placeholders, but `examples/python` has no script users can run with `mlx.launch`. refs ml-explore#2930.
@tylergibbs1
tylergibbs1 marked this pull request as ready for review July 24, 2026 04:50
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant