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

Add advanced dashboard elements for tensor sharding and update exo bench

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-03-31 19:22:17 +01:00
parent ae7ba5d054
commit b30ee156c4
14 changed files with 456 additions and 137 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.mlx_typings/mlx/nn/layers/distributed.pyi
+10
View File
@@ -9,6 +9,12 @@ import mlx.core as mx
from base import Module
from mlx.nn.layers.linear import Linear
def compute_shard_sizes(
dim: int,
N: int,
unit: int = ...,
weights: list[float] | None = ...,
) -> list[int]: ...
@lru_cache
def sum_gradients(
group: mx.distributed.Group,
@@ -20,6 +26,8 @@ def shard_inplace(
*,
segments: Union[int, list[int]] = ...,
group: Optional[mx.distributed.Group] = ...,
unit: int = ...,
weights: list[float] | None = ...,
) -> None:
"""Shard a module in-place by updating its parameter dictionary with the
sharded parameter dictionary.
@@ -51,6 +59,8 @@ def shard_linear(
*,
segments: Union[int, list[int]] = ...,
group: Optional[mx.distributed.Group] = ...,
unit: int = ...,
weights: list[float] | None = ...,
) -> Linear:
"""Create a new linear layer that has its parameters sharded and also
performs distributed communication either in the forward or backward
.mlx_typings/mlx_lm/models/switch_layers.pyi
+5
View File
@@ -8,6 +8,9 @@ import mlx.core as mx
import mlx.nn as nn
class QuantizedSwitchLinear(nn.Module):
weight: mx.array
scales: mx.array
def __init__(
self,
input_dims: int,
@@ -31,6 +34,8 @@ class QuantizedSwitchLinear(nn.Module):
...
class SwitchLinear(nn.Module):
weight: mx.array
def __init__(
self, input_dims: int, output_dims: int, num_experts: int, bias: bool = ...
) -> None: ...
bench/exo_bench.py
+45 -3
View File
@@ -378,8 +378,21 @@ def main() -> int:
default=1.0,
help="System metrics polling interval in seconds (default: 1.0).",
)
ap.add_argument(
"--tensor-strategies",
nargs="+",
default=["Naive"],
help="Tensor shard strategies to benchmark. Choices: Naive, Memory, Compute, Bandwidth, all. Default: Naive.",
)
args = ap.parse_args()
tensor_strategies = []
for s in args.tensor_strategies:
if s.lower() == "all":
tensor_strategies = ["Naive", "Memory", "Compute", "Bandwidth"]
break
tensor_strategies.append(s)
pp_list = parse_int_list(args.pp)
tg_list = parse_int_list(args.tg)
if not pp_list or not tg_list:
@@ -467,20 +480,45 @@ def main() -> int:
all_rows: list[dict[str, Any]] = []
all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
placement_runs: list[tuple[dict[str, Any], str]] = []
for preview in selected:
sharding = str(preview["sharding"])
if sharding == "Tensor":
for ts in tensor_strategies:
placement_runs.append((preview, ts))
else:
placement_runs.append((preview, "Naive"))
for preview, tensor_strategy in placement_runs:
instance = preview["instance"]
instance_id = instance_id_from_instance(instance)
sharding = str(preview["sharding"])
instance_meta = str(preview["instance_meta"])
n_nodes = nodes_used_in_instance(instance)
strategy_label = (
f" / strategy={tensor_strategy}" if sharding == "Tensor" else ""
)
logger.info("=" * 80)
logger.info(
f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes} / instance_id={instance_id}"
f"PLACEMENT: {sharding} / {instance_meta} / nodes={n_nodes}{strategy_label} / instance_id={instance_id}"
)
client.request_json("POST", "/instance", body={"instance": instance})
if sharding == "Tensor" and tensor_strategy != "Naive":
client.request_json(
"POST",
"/place_instance",
body={
"model_id": full_model_id,
"sharding": sharding,
"instance_meta": instance_meta,
"min_nodes": n_nodes,
"tensor_strategy": tensor_strategy,
},
)
else:
client.request_json("POST", "/instance", body={"instance": instance})
try:
wait_for_instance_ready(client, instance_id)
except (RuntimeError, TimeoutError) as e:
@@ -541,6 +579,7 @@ def main() -> int:
"placement_sharding": sharding,
"placement_instance_meta": instance_meta,
"placement_nodes": n_nodes,
"tensor_strategy": tensor_strategy,
"instance_id": instance_id,
"pp_tokens": actual_pp_tokens,
"tg": tg,
@@ -595,6 +634,7 @@ def main() -> int:
"placement_sharding": sharding,
"placement_instance_meta": instance_meta,
"placement_nodes": n_nodes,
"tensor_strategy": tensor_strategy,
"instance_id": instance_id,
"pp_tokens": actual_pp_tokens,
"tg": tg,
@@ -656,7 +696,9 @@ def main() -> int:
finally:
if sampler:
sampler.stop()
placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
placement_label = (
f"{sharding}/{instance_meta}/{n_nodes} nodes{strategy_label}"
)
sampler.print_summary(placement_label)
placement_metrics = sampler.summarize()
if placement_metrics:
dashboard/src/routes/+page.svelte
+59
View File
@@ -886,6 +886,8 @@
}
let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
let selectedTensorStrategy = $state<string>("Naive");
let showAdvancedSettings = $state(false);
type InstanceMeta = "MlxRing" | "MlxJaccl";
// Launch defaults persistence
@@ -1446,6 +1448,8 @@
sharding: selectedSharding,
instance_meta: selectedInstanceType,
min_nodes: 1,
tensor_strategy:
selectedSharding === "Tensor" ? selectedTensorStrategy : "Naive",
}),
});
}
@@ -5874,6 +5878,61 @@
</div>
</div>
</div>
<!-- Advanced Settings -->
{#if selectedSharding === "Tensor"}
<div class="border-t border-exo-medium-gray/30 pt-3 mt-3">
<button
onclick={() =>
(showAdvancedSettings = !showAdvancedSettings)}
class="flex items-center gap-2 text-xs text-white/40 font-mono hover:text-white/60 transition-colors cursor-pointer"
>
<span
class="transition-transform duration-200 {showAdvancedSettings
? 'rotate-90'
: ''}">▶</span
>
Advanced
</button>
{#if showAdvancedSettings}
<div class="mt-3 space-y-3">
<div>
<div class="text-xs text-white/50 font-mono mb-2">
Tensor Strategy:
</div>
<div class="flex gap-2 flex-wrap">
{#each ["Naive", "Memory", "Compute", "Bandwidth"] as strategy}
<button
onclick={() => {
selectedTensorStrategy = strategy;
saveLaunchDefaults();
}}
class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedTensorStrategy ===
strategy
? 'bg-transparent text-exo-yellow border-exo-yellow'
: 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
>
<span
class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedTensorStrategy ===
strategy
? 'border-exo-yellow'
: 'border-exo-medium-gray'}"
>
{#if selectedTensorStrategy === strategy}
<span
class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
></span>
{/if}
</span>
{strategy}
</button>
{/each}
</div>
</div>
</div>
{/if}
</div>
{/if}
{/if}
</div>
src/exo/api/main.py
+1
View File
@@ -383,6 +383,7 @@ class API:
sharding=payload.sharding,
instance_meta=payload.instance_meta,
min_nodes=payload.min_nodes,
tensor_strategy=payload.tensor_strategy,
)
await self._send(command)
src/exo/api/types/api.py
+2 -1
View File
@@ -10,7 +10,7 @@ from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.text_generation import ReasoningEffort
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
from exo.shared.types.worker.shards import Sharding, ShardMetadata, TensorShardStrategy
from exo.utils.pydantic_ext import CamelCaseModel
FinishReason = Literal[
@@ -253,6 +253,7 @@ class PlaceInstanceParams(BaseModel):
sharding: Sharding = Sharding.Pipeline
instance_meta: InstanceMeta = InstanceMeta.MlxRing
min_nodes: int = 1
tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive
class CreateInstanceParams(BaseModel):
src/exo/master/main.py
+1
View File
@@ -298,6 +298,7 @@ class Master:
self.state.instances,
self.state.node_memory,
self.state.node_network,
node_identities=self.state.node_identities,
download_status=self.state.downloads,
)
transition_events = get_transition_events(
src/exo/master/placement.py
+8 -2
View File
@@ -29,7 +29,7 @@ from exo.shared.types.events import (
TaskStatusUpdated,
)
from exo.shared.types.memory import Memory
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo
from exo.shared.types.tasks import Task, TaskId, TaskStatus
from exo.shared.types.worker.downloads import (
DownloadCompleted,
@@ -108,6 +108,7 @@ def place_instance(
current_instances: Mapping[InstanceId, Instance],
node_memory: Mapping[NodeId, MemoryUsage],
node_network: Mapping[NodeId, NodeNetworkInfo],
node_identities: Mapping[NodeId, NodeIdentity] | None = None,
required_nodes: set[NodeId] | None = None,
download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
) -> dict[InstanceId, Instance]:
@@ -197,7 +198,12 @@ def place_instance(
command.sharding = Sharding.Pipeline
shard_assignments = get_shard_assignments(
command.model_card, selected_cycle, command.sharding, node_memory
command.model_card,
selected_cycle,
command.sharding,
node_memory,
tensor_strategy=command.tensor_strategy,
node_identities=node_identities,
)
cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids)
src/exo/master/placement_utils.py
+81 -1
View File
@@ -6,7 +6,7 @@ from exo.shared.models.model_cards import ModelCard
from exo.shared.topology import Topology
from exo.shared.types.common import Host, NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.profiling import MemoryUsage, NodeIdentity, NodeNetworkInfo
from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
from exo.shared.types.worker.runners import RunnerId, ShardAssignments
from exo.shared.types.worker.shards import (
@@ -15,8 +15,54 @@ from exo.shared.types.worker.shards import (
Sharding,
ShardMetadata,
TensorShardMetadata,
TensorShardMode,
TensorShardStrategy,
)
# FP32 TFLOPS by chip_id (used for Compute strategy)
APPLE_SILICON_FLOPS: dict[str, float] = {
"Apple M1": 2.6,
"Apple M1 Pro": 5.3,
"Apple M1 Max": 10.4,
"Apple M1 Ultra": 21.0,
"Apple M2": 3.6,
"Apple M2 Pro": 6.8,
"Apple M2 Max": 13.6,
"Apple M2 Ultra": 27.2,
"Apple M3": 3.5,
"Apple M3 Pro": 5.0,
"Apple M3 Max": 14.2,
"Apple M3 Ultra": 28.4,
"Apple M4": 4.3,
"Apple M4 Pro": 9.2,
"Apple M4 Max": 18.4,
"Apple M5": 4.2,
"Apple M5 Pro": 8.3,
"Apple M5 Max": 19.9,
}
# Memory bandwidth in GB/s by chip_id (used for Bandwidth strategy)
APPLE_SILICON_BANDWIDTH: dict[str, float] = {
"Apple M1": 68,
"Apple M1 Pro": 200,
"Apple M1 Max": 400,
"Apple M1 Ultra": 800,
"Apple M2": 100,
"Apple M2 Pro": 200,
"Apple M2 Max": 400,
"Apple M2 Ultra": 800,
"Apple M3": 100,
"Apple M3 Pro": 150,
"Apple M3 Max": 400,
"Apple M3 Ultra": 800,
"Apple M4": 120,
"Apple M4 Pro": 273,
"Apple M4 Max": 546,
"Apple M5": 154,
"Apple M5 Pro": 307,
"Apple M5 Max": 614,
}
def filter_cycles_by_memory(
cycles: list[Cycle],
@@ -243,12 +289,39 @@ def _get_shard_assignments_for_pure_pipeline(
def get_shard_assignments_for_tensor_parallel(
model_card: ModelCard,
cycle: Cycle,
node_memory: Mapping[NodeId, MemoryUsage] | None = None,
strategy: TensorShardStrategy = TensorShardStrategy.Naive,
node_identities: Mapping[NodeId, NodeIdentity] | None = None,
):
total_layers = model_card.n_layers
world_size = len(cycle)
runner_to_shard: dict[RunnerId, ShardMetadata] = {}
node_to_runner: dict[NodeId, RunnerId] = {}
shard_weights: list[float] | None = None
shard_mode = TensorShardMode.Constant
match strategy:
case TensorShardStrategy.Naive:
pass
case TensorShardStrategy.Memory:
if node_memory is not None:
shard_weights = [
node_memory[node_id].ram_available.in_gb for node_id in cycle
]
shard_mode = TensorShardMode.Greedy
case TensorShardStrategy.Compute:
if node_identities is not None:
shard_weights = [
APPLE_SILICON_FLOPS.get(node_identities[nid].chip_id, 1.0)
for nid in cycle
]
case TensorShardStrategy.Bandwidth:
if node_identities is not None:
shard_weights = [
APPLE_SILICON_BANDWIDTH.get(node_identities[nid].chip_id, 1.0)
for nid in cycle
]
for i, node_id in enumerate(cycle):
shard = TensorShardMetadata(
model_card=model_card,
@@ -257,6 +330,8 @@ def get_shard_assignments_for_tensor_parallel(
start_layer=0,
end_layer=total_layers,
n_layers=total_layers,
shard_weights=shard_weights,
shard_mode=shard_mode,
)
runner_id = RunnerId()
@@ -278,6 +353,8 @@ def get_shard_assignments(
cycle: Cycle,
sharding: Sharding,
node_memory: Mapping[NodeId, MemoryUsage],
tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive,
node_identities: Mapping[NodeId, NodeIdentity] | None = None,
) -> ShardAssignments:
match sharding:
case Sharding.Pipeline:
@@ -290,6 +367,9 @@ def get_shard_assignments(
return get_shard_assignments_for_tensor_parallel(
model_card=model_card,
cycle=cycle,
node_memory=node_memory,
strategy=tensor_strategy,
node_identities=node_identities,
)
src/exo/shared/types/commands.py
+2 -1
View File
@@ -9,7 +9,7 @@ from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.common import CommandId, NodeId, SystemId
from exo.shared.types.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding, ShardMetadata
from exo.shared.types.worker.shards import Sharding, ShardMetadata, TensorShardStrategy
from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -38,6 +38,7 @@ class PlaceInstance(BaseCommand):
sharding: Sharding
instance_meta: InstanceMeta
min_nodes: int
tensor_strategy: TensorShardStrategy = TensorShardStrategy.Naive
class CreateInstance(BaseCommand):
src/exo/shared/types/worker/shards.py
+7
View File
@@ -17,6 +17,13 @@ class TensorShardMode(str, Enum):
Constant = "Constant"
class TensorShardStrategy(str, Enum):
Naive = "Naive"
Memory = "Memory"
Compute = "Compute"
Bandwidth = "Bandwidth"
class BaseShardMetadata(TaggedModel):
"""
Defines a specific shard of the model that is ready to be run on a device.
src/exo/worker/engines/mlx/auto_parallel.py
+228 -127
View File
@@ -14,7 +14,6 @@ from mlx.nn.layers.distributed import (
shard_linear,
sum_gradients,
)
from mlx.utils import tree_flatten
from mlx_lm.models.base import (
scaled_dot_product_attention, # pyright: ignore[reportUnknownVariableType]
)
@@ -497,9 +496,9 @@ def tensor_auto_parallel(
on_timeout: TimeoutCallback | None,
on_layer_loaded: LayerLoadedCallback | None,
shard_weights: list[float] | None = None,
shard_mode: "TensorShardMode | None" = None,
shard_mode: TensorShardMode | None = None,
) -> nn.Module:
shard_mode = shard_mode or "Constant"
resolved_shard_mode = shard_mode or TensorShardMode.Constant
all_to_sharded_linear = partial(
shard_linear,
sharding="all-to-sharded",
@@ -552,7 +551,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(model, (DeepseekV3Model, DeepseekV32Model, KimiK25Model)):
tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
@@ -562,7 +561,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(model, MiniMaxModel):
tensor_parallel_sharding_strategy = MiniMaxShardingStrategy(
@@ -572,7 +571,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(model, GLM4MoeLiteModel):
tensor_parallel_sharding_strategy = GLM4MoeLiteShardingStrategy(
@@ -582,7 +581,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(model, Glm4MoeModel):
tensor_parallel_sharding_strategy = Glm4MoeShardingStrategy(
@@ -592,7 +591,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(
model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
@@ -604,7 +603,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(model, GptOssModel):
tensor_parallel_sharding_strategy = GptOssShardingStrategy(
@@ -614,7 +613,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(model, Step35Model):
tensor_parallel_sharding_strategy = Step35ShardingStrategy(
@@ -624,7 +623,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
elif isinstance(model, NemotronHModel):
tensor_parallel_sharding_strategy = NemotronHShardingStrategy(
@@ -634,7 +633,7 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
shard_weights=shard_weights,
shard_mode=shard_mode,
shard_mode=resolved_shard_mode,
)
else:
raise ValueError(f"Unsupported model type: {type(model)}")
@@ -654,7 +653,7 @@ class TensorParallelShardingStrategy(ABC):
all_to_sharded_linear_in_place: Callable[..., None],
sharded_to_all_linear_in_place: Callable[..., None],
shard_weights: list[float] | None = None,
shard_mode: str = "Constant",
shard_mode: TensorShardMode = TensorShardMode.Constant,
):
self._base_all_to_sharded_linear = all_to_sharded_linear
self._base_sharded_to_all_linear = sharded_to_all_linear
@@ -664,11 +663,13 @@ class TensorParallelShardingStrategy(ABC):
self.shard_mode = shard_mode
self.group = group
self.N = group.size()
self._greedy_trackers: dict[str, list[list[float]]] | None = None
if shard_weights is not None and shard_mode == "Greedy":
self._greedy_trackers: dict[str, list[list[float]] | list[float]] | None = None
if shard_weights is not None and shard_mode == TensorShardMode.Greedy:
self._greedy_trackers = {}
def _greedy_weights_for(self, key: str, dim: int, unit: int = 1) -> list[float] | None:
def _greedy_weights_for(
self, key: str, dim: int, unit: int = 1
) -> list[float] | None:
"""Get adjusted weights for a specific projection type, and record the allocation."""
if self.shard_weights is None or self._greedy_trackers is None:
return self.shard_weights
@@ -677,8 +678,11 @@ class TensorParallelShardingStrategy(ABC):
target = [dim * self.shard_weights[i] / total_w for i in range(n)]
if key not in self._greedy_trackers:
self._greedy_trackers[key] = [[0.0] * n, [0.0] * n, [0] * n]
cum_target, cum_actual, last_sizes = self._greedy_trackers[key]
desired = [target[i] + (cum_target[i] - cum_actual[i]) for i in range(n)]
tracker = cast(list[list[float]], self._greedy_trackers[key])
cum_target, cum_actual, last_sizes = tracker[0], tracker[1], tracker[2]
desired: list[float] = [
target[i] + (cum_target[i] - cum_actual[i]) for i in range(n)
]
min_d = min(desired)
if min_d <= 0:
desired = [d - min_d + 0.01 for d in desired]
@@ -687,21 +691,21 @@ class TensorParallelShardingStrategy(ABC):
cum_target[i] += target[i]
cum_actual[i] += actual_sizes[i]
last_sizes[i] = actual_sizes[i]
self._greedy_trackers[key + "_last_weights"] = desired # type: ignore
self._greedy_trackers[key + "_last_weights"] = desired
return desired
def _greedy_last_sizes(self, key: str) -> list[int]:
"""Get the actual sizes from the last _greedy_weights_for call for this key."""
if self._greedy_trackers is None or key not in self._greedy_trackers:
return []
return self._greedy_trackers[key][2]
tracker = self._greedy_trackers[key]
assert isinstance(tracker, list) and len(tracker) == 3 # noqa: S101
return cast(list[int], tracker[2])
def _greedy_last_weights(self, key: str) -> list[float] | None:
"""Get the weights used in the last _greedy_weights_for call for this key."""
if self._greedy_trackers is None:
return self.shard_weights
w = self._greedy_trackers.get(key + "_last_weights") # type: ignore
return w if w is not None else self.shard_weights
w = self._greedy_trackers.get(key + "_last_weights")
return cast(list[float] | None, w) if w is not None else self.shard_weights
@property
def all_to_sharded_linear(self) -> Callable[..., nn.Linear]:
@@ -750,19 +754,23 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
k_dim = layer.self_attn.k_proj.weight.shape[0]
intermediate = layer.mlp.gate_proj.weight.shape[0]
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=gqa_unit,
layer.self_attn.q_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("q", q_dim, gqa_unit),
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj, unit=head_dim,
layer.self_attn.k_proj,
unit=head_dim,
weights=self._greedy_weights_for("k", k_dim, head_dim),
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj, unit=head_dim,
layer.self_attn.v_proj,
unit=head_dim,
weights=self._greedy_weights_for("v", k_dim, head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=gqa_unit,
layer.self_attn.o_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("o", q_dim, gqa_unit),
)
layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
@@ -843,21 +851,30 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
o_dim = q_dim
if layer.self_attn.q_lora_rank is None:
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=q_head_dim,
layer.self_attn.q_proj,
unit=q_head_dim,
weights=self._greedy_weights_for("q", q_dim, q_head_dim),
)
else:
layer.self_attn.q_b_proj = self.all_to_sharded_linear(
layer.self_attn.q_b_proj, unit=q_head_dim,
layer.self_attn.q_b_proj,
unit=q_head_dim,
weights=self._greedy_weights_for("q", q_dim, q_head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=q_head_dim,
layer.self_attn.o_proj,
unit=q_head_dim,
weights=self._greedy_weights_for("o", o_dim, q_head_dim),
)
q_actual = self._greedy_last_sizes("q")
head_sizes = [s // q_head_dim for s in q_actual] if q_actual else compute_shard_sizes(original_num_heads, self.N, weights=self.shard_weights)
head_sizes = (
[s // q_head_dim for s in q_actual]
if q_actual
else compute_shard_sizes(
original_num_heads, self.N, weights=self.shard_weights
)
)
layer.self_attn.num_heads = head_sizes[self.group.rank()]
# Logic from upstream mlx
@@ -890,23 +907,29 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
else:
if getattr(layer.mlp, "shared_experts", None) is not None:
shared_gate_dim = layer.mlp.shared_experts.gate_proj.weight.shape[0]
shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[-1]
shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[
-1
]
shared_up_dim = layer.mlp.shared_experts.up_proj.weight.shape[0]
self.all_to_sharded_linear_in_place(
layer.mlp.shared_experts.gate_proj,
weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
weights=self._greedy_weights_for(
"shared_gate", shared_gate_dim
),
)
self.sharded_to_all_linear_in_place(
layer.mlp.shared_experts.down_proj,
weights=self._greedy_weights_for("shared_down", shared_down_dim),
weights=self._greedy_weights_for(
"shared_down", shared_down_dim
),
)
self.all_to_sharded_linear_in_place(
layer.mlp.shared_experts.up_proj,
weights=self._greedy_weights_for("shared_up", shared_up_dim),
)
moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
self.all_to_sharded_linear_in_place(
layer.mlp.switch_mlp.gate_proj,
weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -962,35 +985,44 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
timeout_seconds / total,
on_timeout,
)
original_num_heads = layer.self_attn.num_heads # type: ignore
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
if layer.self_attn.q_lora_rank is not None # pyright: ignore[reportUnnecessaryComparison]
else layer.self_attn.q_proj.weight.shape[0] // original_num_heads
) # type: ignore
)
q_dim = (
layer.self_attn.q_proj.weight.shape[0]
if layer.self_attn.q_lora_rank is None
if layer.self_attn.q_lora_rank is None # pyright: ignore[reportUnnecessaryComparison]
else layer.self_attn.q_b_proj.weight.shape[0]
) # type: ignore
)
o_dim = q_dim
if layer.self_attn.q_lora_rank is None: # type: ignore
if layer.self_attn.q_lora_rank is None: # pyright: ignore[reportUnnecessaryComparison]
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=q_head_dim,
layer.self_attn.q_proj,
unit=q_head_dim,
weights=self._greedy_weights_for("q", q_dim, q_head_dim),
)
else:
layer.self_attn.q_b_proj = self.all_to_sharded_linear(
layer.self_attn.q_b_proj, unit=q_head_dim,
layer.self_attn.q_b_proj,
unit=q_head_dim,
weights=self._greedy_weights_for("q", q_dim, q_head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=q_head_dim,
layer.self_attn.o_proj,
unit=q_head_dim,
weights=self._greedy_weights_for("o", o_dim, q_head_dim),
)
q_actual = self._greedy_last_sizes("q")
head_sizes = [s // q_head_dim for s in q_actual] if q_actual else compute_shard_sizes(original_num_heads, self.N, weights=self.shard_weights)
head_sizes = (
[s // q_head_dim for s in q_actual]
if q_actual
else compute_shard_sizes(
original_num_heads, self.N, weights=self.shard_weights
)
)
layer.self_attn.num_heads = head_sizes[self.group.rank()]
# Logic from upstream mlx
@@ -1021,23 +1053,29 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
else:
if getattr(layer.mlp, "shared_experts", None) is not None:
shared_gate_dim = layer.mlp.shared_experts.gate_proj.weight.shape[0]
shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[-1]
shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[
-1
]
shared_up_dim = layer.mlp.shared_experts.up_proj.weight.shape[0]
self.all_to_sharded_linear_in_place(
layer.mlp.shared_experts.gate_proj,
weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
weights=self._greedy_weights_for(
"shared_gate", shared_gate_dim
),
)
self.sharded_to_all_linear_in_place(
layer.mlp.shared_experts.down_proj,
weights=self._greedy_weights_for("shared_down", shared_down_dim),
weights=self._greedy_weights_for(
"shared_down", shared_down_dim
),
)
self.all_to_sharded_linear_in_place(
layer.mlp.shared_experts.up_proj,
weights=self._greedy_weights_for("shared_up", shared_up_dim),
)
moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
self.all_to_sharded_linear_in_place(
layer.mlp.switch_mlp.gate_proj,
weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1158,19 +1196,23 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
q_dim = layer.self_attn.q_proj.weight.shape[0]
k_dim = layer.self_attn.k_proj.weight.shape[0]
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=gqa_unit,
layer.self_attn.q_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("q", q_dim, gqa_unit),
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj, unit=head_dim,
layer.self_attn.k_proj,
unit=head_dim,
weights=self._greedy_weights_for("k", k_dim, head_dim),
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj, unit=head_dim,
layer.self_attn.v_proj,
unit=head_dim,
weights=self._greedy_weights_for("v", k_dim, head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=gqa_unit,
layer.self_attn.o_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("o", q_dim, gqa_unit),
)
@@ -1184,9 +1226,13 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group) # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
# Shard the MoE.
moe_gate_dim = layer.block_sparse_moe.switch_mlp.gate_proj.weight.shape[1]
moe_down_dim = layer.block_sparse_moe.switch_mlp.down_proj.weight.shape[-1]
moe_up_dim = layer.block_sparse_moe.switch_mlp.up_proj.weight.shape[1]
moe_gate_dim = int(
layer.block_sparse_moe.switch_mlp.gate_proj.weight.shape[1]
)
moe_down_dim = int(
layer.block_sparse_moe.switch_mlp.down_proj.weight.shape[-1]
)
moe_up_dim = int(layer.block_sparse_moe.switch_mlp.up_proj.weight.shape[1])
self.all_to_sharded_linear_in_place(
layer.block_sparse_moe.switch_mlp.gate_proj,
weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1232,19 +1278,23 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
q_dim = layer.self_attn.q_proj.weight.shape[0]
k_dim = layer.self_attn.k_proj.weight.shape[0]
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=gqa_unit,
layer.self_attn.q_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("q", q_dim, gqa_unit),
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj, unit=head_dim,
layer.self_attn.k_proj,
unit=head_dim,
weights=self._greedy_weights_for("k", k_dim, head_dim),
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj, unit=head_dim,
layer.self_attn.v_proj,
unit=head_dim,
weights=self._greedy_weights_for("v", k_dim, head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=gqa_unit,
layer.self_attn.o_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("o", q_dim, gqa_unit),
)
layer.self_attn.n_heads = (
@@ -1258,8 +1308,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
if hasattr(layer, "linear_attn"):
linear_attn = layer.linear_attn
k_greedy: list[float] | None = None
v_greedy: list[float] | None = None
if isinstance(linear_attn, Qwen3NextGatedDeltaNet):
# Qwen3-Next: combined projections
qkvz_dim = linear_attn.in_proj_qkvz.weight.shape[0]
ba_dim = linear_attn.in_proj_ba.weight.shape[0]
linear_attn.in_proj_qkvz = self.all_to_sharded_linear(
@@ -1271,19 +1322,18 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
weights=self._greedy_weights_for("linear_ba", ba_dim),
)
else:
# 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
b_dim = linear_attn.in_proj_b.weight.shape[0]
a_dim = linear_attn.in_proj_a.weight.shape[0]
# Compute greedy weights ONCE per dimension — all projections
# sharing the same dim must use the same weights within a layer
k_greedy = self._greedy_weights_for("linear_k_dim", key_dim, head_k_dim)
v_greedy = self._greedy_weights_for("linear_v_dim", value_dim, head_v_dim)
k_greedy = self._greedy_weights_for(
"linear_k_dim", key_dim, head_k_dim
)
v_greedy = self._greedy_weights_for(
"linear_v_dim", value_dim, head_v_dim
)
linear_attn.in_proj_qkv = shard_linear(
linear_attn.in_proj_qkv,
"all-to-sharded",
@@ -1293,7 +1343,8 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
weights=k_greedy,
)
linear_attn.in_proj_z = self.all_to_sharded_linear(
linear_attn.in_proj_z, unit=head_v_dim,
linear_attn.in_proj_z,
unit=head_v_dim,
weights=v_greedy,
)
linear_attn.in_proj_b = self.all_to_sharded_linear(
@@ -1306,22 +1357,26 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
)
is_qwen3next = isinstance(linear_attn, Qwen3NextGatedDeltaNet)
out_dim = linear_attn.out_proj.weight.shape[-1]
out_w = v_greedy if not is_qwen3next else self._greedy_weights_for("linear_out", out_dim, linear_attn.head_v_dim) # pyright: ignore[reportPossiblyUnbound]
out_w = (
v_greedy
if not is_qwen3next
else self._greedy_weights_for(
"linear_out", out_dim, linear_attn.head_v_dim
)
)
linear_attn.out_proj = self.sharded_to_all_linear(
linear_attn.out_proj, unit=linear_attn.head_v_dim,
linear_attn.out_proj,
unit=linear_attn.head_v_dim,
weights=out_w,
)
# Shard conv1d: depthwise conv with non-contiguous channel slicing.
# Channel layout is [q(key_dim), k(key_dim), v(value_dim)].
# Each rank takes its head-slice from each of the three sections.
rank = self.group.rank()
key_dim = linear_attn.key_dim
value_dim = linear_attn.value_dim
head_k_dim = linear_attn.head_k_dim
head_v_dim = linear_attn.head_v_dim
k_w = k_greedy if not is_qwen3next else self.shard_weights # pyright: ignore[reportPossiblyUnbound]
v_w = v_greedy if not is_qwen3next else self.shard_weights # pyright: ignore[reportPossiblyUnbound]
k_w = k_greedy if not is_qwen3next else self.shard_weights
v_w = v_greedy if not is_qwen3next else self.shard_weights
key_shard_sizes = compute_shard_sizes(
key_dim, self.N, unit=head_k_dim, weights=k_w
)
@@ -1379,22 +1434,27 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
)
q_dim = layer.self_attn.q_proj.weight.shape[0]
k_dim = layer.self_attn.k_proj.weight.shape[0]
o_dim = layer.self_attn.o_proj.weight.shape[-1]
qo_greedy = self._greedy_weights_for("qwen_qo", q_dim, kv_head_dim * 2 * gqa_repeat)
qo_greedy = self._greedy_weights_for(
"qwen_qo", q_dim, kv_head_dim * 2 * gqa_repeat
)
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=kv_head_dim * 2 * gqa_repeat,
layer.self_attn.q_proj,
unit=kv_head_dim * 2 * gqa_repeat,
weights=qo_greedy,
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj, unit=kv_head_dim,
layer.self_attn.k_proj,
unit=kv_head_dim,
weights=self._greedy_weights_for("k", k_dim, kv_head_dim),
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj, unit=kv_head_dim,
layer.self_attn.v_proj,
unit=kv_head_dim,
weights=self._greedy_weights_for("v", k_dim, kv_head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=kv_head_dim * gqa_repeat,
layer.self_attn.o_proj,
unit=kv_head_dim * gqa_repeat,
weights=qo_greedy,
)
layer.self_attn.num_attention_heads = (
@@ -1413,9 +1473,9 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
Qwen3_5SparseMoeBlock,
),
):
moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
self.all_to_sharded_linear_in_place(
layer.mlp.switch_mlp.gate_proj,
weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1436,11 +1496,15 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
shared_up_dim = layer.mlp.shared_expert.up_proj.weight.shape[0]
self.all_to_sharded_linear_in_place(
layer.mlp.shared_expert.gate_proj,
weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
weights=self._greedy_weights_for(
"shared_gate", shared_gate_dim
),
)
self.sharded_to_all_linear_in_place(
layer.mlp.shared_expert.down_proj,
weights=self._greedy_weights_for("shared_down", shared_down_dim),
weights=self._greedy_weights_for(
"shared_down", shared_down_dim
),
)
self.all_to_sharded_linear_in_place(
layer.mlp.shared_expert.up_proj,
@@ -1491,19 +1555,23 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
q_dim = layer.self_attn.q_proj.weight.shape[0]
k_dim = layer.self_attn.k_proj.weight.shape[0]
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=gqa_unit,
layer.self_attn.q_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("q", q_dim, gqa_unit),
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj, unit=head_dim,
layer.self_attn.k_proj,
unit=head_dim,
weights=self._greedy_weights_for("k", k_dim, head_dim),
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj, unit=head_dim,
layer.self_attn.v_proj,
unit=head_dim,
weights=self._greedy_weights_for("v", k_dim, head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=gqa_unit,
layer.self_attn.o_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("o", q_dim, gqa_unit),
)
layer.self_attn.n_heads = layer.self_attn.q_proj.weight.shape[0] // head_dim
@@ -1512,9 +1580,9 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
)
if isinstance(layer.mlp, MoE):
moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
self.all_to_sharded_linear_in_place(
layer.mlp.switch_mlp.gate_proj,
weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1529,15 +1597,21 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
)
if getattr(layer.mlp, "shared_experts", None) is not None:
shared_gate_dim = layer.mlp.shared_experts.gate_proj.weight.shape[0]
shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[-1]
shared_down_dim = layer.mlp.shared_experts.down_proj.weight.shape[
-1
]
shared_up_dim = layer.mlp.shared_experts.up_proj.weight.shape[0]
self.all_to_sharded_linear_in_place(
layer.mlp.shared_experts.gate_proj,
weights=self._greedy_weights_for("shared_gate", shared_gate_dim),
weights=self._greedy_weights_for(
"shared_gate", shared_gate_dim
),
)
self.sharded_to_all_linear_in_place(
layer.mlp.shared_experts.down_proj,
weights=self._greedy_weights_for("shared_down", shared_down_dim),
weights=self._greedy_weights_for(
"shared_down", shared_down_dim
),
)
self.all_to_sharded_linear_in_place(
layer.mlp.shared_experts.up_proj,
@@ -1589,19 +1663,23 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
q_dim = layer.self_attn.q_proj.weight.shape[0]
k_dim = layer.self_attn.k_proj.weight.shape[0]
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=gqa_unit,
layer.self_attn.q_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("q", q_dim, gqa_unit),
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj, unit=head_dim,
layer.self_attn.k_proj,
unit=head_dim,
weights=self._greedy_weights_for("k", k_dim, head_dim),
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj, unit=head_dim,
layer.self_attn.v_proj,
unit=head_dim,
weights=self._greedy_weights_for("v", k_dim, head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=gqa_unit,
layer.self_attn.o_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("o", q_dim, gqa_unit),
)
@@ -1618,14 +1696,23 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
rank = self.group.rank()
q_actual = self._greedy_last_sizes("q")
q_head_sizes = [s // head_dim for s in q_actual] if q_actual else compute_shard_sizes(original_num_heads, self.N, unit=gqa_unit // head_dim, weights=self.shard_weights)
q_head_sizes = (
[s // head_dim for s in q_actual]
if q_actual
else compute_shard_sizes(
original_num_heads,
self.N,
unit=gqa_unit // head_dim,
weights=self.shard_weights,
)
)
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]
moe_gate_dim = layer.mlp.experts.gate_proj.weight.shape[1]
moe_down_dim = layer.mlp.experts.down_proj.weight.shape[-1]
moe_up_dim = layer.mlp.experts.up_proj.weight.shape[1]
moe_gate_dim = int(layer.mlp.experts.gate_proj.weight.shape[1])
moe_down_dim = int(layer.mlp.experts.down_proj.weight.shape[-1])
moe_up_dim = int(layer.mlp.experts.up_proj.weight.shape[1])
self.all_to_sharded_linear_in_place(
layer.mlp.experts.gate_proj,
weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1667,19 +1754,23 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
q_dim = layer.self_attn.q_proj.weight.shape[0]
k_dim = layer.self_attn.k_proj.weight.shape[0]
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj, unit=gqa_unit,
layer.self_attn.q_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("q", q_dim, gqa_unit),
)
layer.self_attn.k_proj = self.all_to_sharded_linear(
layer.self_attn.k_proj, unit=head_dim,
layer.self_attn.k_proj,
unit=head_dim,
weights=self._greedy_weights_for("k", k_dim, head_dim),
)
layer.self_attn.v_proj = self.all_to_sharded_linear(
layer.self_attn.v_proj, unit=head_dim,
layer.self_attn.v_proj,
unit=head_dim,
weights=self._greedy_weights_for("v", k_dim, head_dim),
)
layer.self_attn.o_proj = self.sharded_to_all_linear(
layer.self_attn.o_proj, unit=gqa_unit,
layer.self_attn.o_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("o", q_dim, gqa_unit),
)
@@ -1694,7 +1785,8 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
g_dim = layer.self_attn.g_proj.weight.shape[0]
g_unit = gqa_unit // head_dim
layer.self_attn.g_proj = self.all_to_sharded_linear(
layer.self_attn.g_proj, unit=g_unit,
layer.self_attn.g_proj,
unit=g_unit,
weights=self._greedy_weights_for("g", g_dim, g_unit),
)
@@ -1729,9 +1821,9 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.share_expert.down_proj,
weights=self._greedy_weights_for("shared_down", shared_down_dim),
)
moe_gate_dim = layer.mlp.switch_mlp.gate_proj.weight.shape[1]
moe_up_dim = layer.mlp.switch_mlp.up_proj.weight.shape[1]
moe_down_dim = layer.mlp.switch_mlp.down_proj.weight.shape[-1]
moe_gate_dim = int(layer.mlp.switch_mlp.gate_proj.weight.shape[1])
moe_up_dim = int(layer.mlp.switch_mlp.up_proj.weight.shape[1])
moe_down_dim = int(layer.mlp.switch_mlp.down_proj.weight.shape[-1])
self.all_to_sharded_linear_in_place(
layer.mlp.switch_mlp.gate_proj,
weights=self._greedy_weights_for("moe_gate", moe_gate_dim),
@@ -1775,19 +1867,23 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
q_dim = mixer.q_proj.weight.shape[0]
k_dim = mixer.k_proj.weight.shape[0]
mixer.q_proj = self.all_to_sharded_linear(
mixer.q_proj, unit=gqa_unit,
mixer.q_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("q", q_dim, gqa_unit),
)
mixer.k_proj = self.all_to_sharded_linear(
mixer.k_proj, unit=attn_head_dim,
mixer.k_proj,
unit=attn_head_dim,
weights=self._greedy_weights_for("k", k_dim, attn_head_dim),
)
mixer.v_proj = self.all_to_sharded_linear(
mixer.v_proj, unit=attn_head_dim,
mixer.v_proj,
unit=attn_head_dim,
weights=self._greedy_weights_for("v", k_dim, attn_head_dim),
)
mixer.o_proj = self.sharded_to_all_linear(
mixer.o_proj, unit=gqa_unit,
mixer.o_proj,
unit=gqa_unit,
weights=self._greedy_weights_for("o", q_dim, gqa_unit),
)
mixer.num_heads = mixer.q_proj.weight.shape[0] // attn_head_dim
@@ -1800,8 +1896,8 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
elif isinstance(mixer, NemotronHMoE):
# Shard routed experts (SwitchMLP uses fc1/fc2)
moe_fc1_dim = mixer.switch_mlp.fc1.weight.shape[1]
moe_fc2_dim = mixer.switch_mlp.fc2.weight.shape[-1]
moe_fc1_dim = int(mixer.switch_mlp.fc1.weight.shape[1])
moe_fc2_dim = int(mixer.switch_mlp.fc2.weight.shape[-1])
self.all_to_sharded_linear_in_place(
mixer.switch_mlp.fc1,
weights=self._greedy_weights_for("moe_gate", moe_fc1_dim),
@@ -1820,7 +1916,9 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
)
self.sharded_to_all_linear_in_place(
mixer.shared_experts.down_proj,
weights=self._greedy_weights_for("shared_down", shared_down_dim),
weights=self._greedy_weights_for(
"shared_down", shared_down_dim
),
)
mixer = ShardedMoE(mixer) # pyright: ignore[reportArgumentType]
mixer.sharding_group = self.group
@@ -1844,7 +1942,8 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
heads_per_group = num_heads // n_groups
out_unit = heads_per_group * head_dim
mixer.out_proj = self.sharded_to_all_linear(
mixer.out_proj, unit=out_unit,
mixer.out_proj,
unit=out_unit,
weights=self._greedy_weights_for("mamba_out", intermediate_size, out_unit),
)
out_actual = self._greedy_last_sizes("mamba_out")
@@ -1852,7 +1951,9 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
head_sizes = [s // head_dim for s in out_actual]
group_sizes = [s // heads_per_group for s in head_sizes]
else:
group_sizes = compute_shard_sizes(n_groups, world_size, weights=self.shard_weights)
group_sizes = compute_shard_sizes(
n_groups, world_size, weights=self.shard_weights
)
head_sizes = [g * heads_per_group for g in group_sizes]
groups_per_rank = group_sizes[rank]
src/exo/worker/engines/mlx/utils_mlx.py
+4 -1
View File
@@ -271,7 +271,10 @@ def shard_and_load(
match shard_metadata:
case TensorShardMetadata():
logger.info(f"loading model from {model_path} with tensor parallelism")
logger.info(
f"loading model from {model_path} with tensor parallelism "
f"(weights={shard_metadata.shard_weights}, mode={shard_metadata.shard_mode})"
)
model = tensor_auto_parallel(
model,
group,
src/exo/worker/tests/unittests/test_mlx/test_uneven_sharding.py
+3 -1
View File
@@ -12,6 +12,8 @@ import numpy as np
import pytest
from mlx.nn.layers.distributed import compute_shard_sizes
from exo.shared.types.worker.shards import TensorShardMode
RANDOM_SEED = 42
INPUT_TOKENS = [1, 100, 200, 300]
@@ -608,7 +610,7 @@ class TestGreedyShardingTP2:
world_size=2,
base_port=_next_port_block(),
shard_weights=[2.0, 1.0],
shard_mode="Greedy",
shard_mode=TensorShardMode.Greedy,
)
diff = float(np.max(np.abs(single_logits - tp2_logits)))