ModelWeightLoader

spark_model::weight_loader

Trait ModelWeightLoader 

Source 
pub trait ModelWeightLoader {

Show 15 methods    // Required methods
    fn supports_tp(&self) -> bool;
    fn load_layers(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
        layer_kv_dtypes: &[KvCacheDtype],
    ) -> Result<Vec<Box<dyn TransformerLayer>>>;
    fn load_embedding(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight>;
    fn load_final_norm(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight>;
    fn load_lm_head(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<DenseWeight>;
    fn load_mtp_weights(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<Option<MtpWeights>>;

    // Provided methods
    fn prune_after_load(
        &self,
        _store: &mut WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
    ) -> Result<()> { ... }
    fn precision_schedule(&self, _config: &ModelConfig) -> PrecisionSchedule { ... }
    fn load_ngram_embedding(
        &self,
        _store: &WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
        _max_tokens: usize,
    ) -> Result<Option<NgramEmbedding>> { ... }
    fn load_mtp_weights_multi(
        &self,
        store: &WeightStore,
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
    ) -> Result<Vec<MtpWeights>> { ... }
    fn kv_layer_dims(&self, _config: &ModelConfig) -> Vec<(usize, usize)> { ... }
    fn load_dflash_weights(
        &self,
        _drafter_store: &WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
        _tp_size: usize,
    ) -> Result<Option<DflashWeights>> { ... }
    fn load_lora_adapters(
        &self,
        adapters: &[LoraAdapterInput<'_>],
        config: &ModelConfig,
        gpu: &dyn GpuBackend,
        max_loras: usize,
        max_lora_rank: usize,
    ) -> Result<Option<LoraWeights>> { ... }
    fn binds_vision_encoder(&self) -> bool { ... }
    fn load_vision_encoder(
        &self,
        _store: &WeightStore,
        _config: &ModelConfig,
        _gpu: &dyn GpuBackend,
    ) -> Result<Option<VisionEncoder>> { ... }

}
Expand description

Loads weights from a WeightStore into typed layer objects.

Required Methods§

Source

fn supports_tp(&self) -> bool

Whether this loader’s weight slicing is TP-aware. No default — every loader MUST declare this explicitly so adding a new model architecture cannot accidentally inherit a false and silently regress users who pass --tp-size > 1.

Loaders that honour config.tp_world_size / config.tp_rank when loading attention Q/K/V/O, MoE gate/up/down, head-parallel SSM components, and lm_head return true. Loaders that always load full replicated weights return false.

The startup path in spark-server/src/main.rs consults this method to fail-fast at load time when --tp-size > 1 is requested against a TP-unaware loader. Extending TP to a new architecture requires:

  1. Wire slice_for_rank (in crate::tp_shard) per Q/K/V/O, gate/up/down, and any head-parallel SSM tensors.
  2. Divide num_attention_heads / num_key_value_heads per the same axis when constructing layer state.
  3. Return true from this method.

See weight_loader/minimax.rs for the reference implementation.

Source

fn load_layers( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, layer_kv_dtypes: &[KvCacheDtype], ) -> Result<Vec<Box<dyn TransformerLayer>>>

Load all transformer layers from the weight store.

layer_kv_dtypes is indexed by attention layer index (0-based sequential counter over full-attention layers only). Each attention layer receives its own KV cache dtype, enabling mixed-precision KV caching where boundary layers use higher precision.

Source

fn load_embedding( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<DenseWeight>

Source

fn load_final_norm( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<DenseWeight>

Load the final RMSNorm weight used before the LM head.

gpu is passed so model-specific loaders can do on-device weight transforms at load time (e.g. Gemma-4 shifts the learned absolute- scale weight by -1 into the offset-from-1 convention expected by Atlas’s rms_norm kernel). Loaders that don’t need it should ignore the argument.

Source

fn load_lm_head( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<DenseWeight>

Source

fn load_mtp_weights( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<Option<MtpWeights>>

Load MTP head weights (returns None if no MTP weights in store).

Provided Methods§

Source

fn prune_after_load( &self, _store: &mut WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, ) -> Result<()>

Drop store tensors this loader has finished with, after every load_* reader has run and before the buffer arena / KV cache are sized.

Default: keep everything. That is correct for the loaders that bind zero-copy from the store’s device pointers — the store IS the model’s weights, and TransformerModel releases it at teardown.

Override only when the loader uploads its own copies (a TP shard, a host round-trip, a dtype conversion), because then the store’s originals are dead the moment the binder returns. On unified-memory GB10 that duplicate comes straight out of the KV budget.

Source

fn precision_schedule(&self, _config: &ModelConfig) -> PrecisionSchedule

Per-(layer, role) weight precision schedule (C.3, 2026-04-25). Default impl returns the empty schedule (every lookup yields Dtype::Inherit), preserving the existing per-checkpoint dtype logic byte-for-byte. Loader-specific implementations can override to honour MODEL.toml’s [precision] block.

Source

fn load_ngram_embedding( &self, _store: &WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, _max_tokens: usize, ) -> Result<Option<NgramEmbedding>>

Build the n-gram embedding, when this architecture fuses hashed n-gram lookups into the input embedding (LongCat / Qwen3.8-Flash-Next).

Separate from load_embedding because the result is NOT a weight: it is a small engine that needs the sequence’s CONTEXT token ids at forward time, not just the id being embedded. Returning None — the default — leaves the plain embed_tokens gather in place.

Source

fn load_mtp_weights_multi( &self, store: &WeightStore, config: &ModelConfig, gpu: &dyn GpuBackend, ) -> Result<Vec<MtpWeights>>

Load MTP weights for multi-module MTP (DeepSeek-V3 / MiniMax-M2 style: N independent transformer modules, each with its own attention + MoE + KV cache). Returns an empty Vec when the checkpoint has no MTP modules, a 1-element Vec for single-module MTP (Qwen3.5 family), or N elements for multi-module.

Default impl adapts load_mtp_weights so existing single-module loaders don’t need to change. MiniMax overrides this directly.

Source

fn kv_layer_dims(&self, _config: &ModelConfig) -> Vec<(usize, usize)>

Per-layer (num_kv_heads, head_dim) overrides for heterogeneous attention models (e.g. Gemma-4 with sliding 16×256 and full 4×512). Default empty — homogeneous models skip per-layer dims and the KV cache allocator uses the global (num_kv_heads, head_dim). Populated by loaders whose models have different attention geometries per layer. Indexed by attention layer index (same as layer_kv_dtypes).

Source

fn load_dflash_weights( &self, _drafter_store: &WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, _tp_size: usize, ) -> Result<Option<DflashWeights>>

Load DFlash drafter weights from a separate WeightStore pointing at the drafter checkpoint (z-lab/Qwen3.6-{27B,35B-A3B}-DFlash). Default impl returns None so loaders that don’t yet support DFlash silently fall through to the existing MTP path. Override in loaders whose target models pair with a DFlash drafter (Qwen3.5/3.6 family). The same drafter format works across both 27B-dense and 35B-A3B-MoE targets — only the target_hidden_size validated against the drafter’s fc input dimension differs.

Source

fn load_lora_adapters( &self, adapters: &[LoraAdapterInput<'_>], config: &ModelConfig, gpu: &dyn GpuBackend, max_loras: usize, max_lora_rank: usize, ) -> Result<Option<LoraWeights>>

Load one or more startup-static PEFT LoRA adapters from their own WeightStores (the adapter_model.safetensors tensors, already on-device BF16) into the fixed-address rank-padded pool (one slot each).

Unlike load_dflash_weights’ vestigial Ok(None) default, the default here is a WORKING model-agnostic implementation (the remap needs only ModelConfig::layer_type + projection dims); families needing a bespoke key remap override it. Called from factory::build_model BEFORE the buffer arena + KV sizing so the pool bytes are budgeted against the KV cache. A single-element slice is byte-identical to the pre-multi-adapter single-adapter path.

Source

fn binds_vision_encoder(&self) -> bool

Will this loader ever bind a vision encoder for a multimodal checkpoint?

Default true — “load everything” is the safe answer, so a loader that forgets to override this can never lose weights it needs. A loader whose port is deliberately text-only overrides it to false, and the weight loader then skips the tower’s tensors instead of reading a gigabyte of unified memory that nothing will bind. build_model still frees an unbound tower afterwards (keyed off the bind result, not off this), so this is a peak-memory optimisation, not the correctness gate.

Source

fn load_vision_encoder( &self, _store: &WeightStore, _config: &ModelConfig, _gpu: &dyn GpuBackend, ) -> Result<Option<VisionEncoder>>

Load vision encoder weights (returns None for text-only models).

Implementors§