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

Abierto
  • ¿Algo no está bien?
#74 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 1 día

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
20/100
Tipo de issue
Nueva funcionalidad
Claridad
Necesita aclaración
Estado de actividad
Activo
Stack tecnológico
huggingface, python
Área
machine-learning

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Rust
Estrellas
800
Forks
106
Merge medio
1 d 3 h
PR fusionados (30 d)
100

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/hyperflare-inference

Todos los issues de NVIDIA/hyperflare-inference

Issues similares

Más issues de Rust

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.