From ade3ee7ec5c4e8233fd4f04d033c87fc924c7369 Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Tue, 30 Dec 2025 02:20:51 +0000 Subject: [PATCH] fix warmup order. should be rank!=0 then rank=0 --- src/exo/worker/plan.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index e313ca73..6c621e72 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -235,9 +235,8 @@ def _ready_to_warmup( assert device_rank < world_size assert device_rank >= 0 - # TODO: Ensure these align with MLX distributeds expectations. - # Rank < n-1 - accepting_ranks_ready = device_rank < world_size - 1 and all( + # Rank != 0 + accepting_ranks_ready = device_rank > 0 and all( isinstance( all_runners.get(global_runner_id, None), (RunnerLoaded, RunnerWarmingUp), @@ -245,8 +244,8 @@ def _ready_to_warmup( for global_runner_id in shard_assignments.runner_to_shard ) - # Rank = n-1 - connecting_rank_ready = device_rank == world_size - 1 and all( + # Rank = 0 + connecting_rank_ready = device_rank == 0 and all( isinstance(all_runners.get(global_runner_id, None), RunnerWarmingUp) for global_runner_id in shard_assignments.runner_to_shard if global_runner_id != runner_id