fix(mlx): let sharded models re-derive layer indices after pipeline split

[zhast]

The bug

pipeline_auto_parallel() replaces the model's layer stack with a slice after the model has been constructed:

layers = layers[start_layer:end_layer]
...
_set_layers(model, layers)

Any index a model derived from the full stack in __init__ — which cache entry holds the first full-attention layer, which holds the first recurrent layer — now addresses the wrong entry on every shard but the first. The model then builds its attention and SSM masks from the wrong cache.

Nothing raises. Generation runs at full speed and returns fluent-looking tokens with no semantics. That failure mode is expensive to track down: activations stay finite and well-scaled all the way through, so every obvious diagnostic looks healthy.

Why a general hook

auto_parallel already repairs exactly this, by name, for four model types:

if isinstance(inner_model_instance, GptOssMoeModel): ...
if isinstance(inner_model_instance, Step35InnerModel): ...
if isinstance(inner_model_instance, (Qwen3_5TextModelInner, Qwen3NextInnerModel)): ...
if isinstance(inner_model_instance, NemotronHInnerModel): ...

That list has grown once per hybrid model added. The problem is not that the list is incomplete — it is that a stack which is not on it fails silently instead of loudly, including any model loaded from outside the tree.

This adds a small opt-in hook after _set_layers() so a model can re-derive its own state from the sharded layer list. Models that do not define resync_sharded_layers are completely unaffected — no behaviour change for anything in-tree today.

How it was found

Running GLM-5.3-Flash (glm5_next: 34 gated-delta linear-attention layers interleaved with 11 DeepSeek-sparse-attention layers, hyper-connections, 288-expert MoE) across a 4-node Mac Studio pipeline over Thunderbolt/RDMA.

The inner model caches:

self.ssm_idx = next((i for i, l in enumerate(self.layers) if l.is_linear), 0)
self.fa_idx  = next((i for i, l in enumerate(self.layers) if not l.is_linear), 0)

On the full 45-layer stack that gives ssm_idx=0, fa_idx=3. The shard holding global layers 27–44 needs fa_idx=0, ssm_idx=1. With the stale values, create_ssm_mask() receives a full-attention CacheList and the attention path receives the recurrent cache.

Note the , 0) default is a second trap: a shard containing no layer of one kind silently gets index 0 — a valid index into the wrong entry — rather than None.

Testing

src/exo/worker/engines/mlx/tests/test_pipeline_layer_resync.py covers both that the hook fires and corrects indices on a non-first shard, and that a model without the hook shards exactly as before.

uv run pytest src/exo/worker/engines/mlx/tests/  ->  8 passed
uv run ruff check                                ->  All checks passed
uv run ruff format --check                       ->  already formatted
uv run basedpyright <changed files>              ->  0 errors, 0 warnings

Draft: happy to fold the four existing isinstance special-cases onto this hook if you'd prefer that shape, but left them untouched here to keep the change additive.

🤖 Generated with Claude Code