Load pipeline layers sequentially (#1329)

2837225dc7fe6fec740fe9e2caf4078cd7468fee · 2026-01-29 17:08:38 +0000 · rltakashige

## Motivation

Slightly annoyed by needing this change, but same story as for tensor
loading...

Files touched

M src/exo/worker/engines/mlx/auto_parallel.py

Diff

commit 2837225dc7fe6fec740fe9e2caf4078cd7468fee
Author: rltakashige <rl.takashige@gmail.com>
Date:   Thu Jan 29 17:08:38 2026 +0000

    Load pipeline layers sequentially (#1329)
    
    ## Motivation
    
    Slightly annoyed by needing this change, but same story as for tensor
    loading...
---
 src/exo/worker/engines/mlx/auto_parallel.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index ff2052fb..84ed1f29 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -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],

← e4c6a7db nix: add Python packaging with uv2nix · back to Exo · Treat Swift/Xcode build warnings as errors (#1322) 7cc313b2 →