Model Loading¶

Table of Contents¶

• What Happens When You Load a Model
• Loading Pipeline Flowchart
• Step 1: Detection
• The config.json Contract
• ModelType Enum
• Special Detection Cases
• Step 2: Policy Routing
• Directory Route Families
• Weight Route Families
• Step 3: The Four Loaders
• ConfigBacked Loader
• Nonstandard Loader
• Special Loader
• VLM Loader
• The LoadedModel Enum
• Weight Loading Mechanics
• SafeTensors Format
• Tied Embedding Restoration
• Config Sanitization
• Quantization Support
• Load-Time Surgery Hook
• LoRA Adapter Loading

What Happens When You Load a Model¶

When you pass a model path to mlxcel, startup first identifies whether the checkpoint is a generator, embedder, reranker, audio model, detector, or another specialized task. A generation checkpoint then follows three sequential steps:

1. Detection: Read config.json, identify the architecture, produce a ModelType variant.
2. Policy routing: Consult the model_metadata registry to determine which loader family to use.
3. Construction: Execute the selected loader, which reads weights from disk and constructs the model's layer graph in MLX.

The generation result is a LoadedModel enum value and a MlxcelTokenizer. Embedding and reranking checkpoints instead produce their task-specific trait objects and tokenizer, owned by dedicated workers in server mode. This keeps non-generative scoring semantics out of the token-generation enum.

Loading Pipeline Flowchart¶

flowchart TD
    START([Model path]) --> RESOLVE[Resolve directory\nif given a file path]
    RESOLVE --> READ_CONFIG[Read config.json]
    READ_CONFIG --> SANITIZE[Sanitize config JSON\nfix known formatting issues]
    SANITIZE --> KIND{Model kind?}
    KIND -- Embedding --> EMB_LOAD[load_embedding_model\npooling + normalization contract]
    KIND -- Reranker --> RR_LOAD[load_reranker\npair or generative scoring]
    KIND -- Generation --> DETECT[get_model_type\nmap model_type string\nto ModelType enum]
    DETECT --> POLICY[model_load_policy\nconsult model_metadata registry]

    POLICY --> ROUTE{Directory route?}

    ROUTE -- ConfigBacked --> CB_LOAD[try_load_config_backed_model_from_dir\nconfig.json -> ModelArgs -> model.load]
    ROUTE -- Nonstandard --> NS_LOAD[try_load_nonstandard_model_from_dir\nfamily-specific loader]
    ROUTE -- VLM --> VLM_LOAD[try_load_vlm_model_from_dir\nvision encoder + text model]
    ROUTE -- Mistral3Dynamic --> M3_CHECK{inner text_config\nmodel_type?}
    M3_CHECK -- ministral3 --> M3_TEXT[Ministral3Wrapper]
    M3_CHECK -- other --> M3_LLAMA[LlamaFallback]

    CB_LOAD --> LOADED_MODEL
    NS_LOAD --> LOADED_MODEL
    VLM_LOAD --> LOADED_MODEL
    M3_TEXT --> LOADED_MODEL
    M3_LLAMA --> LOADED_MODEL

    LOADED_MODEL[LoadedModel enum variant] --> TOK_LOAD[load_tokenizer\ntokenizer.json / tokenizer.model / .tiktoken]
    TOK_LOAD --> DONE([LoadedModel + MlxcelTokenizer])
    EMB_LOAD --> EMB_DONE([EmbeddingModel + tokenizer])
    RR_LOAD --> RR_DONE([Reranker + tokenizer])

Step 1: Detection¶

The config.json Contract¶

Most HuggingFace checkpoints publish config.json. mlxcel treats model_type as the primary architecture key, then consults architectures, pooling modules, bidirectional flags, classifier heads, and family-specific fields to determine the task and exact variant.

The detection function in src/models/detection.rs reads this file, sanitizes it to work around known formatting quirks in published models, and maps the model_type string to a ModelType enum variant.

config.json → { "model_type": "llama", ... } → ModelType::Llama
config.json → { "model_type": "gemma3", ... } → ModelType::Gemma3
config.json → { "model_type": "qwen2_vl", ... } → ModelType::Qwen2VL

If the available metadata cannot select a supported layout, loading fails with a descriptive error. There is no silent fallback to a default architecture or task.

ModelType Enum¶

ModelType in src/models/mod.rs enumerates every registered architecture. The exact set is printed by mlxcel arch and is organized into families:

Special Detection Cases¶

Several architectures require more than just reading the model_type string.

Embedding and reranker split: Encoder-only families, known embedding architectures, sentence-transformers Pooling modules, and bidirectional layouts select ModelKind::Embedding. A supported ForSequenceClassification architecture selects reranking instead, even if the same BERT or ModernBERT trunk can also embed. Generative Qwen3 and Qwen3-VL rerankers require --reranker-model because their checkpoint config is intentionally indistinguishable from an ordinary generator.

Text vs VLM split: Several model_type values are used by both a text-only and a vision-language variant of the same architecture. For example, gemma3 maps to ModelType::Gemma3 normally but to ModelType::Gemma3VLM when the config contains a vision_config key. The detect_text_or_vlm() helper handles this pattern uniformly. A VLM model_type loads text-only, without a vision tower, when vision_config is absent, null, or an empty object, when top-level language_model_only is true (which wins even if vision weights ship), or when the checkpoint has no vision_tower.* or model.visual.* weights. Image and video requests to such a model fail with model <path> was loaded without a vision tower (text-only checkpoint).

Hunyuan density: The hunyuan model type uses num_experts to distinguish between MoE (HunyuanMoe) and dense (HunyuanV1Dense) variants.

Mistral3 nesting: Mistral3 wraps a nested text_config object. The detection outcome is ModelType::Mistral3, but the policy routing step later inspects the inner text_config.model_type to decide whether to load the Ministral3 text tower or fall back to Llama.

Shared model_type strings: llama and mistral both map to ModelType::Llama because Mistral ½ use the same architecture as Llama ½. Similarly, mamba and falcon_mamba both map to ModelType::Mamba.

Step 2: Policy Routing¶

Once the ModelType is known, the function model_load_policy() in src/model_metadata.rs looks up the StaticModelDescriptor for that type and produces a ModelLoadPolicy. This policy captures two routing decisions:

1. Directory route: which loader to call when constructing the model from a directory.
2. Weight route: which loader to call when constructing the model from pre-loaded weights (used by the LoRA adapter path).

The registry is declared with the for_each_model_registration! macro. Every model type is listed once, with all its routing metadata co-located. Adding a new model type therefore forces both routes to be declared before any loading code exists.

Directory Route Families¶

Weight Route Families¶

VLM model types have weight_route: None because LoRA adapter loading is not supported for VLM families.

Step 3: The Four Loaders¶

ConfigBacked Loader¶

The config_backed loader in src/loading/config_backed.rs handles the majority of text models. It is generated entirely by the for_each_model_registration! macro and requires no hand-written match arms for each new architecture.

The process is: 1. Call the model's load() function with the directory path string. 2. The load() function reads config.json, deserializes it into the model's ModelArgs struct, loads *.safetensors weight files, and constructs the model layer graph. 3. Wrap the result in the appropriate LoadedModel variant.

For models that use wrapper types (Gemma3, ExaOne4, Ministral3, Llama4, GptOss), the wrap closure in the registration entry applies the wrapper before storing the model in LoadedModel.

Nonstandard Loader¶

The nonstandard loader in src/loading/nonstandard.rs handles architectures that have unusual weight naming, non-standard config structures, or preprocessing steps that do not fit the config-backed pattern. Examples include Mamba (which does not use a standard transformer config), Jamba (which combines Mamba and attention layers), RWKV7, and certain hybrid models.

Each nonstandard architecture has its own named loading function.

Special Loader¶

The special loader in src/loading/special.rs is invoked when loading a model from pre-loaded weights rather than from a directory. This path is used exclusively by the LoRA adapter fusion workflow, where base model weights and adapter delta weights are fused in memory before constructing the model. Architectures with unusual weight sanitization or ownership requirements are routed here.

VLM Loader¶

The vlm loader in src/loading/vlm.rs handles vision-language models. Each VLM family has a dedicated function that: 1. Loads the visual encoder (e.g., SigLIPVisionModel, CLIPVisionModel, PatchViT). 2. Loads the text decoder with VLM-specific configuration. 3. Loads the cross-modal projection or connector. 4. Wraps all three components in the family-specific VLM struct.

The try_load_vlm_model_from_dir() function dispatches by ModelType to the correct family loader.

The LoadedModel Enum¶

LoadedModel in src/loaded_model.rs is the central runtime handle for any loaded model. It is a Rust enum with one variant per supported model family.

pub enum LoadedModel {
    Llama(models::Llama3Model),
    Gemma3(models::Gemma3Wrapper),
    Gemma3VLM(vision::VisionLanguageModel),
    Qwen2VL(vision::Qwen2VLModel),
    Mamba(models::MambaModel),
    // ... 80+ variants
}

LoadedModel implements the LanguageModel trait by delegating each method through the delegate_language_model! macro. This macro generates a match arm for every variant, calling the corresponding inner model's trait implementation.

The benefit of this design is that the rest of the codebase (the generation loop, the server scheduler, the CLI) can work with a single &dyn LanguageModel or &LoadedModel without knowing which architecture is loaded. At the same time, capability dispatch (such as checking whether a VLM embedding path is needed) can match on the specific variant when needed.

Weight Loading Mechanics¶

SafeTensors Format¶

mlxcel loads model weights exclusively from the SafeTensors format, the same format used by HuggingFace. The mlxcel_core::weights::load_weights_from_dir() function scans the model directory for *.safetensors files, reads them with MLX's load_safetensors, and exposes the tensors through a WeightMap.

Each generation-model load entry point runs the family loader inside a weight materialization scope (with_weight_materialization in src/loading/materialize.rs, policy in mlxcel_core::weight_materialize). When the loader returns, the scope evaluates and uploads every checkpoint tensor the model kept, so Model loaded in ... reports resident weights and the first token pays no weight I/O. Two groups stay lazy: modality-tower tensors (vision and audio towers load on the first image or clip) and tensors the loader dropped, such as an unused MTP head or a tied lm_head that was sanitized away. Loads outside a scope (pipeline stages, embedding loaders) stay lazy.

MLXCEL_WEIGHT_MATERIALIZATION=lazy restores first-use materialization, where the read happens inside the first evaluation that needs a tensor. On CUDA the eager upload stages tensors through four reused 64 MiB pinned host buffers and falls back to pageable memory with one warning if the pinned allocation fails; MLXCEL_CUDA_PINNED_LOAD=0 forces the pageable path. MLXCEL_PROFILE_TTFT prints a [TTFT] load line per model load.

External Weight Stores¶

Two qwen4_exp structures can live outside MLX parameters when the checkpoint says so.

A checkpoint that declares text_config.ple_storage keeps its PLE n-gram table on disk, about 30 GiB, and reads rows on demand instead of holding the table in parameters. The mlx-vlm v2 manifest is accepted as written: safetensors_ranges, interleaved_q4, affine 4/32, and nvfp4 4/16, with every validation failure naming the field. The manifest and row file must resolve inside the model directory, so an absolute ple_storage.manifest is refused, and each safetensors_ranges tensor must match its file's own SafeTensors header. More than one PLE layer is refused, because one manifest describes one table. MLXCEL_PLE_IO selects mmap (the default), direct, or ram, and a byte-bounded LRU row cache of MLXCEL_PLE_CACHE_BYTES bytes (default 0.5 GiB) serves repeated ids; gather-ahead reads the next prefill piece's rows while the GPU runs the current one. mlxcel prepare-external-ple builds the view without Python. mlxcel inspect and --estimate-memory count the row cache, or the table under ram, rather than the on-disk table, and generation ends with a stats line for hits, misses, bytes, and time. Two limits are known: truncating a file while mmap maps it raises SIGBUS, which direct avoids, and a disk read error panics the forward because that path carries no error channel.

Routed experts can stream from the SSD under a per-layer slot arena while everything else stays resident. The arena is filled on demand: each layer evaluates the router's top-k on the host once, reads misses into LRU victims the window cannot evict, and runs gather_qmm over slot indices. Slot tensors are page-aligned host allocations wrapped as MLX arrays without a copy, and a slot is rewritten only after every command that reads it has finished. Sorted prefill reproduces the resident sort and plans batches so MLX selects the kernel and tile the resident tensor would get; a window that cannot be planned that way logs a one-time warning that its low bits may differ. MLXCEL_EXPERT_STREAM takes auto (the default), 1, or 0. Under auto the mode engages only when the routed experts exceed 60% of Metal's working set, so a large machine stays resident. MLXCEL_EXPERT_ARENA_BYTES fixes the arena size; otherwise it is derived from free memory with 6 GiB kept for the KV cache and transients, at least 64 slots and in multiples of 16 so every plane is whole 16 KiB pages. mlxcel prepare-expert-stream writes a contiguous, hash-checked file that the loader then uses automatically. Startup logs an [Expert arena: ...] line, mlxcel_expert_arena_* metrics report activity, preflight accounting covers the arena, a drafter on a machine below 96 GiB logs a warning, and a malformed or out-of-directory checkpoint is refused. At a 100% hit rate the per-layer sync costs about 15 ms per token on the validated host.

Tied Embedding Restoration¶

Some models use weight tying: the language model head (lm_head) shares the same tensor as the token embedding table (embed_tokens). When such a model is serialized, the weight may be stored only once. mlxcel's sanitize_tied_embeddings() function detects this case by inspecting the tie_word_embeddings flag in the config and copies the embedding tensor to the lm_head key if the head weights are absent.

This step runs before the model is constructed from weights, so model code can always find lm_head weights without special-casing tied embeddings.

Config Sanitization¶

The sanitize_config_json() function in src/models/sanitize.rs preprocesses the raw config string before deserializing it. This corrects known formatting issues in published HuggingFace models, such as null values where numeric fields are expected and architectural quirks introduced by fine-tuners who modify config files without preserving strict types.

Sanitization happens before deserialization, not after, so the type system's guarantees hold once a ModelArgs struct is constructed.

Quantization Support¶

mlxcel-core's weight loading supports quantized models stored in MLX's native quantization format. Quantized weights are detected automatically based on the presence of quantization metadata in the SafeTensors file. The generation loop and layer math in mlxcel-core handle quantized tensors transparently through the MLX computation graph.

Quantization formats¶

Per-layer bit overrides¶

Some checkpoints emit mixed bit widths across a single model: the bulk of the weights are 4-bit, but a handful of selected layers (router gates in MoE models, shared-expert gates, a specific embedding head) are quantized at 8-bit. MLX produces these shapes via nn.quantize(class_predicate=...), and mlx-community/Qwen3.5-35B-A3B-4bit plus the newer Qwen3.6 variants are the first checkpoints in widespread distribution that exercise this pattern.

The loader calls infer_quantization_bits(weight_shape, scales_shape, group_size, caller_bits) for every affine-quantized weight it instantiates. The MLX invariant packed_in * 32 == bits * num_groups * group_size is verified against the tensor's actual shape, and the inferred bit width takes precedence when it disagrees with the caller-supplied bits. The forward path itself relies on QuantizedWeight.bits / group_size as the source of truth consumed by quantized_linear_forward, so the probe affects only weight construction. Pure-format quantizations (mxfp4, nvfp4, mxfp8) have fixed per-mode bits and skip it.

Practical consequence: checkpoints that mix 4-bit and 8-bit tensors load and run without a config override, and qwen3.6-35b-a3b-4bit (which uses the existing qwen3_5_moe architecture with per-layer 8-bit gate overrides) works out of the box. Benchmarks on qwen3.5-35b-a3b-4bit and llama-3.1-8b-4bit show no measurable change in load time or decode throughput; the bit-width probe is a constant-time shape check per weight.

Load-Time Surgery Hook¶

The consolidated text and VLM weight loaders expose a WeightTransform hook that can modify the in-memory WeightMap before the model graph is constructed. --surgery <FILE> and MLXCEL_SURGERY install a process-wide pipeline that implements this hook from a YAML file.

For text models, the hook runs after standard weight sanitization such as tied-embedding restoration and quantized metadata normalization, and before non-quantized bf16 weights are converted to f16 on Apple Silicon. For VLMs, the hook runs in the common VLM weight loader, so rules should normally target checkpoint or common-loader key names rather than later family-specific remaps.

The hook is fail-fast: operations run in YAML order, and model loading aborts if an operation cannot match, reshape, type-check, or safely update the targeted tensors. This makes surgery configurations reproducible and avoids silently serving a partially modified model.

See Runtime Model Surgery for the user-facing YAML schema and examples.

LoRA Adapter Loading¶

mlxcel applies an adapter either fused or as an unfused runtime term. The server's default is the unfused path, which keeps each adapter as a low-rank term with a live scale; see LoRA Adapters. On the fused path (mlxcel generate --adapter, and the server with --lora-fuse or pipeline parallelism), mlxcel fuses weights before constructing the model:

1. Load base model weights from the model directory using load_weights_from_dir().
2. Load adapter weights from the adapter directory.
3. Fuse the two weight sets using lora::apply_lora_adapters(), which adds the low-rank delta matrices to the corresponding base weight tensors.
4. Pass the fused WeightMap to load_model_from_weights(), which uses the weight route (not the directory route) to construct the model.

The fused model is indistinguishable from a standard model at runtime. The LoRA structure is dissolved during fusion; the merged model has standard weight shapes.

VLM model types reject adapter loading with an explicit error message. This restriction is noted in the adapter_unsupported_message field of each VLM's StaticModelDescriptor.