cp: fix(moe): preserve fp32 A_log in Qwen3.5-{MoE,Next GatedDeltaNet} (2484) into r0.5.0#2664
Merged
Conversation
Contributor
Author
|
/ok to test 67e0f6e |
…aNet (#2484) * fix(moe): preserve fp32 A_log in Qwen3.5-MoE and Qwen3-Next GatedDeltaNet GatedDeltaNet A_log/dt_bias are intrinsically fp32, but the native MoE Qwen3.5-MoE and Qwen3-Next models sharded through moe/parallelizer.py with a single bf16 mixed-precision policy, silently rounding these params to bf16 and degrading the exponentiated decay rate. Generalize the dense Qwen3.5 _fp32_params holder mechanism into a shared helper (components/models/common/gated_delta_net_fp32.py), isolate the fp32 GDN params into per-module holders at parallelize time, and shard each holder as a dedicated fp32 FSDP unit in apply_fsdp (preserving EP ignored_params). Add an fp32-aware Qwen3NextGatedDeltaNet subclass that routes the gate through the holder, and strip the _fp32_params prefix on save so checkpoints stay HF-clean. Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(moe): avoid pre-reading holder-owned dt_bias in Qwen3.5-MoE gate CPAwareGatedDeltaNet._compute_gate evaluated self.dt_bias before entering the holder forward. Once isolate_fp32_params moves an intrinsically-fp32 dt_bias into the _fp32_params holder, the patched __getattr__ makes self.dt_bias resolve to the holder-owned param, reading it outside the holder forward (before FSDP unshards), which defeats the fp32 isolation. The holder forward already prefers its own dt_bias, so route through holder(a) when it owns dt_bias and only pass self.dt_bias when dt_bias stays a bare (bf16) param -- mirroring Qwen3-Next. Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * test(moe): cover apply_fsdp fp32 GDN branch; exclude GPU-only Next forward Add a CPU unit test that drives the apply_fsdp fp32 branch (isolation reported, fp32 holder policy built, per-block holder sharded) to lift Codecov patch coverage reliably on every PR. Mark Qwen3NextFp32GatedDeltaNet.forward with ``# pragma: no cover`` since it is a verbatim HF GDN forward that requires CUDA conv1d/FLA kernels and is only exercised by GPU/functional runs. Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(qwen): preserve GDN fp32 params under FSDP2 Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(qwen): scope GDN fp32 isolation to linear attention Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(qwen): mark fp32 GDN modules explicitly Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * docs(model): clarify GDN fp32 marker contract Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(nemotron): keep router correction bias fp32 Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(dtype): preserve fp32 values during model casts Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * refactor(models): make qwen-next own fp32 gdn holder Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * docs(skills): clarify fp32 holder ownership Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(checkpoint): tolerate missing TE extra state Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> * fix(models): preserve duplicate fp32 buffers Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> --------- Signed-off-by: Yuhe Zhang <yuhez@nvidia.com> Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> (cherry picked from commit 174c5d1) Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
akoumpa
force-pushed
the
cherry-pick-2484-r0.5.0
branch
from
67e0f6e to
907e47f
Compare
Contributor
Author
|
/ok to test 907e47f |
akoumpa
changed the title
fix(moe): preserve fp32 A_log in Qwen3.5-MoE and Qwen3-Next GatedDeltaNet (2484) into r0.5.0fix(moe): preserve fp32 A_log in Qwen3.5-{MoE,Next GatedDeltaNet} (2484) into r0.5.0
Contributor
Author
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.
Cherry-pick of #2484 (174c5d1) onto
r0.5.0.Release-branch resolutions
gated_delta_net_fp32.py,moe/parallelizer.pyfp32-holder sharding,model_init,nemotron_v3).checkpointing.py"tolerate missing TE_extra_state" hunk (theextra_state_keysblock from thefix(checkpoint): tolerate missing TE extra statesub-commit). It is stacked on the main-only feat(checkpoint): export torch_save checkpoints to hf #2487allow_checkpoint_key_subsetfeature (labeledcommunity-request, never slated for r0.5.0) and is a robustness tangent separate from the fp32 fix. The fp32 checkpoint round-trip itself is handled by the state-dict adapters (strip _fp32_params on save), which are included — so the core fix is intact. Thetest_checkpointing.pyconflict was likewise just r0.5.0 lacking feat(checkpoint): export torch_save checkpoints to hf #2487'sTestLoadModelCheckpointKeySubset(fix(moe): preserve fp32 A_log in Qwen3.5-MoE and Qwen3-Next GatedDeltaNet #2484 makes no net change to that file).test_shard_fp32_param_holders_shards_each_holder,test_apply_fsdp_shards_model_owned_fp32_holders) inmoe/test_parallelizer.py, rebased onto feat(qwen3_5): port dense Qwen3.5 to a native custom-model implementation #2557's version of that file.If the missing-TE-
_extra_statetolerance is wanted on r0.5.0, it should be backported separately (it needs an r0.5.0-appropriate adaptation that does not depend on #2487).Validation (CPU/local):
py_compileclean (25 files); fp32 + parallelizer + qwen3_5_moe/qwen3_next/common unit tests pass. GPU parity not re-run locally — see #2484's own 8×H100 validation.