Skip to content

[NNX] Delete Linen (12.1b): collapse dispatch in sharding and muon - #4684

Open
ecnal-cienet wants to merge 1 commit into
feat/nnx-del-linen-pretrain-corefrom
feat/nnx-del-linen-pretrain-sharding-opt
Open

[NNX] Delete Linen (12.1b): collapse dispatch in sharding and muon#4684
ecnal-cienet wants to merge 1 commit into
feat/nnx-del-linen-pretrain-corefrom
feat/nnx-del-linen-pretrain-sharding-opt

Conversation

@ecnal-cienet

@ecnal-cienet ecnal-cienet commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

  • why is this change being made,
  • the problem being solved and any relevant context,
  • why this is a good solution,
  • some information about the specific implementation,
  • shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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.

@gemini-code-assist

gemini-code-assist Bot commented Jul 31, 2026

Copy link
Copy Markdown

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

@codecov

codecov Bot commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.90909% with 1 line in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/utils/muon_utils.py 90.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@ecnal-cienet ecnal-cienet changed the title Feat/nnx del linen pretrain sharding opt [NNX] Delete Linen (12.1b): collapse dispatch in sharding and muon Jul 31, 2026
@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-sharding-opt branch 4 times, most recently from fd2bc25 to 0279059 Compare August 20, 2026 21:38
@ecnal-cienet
ecnal-cienet changed the base branch from main to feat/nnx-del-linen-pretrain-core August 24, 2026 20:10
@ecnal-cienet
ecnal-cienet marked this pull request as ready for review August 24, 2026 20:15
@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-core branch from d095c8f to 70a1185 Compare August 25, 2026 18:38
@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-sharding-opt branch from 0279059 to 89897a7 Compare August 25, 2026 18:38
@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-core branch from 70a1185 to 94526bc Compare August 31, 2026 14:13
@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-sharding-opt branch from 89897a7 to b51c644 Compare August 31, 2026 14:13
@ecnal-cienet
[NNX] Delete Linen (pre-train 2/3): collapse dispatch in sharding and…
dd76238 
… muon utils

- sharding.maybe_update_params_sharding_with_opt delegates to the _nnx variant;
  build_zero1_input_state_mesh_shardings drops its Linen early return.
- muon_utils.get_muon_weight_dimension_numbers drops the isinstance(nnx.Module)
  test and the Linen get_abstract_param path; get_model_mdn loses its pure_nnx
  parameter and always builds the abstract NNX model.
- run_sharding_dump loses its --pure_nnx flag.

Tests follow: the Linen branch test in muon_utils_test goes away, optimizers_test
drops the pure_nnx argument and the dual-shape comparison, sharding_nnx_test
drops the flag from its fake config.
@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-sharding-opt branch from b51c644 to dd76238 Compare August 31, 2026 15:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@bvandermoon bvandermoon bvandermoon approved these changes
@NuojCheng NuojCheng NuojCheng 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
@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
@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner
@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner
@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
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ecnal-cienet @bvandermoon @NuojCheng