Load pipeline layers sequentially (#1329)

## Motivation Slightly annoyed by needing this change, but same story as for tensor loading...
2026-01-29 17:08:38 +00:00
parent e4c6a7dbb4
commit 2837225dc7
1 changed files with 3 additions and 0 deletions
@@ -201,6 +201,9 @@ def pipeline_auto_parallel(
    device_rank, world_size = model_shard_meta.device_rank, model_shard_meta.world_size

    layers = layers[start_layer:end_layer]
+    for layer in layers:
+        mx.eval(layer)  # type: ignore
+
    layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
    layers[-1] = PipelineLastLayer(
        layers[-1],