Skip to content

[NNX] Delete Linen (12.1c): collapse dispatch in quantization and model creation - #4685

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

[NNX] Delete Linen (12.1c): collapse dispatch in quantization and model creation#4685
ecnal-cienet wants to merge 1 commit into
feat/nnx-del-linen-pretrain-sharding-optfrom
feat/nnx-del-linen-pretrain-quant

Conversation

@ecnal-cienet

@ecnal-cienet ecnal-cienet commented Jul 31, 2026 •
edited
Loading

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.

@ecnal-cienet ecnal-cienet changed the title NNX Migration 12.1c (pre-train 3/3) — feat/nnx-del-linen-pretrain-quant (base: feat/nnx-del-linen-pretrain-sharding-opt [NNX] Delete Linen (12.1c): collapse dispatch in quantization and model creation Jul 31, 2026
@codecov

codecov Bot commented Jul 31, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-quant branch 4 times, most recently from 87c0371 to 125109f Compare August 20, 2026 21:38
@ecnal-cienet
ecnal-cienet changed the base branch from main to feat/nnx-del-linen-pretrain-sharding-opt 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-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-quant branch from 125109f to df096b4 Compare August 25, 2026 18:38
@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
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-quant branch from df096b4 to d482492 Compare August 31, 2026 14:13
@ecnal-cienet
[NNX] Delete Linen (pre-train 3/3): collapse dispatch in quantization…
7da8b8a 
… and model creation

- quantizations.maybe_quantize_model always runs the qwix forward pass with the
  dummy tokens/positions/segment ids (and the MTP decoder targets when
  mtp_num_layers > 0), then pops the transient nnx.Intermediate variables the
  traced forward sows.
- model_creation_utils.from_pretrained always builds the sharded model through
  maxtext_utils_nnx.create_nnx_sharded_model.

Tests follow: quantizations_test and nnx_quant_guard_test drop their flag
arguments and Linen expectations, correctness_tests_nnx_dispatch_test keeps only
the NNX case, and forward_pass_logit_checker always loads via from_pretrained.
@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
@ecnal-cienet
ecnal-cienet force-pushed the feat/nnx-del-linen-pretrain-quant branch from d482492 to 7da8b8a 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
@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner
@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner
@jshin1394 jshin1394 Awaiting requested review from jshin1394 jshin1394 is a code owner
@liudangyi liudangyi Awaiting requested review from liudangyi liudangyi is a code owner
@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner
@shralex shralex Awaiting requested review from shralex shralex is a code owner
@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner
@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner
@zxhe-sean zxhe-sean Awaiting requested review from zxhe-sean zxhe-sean is a code owner
@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner
@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk 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
@gagika gagika Awaiting requested review from gagika gagika 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
@xibinliu xibinliu Awaiting requested review from xibinliu xibinliu is a code owner
@huytransformer huytransformer Awaiting requested review from huytransformer huytransformer 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