feat(neuron): qwen3_next config + MoE FFN block, single-GPU path with exact HF parity (#92) #108

Merged

grenade merged 2 commits from feat/92-qwen3_5-moe into main 2026-07-01 21:23:24 +00:00

Conversation 0 Commits 2 Files Changed 17 +1481 -35

grenade commented 2026-07-01 21:23:15 +00:00

Owner

Copy Link

Part of #92 (F1, epic #84) — slices 1+2 of the MoE integration: everything except TP expert sharding (slice 3) and the fused CUDA grouped-GEMM (slice 4), which follow as separate PRs.

Slice 1 — config + normalization

• TextConfig gains the MoE hyperparameters (num_experts, num_experts_per_tok, moe_intermediate_size, shared_expert_intermediate_size, decoder_sparse_step, mlp_only_layers, norm_topk_prob), all defaulting to the dense case, plus layer_uses_moe() implementing the upstream sparse-layer selection.
• Config::from_config_json normalises the flat qwen3_next config layout (Qwen3-Next-80B-A3B family) into the nested qwen3_5 shape: flat rope fields → rope_parameters, attn_output_gate forced on (unconditional upstream), layer_types derived from full_attention_interval via (i+1) % interval.
• All four config parse sites (dense CPU, dense CUDA worker, TP leader, TP worker subprocess) route through the normalizer; qwen3_next added to the supported-type gates and dispatch arms.

Slice 2 — the MoE block, single-GPU

• Qwen3_5MoeBlock: top-k router (upstream softmax-then-topk order; renorm iff norm_topk_prob), per-expert SwiGLU at moe_intermediate_size, and the always-on shared expert mixed via sigmoid(shared_expert_gate(x)) — a path candle-transformers' qwen3_moe lacks entirely. Correctness-first scatter dispatch; grouped-GEMM lands in slice 4 behind the same forward signature.
• MlpKind { Dense, Moe } dispatch in the decoder layer, mirroring AttentionKind.
• Fused-checkpoint loading: qwen3_next stores in_proj_qkvz/in_proj_ba interleaved per key-head group (upstream fix_query_key_value_ordering layout); split_fused_qkvz/split_fused_ba de-interleave once at load into the contiguous layout the forward path already uses. Auto-detected via contains_tensor — Qwen3.6 checkpoints untouched.
• text_weight_prefix(): qwen3_next checkpoints keep the text core at model.* vs Qwen3.6's model.language_model.*.

Validation

• Exact HF parity: the new qwen3_next_parity test replays a committed tiny random-weight Qwen3NextForCausalLM checkpoint (generated by script/dump_qwen3_next_tiny.py on beast — transformers 5.9.0, torch 2.9.1) through neuron's full load path: max_abs 0.000000, cosine 1.00000000 at f32. Pins config normalisation, weight prefix, qkvz/ba de-interleave, hybrid layer interleaving, and the whole MoE block against upstream, in every CI run.
• Unit tests: scatter forward vs an independent per-token dense reference, no-shared-expert/no-renorm semantics, fused-split round-trip, flat-layout end-to-end structural load.

Inert for the served 27B (dense path unchanged, parity fixtures still pass); a TP guard fails any premature 80B TP load with a clear "#92 slice 3" message instead of a missing-tensor error.

🤖 Generated with Claude Code

https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

Part of #92 (F1, epic #84) — slices 1+2 of the MoE integration: everything except TP expert sharding (slice 3) and the fused CUDA grouped-GEMM (slice 4), which follow as separate PRs. ## Slice 1 — config + normalization - `TextConfig` gains the MoE hyperparameters (`num_experts`, `num_experts_per_tok`, `moe_intermediate_size`, `shared_expert_intermediate_size`, `decoder_sparse_step`, `mlp_only_layers`, `norm_topk_prob`), all defaulting to the dense case, plus `layer_uses_moe()` implementing the upstream sparse-layer selection. - `Config::from_config_json` normalises the flat `qwen3_next` config layout (Qwen3-Next-80B-A3B family) into the nested qwen3_5 shape: flat rope fields → `rope_parameters`, `attn_output_gate` forced on (unconditional upstream), `layer_types` derived from `full_attention_interval` via `(i+1) % interval`. - All four config parse sites (dense CPU, dense CUDA worker, TP leader, TP worker subprocess) route through the normalizer; `qwen3_next` added to the supported-type gates and dispatch arms. ## Slice 2 — the MoE block, single-GPU - `Qwen3_5MoeBlock`: top-k router (upstream softmax-then-topk order; renorm iff `norm_topk_prob`), per-expert SwiGLU at `moe_intermediate_size`, and the always-on **shared expert** mixed via `sigmoid(shared_expert_gate(x))` — a path candle-transformers' qwen3_moe lacks entirely. Correctness-first scatter dispatch; grouped-GEMM lands in slice 4 behind the same forward signature. - `MlpKind { Dense, Moe }` dispatch in the decoder layer, mirroring `AttentionKind`. - Fused-checkpoint loading: qwen3_next stores `in_proj_qkvz`/`in_proj_ba` interleaved **per key-head group** (upstream `fix_query_key_value_ordering` layout); `split_fused_qkvz`/`split_fused_ba` de-interleave once at load into the contiguous layout the forward path already uses. Auto-detected via `contains_tensor` — Qwen3.6 checkpoints untouched. - `text_weight_prefix()`: qwen3_next checkpoints keep the text core at `model.*` vs Qwen3.6's `model.language_model.*`. ## Validation - **Exact HF parity**: the new `qwen3_next_parity` test replays a committed tiny random-weight `Qwen3NextForCausalLM` checkpoint (generated by `script/dump_qwen3_next_tiny.py` on beast — transformers 5.9.0, torch 2.9.1) through neuron's full load path: **max_abs 0.000000, cosine 1.00000000** at f32. Pins config normalisation, weight prefix, qkvz/ba de-interleave, hybrid layer interleaving, and the whole MoE block against upstream, in every CI run. - Unit tests: scatter forward vs an independent per-token dense reference, no-shared-expert/no-renorm semantics, fused-split round-trip, flat-layout end-to-end structural load. Inert for the served 27B (dense path unchanged, parity fixtures still pass); a TP guard fails any premature 80B TP load with a clear "#92 slice 3" message instead of a missing-tensor error. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

grenade added 2 commits 2026-07-01 21:23:15 +00:00

feat(neuron): accept qwen3_next configs into the qwen3_5 arch (#92) ... a1426f177c

F1 slice 1 — config + normalization, no kernel work yet:

- TextConfig gains the MoE hyperparameters (num_experts,
  num_experts_per_tok, moe_intermediate_size,
  shared_expert_intermediate_size, decoder_sparse_step,
  mlp_only_layers, norm_topk_prob), all defaulting to the dense case
  so existing Qwen3.6 configs parse unchanged, plus layer_uses_moe()
  implementing the upstream sparse-layer selection.
- Config::from_config_json normalises the flat qwen3_next layout
  (Qwen3-Next-80B-A3B family) into the nested qwen3_5 shape: wraps
  flat hyperparameters, nests flat rope fields into rope_parameters,
  forces attn_output_gate (unconditional in upstream qwen3_next), and
  derives layer_types from full_attention_interval via the upstream
  (i+1) % interval convention when absent.
- All four config parse sites (dense CPU, dense CUDA worker, TP
  leader, TP worker subprocess) route through from_config_json;
  "qwen3_next" added to DENSE/TP_SUPPORTED_MODEL_TYPES and the
  dispatch match arms.
- Explicit num_experts > 0 guards in Qwen3_5Model::load and
  TpQwen3_5Model::load so a premature 80B load fails with a clear
  "MoE not implemented yet (#92)" instead of a cryptic missing-tensor
  error. Slices 2/3 remove them.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

feat(neuron): qwen3_next MoE FFN block, single-GPU path + HF parity (#92) ... 9bf13f09dd

F1 slice 2 — the MoE block itself, CPU/single-GPU:

- arch/qwen3_5/moe.rs: Qwen3_5MoeBlock — top-k router (upstream
  softmax-then-topk order, renorm iff norm_topk_prob), per-expert
  SwiGLU (reusing Qwen3_5MLP at moe_intermediate_size), and the
  always-on shared expert mixed via sigmoid(shared_expert_gate).
  Correctness-first host-side scatter dispatch; the fused grouped-GEMM
  path is slice 4 behind the same forward signature.
- decoder.rs: MlpKind { Dense, Moe } dispatch on layer_uses_moe,
  mirroring AttentionKind.
- linear_attn.rs: fused-checkpoint support — qwen3_next stores
  in_proj_qkvz / in_proj_ba interleaved per key-head group (upstream
  fix_query_key_value_ordering layout); split_fused_qkvz/ba
  de-interleave once at load into the contiguous [Q|K|V] + Z / B + A
  layout the forward path (incl. the conv channels) already uses.
  Auto-detected via contains_tensor, so Qwen3.6 checkpoints are
  untouched.
- mod.rs: text_weight_prefix() — qwen3_next checkpoints put the text
  core at `model.*`, Qwen3.6 at `model.language_model.*`; the slice-1
  single-GPU MoE guard is removed (TP guard stays until slice 3).

Validation:
- qwen3_next_parity integration test replays a committed tiny
  random-weight HF Qwen3NextForCausalLM checkpoint (generated by
  script/dump_qwen3_next_tiny.py on beast: transformers 5.9.0,
  torch 2.9.1) through neuron's full load path: max_abs 0.000000,
  cosine 1.00000000 at f32 — exact parity, pinning the config
  normalisation, weight prefix, qkvz/ba de-interleave, hybrid layer
  interleaving, and the whole MoE block against upstream.
- Unit tests: scatter forward vs per-token dense reference,
  no-shared-expert/no-renorm behaviour, fused-split round-trip, and a
  flat-layout end-to-end structural load.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

grenade merged commit 781d92d9a3 into main 2026-07-01 21:23:24 +00:00

grenade referenced this issue from a commit 2026-07-01 21:23:24 +00:00

Merge pull request 'feat(neuron): qwen3_next config + MoE FFN block, single-GPU path with exact HF parity (#92)' (#108) from feat/92-qwen3_5-moe into main

grenade referenced this pull request 2026-07-01 21:23:55 +00:00

F1 — neuron: MoE FFN into the qwen3_5 hybrid layer stack + TP expert sharding #92

grenade referenced this pull request 2026-07-01 21:54:03 +00:00

feat(neuron): TP expert sharding for the qwen3_next MoE block (#92) #109

Sign in to join this conversation.

Reviewers

No Reviewers

Labels

Clear labels

out-of-scope

Outside helexa's constraining principles: near-frontier models on consumer hardware, served transparently and predictably

p1-now

Active pain or unblocking work — pick up first

p2-next

Engine levers — start once P1 baselines exist

p3-later

Dependent on earlier work, or demand-driven

No Label

Milestone

No items

No Milestone

Projects

Clear projects

No project

Assignees

Clear assignees

grenade

No Assignees

1 Participants

Notifications

Subscribe

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: helexa/helexa#108