← back to Exo
cehckpoint
29ca6ad1ed2341f6e1ad711ff6d3165363fa69e2 · 2024-10-29 00:07:02 -0400 · thenatlog
Files touched
M exo/main.pyM exo/models.py
Diff
commit 29ca6ad1ed2341f6e1ad711ff6d3165363fa69e2
Author: thenatlog <logan.w.omalley@gmail.com>
Date: Tue Oct 29 00:07:02 2024 -0400
cehckpoint
---
exo/main.py | 337 ++++++++++++++++++++++++++++++++++++++--------------------
exo/models.py | 290 ++++++++++++++++++++++++++++++++++++++------------
2 files changed, 440 insertions(+), 187 deletions(-)
diff --git a/exo/main.py b/exo/main.py
index 90bcc73c..8cda941d 100644
--- a/exo/main.py
+++ b/exo/main.py
@@ -17,7 +17,15 @@ from exo.topology.ring_memory_weighted_partitioning_strategy import RingMemoryWe
from exo.api import ChatGPTAPI
from exo.download.shard_download import ShardDownloader, RepoProgressEvent, NoopShardDownloader
from exo.download.hf.hf_shard_download import HFShardDownloader
-from exo.helpers import print_yellow_exo, find_available_port, DEBUG, get_system_info, get_or_create_node_id, get_all_ip_addresses, terminal_link
+from exo.helpers import (
+ print_yellow_exo,
+ find_available_port,
+ DEBUG,
+ get_system_info,
+ get_or_create_node_id,
+ get_all_ip_addresses,
+ terminal_link,
+)
from exo.inference.shard import Shard
from exo.inference.inference_engine import get_inference_engine, InferenceEngine
from exo.inference.dummy_inference_engine import DummyInferenceEngine
@@ -26,9 +34,9 @@ from exo.orchestration.node import Node
from exo.models import model_base_shards
from exo.viz.topology_viz import TopologyViz
-# parse args
-parser = argparse.ArgumentParser(description="Initialize GRPC Discovery")
-parser.add_argument("command", nargs="?", choices=["run"], help="Command to run")
+# Parse arguments
+parser = argparse.ArgumentParser(description="Initialize GRPC Discovery and Run Inference")
+parser.add_argument("command", nargs="?", choices=["run"], default="run", help="Command to run (default: run)")
parser.add_argument("model_name", nargs="?", help="Model name to run")
parser.add_argument("--node-id", type=str, default=None, help="Node ID")
parser.add_argument("--node-host", type=str, default="0.0.0.0", help="Node host")
@@ -38,97 +46,177 @@ parser.add_argument("--download-quick-check", action="store_true", help="Quick c
parser.add_argument("--max-parallel-downloads", type=int, default=4, help="Max parallel downloads for model shards download")
parser.add_argument("--prometheus-client-port", type=int, default=None, help="Prometheus client port")
parser.add_argument("--broadcast-port", type=int, default=5678, help="Broadcast port for discovery")
-parser.add_argument("--discovery-module", type=str, choices=["udp", "tailscale", "manual"], default="udp", help="Discovery module to use")
+parser.add_argument(
+ "--discovery-module",
+ type=str,
+ choices=["udp", "tailscale", "manual"],
+ default="udp",
+ help="Discovery module to use",
+)
parser.add_argument("--discovery-timeout", type=int, default=30, help="Discovery timeout in seconds")
parser.add_argument("--discovery-config-path", type=str, default=None, help="Path to discovery config json file")
parser.add_argument("--wait-for-peers", type=int, default=0, help="Number of peers to wait to connect to before starting")
parser.add_argument("--chatgpt-api-port", type=int, default=8000, help="ChatGPT API port")
-parser.add_argument("--chatgpt-api-response-timeout", type=int, default=90, help="ChatGPT API response timeout in seconds")
+parser.add_argument(
+ "--chatgpt-api-response-timeout",
+ type=int,
+ default=90,
+ help="ChatGPT API response timeout in seconds",
+)
parser.add_argument("--max-generate-tokens", type=int, default=10000, help="Max tokens to generate in each request")
-parser.add_argument("--inference-engine", type=str, default=None, help="Inference engine to use (mlx, tinygrad, or dummy)")
+parser.add_argument(
+ "--inference-engine",
+ type=str,
+ choices=["mlx", "tinygrad", "dummy"],
+ default=None,
+ help="Inference engine to use (mlx, tinygrad, or dummy)",
+)
parser.add_argument("--disable-tui", action=argparse.BooleanOptionalAction, help="Disable TUI")
parser.add_argument("--run-model", type=str, help="Specify a model to run directly")
-parser.add_argument("--prompt", type=str, help="Prompt for the model when using --run-model", default="Who are you?")
+parser.add_argument(
+ "--prompt",
+ type=str,
+ help="Prompt for the model when using --run-model",
+ default="Who are you?",
+)
parser.add_argument("--tailscale-api-key", type=str, default=None, help="Tailscale API key")
parser.add_argument("--tailnet-name", type=str, default=None, help="Tailnet name")
args = parser.parse_args()
print(f"Selected inference engine: {args.inference_engine}")
-
print_yellow_exo()
-
system_info = get_system_info()
print(f"Detected system: {system_info}")
-shard_downloader: ShardDownloader = HFShardDownloader(quick_check=args.download_quick_check, max_parallel_downloads=args.max_parallel_downloads) if args.inference_engine != "dummy" else NoopShardDownloader()
+# Determine shard downloader based on inference engine
+shard_downloader: ShardDownloader = (
+ HFShardDownloader(quick_check=args.download_quick_check, max_parallel_downloads=args.max_parallel_downloads)
+ if args.inference_engine != "dummy"
+ else NoopShardDownloader()
+)
+
+# Determine inference engine name
inference_engine_name = args.inference_engine or ("mlx" if system_info == "Apple Silicon Mac" else "tinygrad")
print(f"Inference engine name after selection: {inference_engine_name}")
inference_engine = get_inference_engine(inference_engine_name, shard_downloader)
-print(f"Using inference engine: {inference_engine.__class__.__name__} with shard downloader: {shard_downloader.__class__.__name__}")
+print(
+ f"Using inference engine: {inference_engine.__class__.__name__} with shard downloader: {shard_downloader.__class__.__name__}"
+)
+
+# Select the first model as default if no model is specified
+available_models = list(model_base_shards.keys())
+
+if not available_models:
+ raise ValueError("No models available in model_base_shards.")
+
+default_model = available_models[0] # Retrieve the first model
+
+if DEBUG >= 1:
+ print(f"Available models: {available_models}")
+ print(f"Default model selected: {default_model}")
if args.node_port is None:
- args.node_port = find_available_port(args.node_host)
- if DEBUG >= 1: print(f"Using available port: {args.node_port}")
+ args.node_port = find_available_port(args.node_host)
+ if DEBUG >= 1:
+ print(f"Using available port: {args.node_port}")
args.node_id = args.node_id or get_or_create_node_id()
chatgpt_api_endpoints = [f"http://{ip}:{args.chatgpt_api_port}/v1/chat/completions" for ip in get_all_ip_addresses()]
web_chat_urls = [f"http://{ip}:{args.chatgpt_api_port}" for ip in get_all_ip_addresses()]
if DEBUG >= 0:
- print("Chat interface started:")
- for web_chat_url in web_chat_urls:
- print(f" - {terminal_link(web_chat_url)}")
- print("ChatGPT API endpoint served at:")
- for chatgpt_api_endpoint in chatgpt_api_endpoints:
- print(f" - {terminal_link(chatgpt_api_endpoint)}")
+ print("Chat interface started:")
+ for web_chat_url in web_chat_urls:
+ print(f" - {terminal_link(web_chat_url)}")
+ print("ChatGPT API endpoint served at:")
+ for chatgpt_api_endpoint in chatgpt_api_endpoints:
+ print(f" - {terminal_link(chatgpt_api_endpoint)}")
+# Initialize discovery based on the selected discovery module
if args.discovery_module == "udp":
- discovery = UDPDiscovery(args.node_id, args.node_port, args.listen_port, args.broadcast_port, lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities), discovery_timeout=args.discovery_timeout)
+ discovery = UDPDiscovery(
+ args.node_id,
+ args.node_port,
+ args.listen_port,
+ args.broadcast_port,
+ lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities),
+ discovery_timeout=args.discovery_timeout,
+ )
elif args.discovery_module == "tailscale":
- discovery = TailscaleDiscovery(args.node_id, args.node_port, lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities), discovery_timeout=args.discovery_timeout, tailscale_api_key=args.tailscale_api_key, tailnet=args.tailnet_name)
+ discovery = TailscaleDiscovery(
+ args.node_id,
+ args.node_port,
+ lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities),
+ discovery_timeout=args.discovery_timeout,
+ tailscale_api_key=args.tailscale_api_key,
+ tailnet=args.tailnet_name,
+ )
elif args.discovery_module == "manual":
- if not args.discovery_config_path:
- raise ValueError(f"--discovery-config-path is required when using manual discovery. Please provide a path to a config json file.")
- discovery = ManualDiscovery(args.discovery_config_path, args.node_id, create_peer_handle=lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities))
+ if not args.discovery_config_path:
+ raise ValueError(
+ "--discovery-config-path is required when using manual discovery. Please provide a path to a config json file."
+ )
+ discovery = ManualDiscovery(
+ args.discovery_config_path,
+ args.node_id,
+ create_peer_handle=lambda peer_id, address, device_capabilities: GRPCPeerHandle(peer_id, address, device_capabilities),
+ )
+
+# Initialize topology visualization if not disabled
topology_viz = TopologyViz(chatgpt_api_endpoints=chatgpt_api_endpoints, web_chat_urls=web_chat_urls) if not args.disable_tui else None
+
+# Initialize the node
node = StandardNode(
- args.node_id,
- None,
- inference_engine,
- discovery,
- partitioning_strategy=RingMemoryWeightedPartitioningStrategy(),
- max_generate_tokens=args.max_generate_tokens,
- topology_viz=topology_viz,
- shard_downloader=shard_downloader
+ args.node_id,
+ None,
+ inference_engine,
+ discovery,
+ partitioning_strategy=RingMemoryWeightedPartitioningStrategy(),
+ max_generate_tokens=args.max_generate_tokens,
+ topology_viz=topology_viz,
+ shard_downloader=shard_downloader,
)
+
+# Initialize the GRPC server
server = GRPCServer(node, args.node_host, args.node_port)
node.server = server
+
+# Initialize the ChatGPT API
api = ChatGPTAPI(
- node,
- inference_engine.__class__.__name__,
- response_timeout=args.chatgpt_api_response_timeout,
- on_chat_completion_request=lambda req_id, __, prompt: topology_viz.update_prompt(req_id, prompt) if topology_viz else None
+ node,
+ inference_engine.__class__.__name__,
+ response_timeout=args.chatgpt_api_response_timeout,
+ on_chat_completion_request=lambda req_id, __, prompt: topology_viz.update_prompt(req_id, prompt) if topology_viz else None,
)
+
+# Register token update handler for topology visualization
node.on_token.register("update_topology_viz").on_next(
- lambda req_id, tokens, __: topology_viz.update_prompt_output(req_id, inference_engine.tokenizer.decode(tokens)) if topology_viz and hasattr(inference_engine, "tokenizer") else None
+ lambda req_id, tokens, __: topology_viz.update_prompt_output(req_id, inference_engine.tokenizer.decode(tokens))
+ if topology_viz and hasattr(inference_engine, "tokenizer")
+ else None
)
+
def preemptively_start_download(request_id: str, opaque_status: str):
- try:
- status = json.loads(opaque_status)
- if status.get("type") == "node_status" and status.get("status") == "start_process_prompt":
- current_shard = node.get_current_shard(Shard.from_dict(status.get("shard")))
- if DEBUG >= 2: print(f"Preemptively starting download for {current_shard}")
- asyncio.create_task(shard_downloader.ensure_shard(current_shard))
- except Exception as e:
- if DEBUG >= 2:
- print(f"Failed to preemptively start download: {e}")
- traceback.print_exc()
+ try:
+ status = json.loads(opaque_status)
+ if status.get("type") == "node_status" and status.get("status") == "start_process_prompt":
+ current_shard = node.get_current_shard(Shard.from_dict(status.get("shard")))
+ if DEBUG >= 2:
+ print(f"Preemptively starting download for {current_shard}")
+ asyncio.create_task(shard_downloader.ensure_shard(current_shard))
+ except Exception as e:
+ if DEBUG >= 2:
+ print(f"Failed to preemptively start download: {e}")
+ traceback.print_exc()
+
node.on_opaque_status.register("start_download").on_next(preemptively_start_download)
+# Start Prometheus metrics server if specified
if args.prometheus_client_port:
- from exo.stats.metrics import start_metrics_server
- start_metrics_server(node, args.prometheus_client_port)
+ from exo.stats.metrics import start_metrics_server
+
+ start_metrics_server(node, args.prometheus_client_port)
last_broadcast_time = 0
@@ -137,88 +225,101 @@ def throttled_broadcast(shard: Shard, event: RepoProgressEvent):
current_time = time.time()
if event.status == "complete" or current_time - last_broadcast_time >= 0.1:
last_broadcast_time = current_time
- asyncio.create_task(node.broadcast_opaque_status("", json.dumps({
- "type": "download_progress",
- "node_id": node.id,
- "progress": event.to_dict()
- })))
+ asyncio.create_task(
+ node.broadcast_opaque_status(
+ "",
+ json.dumps(
+ {
+ "type": "download_progress",
+ "node_id": node.id,
+ "progress": event.to_dict(),
+ }
+ ),
+ )
+ )
shard_downloader.on_progress.register("broadcast").on_next(throttled_broadcast)
-
-async def shutdown(signal, loop):
- """Gracefully shutdown the server and close the asyncio loop."""
- print(f"Received exit signal {signal.name}...")
- print("Thank you for using exo.")
- print_yellow_exo()
- server_tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
- [task.cancel() for task in server_tasks]
- print(f"Cancelling {len(server_tasks)} outstanding tasks")
- await asyncio.gather(*server_tasks, return_exceptions=True)
- await server.stop()
- loop.stop()
+async def shutdown(signal_name, loop):
+ """Gracefully shutdown the server and close the asyncio loop."""
+ print(f"Received exit signal {signal_name}...")
+ print("Thank you for using exo.")
+ print_yellow_exo()
+ server_tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
+ [task.cancel() for task in server_tasks]
+ print(f"Cancelling {len(server_tasks)} outstanding tasks")
+ await asyncio.gather(*server_tasks, return_exceptions=True)
+ await server.stop()
+ loop.stop()
async def run_model_cli(node: Node, inference_engine: InferenceEngine, model_name: str, prompt: str):
- shard = model_base_shards.get(model_name, {}).get(inference_engine.__class__.__name__)
- if not shard:
- print(f"Error: Unsupported model '{model_name}' for inference engine {inference_engine.__class__.__name__}")
- return
- tokenizer = await resolve_tokenizer(shard.model_id)
- request_id = str(uuid.uuid4())
- callback_id = f"cli-wait-response-{request_id}"
- callback = node.on_token.register(callback_id)
- if topology_viz:
- topology_viz.update_prompt(request_id, prompt)
- prompt = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True)
-
- try:
- print(f"Processing prompt: {prompt}")
- await node.process_prompt(shard, prompt, None, request_id=request_id)
-
- _, tokens, _ = await callback.wait(lambda _request_id, tokens, is_finished: _request_id == request_id and is_finished, timeout=300)
-
- print("\nGenerated response:")
- print(tokenizer.decode(tokens))
- except Exception as e:
- print(f"Error processing prompt: {str(e)}")
- traceback.print_exc()
- finally:
- node.on_token.deregister(callback_id)
+ shard = model_base_shards.get(model_name, {}).get(inference_engine.__class__.__name__)
+ if not shard:
+ print(f"Error: Unsupported model '{model_name}' for inference engine {inference_engine.__class__.__name__}")
+ return
+ tokenizer = await resolve_tokenizer(shard.model_id)
+ request_id = str(uuid.uuid4())
+ callback_id = f"cli-wait-response-{request_id}"
+ callback = node.on_token.register(callback_id)
+ if topology_viz:
+ topology_viz.update_prompt(request_id, prompt)
+ prompt_formatted = tokenizer.apply_chat_template(
+ [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True
+ )
+ try:
+ print(f"Processing prompt: {prompt_formatted}")
+ await node.process_prompt(shard, prompt_formatted, None, request_id=request_id)
+
+ _, tokens, _ = await callback.wait(
+ lambda _request_id, tokens, is_finished: _request_id == request_id and is_finished,
+ timeout=300,
+ )
+
+ print("\nGenerated response:")
+ print(tokenizer.decode(tokens))
+ except Exception as e:
+ print(f"Error processing prompt: {str(e)}")
+ traceback.print_exc()
+ finally:
+ node.on_token.deregister(callback_id)
async def main():
- loop = asyncio.get_running_loop()
+ loop = asyncio.get_running_loop()
+
+ # Handle exit signals
+ def handle_exit():
+ asyncio.ensure_future(shutdown(signal.SIGTERM, loop))
- # Use a more direct approach to handle signals
- def handle_exit():
- asyncio.ensure_future(shutdown(signal.SIGTERM, loop))
+ for s in [signal.SIGINT, signal.SIGTERM]:
+ loop.add_signal_handler(s, handle_exit)
- for s in [signal.SIGINT, signal.SIGTERM]:
- loop.add_signal_handler(s, handle_exit)
+ await node.start(wait_for_peers=args.wait_for_peers)
- await node.start(wait_for_peers=args.wait_for_peers)
+ if args.command == "run" or args.run_model:
+ # Use the provided model name or default to the first model
+ model_name = args.model_name or args.run_model or default_model
- if args.command == "run" or args.run_model:
- model_name = args.model_name or args.run_model
- if not model_name:
- print("Error: Model name is required when using 'run' command or --run-model")
- return
- await run_model_cli(node, inference_engine, model_name, args.prompt)
- else:
- asyncio.create_task(api.run(port=args.chatgpt_api_port)) # Start the API server as a non-blocking task
- await asyncio.Event().wait()
+ # Inform the user about the default selection if no model was specified
+ if not args.model_name and not args.run_model:
+ print(f"No model specified. Defaulting to the first available model: '{default_model}'")
+ await run_model_cli(node, inference_engine, model_name, args.prompt)
+ else:
+ # Start the ChatGPT API server as a non-blocking task
+ asyncio.create_task(api.run(port=args.chatgpt_api_port))
+ await asyncio.Event().wait()
def run():
- loop = asyncio.new_event_loop()
- asyncio.set_event_loop(loop)
- try:
- loop.run_until_complete(main())
- except KeyboardInterrupt:
- print("Received keyboard interrupt. Shutting down...")
- finally:
- loop.run_until_complete(shutdown(signal.SIGTERM, loop))
- loop.close()
+ loop = asyncio.new_event_loop()
+ asyncio.set_event_loop(loop)
+ try:
+ loop.run_until_complete(main())
+ except KeyboardInterrupt:
+ print("Received keyboard interrupt. Shutting down...")
+ finally:
+ loop.run_until_complete(shutdown(signal.SIGTERM, loop))
+ loop.close()
if __name__ == "__main__":
- run()
+ run()
diff --git a/exo/models.py b/exo/models.py
index a462f70d..6f3f3e58 100644
--- a/exo/models.py
+++ b/exo/models.py
@@ -1,73 +1,225 @@
from exo.inference.shard import Shard
model_base_shards = {
- ### llama
- "llama-3.2-1b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Llama-3.2-1B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=16),
- },
- "llama-3.2-3b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Llama-3.2-3B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
- },
- "llama-3.1-8b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=32),
- "TinygradDynamicShardInferenceEngine": Shard(model_id="mlabonne/Meta-Llama-3.1-8B-Instruct-abliterated", start_layer=0, end_layer=0, n_layers=32),
- },
- "llama-3.1-70b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=80),
- "TinygradDynamicShardInferenceEngine": Shard(model_id="NousResearch/Meta-Llama-3.1-70B-Instruct", start_layer=0, end_layer=0, n_layers=80),
- },
- "llama-3.1-70b-bf16": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-bf16-CORRECTED", start_layer=0, end_layer=0, n_layers=80),
- "TinygradDynamicShardInferenceEngine": Shard(model_id="NousResearch/Meta-Llama-3.1-70B-Instruct", start_layer=0, end_layer=0, n_layers=80),
- },
- "llama-3.1-405b": {"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3.1-405B-4bit", start_layer=0, end_layer=0, n_layers=126),},
- "llama-3-8b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3-8B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=32),
- "TinygradDynamicShardInferenceEngine": Shard(model_id="TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R", start_layer=0, end_layer=0, n_layers=32),
- },
- "llama-3-70b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Meta-Llama-3-70B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=80),
- "TinygradDynamicShardInferenceEngine": Shard(model_id="TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-70B-R", start_layer=0, end_layer=0, n_layers=80),
- },
- ### mistral
- "mistral-nemo": {"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Mistral-Nemo-Instruct-2407-4bit", start_layer=0, end_layer=0, n_layers=40),},
- "mistral-large": {"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Mistral-Large-Instruct-2407-4bit", start_layer=0, end_layer=0, n_layers=88),},
- ### deepseek
- "deepseek-coder-v2-lite": {"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx", start_layer=0, end_layer=0, n_layers=27),},
- "deepseek-coder-v2.5": {"MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/DeepSeek-V2.5-MLX-AQ4_1_64", start_layer=0, end_layer=0, n_layers=60),},
- ### llava
- "llava-1.5-7b-hf": {"MLXDynamicShardInferenceEngine": Shard(model_id="llava-hf/llava-1.5-7b-hf", start_layer=0, end_layer=0, n_layers=32),},
- ### qwen
- "qwen-2.5-coder-1.5b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-Coder-1.5B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
- },
- "qwen-2.5-coder-7b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-Coder-7B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
- },
- "qwen-2.5-7b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-7B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
- },
- "qwen-2.5-math-7b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-Math-7B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=28),
- },
- "qwen-2.5-14b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-14B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=48),
- },
- "qwen-2.5-72b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-72B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=80),
- },
- "qwen-2.5-math-72b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Qwen2.5-Math-72B-Instruct-4bit", start_layer=0, end_layer=0, n_layers=80),
- },
- ### nemotron
- "nemotron-70b": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/nvidia_Llama-3.1-Nemotron-70B-Instruct-HF_4bit", start_layer=0, end_layer=0, n_layers=80),
- },
- "nemotron-70b-bf16": {
- "MLXDynamicShardInferenceEngine": Shard(model_id="mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-bf16", start_layer=0, end_layer=0, n_layers=80),
- },
- # dummy
- "dummy": {
- "DummyInferenceEngine": Shard(model_id="dummy", start_layer=0, end_layer=7, n_layers=8),
- },
+ ### LLaMA Models
+ "llama-3.2-1b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Llama-3.2-1B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=16,
+ ),
+ },
+ "llama-3.2-3b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Llama-3.2-3B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=28,
+ ),
+ },
+ "llama-3.1-8b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Meta-Llama-3.1-8B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=32,
+ ),
+ "TinygradDynamicShardInferenceEngine": Shard(
+ model_id="mlabonne/Meta-Llama-3.1-8B-Instruct-abliterated",
+ start_layer=0,
+ end_layer=0,
+ n_layers=32,
+ ),
+ },
+ "llama-3.1-70b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ "TinygradDynamicShardInferenceEngine": Shard(
+ model_id="NousResearch/Meta-Llama-3.1-70B-Instruct",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ },
+ "llama-3.1-70b-bf16": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Meta-Llama-3.1-70B-Instruct-bf16-CORRECTED",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ "TinygradDynamicShardInferenceEngine": Shard(
+ model_id="NousResearch/Meta-Llama-3.1-70B-Instruct",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ },
+ "llama-3.1-405b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Meta-Llama-3.1-405B-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=126,
+ ),
+ },
+ "llama-3-8b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Meta-Llama-3-8B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=32,
+ ),
+ "TinygradDynamicShardInferenceEngine": Shard(
+ model_id="TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R",
+ start_layer=0,
+ end_layer=0,
+ n_layers=32,
+ ),
+ },
+ "llama-3-70b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Meta-Llama-3-70B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ "TinygradDynamicShardInferenceEngine": Shard(
+ model_id="TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-70B-R",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ },
+ ### Mistral Models
+ "mistral-nemo": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Mistral-Nemo-Instruct-2407-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=40,
+ ),
+ },
+ "mistral-large": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Mistral-Large-Instruct-2407-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=88,
+ ),
+ },
+ ### DeepSeek Models
+ "deepseek-coder-v2-lite": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx",
+ start_layer=0,
+ end_layer=0,
+ n_layers=27,
+ ),
+ },
+ "deepseek-coder-v2.5": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/DeepSeek-V2.5-MLX-AQ4_1_64",
+ start_layer=0,
+ end_layer=0,
+ n_layers=60,
+ ),
+ },
+ ### Llava Models
+ "llava-1.5-7b-hf": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="llava-hf/llava-1.5-7b-hf",
+ start_layer=0,
+ end_layer=0,
+ n_layers=32,
+ ),
+ },
+ ### Qwen Models
+ "qwen-2.5-coder-1.5b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Qwen2.5-Coder-1.5B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=28,
+ ),
+ },
+ "qwen-2.5-coder-7b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Qwen2.5-Coder-7B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=28,
+ ),
+ },
+ "qwen-2.5-7b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Qwen2.5-7B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=28,
+ ),
+ },
+ "qwen-2.5-math-7b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Qwen2.5-Math-7B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=28,
+ ),
+ },
+ "qwen-2.5-14b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Qwen2.5-14B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=48,
+ ),
+ },
+ "qwen-2.5-72b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Qwen2.5-72B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ },
+ "qwen-2.5-math-72b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Qwen2.5-Math-72B-Instruct-4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ },
+ ### Nemotron Models
+ "nemotron-70b": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/nvidia_Llama-3.1-Nemotron-70B-Instruct-HF_4bit",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ },
+ "nemotron-70b-bf16": {
+ "MLXDynamicShardInferenceEngine": Shard(
+ model_id="mlx-community/Llama-3.1-Nemotron-70B-Instruct-HF-bf16",
+ start_layer=0,
+ end_layer=0,
+ n_layers=80,
+ ),
+ },
+ ### Dummy Model
+ "dummy": {
+ "DummyInferenceEngine": Shard(
+ model_id="dummy",
+ start_layer=0,
+ end_layer=7,
+ n_layers=8,
+ ),
+ },
}
← 9dc93fd5 add traceback.print_exc on topology collection errors from p
·
back to Exo
·
remove buggy select logic 7dfe856f →