refactor(quant): hand the checkpoint QuantConfig to the weight readers - #427
Merged
Merged
Conversation
jason-fxz
added this pull request to stack #432
September 10, 2026 08:44
gdevenyi
added a commit
to gdevenyi/FreeToken
that referenced
this pull request
Sep 10, 2026
Brings the quantization refactor (FlashML-org#418 config/scheme/method layers, FlashML-org#427 the QuantConfig reaching the weight readers) and FlashML-org#426, which takes qwen4_exp's expert quant kind from the checkpoint QuantConfig. Resolutions, all taken from the rebased PR branches so the deploy tree and the upstream PRs stay identical where they overlap (rb/tp, rb/vision, rb/fp8 -> rb/all): - qwen4_exp attention/gdn/moe/config/model, layers/linear.py, layers/embedding.py, models/nvfp4_banks.py: the rebased versions. Notably o_proj stays LinearOProj (row-parallel); main's LinearReplicated is correct only at TP=1. - NVFP4 expert TP sharding moved to the new seam: the piece stream is sliced along the intermediate axis in nvfp4_banks, and the kernel sizes its banks from MoEConfig.local_intermediate. The old _Placer/_alloc_nvfp4_host_banks are gone with the functions they lived in. - models/quant_linear.py: deleted, as upstream did; nothing imports it. - DROPPED, superseded upstream: the FlashML-org#320 cherry-pick (_dense_is_block_fp8, _block_fp8_fusions, _load_maybe_block_fp8). FlashML-org#320 is CLOSED upstream because FlashML-org#426 does the same job through QuantConfig. This checkpoint's dense side is bf16 (a 292-entry modelopt ignore list), so none of those paths were exercised here. - KEPT: the --num-tokens/--num-pages KV reserve fix (issue FlashML-org#383, still unfixed on main), FlashML-org#231's collect_decode_freq, FlashML-org#169's prefill_warmup, and every other deploy-only commit -- they merged clean. - tests/engine/test_cache_budget.py keeps both the KV-reserve test and upstream's new slot_limit test; _resolve_auto_moe_cache_size's new method arg is optional. deploy/chatdnp stays at adc32da as the rollback point. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
This was referenced Sep 10, 2026
Closed
This was referenced Sep 10, 2026
This was referenced Sep 10, 2026
This was referenced Sep 10, 2026
This was referenced Sep 11, 2026
tomasuz
pushed a commit
to tomasuz/FreeToken
that referenced
this pull request
Sep 11, 2026
Upstream FlashML-org#418/FlashML-org#427/FlashML-org#438 moved expert quantization into MoEMethod/MoEKernel (layers/quantization/moe): bf16, NVFP4 and fp8 experts are packed by build_expert_banks and applied through quant_method.apply; only GGUF still loads through a format provider. --moe-backend became --moe-strategy and EngineConfig.moe_backend is folded into moe_strategy and left None. Resolution keeps every feature of this branch on top of that design: - engine: take upstream's shared_offload_method / _check_pin_budget / PinFailed / layout+max_slots path; keep the resident tier, worker executors, in-place worker layers, shared banks and decode-frequency stats. load_expert_banks is called once, inside shared_banks(), with both method= and resident_layers=. Our helpers read moe_strategy. - expert_banks: upstream build_expert_banks plus ResidentUploader; the resident upload rides host_banks' PinPipeline sink, which upstream's builder already drives. The bf16/nvfp4/ds_fp4 providers are gone as upstream intended; one GGUF provider serves every native ggml quant. - layers/moe: the format-tag dispatch keeps only the GGUF branch; the split/worker decode path already goes through _expert_gemm and so through quant_method.apply. - models/weight: load_gguf_moe_expert_sources, with upstream's load_q4_0_moe_expert_sources kept as an alias. - moe/fused_q4_0: add upstream's fused_experts_gguf_q4_0 entry point. - moe/_worker_main: take NVFP4 banks by kernel role (gate_up, down, ...), falling back to the FTW names. - tests: fake configs in test_resident / test_worker_placement use moe_strategy. Checked on tm (ROCm 7.2, torch 2.11): no undefined names in the 180 changed files; tests/engine + tests/moe show no failure that is not also present on the pre-merge branch or on pure upstream main in the same environment (51 failed / 285 passed; test_nvfp4_moe_vec.py and test_warp_width.py excluded -- both segfault identically before the merge). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014Tt8JH1wvkbStVdTSzHNPE
nomanoma121
pushed a commit
to nomanoma121/My-FreeToken
that referenced
this pull request
Sep 16, 2026
… the readers Upstream FlashML-org#426 stops qwen4_exp guessing what a checkpoint quantized. It had been matching the modelopt ignore list against probe names by hand, and the new nvidia/Qwen3.8-Flash-Next-NVFP4 build does not fit that guess. The expert kind now comes from QuantConfig.scheme_for_name, and attn_quant / dense_quant / lm_head_quant are no longer parsed for this family -- every module takes its own scheme from the QuantConfig when it is built. Nothing in the fork read those three for Flash-Next. FlashML-org#427 is the groundwork for it: checkpoint_quant_config moved to models/register.py, and EngineConfig installs the result in a module global (set_quant_config), because the weight readers are handed the model path and nothing else. Upstream main has no caller yet -- the consumer is its qwen4_exp block-fp8 dense reader, which is not merged here. --dense-quant is what makes the order matter. It wraps the checkpoint's config in LoadTimeFp8Config, so the install has to come after the wrap, or a reader would quantize against schemes the layers were not built from. Nothing calls get_quant_config() yet, so no test would have noticed the difference: test_engine_config_dense_quant_override now asserts the identity, and that assert is what will catch it when the fp8-dense reader lands. FlashML-org#424 adds scripts/ftw_hotfix.py, which repairs FTW checkpoints written by older builds. It does not apply to this machine or the 3060 -- both serve the HF checkpoint dirs, not FTW. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
nomanoma121
pushed a commit
to nomanoma121/My-FreeToken
that referenced
this pull request
Sep 16, 2026
…layer it will become FlashML-org#428's dense fuser asks the QuantConfig what the model built a projection from and refuses the tensor whose dtype disagrees. Under --dense-quant those are two different answers on purpose: LoadTimeFp8Config reports every projection it is going to convert as per-row fp8, because that is what the layer becomes, while the shard still holds bf16 and the engine converts it in _quantize_at_load -- after the reader has handed it over. So the first GDN projection of the official Flash-Next checkpoint stopped the boot on both ranks: model.layers.0.linear_attn.in_proj_qkv.weight is torch.bfloat16 but the checkpoint's quant config declares model.layers.0.linear_attn.in_proj_qkvz QuantScheme(fp8_tensor, ...) The check now asks what the checkpoint stores that part under (scheme_for_name on the part's own checkpoint name, the same question qwen3_5_moe's _DenseReader.stored() asks), while the GDN in_proj layout keeps coming from the built scheme -- that one has to match gdn.py, which splits qkv|z from b|a exactly when scheme_for says the qkvz half is quantized, at load or in the file. This is the collision the FlashML-org#427 sync predicted and the assert in test_engine_config_dense_quant_override was placed for: get_quant_config() hands the readers the wrapped config, and FlashML-org#428 is its first consumer. Upstream has no --dense-quant, so upstream cannot see it; the new case reads a RadixArk-shaped bf16 checkpoint through the wrapper and pins both halves -- as-stored dtypes and the split layout. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
MT-z
added a commit
to MT-z/FreeToken
that referenced
this pull request
Sep 16, 2026
…off and FlashML-org#424's FTW script fb7f732 (FlashML-org#427) refactor(quant): hand the checkpoint QuantConfig to the weight readers 505477a (FlashML-org#424) chore(scripts): ftw_hotfix.py FlashML-org#427 is what makes this worth merging now rather than later. It moves checkpoint_quant_config from engine/config.py to models/register.py and installs the resolved QuantConfig in a module-level global from EngineConfig.model_config. Nothing in upstream main reads that global yet, so today it changes no behaviour -- but it is groundwork on the same path FlashML-org#418 and FlashML-org#426 already moved, and this box's model (Ornith, qwen3_5_moe, MIXED_PRECISION) loads through exactly that code. Divergence here costs more at the next commit than the merge costs now. FlashML-org#424 comes along because it sits between us and FlashML-org#427 on main. It is a standalone repair script plus docs, no engine change, and there are no FTW checkpoints on this box, so it does nothing here either way. Verified on this box (results/20260910-merge-427.txt in freetoken-systest): same 8 failures by node ID, 1,910 passed; greedy output bit-identical over 399 tokens at temperature 0.0; same startup warnings, same expert path. Assisted-by: Claude Opus 5
trcwebdesign
pushed a commit
to trcwebdesign/FreeToken
that referenced
this pull request
Sep 21, 2026
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
1 participant
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.
Weight readers only get the model path and cannot reach the QuantConfig the engine builds. This hands it over without touching reader signatures:
checkpoint_quant_configmoves fromengine/config.pytomodels/register.py, next to theModelSpecfields it interprets.set_quant_config/get_quant_config:EngineConfig.model_configinstalls the checkpoint's QuantConfig before parsing the family config; readers read it. Unset raises;Noneis the GGUF value. One per process, likeset_quant_backend.No reader uses it yet, so nothing loads differently.
Prerequisite for the qwen4_exp block-fp8 dense reader.