feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend) - #385
feat(qwen4_exp): tensor parallelism for Qwen3.8-Flash-Next (offload backend)#385gdevenyi wants to merge 5 commits into
Conversation
|
|
Both from @gdevenyi's review on 2 x RTX 6000 Ada, where this is carried on a deploy branch alongside FlashML-org#385 (TP). 1. Attention routed its bf16 fallback through the quantized factories too, which swaps in their generic fallback and drops the tensor-parallel classes FlashML-org#385 needs (per-rank local_output_sizes, row-parallel o_proj). That is exactly the path a rank takes under TP>1, since the block-FP8 linears have no parallel variant. The factories are now used only on the fp8_block branch; every other case keeps LinearColParallelMerged / LinearReplicated as before. 2. config.parse_config and weight._dense_is_block_fp8 read the same declaration through two independent code paths, each with its own copy of _FP8_BLOCK_ALGOS. That is safe only while they cannot disagree, and they can: a rank downgrading under TP>1 must have the modules it BUILDS and the buffers it LOADS downgrade together, or the buffers will not match. Both now resolve through one helper, config.dense_quant_mode, which owns the declaration test and the TP downgrade. The duplicate constant is gone. It reads TP through try_get_tp_info, not get_tp_info: Engine.__init__ sets TP info as its first statement so a rank always knows its size by the time this matters, but config parsing also happens with no engine at all (checkpoint conversion, tooling, tests) where get_tp_info raises. Verified on the modelopt checkpoint: parse_config still yields nvfp4/fp8_block; the two sides agree at TP=1 (both fp8_block) and at TP=2 (both downgraded); attention builds LinearColParallelMerged/LinearReplicated under bf16 and Fp8BlockColMerged/Fp8BlockLinear under fp8_block. tests/models/qwen4_exp/test_config.py + test_weight.py: 30 passed. The whole qwen4_exp suite reports 47 failed / 46 passed / 50 skipped both at the merge-base and with these fixes - identical sets, no regressions. Those failures are pre-existing and are an artefact of this box rather than the code: its single 24 GB card is 23.6 GB occupied serving a model, so the GPU-dependent tests cannot allocate. I have not been able to run them on a free card.
Follow-up to @gdevenyi's note on the previous fix. The bf16 branch kept o_proj as LinearReplicated, which is what main does today and is correct at TP=1, but it is the path a rank falls back to under TP>1 (FlashML-org#385), and there a replicated o_proj is wrong three ways at once: qkv_proj is column-parallel so each rank's attention output is its local head slice, o_proj therefore needs the sharded input dim, and the partial sums need an all-reduce. It also fails quietly, since a missing reduction still decodes to fluent-looking text. LinearOProj does all of that and degenerates to LinearReplicated at TP=1: div_even(x, 1) == x, and the all-reduce is skipped when tp_size == 1. So this is a no-op for main and only changes what FlashML-org#385 finds when the two meet, whichever lands second. It does mean get_tp_info() runs in __init__, but the same branch already does that two lines up through LinearColParallelMerged, so there is no new constraint: this path was engine-only before and still is. Config parsing, the one no-engine path that mattered, stays on try_get_tp_info. The comment above the branch also claimed a row-parallel o_proj that the code did not build; it now describes what is built. Verified on ailab1 (single RTX PRO 4000 Blackwell, TP=1): - tests/models/qwen4_exp/{test_config,test_weight,test_skeleton, test_qsa_backend}.py: 49 passed at 72773b0 and 49 passed with this change. That set includes test_qsa_layer_matches_hf_dense, which builds Qwen4ExpAttention on the card and checks the whole layer against the HF dense reference, so the new o_proj is exercised through a real forward. - A direct check on CPU and CUDA: the built o_proj is a LinearOProj with weight [hidden, qo_attn_dim] and local_input_size == qo_attn_dim, and its forward is bit-identical (max |diff| = 0.0) to a LinearReplicated carrying the same weight. Correction to the previous message: the 47 failures I attributed to VRAM contention were, for these four files at least, a missing ninja on PATH in my throwaway test venv. With the serving venv's bin on PATH the files pass in full.
qwen4_exp is the only family that builds its attention o_proj as LinearReplicated; llama, gpt_oss and minimax_m2 all use LinearOProj. That is correct at TP=1 and wrong under TP>1 three ways at once: qkv_proj is column-parallel, so a rank's attention output is its local head slice rather than the full qo_attn_dim, o_proj therefore has to take the sharded input dim, and the partial sums need an all-reduce. LinearReplicated keeps the full [hidden, qo_attn_dim] weight, expects the unsharded input and reduces nothing. It also fails quietly: a missing all-reduce leaves each rank holding a partial sum that still decodes to fluent-looking text. LinearOProj degenerates to exactly LinearReplicated at TP=1 -- div_even(x, 1) == x, and the all-reduce is skipped when tp_size == 1 -- so this is a no-op for main as it stands and only changes what FlashML-org#385 finds when the two meet. It adds no constraint from calling get_tp_info() in __init__ either, since the same constructor already reaches it two lines up through LinearColParallelMerged. The comment above the branch now describes what is built. Raised by @gdevenyi against the earlier form of this work in FlashML-org#392.
…ackend) Shard the dense weights per rank at load (attention qkv by head, GDN in_proj as its six parts with the matching conv1d channels and A_log/dt_bias, shared-expert gate_up per part; o_proj/out_proj/down_proj row-parallel; embed/lm_head by vocab rows) and the NVFP4 expert banks along the intermediate axis, so every rank holds half the experts and each MoE layer needs one all-reduce (routed + gate * shared are combined before the reduce). Router, QSA indexer, norms, hyper-connections and PLE stay replicated so all ranks select the same blocks and n-gram rows. Also: LinearColParallelMerged(local_output_sizes=) for the kv-replicated case and distributed_timeout 60 -> 1800 s (ranks reach their first collective minutes apart behind a 100+ GiB load). Limits: offload backend with bf16 dense projections; fp8_block / nvfp4 dense checkpoints raise under TP. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
tests/models/qwen4_exp/test_weight.py feeds iter_weights a synthetic checkpoint whose config.json has no model_type; at TP=1 nothing is sharded, so do not touch the config. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
d16d634 to
a99433c
Compare
|
Rebased onto What changed and why. The original sharded inside So the sharding moved to the two places the new architecture actually owns:
Two other resolutions worth flagging:
Testing. Full 🤖 Generated with Claude Code |
|
Re-tested against current main ( Method. This PR's head merged onto main, then the full Result: 1214 passed, 350 skipped, no new failures. The +9 over main are this PR's own tests, and they ran (not skipped). 🤖 Generated with Claude Code |
…build
Qwen4ExpDecoderLayer builds its MoE as `Qwen4ExpMoE(config, layer_id, prefix=...)`,
but this PR's override of __init__ (added to hold the TP communicator) took only
(config, layer_id), so a server boot died with
TypeError: Qwen4ExpMoE.__init__() got an unexpected keyword argument 'prefix'
The whole CPU test suite was green with that bug in place, because every test that
builds a decoder layer is behind requires_cuda -- nothing without a GPU ever
constructed the model. tests/models/qwen4_exp/test_build_cpu.py closes that: it
builds the full model on the meta device (no GPU, no memory) and asserts the state
dict has both layer families, an lm_head, and MoE weights on more than one layer,
so a dropped or shared prefix fails too.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…el shard
With the expert piece stream sliced per rank and the banks sized from
MoEConfig.local_intermediate, a rank holds exactly its half of every expert, so this
kernel can serve TP>1 -- the routed output is a partial sum and the MoE layer already
reduces it (_maybe_all_reduce, or the single combined all-reduce in qwen4_exp's block).
Without this the whole selection table is empty under TP=2 on sm_89 and the server
refuses to start:
KernelSelectionError: no usable kernel in table;
triton: TP > 1 is not supported for this expert format;
marlin: vLLM is not installed;
b12x: b12x requires sm_120+, got sm_89
marlin and b12x keep tp_ok=False deliberately: their pack() repacks the native rows and
neither has been verified against a per-rank bank.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
|
Pushed Integrating this branch with the NVFP4 expert path after the #418/#427 quantization refactor, the engine refused to start: Every NVFP4 MoE kernel carried Worth a second pair of eyes on one thing: this made the bank-sizing bug below silent rather than fatal, so I want the sizing assertion reviewed as much as the flag. Related, and the more interesting failure. A merge of upstream main into the deployment line silently dropped this same Also on this branch since the last comment: 🤖 Generated with Claude Code |
|
End-to-end validation of this PR at TP=2, as a controlled A/B against the branch it replaces. Our deployment branch (upstream
Throughput and residency are identical. TTFT is ~7% better on the mean, and more interestingly it is flat at 0.68 from the first measurement where the old branch drifts 0.79 → 0.70 as its expert cache settles. Accuracy is one question higher with two fewer genuine misses — both inside noise at n=300, but nothing regressed. Two methodology notes, because they cost me real time and may save yours: Do not benchmark this box inside ~10 minutes of start. The aggregate reads 2–3% low while the expert cache fills. My first run on the new branch showed 343.7 against a 352.3 baseline and I nearly reported a 2.4% regression that did not exist; settled, it is 354–357. GSM8K at the default 768-token cap measures truncation as much as reasoning. At 768 the two branches score 95.67% and 96.33%, which looks like a regression; the difference is entirely that one truncated 10 answers and the other 6. At 1536 the ordering reverses. Per-question flip analysis: 9 of 10 truncated answers become correct when the cap is lifted. Since the last comment this branch also carries 🤖 Generated with Claude Code |
1 participant
What this adds
ft serve --tp-size 2forqwen4_exp(Qwen3.8-Flash-Next) on the offload MoE backend. Upstream refuses TP>1 for this architecture; with two 48 GiB cards the model runs as two independent TP=1 instances, each keeping 37% of the NVFP4 experts resident and streaming the rest over PCIe every step.Per rank, the patch:
weight.py::_shard): attentionqkv_projby head ([q|gate]per head; kv heads split, or replicated when there are fewer than ranks), GDNin_projas its six parts ([q | k | v | z | b | a]) plus the matchingconv1dchannels andA_log/dt_bias, shared-expertgate_up_projper part;o_proj,out_projand the shareddown_projrow-parallel with the all-reduce inside;embed_tokens/lm_headby vocab rows. Router, QSA indexer, norms, hyper-connections and PLE stay replicated, so every rank selects the same sparse blocks and n-gram rows.nvfp4_banks.py, I=640 -> 320 per rank): packed codes, the 16-wide scale blocks and the per-row globals for the gate/up rows and the down columns. The offload cache then holds half the experts per rank.moe.py): the routed and shared partial sums are combined asrouted + sigmoid(gate) * sharedbefore a single reduce, instead of one reduce each.LinearColParallelMerged(local_output_sizes=)for the kv-replicated case (the same shape as feat(models): support TP for qwen3_5_moe #104's hunk), anddistributed_timeout60 s -> 1800 s: behind a 100+ GiB load the ranks reach their first collective minutes apart, and 60 s kills the launch.Everything else in the engine (scheduler, KV pool, PLE table, CUDA graphs) is untouched.
Measurements
2 x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4 x16, no NVLink), 2 x Xeon Gold 6526Y, 503 GiB RAM.
RadixArk/Qwen3.8-Flash-Next-NVFP4,--moe-backend offload --ple-backend pinned --num-tokens 262144 --memory-ratio 0.94 --moe-prefill-hit-d2d.A 262,144-token prompt reaches its first token in 74 s at TP=2 (116 s at TP=1). The decode step is dominated by the bf16 dense read; TP=2 halves it per GPU and removes the PCIe expert gather because everything fits.
Correctness. An 8-question probe gives identical answers at TP=1 and TP=2. Three raw prompts decoded greedily for 256 tokens: the ~1k-token prompt (QSA over many blocks, GDN state, PLE context) is word-for-word identical across every run; the two short prompts diverge after 13 and 29 words between TP=1 and TP=2, but the TP=2 server diverges from itself at the same points on a second pass (bf16 atomics in the expert kernels), so that is run-to-run noise rather than a sharding error. Greedy output at TP=2 is not bit-exact between passes.
Limits
Offload backend with bf16 dense projections only:
fp8_block/nvfp4dense checkpoints raise under TP (row-parallel FP8 / NVFP4 linears do not exist yet, the same gap #104 has). The hybrid / CPU MoE backends are not sharded.Related: #62, #29 (TP for offloaded MoE), #104 (TP for qwen3_5_moe, which this reuses the merged-linear hunk from).
Testing
tests/models/qwen4_exp/test_tp_shard.py: the per-head / per-part row sharding of every fused projection reassembles to the original (CPU).tests/models/test_nvfp4_banks_tp.py: the bank placer's per-rank slices of codes, scales and globals cover the intermediate axis exactly once (CPU).tests/models/qwen4_expon one of its GPUs: 97 passed, 3 failed; the same 3 (test_chunked_prefill_matches_one_shot[*], a bit-exact assertion off by bf16 noise on this torch 2.11 / flashinfer 0.6.18 / triton 3.6 stack) fail on plainmainthere too.🤖 Generated with Claude Code
https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt