feat: support qwen3 next with muon#2875
feat: support qwen3 next with muon#2875Prayer3th wants to merge 1 commit intoAI-Hypercomputer:mainfrom
Conversation
Prayer3th
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
jiangjy1982,
khatwanimohit,
richjames0,
shralex,
suexu1025 and
vipannalla
as code owners
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
LGTM! A few minor comments. Thank you for the change!
| model_name_arg = sys.argv[1] | ||
| scan_layers_arg = sys.argv[2].lower() == "true" | ||
| get_model_mdn(model_name_arg, scan_layers_arg, verbose=True) | ||
| get_model_mdn(model_name_arg, scan_layers_arg, verbose=True) No newline at end of file |
There was a problem hiding this comment.
nit: extra line, similar for other files.
| An instance of `MuonDimensionNumbers` if a specific mapping is found, | ||
| `None` for excluded parameters, or a default `mdn` for standard weights. | ||
| 1. Exclusions: Skip vectors/biases/embeddings (AdamW). | ||
| 2. MoE: Handle both DeepSeek style (MoeBlock_0) and Qwen3-Next style (routed_experts). |
There was a problem hiding this comment.
Nit: quote the checkpoint name MoE: Handle both DeepSeek style ("MoeBlock_0") and Qwen3-Next style ("routed_experts").
| @@ -48,51 +48,69 @@ def _is_path_contain_any(tuples, path): | |||
| def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]: | |||
There was a problem hiding this comment.
@shuningjin Could we have a refactor PR to pass in model name as follow up? I have seen checkpoint name divergence. It will be better we transform weights based on model, instead of checkpoint path.
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Could you pass the screenshot or log, to train Qwen3 with Muon optimizer end-to-end?
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions. |
|
This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it. |
5 participants
Description
support qwen3-next use muon optimizer
Tests
tests/muon_test.py
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.