Per-Component Quantized Checkpoint Loading

Created: · Last updated:

Status: Proposed Date: 2026-08-26

Implemented precedents

The complete component manifest, codec registry, and ModelWeightAdapter design below remain proposed. Native GPT-OSS MXFP4 support has since landed several narrower building blocks that should be reused by the migration:

  • QuantizedWeightFormat distinguishes integer-affine storage from MXFP4 instead of inferring semantics from bit width and group size.

  • StreamingWeightPlan supports direct, expert-bank, and transformed lazy sources with fail-closed source/target metadata validation.

  • Architecture-specific planners can map checkpoint headers to graph initializers without materializing a complete state dict.

  • The native GPT-OSS loader registers its lazy source artifacts and marks its report for transactional ModelPackage.save() publication. These lifecycle protections are opt-in requirements, not yet automatic for every streaming loader.

These are precedents, not an implementation of the proposed QuantizationCodec or ModelWeightAdapter registries. The current Transformers builder still dispatches exceptional loaders such as GPT-OSS and Qwen4 explicitly.

Executive summary

Mobius must load pre-quantized multi-component checkpoints in which each exported component can use a different weight layout. For example:

decoder        Olive INT4, group size 32
vision_encoder Olive INT8, group size 64
audio_encoder  floating point
embedding      floating point

Mobius does not quantize floating-point weights. Quantization remains the responsibility of Olive, GPTQ, AWQ, or another upstream producer. Mobius reads the producer’s metadata and packed tensors, builds matching ONNX parameter scaffolding, converts supported packed layouts into the canonical ORT layout, and binds the existing quantized values.

The current architecture was designed around one model-wide quantization_config. Multi-component tasks later split the ONNX graphs, but configuration and weight preprocessing remained model-wide. As a result, model-specific preprocess_weights() methods now mix two different concerns:

  1. architecture-specific semantic transforms such as rename, transpose, split, fuse, and tied-weight synthesis; and

  2. generic checkpoint mechanics such as component routing, quantization format dispatch, packed-sidecar grouping, validation, and binding.

This proposal separates those concerns through four reusable abstractions:

  1. ComponentManifest

  2. WeightRecord and WeightBundle

  3. QuantizationCodec

  4. ModelWeightAdapter

Goals

  • Represent one quantization configuration per ModelPackage component.

  • Allow an authoritative component configuration to keep selected modules floating point through exact or regex exclusions.

  • Determine component configuration before component module construction.

  • Route checkpoint weights with one canonical component manifest.

  • Preserve packed qweight, scale, and zero-point tensors as one logical record throughout name mapping.

  • Keep float-to-integer quantization outside Mobius.

  • Keep architecture-specific transforms available without making them own the generic loading pipeline.

  • Fail before publication when metadata, graph parameter shapes, and packed checkpoint values disagree.

  • Preserve current single-component and global-configuration behavior during migration.

Non-goals

  • Calibration, range collection, or quantization parameter estimation.

  • Floating-point to INT2/INT4/INT8 conversion.

  • A universal representation for native non-affine GGUF codebooks. Those formats retain their existing typed import path.

  • Eliminating every model-specific weight transform. Some architectures inherently require semantic transforms that cannot be inferred safely.

Current architecture and failure mode

The current Transformers path is effectively:

HF config
  -> one ArchitectureConfig.quantization
  -> instantiate one top-level module
  -> task splits module into component graphs
  -> download one flat state_dict
  -> model.preprocess_weights(state_dict)
  -> ModelPackage.apply_weights(state_dict)

This is coherent when all graph components share one layout. It becomes ambiguous when the decoder is INT4 and the vision encoder is INT8:

  • component identity exists in the task;

  • HuggingFace source ownership exists on the model class;

  • optimization roles exist in model_roles;

  • but no single object combines those facts;

  • and preprocess_weights() receives the entire checkpoint plus one global quantization configuration.

The result is either applying decoder packing parameters to every component or adding model-specific conditionals that reimplement component routing.

Proposed architecture

Checkpoint metadata
        |
        v
ComponentManifest + ComponentQuantizationManifest
        |
        v
Component-specific module construction
        |
        v
CheckpointReader -> WeightRecord stream
        |
        v
ComponentRouter
        |
        v
QuantizationCodec
        |
        v
ModelWeightAdapter
        |
        v
WeightBinder + publication validation

1. ComponentManifest

ComponentManifest is the authoritative description of package components. It initially resolves existing declarations without requiring every model to migrate in one change.

@dataclasses.dataclass(frozen=True)
class ComponentDescriptor:
    name: str
    # Dotted Python attribute path from the root nn.Module passed to task.build().
    # This is neither a ModelPackage key nor a checkpoint key prefix.
    module_attribute_path: str
    role: str
    source_paths: tuple[str, ...]

    def source_module_names(self, local_module_path: str) -> tuple[str, ...]: ...


@dataclasses.dataclass(frozen=True)
class ComponentManifest:
    components: tuple[ComponentDescriptor, ...]

    def by_name(self, name: str) -> ComponentDescriptor: ...

The resolver combines:

  • ModelTask.components

  • ModelTask.model_roles

  • model HF_COMPONENT_SOURCES

Every consumer uses the same manifest:

  • component inspection;

  • module construction;

  • optimization role selection;

  • weight routing;

  • package persistence;

  • component quantization selection.

An unresolved component is a typed error. Consumers must not independently guess aliases or state-dict prefixes.

Dynamic get_hf_component_sources(model_type=..., hf_config=...) hooks use the explicit model type, or hf_config.model_type when it is not supplied. If neither is available, the dynamic hook is skipped without guessing static source paths. Classes without a dynamic hook still publish their HF_COMPONENT_SOURCES for manually constructed configs that omit the model type.

Mapping local module paths to HuggingFace names

Quantizer metadata names runtime HuggingFace modules, while graph construction walks paths inside a Mobius component. For example:

HF:     model.language_model.layers.0.per_layer_input_gate
Mobius: model.layers.0.per_layer_input_gate

ComponentDescriptor.source_module_names() derives candidate HuggingFace names by anchoring the local path at each declared source root. The shared matcher then applies the producer’s exact/regex semantics to those candidates. Architectures whose names cannot be aligned structurally declare explicit path aliases on their component descriptor; individual loaders must not invent their own mapping.

Component-owned source roots also prefix local descendants without a shared anchor segment, such as T5’s encoder and decoder stacks. Models can publish static HF_COMPONENT_MODULE_ALIASES or implement get_hf_component_module_aliases(hf_config=...) when aliases depend on configuration, such as distinct encoder and decoder layer counts. These config-based aliases do not require a model type and remain available for manually constructed T5 configs. An existing alias or separately declared source takes precedence over synthesizing a component-root name; a top-level lm_head must not also become decoder.lm_head for rule matching.

2. Component quantization manifest

The parsed architecture config keeps legacy compatibility while exposing an authoritative component mapping:

@dataclasses.dataclass(frozen=True)
class ComponentQuantizationManifest:
    components: Mapping[str, QuantizationConfig | None]

    def for_component(self, name: str) -> QuantizationConfig | None: ...

Accepted metadata forms can include:

  • explicit component_quantization;

  • quantization_config.components;

  • nested vision_config.quantization_config;

  • nested audio_config.quantization_config;

  • a legacy model-wide configuration converted at the parser boundary.

Missing entries mean floating-point storage. A component entry must describe one affine layout for its quantized projections, but may include an authoritative selection policy:

@dataclasses.dataclass(frozen=True)
class ComponentQuantizationPlan:
    layout: QuantizationConfig
    modules_to_not_convert: tuple[str, ...]
    overrides: Mapping[str, QuantizationOverride]

modules_to_not_convert is evaluated for each candidate HuggingFace module name while the component is constructed. A matching module remains a normal Linear; other eligible projections use the component’s packed layout. Literal exclusions and overrides match the complete module path or a dot-separated descendant, never an arbitrary substring: model.layers.1 matches model.layers.1.self_attn but not model.layers.10. Patterns prefixed with re: continue to use full-match regular expressions. Different packed layouts inside one component require per-module overrides and must be represented explicitly rather than inferred from a root plan.

Legacy model-wide module rules apply when the resolved manifest owns one component, regardless of its package key. Graph construction and normalization use the same manifest-based decision; passing a one-component subset of a multi-component package does not enable those rules.

3. Component-specific construction

The component configuration is selected before creating its parameters:

decoder = DecoderModel(config.for_component("decoder"))
vision = VisionModel(config.for_component("vision_encoder"))

Each projection additionally asks the component plan whether its source module is quantized:

linear_class = component_plan.linear_class_for(
    local_module_path="model.layers.0.per_layer_input_gate",
    descriptor=decoder_descriptor,
)

This keeps per_layer_input_gate, per_layer_projection, and other explicit float exceptions as ordinary Linear modules without weakening the rest of a quantized decoder.

Graph construction and packed-weight normalization use the same effective module policy. Model-declared COMPONENT_OUTPUT_HEADS stay float unless quantize_lm_head is enabled. A module’s component_quantization_excluded_methods applies to that module and all its descendants; selecting a quantized component must not override this invariant. This includes the deliberately floating-point shared-expert gates in Qwen2-MoE and Qwen3-Next. quantize_embeddings converts input token tables identified by their local or aliased source names, not learned positional embeddings.

This replaces post-construction scanning that tries to swap Linear instances after a model has already encoded architecture-specific choices.

During migration, top-level model classes can expose component factories:

class ModelComponents:
    def build_component(self, descriptor, config):
        ...

The task calls those factories using the resolved manifest. Existing top-level modules remain supported behind a compatibility adapter until their components migrate.

4. WeightRecord

Checkpoint sidecars are grouped into one typed logical weight immediately after reading metadata:

@dataclasses.dataclass(frozen=True)
class FloatWeight:
    value: torch.Tensor


@dataclasses.dataclass(frozen=True)
class AffinePackedWeight:
    qweight: torch.Tensor
    scales: torch.Tensor
    zero_points: torch.Tensor | None
    bits: int
    group_size: int
    symmetric: bool
    method: str


@dataclasses.dataclass(frozen=True)
class WeightRecord:
    source_name: str
    component: str
    storage: FloatWeight | AffinePackedWeight

Grouping prevents several raw sidecar keys from being renamed, split, or overwritten independently. Missing scales are rejected during grouping. Zero-point requirements are validated after resolving the module’s effective symmetry, since an override may make one projection symmetric inside an otherwise asymmetric component.

WeightBundle owns records for one component:

@dataclasses.dataclass
class WeightBundle:
    component: ComponentDescriptor
    records: dict[str, WeightRecord]

5. QuantizationCodec

Format-specific layout handling lives in a registry:

class QuantizationCodec(Protocol):
    method: str

    def group(
        self,
        component: ComponentDescriptor,
        state_dict: Mapping[str, torch.Tensor],
        config: QuantizationConfig,
    ) -> WeightBundle: ...

    def normalize(
        self,
        record: WeightRecord,
        config: QuantizationConfig,
        *,
        kind: Literal["linear", "embedding"] = "linear",
    ) -> dict[str, torch.Tensor]: ...

Initial codecs:

  • Olive affine packed weights

  • GPTQ

  • AWQ

The codec:

  • validates the packed payload and quantization parameters;

  • checks logical and physical shapes;

  • repacks supported source layouts into canonical ORT layout;

  • never computes quantization parameters from float values.

Model names and model-specific parameter paths are prohibited inside codecs.

The target module supplies the weight kind: linear projections use the three-dimensional MatMulNBits layout, while embedding tables use the two-dimensional GatherBlockQuantized layout. A shared .scales suffix alone does not identify a canonical group; the packed tensor, scales, and required zero points must all match the target parameters before normalization is skipped. The component normalizer also retains complete canonical QMoE expert groups. When component plans or single-component module rules are active, it rejects unconsumed raw sidecars even when no qweight is present or no component owns them; valid canonical scales must not be mistaken for orphan checkpoint data.

6. ModelWeightAdapter

The current implementation retains a raw-state compatibility adapter for architecture-specific name mapping, tying, and expert transformations:

@dataclasses.dataclass(frozen=True)
class WeightAdapterContext:
    config: BaseModelConfig
    manifest: ComponentManifest


class ModelWeightAdapter(Protocol):
    def adapt(
        self,
        module: nn.Module,
        state_dict: Mapping[str, torch.Tensor],
        context: WeightAdapterContext,
    ) -> dict[str, torch.Tensor]: ...

Routing and codec grouping happen after this adapter. With component plans or single-component module rules, ordinary packed sidecars stay raw until the shared loader resolves their per-projection layouts. The shared legacy preprocessor can still pack QMoE expert banks and tie floating-point tables. A per-component WeightBundle adapter is a future migration target, not the currently implemented interface.

Native QMoE additionally has a projection-specific adapter for Olive fused K-last routed experts. Under a model-wide INT4 fallback, supported FC1/FC2 expert widths are (2,4), (4,8), (8,4), and (8,8), with one common power-of-two group size of at least 16. A model-wide INT8 fallback is not enabled. Exact and re: Olive overrides are resolved against the original Hugging Face paths; plain Olive exclusions retain their producer-defined substring matching. GPTQ and AWQ instead use the generic plain-path subtree and re: full-match rules. FC1 and FC2 buffers are then validated independently. The adapter requires every routed layer and validates rank, dtype, expert count, packed byte width, scale geometry, and optional zero-point geometry before binding. Qwen3.5-VL keeps the decoder’s full module plan (preserve_module_plan=True) while constructing the split package: expert overrides are resolved against authoritative model.language_model source paths, while validation and binding use the decoder-prefixed checkpoint roots created by the VL weight router. It emits QMoE’s FC-specific bit attributes for differing widths (including FC3=FC1 for fused SwiGLU). Uniform INT4 retains the legacy payload path; uniform 8/8 uses global expert_weight_bits=8 without FC-specific attributes. ORT’s merged CUDA fallback in microsoft/onnxruntime#32743 executes integer mixed-width QMoE, including tested (4,8) and (8,4) widths. It dequantizes both expert banks into FP16/BF16 scratch, subject to canonical raw weights, 3-D blockwise scales, block size 16–256, divisible reduction dimensions, and a scratch-memory limit. ORT tests cover these INT8 combinations without fused SwiGLU; only (2,4) has a fused-SwiGLU test. The opt-in tests/integration/qmoe_cuda_test.py runs a synthetic, nontrivial Qwen3-MoE export with adjacent (4,8)/(8,4) layers and uniform (8,8) on an ORT build containing #32743. It checks full logits against an independently dequantized reference (FP16 rtol=atol=0.01) and profiles QMoE on CUDA. The same opt-in test also reads layer 0, experts 0–3 of the pinned Qwen3-30B-A3B checkpoint, quantizes those real weights with Olive’s RTN tensor producer, saves and reloads the sidecars, and exports one reduced four-expert decoder layer. On ORT CUDA at ee5f6e7, both (4,8) and (8,4) execute QMoE on CUDA and match an independently unpacked full-logits reference within rtol=atol=0.01 (observed maximum absolute differences 0.005148 and 0.004774, respectively). The checkpoint embedding and router are sliced; attention is zeroed, and the output head and normalization weights are artificial. Set MOBIUS_QMOE_REAL_SLICE_CUDA_TEST=1 to run this local-only test with the pinned checkpoint, Olive producer source, and compatible ORT CUDA build; it saves the selected tensor hashes, serialized sidecars, and profiles under /datadisks/disk5/titaiwang/qmoe-real-slice-744/. This demonstrates real expert-weight loading and reduced routing, not full-model Hugging Face parity or complete Olive pass orchestration. Uniform (8,8) takes the existing packed INT8 path, which also requires a block size divisible by 32 and both reduction dimensions divisible by 64; the mixed dequantization fallback accepts a block size of 16. These tiny synthetic and reduced real-weight checks do not qualify production-size throughput or peak VRAM. For Qwen3-30B-A3B, the FP16 FC1+FC2 scratch alone is approximately 1.125 GiB per MoE layer, above the fallback’s default 1 GiB limit; adjust ep.cuda.qmoe_int_dequant_max_scratch_bytes and measure the actual memory cost before deployment. CPU mixed-width execution remains unverified. The packed decode path in microsoft/onnxruntime#32761 is open (not merged) and covers only the 2/4-bit set, not INT8. Do not silently replace an unsupported QMoE configuration with the dense expert-loop exporter.

Appropriate model-specific operations include:

  • HuggingFace-to-ONNX name alignment;

  • transpose;

  • fused QKV split or merge;

  • expert-bank stacking;

  • model-specific tied-weight synthesis;

  • special projector or wrapper removal.

The adapter must not:

  • override the producer’s per-projection quantization policy;

  • eagerly repack ordinary sidecars using one model-wide or component-wide layout;

  • duplicate generic Olive/GPTQ/AWQ format conversion;

  • bind graph initializers;

  • silently drop unrecognized packed tensors.

Most adapters should become declarative rename/transform rules. Code hooks remain for QMoE and other transforms whose semantics are genuinely architecture-specific.

7. Generic binder and publication gate

The binder consumes normalized, adapted records and graph initializers.

Required invariants:

  • every required graph initializer is bound;

  • every packed checkpoint record is consumed;

  • no qweight sidecar remains after normalization;

  • bits and group size match the component graph;

  • scale and zero-point shapes match the canonical parameter schema;

  • a component declared quantized cannot bind float weights to packed parameters;

  • a component declared float cannot contain packed checkpoint weights;

  • unknown checkpoint records are rejected unless explicitly allowlisted.

These checks happen before save or runtime metadata generation.

The loading-time binding gate scans emitted MatMulNBits and GatherBlockQuantized nodes regardless of whether configuration came from a component mapping or a legacy root plan. build(load_weights=False) still returns an unbound skeleton. Independently, ModelPackage.save() checks all initializers by default.

Ordering of normalization and model transforms

Some transforms operate on logical tensors while others depend on a source packing layout. The target typed pipeline distinguishes:

  1. sidecar grouping;

  2. semantic routing/name mapping — typed records remain packed;

  3. architecture transform planning — split/fuse intent is declared;

  4. codec normalization — source layout becomes canonical target layout;

  5. planned transform execution;

  6. binding.

A transform unsupported for packed storage fails explicitly. It must not silently dequantize and requantize.

The current compatibility path performs raw semantic name mapping before grouping, as described above. It preserves ordinary sidecars rather than claiming that the full typed transform-planning pipeline is already implemented.

Compatibility plan

The migration keeps current APIs operational:

  • a global quantization_config maps to the legacy component;

  • existing preprocess_weights() methods run through a legacy adapter;

  • existing ComponentSpec, model_roles, and HF_COMPONENT_SOURCES feed the manifest resolver;

  • loaders without a component manifest retain current single-model behavior.

Local Transformers --config builds use the same shared build() entry point as --model. The CLI forwards loading and quantization policy options; the shared builder owns component adaptation, normalization, and format-specific loading, including native GPT-OSS MXFP4 streaming. These steps must not be duplicated in a local CLI loader.

Component-aware compatibility adapters preserve raw sidecars through their architecture-specific renames. They must not normalize an entire component with one root-resolved layout before the generic loader evaluates projection overrides.

Compatibility paths emit no behavior changes in the first two implementation PRs. Removal occurs only after model migrations and parity coverage.

Migration stack

PR 1: Component manifest

  • Add ComponentDescriptor, ComponentManifest, and the resolver.

  • Convert inspection and validation to consume the manifest.

  • Keep graph output and weight behavior unchanged.

PR 2: Typed weight pipeline

  • Add WeightRecord, WeightBundle, and codec interfaces.

  • Group packed sidecars and validate metadata.

  • Wrap existing preprocess functions behind compatibility codecs.

  • Keep existing loaders as the active path.

PR 3: Per-component loader

  • Parse the component quantization manifest.

  • Instantiate components with their effective config.

  • Route records through manifest, codec, adapter, and binder.

  • Add broad multi-component graph and synthetic packed-weight coverage.

PR 4: Model adapter migration

  • Migrate T5 and conventional VLMs first.

  • Migrate Gemma4 clipped vision/audio projections.

  • Migrate Qwen3.5/QMoE specialized transforms.

  • Remove post-construction module scanning and migrated legacy branches.

Testing strategy

Unit tests

  • manifest resolution from every declaration source;

  • sidecar grouping for Olive/GPTQ/AWQ;

  • partial sidecar rejection;

  • component routing with module path differing from package name;

  • exact and regex float exclusions mapped from HuggingFace names to local component module paths;

  • codec shape and packing validation;

  • binder completeness and leftover-record rejection.

Broad graph tests

For every registered multi-component configuration:

  • build with distinct component layouts;

  • assert eligible component projections use the expected bits/group size;

  • assert omitted components remain float;

  • assert single-component and text-only aliases preserve legacy behavior.

Weight-level tests

  • T5 encoder INT8 plus decoder INT4;

  • VLM decoder INT4 plus vision INT8 plus float embedding;

  • decoder INT4 with float per_layer_input_gate and per_layer_projection;

  • speech decoder INT4 plus audio INT8;

  • component path differing from package name;

  • tied embedding/head behavior;

  • QMoE component specialization.

Runtime evidence

At least one real-weight package must:

  • load every ONNX component directly with ONNX Runtime;

  • consume all packed checkpoint records;

  • run deterministic logits or component parity;

  • run multi-token generation through the declared runtime when applicable.

Acceptance criteria

The migration is complete when:

  • one manifest is the source of truth for every component consumer;

  • no generic loader uses model-name conditionals;

  • codecs contain no model-specific paths;

  • migrated adapters receive only one component bundle;

  • every packed record is either bound or rejected;

  • Mobius never converts float weights into quantized weights;

  • per-component real-weight parity is demonstrated;

  • legacy global-config tests remain unchanged.