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 <[email protected]>

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
 

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