Skip to content

ModelOpt checkpoints: the weight source stops at the projection; detection, shards and config are still on every consumer #3866

Open

Description

@astorise

Summary

#3846 asked for a ModelOpt weight source. What landed covers the projection: given a VarBuilder already positioned at a linear layer, quantized_nvfp4::auto_linear reads the packed weight, the block scales and the tensor scale, and picks the operator. That is the hard half, and it works.

The other half is still on every consumer: deciding that a directory is a ModelOpt NVFP4 checkpoint, and reading the config it ships.

Concretely, before auto_linear can be called even once, a caller must:

  • recognise NVFP4 from W4A16_NVFP4 metadata in config.json or hf_quant_config.json, or from ModelOpt tensor names such as .weight_scale_2;
  • read model.safetensors.index.json, resolve the shard set, and hand the right paths to VarBuilder::from_mmaped_safetensors;
  • deserialize the checkpoint's own config.json into the model's Config.

None of that is application logic, and all of it is identical for every consumer.

The sharpest symptom: two config types for one model

qwen3_5::Config and TextConfig both derive serde::Deserialize, so they look ready to read a checkpoint. They are not, for a ModelOpt export: the nesting and several field names differ from what the file carries. So a downstream ends up with a second config type that deserializes the real file, plus a translation function into candle's.

That is what we have. The translation is 35 lines of field-by-field copying whose only reason to exist is that the two shapes disagree:

pub(crate) fn upstream_config(config: &Qwen35MoeConfig) -> UpstreamConfig {
    UpstreamConfig { text_config: TextConfig {
        vocab_size: config.vocab_size,
        head_dim: Some(config.head_dim),
        hidden_act: candle_nn::Activation::Silu,   // not in the file
        attention_bias: false,                     // not in the file
        // … 20 more
    }}
}

Two of those fields are not in the checkpoint at all and have to be pinned by the caller from knowledge of the architecture — which is precisely the knowledge candle has and the caller should not need.

Any second consumer will write this same struct and this same function, and the two copies will disagree the first time a field is added upstream.

What this asks for

A directory-level entry point beside the projection-level one. Roughly:

// Detects the format and reads what the directory states about itself.
let ckpt = ModelOptCheckpoint::open(path)?;          // None if not ModelOpt/NVFP4
let cfg: qwen3_5::Config = ckpt.model_config()?;     // the file's own config.json
let vb = ckpt.var_builder(dtype, &device)?;          // index + shards resolved

let model = qwen3_5::ModelForCausalLM::new_with_projections(
    &cfg, vb,
    |name, vb, in_dim, out_dim, bias| {
        Ok(Box::new(auto_linear(in_dim, out_dim, bias, vb.pp(name), cfg)?))
    },
)?;

The shape matters more than the names. Three things are being asked for, and they are separable:

  1. Detection — is this directory a ModelOpt NVFP4 checkpoint? Today every consumer re-derives it from metadata keys and tensor-name suffixes.
  2. Shard resolution — index to shard paths. Mechanical, and identical everywhere.
  3. Config deserialization — the checkpoint's config.json into the model's Config, including the fields the file does not carry.

(3) is the one that removes a whole parallel type from downstreams. (1) and (2) are smaller but are pure duplication.

Why this is candle's and not the application's

The line we have been drawing: candle owns the model, the kernels and the checkpoint format; the application owns policy. Detection, shard resolution and config parsing are all format facts. What legitimately stays with a consumer is which producer versions it has qualified, what it does when a checkpoint fails validation, and where it will run — none of which candle should decide.

Our own modelopt_nvfp4.rs is ~2.6k lines. Most of it is (1), (2) and the tensor grouping that auto_linear already made redundant; the part that is genuinely ours — a fail-closed compatibility profile that refuses a producer version we have not qualified — is a few hundred lines and belongs where it is.

Notes from having written it downstream

  • Detection has to accept both spellings of the packed weight (weight for ModelOpt, weight_packed for compressed-tensors) — the same distinction ModelOpt/NVFP4: nothing connects candle-nvfp4-kernels to the models that can use them #3846 settled at the projection level applies to the directory-level probe.
  • hf_quant_config.json is not always present; some exports carry the quantization block inside config.json. Both are real.
  • A checkpoint is mixed: dense BF16, FP8 and W4A16 NVFP4 tensors in one graph. Detection must not require that every linear be NVFP4.
  • Sharded and single-file checkpoints both occur; model.safetensors.index.json is absent in the second case.
  • Fields like hidden_act and attention_bias are architecture facts rather than file contents. If Config deserialization owns them via #[serde(default = …)], no consumer has to guess them.

Context

Follows #3846 (projection-level weight source) and #3857 (packed CPU operator). Downstream in Tachyon-Mesh we have just deleted a local reimplementation of Qwen 3.5 in favour of candle's; what remains model-specific on our side is almost entirely this loading gap and the policy around it.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions