From bfbbea4e48c88d143e1769f6a921ba1ffc43a5db Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Thu, 12 Feb 2026 11:24:31 -0800 Subject: [PATCH] fix: use lazy=False and exclude safetensors index for weight broadcast receivers When receiving weights via MLX distributed broadcast, the receiver node has no .safetensors files on disk. Two issues caused rms_norm shape mismatch during warmup: 1. model.safetensors.index.json was transferred as metadata (has .json ext), causing load_model to create lazy tensor refs to nonexistent files 2. lazy=True created dangling references even without the index file Fix: exclude *.safetensors.index.json from metadata transfer, and use lazy=False when receiver has no local weight files. Co-Authored-By: Claude Opus 4.6 --- src/exo/worker/engines/mlx/model_transfer.py | 8 +++++++- src/exo/worker/engines/mlx/utils_mlx.py | 7 +++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/src/exo/worker/engines/mlx/model_transfer.py b/src/exo/worker/engines/mlx/model_transfer.py index 7e4f1116..553a8a61 100644 --- a/src/exo/worker/engines/mlx/model_transfer.py +++ b/src/exo/worker/engines/mlx/model_transfer.py @@ -52,7 +52,13 @@ def _all_sum_cpu(x: mx.array, group: Group) -> mx.array: def _is_metadata_file(filename: str) -> bool: - """Check if a file is a metadata file (not a weight file).""" + """Check if a file is a metadata file (not a weight file). + + Excludes safetensors index files (e.g. model.safetensors.index.json) since + they reference .safetensors shard files that won't exist on the receiver. + """ + if filename.endswith(".safetensors.index.json"): + return False _, ext = os.path.splitext(filename) return ext.lower() in _METADATA_EXTENSIONS diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 83c75a58..133e448f 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -230,8 +230,11 @@ def shard_and_load( model_path, group, has_local_weights ) - # Create model architecture (all ranks have config.json on disk now) - model, _ = load_model(model_path, lazy=True, strict=False) + # Create model architecture (all ranks have config.json on disk now). + # Use lazy=False when receiver has no local weights: lazy=True would create + # dangling references to nonexistent safetensors files. + use_lazy = has_local_weights or not needs_transfer + model, _ = load_model(model_path, lazy=use_lazy, strict=False) logger.debug(model) if hasattr(model, "model") and isinstance(model.model, DeepseekV3Model): # type: ignore pass