← back to Exo
Simplify model cards (#1204)
d4f551c602205c95941b872a4af29cfc7abb538d · 2026-01-20 11:01:19 +0000 · Evan Quiney
## Motivation
We have a lot of unneeded data in the model card - lets just keep the
necessary stuff and add back more data when we need it
## Test Plan
EXO still runs! (pipeline on 2)
Co-authored-by: rltakashige <rl.takashige@gmail.com>
Files touched
M dashboard/src/routes/+page.svelteM dashboard/src/routes/downloads/+page.svelteM src/exo/master/api.pyM src/exo/master/placement.pyM src/exo/master/placement_utils.pyM src/exo/master/tests/test_master.pyM src/exo/master/tests/test_placement.pyM src/exo/master/tests/test_placement_utils.pyM src/exo/shared/models/model_cards.pyM src/exo/shared/models/model_meta.pyM src/exo/shared/tests/conftest.pyM src/exo/shared/types/api.pyM src/exo/shared/types/chunks.pyM src/exo/shared/types/commands.pyM src/exo/shared/types/common.pyD src/exo/shared/types/models.pyM src/exo/shared/types/worker/runners.pyM src/exo/shared/types/worker/shards.pyM src/exo/worker/download/download_utils.pyM src/exo/worker/download/impl_shard_downloader.pyM src/exo/worker/download/shard_downloader.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/runner/runner.pyM src/exo/worker/tests/constants.pyM src/exo/worker/tests/unittests/conftest.pyM src/exo/worker/tests/unittests/test_mlx/conftest.pyM src/exo/worker/tests/unittests/test_mlx/test_tokenizers.pyM src/exo/worker/tests/unittests/test_plan/test_download_and_loading.pyM tests/headless_runner.py
Diff
commit d4f551c602205c95941b872a4af29cfc7abb538d
Author: Evan Quiney <evanev7@gmail.com>
Date: Tue Jan 20 11:01:19 2026 +0000
Simplify model cards (#1204)
## Motivation
We have a lot of unneeded data in the model card - lets just keep the
necessary stuff and add back more data when we need it
## Test Plan
EXO still runs! (pipeline on 2)
Co-authored-by: rltakashige <rl.takashige@gmail.com>
---
dashboard/src/routes/+page.svelte | 4 +-
dashboard/src/routes/downloads/+page.svelte | 4 +-
src/exo/master/api.py | 72 ++-
src/exo/master/placement.py | 16 +-
src/exo/master/placement_utils.py | 27 +-
src/exo/master/tests/test_master.py | 8 +-
src/exo/master/tests/test_placement.py | 46 +-
src/exo/master/tests/test_placement_utils.py | 12 +-
src/exo/shared/models/model_cards.py | 645 ++++++---------------
src/exo/shared/models/model_meta.py | 26 +-
src/exo/shared/tests/conftest.py | 5 +-
src/exo/shared/types/api.py | 4 +-
src/exo/shared/types/chunks.py | 2 +-
src/exo/shared/types/commands.py | 4 +-
src/exo/shared/types/common.py | 4 +-
src/exo/shared/types/models.py | 18 -
src/exo/shared/types/worker/runners.py | 2 +-
src/exo/shared/types/worker/shards.py | 6 +-
src/exo/worker/download/download_utils.py | 30 +-
src/exo/worker/download/impl_shard_downloader.py | 18 +-
src/exo/worker/download/shard_downloader.py | 5 +-
src/exo/worker/engines/mlx/utils_mlx.py | 8 +-
src/exo/worker/main.py | 16 +-
src/exo/worker/plan.py | 6 +-
src/exo/worker/runner/runner.py | 4 +-
src/exo/worker/tests/constants.py | 2 +-
src/exo/worker/tests/unittests/conftest.py | 5 +-
.../worker/tests/unittests/test_mlx/conftest.py | 8 +-
.../tests/unittests/test_mlx/test_tokenizers.py | 6 +-
.../test_plan/test_download_and_loading.py | 2 +-
tests/headless_runner.py | 24 +-
31 files changed, 351 insertions(+), 688 deletions(-)
diff --git a/dashboard/src/routes/+page.svelte b/dashboard/src/routes/+page.svelte
index e013688a..c6dc8d14 100644
--- a/dashboard/src/routes/+page.svelte
+++ b/dashboard/src/routes/+page.svelte
@@ -434,8 +434,8 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
const shardData = shardObj[shardKeys[0]] as Record<string, unknown>;
if (!shardData) return null;
- // Model meta is nested: shard.model_meta.model_id
- const modelMeta = shardData.model_meta ?? shardData.modelMeta;
+ // Model meta is nested: shard.model_card.model_id
+ const modelMeta = shardData.model_card ?? shardData.modelCard;
if (!modelMeta || typeof modelMeta !== 'object') return null;
const meta = modelMeta as Record<string, unknown>;
diff --git a/dashboard/src/routes/downloads/+page.svelte b/dashboard/src/routes/downloads/+page.svelte
index 6d8ad418..bfdef5ca 100644
--- a/dashboard/src/routes/downloads/+page.svelte
+++ b/dashboard/src/routes/downloads/+page.svelte
@@ -98,7 +98,7 @@
const shardData = shardObj[shardKeys[0]] as Record<string, unknown>;
if (!shardData) return null;
- const modelMeta = shardData.model_meta ?? shardData.modelMeta;
+ const modelMeta = shardData.model_card ?? shardData.modelCard;
if (!modelMeta || typeof modelMeta !== 'object') return null;
const meta = modelMeta as Record<string, unknown>;
@@ -190,7 +190,7 @@
const shardKeys = Object.keys(shardObj);
if (shardKeys.length !== 1) return null;
const shardData = shardObj[shardKeys[0]] as Record<string, unknown>;
- const modelMeta = shardData?.model_meta ?? shardData?.modelMeta;
+ const modelMeta = shardData?.model_card ?? shardData?.modelCard;
if (!modelMeta || typeof modelMeta !== 'object') return null;
const meta = modelMeta as Record<string, unknown>;
return (meta.prettyName as string) ?? null;
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index 5ea958ed..e8a4e62b 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -19,8 +19,8 @@ from exo.master.placement import place_instance as get_instance_placements
from exo.shared.apply import apply
from exo.shared.election import ElectionMessage
from exo.shared.logging import InterceptLogger
-from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.models.model_meta import get_model_meta
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
+from exo.shared.models.model_meta import get_model_card
from exo.shared.types.api import (
BenchChatCompletionResponse,
BenchChatCompletionTaskParams,
@@ -59,7 +59,6 @@ from exo.shared.types.events import (
IndexedEvent,
)
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.state import State
from exo.shared.types.tasks import ChatCompletionTaskParams
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
@@ -87,12 +86,12 @@ def chunk_to_response(
)
-async def resolve_model_meta(model_id: str) -> ModelMetadata:
+async def resolve_model_card(model_id: str) -> ModelCard:
if model_id in MODEL_CARDS:
model_card = MODEL_CARDS[model_id]
- return model_card.metadata
+ return model_card
else:
- return await get_model_meta(model_id)
+ return await get_model_card(model_id)
class API:
@@ -197,7 +196,7 @@ class API:
async def place_instance(self, payload: PlaceInstanceParams):
command = PlaceInstance(
- model_meta=await resolve_model_meta(payload.model_id),
+ model_card=await resolve_model_card(payload.model_id),
sharding=payload.sharding,
instance_meta=payload.instance_meta,
min_nodes=payload.min_nodes,
@@ -207,15 +206,15 @@ class API:
return CreateInstanceResponse(
message="Command received.",
command_id=command.command_id,
- model_meta=command.model_meta,
+ model_card=command.model_card,
)
async def create_instance(
self, payload: CreateInstanceParams
) -> CreateInstanceResponse:
instance = payload.instance
- model_meta = await resolve_model_meta(instance.shard_assignments.model_id)
- required_memory = model_meta.storage_size
+ model_card = await resolve_model_card(instance.shard_assignments.model_id)
+ required_memory = model_card.storage_size
available_memory = self._calculate_total_available_memory()
if required_memory > available_memory:
@@ -232,7 +231,7 @@ class API:
return CreateInstanceResponse(
message="Command received.",
command_id=command.command_id,
- model_meta=model_meta,
+ model_card=model_card,
)
async def get_placement(
@@ -242,12 +241,12 @@ class API:
instance_meta: InstanceMeta = InstanceMeta.MlxRing,
min_nodes: int = 1,
) -> Instance:
- model_meta = await resolve_model_meta(model_id)
+ model_card = await resolve_model_card(model_id)
try:
placements = get_instance_placements(
PlaceInstance(
- model_meta=model_meta,
+ model_card=model_card,
sharding=sharding,
instance_meta=instance_meta,
min_nodes=min_nodes,
@@ -280,7 +279,7 @@ class API:
if len(list(self.state.topology.list_nodes())) == 0:
return PlacementPreviewResponse(previews=[])
- cards = [card for card in MODEL_CARDS.values() if card.short_id == model_id]
+ cards = [card for card in MODEL_CARDS.values() if card.model_id == model_id]
if not cards:
raise HTTPException(status_code=404, detail=f"Model {model_id} not found")
@@ -298,13 +297,12 @@ class API:
# TODO: PDD
# instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
- for card in cards:
- model_meta = card.metadata
+ for model_card in cards:
for sharding, instance_meta, min_nodes in instance_combinations:
try:
placements = get_instance_placements(
PlaceInstance(
- model_meta=model_meta,
+ model_card=model_card,
sharding=sharding,
instance_meta=instance_meta,
min_nodes=min_nodes,
@@ -315,17 +313,17 @@ class API:
current_instances=self.state.instances,
)
except ValueError as exc:
- if (card.model_id, sharding, instance_meta, 0) not in seen:
+ if (model_card.model_id, sharding, instance_meta, 0) not in seen:
previews.append(
PlacementPreview(
- model_id=card.model_id,
+ model_id=model_card.model_id,
sharding=sharding,
instance_meta=instance_meta,
instance=None,
error=str(exc),
)
)
- seen.add((card.model_id, sharding, instance_meta, 0))
+ seen.add((model_card.model_id, sharding, instance_meta, 0))
continue
current_ids = set(self.state.instances.keys())
@@ -336,17 +334,17 @@ class API:
]
if len(new_instances) != 1:
- if (card.model_id, sharding, instance_meta, 0) not in seen:
+ if (model_card.model_id, sharding, instance_meta, 0) not in seen:
previews.append(
PlacementPreview(
- model_id=card.model_id,
+ model_id=model_card.model_id,
sharding=sharding,
instance_meta=instance_meta,
instance=None,
error="Expected exactly one new instance from placement",
)
)
- seen.add((card.model_id, sharding, instance_meta, 0))
+ seen.add((model_card.model_id, sharding, instance_meta, 0))
continue
instance = new_instances[0]
@@ -355,7 +353,7 @@ class API:
memory_delta_by_node: dict[str, int] = {}
if node_ids:
- total_bytes = model_meta.storage_size.in_bytes
+ total_bytes = model_card.storage_size.in_bytes
per_node = total_bytes // len(node_ids)
remainder = total_bytes % len(node_ids)
for index, node_id in enumerate(sorted(node_ids, key=str)):
@@ -363,14 +361,14 @@ class API:
memory_delta_by_node[str(node_id)] = per_node + extra
if (
- card.model_id,
+ model_card.model_id,
sharding,
instance_meta,
len(node_ids),
) not in seen:
previews.append(
PlacementPreview(
- model_id=card.model_id,
+ model_id=model_card.model_id,
sharding=sharding,
instance_meta=instance_meta,
instance=instance,
@@ -378,7 +376,7 @@ class API:
error=None,
)
)
- seen.add((card.model_id, sharding, instance_meta, len(node_ids)))
+ seen.add((model_card.model_id, sharding, instance_meta, len(node_ids)))
return PlacementPreviewResponse(previews=previews)
@@ -553,8 +551,8 @@ class API:
self, payload: ChatCompletionTaskParams
) -> ChatCompletionResponse | StreamingResponse:
"""Handle chat completions, supporting both streaming and non-streaming responses."""
- model_meta = await resolve_model_meta(payload.model)
- payload.model = model_meta.model_id
+ model_card = await resolve_model_card(payload.model)
+ payload.model = model_card.model_id
if not any(
instance.shard_assignments.model_id == payload.model
@@ -580,8 +578,8 @@ class API:
async def bench_chat_completions(
self, payload: BenchChatCompletionTaskParams
) -> BenchChatCompletionResponse:
- model_meta = await resolve_model_meta(payload.model)
- payload.model = model_meta.model_id
+ model_card = await resolve_model_card(payload.model)
+ payload.model = model_card.model_id
if not any(
instance.shard_assignments.model_id == payload.model
@@ -614,13 +612,13 @@ class API:
return ModelList(
data=[
ModelListModel(
- id=card.short_id,
+ id=card.model_id,
hugging_face_id=card.model_id,
- name=card.name,
- description=card.description,
- tags=card.tags,
- storage_size_megabytes=int(card.metadata.storage_size.in_mb),
- supports_tensor=card.metadata.supports_tensor,
+ name=card.model_id.short(),
+ description="",
+ tags=[],
+ storage_size_megabytes=int(card.storage_size.in_mb),
+ supports_tensor=card.supports_tensor,
)
for card in MODEL_CARDS.values()
]
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index 3448cd13..5bf7c829 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -14,6 +14,7 @@ from exo.master.placement_utils import (
get_shard_assignments,
get_smallest_cycles,
)
+from exo.shared.models.model_cards import ModelId
from exo.shared.topology import Topology
from exo.shared.types.commands import (
CreateInstance,
@@ -23,7 +24,6 @@ from exo.shared.types.commands import (
from exo.shared.types.common import NodeId
from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.worker.instances import (
Instance,
@@ -60,27 +60,27 @@ def place_instance(
cycles = topology.get_cycles()
candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
cycles_with_sufficient_memory = filter_cycles_by_memory(
- candidate_cycles, node_memory, command.model_meta.storage_size
+ candidate_cycles, node_memory, command.model_card.storage_size
)
if len(cycles_with_sufficient_memory) == 0:
raise ValueError("No cycles found with sufficient memory")
if command.sharding == Sharding.Tensor:
- if not command.model_meta.supports_tensor:
+ if not command.model_card.supports_tensor:
raise ValueError(
- f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_meta.model_id}"
+ f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}"
)
# TODO: the condition here for tensor parallel is not correct, but it works good enough for now.
cycles_with_sufficient_memory = [
cycle
for cycle in cycles_with_sufficient_memory
- if command.model_meta.hidden_size % len(cycle) == 0
+ if command.model_card.hidden_size % len(cycle) == 0
]
if not cycles_with_sufficient_memory:
raise ValueError(
- f"No tensor sharding found for model with hidden_size {command.model_meta.hidden_size} candidate cycles"
+ f"No tensor sharding found for model with hidden_size {command.model_card.hidden_size} candidate cycles"
)
- if command.sharding == Sharding.Pipeline and command.model_meta.model_id == ModelId(
+ if command.sharding == Sharding.Pipeline and command.model_card.model_id == ModelId(
"mlx-community/DeepSeek-V3.1-8bit"
):
raise ValueError(
@@ -111,7 +111,7 @@ def place_instance(
)
shard_assignments = get_shard_assignments(
- command.model_meta, selected_cycle, command.sharding, node_memory
+ command.model_card, selected_cycle, command.sharding, node_memory
)
cycle_digraph: Topology = topology.get_subgraph_from_nodes(selected_cycle.node_ids)
diff --git a/src/exo/master/placement_utils.py b/src/exo/master/placement_utils.py
index eb026135..73d857e9 100644
--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -2,10 +2,10 @@ from collections.abc import Generator, Mapping
from loguru import logger
+from exo.shared.models.model_cards import ModelCard
from exo.shared.topology import Topology
from exo.shared.types.common import Host, NodeId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelMetadata
from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
from exo.shared.types.worker.runners import RunnerId, ShardAssignments
@@ -75,7 +75,7 @@ def allocate_layers_proportionally(
def get_shard_assignments_for_pipeline_parallel(
- model_meta: ModelMetadata,
+ model_card: ModelCard,
cycle: Cycle,
node_memory: Mapping[NodeId, MemoryUsage],
):
@@ -86,11 +86,10 @@ def get_shard_assignments_for_pipeline_parallel(
(node_memory[node_id].ram_available for node_id in cycle.node_ids),
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
+ total_layers = model_card.n_layers
world_size = len(cycle)
runner_to_shard: dict[RunnerId, ShardMetadata] = {}
node_to_runner: dict[NodeId, RunnerId] = {}
@@ -104,7 +103,7 @@ def get_shard_assignments_for_pipeline_parallel(
)
# Validate each node has sufficient memory for its assigned layers
- memory_per_layer = model_meta.storage_size.in_bytes / total_layers
+ memory_per_layer = model_card.storage_size.in_bytes / total_layers
for i, (node_id, node_layers) in enumerate(
zip(cycle.node_ids, layer_allocations, strict=True)
):
@@ -124,7 +123,7 @@ def get_shard_assignments_for_pipeline_parallel(
runner_id = RunnerId()
shard = PipelineShardMetadata(
- model_meta=model_meta,
+ model_card=model_card,
device_rank=i,
world_size=world_size,
start_layer=layers_assigned,
@@ -137,7 +136,7 @@ def get_shard_assignments_for_pipeline_parallel(
layers_assigned += node_layers
shard_assignments = ShardAssignments(
- model_id=model_meta.model_id,
+ model_id=model_card.model_id,
runner_to_shard=runner_to_shard,
node_to_runner=node_to_runner,
)
@@ -146,17 +145,17 @@ def get_shard_assignments_for_pipeline_parallel(
def get_shard_assignments_for_tensor_parallel(
- model_meta: ModelMetadata,
+ model_card: ModelCard,
cycle: Cycle,
):
- total_layers = model_meta.n_layers
+ total_layers = model_card.n_layers
world_size = len(cycle)
runner_to_shard: dict[RunnerId, ShardMetadata] = {}
node_to_runner: dict[NodeId, RunnerId] = {}
for i, node_id in enumerate(cycle):
shard = TensorShardMetadata(
- model_meta=model_meta,
+ model_card=model_card,
device_rank=i,
world_size=world_size,
start_layer=0,
@@ -170,7 +169,7 @@ def get_shard_assignments_for_tensor_parallel(
node_to_runner[node_id] = runner_id
shard_assignments = ShardAssignments(
- model_id=model_meta.model_id,
+ model_id=model_card.model_id,
runner_to_shard=runner_to_shard,
node_to_runner=node_to_runner,
)
@@ -179,7 +178,7 @@ def get_shard_assignments_for_tensor_parallel(
def get_shard_assignments(
- model_meta: ModelMetadata,
+ model_card: ModelCard,
cycle: Cycle,
sharding: Sharding,
node_memory: Mapping[NodeId, MemoryUsage],
@@ -187,13 +186,13 @@ def get_shard_assignments(
match sharding:
case Sharding.Pipeline:
return get_shard_assignments_for_pipeline_parallel(
- model_meta=model_meta,
+ model_card=model_card,
cycle=cycle,
node_memory=node_memory,
)
case Sharding.Tensor:
return get_shard_assignments_for_tensor_parallel(
- model_meta=model_meta,
+ model_card=model_card,
cycle=cycle,
)
diff --git a/src/exo/master/tests/test_master.py b/src/exo/master/tests/test_master.py
index 499814b2..0c1cbea8 100644
--- a/src/exo/master/tests/test_master.py
+++ b/src/exo/master/tests/test_master.py
@@ -7,6 +7,7 @@ from loguru import logger
from exo.master.main import Master
from exo.routing.router import get_node_id_keypair
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
from exo.shared.types.commands import (
ChatCompletion,
@@ -23,7 +24,6 @@ from exo.shared.types.events import (
TaskCreated,
)
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.profiling import (
MemoryUsage,
)
@@ -109,9 +109,8 @@ async def test_master():
command=(
PlaceInstance(
command_id=CommandId(),
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=ModelId("llama-3.2-1b"),
- pretty_name="Llama 3.2 1B",
n_layers=16,
storage_size=Memory.from_bytes(678948),
hidden_size=7168,
@@ -167,9 +166,8 @@ async def test_master():
start_layer=0,
end_layer=16,
n_layers=16,
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=ModelId("llama-3.2-1b"),
- pretty_name="Llama 3.2 1B",
n_layers=16,
storage_size=Memory.from_bytes(678948),
hidden_size=7168,
diff --git a/src/exo/master/tests/test_placement.py b/src/exo/master/tests/test_placement.py
index 2a32892b..767f08a3 100644
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -10,12 +10,12 @@ from exo.master.tests.conftest import (
create_rdma_connection,
create_socket_connection,
)
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.topology import Topology
from exo.shared.types.commands import PlaceInstance
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.events import InstanceCreated, InstanceDeleted
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
from exo.shared.types.topology import Connection, SocketConnection
@@ -43,21 +43,20 @@ def instance() -> Instance:
@pytest.fixture
-def model_meta() -> ModelMetadata:
- return ModelMetadata(
+def model_card() -> ModelCard:
+ return ModelCard(
model_id=ModelId("test-model"),
storage_size=Memory.from_kb(1000),
- pretty_name="Test Model",
n_layers=10,
hidden_size=30,
supports_tensor=True,
)
-def place_instance_command(model_meta: ModelMetadata) -> PlaceInstance:
+def place_instance_command(model_card: ModelCard) -> PlaceInstance:
return PlaceInstance(
command_id=CommandId(),
- model_meta=model_meta,
+ model_card=model_card,
sharding=Sharding.Pipeline,
instance_meta=InstanceMeta.MlxRing,
min_nodes=1,
@@ -76,16 +75,16 @@ def test_get_instance_placements_create_instance(
available_memory: tuple[int, int, int],
total_layers: int,
expected_layers: tuple[int, int, int],
- model_meta: ModelMetadata,
+ model_card: ModelCard,
):
# arrange
- model_meta.n_layers = total_layers
- model_meta.storage_size.in_bytes = sum(
+ model_card.n_layers = total_layers
+ model_card.storage_size.in_bytes = sum(
available_memory
) # make it exactly fit across all nodes
topology = Topology()
- cic = place_instance_command(model_meta)
+ cic = place_instance_command(model_card)
node_id_a = NodeId()
node_id_b = NodeId()
node_id_c = NodeId()
@@ -137,7 +136,7 @@ def test_get_instance_placements_create_instance(
assert len(placements) == 1
instance_id = list(placements.keys())[0]
instance = placements[instance_id]
- assert instance.shard_assignments.model_id == model_meta.model_id
+ assert instance.shard_assignments.model_id == model_card.model_id
runner_id_a = instance.shard_assignments.node_to_runner[node_id_a]
runner_id_b = instance.shard_assignments.node_to_runner[node_id_b]
@@ -164,10 +163,9 @@ def test_get_instance_placements_one_node_exact_fit() -> None:
node_memory = {node_id: create_node_memory(1000 * 1024)}
node_network = {node_id: create_node_network()}
cic = place_instance_command(
- ModelMetadata(
+ ModelCard(
model_id=ModelId("test-model"),
storage_size=Memory.from_kb(1000),
- pretty_name="Test Model",
n_layers=10,
hidden_size=1000,
supports_tensor=True,
@@ -191,10 +189,9 @@ def test_get_instance_placements_one_node_fits_with_extra_memory() -> None:
node_memory = {node_id: create_node_memory(1001 * 1024)}
node_network = {node_id: create_node_network()}
cic = place_instance_command(
- ModelMetadata(
+ ModelCard(
model_id=ModelId("test-model"),
storage_size=Memory.from_kb(1000),
- pretty_name="Test Model",
n_layers=10,
hidden_size=1000,
supports_tensor=True,
@@ -218,10 +215,9 @@ def test_get_instance_placements_one_node_not_fit() -> None:
node_memory = {node_id: create_node_memory(1000 * 1024)}
node_network = {node_id: create_node_network()}
cic = place_instance_command(
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=ModelId("test-model"),
storage_size=Memory.from_kb(1001),
- pretty_name="Test Model",
n_layers=10,
hidden_size=1000,
supports_tensor=True,
@@ -275,12 +271,14 @@ def test_get_transition_events_delete_instance(instance: Instance):
def test_placement_selects_leaf_nodes(
- model_meta: ModelMetadata,
+ model_card: ModelCard,
):
# arrange
topology = Topology()
- model_meta.storage_size = Memory.from_bytes(1000)
+ # Model requires more than any single node but fits within a 3-node cycle
+ model_card.storage_size.in_bytes = 1500
+ model_card.n_layers = 12
node_id_a = NodeId()
node_id_b = NodeId()
@@ -325,7 +323,7 @@ def test_placement_selects_leaf_nodes(
Connection(source=node_id_d, sink=node_id_c, edge=create_socket_connection(1))
)
- cic = place_instance_command(model_meta=model_meta)
+ cic = place_instance_command(model_card=model_card)
# act
placements = place_instance(cic, topology, {}, node_memory, node_network)
@@ -344,12 +342,12 @@ def test_placement_selects_leaf_nodes(
def test_tensor_rdma_backend_connectivity_matrix(
- model_meta: ModelMetadata,
+ model_card: ModelCard,
):
# arrange
topology = Topology()
- model_meta.n_layers = 12
- model_meta.storage_size.in_bytes = 1500
+ model_card.n_layers = 12
+ model_card.storage_size.in_bytes = 1500
node_a = NodeId()
node_b = NodeId()
@@ -411,7 +409,7 @@ def test_tensor_rdma_backend_connectivity_matrix(
sharding=Sharding.Tensor,
instance_meta=InstanceMeta.MlxJaccl,
command_id=CommandId(),
- model_meta=model_meta,
+ model_card=model_card,
min_nodes=1,
)
diff --git a/src/exo/master/tests/test_placement_utils.py b/src/exo/master/tests/test_placement_utils.py
index 867cfdd3..d26cf0c6 100644
--- a/src/exo/master/tests/test_placement_utils.py
+++ b/src/exo/master/tests/test_placement_utils.py
@@ -12,10 +12,10 @@ from exo.master.tests.conftest import (
create_node_memory,
create_socket_connection,
)
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.topology import Topology
from exo.shared.types.common import Host, NodeId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.profiling import (
NetworkInterfaceInfo,
NodeNetworkInfo,
@@ -232,9 +232,8 @@ def test_get_shard_assignments(
node_c_id: node_c_mem,
}
- model_meta = ModelMetadata(
+ model_card = ModelCard(
model_id=ModelId("test-model"),
- pretty_name="Test Model",
n_layers=total_layers,
storage_size=Memory.from_kb(1000),
hidden_size=1000,
@@ -248,7 +247,7 @@ def test_get_shard_assignments(
# act
shard_assignments = get_shard_assignments(
- model_meta, selected_cycle, Sharding.Pipeline, node_memory=node_memory
+ model_card, selected_cycle, Sharding.Pipeline, node_memory=node_memory
)
# assert
@@ -512,9 +511,8 @@ def test_get_shard_assignments_insufficient_memory_raises():
node_c_id: node_c_mem,
}
- model_meta = ModelMetadata(
+ model_card = ModelCard(
model_id=ModelId("test-model"),
- pretty_name="Test Model",
n_layers=20,
storage_size=Memory.from_kb(1000),
hidden_size=1000,
@@ -525,5 +523,5 @@ def test_get_shard_assignments_insufficient_memory_raises():
with pytest.raises(ValueError, match="insufficient memory"):
get_shard_assignments(
- model_meta, selected_cycle, Sharding.Pipeline, node_memory
+ model_card, selected_cycle, Sharding.Pipeline, node_memory
)
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index 4d091dca..e93d322f 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -1,613 +1,310 @@
+from pydantic import PositiveInt
+
+from exo.shared.types.common import Id
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.utils.pydantic_ext import CamelCaseModel
+class ModelId(Id):
+ def normalize(self) -> str:
+ return self.replace("/", "--")
+
+ def short(self) -> str:
+ return self.split("/")[-1]
+
+
class ModelCard(CamelCaseModel):
- short_id: str
model_id: ModelId
- name: str
- description: str
- tags: list[str]
- metadata: ModelMetadata
+ storage_size: Memory
+ n_layers: PositiveInt
+ hidden_size: PositiveInt
+ supports_tensor: bool
MODEL_CARDS: dict[str, ModelCard] = {
# deepseek v3
"deepseek-v3.1-4bit": ModelCard(
- short_id="deepseek-v3.1-4bit",
model_id=ModelId("mlx-community/DeepSeek-V3.1-4bit"),
- name="DeepSeek V3.1 (4-bit)",
- description="""DeepSeek V3.1 is a large language model trained on the DeepSeek V3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/DeepSeek-V3.1-4bit"),
- pretty_name="DeepSeek V3.1 (4-bit)",
- storage_size=Memory.from_gb(378),
- n_layers=61,
- hidden_size=7168,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(378),
+ n_layers=61,
+ hidden_size=7168,
+ supports_tensor=True,
),
"deepseek-v3.1-8bit": ModelCard(
- short_id="deepseek-v3.1-8bit",
model_id=ModelId("mlx-community/DeepSeek-V3.1-8bit"),
- name="DeepSeek V3.1 (8-bit)",
- description="""DeepSeek V3.1 is a large language model trained on the DeepSeek V3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/DeepSeek-V3.1-8bit"),
- pretty_name="DeepSeek V3.1 (8-bit)",
- storage_size=Memory.from_gb(713),
- n_layers=61,
- hidden_size=7168,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(713),
+ n_layers=61,
+ hidden_size=7168,
+ supports_tensor=True,
),
# kimi k2
"kimi-k2-instruct-4bit": ModelCard(
- short_id="kimi-k2-instruct-4bit",
model_id=ModelId("mlx-community/Kimi-K2-Instruct-4bit"),
- name="Kimi K2 Instruct (4-bit)",
- description="""Kimi K2 is a large language model trained on the Kimi K2 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Kimi-K2-Instruct-4bit"),
- pretty_name="Kimi K2 Instruct (4-bit)",
- storage_size=Memory.from_gb(578),
- n_layers=61,
- hidden_size=7168,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(578),
+ n_layers=61,
+ hidden_size=7168,
+ supports_tensor=True,
),
"kimi-k2-thinking": ModelCard(
- short_id="kimi-k2-thinking",
model_id=ModelId("mlx-community/Kimi-K2-Thinking"),
- name="Kimi K2 Thinking (4-bit)",
- description="""Kimi K2 Thinking is the latest, most capable version of open-source thinking model.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Kimi-K2-Thinking"),
- pretty_name="Kimi K2 Thinking (4-bit)",
- storage_size=Memory.from_gb(658),
- n_layers=61,
- hidden_size=7168,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(658),
+ n_layers=61,
+ hidden_size=7168,
+ supports_tensor=True,
),
# llama-3.1
"llama-3.1-8b": ModelCard(
- short_id="llama-3.1-8b",
model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"),
- name="Llama 3.1 8B (4-bit)",
- description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"),
- pretty_name="Llama 3.1 8B (4-bit)",
- storage_size=Memory.from_mb(4423),
- n_layers=32,
- hidden_size=4096,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(4423),
+ n_layers=32,
+ hidden_size=4096,
+ supports_tensor=True,
),
"llama-3.1-8b-8bit": ModelCard(
- short_id="llama-3.1-8b-8bit",
model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"),
- name="Llama 3.1 8B (8-bit)",
- description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-8bit"),
- pretty_name="Llama 3.1 8B (8-bit)",
- storage_size=Memory.from_mb(8540),
- n_layers=32,
- hidden_size=4096,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(8540),
+ n_layers=32,
+ hidden_size=4096,
+ supports_tensor=True,
),
"llama-3.1-8b-bf16": ModelCard(
- short_id="llama-3.1-8b-bf16",
model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"),
- name="Llama 3.1 8B (BF16)",
- description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-bf16"),
- pretty_name="Llama 3.1 8B (BF16)",
- storage_size=Memory.from_mb(16100),
- n_layers=32,
- hidden_size=4096,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(16100),
+ n_layers=32,
+ hidden_size=4096,
+ supports_tensor=True,
),
"llama-3.1-70b": ModelCard(
- short_id="llama-3.1-70b",
model_id=ModelId("mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"),
- name="Llama 3.1 70B (4-bit)",
- description="""Llama 3.1 is a large language model trained on the Llama 3.1 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Meta-Llama-3.1-70B-Instruct-4bit"),
- pretty_name="Llama 3.1 70B (4-bit)",
- storage_size=Memory.from_mb(38769),
- n_layers=80,
- hidden_size=8192,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(38769),
+ n_layers=80,
+ hidden_size=8192,
+ supports_tensor=True,
),
# llama-3.2
"llama-3.2-1b": ModelCard(
- short_id="llama-3.2-1b",
model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
- name="Llama 3.2 1B (4-bit)",
- description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Llama-3.2-1B-Instruct-4bit"),
- pretty_name="Llama 3.2 1B (4-bit)",
- storage_size=Memory.from_mb(696),
- n_layers=16,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(696),
+ n_layers=16,
+ hidden_size=2048,
+ supports_tensor=True,
),
"llama-3.2-3b": ModelCard(
- short_id="llama-3.2-3b",
model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-4bit"),
- name="Llama 3.2 3B (4-bit)",
- description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-4bit"),
- pretty_name="Llama 3.2 3B (4-bit)",
- storage_size=Memory.from_mb(1777),
- n_layers=28,
- hidden_size=3072,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(1777),
+ n_layers=28,
+ hidden_size=3072,
+ supports_tensor=True,
),
"llama-3.2-3b-8bit": ModelCard(
- short_id="llama-3.2-3b-8bit",
model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-8bit"),
- name="Llama 3.2 3B (8-bit)",
- description="""Llama 3.2 is a large language model trained on the Llama 3.2 dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Llama-3.2-3B-Instruct-8bit"),
- pretty_name="Llama 3.2 3B (8-bit)",
- storage_size=Memory.from_mb(3339),
- n_layers=28,
- hidden_size=3072,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(3339),
+ n_layers=28,
+ hidden_size=3072,
+ supports_tensor=True,
),
# llama-3.3
"llama-3.3-70b": ModelCard(
- short_id="llama-3.3-70b",
model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-4bit"),
- name="Llama 3.3 70B (4-bit)",
- description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-4bit"),
- pretty_name="Llama 3.3 70B",
- storage_size=Memory.from_mb(38769),
- n_layers=80,
- hidden_size=8192,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(38769),
+ n_layers=80,
+ hidden_size=8192,
+ supports_tensor=True,
),
"llama-3.3-70b-8bit": ModelCard(
- short_id="llama-3.3-70b-8bit",
model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-8bit"),
- name="Llama 3.3 70B (8-bit)",
- description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Llama-3.3-70B-Instruct-8bit"),
- pretty_name="Llama 3.3 70B (8-bit)",
- storage_size=Memory.from_mb(73242),
- n_layers=80,
- hidden_size=8192,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(73242),
+ n_layers=80,
+ hidden_size=8192,
+ supports_tensor=True,
),
"llama-3.3-70b-fp16": ModelCard(
- short_id="llama-3.3-70b-fp16",
model_id=ModelId("mlx-community/llama-3.3-70b-instruct-fp16"),
- name="Llama 3.3 70B (FP16)",
- description="""The Meta Llama 3.3 multilingual large language model (LLM) is an instruction tuned generative model in 70B (text in/text out)""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/llama-3.3-70b-instruct-fp16"),
- pretty_name="Llama 3.3 70B (FP16)",
- storage_size=Memory.from_mb(137695),
- n_layers=80,
- hidden_size=8192,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(137695),
+ n_layers=80,
+ hidden_size=8192,
+ supports_tensor=True,
),
# qwen3
"qwen3-0.6b": ModelCard(
- short_id="qwen3-0.6b",
model_id=ModelId("mlx-community/Qwen3-0.6B-4bit"),
- name="Qwen3 0.6B (4-bit)",
- description="""Qwen3 0.6B is a large language model trained on the Qwen3 0.6B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-0.6B-4bit"),
- pretty_name="Qwen3 0.6B (4-bit)",
- storage_size=Memory.from_mb(327),
- n_layers=28,
- hidden_size=1024,
- supports_tensor=False,
- ),
+ storage_size=Memory.from_mb(327),
+ n_layers=28,
+ hidden_size=1024,
+ supports_tensor=False,
),
"qwen3-0.6b-8bit": ModelCard(
- short_id="qwen3-0.6b-8bit",
model_id=ModelId("mlx-community/Qwen3-0.6B-8bit"),
- name="Qwen3 0.6B (8-bit)",
- description="""Qwen3 0.6B is a large language model trained on the Qwen3 0.6B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-0.6B-8bit"),
- pretty_name="Qwen3 0.6B (8-bit)",
- storage_size=Memory.from_mb(666),
- n_layers=28,
- hidden_size=1024,
- supports_tensor=False,
- ),
+ storage_size=Memory.from_mb(666),
+ n_layers=28,
+ hidden_size=1024,
+ supports_tensor=False,
),
"qwen3-30b": ModelCard(
- short_id="qwen3-30b",
model_id=ModelId("mlx-community/Qwen3-30B-A3B-4bit"),
- name="Qwen3 30B A3B (4-bit)",
- description="""Qwen3 30B is a large language model trained on the Qwen3 30B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-30B-A3B-4bit"),
- pretty_name="Qwen3 30B A3B (4-bit)",
- storage_size=Memory.from_mb(16797),
- n_layers=48,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(16797),
+ n_layers=48,
+ hidden_size=2048,
+ supports_tensor=True,
),
"qwen3-30b-8bit": ModelCard(
- short_id="qwen3-30b-8bit",
model_id=ModelId("mlx-community/Qwen3-30B-A3B-8bit"),
- name="Qwen3 30B A3B (8-bit)",
- description="""Qwen3 30B is a large language model trained on the Qwen3 30B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-30B-A3B-8bit"),
- pretty_name="Qwen3 30B A3B (8-bit)",
- storage_size=Memory.from_mb(31738),
- n_layers=48,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(31738),
+ n_layers=48,
+ hidden_size=2048,
+ supports_tensor=True,
),
"qwen3-80b-a3B-4bit": ModelCard(
- short_id="qwen3-80b-a3B-4bit",
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"),
- name="Qwen3 80B A3B (4-bit)",
- description="""Qwen3 80B""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-4bit"),
- pretty_name="Qwen3 80B A3B (4-bit)",
- storage_size=Memory.from_mb(44800),
- n_layers=48,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(44800),
+ n_layers=48,
+ hidden_size=2048,
+ supports_tensor=True,
),
"qwen3-80b-a3B-8bit": ModelCard(
- short_id="qwen3-80b-a3B-8bit",
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"),
- name="Qwen3 80B A3B (8-bit)",
- description="""Qwen3 80B""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Instruct-8bit"),
- pretty_name="Qwen3 80B A3B (8-bit)",
- storage_size=Memory.from_mb(84700),
- n_layers=48,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(84700),
+ n_layers=48,
+ hidden_size=2048,
+ supports_tensor=True,
),
"qwen3-80b-a3B-thinking-4bit": ModelCard(
- short_id="qwen3-80b-a3B-thinking-4bit",
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"),
- name="Qwen3 80B A3B Thinking (4-bit)",
- description="""Qwen3 80B Reasoning model""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-4bit"),
- pretty_name="Qwen3 80B A3B (4-bit)",
- storage_size=Memory.from_mb(84700),
- n_layers=48,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(84700),
+ n_layers=48,
+ hidden_size=2048,
+ supports_tensor=True,
),
"qwen3-80b-a3B-thinking-8bit": ModelCard(
- short_id="qwen3-80b-a3B-thinking-8bit",
model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"),
- name="Qwen3 80B A3B Thinking (8-bit)",
- description="""Qwen3 80B Reasoning model""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-Next-80B-A3B-Thinking-8bit"),
- pretty_name="Qwen3 80B A3B (8-bit)",
- storage_size=Memory.from_mb(84700),
- n_layers=48,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_mb(84700),
+ n_layers=48,
+ hidden_size=2048,
+ supports_tensor=True,
),
"qwen3-235b-a22b-4bit": ModelCard(
- short_id="qwen3-235b-a22b-4bit",
model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"),
- name="Qwen3 235B A22B (4-bit)",
- description="""Qwen3 235B (Active 22B) is a large language model trained on the Qwen3 235B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-4bit"),
- pretty_name="Qwen3 235B A22B (4-bit)",
- storage_size=Memory.from_gb(132),
- n_layers=94,
- hidden_size=4096,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(132),
+ n_layers=94,
+ hidden_size=4096,
+ supports_tensor=True,
),
"qwen3-235b-a22b-8bit": ModelCard(
- short_id="qwen3-235b-a22b-8bit",
model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"),
- name="Qwen3 235B A22B (8-bit)",
- description="""Qwen3 235B (Active 22B) is a large language model trained on the Qwen3 235B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-235B-A22B-Instruct-2507-8bit"),
- pretty_name="Qwen3 235B A22B (8-bit)",
- storage_size=Memory.from_gb(250),
- n_layers=94,
- hidden_size=4096,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(250),
+ n_layers=94,
+ hidden_size=4096,
+ supports_tensor=True,
),
"qwen3-coder-480b-a35b-4bit": ModelCard(
- short_id="qwen3-coder-480b-a35b-4bit",
model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"),
- name="Qwen3 Coder 480B A35B (4-bit)",
- description="""Qwen3 Coder 480B (Active 35B) is a large language model trained on the Qwen3 Coder 480B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-4bit"),
- pretty_name="Qwen3 Coder 480B A35B (4-bit)",
- storage_size=Memory.from_gb(270),
- n_layers=62,
- hidden_size=6144,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(270),
+ n_layers=62,
+ hidden_size=6144,
+ supports_tensor=True,
),
"qwen3-coder-480b-a35b-8bit": ModelCard(
- short_id="qwen3-coder-480b-a35b-8bit",
model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit"),
- name="Qwen3 Coder 480B A35B (8-bit)",
- description="""Qwen3 Coder 480B (Active 35B) is a large language model trained on the Qwen3 Coder 480B dataset.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/Qwen3-Coder-480B-A35B-Instruct-8bit"),
- pretty_name="Qwen3 Coder 480B A35B (8-bit)",
- storage_size=Memory.from_gb(540),
- n_layers=62,
- hidden_size=6144,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(540),
+ n_layers=62,
+ hidden_size=6144,
+ supports_tensor=True,
),
# gpt-oss
"gpt-oss-120b-MXFP4-Q8": ModelCard(
- short_id="gpt-oss-120b-MXFP4-Q8",
model_id=ModelId("mlx-community/gpt-oss-120b-MXFP4-Q8"),
- name="GPT-OSS 120B (MXFP4-Q8, MLX)",
- description="""OpenAI's GPT-OSS 120B is a 117B-parameter Mixture-of-Experts model designed for high-reasoning and general-purpose use; this variant is a 4-bit MLX conversion for Apple Silicon.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/gpt-oss-120b-MXFP4-Q8"),
- pretty_name="GPT-OSS 120B (MXFP4-Q8, MLX)",
- storage_size=Memory.from_kb(68_996_301),
- n_layers=36,
- hidden_size=2880,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_kb(68_996_301),
+ n_layers=36,
+ hidden_size=2880,
+ supports_tensor=True,
),
"gpt-oss-20b-MXFP4-Q8": ModelCard(
- short_id="gpt-oss-20b-MXFP4-Q8",
model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
- name="GPT-OSS 20B (MXFP4-Q8, MLX)",
- description="""OpenAI's GPT-OSS 20B is a medium-sized MoE model for lower-latency and local or specialized use cases; this variant is a 4-bit MLX conversion for Apple Silicon.""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
- pretty_name="GPT-OSS 20B (MXFP4-Q8, MLX)",
- storage_size=Memory.from_kb(11_744_051),
- n_layers=24,
- hidden_size=2880,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_kb(11_744_051),
+ n_layers=24,
+ hidden_size=2880,
+ supports_tensor=True,
),
# glm 4.5
"glm-4.5-air-8bit": ModelCard(
# Needs to be quantized g32 or g16 to work with tensor parallel
- short_id="glm-4.5-air-8bit",
model_id=ModelId("mlx-community/GLM-4.5-Air-8bit"),
- name="GLM 4.5 Air 8bit",
- description="""GLM 4.5 Air 8bit""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.5-Air-8bit"),
- pretty_name="GLM 4.5 Air 8bit",
- storage_size=Memory.from_gb(114),
- n_layers=46,
- hidden_size=4096,
- supports_tensor=False,
- ),
+ storage_size=Memory.from_gb(114),
+ n_layers=46,
+ hidden_size=4096,
+ supports_tensor=False,
),
"glm-4.5-air-bf16": ModelCard(
- short_id="glm-4.5-air-bf16",
model_id=ModelId("mlx-community/GLM-4.5-Air-bf16"),
- name="GLM 4.5 Air bf16",
- description="""GLM 4.5 Air bf16""",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.5-Air-bf16"),
- pretty_name="GLM 4.5 Air bf16",
- storage_size=Memory.from_gb(214),
- n_layers=46,
- hidden_size=4096,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(214),
+ n_layers=46,
+ hidden_size=4096,
+ supports_tensor=True,
),
# glm 4.7
"glm-4.7-4bit": ModelCard(
- short_id="glm-4.7-4bit",
model_id=ModelId("mlx-community/GLM-4.7-4bit"),
- name="GLM 4.7 4bit",
- description="GLM 4.7 4bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.7-4bit"),
- pretty_name="GLM 4.7 4bit",
- storage_size=Memory.from_bytes(198556925568),
- n_layers=91,
- hidden_size=5120,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_bytes(198556925568),
+ n_layers=91,
+ hidden_size=5120,
+ supports_tensor=True,
),
"glm-4.7-6bit": ModelCard(
- short_id="glm-4.7-6bit",
model_id=ModelId("mlx-community/GLM-4.7-6bit"),
- name="GLM 4.7 6bit",
- description="GLM 4.7 6bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.7-6bit"),
- pretty_name="GLM 4.7 6bit",
- storage_size=Memory.from_bytes(286737579648),
- n_layers=91,
- hidden_size=5120,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_bytes(286737579648),
+ n_layers=91,
+ hidden_size=5120,
+ supports_tensor=True,
),
"glm-4.7-8bit-gs32": ModelCard(
- short_id="glm-4.7-8bit-gs32",
model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"),
- name="GLM 4.7 8bit (gs32)",
- description="GLM 4.7 8bit (gs32)",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"),
- pretty_name="GLM 4.7 8bit (gs32)",
- storage_size=Memory.from_bytes(396963397248),
- n_layers=91,
- hidden_size=5120,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_bytes(396963397248),
+ n_layers=91,
+ hidden_size=5120,
+ supports_tensor=True,
),
# glm 4.7 flash
"glm-4.7-flash-4bit": ModelCard(
- short_id="glm-4.7-flash-4bit",
model_id=ModelId("mlx-community/GLM-4.7-Flash-4bit"),
- name="GLM 4.7 Flash 4bit",
- description="GLM 4.7 Flash 4bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.7-Flash-4bit"),
- pretty_name="GLM 4.7 Flash 4bit",
- storage_size=Memory.from_gb(18),
- n_layers=47,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(18),
+ n_layers=47,
+ hidden_size=2048,
+ supports_tensor=True,
),
"glm-4.7-flash-5bit": ModelCard(
- short_id="glm-4.7-flash-5bit",
model_id=ModelId("mlx-community/GLM-4.7-Flash-5bit"),
- name="GLM 4.7 Flash 5bit",
- description="GLM 4.7 Flash 5bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.7-Flash-5bit"),
- pretty_name="GLM 4.7 Flash 5bit",
- storage_size=Memory.from_gb(21),
- n_layers=47,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(21),
+ n_layers=47,
+ hidden_size=2048,
+ supports_tensor=True,
),
"glm-4.7-flash-6bit": ModelCard(
- short_id="glm-4.7-flash-6bit",
model_id=ModelId("mlx-community/GLM-4.7-Flash-6bit"),
- name="GLM 4.7 Flash 6bit",
- description="GLM 4.7 Flash 6bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.7-Flash-6bit"),
- pretty_name="GLM 4.7 Flash 6bit",
- storage_size=Memory.from_gb(25),
- n_layers=47,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(25),
+ n_layers=47,
+ hidden_size=2048,
+ supports_tensor=True,
),
"glm-4.7-flash-8bit": ModelCard(
- short_id="glm-4.7-flash-8bit",
model_id=ModelId("mlx-community/GLM-4.7-Flash-8bit"),
- name="GLM 4.7 Flash 8bit",
- description="GLM 4.7 Flash 8bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/GLM-4.7-Flash-8bit"),
- pretty_name="GLM 4.7 Flash 8bit",
- storage_size=Memory.from_gb(32),
- n_layers=47,
- hidden_size=2048,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_gb(32),
+ n_layers=47,
+ hidden_size=2048,
+ supports_tensor=True,
),
# minimax-m2
"minimax-m2.1-8bit": ModelCard(
- short_id="minimax-m2.1-8bit",
model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"),
- name="MiniMax M2.1 8bit",
- description="MiniMax M2.1 8bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"),
- pretty_name="MiniMax M2.1 8bit",
- storage_size=Memory.from_bytes(242986745856),
- n_layers=61,
- hidden_size=3072,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_bytes(242986745856),
+ n_layers=61,
+ hidden_size=3072,
+ supports_tensor=True,
),
"minimax-m2.1-3bit": ModelCard(
- short_id="minimax-m2.1-3bit",
model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"),
- name="MiniMax M2.1 3bit",
- description="MiniMax M2.1 3bit",
- tags=[],
- metadata=ModelMetadata(
- model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"),
- pretty_name="MiniMax M2.1 3bit",
- storage_size=Memory.from_bytes(100086644736),
- n_layers=61,
- hidden_size=3072,
- supports_tensor=True,
- ),
+ storage_size=Memory.from_bytes(100086644736),
+ n_layers=61,
+ hidden_size=3072,
+ supports_tensor=True,
),
}
diff --git a/src/exo/shared/models/model_meta.py b/src/exo/shared/models/model_meta.py
index 43f84c88..87b4594b 100644
--- a/src/exo/shared/models/model_meta.py
+++ b/src/exo/shared/models/model_meta.py
@@ -6,9 +6,8 @@ from huggingface_hub import model_info
from loguru import logger
from pydantic import BaseModel, Field
-from exo.shared.models.model_cards import MODEL_CARDS
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.worker.download.download_utils import (
ModelSafetensorsIndex,
download_file_with_retry,
@@ -92,18 +91,18 @@ async def get_safetensors_size(model_id: str) -> Memory:
return Memory.from_bytes(info.safetensors.total)
-_model_meta_cache: dict[str, ModelMetadata] = {}
+_model_card_cache: dict[str, ModelCard] = {}
-async def get_model_meta(model_id: str) -> ModelMetadata:
- if model_id in _model_meta_cache:
- return _model_meta_cache[model_id]
- model_meta = await _get_model_meta(model_id)
- _model_meta_cache[model_id] = model_meta
- return model_meta
+async def get_model_card(model_id: str) -> ModelCard:
+ if model_id in _model_card_cache:
+ return _model_card_cache[model_id]
+ model_card = await _get_model_card(model_id)
+ _model_card_cache[model_id] = model_card
+ return model_card
-async def _get_model_meta(model_id: str) -> ModelMetadata:
+async def _get_model_card(model_id: str) -> ModelCard:
"""Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
config_data = await get_config_data(model_id)
num_layers = config_data.layer_count
@@ -113,14 +112,11 @@ async def _get_model_meta(model_id: str) -> ModelMetadata:
None,
)
- return ModelMetadata(
+ return ModelCard(
model_id=ModelId(model_id),
- pretty_name=model_card.name if model_card is not None else model_id,
storage_size=mem_size_bytes,
n_layers=num_layers,
hidden_size=config_data.hidden_size or 0,
# TODO: all custom models currently do not support tensor. We could add a dynamic test for this?
- supports_tensor=model_card.metadata.supports_tensor
- if model_card is not None
- else False,
+ supports_tensor=model_card.supports_tensor if model_card is not None else False,
)
diff --git a/src/exo/shared/tests/conftest.py b/src/exo/shared/tests/conftest.py
index 760e3452..7ed28ad9 100644
--- a/src/exo/shared/tests/conftest.py
+++ b/src/exo/shared/tests/conftest.py
@@ -7,8 +7,8 @@ import pytest
from _pytest.logging import LogCaptureFixture
from loguru import logger
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
@@ -31,9 +31,8 @@ def get_pipeline_shard_metadata(
model_id: ModelId, device_rank: int, world_size: int = 1
) -> ShardMetadata:
return PipelineShardMetadata(
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=model_id,
- pretty_name=str(model_id),
storage_size=Memory.from_mb(100000),
n_layers=32,
hidden_size=1000,
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 3ae6917a..eba83de0 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -4,9 +4,9 @@ from typing import Any, Literal
from pydantic import BaseModel, Field, field_validator
from pydantic_core import PydanticUseDefault
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.common import CommandId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
@@ -206,7 +206,7 @@ class DeleteInstanceTaskParams(BaseModel):
class CreateInstanceResponse(BaseModel):
message: str
command_id: CommandId
- model_meta: ModelMetadata
+ model_card: ModelCard
class DeleteInstanceResponse(BaseModel):
diff --git a/src/exo/shared/types/chunks.py b/src/exo/shared/types/chunks.py
index 4785b172..5e5cfed1 100644
--- a/src/exo/shared/types/chunks.py
+++ b/src/exo/shared/types/chunks.py
@@ -1,10 +1,10 @@
from enum import Enum
+from exo.shared.models.model_cards import ModelId
from exo.shared.types.api import GenerationStats
from exo.utils.pydantic_ext import TaggedModel
from .api import FinishReason
-from .models import ModelId
class ChunkType(str, Enum):
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index 5d8a5026..c264a592 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -1,8 +1,8 @@
from pydantic import Field
+from exo.shared.models.model_cards import ModelCard
from exo.shared.types.api import ChatCompletionTaskParams
from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.models import ModelMetadata
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -21,7 +21,7 @@ class ChatCompletion(BaseCommand):
class PlaceInstance(BaseCommand):
- model_meta: ModelMetadata
+ model_card: ModelCard
sharding: Sharding
instance_meta: InstanceMeta
min_nodes: int
diff --git a/src/exo/shared/types/common.py b/src/exo/shared/types/common.py
index 42b682dc..774210f5 100644
--- a/src/exo/shared/types/common.py
+++ b/src/exo/shared/types/common.py
@@ -16,7 +16,9 @@ class Id(str):
cls, _source: type, handler: GetCoreSchemaHandler
) -> core_schema.CoreSchema:
# Just use a plain string schema
- return core_schema.str_schema()
+ return core_schema.no_info_after_validator_function(
+ cls, core_schema.str_schema()
+ )
class NodeId(Id):
diff --git a/src/exo/shared/types/models.py b/src/exo/shared/types/models.py
deleted file mode 100644
index e99a7667..00000000
--- a/src/exo/shared/types/models.py
+++ /dev/null
@@ -1,18 +0,0 @@
-from pydantic import PositiveInt
-
-from exo.shared.types.common import Id
-from exo.shared.types.memory import Memory
-from exo.utils.pydantic_ext import CamelCaseModel
-
-
-class ModelId(Id):
- pass
-
-
-class ModelMetadata(CamelCaseModel):
- model_id: ModelId
- pretty_name: str
- storage_size: Memory
- n_layers: PositiveInt
- hidden_size: PositiveInt
- supports_tensor: bool
diff --git a/src/exo/shared/types/worker/runners.py b/src/exo/shared/types/worker/runners.py
index ff5026e3..7746949d 100644
--- a/src/exo/shared/types/worker/runners.py
+++ b/src/exo/shared/types/worker/runners.py
@@ -2,8 +2,8 @@ from collections.abc import Mapping
from pydantic import model_validator
+from exo.shared.models.model_cards import ModelId
from exo.shared.types.common import Id, NodeId
-from exo.shared.types.models import ModelId
from exo.shared.types.worker.shards import ShardMetadata
from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
diff --git a/src/exo/shared/types/worker/shards.py b/src/exo/shared/types/worker/shards.py
index e8e86730..8bb23a57 100644
--- a/src/exo/shared/types/worker/shards.py
+++ b/src/exo/shared/types/worker/shards.py
@@ -2,7 +2,7 @@ from enum import Enum
from pydantic import Field
-from exo.shared.types.models import ModelMetadata
+from exo.shared.models.model_cards import ModelCard
from exo.utils.pydantic_ext import TaggedModel
@@ -17,7 +17,7 @@ class BaseShardMetadata(TaggedModel):
Replaces previous `Shard` object.
"""
- model_meta: ModelMetadata
+ model_card: ModelCard
device_rank: int
world_size: int
@@ -41,7 +41,7 @@ class BaseShardMetadata(TaggedModel):
def __hash__(self) -> int:
return hash(
(
- self.model_meta.model_id,
+ self.model_card.model_id,
self.start_layer,
self.end_layer,
self.n_layers,
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index 05d2ff73..0165e59b 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -460,10 +460,10 @@ async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
# (iii) Tensor parallel requires all files.
return ["*"]
try:
- weight_map = await get_weight_map(str(shard.model_meta.model_id))
+ weight_map = await get_weight_map(str(shard.model_card.model_id))
return get_allow_patterns(weight_map, shard)
except Exception:
- logger.error(f"Error getting weight map for {shard.model_meta.model_id=}")
+ logger.error(f"Error getting weight map for {shard.model_card.model_id=}")
logger.error(traceback.format_exc())
return ["*"]
@@ -532,18 +532,18 @@ async def download_shard(
allow_patterns: list[str] | None = None,
) -> tuple[Path, RepoDownloadProgress]:
if not skip_download:
- logger.info(f"Downloading {shard.model_meta.model_id=}")
+ logger.info(f"Downloading {shard.model_card.model_id=}")
# Handle local paths
- if await aios.path.exists(str(shard.model_meta.model_id)):
- logger.info(f"Using local model path {shard.model_meta.model_id}")
- local_path = Path(str(shard.model_meta.model_id))
+ if await aios.path.exists(str(shard.model_card.model_id)):
+ logger.info(f"Using local model path {shard.model_card.model_id}")
+ local_path = Path(str(shard.model_card.model_id))
return local_path, await download_progress_for_local_path(
- str(shard.model_meta.model_id), shard, local_path
+ str(shard.model_card.model_id), shard, local_path
)
revision = "main"
- target_dir = await ensure_models_dir() / str(shard.model_meta.model_id).replace(
+ target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
"/", "--"
)
if not skip_download:
@@ -552,13 +552,13 @@ async def download_shard(
if not allow_patterns:
allow_patterns = await resolve_allow_patterns(shard)
- logger.info(f"Downloading {shard.model_meta.model_id=} with {allow_patterns=}")
+ logger.info(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
all_start_time = time.time()
# TODO: currently not recursive. Some models might require subdirectories - thus this will need to be changed.
# Update: <- This does not seem to be the case. Yay?
file_list = await fetch_file_list_with_cache(
- str(shard.model_meta.model_id), revision, recursive=True
+ str(shard.model_card.model_id), revision, recursive=True
)
filtered_file_list = list(
filter_repo_objects(
@@ -592,7 +592,7 @@ async def download_shard(
else timedelta(seconds=0)
)
file_progress[file.path] = RepoFileDownloadProgress(
- repo_id=str(shard.model_meta.model_id),
+ repo_id=str(shard.model_card.model_id),
repo_revision=revision,
file_path=file.path,
downloaded=Memory.from_bytes(curr_bytes),
@@ -609,7 +609,7 @@ async def download_shard(
shard,
calculate_repo_progress(
shard,
- str(shard.model_meta.model_id),
+ str(shard.model_card.model_id),
revision,
file_progress,
all_start_time,
@@ -619,7 +619,7 @@ async def download_shard(
for file in filtered_file_list:
downloaded_bytes = await get_downloaded_size(target_dir / file.path)
file_progress[file.path] = RepoFileDownloadProgress(
- repo_id=str(shard.model_meta.model_id),
+ repo_id=str(shard.model_card.model_id),
repo_revision=revision,
file_path=file.path,
downloaded=Memory.from_bytes(downloaded_bytes),
@@ -643,7 +643,7 @@ async def download_shard(
async def download_with_semaphore(file: FileListEntry) -> None:
async with semaphore:
await download_file_with_retry(
- str(shard.model_meta.model_id),
+ str(shard.model_card.model_id),
revision,
file.path,
target_dir,
@@ -657,7 +657,7 @@ async def download_shard(
*[download_with_semaphore(file) for file in filtered_file_list]
)
final_repo_progress = calculate_repo_progress(
- shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
+ shard, str(shard.model_card.model_id), revision, file_progress, all_start_time
)
await on_progress(shard, final_repo_progress)
if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
diff --git a/src/exo/worker/download/impl_shard_downloader.py b/src/exo/worker/download/impl_shard_downloader.py
index 1ca2e4e7..1610ad45 100644
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -4,7 +4,7 @@ from pathlib import Path
from typing import AsyncIterator, Callable
from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.models.model_meta import get_model_meta
+from exo.shared.models.model_meta import get_model_card
from exo.shared.types.worker.shards import (
PipelineShardMetadata,
ShardMetadata,
@@ -20,21 +20,21 @@ def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
async def build_base_shard(model_id: str) -> ShardMetadata:
- model_meta = await get_model_meta(model_id)
+ model_card = await get_model_card(model_id)
return PipelineShardMetadata(
- model_meta=model_meta,
+ model_card=model_card,
device_rank=0,
world_size=1,
start_layer=0,
- end_layer=model_meta.n_layers,
- n_layers=model_meta.n_layers,
+ end_layer=model_card.n_layers,
+ n_layers=model_card.n_layers,
)
async def build_full_shard(model_id: str) -> PipelineShardMetadata:
base_shard = await build_base_shard(model_id)
return PipelineShardMetadata(
- model_meta=base_shard.model_meta,
+ model_card=base_shard.model_card,
device_rank=base_shard.device_rank,
world_size=base_shard.world_size,
start_layer=base_shard.start_layer,
@@ -93,11 +93,11 @@ class CachedShardDownloader(ShardDownloader):
async def ensure_shard(
self, shard: ShardMetadata, config_only: bool = False
) -> Path:
- if (shard.model_meta.model_id, shard) in self.cache:
- return self.cache[(shard.model_meta.model_id, shard)]
+ if (shard.model_card.model_id, shard) in self.cache:
+ return self.cache[(shard.model_card.model_id, shard)]
target_dir = await self.shard_downloader.ensure_shard(shard, config_only)
- self.cache[(shard.model_meta.model_id, shard)] = target_dir
+ self.cache[(shard.model_card.model_id, shard)] = target_dir
return target_dir
async def get_shard_download_status(
diff --git a/src/exo/worker/download/shard_downloader.py b/src/exo/worker/download/shard_downloader.py
index 60a4692e..d964196a 100644
--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -5,8 +5,8 @@ from datetime import timedelta
from pathlib import Path
from typing import AsyncIterator, Callable
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.worker.shards import (
PipelineShardMetadata,
ShardMetadata,
@@ -86,9 +86,8 @@ NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(
repo_id="noop",
repo_revision="noop",
shard=PipelineShardMetadata(
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=ModelId("noop"),
- pretty_name="noope",
storage_size=Memory.from_bytes(0),
n_layers=1,
hidden_size=1,
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 192e5f1f..3b11ebc9 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -75,7 +75,7 @@ def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
return Memory.from_float_kb(
(model_shard_meta.end_layer - model_shard_meta.start_layer)
/ model_shard_meta.n_layers
- * model_shard_meta.model_meta.storage_size.in_kb
+ * model_shard_meta.model_card.storage_size.in_kb
/ (
1
if isinstance(model_shard_meta, PipelineShardMetadata)
@@ -206,7 +206,7 @@ def load_mlx_items(
) -> tuple[Model, TokenizerWrapper]:
if group is None:
logger.info(f"Single device used for {bound_instance.instance}")
- model_path = build_model_path(bound_instance.bound_shard.model_meta.model_id)
+ model_path = build_model_path(bound_instance.bound_shard.model_card.model_id)
start_time = time.perf_counter()
model, _ = load_model(model_path, strict=True)
end_time = time.perf_counter()
@@ -234,7 +234,7 @@ def shard_and_load(
group: Group,
on_timeout: TimeoutCallback | None = None,
) -> tuple[nn.Module, TokenizerWrapper]:
- model_path = build_model_path(shard_metadata.model_meta.model_id)
+ model_path = build_model_path(shard_metadata.model_card.model_id)
model, _ = load_model(model_path, lazy=True, strict=False)
logger.debug(model)
@@ -293,7 +293,7 @@ def shard_and_load(
def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
"""Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
- return load_tokenizer_for_model_id(shard_metadata.model_meta.model_id, model_path)
+ return load_tokenizer_for_model_id(shard_metadata.model_card.model_id, model_path)
def get_eos_token_ids_for_model(model_id: str) -> list[int] | None:
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index f389f920..d58cf47c 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -8,6 +8,7 @@ from loguru import logger
from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelId
from exo.shared.types.commands import ForwarderCommand, RequestEventLog
from exo.shared.types.common import NodeId, SessionId
from exo.shared.types.events import (
@@ -22,7 +23,6 @@ from exo.shared.types.events import (
TopologyEdgeCreated,
TopologyEdgeDeleted,
)
-from exo.shared.types.models import ModelId
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.state import State
from exo.shared.types.tasks import (
@@ -186,11 +186,11 @@ class Worker:
)
)
case DownloadModel(shard_metadata=shard):
- if shard.model_meta.model_id not in self.download_status:
+ if shard.model_card.model_id not in self.download_status:
progress = DownloadPending(
shard_metadata=shard, node_id=self.node_id
)
- self.download_status[shard.model_meta.model_id] = progress
+ self.download_status[shard.model_card.model_id] = progress
await self.event_sender.send(
NodeDownloadProgress(download_progress=progress)
)
@@ -205,7 +205,7 @@ class Worker:
node_id=self.node_id,
total_bytes=initial_progress.total_bytes,
)
- self.download_status[shard.model_meta.model_id] = progress
+ self.download_status[shard.model_card.model_id] = progress
await self.event_sender.send(
NodeDownloadProgress(download_progress=progress)
)
@@ -339,7 +339,7 @@ class Worker:
initial_progress
),
)
- self.download_status[task.shard_metadata.model_meta.model_id] = status
+ self.download_status[task.shard_metadata.model_card.model_id] = status
self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
last_progress_time = 0.0
@@ -356,7 +356,7 @@ class Worker:
node_id=self.node_id,
total_bytes=progress.total_bytes,
)
- self.download_status[shard.model_meta.model_id] = status
+ self.download_status[shard.model_card.model_id] = status
await self.event_sender.send(
NodeDownloadProgress(download_progress=status)
)
@@ -376,7 +376,7 @@ class Worker:
progress
),
)
- self.download_status[shard.model_meta.model_id] = status
+ self.download_status[shard.model_card.model_id] = status
await self.event_sender.send(
NodeDownloadProgress(download_progress=status)
)
@@ -478,7 +478,7 @@ class Worker:
else:
continue
- self.download_status[progress.shard.model_meta.model_id] = status
+ self.download_status[progress.shard.model_card.model_id] = status
await self.event_sender.send(
NodeDownloadProgress(download_progress=status)
)
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 70b2e02d..54ca1106 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -2,8 +2,8 @@
from collections.abc import Mapping, Sequence
+from exo.shared.models.model_cards import ModelId
from exo.shared.types.common import NodeId
-from exo.shared.types.models import ModelId
from exo.shared.types.tasks import (
ChatCompletion,
ConnectToGroup,
@@ -114,7 +114,7 @@ def _model_needs_download(
download_status: Mapping[ModelId, DownloadProgress],
) -> DownloadModel | None:
for runner in runners.values():
- model_id = runner.bound_instance.bound_shard.model_meta.model_id
+ model_id = runner.bound_instance.bound_shard.model_card.model_id
if isinstance(runner.status, RunnerIdle) and (
model_id not in download_status
or not isinstance(
@@ -191,7 +191,7 @@ def _load_model(
nid in global_download_status
and any(
isinstance(dp, DownloadCompleted)
- and dp.shard_metadata.model_meta.model_id == shard_assignments.model_id
+ and dp.shard_metadata.model_card.model_id == shard_assignments.model_id
for dp in global_download_status[nid]
)
for nid in shard_assignments.node_to_runner
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 792b9f30..dead1191 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -213,7 +213,7 @@ def main(
command_id=command_id,
chunk=TokenChunk(
idx=response.token,
- model=shard_metadata.model_meta.model_id,
+ model=shard_metadata.model_card.model_id,
text=response.text,
token_id=response.token,
finish_reason=response.finish_reason,
@@ -230,7 +230,7 @@ def main(
command_id=command_id,
chunk=TokenChunk(
idx=0,
- model=shard_metadata.model_meta.model_id,
+ model=shard_metadata.model_card.model_id,
text="",
token_id=0,
finish_reason="error",
diff --git a/src/exo/worker/tests/constants.py b/src/exo/worker/tests/constants.py
index 55b26177..05217a32 100644
--- a/src/exo/worker/tests/constants.py
+++ b/src/exo/worker/tests/constants.py
@@ -1,7 +1,7 @@
from typing import Final
+from exo.shared.models.model_cards import ModelId
from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.models import ModelId
from exo.shared.types.tasks import TaskId
from exo.shared.types.worker.instances import InstanceId, RunnerId
diff --git a/src/exo/worker/tests/unittests/conftest.py b/src/exo/worker/tests/unittests/conftest.py
index dea085fe..67b19f3c 100644
--- a/src/exo/worker/tests/unittests/conftest.py
+++ b/src/exo/worker/tests/unittests/conftest.py
@@ -1,8 +1,8 @@
from dataclasses import dataclass, field
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.common import NodeId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.tasks import BaseTask, TaskId
from exo.shared.types.worker.instances import (
BoundInstance,
@@ -32,9 +32,8 @@ def get_pipeline_shard_metadata(
model_id: ModelId, device_rank: int, world_size: int = 1
) -> ShardMetadata:
return PipelineShardMetadata(
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=model_id,
- pretty_name=str(model_id),
storage_size=Memory.from_mb(100000),
n_layers=32,
hidden_size=2048,
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index daafc3e0..16eeea6a 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -11,9 +11,9 @@ import mlx.core as mx
import mlx.nn as nn
from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.api import ChatCompletionMessage
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
from exo.shared.types.tasks import ChatCompletionTaskParams
from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
from exo.worker.engines.mlx import Model
@@ -81,9 +81,8 @@ def run_gpt_oss_pipeline_device(
start_layer, end_layer = layer_splits[rank]
shard_meta = PipelineShardMetadata(
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
- pretty_name="GPT-OSS 20B",
storage_size=Memory.from_gb(12),
n_layers=24,
hidden_size=2880,
@@ -151,9 +150,8 @@ def run_gpt_oss_tensor_parallel_device(
# For tensor parallelism, all devices run all layers
shard_meta = TensorShardMetadata(
- model_meta=ModelMetadata(
+ model_card=ModelCard(
model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
- pretty_name="GPT-OSS 20B",
storage_size=Memory.from_gb(12),
n_layers=24,
hidden_size=2880,
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
index 5eef4ede..ba6eacc2 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
@@ -76,13 +76,13 @@ def get_test_models() -> list[tuple[str, ModelCard]]:
"""Get a representative sample of models to test."""
# Pick one model from each family to test
families: dict[str, tuple[str, ModelCard]] = {}
- for short_id, card in MODEL_CARDS.items():
+ for _, card in MODEL_CARDS.items():
# Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
- parts = short_id.split("-")
+ parts = card.model_id.short().split("-")
family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
if family not in families:
- families[family] = (short_id, card)
+ families[family] = (card.model_id.short(), card)
return list(families.values())
diff --git a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
index 69211769..8d2e04a8 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
@@ -1,7 +1,7 @@
import exo.worker.plan as plan_mod
+from exo.shared.models.model_cards import ModelId
from exo.shared.types.common import NodeId
from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId
from exo.shared.types.tasks import LoadModel
from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress
from exo.shared.types.worker.instances import BoundInstance
diff --git a/tests/headless_runner.py b/tests/headless_runner.py
index 04ba78a6..354186a5 100644
--- a/tests/headless_runner.py
+++ b/tests/headless_runner.py
@@ -82,7 +82,7 @@ async def tb_detection():
send, recv = channel[GatheredInfo]()
ig = InfoGatherer(send)
with anyio.move_on_after(1):
- await ig._monitor_system_profiler() # pyright: ignore[reportPrivateUsage]
+ await ig._monitor_system_profiler_thunderbolt_data() # pyright: ignore[reportPrivateUsage]
with recv:
return recv.collect()
@@ -135,7 +135,7 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
else:
raise ValueError(f"{hn} not in {test.devs}")
- meta = MODEL_CARDS[test.model_id].metadata
+ card = MODEL_CARDS[test.model_id]
instance = MlxRingInstance(
instance_id=iid,
ephemeral_port=52416,
@@ -145,15 +145,15 @@ def ring_instance(test: Tests, iid: InstanceId, hn: str) -> Instance:
node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
runner_to_shard={
RunnerId(test.devs[i][0]): PipelineShardMetadata(
- model_meta=meta,
+ model_card=card,
device_rank=i,
world_size=world_size,
- start_layer=(meta.n_layers // world_size) * i,
+ start_layer=(card.n_layers // world_size) * i,
end_layer=min(
- meta.n_layers, (meta.n_layers // world_size) * (i + 1)
+ card.n_layers, (card.n_layers // world_size) * (i + 1)
),
- n_layers=min(meta.n_layers, (meta.n_layers // world_size) * (i + 1))
- - (meta.n_layers // world_size) * i,
+ n_layers=min(card.n_layers, (card.n_layers // world_size) * (i + 1))
+ - (card.n_layers // world_size) * i,
)
for i in range(world_size)
},
@@ -224,7 +224,7 @@ async def jaccl_backend(test: Tests):
def jaccl_instance(test: Tests, iid: InstanceId):
- meta = MODEL_CARDS[test.model_id].metadata
+ card = MODEL_CARDS[test.model_id]
world_size = len(test.devs)
return MlxJacclInstance(
@@ -239,12 +239,12 @@ def jaccl_instance(test: Tests, iid: InstanceId):
node_to_runner={NodeId(host[0]): RunnerId(host[0]) for host in test.devs},
runner_to_shard={
RunnerId(test.devs[i][0]): TensorShardMetadata(
- model_meta=meta,
+ model_card=card,
device_rank=i,
world_size=world_size,
- start_layer=meta.n_layers,
- end_layer=meta.n_layers,
- n_layers=meta.n_layers,
+ start_layer=card.n_layers,
+ end_layer=card.n_layers,
+ n_layers=card.n_layers,
)
for i in range(world_size)
},
← 176ab5ba Add GLM-4.7-Flash model cards (4bit, 5bit, 6bit, 8bit) (#121
·
back to Exo
·
don't clear mdns discovered connections 11655883 →