Distributed support - #482
ronaldmannak wants to merge 9 commits into
Conversation
Compile mlx-c/mlx/c/distributed.cpp and distributed_group.cpp into the Cmlx target. No backend is enabled yet, so every backend reports unavailable and distributed::init() returns an empty group of size 1. Adds MLXDistributed.isAvailable() and a test, which together verify the wrappers both compile and link.
Build mlx/distributed/ring/ring.cpp instead of the no_ring.cpp stub. Ring needs only TCP sockets, so it builds on every platform without additional include paths or SDK requirements. Without MLX_HOSTFILE and MLX_RANK, ring::init() returns null for a non-strict init, so the default group remains an empty group of size 1 and single-process tests are unaffected.
Port mlx.core.distributed to Swift: MLXDistributed.Group with rank, size and split, plus the eight collectives and point to point operations (allSum, allMax, allMin, allGather, sumScatter, send, recv, recvLike). Following the rest of the Swift API none of these throw -- errors raised by MLX are reported through mlx_set_error_handler and surface via withError()/checkedEval(). Tests run single process: at size one MLX short circuits the collectives to the identity and rejects the point to point operations, which covers the whole API surface without a launcher.
Port the core of mlx.nn.layers.distributed: AllToShardedLinear and ShardedToAllLinear, the shardLinear/shardInPlace helpers, sumGradients and averageGradients. Python's segments argument accepts an int or a list of either indices or fractions; that becomes the Segments enum here. The from_linear class methods become convenience initializers taking a Linear, matching how QuantizedLinear converts. The quantized layers and fully_shard are not included.
The distributed primitives have no GPU implementation on Metal -- AllReduce::eval_gpu and friends in backend/metal/distributed.cpp throw unconditionally -- so a collective submitted to the default GPU stream always fails. MLX handles this by letting each backend choose its device when no stream is given: Group::communication_stream returns to_stream(s, Device::cpu) for ring, mpi and jaccl, and Device::gpu for nccl. That path is unreachable through mlx-c, whose mlx_stream_get_() rejects a null handle, so Swift has to pass a concrete stream. Default to the CPU stream, which is correct for every backend that can be built on Apple platforms. This will need revisiting if nccl is ever enabled; the real fix is for mlx-c to accept an unspecified stream.
The test process is itself a rank: with MLX_TEST_DISTRIBUTED=1 the first
process reserves two loopback ports, writes a temporary hostfile,
re-executes the test bundle as rank 1 and runs as rank 0. Both ranks
execute the same body, so no separate worker executable is needed --
this mirrors how the Python distributed tests are run under mlx.launch.
Skipped unless MLX_TEST_DISTRIBUTED=1, matching the Python tests, which
are likewise not part of CI. It can be enabled by adding a second
xctest invocation with that variable set.
There is a single test method because MLX caches the group per process,
and because every rank must issue the same collectives in the same
order. The single process tests skip during a multi process run, since
their assertions describe a group of size one.
Run with:
xcodebuild build-for-testing -scheme mlx-swift-Package -destination 'platform=macOS'
MLX_TEST_DISTRIBUTED=1 xcrun xctest -XCTest DistributedRingTests \
~/Library/Developer/Xcode/DerivedData/mlx-swift-*/Build/Products/Debug/MLXTests.xctest
The single process tests only exercise the sharded layers in a group of size one, where sharding is the identity and the interesting paths -- splitting the weight, selecting this rank's slice, reducing inside the layer -- never run. Every rank seeds identically and builds the same Linear, then: AllToShardedLinear holds half the output dimensions; gathering the transposed partial results reassembles the columns in rank order and must equal the unsharded result. ShardedToAllLinear holds half the input dimensions and is fed the matching slice; its internal allSum must reproduce the unsharded result on both ranks.
mlx_distributed_init reports failure through the error handler and leaves the group handle null. Wrapping that in a Group produced an object whose rank and size silently read as zero, so a caller that did not wrap the call tightly in withError got a plausible looking group that was quietly broken. Make both initialize and split failable, and add globalGroup for the common case where a non-strict init cannot meaningfully fail. NOT BUILT.
|
Pushed one fix found while enabling JACCL: |
| public func shardLinear( | ||
| _ layer: Linear, sharding: ShardingType, segments: Segments = .count(1), | ||
| group: MLXDistributed.Group? = nil | ||
| ) -> Module { | ||
| switch sharding { | ||
| case .allToSharded: | ||
| AllToShardedLinear(layer, segments: segments, group: group) |
There was a problem hiding this comment.
If passed a QuantizedLinear this will (I think) not do anything good. It looks like the python side has QuantizedAllToShardedLinear etc.
|
Doesn't have to be fixed right away, but the new test files have to be added in the xcode/MLX.xcodeproj. I think just opening that file in Xcode should populate it. |
| /// the result is sharded across the group. | ||
| /// | ||
| /// The gradients are automatically aggregated from each member of the group. | ||
| open class AllToShardedLinear: Module, UnaryLayer { |
There was a problem hiding this comment.
Is this something that would be substituted inside the model, like quantization does today? If so, I wonder:
- do we need marker protocols like Quantizable/Quantized?
- should this be a subtype of Linear?
Still working on understanding how you use this in practice so these questions might be silly.
There was a problem hiding this comment.
Reading the python code it isn't generic like quantization. There is an optional method at the LanguageModel layer:
def shard(self, group: Optional[mx.distributed.Group] = None):
group = group or mx.distributed.init()
N = group.size()
for layer in self.model.layers:
# Shard the self attention
layer.self_attn.q_proj = shard_linear(
layer.self_attn.q_proj, "all-to-sharded", group=group
)
layer.self_attn.k_proj = shard_linear(
layer.self_attn.k_proj, "all-to-sharded", group=group
)So I think Quantizable/Quantized isn't strictly required. We could add it, but there is no clear use for it.
Making this a subtype might be useful. Otherwise we would have to change the swift models from declaring modules as Linear -> modules as UnaryLayer. I think it loses some of the documentation of the actual types but since authors have to write the shard method by hand anyway, it could be ok.
So combining the two ideas, perhaps a protocol like ShardableLinear would be useful -- both Linear and AllToShardedLinear could conform. This keeps the documentation of the type and prevents accidental sharding of the wrong layers.
Anyway, something to think of -- this PR isn't adding the adoption just yet, only the underpinnings.
| private func shard( | ||
| _ parameters: ModuleParameters, group: MLXDistributed.Group, | ||
| _ sharding: ShardingPredicate | ||
| ) -> ModuleParameters { |
There was a problem hiding this comment.
Looking at how this works on the python side:
# mlx-lm def sharded_load
model, _ = load_model(
model_path, lazy=True, strict=False, trust_remote_code=trust_remote_code
)
# dkoski -- still lazy at this point
if tensor_group is not None:
model.shard(tensor_group)
if pipeline_group is not None:
model.model.pipeline(pipeline_group)
mx.eval(model.parameters())We have recently made some changes that make loading faster but non-lazy -- that may defeat some of this behavior (meaning, I think, very large models will require full memory to load at peak).
Well, this is special loading code for distributed in python, so I guess we need something similar in swift. Anyway, just highlighting the point that swift is maybe diverging slightly from the python core in loading.
@aleroot FYI
| /// | ||
| /// - Parameters: | ||
| /// - backend: the backend to use, defaulting to ``Backend/any`` | ||
| /// - strict: if `true` report an error when no backend can be initialized |
There was a problem hiding this comment.
How is the error reported?
I wonder if this would be easier if we had two calls: one that throws and one that does not?
| public func shardLinear( | ||
| _ layer: Linear, sharding: ShardingType, segments: Segments = .count(1), | ||
| group: MLXDistributed.Group? = nil | ||
| ) -> Module { |
There was a problem hiding this comment.
See the question below on AllToShardedLinear -- this probably needs to return something that can be called. A plain Module type doesn't have callAsFunction()
2 participants
Proposed changes
Distributed support for MLX Swift
Ports
mlx.core.distributedand the core ofmlx.nn.layers.distributed, andenables the ring backend. Supersedes #371, which was written against mlx 0.31.1.
Not included, deliberately: the quantized sharded layers,
fully_shard, and theJACCL backend, which will be added in a separate PR.
mlx-c gaps
Three things in
mlx::core::distributedaren't reachable through mlx-c. Noneblock this PR, but (1) and (2) would need to land before a follow-up.
1. No way to pass an unspecified stream.
mlx_stream_get_()rejects a nullhandle, so
Group::communication_streamnever applies and every caller musthardcode a device, CPU here, which is wrong the moment NCCL is enabled, since
NCCL is the one backend that wants the GPU.
The distributed entry points already take
const mlx_distributed_group group /* may be null */and read null as "the global group". Letting
mlx_streamfollow the sameconvention would resolve it without an API change:
Blocks: correct NCCL support, and full parity with Python's
stream=None.2. The JACCL side-channel
initoverload isn't exposed.Python offers this as
all_gather_factory. There's no C equivalent, so Swiftcan't supply a custom rendezvous.
3.
clear_backends()isn't exposed. Groups are cached per backend for theprocess lifetime with no way to release them. Python registers it with
atexit.Blocks: nothing outright, but it makes test isolation awkward, a process that
has formed a real group can't go back to a singleton one.
Tests
DistributedRingTestsruns two real ranks on loopback, including the shardedlayers. Set
MLX_TEST_DISTRIBUTED=1to run it:Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes