[BUG] shard_tensor() chunks 0-dim ModelOpt scales; kv-cache scales break the qwen3_moe qkv merge #496

Before you start

• I have read the FAQ and my problem is not answered there.
• I have read the Roadmap and this is not already planned there.
• I have searched existing issues and found no duplicate.
• I am on the latest release, or on a freshly rebuilt main when building from source.

What happened

Up front on the duplicate check, because it's not a clean "no": #425 already shows Bug 1's traceback — same chunk expects at least a 1-dimensional tensor. I'm not claiming that crash is news.

What I think earns a separate issue is that the crash isn't really about qwen3_moe's missing NVFP4 expert wiring. shard_tensor() will happily chunk any 0-dim tensor whose key substring-matches, and some of the ones ModelOpt ships aren't expert tensors at all — so they'd still crash after the fix #425 is asking for. And there's a second bug sitting right behind it that nobody's mentioned yet.

Both are two-liners, both turned up loading nvidia/Qwen3-235B-A22B-Instruct-2507-NVFP4, both were run on real hardware. Happy to send them as PRs (one each).

Line numbers below are main @ cac247a; the run itself was on the 0.1.2 wheel, where these two code paths are identical.

Bug 1 — shard_tensor() chunks 0-dim quantization scales

python/freetoken/models/loader.py routes by substring:

SPLIT_DIM_0 = (".q_proj", ".k_proj", ".v_proj", ".gate_proj", ".up_proj")
SPLIT_DIM_1 = (".o_proj", ".down_proj")
...
        return value.chunk(world_size, dim=0)[rank].clone()   # L39
    if any(key.count(sub) for sub in SPLIT_DIM_1):
        return value.chunk(world_size, dim=1)[rank].clone()   # L41

ModelOpt exports ship per-tensor scalar companions whose keys contain those exact substrings — ...gate_proj.input_scale, ...gate_proj.weight_scale_2, ...o_proj.input_scale, ...o_proj.weight_scale_2. They're 0-dim, and chunk() doesn't care about world_size, it rejects them outright:

RuntimeError: chunk expects at least a 1-dimensional tensor
  File ".../freetoken/models/loader.py", in shard_tensor
    return value.chunk(world_size, dim=0)[rank].clone()

Loading dies at 0%. Nothing completes.

Why this is worth fixing on its own, separate from #425. #425's ask is to route routed-experts to load_moe_expert_sources the way qwen4_exp / glm5_next do, which would keep the expert scalars away from shard_tensor. That's the right fix for that issue and I'm not proposing an alternative. But the attention-side scalars — self_attn.o_proj.input_scale and .weight_scale_2 — are not expert tensors. _EXPERT_PATTERN doesn't match them, so they go down the include_non_moe path into shard_tensor() regardless, hit the .o_proj arm of SPLIT_DIM_1, and crash exactly the same way. Wiring up expert sources doesn't touch them.

The TP>1 path has a second flavour of the same thing: a 0-dim k_scale/v_scale reaching the num_kv_heads < world_size branch does value.shape[0] and gets an IndexError instead.

Fix. A per-tensor scale describes the whole tensor — there's nothing to shard, every rank wants the same number. A 0-dim tensor belongs on the function's existing "return it unchanged" path (L48), so guard at the top:

if value.dim() == 0:
    return value

No-op for the TP=1 path Qwen3-MoE is gated to today (#29), correct for TP>1. With it, all 28 shards load.

Bug 2 — kv-cache scales can never complete a qkv merge group

This one I haven't seen reported. With Bug 1 patched, weights load fully and the merge stage asserts:

AssertionError: qwen3_moe: Incomplete merge groups in checkpoint:
['model.layers.17.self_attn.qkv_proj.k_scale', 'model.layers.17.self_attn.qkv_proj.v_scale',
 'model.layers.18.self_attn.qkv_proj.k_scale', ...]

2 entries per layer, all 94 layers. _MERGE_RULES maps q/k/v → .qkv_proj and iter_merged_tensors wants all three slots filled for every suffix. But ModelOpt's kv-cache quantization scales are inherently single-sided — only k_proj carries k_scale, only v_proj carries v_scale ("kv_cache_quant_algo": "FP8" in hf_quant_config.json). There is no q_proj.k_scale, and there never will be. Those groups can't complete by construction.

Fix. The qwen3_moe engine path has no quantized-KV-cache consumer, so dropping the keys during weight iteration is correct — the KV cache just runs at native precision:

for raw_name in f.keys():
    name = raw_name.removeprefix("language_model.")
    if name.endswith((".k_scale", ".v_scale")):
        continue

iter_weights_parallel needs no equivalent change — it filters to _is_expert first, so attention tensors never reach it.

Where it ends up

With both patched, the merge stage passes and loading gets as far as load_state_dict, where this checkpoint runs into the structural gap already tracked in #425 — qwen3_moe declares every param as full-width BF16, and this export is mixed precision. I've dropped the layer-0 layout for the 235B export in a comment over there.

Not a duplicate of #20

#20 has the same outer Backend worker is gone supervisor message, but on an unquantized Qwen3.6-35B-A3B — no scale tensors in play at all, and the thread landed on a prefill-workspace OOM. Same wrapper, different root cause.

How did you install FreeToken

pip / uv wheel

FreeToken version

0.1.2 (PyPI, freetoken[accel]==0.1.2)

OS

Ubuntu

OS details

Linux x86_64, Python 3.12, CUDA toolkit 13.0, torch 2.11.0+cu130

GPU and driver

RTX PRO 6000 Blackwell Max-Q, 96 GB, driver 595.58.03

CPU and system RAM

Dual-socket EPYC, 1.1 TiB RAM

Checkpoint

nvidia/Qwen3-235B-A22B-Instruct-2507-NVFP4

Command

ft serve --model <local dir>

Full log

See `What happened` part

Anything else

ft bench bw profile present; MoE backend auto-selected hybrid. Can pull full logs if you want them.