← back to Exo
Fix several issues with placement (#1200)
f82f862fd7a368d6ee563c2aaa07e3c063e06661 · 2026-01-19 11:52:35 +0000 · rltakashige
## Motivation
Uneven placements were causing issues for some users with lopsided
setups. While fixing, I ran into another issue with impossible
allocation of memory.
## Changes
- Allocate at least 1 layer per device.
- Catch overallocation of memory with an error.
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
Tested that GPT OSS is placed correctly.
### Automated Testing
Added breaking tests in the first commit. Resolved with new placement
algorithm in the second one.
Files touched
M src/exo/master/placement_utils.pyM src/exo/master/tests/test_placement.pyM src/exo/master/tests/test_placement_utils.py
Diff
commit f82f862fd7a368d6ee563c2aaa07e3c063e06661
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Jan 19 11:52:35 2026 +0000
Fix several issues with placement (#1200)
## Motivation
Uneven placements were causing issues for some users with lopsided
setups. While fixing, I ran into another issue with impossible
allocation of memory.
## Changes
- Allocate at least 1 layer per device.
- Catch overallocation of memory with an error.
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
Tested that GPT OSS is placed correctly.
### Automated Testing
Added breaking tests in the first commit. Resolved with new placement
algorithm in the second one.
---
src/exo/master/placement_utils.py | 74 +++++++++++++++++----
src/exo/master/tests/test_placement.py | 2 +-
src/exo/master/tests/test_placement_utils.py | 97 ++++++++++++++++++++++++++++
3 files changed, 160 insertions(+), 13 deletions(-)
diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py
index 04f5f7a9..a5d70f7a 100644
--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -49,33 +49,83 @@ def get_smallest_cycles(cycles: list[list[NodeInfo]]) -> list[list[NodeInfo]]:
return [cycle for cycle in cycles if len(cycle) == min_nodes]
+def allocate_layers_proportionally(
+ total_layers: int,
+ memory_fractions: list[float],
+) -> list[int]:
+ n = len(memory_fractions)
+ if n == 0:
+ raise ValueError("Cannot allocate layers to an empty node list")
+ if total_layers < n:
+ raise ValueError(
+ f"Cannot distribute {total_layers} layers across {n} nodes "
+ "(need at least 1 layer per node)"
+ )
+
+ # Largest remainder: floor each, then distribute remainder by fractional part
+ raw = [f * total_layers for f in memory_fractions]
+ result = [int(r) for r in raw]
+ by_remainder = sorted(range(n), key=lambda i: raw[i] - result[i], reverse=True)
+ for i in range(total_layers - sum(result)):
+ result[by_remainder[i]] += 1
+
+ # Ensure minimum 1 per node by taking from the largest
+ for i in range(n):
+ if result[i] == 0:
+ max_idx = max(range(n), key=lambda j: result[j])
+ assert result[max_idx] > 1
+ result[max_idx] -= 1
+ result[i] = 1
+
+ return result
+
+
def get_shard_assignments_for_pipeline_parallel(
model_meta: ModelMetadata,
selected_cycle: list[NodeWithProfile],
):
+ if not selected_cycle:
+ raise ValueError("Cannot create shard assignments for empty node cycle")
+
cycle_memory = sum(
(node.node_profile.memory.ram_available for node in selected_cycle),
start=Memory(),
)
+
+ if cycle_memory.in_bytes == 0:
+ raise ValueError("Cannot create shard assignments: total available memory is 0")
+
total_layers = model_meta.n_layers
world_size = len(selected_cycle)
runner_to_shard: dict[RunnerId, ShardMetadata] = {}
node_to_runner: dict[NodeId, RunnerId] = {}
- layers_assigned = 0
- for i, node in enumerate(selected_cycle):
- if i == len(selected_cycle) - 1:
- node_layers = total_layers - layers_assigned
- else:
- node_layers = round(
- total_layers
- * (
- node.node_profile.memory.ram_available.in_bytes
- / cycle_memory.in_bytes
- )
+ layer_allocations = allocate_layers_proportionally(
+ total_layers=total_layers,
+ memory_fractions=[
+ node.node_profile.memory.ram_available.in_bytes / cycle_memory.in_bytes
+ for node in selected_cycle
+ ],
+ )
+
+ # Validate each node has sufficient memory for its assigned layers
+ memory_per_layer = model_meta.storage_size.in_bytes / total_layers
+ for i, (node, node_layers) in enumerate(
+ zip(selected_cycle, layer_allocations, strict=True)
+ ):
+ required_memory = node_layers * memory_per_layer
+ available_memory = node.node_profile.memory.ram_available.in_bytes
+ if required_memory > available_memory:
+ raise ValueError(
+ f"Node {i} ({node.node_id}) has insufficient memory: "
+ f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
+ f"but only has {available_memory / (1024**3):.2f} GB available"
)
- node_layers = max(1, node_layers)
+ layers_assigned = 0
+ for i, (node, node_layers) in enumerate(
+ zip(selected_cycle, layer_allocations, strict=True)
+ ):
runner_id = RunnerId()
shard = PipelineShardMetadata(
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index f32ac7b1..3b8be325 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -70,7 +70,7 @@ def place_instance_command(model_meta: ModelMetadata) -> PlaceInstance:
[
((500, 500, 1000), 12, (3, 3, 6)),
((500, 500, 500), 12, (4, 4, 4)),
- ((312, 518, 1024), 12, (2, 3, 7)),
+ ((312, 468, 1092), 12, (2, 3, 7)),
],
)
def test_get_instance_placements_create_instance(
diff --git a/src/exo/master/tests/test_placement_utils.py b/src/exo/master/tests/test_placement_utils.py
index 9a9466ef..fe79ef2e 100644
--- a/src/exo/master/tests/test_placement_utils.py
+++ b/src/exo/master/tests/test_placement_utils.py
@@ -3,6 +3,7 @@ from typing import Callable
import pytest
from exo.master.placement_utils import (
+ allocate_layers_proportionally,
filter_cycles_by_memory,
get_hosts_from_subgraph,
get_mlx_jaccl_coordinators,
@@ -165,6 +166,9 @@ def test_get_smallest_cycles(
((500, 500, 1000), 12, (3, 3, 6)),
((500, 500, 500), 12, (4, 4, 4)),
((312, 518, 1024), 12, (2, 3, 7)),
+ # Edge case: one node has ~90% of memory - should not over-allocate.
+ # Each node must have enough memory for at least 1 layer (50 KB = 1000/20).
+ ((900, 50, 50), 20, (18, 1, 1)),
],
)
def test_get_shard_assignments(
@@ -397,3 +401,96 @@ def test_get_mlx_jaccl_coordinators(
assert coordinators[node_c_id] == (
f"{conn_c_a.send_back_multiaddr.ip_address}:5000"
), "node_c should use the IP from conn_c_a"
+
+
+class TestAllocateLayersProportionally:
+ def test_empty_node_list_raises(self):
+ with pytest.raises(ValueError, match="empty node list"):
+ allocate_layers_proportionally(total_layers=10, memory_fractions=[])
+
+ def test_zero_layers_raises(self):
+ with pytest.raises(ValueError, match="need at least 1 layer per node"):
+ allocate_layers_proportionally(total_layers=0, memory_fractions=[0.5, 0.5])
+
+ def test_negative_layers_raises(self):
+ with pytest.raises(ValueError, match="need at least 1 layer per node"):
+ allocate_layers_proportionally(total_layers=-1, memory_fractions=[0.5, 0.5])
+
+ def test_fewer_layers_than_nodes_raises(self):
+ with pytest.raises(ValueError, match="need at least 1 layer per node"):
+ allocate_layers_proportionally(
+ total_layers=2, memory_fractions=[0.33, 0.33, 0.34]
+ )
+
+ def test_equal_distribution(self):
+ result = allocate_layers_proportionally(
+ total_layers=12, memory_fractions=[0.25, 0.25, 0.25, 0.25]
+ )
+ assert result == [3, 3, 3, 3]
+ assert sum(result) == 12
+
+ def test_proportional_distribution(self):
+ result = allocate_layers_proportionally(
+ total_layers=12, memory_fractions=[0.25, 0.25, 0.50]
+ )
+ assert result == [3, 3, 6]
+ assert sum(result) == 12
+
+ def test_extreme_imbalance_ensures_minimum(self):
+ result = allocate_layers_proportionally(
+ total_layers=20, memory_fractions=[0.975, 0.0125, 0.0125]
+ )
+ assert all(layers >= 1 for layers in result)
+ assert sum(result) == 20
+ # Small nodes get minimum 1 layer
+ assert result == [18, 1, 1]
+
+ def test_single_node_gets_all_layers(self):
+ result = allocate_layers_proportionally(total_layers=10, memory_fractions=[1.0])
+ assert result == [10]
+
+ def test_minimum_viable_allocation(self):
+ result = allocate_layers_proportionally(
+ total_layers=3, memory_fractions=[0.33, 0.33, 0.34]
+ )
+ assert result == [1, 1, 1]
+ assert sum(result) == 3
+
+
+def test_get_shard_assignments_insufficient_memory_raises(
+ topology: Topology,
+ create_node: Callable[[int, NodeId | None], NodeInfo],
+ create_connection: Callable[[NodeId, NodeId], Connection],
+):
+ """Test that ValueError is raised when a node has insufficient memory for its layers."""
+ node_a_id = NodeId()
+ node_b_id = NodeId()
+ node_c_id = NodeId()
+
+ # Node C has only 10 KB but would need 50 KB for 1 layer (1000 KB / 20 layers)
+ node_a = create_node(900 * 1024, node_a_id)
+ node_b = create_node(50 * 1024, node_b_id)
+ node_c = create_node(10 * 1024, node_c_id) # Insufficient memory
+
+ topology.add_node(node_a)
+ topology.add_node(node_b)
+ topology.add_node(node_c)
+
+ topology.add_connection(create_connection(node_a_id, node_b_id))
+ topology.add_connection(create_connection(node_b_id, node_c_id))
+ topology.add_connection(create_connection(node_c_id, node_a_id))
+ topology.add_connection(create_connection(node_b_id, node_a_id))
+
+ model_meta = ModelMetadata(
+ model_id=ModelId("test-model"),
+ pretty_name="Test Model",
+ n_layers=20,
+ storage_size=Memory.from_kb(1000),
+ hidden_size=1000,
+ supports_tensor=True,
+ )
+ cycles = topology.get_cycles()
+ selected_cycle = cycles[0]
+
+ with pytest.raises(ValueError, match="insufficient memory"):
+ get_shard_assignments(model_meta, selected_cycle, Sharding.Pipeline)
← 7ff937d8 Add dashboard screenshots to README (#1185)
·
back to Exo
·
tidy: remove context manager from api (#1199) 746589ba →