Skip to content

DeepSeek V4 Muon supprt#4065

Open
snehalv2002 wants to merge 1 commit into
mainfrom
snehalv-dsv4-muon
Open

DeepSeek V4 Muon supprt#4065
snehalv2002 wants to merge 1 commit into
mainfrom
snehalv-dsv4-muon

Conversation

@snehalv2002

@snehalv2002 snehalv2002 commented Jun 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Add support for Muon optimizer for DeepSeek V4. This was the parameter mapping for parameter -> Muon/AdamW used, following quote from the paper: "We maintain the AdamW (Loshchilov and Hutter, 2017) optimizer for the embedding module, the prediction head module, the static biases and gating factors of mHC modules, and the weights of all RMSNorm modules. All other modules are updated with Muon."

Tests

Can see the the perplexity/lm_loss decreasing in this run using this train command where we train with Muon.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@github-actions

github-actions Bot commented Jul 4, 2026

Copy link
Copy Markdown
Contributor

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.

@github-actions github-actions Bot added the stale Automatically applied to stale PRs. label Jul 4, 2026
@snehalv2002
snehalv2002 changed the base branch from main to dsv4-bug-fixes July 9, 2026 21:52
@snehalv2002 snehalv2002 changed the title Snehalv dsv4 muon DeepSeek V4 Muon supprt Jul 10, 2026
@snehalv2002
snehalv2002 marked this pull request as ready for review July 16, 2026 00:28
@snehalv2002 snehalv2002 removed the stale Automatically applied to stale PRs. label Jul 16, 2026
@snehalv2002
snehalv2002 force-pushed the snehalv-dsv4-muon branch 2 times, most recently from 9b98b4c to a3f7c25 Compare July 16, 2026 23:05
Base automatically changed from dsv4-bug-fixes to main July 18, 2026 01:23
@github-actions

github-actions Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

🤖 Hi @shuningjin, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

@github-actions

github-actions Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

🤖 I'm sorry @shuningjin, but I was unable to process your request. Please see the logs for more details.

@github-actions

github-actions Bot commented Jul 20, 2026

Copy link
Copy Markdown
Contributor

🤖 Hi @shuningjin, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

shuningjin
shuningjin reviewed Jul 21, 2026
View reviewed changes

@shuningjin shuningjin left a comment •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Greak work! Thanks for incorporating hybrid Newton Schultz iterations and adding parameter transformation for DeepSeek4!

The mapping looks reasonable. Only minor comments.

Comment thread src/maxtext/optimizers/optimizers.py Outdated
Comment thread src/maxtext/utils/muon_utils.py
Comment thread src/maxtext/utils/muon_utils.py
Comment thread src/maxtext/utils/muon_utils.py Outdated
Comment thread tests/unit/optimizers_test.py Outdated
shuningjin
shuningjin approved these changes Jul 22, 2026
View reviewed changes

@shuningjin shuningjin left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. Thank you!

parambole
parambole approved these changes Jul 22, 2026
View reviewed changes

@parambole parambole left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@parambole parambole parambole approved these changes
@aireenmei aireenmei aireenmei approved these changes
@shuningjin shuningjin shuningjin approved these changes
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner
@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner
@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@gagika gagika Awaiting requested review from gagika gagika is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner
@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner
@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner
@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner
@NicoGrande NicoGrande Awaiting requested review from NicoGrande
@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu is a code owner
@michelle-yooh michelle-yooh Awaiting requested review from michelle-yooh
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer is a code owner

Assignees

No one assigned

Labels

gemini-review pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@snehalv2002 @parambole @aireenmei @shuningjin