feat(models): complete TP for qwen3_5_moe quantized checkpoints (FP8/NVFP4) - #1
Open
CraigStone-Dev wants to merge 9 commits into
Open
Conversation
…NVFP4) Finishes the work started in this branch (which only covered BF16): - Row-parallel quantized linears: Fp8PerTensorRowParallel, Fp8BlockRowParallel, Nvfp4DenseRowParallel (sharded-K GEMM + all-reduce, bias after reduce). - Column-merged quantized linears accept local_output_sizes (GQA KV-head replication), removing the make_col_merged_quant guard. - gdn.py: in_proj_qkvz now declares TP-local sizes (was full-size while forward split by local sizes). - weight.py: _iter_weights_attn_fp8 (mixed modelopt checkpoints) ungated and TP-sharded: fused fp8 qkv/in_proj_qkvz per sub-part with KV replication, standalone fp8 o_proj/out_proj row-shard, in_proj_ba per-head, conv1d / A_log / dt_bias / embed_tokens sharded; NVFP4 dense (shared_expert, lm_head) stays replicated. _iter_weights_fp8 (block-fp8) and the expert bank paths ungated with the same sharding. - tests: CPU unit tests for the shard math and row-parallel numerics (pure-torch reference path).
Two small dependency-free clients for the E2E check used to validate the quantized TP port: ab_suite.py (10 short deterministic prompts, JSON out, --diff mode) and ab_long.py (5 long-form prompts, full reasoning+answer capture, --diff with first-divergence reporting). Used to compare TP=1 vs TP=2 outputs of nvidia/Qwen3.6-35B-A3B-NVFP4 on 2x RTX PRO 4000.
…TP>1 The block-fp8 pass (_iter_weights_fp8) sharded embed_tokens but not lm_head, so at TP>1 the full-width (vocab, hidden) tensor was yielded against the vocab-sharded ParallelLMHead buffer and load_state_dict failed with a shape assert (reported on PR FlashML-org#104: 124160 vs 248320 rows at TP=2). Extract the per-tensor sharding into _fp8_shard_standalone (behavior unchanged for o_proj/out_proj, embed_tokens, A_log/dt_bias, conv1d) and add lm_head.weight to the vocab-parallel dim-0 branch. Nvfp4LMHead (mixed checkpoint) stays replicated, as before. CPU test: test_fp8_shard_standalone_vocab_row_conv_tp2.
make_tiny_fp8_ckpt() shrinks the Qwen3.6-35B-A3B-FP8 geometry (4 layers: 3 GDN + 1 full-attn, few experts) to a few hundred MB while keeping hidden 2048, all per-layer dims, and by default the full 248320 vocab -- so lm_head / embed_tokens exercise the exact vocab-parallel shapes the TP=2 loader bug hit. Random weights: for load/boot/serve validation, not coherent generation. test_fp8_iter_weights_tp2_sharding runs the full _iter_weights_fp8 at TP=2 over the synthetic checkpoint and asserts the per-rank shapes (lm_head / embed_tokens vocab-sharded, o_proj / out_proj input-sharded, qkv / in_proj_qkvz per-part local rows, conv1d sub-parts, replicated router / shared expert). Fails on the pre-fix loader (lm_head full-width), passes with the lm_head vocab-parallel shard fix. Also validated E2E: with the full-vocab checkpoint, ft serve --tensor-parallel-size 2 --disable-pynccl boots and serves on 2x RTX PRO 4000 (the pre-fix loader dies in load_state_dict with the same LOAD_MISMATCH reported on PR FlashML-org#104).
…4-native emits Found by a post-merge correctness review of the TP port (two latent loader bugs, same class as the block-fp8 lm_head fix 6197d63): 1. _iter_weights_attn_fp8 (mixed modelopt path): a bf16 lm_head (lm_head_quant=='none') was yielded full-width against the vocab-sharded ParallelLMHead at TP>1. An NVFP4 lm_head never reaches the sharding block (emitted native above), so adding lm_head.weight to the vocab-parallel branch is safe in both cases. 2. Main iter_weights path: the NVFP4-native emit (shared_expert / dense MLP / lm_head) ran through _maybe_shard, which sharded down_proj (dim 1) / lm_head (dim 0) while leaving weight_scale / weight_global full-width -- but the model holds every NVFP4-native layer replicated (Nvfp4Dense* / Nvfp4LMHead). Yield it unsharded. Test: make_tiny_mixed_ckpt (per-tensor FP8 attn/GDN + NVFP4 shared_expert + bf16 lm_head) + test_mixed_iter_weights_tp2_sharding -- fails on the pre-fix loader (full-width lm_head), passes with the fix. Full CPU suite 11/11.
Under TP>1 every rank process auto-pinned its CPU MoE worker pool to the whole machine (one thread per physical core), so the rank pools time-sliced on the same cores with constant cache thrashing. Measured on 2x RTX 3090 (alogotron) and 2x RTX PRO 4000 (Rudi): hybrid TP=2 decode collapsed to ~2 tok/s; disjoint per-rank slices recover ~28-29 tok/s on Rudi (13.5x), and the --moe-cpu-threads 8 cap only ~2x because both ranks still take the first N cores. - resolve_threads_and_affinity(rank, world_size): each rank gets a disjoint round-robin slice of the physical cores (reps[rank::world_size]); an explicit --moe-cpu-threads count stays inside the rank's slice. Round-robin keeps a fair P/E-core mix on heterogeneous CPUs. - CpuMoeExecutor reads try_get_tp_info() and passes rank/world through; the torch intra-op thread clamp now counts the rank's slice, not the whole machine. - The 'CPU MoE executor ready' line now logs on ALL ranks with tp_rank=N/M (it was info_rank0, which made a pinning collision invisible in the logs). - tests/moe/test_cpu_executor_affinity.py: partition invariants (disjoint slices, explicit count stays in slice, TP=1 unchanged).
…patch) Adds the 'nvfp4' weight_format to the resident (fused) MoELayer: - _alloc_resident_experts: 6 native ModelOpt NVFP4 banks (gate_up_packed/scale/ global, down_packed/scale/global) on GPU, TP-sharded on the intermediate dim (gate_up column-parallel, down row-parallel; down global keeps full H rows). - _resident_gemm: dispatch to the existing plain *_nvfp4 Triton kernels (fused_experts_nvfp4 prefill / fused_experts_decode_nvfp4_marlin decode), mirroring OffloadMoELayer._expert_gemm's 'nvfp4' branch. - Qwen3_5MoE selects weight_format='nvfp4' for expert_quant=='nvfp4'. Unwired: iter_weights still skips per-expert NVFP4 tensors (next commit builds + yields the resident sharded banks).
Wires the resident (fused) NVFP4 expert path end to end: - nvfp4_banks.load_nvfp4_expert_resident_banks: the resident mirror of load_nvfp4_expert_source_banks -- allocates the 6 native NVFP4 banks on host, TP-sharded on the intermediate dim (gate_up column-parallel on the 2*I output rows, down row-parallel on the I input cols: //2 for the FP4 bytes, //16 for the per-16 block scale; down global keeps full H rows). Never pinned. - qwen3_5_moe/weight._iter_resident_nvfp4_experts: yields the per-layer sharded banks as state-dict views (keys match the MoELayer's resident NVFP4 attrs). Called from both dense passes (_iter_weights_attn_fp8 mixed + the pure-NVFP4 path) when include_moe_experts=True (the fused backend). - tests: tiny_fp8_ckpt.make_tiny_mixed_ckpt gains routed_experts=True to emit the per-expert NVFP4 tensors; two CPU tests verify the TP=2 sharding (each rank's slice == the checkpoint slice) and the TP=1 full banks. CPU-verified on freetoken:local (TP=1 + TP=2 sharding). GPU E2E next.
_adjust_config rejected expert_quant=='nvfp4' for any non-offload backend. Now that MoELayer has a resident NVFP4 allocation + kernel dispatch (TP-sharded on the intermediate dim), 'nvfp4' joins 'none'/'fp8_block' as a valid resident format. The other quant formats (mxfp4, q4_0, ds_fp4) still require offload/cpu.
Owner
|
@CraigStone-Dev Great job， find one critical issue:
TP=2 → output = 2× correct value, no error. Offload path is unaffected (OffloadMoELayer._maybe_all_reduce is a no-op). |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This branch completes the TP port for the quantized
qwen3_5_moecheckpoints — the piece missing from this PR, which covers BF16. Community testing (alogotron, 2x RTX 3090) showed FP8/NVFP4 checkpoints die atmake_row_parallel_quantunder TP>1.Changes (on top of 3c8b281)
2eae9a8— feat(models): complete TP for qwen3_5_moe quantized checkpoints (FP8/NVFP4)Fp8PerTensorRowParallel,Fp8BlockRowParallel,Nvfp4DenseRowParallel— sharded-K GEMM + all-reduce, bias after the reduce. The existing Triton GEMMs are K-generic, so no new kernels were needed.local_output_sizes(GQA KV-head replication); themake_col_merged_quantguard is gone.gdn.py:in_proj_qkvznow declares TP-local sizes (the PR left it full-size while forward splits by local sizes — a latent TP>1 bug).tests/models/test_qwen3_5_tp_quant.py(12 tests, all pass).ab82a3f— test: add TP1/TP2 A/B verification scriptstests/ab_suite.py/tests/ab_long.py: deterministic TP1-vs-TP2 comparison clients for a served endpoint (used for the E2E check below).6197d63— fix(models): shard lm_head vocab-parallel in the block-fp8 loader at TP>1Qwen/Qwen3.6-35B-A3B-FP8) dies at TP=2: the loader yieldedlm_headfull-width (248320x2048) against the vocab-shardedParallelLMHead(124160 rows). Per-tensor sharding is now in_fp8_shard_standalone, withlm_head.weightadded to the vocab-parallel branch.a77453f— test: tiny block-FP8 checkpoint generator + TP=2 loader regression testtests/models/tiny_fp8_ckpt.py: shrinks the real Qwen3.6-35B-A3B-FP8 geometry (4 layers: 3 GDN + 1 full-attn, few experts) to a few hundred MB while keeping hidden 2048, all per-layer dims, and by default the full 248320 vocab (random weights -- for load/boot validation, not coherent text).test_fp8_iter_weights_tp2_sharding: runs the full_iter_weights_fp8at TP=2 over the synthetic checkpoint (lm_head/embed vocab-sharded, o_proj/out_proj input-sharded, qkv/in_proj_qkvz per-part local rows, conv1d sub-parts, replicated router/shared expert); fails on the pre-fix loader, passes with the fix.ft serve --tensor-parallel-size 2 --disable-pyncclboots and serves on 2x RTX PRO 4000 (the pre-fix loader dies inload_state_dictwith the reported mismatch).f1ff8e2— fix(models): shard bf16 lm_head in the mixed loader; never shard NVFP4-native emits_iter_weights_attn_fp8) yielded a bf16 lm_head (lm_head_quant=="none") full-width against the vocab-shardedParallelLMHead-- same class as6197d63(an NVFP4 lm_head is emitted native above the sharding block and correctly stays replicated); (2) the mainiter_weightspath ran the NVFP4-native emit (shared_expert / dense MLP / lm_head) through_maybe_shard, which shardeddown_proj/lm_headwhile leaving their scales full-width -- but the model holds every NVFP4-native layer replicated.make_tiny_mixed_ckpt(per-tensor FP8 attn/GDN + NVFP4 shared_expert + bf16 lm_head) +test_mixed_iter_weights_tp2_shardingguard (1); fails pre-fix, passes post-fix. Full CPU suite: 12/12 intest_qwen3_5_tp_quant.py+ 3/3 intest_qwen3_5_tp.py.Validation (2x RTX PRO 4000 Blackwell, PCIe, no NVLink)
Model:
nvidia/Qwen3.6-35B-A3B-NVFP4(mixed: FP8 per-tensor attention + NVFP4 experts/shared_expert/lm_head) — the target checkpoint for this port.ft serve --tensor-parallel-size 2 --disable-pynccl --moe-cache-size 4096 --expert-load serialboots clean; CUDA-graph capture (bs 1/2/4) with NCCL all-reduce works.KV capacity: 1,223,892 tokens/rank at TP=2 vs 517,045 at TP=1 -> ~2.4x (about 4-5 concurrent 256K contexts vs about 2).
Correctness A/B (temperature 0): 10/10 identical on the short suite; on the 512-token long suite 1/5 byte-identical and the rest semantically identical with only last-ulp drift (all-reduce changes FP reduction order, flipping a few near-tie greedy tokens). Final answers matched.
3-way concurrent requests batch correctly (warm wall 11s, all correct).
Note:
--expert-load serialis required on a 54 GB-RAM box (expert banks are 21.8 GB per rank, replicated per Offloaded MoE ignores tensor parallelism: TP=2 is a regression, and the fix is not expert parallelism FlashML-org/FreeToken#62 -> ~44 GB pinned; the parallel O_DIRECT build OOM-kills rank 0).Re-validated after all fixes (2026-08-26):
freetoken:localimage rebuilt with the full stack; both--tensor-parallel-size 2and1boot clean and serve correct completions on the same hardware (KV pools unchanged: 1,223,892 tokens/rank at TP=2, 517,045 at TP=1).Disclosure
This contribution was developed with an AI coding agent (implementation and tests are AI-generated); all validation above was run on real hardware (2x RTX PRO 4000 Blackwell) and the numbers are from those runs. Happy to adjust scope or split the commits however is most useful for review.