qwen38: support compressed-tensors NVFP4 expert-only checkpoints such as ThinkingCap-Qwen3.8-Flash-Next-NVFP4

Aperta
  • Qualcosa non va?
#74 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 1 giorno

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
20/100
Tipo di issue
Funzionalità
Chiarezza
Da chiarire
Stato di attività
Attiva
Stack tecnologico
huggingface, python
Ambito
machine-learning

Direzione di ricerca

Read AGENTS.md for the maintainer-decision rule, then trace checkpoint validation in ignition/preparation/quantization.py and ignition/preparation/checkpoint.py. Study gemma4/config.py for compressed-tensors parsing and qwen38/weights.py for expert loading. Done means qwen38 serves the NVFP4 expert-only checkpoint after the three open questions are resolved.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

model:qwen3.8-flash-next triage:design-review

Would serving compressed-tensors quantizations of Qwen3.8 Flash Next fine-tunes be in scope for the qwen38 recipe? If it is, would you accept a PR for the checkpoint-loading changes below? I'm asking first because this is a feature request, and AGENTS.md asks for a maintainer decision before that kind of work.

The checkpoint is bottlecapai/ThinkingCap-Qwen3.8-Flash-Next-NVFP4@df3922b, an NVFP4 quantization of the ThinkingCap fine-tune of Qwen3.8 Flash Next. Its architecture, layer types and every TEXT_GEOMETRY and VISION_GEOMETRY field except ple_embedding_dtype match the recipe. Its tokenizer, tokenizer config, chat template and preprocessor config match FRONTEND_SHA256 byte for byte. Only the storage differs. The quantization_config is compressed-tensors nvfp4-pack-quantized with one weight-only group that targets the routed and MTP experts (re:.*mlp\.experts\.\d+\.(gate_proj|up_proj|down_proj)$) and an empty ignore list. Every other projection and norm is stored in BF16, and the PLE table is FP8 E4M3 with one per-tensor scale.

What fails today

I ran Flare's CPU-only checkpoint resolution and inspect_qwen38_checkpoint from main 9a45b0b against the checkpoint's real config, index and shard headers. After each failure I patched the metadata to reach the next check:

  1. Resolution stops in _stored_bytes with UsageError: ... has no readable model.safetensors.index.json total_size, because the index has "metadata": {}. The directory and repository-id paths of resolve() both call it.
  2. CheckpointError: text_config.ple_embedding_dtype: expected 'nvfp4', got 'float8_e4m3fn'
  3. CheckpointError: expected modelopt MIXED_PRECISION with unquantized KV, from quantized_layers. test_quantization_config_requires_supported_schema pins this rejection.
  4. With a hand-written ModelOpt declaration, CheckpointError: index total_size does not match tensor payload bytes: 131538245626 (checkpoint.py). This is a second check that needs total_size.
  5. With total_size added, CheckpointError: model.language_model.hyper_connection_mixer.hc_norm.weight: expected F32 (10240,), got BF16 (10240,)

With the decoder experts declared W4A16_NVFP4, which expert_format already accepts as mode a16, 76,102 of the checkpoint's 227,353 tensors match the qwen38 tensor contract. The rest differ as follows:

Tensors Checkpoint Contract
150,528 Decoder and MTP expert weight_packed and weight_global_scale weight and weight_scale_2
438 BF16 attention, GDN, shared-expert and vision-attention projections MXFP8 with weight_scale
56 BF16 vision MLP and merger projections W4A16 NVFP4
100 BF16 hyper-connection and PLE norms FP32
129 FP8 PLE shards [2500012, 160] and one BF16 weight_scale NVFP4 shards [2500012, 80] with block and global scales

The expert block scales already match the contract (FP8 E4M3, group 16). The global scales use the compressed-tensors reciprocal form. For example, layer 0 expert 0 gate_proj stores 15,123.7, which is 448 × 6 / amax for an amax of 0.178.

What a PR could change

Most of the loading code already exists:

  • gemma4 parses compressed-tensors config_groups and ignore, including nvfp4-pack-quantized, and leaves untargeted layers unquantized (config.py).
  • DirectWeightLoader.load_fp4_experts takes weight_name, global_scale_name and global_scale_kind="reciprocal" (weights.py). minimax_m3 already calls it with the compressed-tensors names (weights.py).
  • QwenWeightLoader.load_experts already sends non-CSF experts to load_fp4_experts with their declared activation mode (weights.py).
  • DirectWeightLoader.linear_format accepts undeclared plain BF16 projections (weights.py), and PreparedSharedExpert accepts stored BF16 (moe.py).

The change I have in mind stays within checkpoint loading and the storage contract:

  1. Map the compressed-tensors config onto the quantized_layers form that _validate_config checks, with untargeted projections treated as stored BF16. Pass the expert tensor names and the reciprocal scale kind through QwenWeightLoader.load_experts.
  2. Accept BF16 hyper-connection and PLE norms and upcast them to FP32 at load. The conversion is exact, and ple.py keeps its FP32 requirement.
  3. Accept an index without total_size in both checks, using the payload size from the shard headers that Flare already reads.
  4. Optionally let qwen38 declare online_mxfp8 (planning.py), as kimi3 does, so --online-mxfp8 can store the BF16 dense projections as MXFP8 at load. Without the flag, BF16 weights stay BF16.
  5. Correct the recipe catalog docstring, "A recipe serves any checkpoint of its architecture" (recipes/init.py), and document which checkpoint storage formats qwen38 accepts.

Three questions would decide the shape:

  • Concurrent attention projections require MXFP8 or MXFP4 weights (attention_overlap.py). With BF16 attention, should export require --no-stream-overlap or --online-mxfp8, or reject the checkpoint? I'm not proposing changes to the overlap path.
  • Is the norm upcast acceptable, or should the contract accept BF16 norms directly?
  • Should total_size stay required instead, with the docs saying so?

Serving ThinkingCap would also need FP8 per-tensor PLE tables. That touches the qwen38 PLE recipe, EngramTable.Encoding and Helios' two-plane table sources, although b12x already implements fp8_e4m3_per_tensor. If this issue is in scope, I'd raise that as a separate design question.

Evidence and limits

Every result above comes from Flare's own CPU-only code. The shard files were sparse files holding the real headers, so no tensor payload was read. I have not run ignition export, any GPU step, the qwen38 graphs with BF16 dense projections, or a16 decoder experts on SM121, so I make no performance or accuracy claims. The script below reproduces steps 1 to 5 without a GPU or PyTorch.

Repro script 
# Save in a Flare checkout and run: python3 repro_thinkingcap.py /tmp/thinkingcap-meta
# Needs Python 3.11+ and network access; no GPU, PyTorch or tensor payloads.
import json, struct, sys, urllib.request
from pathlib import Path

from ignition.errors import FlareError
from ignition.preparation import checkpoint, qwen38
from ignition.recipes import resolve

REPO = "bottlecapai/ThinkingCap-Qwen3.8-Flash-Next-NVFP4"
REV = "df3922b21fc28a8604cc9b31451495690d155194"
out = Path(sys.argv[1])
out.mkdir(parents=True, exist_ok=True)


def fetch(name, first=None, last=None):
    request = urllib.request.Request(f"https://huggingface.co/{REPO}/resolve/{REV}/{name}")
    if first is not None:
        request.add_header("Range", f"bytes={first}-{last}")
    return urllib.request.urlopen(request).read()


for name in ("config.json", "model.safetensors.index.json", "generation_config.json", "tokenizer.json",
             "tokenizer_config.json", "chat_template.jinja", "preprocessor_config.json"):
    (out / name).write_bytes(fetch(name))
index = json.loads((out / "model.safetensors.index.json").read_text())
payload = 0
# Each shard keeps its real header and file size; the payload stays a sparse hole.
for shard in sorted(set(index["weight_map"].values())):
    length = struct.unpack("<Q", fetch(shard, 0, 7))[0]
    header = fetch(shard, 8, 7 + length)
    offsets = [v["data_offsets"] for k, v in json.loads(header).items() if k != "__metadata__"]
    end = max(last for _, last in offsets)
    payload += sum(last - first for first, last in offsets)
    with open(out / shard, "wb") as file:
        file.write(struct.pack("<Q", length) + header)
        file.truncate(8 + length + end)


def attempt(step, call):
    try:
        call()
        print(f"{step}: passed")
    except (FlareError, ValueError) as error:
        print(f"{step}: {type(error).__name__}: {error}")


def rewrite(name, change):
    path = out / name
    data = json.loads(path.read_text())
    change(data)
    path.write_text(json.dumps(data))


inspect = lambda: checkpoint.inspect_qwen38_checkpoint(out)
attempt("1 resolve", lambda: resolve(str(out)))
attempt("2 inspect as published", inspect)
rewrite("config.json", lambda c: c["text_config"].update(ple_embedding_dtype="nvfp4"))
attempt("3 with ple_embedding_dtype=nvfp4", inspect)
layers = qwen38.expected_quantization()
for name in layers:
    if name.startswith(qwen38.TEXT_PREFIX) and name.endswith(".mlp.experts"):
        layers[name] = {"quant_algo": "W4A16_NVFP4", "group_size": 16}
rewrite("config.json", lambda c: c.update(quantization_config={
    "quant_method": "modelopt", "quant_algo": "MIXED_PRECISION", "kv_cache_quant_algo": None,
    "quantized_layers": layers}))
attempt("4 with a ModelOpt declaration (W4A16 decoder experts)", inspect)
rewrite("model.safetensors.index.json", lambda i: i.update(metadata={"total_size": payload}))
attempt("5 with index total_size from the shard headers", inspect)

I investigated this with Claude Code (Claude Opus 5.5), which ran the checks and drafted this issue. I reviewed it before posting. Following AGENTS.md, it would not write kernel, fusion or stream-overlap code for a PR.

Lingua principale
Rust
Stelle
800
Fork
106
Merge medio
1g 3h
PR unite (30g)
100

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/hyperflare-inference

Tutte le issue di NVIDIA/hyperflare-inference

Issue simili

Altre issue su Rust

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.