electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add uneven sharding

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-03-31 13:22:05 +01:00
parent 4688adb5d2
commit 55fa5362bb
4 changed files with 669 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
+1 -1
View File
@@ -62,7 +62,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
[tool.uv.sources]
exo_pyo3_bindings = { workspace = true }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "leo/add-uneven-sharding", marker = "sys_platform == 'darwin'" }
mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-arrayscache-leak" }
# Uncomment to use local mlx/mlx-lm development versions:
# mlx = { path = "/Users/Shared/mlx", editable=true }
src/exo/worker/engines/mlx/auto_parallel.py
+144 -104
View File
@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Literal, Protocol, cast
import mlx.core as mx
import mlx.nn as nn
from mlx.nn.layers.distributed import (
compute_shard_sizes,
shard_inplace,
shard_linear,
sum_gradients,
@@ -657,13 +658,15 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
for i, layer in enumerate(model.layers):
# Force load weights before sharding to avoid FAST_SYNCH deadlock
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
layer.self_attn.n_heads //= self.N
if layer.self_attn.n_kv_heads is not None:
layer.self_attn.n_kv_heads //= self.N
head_dim = layer.self_attn.head_dim
n_kv = layer.self_attn.n_kv_heads or layer.self_attn.n_heads
gqa_unit = head_dim * (layer.self_attn.n_heads // n_kv)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
layer.self_attn.n_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
@@ -715,22 +718,24 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
original_num_heads = layer.self_attn.num_heads
q_head_dim = layer.self_attn.q_b_proj.weight.shape[0] // original_num_heads if layer.self_attn.q_lora_rank is not None else layer.self_attn.q_proj.weight.shape[0] // original_num_heads
if layer.self_attn.q_lora_rank is None:
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
layer.self_attn.q_proj, unit=q_head_dim
)
else:
layer.self_attn.q_b_proj = self.all_to_sharded_linear(
layer.self_attn.q_b_proj
layer.self_attn.q_b_proj, unit=q_head_dim
)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
layer.self_attn.num_heads //= self.N
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=q_head_dim)
head_sizes = compute_shard_sizes(original_num_heads, self.N)
layer.self_attn.num_heads = head_sizes[self.group.rank()]
# Logic from upstream mlx
num_heads = layer.self_attn.num_heads
sh = self.group.rank() * num_heads
eh = sh + num_heads
sh = sum(head_sizes[:self.group.rank()])
eh = sh + head_sizes[self.group.rank()]
def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
return w[sh:eh]
@@ -802,22 +807,24 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
timeout_seconds / total,
on_timeout,
)
original_num_heads = layer.self_attn.num_heads # type: ignore
q_head_dim = layer.self_attn.q_b_proj.weight.shape[0] // original_num_heads if layer.self_attn.q_lora_rank is not None else layer.self_attn.q_proj.weight.shape[0] // original_num_heads # type: ignore
if layer.self_attn.q_lora_rank is None: # type: ignore
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
layer.self_attn.q_proj, unit=q_head_dim
)
else:
layer.self_attn.q_b_proj = self.all_to_sharded_linear(
layer.self_attn.q_b_proj
layer.self_attn.q_b_proj, unit=q_head_dim
)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
layer.self_attn.num_heads //= self.N
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=q_head_dim)
head_sizes = compute_shard_sizes(original_num_heads, self.N)
layer.self_attn.num_heads = head_sizes[self.group.rank()]
# Logic from upstream mlx
num_heads = layer.self_attn.num_heads
sh = self.group.rank() * num_heads
eh = sh + num_heads
sh = sum(head_sizes[:self.group.rank()])
eh = sh + head_sizes[self.group.rank()]
def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
return w[sh:eh]
@@ -944,13 +951,15 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
head_dim = layer.self_attn.head_dim
gqa_unit = head_dim * (layer.self_attn.num_attention_heads // layer.self_attn.num_key_value_heads)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
layer.self_attn.num_attention_heads //= self.N
layer.self_attn.num_key_value_heads //= self.N
layer.self_attn.num_attention_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
layer.self_attn.num_key_value_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group) # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
@@ -988,20 +997,22 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
if isinstance(layer, Qwen3MoeDecoderLayer):
head_dim = layer.self_attn.q_proj.weight.shape[0] // layer.self_attn.n_heads
gqa_unit = head_dim * (layer.self_attn.n_heads // layer.self_attn.n_kv_heads)
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
layer.self_attn.q_proj, unit=gqa_unit
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj
layer.self_attn.k_proj, unit=head_dim
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj
layer.self_attn.v_proj, unit=head_dim
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj
layer.self_attn.o_proj, unit=gqa_unit
)
layer.self_attn.n_heads //= self.N
layer.self_attn.n_kv_heads //= self.N
layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
layer.self_attn.n_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
else:
assert isinstance(layer, (Qwen3NextDecoderLayer, Qwen3_5DecoderLayer))
if hasattr(layer, "linear_attn"):
@@ -1019,16 +1030,19 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
# Qwen3.5: separate projections
# in_proj_qkv has sections [q(key_dim), k(key_dim), v(value_dim)]
# that must be split section-aware, not as a contiguous block
head_k_dim = linear_attn.head_k_dim
head_v_dim = linear_attn.head_v_dim
key_dim = linear_attn.key_dim
value_dim = linear_attn.value_dim
linear_attn.in_proj_qkv = shard_linear(
linear_attn.in_proj_qkv,
"all-to-sharded",
segments=[key_dim, key_dim + key_dim],
unit=head_k_dim,
group=self.group,
)
linear_attn.in_proj_z = self.all_to_sharded_linear(
linear_attn.in_proj_z
linear_attn.in_proj_z, unit=head_v_dim
)
linear_attn.in_proj_b = self.all_to_sharded_linear(
linear_attn.in_proj_b
@@ -1037,7 +1051,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
linear_attn.in_proj_a
)
linear_attn.out_proj = self.sharded_to_all_linear(
linear_attn.out_proj
linear_attn.out_proj, unit=linear_attn.head_v_dim
)
# Shard conv1d: depthwise conv with non-contiguous channel slicing.
@@ -1046,31 +1060,37 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
rank = self.group.rank()
key_dim = linear_attn.key_dim
value_dim = linear_attn.value_dim
key_dim_shard = key_dim // self.N
value_dim_shard = value_dim // self.N
head_k_dim = linear_attn.head_k_dim
head_v_dim = linear_attn.head_v_dim
key_shard_sizes = compute_shard_sizes(key_dim, self.N, unit=head_k_dim)
value_shard_sizes = compute_shard_sizes(value_dim, self.N, unit=head_v_dim)
key_dim_shard = key_shard_sizes[rank]
value_dim_shard = value_shard_sizes[rank]
key_dim_offset = sum(key_shard_sizes[:rank])
value_dim_offset = sum(value_shard_sizes[:rank])
q_idx = mx.arange(rank * key_dim_shard, (rank + 1) * key_dim_shard)
q_idx = mx.arange(key_dim_offset, key_dim_offset + key_dim_shard)
k_idx = mx.arange(
key_dim + rank * key_dim_shard,
key_dim + (rank + 1) * key_dim_shard,
key_dim + key_dim_offset,
key_dim + key_dim_offset + key_dim_shard,
)
v_idx = mx.arange(
2 * key_dim + rank * value_dim_shard,
2 * key_dim + (rank + 1) * value_dim_shard,
2 * key_dim + value_dim_offset,
2 * key_dim + value_dim_offset + value_dim_shard,
)
conv_indices = mx.concatenate([q_idx, k_idx, v_idx])
linear_attn.conv1d.weight = linear_attn.conv1d.weight[conv_indices]
new_conv_dim = key_dim_shard * 2 + value_dim_shard
linear_attn.conv1d.groups = new_conv_dim
num_v_shard = linear_attn.num_v_heads // self.N
v_start = rank * num_v_shard
v_end = v_start + num_v_shard
linear_attn.A_log = linear_attn.A_log[v_start:v_end]
linear_attn.dt_bias = linear_attn.dt_bias[v_start:v_end]
num_k_per_rank = key_dim_shard // head_k_dim
num_v_per_rank = value_dim_shard // head_v_dim
v_offset = value_dim_offset // head_v_dim
linear_attn.A_log = linear_attn.A_log[v_offset:v_offset + num_v_per_rank]
linear_attn.dt_bias = linear_attn.dt_bias[v_offset:v_offset + num_v_per_rank]
linear_attn.num_k_heads //= self.N
linear_attn.num_v_heads //= self.N
linear_attn.num_k_heads = num_k_per_rank
linear_attn.num_v_heads = num_v_per_rank
linear_attn.key_dim = (
linear_attn.head_k_dim * linear_attn.num_k_heads
)
@@ -1081,20 +1101,22 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
linear_attn.key_dim * 2 + linear_attn.value_dim
)
else:
kv_head_dim = layer.self_attn.k_proj.weight.shape[0] // layer.self_attn.num_key_value_heads
gqa_repeat = layer.self_attn.num_attention_heads // layer.self_attn.num_key_value_heads
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
layer.self_attn.q_proj, unit=kv_head_dim * 2 * gqa_repeat
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj
layer.self_attn.k_proj, unit=kv_head_dim
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj
layer.self_attn.v_proj, unit=kv_head_dim
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj
layer.self_attn.o_proj, unit=kv_head_dim * gqa_repeat
)
layer.self_attn.num_attention_heads //= self.N
layer.self_attn.num_key_value_heads //= self.N
layer.self_attn.num_attention_heads = layer.self_attn.q_proj.weight.shape[0] // (kv_head_dim * 2)
layer.self_attn.num_key_value_heads = layer.self_attn.k_proj.weight.shape[0] // kv_head_dim
# Shard the MoE.
if isinstance(
@@ -1146,12 +1168,14 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
layer.self_attn.n_heads //= self.N
layer.self_attn.n_kv_heads //= self.N
head_dim = layer.self_attn.q_proj.weight.shape[0] // layer.self_attn.n_heads
gqa_unit = head_dim * (layer.self_attn.n_heads // layer.self_attn.n_kv_heads)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
layer.self_attn.n_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
if isinstance(layer.mlp, MoE):
self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
@@ -1194,23 +1218,26 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
head_dim = layer.self_attn.head_dim
original_num_heads = layer.self_attn.num_attention_heads
gqa_unit = head_dim * (layer.self_attn.num_attention_heads // layer.self_attn.num_key_value_heads)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
layer.self_attn.num_attention_heads //= self.N
layer.self_attn.num_key_value_heads //= self.N
layer.self_attn.num_attention_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
layer.self_attn.num_key_value_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
layer.self_attn.num_key_value_groups = (
layer.self_attn.num_attention_heads
// layer.self_attn.num_key_value_heads
)
layer.self_attn.sinks = layer.self_attn.sinks[
layer.self_attn.num_attention_heads
* self.group.rank() : layer.self_attn.num_attention_heads
* (self.group.rank() + 1)
]
rank = self.group.rank()
q_head_sizes = compute_shard_sizes(original_num_heads, self.N, unit=gqa_unit // head_dim)
sink_start = sum(q_head_sizes[:rank])
sink_end = sink_start + q_head_sizes[rank]
layer.self_attn.sinks = layer.self_attn.sinks[sink_start:sink_end]
self.all_to_sharded_linear_in_place(layer.mlp.experts.gate_proj)
self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj)
@@ -1237,17 +1264,19 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
head_dim = layer.self_attn.head_dim
gqa_unit = head_dim * (layer.self_attn.num_heads // layer.self_attn.num_kv_heads)
layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj, unit=gqa_unit)
layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj, unit=head_dim)
layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj, unit=head_dim)
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj, unit=gqa_unit)
layer.self_attn.num_heads //= self.N
layer.self_attn.num_kv_heads //= self.N
layer.self_attn.num_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
layer.self_attn.num_kv_heads = layer.self_attn.k_proj.weight.shape[0] // head_dim
if getattr(layer.self_attn, "use_head_wise_attn_gate", False):
layer.self_attn.g_proj = self.all_to_sharded_linear(
layer.self_attn.g_proj
layer.self_attn.g_proj, unit=gqa_unit // head_dim
)
if isinstance(layer.mlp, Step35MLP):
@@ -1286,12 +1315,14 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
mixer = layer.mixer
if isinstance(mixer, NemotronHAttention):
mixer.q_proj = self.all_to_sharded_linear(mixer.q_proj)
mixer.k_proj = self.all_to_sharded_linear(mixer.k_proj)
mixer.v_proj = self.all_to_sharded_linear(mixer.v_proj)
mixer.o_proj = self.sharded_to_all_linear(mixer.o_proj)
mixer.num_heads //= self.N
mixer.num_key_value_heads //= self.N
attn_head_dim = mixer.head_dim
gqa_unit = attn_head_dim * (mixer.num_heads // mixer.num_key_value_heads)
mixer.q_proj = self.all_to_sharded_linear(mixer.q_proj, unit=gqa_unit)
mixer.k_proj = self.all_to_sharded_linear(mixer.k_proj, unit=attn_head_dim)
mixer.v_proj = self.all_to_sharded_linear(mixer.v_proj, unit=attn_head_dim)
mixer.o_proj = self.sharded_to_all_linear(mixer.o_proj, unit=gqa_unit)
mixer.num_heads = mixer.q_proj.weight.shape[0] // attn_head_dim
mixer.num_key_value_heads = mixer.k_proj.weight.shape[0] // attn_head_dim
elif isinstance(mixer, NemotronHMamba2Mixer):
self._shard_mamba2_mixer(mixer, rank)
@@ -1322,12 +1353,22 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
ssm_state_size = mixer.ssm_state_size
intermediate_size = mixer.intermediate_size # = num_heads * head_dim
# Per-rank sizes
heads_per_rank = num_heads // world_size
groups_per_rank = n_groups // world_size
# Distribute groups first, derive heads from groups
heads_per_group = num_heads // n_groups
group_sizes = compute_shard_sizes(n_groups, world_size)
head_sizes = [g * heads_per_group for g in group_sizes]
# Per-rank sizes from uneven distribution
groups_per_rank = group_sizes[rank]
heads_per_rank = head_sizes[rank]
is_per_rank = heads_per_rank * head_dim
bc_per_rank = groups_per_rank * ssm_state_size
# Cumulative offsets
is_offset = sum(head_sizes[:rank]) * head_dim
bc_offset = sum(group_sizes[:rank]) * ssm_state_size
head_offset = sum(head_sizes[:rank])
# === in_proj: output layout is [gate:IS | conv_ssm:IS | B:NG*SS | C:NG*SS | dt:NH] ===
gate_start = 0
conv_ssm_start = intermediate_size
@@ -1337,38 +1378,38 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
# Build index tensor for this rank's slice of each section
gate_idx = mx.arange(
gate_start + rank * is_per_rank, gate_start + (rank + 1) * is_per_rank
gate_start + is_offset, gate_start + is_offset + is_per_rank
)
conv_ssm_idx = mx.arange(
conv_ssm_start + rank * is_per_rank,
conv_ssm_start + (rank + 1) * is_per_rank,
conv_ssm_start + is_offset,
conv_ssm_start + is_offset + is_per_rank,
)
b_idx = mx.arange(
b_start + rank * bc_per_rank, b_start + (rank + 1) * bc_per_rank
b_start + bc_offset, b_start + bc_offset + bc_per_rank
)
c_idx = mx.arange(
c_start + rank * bc_per_rank, c_start + (rank + 1) * bc_per_rank
c_start + bc_offset, c_start + bc_offset + bc_per_rank
)
dt_idx = mx.arange(
dt_start + rank * heads_per_rank, dt_start + (rank + 1) * heads_per_rank
dt_start + head_offset, dt_start + head_offset + heads_per_rank
)
indices = mx.concatenate([gate_idx, conv_ssm_idx, b_idx, c_idx, dt_idx])
mixer.in_proj.weight = mixer.in_proj.weight[indices]
# === out_proj: input is intermediate_size (sharded) → hidden_size (reduce) ===
mixer.out_proj = self.sharded_to_all_linear(mixer.out_proj)
mixer.out_proj = self.sharded_to_all_linear(mixer.out_proj, unit=heads_per_group * head_dim)
# === conv1d: depthwise conv on conv_dim channels ===
# conv_dim layout: [ssm_hidden:IS | B:NG*SS | C:NG*SS]
conv_ssm_idx_local = mx.arange(rank * is_per_rank, (rank + 1) * is_per_rank)
conv_ssm_idx_local = mx.arange(is_offset, is_offset + is_per_rank)
conv_b_idx = mx.arange(
intermediate_size + rank * bc_per_rank,
intermediate_size + (rank + 1) * bc_per_rank,
intermediate_size + bc_offset,
intermediate_size + bc_offset + bc_per_rank,
)
conv_c_idx = mx.arange(
intermediate_size + n_groups * ssm_state_size + rank * bc_per_rank,
intermediate_size + n_groups * ssm_state_size + (rank + 1) * bc_per_rank,
intermediate_size + n_groups * ssm_state_size + bc_offset,
intermediate_size + n_groups * ssm_state_size + bc_offset + bc_per_rank,
)
conv_indices = mx.concatenate([conv_ssm_idx_local, conv_b_idx, conv_c_idx])
mixer.conv1d.weight = mixer.conv1d.weight[conv_indices]
@@ -1378,16 +1419,15 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
mixer.conv1d.bias = mixer.conv1d.bias[conv_indices]
# === Per-head parameters ===
h_start = rank * heads_per_rank
h_start = head_offset
h_end = h_start + heads_per_rank
mixer.dt_bias = mixer.dt_bias[h_start:h_end]
mixer.A_log = mixer.A_log[h_start:h_end]
mixer.D = mixer.D[h_start:h_end]
# === Norm: weight is intermediate_size ===
mixer.norm.weight = mixer.norm.weight[
rank * is_per_rank : (rank + 1) * is_per_rank
]
mixer.norm.weight = mixer.norm.weight[is_offset : is_offset + is_per_rank]
mixer.norm.group_size = is_per_rank // groups_per_rank
# === Update dimensions ===
mixer.num_heads = heads_per_rank
src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py
+517
View File
@@ -0,0 +1,517 @@
# type: ignore
import importlib
import itertools
import json
import multiprocessing as mp
import os
import tempfile
import traceback
import mlx.core as mx
import mlx.nn as nn
import numpy as np
import pytest
from mlx.nn.layers.distributed import compute_shard_sizes
RANDOM_SEED = 42
INPUT_TOKENS = [1, 100, 200, 300]
REDUCED_CONFIGS = {
"llama": {
"model_type": "llama",
"num_hidden_layers": 2,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"vocab_size": 1024,
"intermediate_size": 512,
"rms_norm_eps": 1e-5,
"rope_theta": 10000.0,
"tie_word_embeddings": True,
},
"gpt_oss": {
"model_type": "gpt_oss",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"head_dim": 64,
"vocab_size": 1024,
"intermediate_size": 256,
"num_local_experts": 4,
"num_experts_per_tok": 2,
"sliding_window": 64,
"rope_theta": 10000.0,
"tie_word_embeddings": True,
"rms_norm_eps": 1e-5,
},
"deepseek_v3": {
"model_type": "deepseek_v3",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"vocab_size": 1024,
"intermediate_size": 512,
"moe_intermediate_size": 128,
"n_routed_experts": 4,
"n_shared_experts": 1,
"num_experts_per_tok": 2,
"moe_layer_freq": 2,
"qk_rope_head_dim": 32,
"qk_nope_head_dim": 32,
"v_head_dim": 64,
"q_lora_rank": None,
"kv_lora_rank": 64,
"rope_theta": 10000.0,
"rms_norm_eps": 1e-5,
"tie_word_embeddings": True,
},
"step3p5": {
"model_type": "step3p5",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 8,
"num_attention_groups": 4,
"head_dim": 32,
"vocab_size": 1024,
"intermediate_size": 512,
"moe_num_experts": 4,
"moe_top_k": 2,
"moe_intermediate_size": 128,
"share_expert_dim": 128,
"sliding_window": 64,
"rms_norm_eps": 1e-5,
"rope_theta": 10000.0,
"tie_word_embeddings": True,
},
"minimax": {
"model_type": "minimax",
"num_hidden_layers": 2,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"vocab_size": 1024,
"intermediate_size": 512,
"num_local_experts": 4,
"num_experts_per_tok": 2,
"shared_intermediate_size": 256,
"max_position_embeddings": 1024,
"rms_norm_eps": 1e-5,
"rope_theta": 10000.0,
"rotary_dim": 32,
# QK norm's all_gather pattern requires equal shard sizes across ranks,
# incompatible with uneven tp — needs separate fix for padded all_gather
"use_qk_norm": False,
},
"qwen3_moe": {
"model_type": "qwen3_moe",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"head_dim": 32,
"vocab_size": 1024,
"intermediate_size": 512,
"num_experts": 4,
"num_experts_per_tok": 2,
"decoder_sparse_step": 2,
"mlp_only_layers": [],
"moe_intermediate_size": 128,
"rms_norm_eps": 1e-5,
"rope_theta": 10000.0,
"tie_word_embeddings": True,
"max_position_embeddings": 1024,
"norm_topk_prob": True,
},
"qwen3_5": {
"model_type": "qwen3_5",
"text_config": {
"model_type": "qwen3_5",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"vocab_size": 1024,
"intermediate_size": 512,
"num_experts": 4,
"num_experts_per_tok": 2,
"shared_expert_intermediate_size": 256,
"moe_intermediate_size": 128,
"linear_num_key_heads": 4,
"linear_num_value_heads": 4,
"linear_key_head_dim": 32,
"linear_value_head_dim": 32,
"linear_conv_kernel_dim": 4,
"full_attention_interval": 2,
"rms_norm_eps": 1e-5,
"rope_theta": 10000.0,
"tie_word_embeddings": True,
"max_position_embeddings": 1024,
"head_dim": 32,
},
},
"glm4_moe": {
"model_type": "glm4_moe",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"head_dim": 32,
"vocab_size": 1024,
"intermediate_size": 512,
"moe_intermediate_size": 128,
"n_routed_experts": 4,
"n_shared_experts": 1,
"n_group": 1,
"topk_group": 1,
"num_experts_per_tok": 2,
"first_k_dense_replace": 1,
"routed_scaling_factor": 1.0,
"max_position_embeddings": 1024,
"rms_norm_eps": 1e-5,
"rope_theta": 10000.0,
"rope_scaling": None,
"use_qk_norm": False,
"tie_word_embeddings": True,
"attention_bias": False,
"partial_rotary_factor": 1.0,
"norm_topk_prob": True,
},
"nemotron_h": {
"model_type": "nemotron_h",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 8,
"num_key_value_heads": 4,
"vocab_size": 1024,
"intermediate_size": 512,
"max_position_embeddings": 1024,
"attention_bias": False,
"mamba_num_heads": 8,
"mamba_head_dim": 32,
"mamba_proj_bias": False,
"ssm_state_size": 16,
"conv_kernel": 4,
"n_groups": 4,
"mlp_bias": False,
"layer_norm_epsilon": 1e-5,
"use_bias": False,
"use_conv_bias": True,
"head_dim": 32,
"hybrid_override_pattern": "M-*E",
"n_routed_experts": 4,
"num_experts_per_tok": 2,
"moe_intermediate_size": 128,
"n_group": 1,
"topk_group": 1,
"norm_topk_prob": True,
"routed_scaling_factor": 1.0,
},
"glm4_moe_lite": {
"model_type": "glm4_moe_lite",
"num_hidden_layers": 4,
"hidden_size": 256,
"num_attention_heads": 4,
"num_key_value_heads": 4,
"vocab_size": 1024,
"intermediate_size": 512,
"moe_intermediate_size": 128,
"n_routed_experts": 4,
"n_shared_experts": 1,
"num_experts_per_tok": 2,
"moe_layer_freq": 2,
"qk_rope_head_dim": 32,
"qk_nope_head_dim": 32,
"v_head_dim": 64,
"q_lora_rank": None,
"kv_lora_rank": 64,
"rope_theta": 10000.0,
"rms_norm_eps": 1e-5,
"tie_word_embeddings": True,
},
}
def _build_model(config):
mx.random.seed(RANDOM_SEED)
mod = importlib.import_module(f"mlx_lm.models.{config['model_type']}")
args = mod.ModelArgs.from_dict(config)
model = mod.Model(args)
mx.eval(model.parameters())
return model
def _forward(model, tokens):
x = mx.array([tokens])
logits = model(x)
mx.eval(logits)
return np.array(logits[0, -1, :])
def _create_hostfile(world_size, base_port):
hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
f = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
json.dump(hosts, f)
f.close()
return f.name
def _run_single_device(config, result_queue):
try:
model = _build_model(config)
logits = _forward(model, INPUT_TOKENS)
result_queue.put((0, True, logits))
except Exception as e:
result_queue.put((0, False, f"{e}\n{traceback.format_exc()}"))
def _run_tensor_device(rank, world_size, hostfile_path, config, result_queue):
os.environ["MLX_HOSTFILE"] = hostfile_path
os.environ["MLX_RANK"] = str(rank)
try:
group = mx.distributed.init(backend="ring", strict=True)
model = _build_model(config)
from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
model = tensor_auto_parallel(
model, group, timeout_seconds=60.0, on_timeout=None, on_layer_loaded=None
)
logits = _forward(model, INPUT_TOKENS)
result_queue.put((rank, True, logits))
except Exception as e:
result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))
def _run_single(config):
ctx = mp.get_context("spawn")
result_queue = ctx.Queue()
p = ctx.Process(target=_run_single_device, args=(config, result_queue))
p.start()
p.join(timeout=60)
if p.is_alive():
p.terminate()
p.join(timeout=5)
raise TimeoutError("Single device timed out")
rank, success, value = result_queue.get()
assert success, f"Single device failed: {value}"
return value
def _run_tensor(config, world_size, base_port):
ctx = mp.get_context("spawn")
hostfile_path = _create_hostfile(world_size, base_port)
try:
result_queue = ctx.Queue()
processes = []
for rank in range(world_size):
p = ctx.Process(
target=_run_tensor_device,
args=(rank, world_size, hostfile_path, config, result_queue),
)
p.start()
processes.append(p)
for p in processes:
p.join(timeout=120)
timed_out = any(p.is_alive() for p in processes)
for p in processes:
if p.is_alive():
p.terminate()
p.join(timeout=5)
assert not timed_out, "Tensor parallel timed out"
results = {}
while not result_queue.empty():
rank, success, value = result_queue.get()
results[rank] = (success, value)
assert len(results) == world_size, f"Missing results: got {list(results.keys())}"
for rank, (success, value) in results.items():
assert success, f"Rank {rank} failed: {value}"
return results[0][1]
finally:
os.unlink(hostfile_path)
class TestComputeShardSizes:
def test_even_division(self):
assert compute_shard_sizes(64, 2) == [32, 32]
assert compute_shard_sizes(64, 4) == [16, 16, 16, 16]
def test_uneven_division(self):
assert compute_shard_sizes(8, 3) == [3, 3, 2]
assert compute_shard_sizes(64, 3) == [22, 21, 21]
assert compute_shard_sizes(10, 3) == [4, 3, 3]
def test_sum_invariant(self):
for total in [7, 8, 64, 100, 255, 2880]:
for n in [2, 3, 5, 7]:
sizes = compute_shard_sizes(total, n)
assert sum(sizes) == total, f"sum({sizes}) != {total}"
class TestWeightSplitMath:
def test_all_to_sharded_unquantized(self):
mx.random.seed(RANDOM_SEED)
weight = mx.random.normal((64, 256))
x = mx.random.normal((1, 4, 256))
full_output = x @ weight.T
mx.eval(full_output)
for n in [2, 3, 5, 7]:
sizes = compute_shard_sizes(64, n)
indices = list(itertools.accumulate(sizes[:-1]))
shards = mx.split(weight, indices, axis=0)
reconstructed = mx.concatenate([x @ s.T for s in shards], axis=-1)
mx.eval(reconstructed)
diff = float(mx.max(mx.abs(full_output - reconstructed)))
assert diff < 1e-5, f"all-to-sharded N={n}: diff={diff}"
def test_sharded_to_all_unquantized(self):
mx.random.seed(RANDOM_SEED)
weight = mx.random.normal((128, 256))
x = mx.random.normal((1, 4, 256))
full_output = x @ weight.T
mx.eval(full_output)
for n in [2, 3, 5, 7]:
sizes = compute_shard_sizes(256, n)
w_indices = list(itertools.accumulate(sizes[:-1]))
x_indices = list(itertools.accumulate(sizes[:-1]))
w_shards = mx.split(weight, w_indices, axis=-1)
x_shards = mx.split(x, x_indices, axis=-1)
partial_outputs = [xs @ ws.T for xs, ws in zip(x_shards, w_shards)]
reconstructed = sum(partial_outputs)
mx.eval(reconstructed)
diff = float(mx.max(mx.abs(full_output - reconstructed)))
assert diff < 5e-5, f"sharded-to-all N={n}: diff={diff}"
def test_all_to_sharded_quantized(self):
mx.random.seed(RANDOM_SEED)
weight = mx.random.normal((64, 256))
group_size = 32
bits = 4
qw, scales, biases = mx.quantize(weight, group_size=group_size, bits=bits)
x = mx.random.normal((1, 4, 256))
full_output = mx.quantized_matmul(
x, qw, scales=scales, biases=biases, transpose=True, group_size=group_size, bits=bits
)
mx.eval(full_output)
for n in [2, 3]:
sizes = compute_shard_sizes(64, n)
indices = list(itertools.accumulate(sizes[:-1]))
qw_shards = mx.split(qw, indices, axis=0)
scales_shards = mx.split(scales, indices, axis=0)
biases_shards = mx.split(biases, indices, axis=0)
partial = [
mx.quantized_matmul(
x, qw_s, scales=sc_s, biases=bi_s,
transpose=True, group_size=group_size, bits=bits,
)
for qw_s, sc_s, bi_s in zip(qw_shards, scales_shards, biases_shards)
]
reconstructed = mx.concatenate(partial, axis=-1)
mx.eval(reconstructed)
diff = float(mx.max(mx.abs(full_output - reconstructed)))
assert diff < 1e-5, f"quantized all-to-sharded N={n}: diff={diff}"
def test_sharded_to_all_quantized(self):
mx.random.seed(RANDOM_SEED)
weight = mx.random.normal((128, 256))
group_size = 32
bits = 4
qw, scales, biases = mx.quantize(weight, group_size=group_size, bits=bits)
x = mx.random.normal((1, 4, 256))
full_output = mx.quantized_matmul(
x, qw, scales=scales, biases=biases, transpose=True, group_size=group_size, bits=bits
)
mx.eval(full_output)
pack_factor = 32 // bits
num_quant_groups = scales.shape[-1]
for n in [2]:
# Split in quantization-group space (same as _shard_quantized_s2a)
group_counts = compute_shard_sizes(num_quant_groups, n)
weight_ppg = group_size * bits // 32
packed_sizes = [gc * weight_ppg for gc in group_counts]
packed_indices = list(itertools.accumulate(packed_sizes[:-1]))
qw_shards = mx.split(qw, packed_indices, axis=-1)
scale_indices = list(itertools.accumulate(group_counts[:-1]))
scales_shards = mx.split(scales, scale_indices, axis=-1)
biases_shards = mx.split(biases, scale_indices, axis=-1)
logical_sizes = [gc * group_size for gc in group_counts]
x_indices = list(itertools.accumulate(logical_sizes[:-1]))
x_shards = mx.split(x, x_indices, axis=-1)
partial = [
mx.quantized_matmul(
xs, qw_s, scales=sc_s, biases=bi_s,
transpose=True, group_size=group_size, bits=bits,
)
for xs, qw_s, sc_s, bi_s in zip(x_shards, qw_shards, scales_shards, biases_shards)
]
reconstructed = sum(partial)
mx.eval(reconstructed)
diff = float(mx.max(mx.abs(full_output - reconstructed)))
assert diff < 1e-4, f"quantized sharded-to-all N={n}: diff={diff}"
# Port allocation: 31200-31999 (non-colliding with conftest 29600-29800 and qwen35 29950-31100)
_BASE_PORT = 31200
_port_counter = 0
def _next_port_block():
global _port_counter
port = _BASE_PORT + _port_counter * 100
_port_counter += 1
return port
@pytest.mark.slow
class TestTensorParallelTP2:
@pytest.mark.parametrize("model_name", list(REDUCED_CONFIGS.keys()))
def test_tp2_matches_single(self, model_name):
config = REDUCED_CONFIGS[model_name]
single_logits = _run_single(config)
tp2_logits = _run_tensor(config, world_size=2, base_port=_next_port_block())
diff = float(np.max(np.abs(single_logits - tp2_logits)))
assert diff < 3e-6, f"{model_name} tp=2 logit diff: {diff}"
@pytest.mark.slow
class TestTensorParallelTP3:
@pytest.mark.parametrize("model_name", list(REDUCED_CONFIGS.keys()))
def test_tp3_matches_single(self, model_name):
config = REDUCED_CONFIGS[model_name]
single_logits = _run_single(config)
tp3_logits = _run_tensor(config, world_size=3, base_port=_next_port_block())
diff = float(np.max(np.abs(single_logits - tp3_logits)))
assert diff < 3e-6, f"{model_name} tp=3 logit diff: {diff}"
uv.lock
Generated
+7 -7
View File
@@ -558,7 +558,7 @@ dependencies = [
{ name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
{ name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -597,7 +597,7 @@ requires-dist = [
{ name = "hypercorn", specifier = ">=0.18.0" },
{ name = "loguru", specifier = ">=0.7.3" },
{ name = "mflux", specifier = "==0.17.2" },
{ name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
{ name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding" },
{ name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
{ name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak" },
{ name = "mlx-vlm", specifier = ">=0.3.11" },
@@ -1436,7 +1436,7 @@ dependencies = [
{ name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
{ name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1494,8 +1494,8 @@ cuda13 = [
[[package]]
name = "mlx"
version = "0.31.2.dev20260324+e5e64331"
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }
version = "0.31.2.dev20260331+71bcd7a2"
source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }
resolution-markers = [
"python_full_version >= '3.14' and sys_platform == 'darwin'",
"python_full_version < '3.14' and sys_platform == 'darwin'",
@@ -1531,7 +1531,7 @@ version = "0.31.2"
source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#d36e9b661e55a5fc0f77fb6f17ea643aa2dc87aa" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1548,7 +1548,7 @@ dependencies = [
{ name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" },
{ name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" },
{ name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },