From 81a732593fa646c5b4d28aa576c80e056b223d52 Mon Sep 17 00:00:00 2001 From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> Date: Wed, 1 Jul 2026 22:34:52 -0700 Subject: [PATCH 1/2] mmap support for ios --- .../src/coreai_models/export/compression.py | 9 ++- python/src/coreai_models/export/pipeline.py | 65 ++++++++----------- python/src/coreai_models/models/base.py | 53 +++++++++++---- .../src/coreai_models/models/ios/mistral.py | 65 +++++++++++-------- python/src/coreai_models/models/ios/qwen2.py | 65 +++++++++++-------- python/src/coreai_models/models/ios/qwen3.py | 65 +++++++++++-------- 6 files changed, 190 insertions(+), 132 deletions(-) diff --git a/python/src/coreai_models/export/compression.py b/python/src/coreai_models/export/compression.py index 4f5cfd4..38dc818 100644 --- a/python/src/coreai_models/export/compression.py +++ b/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 diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 292c1d1..f9cbef2 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -11,7 +11,6 @@ """ import asyncio -import contextlib import logging import os import re @@ -159,38 +158,21 @@ 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, - ) + # Both variants load weights layer-by-layer (mmap-backed) and disk-checkpoint + # the compressor: macOS streams through quantization, iOS through palettization + # (both finalize with mmap_dir on the CoreAI backend). + 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: @@ -248,11 +230,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. @@ -300,7 +279,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": diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index ad04c99..30c9fa3 100644 --- a/python/src/coreai_models/models/base.py +++ b/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. @@ -465,6 +487,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") @@ -476,29 +504,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) @@ -578,6 +599,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( diff --git a/python/src/coreai_models/models/ios/mistral.py b/python/src/coreai_models/models/ios/mistral.py index 454d78f..75ecdbb 100644 --- a/python/src/coreai_models/models/ios/mistral.py +++ b/python/src/coreai_models/models/ios/mistral.py @@ -257,47 +257,60 @@ def forward( ) def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: - max_layer = -1 + 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])) - - if max_layer < 0: - err = "invalid state_dict" + present_layers.add(int(name_split[2])) + + # This method runs on the full dict, on one per-layer slice at a time, and + # (unlike the macOS models) on the shared-params slice. Iterate only the + # layers actually present, and access each layer's keys unconditionally so + # a malformed layer fails loudly rather than being silently skipped. A + # layer-less slice is valid only when 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 +326,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) diff --git a/python/src/coreai_models/models/ios/qwen2.py b/python/src/coreai_models/models/ios/qwen2.py index 84ff749..e39f6c6 100644 --- a/python/src/coreai_models/models/ios/qwen2.py +++ b/python/src/coreai_models/models/ios/qwen2.py @@ -257,47 +257,60 @@ def forward( ) def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: - max_layer = -1 + 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])) - - if max_layer < 0: - err = "invalid state_dict" + present_layers.add(int(name_split[2])) + + # This method runs on the full dict, on one per-layer slice at a time, and + # (unlike the macOS models) on the shared-params slice. Iterate only the + # layers actually present, and access each layer's keys unconditionally so + # a malformed layer fails loudly rather than being silently skipped. A + # layer-less slice is valid only when 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") # Qwen2Model is held inside Qwen2Extend — add "extend." prefix new_state_dict = {} @@ -313,7 +326,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) diff --git a/python/src/coreai_models/models/ios/qwen3.py b/python/src/coreai_models/models/ios/qwen3.py index 95788e8..4183828 100644 --- a/python/src/coreai_models/models/ios/qwen3.py +++ b/python/src/coreai_models/models/ios/qwen3.py @@ -263,47 +263,60 @@ def forward( ) def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: - max_layer = -1 + 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])) - - if max_layer < 0: - err = "invalid state_dict" + present_layers.add(int(name_split[2])) + + # This method runs on the full dict, on one per-layer slice at a time, and + # (unlike the macOS models) on the shared-params slice. Iterate only the + # layers actually present, and access each layer's keys unconditionally so + # a malformed layer fails loudly rather than being silently skipped. A + # layer-less slice is valid only when 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. will be 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") # Qwen3Model is held inside Qwen3Extend — add "extend." prefix new_state_dict = {} @@ -319,7 +332,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) From 1019a6dac6f03b314069f94d26dfd4ffa14811cb Mon Sep 17 00:00:00 2001 From: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> Date: Tue, 14 Jul 2026 11:36:27 -0700 Subject: [PATCH 2/2] comments and docstrings Signed-off-by: Prathamesh Mandke <46148373+pkmandke@users.noreply.github.com> --- python/src/coreai_models/export/pipeline.py | 4 +--- python/src/coreai_models/models/ios/mistral.py | 15 ++++++++++----- python/src/coreai_models/models/ios/qwen2.py | 15 ++++++++++----- python/src/coreai_models/models/ios/qwen3.py | 15 ++++++++++----- 4 files changed, 31 insertions(+), 18 deletions(-) diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 4713695..1a87b29 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -178,9 +178,7 @@ async def _async_export_model(config: ExportConfig) -> str: logger.info(f"Loading {config.hf_model_id} ({config.variant}, dtype={target_dtype})...") - # Both variants load weights layer-by-layer (mmap-backed) and disk-checkpoint - # the compressor: macOS streams through quantization, iOS through palettization - # (both finalize with mmap_dir on the CoreAI backend). + # 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) diff --git a/python/src/coreai_models/models/ios/mistral.py b/python/src/coreai_models/models/ios/mistral.py index 75ecdbb..4256e05 100644 --- a/python/src/coreai_models/models/ios/mistral.py +++ b/python/src/coreai_models/models/ios/mistral.py @@ -257,6 +257,13 @@ def forward( ) def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: + """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(".") @@ -266,11 +273,9 @@ def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: continue present_layers.add(int(name_split[2])) - # This method runs on the full dict, on one per-layer slice at a time, and - # (unlike the macOS models) on the shared-params slice. Iterate only the - # layers actually present, and access each layer's keys unconditionally so - # a malformed layer fails loudly rather than being silently skipped. A - # layer-less slice is valid only when it carries shared params. + # 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 ) diff --git a/python/src/coreai_models/models/ios/qwen2.py b/python/src/coreai_models/models/ios/qwen2.py index e39f6c6..6f7f0b8 100644 --- a/python/src/coreai_models/models/ios/qwen2.py +++ b/python/src/coreai_models/models/ios/qwen2.py @@ -257,6 +257,13 @@ def forward( ) def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: + """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(".") @@ -266,11 +273,9 @@ def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: continue present_layers.add(int(name_split[2])) - # This method runs on the full dict, on one per-layer slice at a time, and - # (unlike the macOS models) on the shared-params slice. Iterate only the - # layers actually present, and access each layer's keys unconditionally so - # a malformed layer fails loudly rather than being silently skipped. A - # layer-less slice is valid only when it carries shared params. + # 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 ) diff --git a/python/src/coreai_models/models/ios/qwen3.py b/python/src/coreai_models/models/ios/qwen3.py index 4183828..951343b 100644 --- a/python/src/coreai_models/models/ios/qwen3.py +++ b/python/src/coreai_models/models/ios/qwen3.py @@ -263,6 +263,13 @@ def forward( ) def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: + """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(".") @@ -272,11 +279,9 @@ def _mutate_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: continue present_layers.add(int(name_split[2])) - # This method runs on the full dict, on one per-layer slice at a time, and - # (unlike the macOS models) on the shared-params slice. Iterate only the - # layers actually present, and access each layer's keys unconditionally so - # a malformed layer fails loudly rather than being silently skipped. A - # layer-less slice is valid only when it carries shared params. + # 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 )