Recognize Qwen3.5's RMSNorm variants in AutoTP module loading - #8306
Conversation
Loading.is_load_module() in auto_tp.py gates whether a leaf module's parameters get loaded from a raw state dict during AutoTP-based checkpoint loading (both call sites in auto_tp.py and replace_module.py use it the same way). It matches by an exact class-name allowlist that gets a new entry each time a model family ships its own RMSNorm class, but Qwen3.5's classes were never added. Qwen3_5RMSNorm, Qwen3_5RMSNormGated (dense), Qwen3_5MoeRMSNorm and Qwen3_5MoeRMSNormGated (MoE) each own a `weight` nn.Parameter, same shape as every other listed *RMSNorm class - without an allowlist entry, is_load_module() returns False for them and that weight is never loaded from the checkpoint, silently left at its random init value instead. (deepspeedai#7947) Verified against the real modeling code (transformers 5.3.0, the version reported in the issue): built a small Qwen3_5TextModel and confirmed Qwen3_5RMSNorm/Qwen3_5RMSNormGated instances only carry a `weight` parameter, matching the pattern of the already-listed LlamaRMSNorm et al. Confirmed the MoE variants (Qwen3_5MoeRMSNorm, Qwen3_5MoeRMSNormGated) carry the same single `weight` parameter. Deliberately not adding Qwen3_5TextRotaryEmbedding: it carries no parameters, only two buffers (inv_freq, original_inv_freq), and both call sites already load any child's buffers unconditionally (`if len(child._buffers) != 0`) regardless of is_load_module() - so listing it would be a no-op, not a functional fix. Test: tests/unit/module_inject/test_auto_tp_is_load_module.py (new). Uses name-matched dummy nn.Module stand-ins rather than real transformers classes, since is_load_module() matches by class name only and the dev requirement (transformers>=4.51.3) predates Qwen3.5 - importing the real classes would make the test depend on a newer transformers than the repo's own pinned minimum. Covers all four new allowlist entries plus one negative case for an unrelated class name. 5 passed locally. pre-commit (yapf, flake8, check-license, codespell) clean on both changed files. Signed-off-by: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a348922b39
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| "Phi3RMSNorm", "YuanRMSNorm", "YuanRotaryEmbedding", "Phi3LongRoPEScaledRotaryEmbedding", "Qwen2RMSNorm", | ||
| "Qwen3RMSNorm", "Qwen3MoeRMSNorm", "DeepseekV2RMSNorm", "DeepseekV3RMSNorm", | ||
| "DeepseekV2YarnRotaryEmbedding", "DeepseekV3YarnRotaryEmbedding", "MoEGate" | ||
| "Qwen3RMSNorm", "Qwen3MoeRMSNorm", "Qwen3_5RMSNorm", "Qwen3_5RMSNormGated", "Qwen3_5MoeRMSNorm", |
There was a problem hiding this comment.
This non-merge commit has no Signed-off-by trailer, so it violates the repository's mandatory commit policy. Recreate the commit with --signoff using the configured Git name and email before merging.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
|
Hi @promptsmith1990 , thanks for your fix. Does this PR fix issue #7947, so that finetuning Qwen 3.5 family with AutoTP is possible, or it just fix the stated Qwen3.5 RMSNorm variants issue? Thanks! |
|
|
||
| class TestIsLoadModule: | ||
|
|
||
| @pytest.mark.parametrize("class_name", [ |
There was a problem hiding this comment.
I think we should either remove this test or verify all elements in is_load_module. Given that it is unlikely that an item be taken out of is_load_module with out causing attention, I suggest not having this test at all.
There was a problem hiding this comment.
Removed the test file entirely per your suggestion — it only covered the four new Qwen3.5 entries rather than the full allow-list, and as you noted, removals from is_load_module would be obvious in code review. The functional change in auto_tp.py stands unchanged.
tohtana
left a comment
tohtana
left a comment
There was a problem hiding this comment.
Please use the new copyright header showing DeepSpeed team for new files.
# Copyright (c) DeepSpeed Team.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
|
@promptsmith1990 any updates? Thanks! |
Hi! Thanks for your review, really appreciate it. Working on your comments. |
The test only exercised the four newly-added Qwen3.5 class names, not the full allow-list in is_load_module(). Rather than expand it into an exhaustive list (which would need updating every time a new family is added), drop the test entirely per delock's suggestion — removals from is_load_module() would be immediately visible in code review anyway. Signed-off-by: promptsmith1990 <319963136+promptsmith1990@users.noreply.github.com>
|
Thank you for your review! Reviewer feedback addressed:
|
Hi @thuanaislab — good question. This PR specifically fixes Loading.is_load_module() in auto_tp.py, which controls whether a module's parameters get loaded from the checkpoint during AutoTP-based model loading. Without this fix, Qwen3.5's RMSNorm layers (Qwen3_5RMSNorm, Qwen3_5RMSNormGated, Qwen3_5MoeRMSNorm, Qwen3_5MoeRMSNormGated) silently start with randomly initialized weights instead of checkpoint values — which is the root cause of the failure reported in #7947 when using transformers ≥ 5.3.0. So yes, it addresses the is_load_module gap that was blocking Qwen3.5 + AutoTP. |
tohtana
left a comment
tohtana
left a comment
There was a problem hiding this comment.
Thank you for the update, @promptsmith1990! Looks good to me.
3 participants
Fixes #7947.
Problem
Loading.is_load_module()inauto_tp.pygates whether a leaf module's parameters get loaded from a raw state dict during AutoTP-based checkpoint loading (both call sites —auto_tp.pyandreplace_module.py— use it the same way). It matches by an exact class-name allowlist that gets a new entry each time a model family ships its own RMSNorm class, but Qwen3.5's classes were never added.Qwen3_5RMSNorm,Qwen3_5RMSNormGated(dense),Qwen3_5MoeRMSNormandQwen3_5MoeRMSNormGated(MoE) each own aweightnn.Parameter, same shape as every other listed*RMSNormclass — without an allowlist entry,is_load_module()returnsFalsefor them and that weight is never loaded from the checkpoint, silently left at its random init value instead.Verification
Verified against the real modeling code (transformers 5.3.0, the version reported in the issue): built a small
Qwen3_5TextModeland confirmedQwen3_5RMSNorm/Qwen3_5RMSNormGatedinstances only carry aweightparameter, matching the pattern of the already-listedLlamaRMSNormet al. Confirmed the MoE variants carry the same singleweightparameter.Deliberately not adding
Qwen3_5TextRotaryEmbedding: it carries no parameters, only two buffers (inv_freq,original_inv_freq), and both call sites already load any child's buffers unconditionally (if len(child._buffers) != 0) regardless ofis_load_module()— so listing it would be a no-op, not a functional fix. (Happy to add it anyway for consistency with the codebase's existingPhi3RotaryEmbedding/YuanRotaryEmbeddingentries if maintainers prefer explicitness over strict minimality here — wasn't sure which this repo's convention favors.)Test
tests/unit/module_inject/test_auto_tp_is_load_module.py(new). Uses name-matched dummynn.Modulestand-ins rather than realtransformersclasses, sinceis_load_module()matches by class name only and the dev requirement (transformers>=4.51.3) predates Qwen3.5 — importing the real classes would make the test depend on a newer transformers than the repo's own pinned minimum. Covers all four new allowlist entries plus one negative case for an unrelated class name.pre-commit run --files deepspeed/module_inject/auto_tp.py tests/unit/module_inject/test_auto_tp_is_load_module.py(yapf, flake8, check-license, codespell) is clean.Prepared with AI assistance under my review; the root cause, the parameter-vs-buffer distinction behind the RotaryEmbedding decision, and the test were verified locally before opening this PR.