style: apply nix fmt formatting fixes

Co-Authored-By: Claude Opus 4.6 <[email protected]>
This commit is contained in:
Alex Cheema
2026-02-21 11:08:02 -08:00
co-authored by Claude Opus 4.6
parent ee7db70d6c
commit c3da520a84
7 changed files with 53 additions and 35 deletions
@@ -185,8 +185,7 @@
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (instanceTag === "MlxJacclInstance")
instanceType = "MLX RDMA";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
let sharding: string | null = null;
const inst = instance as {
+19 -3
View File
@@ -247,15 +247,31 @@
<!-- Keyboard base (trapezoidal) -->
<path
d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX + mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY + mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
fill="#2c2c2c"
stroke={wireColor}
stroke-width="1"
/>
<!-- Keyboard area -->
<rect x={mbKbX} y={mbKbY} width={mbKbW} height={mbKbH} fill="rgba(0,0,0,0.2)" rx="2" />
<rect
x={mbKbX}
y={mbKbY}
width={mbKbW}
height={mbKbH}
fill="rgba(0,0,0,0.2)"
rx="2"
/>
<!-- Trackpad -->
<rect x={mbTpX} y={mbTpY} width={mbTpW} height={mbTpH} fill="rgba(255,255,255,0.08)" rx="2" />
<rect
x={mbTpX}
y={mbTpY}
width={mbTpW}
height={mbTpH}
fill="rgba(255,255,255,0.08)"
rx="2"
/>
{/if}
+1 -2
View File
@@ -917,8 +917,7 @@ class AppStore {
let instanceType: string | null = null;
if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
else if (instanceTag === "MlxJacclInstance")
instanceType = "MLX RDMA";
else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
let sharding: string | null = null;
const inst = instance as {
+7 -21
View File
@@ -511,9 +511,7 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
total = len(model.layers)
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
)
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)
@@ -569,9 +567,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(
layer.parameters(), timeout_seconds / total, on_timeout
)
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
if layer.self_attn.q_lora_rank is None:
@@ -801,9 +797,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
model = cast(MiniMaxModel, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(
layer.parameters(), timeout_seconds / total, on_timeout
)
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)
@@ -844,9 +838,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
model = cast(Qwen3MoeModel | Qwen3NextModel, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(
layer.parameters(), timeout_seconds / total, on_timeout
)
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
if isinstance(layer, Qwen3DecoderLayer):
layer.self_attn.q_proj = self.all_to_sharded_linear(
@@ -971,9 +963,7 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
model = cast(Glm4MoeModel, model)
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(
layer.parameters(), timeout_seconds / total, on_timeout
)
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)
@@ -1022,9 +1012,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(
layer.parameters(), timeout_seconds / total, on_timeout
)
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)
@@ -1067,9 +1055,7 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy):
total = len(model.layers)
for i, layer in enumerate(model.layers):
eval_with_timeout(
layer.parameters(), timeout_seconds / total, on_timeout
)
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)
+10 -3
View File
@@ -202,7 +202,10 @@ def load_mlx_items(
logger.info("Starting distributed init")
start_time = time.perf_counter()
model, tokenizer = shard_and_load(
bound_instance.bound_shard, group=group, on_timeout=on_timeout, on_layer_loaded=on_layer_loaded
bound_instance.bound_shard,
group=group,
on_timeout=on_timeout,
on_layer_loaded=on_layer_loaded,
)
end_time = time.perf_counter()
logger.info(
@@ -259,10 +262,14 @@ def shard_and_load(
match shard_metadata:
case TensorShardMetadata():
logger.info(f"loading model from {model_path} with tensor parallelism")
model = tensor_auto_parallel(model, group, timeout_seconds, on_timeout, on_layer_loaded)
model = tensor_auto_parallel(
model, group, timeout_seconds, on_timeout, on_layer_loaded
)
case PipelineShardMetadata():
logger.info(f"loading model from {model_path} with pipeline parallelism")
model = pipeline_auto_parallel(model, group, shard_metadata, on_layer_loaded=on_layer_loaded)
model = pipeline_auto_parallel(
model, group, shard_metadata, on_layer_loaded=on_layer_loaded
)
eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
case CfgShardMetadata():
raise ValueError(
+11 -3
View File
@@ -151,7 +151,9 @@ def main(
isinstance(current_status, RunnerConnected) and group is not None
) or (isinstance(current_status, RunnerIdle) and group is None):
total_layers = shard_metadata.end_layer - shard_metadata.start_layer
current_status = RunnerLoading(layers_loaded=0, total_layers=total_layers)
current_status = RunnerLoading(
layers_loaded=0, total_layers=total_layers
)
logger.info("runner loading")
event_sender.send(
RunnerStatusUpdated(
@@ -174,7 +176,10 @@ def main(
def on_layer_loaded(layers_loaded: int, total: int) -> None:
event_sender.send(
RunnerStatusUpdated(
runner_id=runner_id, runner_status=RunnerLoading(layers_loaded=layers_loaded, total_layers=total)
runner_id=runner_id,
runner_status=RunnerLoading(
layers_loaded=layers_loaded, total_layers=total
),
)
)
@@ -182,7 +187,10 @@ def main(
ModelTask.TextGeneration in shard_metadata.model_card.tasks
), f"Incorrect model task(s): {shard_metadata.model_card.tasks}"
inference_model, tokenizer = load_mlx_items(
bound_instance, group, on_timeout=on_model_load_timeout, on_layer_loaded=on_layer_loaded
bound_instance,
group,
on_timeout=on_model_load_timeout,
on_layer_loaded=on_layer_loaded,
)
logger.info(
f"model has_tool_calling={tokenizer.has_tool_calling} using tokens {tokenizer.tool_call_start}, {tokenizer.tool_call_end}"
@@ -224,7 +224,10 @@ def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):
),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerConnected()),
TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Running),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoading(layers_loaded=0, total_layers=32)),
RunnerStatusUpdated(
runner_id=RUNNER_1_ID,
runner_status=RunnerLoading(layers_loaded=0, total_layers=32),
),
TaskAcknowledged(task_id=LOAD_TASK_ID),
TaskStatusUpdated(task_id=LOAD_TASK_ID, task_status=TaskStatus.Complete),
RunnerStatusUpdated(runner_id=RUNNER_1_ID, runner_status=RunnerLoaded()),