add runnable data parallel example#3905
Open
tylergibbs1 wants to merge 1 commit into
Open
Conversation
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
marked this pull request as ready for review
July 24, 2026 04:50
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
proposed changes
#2930 asks for standalone distributed examples. #2973 added the explanation, but
examples/pythonstill 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:
on my m4 mac, the one rank run finished with an l2 distance of
0.00971. the two rank run finished at0.00464.checks
pre-commit run --all-filespython examples/python/data_parallel.pymlx.launch -n 2 python examples/python/data_parallel.pychecklist
pre-commit run --all-files