Add support for memory efficient iOS exports for large models

python/src/coreai_models/export/compression.py

@@ -207,6 +207,7 @@ def palettize_pytorch_model(
     model: nn.Module,
     example_inputs: tuple,
     palettization_config: "dict | KMeansPalettizerConfig",
+    mmap_dir: str | None = None,
 ) -> nn.Module:
     """
     Palettize a PyTorch model using post-training palettization with coreai-opt.
@@ -218,6 +219,10 @@ def palettize_pytorch_model(
         palettization_config: Either a configuration dictionary (matching the
             inner shape coreai-opt expects under `kmeans_palettization_config`)
             or a prebuilt KMeansPalettizerConfig instance.
+        mmap_dir: Optional directory for memory-efficient finalization.
+            When provided, each finalized layer is saved to a per-layer
+            safetensors file and reloaded mmap-backed.
+            When None, finalization keeps weights in RAM.
 
     Returns:
         Palettized model ready for export.
@@ -240,7 +245,9 @@ def palettize_pytorch_model(
     palettizer = KMeansPalettizer(model, config)
     prepared_model = palettizer.prepare(example_inputs=example_inputs, num_workers=32)
 
-    finalized_model = palettizer.finalize(prepared_model, backend=ExportBackend.CoreAI)
+    finalized_model = palettizer.finalize(
+        prepared_model, backend=ExportBackend.CoreAI, mmap_dir=mmap_dir
+    )
 
     logger.info("Palettization with coreai-opt complete")
     return finalized_model

python/src/coreai_models/export/pipeline.py

@@ -11,7 +11,6 @@
 """
 
 import asyncio
-import contextlib
 import logging
 import os
 import re
@@ -181,39 +180,19 @@ async def _async_export_model(config: ExportConfig) -> str:
 
     logger.info(f"Loading {config.hf_model_id} ({config.variant}, dtype={target_dtype})...")
 
-    # Memory-efficient layer-by-layer loading + quantizer disk-checkpointing
-    # is macOS-only for now. The iOS variant keeps the legacy full-RAM path
-    # since its palettization flow has not been validated against streaming
-    # weight loading.
-    use_memory_efficient = config.variant == "macOS"
-    temp_dir_ctx: contextlib.AbstractContextManager[str | None] = (
-        tempfile.TemporaryDirectory(prefix="coreai_export_")
-        if use_memory_efficient
-        else contextlib.nullcontext(None)
-    )
-
-    with temp_dir_ctx as temp_dir:
-        if use_memory_efficient:
-            assert temp_dir is not None  # nullcontext yields None only when not memory-efficient
-            layer_mmap_dir = os.path.join(temp_dir, "layers")
-            os.makedirs(layer_mmap_dir, exist_ok=True)
-            model = model_class.from_hf_memory_efficient(
-                config.hf_model_id,
-                max_context_length=max_context_length,
-                target_dtype=target_dtype,
-                mmap_path=layer_mmap_dir,
-                num_layers=config.num_layers,
-                hf_config_attr=entry.hf_config_attr,
-                hf_state_dict_prefix=entry.hf_state_dict_prefix,
-            )
-        else:
-            model = model_class.from_hf(
-                config.hf_model_id,
-                max_context_length=max_context_length,
-                target_dtype=target_dtype,
-                num_layers=config.num_layers,
-                disable_embedding_quantization=config.disable_embedding_quantization,
-            )
+    # Both variants load weights layer-by-layer (mmap-backed).
+    with tempfile.TemporaryDirectory(prefix="coreai_export_") as temp_dir:
+        layer_mmap_dir = os.path.join(temp_dir, "layers")
+        os.makedirs(layer_mmap_dir, exist_ok=True)
+        model = model_class.from_hf_memory_efficient(
+            config.hf_model_id,
+            max_context_length=max_context_length,
+            target_dtype=target_dtype,
+            mmap_path=layer_mmap_dir,
+            num_layers=config.num_layers,
+            hf_config_attr=entry.hf_config_attr,
+            hf_state_dict_prefix=entry.hf_state_dict_prefix,
+        )
         model = model.eval()
         # ---- 3. Resolve compression preset ----
         if config.compression_config_object is not None:
@@ -271,11 +250,8 @@ def get_calibration_data():  # type: ignore[no-untyped-def]
                 tokenizer = AutoTokenizer.from_pretrained(config.hf_model_id)
                 return get_c4(tokenizer)
 
-            quantizer_mmap_dir: str | None = None
-            if use_memory_efficient:
-                assert temp_dir is not None
-                quantizer_mmap_dir = os.path.join(temp_dir, "quantized")
-                os.makedirs(quantizer_mmap_dir, exist_ok=True)
+            quantizer_mmap_dir = os.path.join(temp_dir, "quantized")
+            os.makedirs(quantizer_mmap_dir, exist_ok=True)
 
             # Pass-through prebuilt QuantizerConfig objects.
             # copy dicts so we don't mutate the shared preset.
@@ -323,7 +299,15 @@ def get_calibration_data():  # type: ignore[no-untyped-def]
                 key_cache,
                 value_cache,
             )
-            model = palettize_pytorch_model(model, palettization_inputs, torch_palettization_config)
+            palettization_mmap_dir = os.path.join(temp_dir, "palettized")
+            os.makedirs(palettization_mmap_dir, exist_ok=True)
+
+            model = palettize_pytorch_model(
+                model,
+                palettization_inputs,
+                torch_palettization_config,
+                mmap_dir=palettization_mmap_dir,
+            )
 
         # ---- 4. Variant-specific export ----
         if config.variant == "macOS":

python/src/coreai_models/models/base.py

@@ -117,7 +117,23 @@ def _save_and_mmap_safetensors(
 
     param_names = {name for name, _ in module.named_parameters()}
 
-    save_file({k: v.contiguous() for k, v in tensors.items()}, path)
+    # safetensors `save_file` rejects entries that share storage (e.g. the iOS
+    # mutate ties `extend.emb_scale` to `gather_embeddings.scale`). Clone any
+    # entry that aliases an earlier one so each key is written independently.
+    # The mmap round-trip below returns independent tensors per key regardless,
+    # so this does not change the reloaded model.
+    seen_storage: dict[int, str] = {}
+    to_save: dict[str, torch.Tensor] = {}
+    for key, value in tensors.items():
+        value = value.contiguous()
+        storage_ptr = value.untyped_storage().data_ptr()
+        if storage_ptr in seen_storage:
+            value = value.clone()
+        else:
+            seen_storage[storage_ptr] = key
+        to_save[key] = value
+
+    save_file(to_save, path)
 
     new_sd: dict[str, torch.Tensor] = {}
     with safe_open(path, framework="pt", device="cpu") as f:
@@ -230,6 +246,12 @@ class BaseForCausalLM(torch.nn.Module):
     # Subclasses must override this with their specific HuggingFace model class
     _HF_MODEL_CLASS: type | None = None
 
+    # When True, the memory-efficient streaming loader runs `_mutate_state_dict`
+    # on the shared (non-layer) params too, so subclasses that transform
+    # embeddings / lm_head / norms (e.g. the iOS models) produce their expected
+    # keys. macOS models load shared params as-is and leave this False.
+    _MUTATE_SHARED_PARAMS_ON_STREAM: bool = False
+
     @staticmethod
     def cast_logits_bfloat16_to_float16(forward_fn: Callable) -> Callable:
         """Decorator to cast torch.bfloat16 logits outputs to float16.
@@ -481,6 +503,12 @@ def from_hf_memory_efficient(
         shared_dict = {k.removeprefix(hf_state_dict_prefix): v for k, v in shared_dict.items()}
         del shared_index
 
+        # Models that transform shared params (iOS embedding quantization,
+        # lm_head / norm renaming) mutate the shared slice here. macOS models
+        # load shared params unchanged.
+        if model._MUTATE_SHARED_PARAMS_ON_STREAM:
+            model._mutate_state_dict(shared_dict)
+
         if mmap_path is not None:
             os.makedirs(mmap_path, exist_ok=True)
             shared_path = os.path.join(mmap_path, "shared.safetensors")
@@ -492,29 +520,22 @@ def from_hf_memory_efficient(
         gc.collect()
 
         # One transformer layer at a time.
-        exclude_buffers = {KVCache.HF_K_BUFFER_NAME, KVCache.HF_V_BUFFER_NAME}
         for layer_idx in sorted(per_layer_index.keys()):
             layer_key_to_file = per_layer_index.pop(layer_idx)
             layer_sd = _load_tensors_for_keys(layer_key_to_file, target_dtype)
             layer_sd = {k.removeprefix(hf_state_dict_prefix): v for k, v in layer_sd.items()}
             del layer_key_to_file
 
-            # Per-model fusion (qkv, qk_norm, MoE expert stacking, ...).
-            # Subclass `_mutate_state_dict` is layer-keyed and safe on a
-            # single-layer slice.
+            # Per-model fusion (qkv, qk_norm, MoE expert stacking, Conv2d
+            # reshapes, ...). Subclass `_mutate_state_dict` is layer-keyed and
+            # safe on a single-layer slice; it rewrites keys into the model's
+            # own namespace (e.g. iOS wraps layers under `extend.`), so the
+            # mutated slice is assigned against the whole model.
             model._mutate_state_dict(layer_sd)
 
             if mmap_path is not None:
-                layer_prefix = f"model.layers.{layer_idx}."
-                relative_sd = {
-                    k.removeprefix(layer_prefix): v
-                    for k, v in layer_sd.items()
-                    if k.removeprefix(layer_prefix) not in exclude_buffers
-                }
-                layer_module = model.model.layers[layer_idx]
                 layer_path = os.path.join(mmap_path, f"layer_{layer_idx}.safetensors")
-                _save_and_mmap_safetensors(layer_module, relative_sd, layer_path)
-                del relative_sd
+                _save_and_mmap_safetensors(model, layer_sd, layer_path)
             else:
                 model.load_state_dict(layer_sd, assign=True, strict=False)
 
@@ -594,6 +615,10 @@ def to(self: Self, dtype) -> Self:
 
 
 class BaseForCausalLMForiOS(BaseForCausalLM):
+    # iOS models rewrite the shared embedding / lm_head / norm keys inside
+    # `_mutate_state_dict`, so the streaming loader must mutate the shared slice.
+    _MUTATE_SHARED_PARAMS_ON_STREAM: bool = True
+
     def __init__(self: Self, config, model_device: str, disable_embedding_quantization=False):
         super().__init__(config, model_device)
         self.load_embeddings = LoadEmbeddings(

python/src/coreai_models/models/ios/mistral.py

@@ -257,47 +257,65 @@ def forward(
         )
 
     def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
-        max_layer = -1
+        """Rewrite HF weight keys into the iOS module layout, in place.
+
+        Called repeatedly on partial state_dicts: the full dict, one
+        single-layer slice at a time, and a shared-params-only slice
+        (embeddings / lm_head). Handles any subset of keys and does not assume
+        all layers are present.
+        """
+        present_layers = set()
         for k in state_dict:
             name_split = k.split(".")
             if len(name_split) != 6:
                 continue
             if not k.startswith("model.layers."):
                 continue
-            max_layer = max(max_layer, int(name_split[2]))
+            present_layers.add(int(name_split[2]))
 
-        if max_layer < 0:
-            err = "invalid state_dict"
+        # Access present layers' keys unconditionally so a malformed layer fails
+        # loudly instead of being silently skipped; a layer-less slice is valid
+        # only if it carries shared params.
+        has_shared_keys = (
+            "model.embed_tokens.weight" in state_dict or "lm_head.weight" in state_dict
+        )
+        if not present_layers and not has_shared_keys:
+            err = (
+                "state_dict has no recognizable keys: expected per-layer weights "
+                "('model.layers.N.*') and/or shared params "
+                "('model.embed_tokens.weight', 'lm_head.weight')."
+            )
             raise ValueError(err)
 
-        for i in range(max_layer + 1):
+        for layer_idx in sorted(present_layers):
             # Reshape attention weights for Conv2d
             for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]:
-                weight_key = f"model.layers.{i}.self_attn.{proj}.weight"
+                weight_key = f"model.layers.{layer_idx}.self_attn.{proj}.weight"
                 state_dict[weight_key] = state_dict[weight_key].unsqueeze(-1).unsqueeze(-1)
 
             # Reshape MLP weights for Conv2d
             for proj in ["up_proj", "gate_proj", "down_proj"]:
-                weight_key = f"model.layers.{i}.mlp.{proj}.weight"
+                weight_key = f"model.layers.{layer_idx}.mlp.{proj}.weight"
                 state_dict[weight_key] = state_dict[weight_key].unsqueeze(-1).unsqueeze(-1)
 
-        # Handle embeddings
-        embedding_table = state_dict["model.embed_tokens.weight"].unsqueeze(1)
-        if not self.disable_embedding_quantization:
-            embedding_table, scale, zero_point = quantize_per_tensor(
-                embedding_table, nbits=8, symmetric=True
-            )
-        else:
-            scale = torch.tensor(1.0, dtype=embedding_table.dtype)
-            zero_point = torch.tensor(0, dtype=torch.int8)
-
-        state_dict["load_embeddings.embedding_table"] = embedding_table
-        state_dict["gather_embeddings.scale"] = scale
-        state_dict["gather_embeddings.zero_point"] = zero_point
-        state_dict["extend.emb_scale"] = scale
-        state_dict["extend.emb_zero_point"] = zero_point
-
-        state_dict.pop("model.embed_tokens.weight")
+        # Handle embeddings (shared param; absent from per-layer slices)
+        if "model.embed_tokens.weight" in state_dict:
+            embedding_table = state_dict["model.embed_tokens.weight"].unsqueeze(1)
+            if not self.disable_embedding_quantization:
+                embedding_table, scale, zero_point = quantize_per_tensor(
+                    embedding_table, nbits=8, symmetric=True
+                )
+            else:
+                scale = torch.tensor(1.0, dtype=embedding_table.dtype)
+                zero_point = torch.tensor(0, dtype=torch.int8)
+
+            state_dict["load_embeddings.embedding_table"] = embedding_table
+            state_dict["gather_embeddings.scale"] = scale
+            state_dict["gather_embeddings.zero_point"] = zero_point
+            state_dict["extend.emb_scale"] = scale
+            state_dict["extend.emb_zero_point"] = zero_point
+
+            state_dict.pop("model.embed_tokens.weight")
 
         # MistralModel is held inside MistralExtend — add "extend." prefix
         new_state_dict = {}
@@ -313,7 +331,7 @@ def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
             state_dict.pop(k)
         state_dict.update(new_state_dict)
 
-        if not self.config.tie_word_embeddings:
+        if not self.config.tie_word_embeddings and "lm_head.weight" in state_dict:
             state_dict["extend.lm_head.weight"] = state_dict["lm_head.weight"]
 
         state_dict.pop("lm_head.weight", None)