From 2837225dc7fe6fec740fe9e2caf4078cd7468fee Mon Sep 17 00:00:00 2001
From: rltakashige <rl.takashige@gmail.com>
Date: Thu, 29 Jan 2026 17:08:38 +0000
Subject: [PATCH] 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],