Okay we should probably await the update

3e869051f618c3b38bb8156a965b11202e71cc9b · 2024-11-19 07:59:53 -0800 · Nel Nibcord

Files touched

M exo/inference/mlx/sharded_inference_engine.py

Diff

commit 3e869051f618c3b38bb8156a965b11202e71cc9b
Author: Nel Nibcord <blindcrone@tuta.io>
Date:   Tue Nov 19 07:59:53 2024 -0800

    Okay we should probably await the update
---
 exo/inference/mlx/sharded_inference_engine.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/inference/mlx/sharded_inference_engine.py b/exo/inference/mlx/sharded_inference_engine.py
index 00c07f9f..3c5de55b 100644
--- a/exo/inference/mlx/sharded_inference_engine.py
+++ b/exo/inference/mlx/sharded_inference_engine.py
@@ -93,8 +93,8 @@ class MLXDynamicShardInferenceEngine(InferenceEngine):
     l = mx.array(lengths)
     loop = asyncio.get_running_loop()
     score, grad = await loop.run_in_executor(self.executor, self.session['LVaG'], self.model, x, y, l)
-    loop.run_in_executor(self.executor, self.update_model, grad, score)
     layers = [{k: v["weight"].shape for k,v in l.items() if 'weight' in v} for l in grad['model']['model']['layers'] if l]
+    await loop.run_in_executor(self.executor, self.update_model, grad, score)
 
     return np.array(score).reshape(inputs.shape[0], -1), np.array(layers[0]['input_layernorm']).reshape(inputs.shape[0], -1)

← 75c8650f Naive network-propagated loss implementation on MLX · back to Exo · Coordination biz 175ebc1c →