feat: convert `SimpleDecoderLayer` and `SimpleMlpDecoderLayer` to use nnx by mesakhcienet · Pull Request #1974 · AI-Hypercomputer/maxtext

mesakhcienet · 2025-07-15T03:01:43Z

Description

Use NNX for the SimpleDecodeLayer and SimpleMlpDecodeLayer modules.

We also implement a function to return class of the linen instead of an instantiated Linen object.

Tests

It should work.

Environment

Machine type : v5litepod-8 TPU VM
Software machine : tpu-ubuntu2204-base
Python Version : 3.10.12

Changes

`SimpleDecodeLayer`

Before change to NNX log : https://paste.googleplex.com/4892995450830848
After change to NNX log (train from scratch) : https://paste.googleplex.com/6179580754526208
After change to NNX log (continue train from main branch) : https://paste.googleplex.com/6264063394578432

Run command :

python -m MaxText.train MaxText/configs/base.yml \
  run_name=mesa-simple-test-small \
  base_output_directory=gs://mesa-maxtext/simple_test_small_1/  \
  decoder_block=simple dataset_type=synthetic \
  steps=10

`SimpleMlpDecodeLayer`

Before change to NNX log : https://paste.googleplex.com/5737594159366144
After change to NNX log (train from scratch) : https://paste.googleplex.com/5264591793422336
After change to NNX log (continue train from main branch) : https://paste.googleplex.com/4805767982743552

Run command :

python -m MaxText.train MaxText/configs/base.yml \
  run_name=mesa-simple-test-small-mlp  \
 base_output_directory=gs://mesa-maxtext/simple_test_small_2/ \ 
  decoder_block=simple_mlp \
  dataset_type=synthetic  \ 
  steps=10

Checklist

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

I have performed a self-review of my code.
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.

bvandermoon

It looks like the SimpleMlpDecoder test runs in the description are getting different perf before/after this change:

mesakhcienet · 2025-07-29T08:09:10Z

It looks like the SimpleMlpDecoder test runs in the description are getting different perf before/after this change:

my bad, I used the wrong function to call SimpleMlpDecoderLayer by using simple_decoder_layer_class() instead of simple_mlp_decoder_layer_class.

Both code and the log are updated, thank you.

bvandermoon · 2025-10-03T18:12:45Z

-        "weights",
-        nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-        (self.config.emb_dim, self.config.emb_dim),
+    init_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), sharding=("embed", "mlp"), mesh=self.mesh)


Is this needed? Not seeing it in the previous Linen version

Thank you for asking this part. We followed the previous Linen version as a reference to current change. For example:

self.weight_mat = self.param( "weights", nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")), (self.config.emb_dim, self.config.emb_dim), )

Since self.param should now be replaced with nnx.Param in NNX, we updated the implementation accordingly.
We also replaced nn.with_logical_partitioning from Linen with nnx.with_partitioning from NNX. And for clarity, we moved the nnx.with_partitioning initializer wrapper ( in current case -- init_fn) as an independent variable.

If I am correct, we also have similar implemntation for self.param at different places, such as gpt3.py

Feel free to tell me if the change needs to be updated, thank you

bvandermoon · 2025-10-03T18:15:23Z

-        "ff_1",
-        nn.with_logical_partitioning(nn.initializers.lecun_normal(), ("embed", "mlp")),
-        (self.config.emb_dim, self.config.mlp_dim),
+    init_ff1_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), sharding=("embed", "mlp"), mesh=self.mesh)


Why is this and the nnx.with_partitioning below added?

similar with answer above, thank you

NuojCheng

Overall LGTM, the memory usages match for both cases.

mesakhcienet · 2025-10-14T07:01:30Z

+      rngs = nnx.Rngs(params=0)
+      single_pipeline_stage = simple_layer.SimpleDecoderLayerToLinen(
+          config=config, mesh=mesh, model_mode=model_mode, rngs=rngs
+      )


since we update the implementation from Linen to NNX, we also need to update its related unit test

mesakhcienet · 2025-10-14T07:01:35Z

+    rngs = nnx.Rngs(params=0)
+    single_pipeline_stage = simple_layer.SimpleDecoderLayerToLinen(
+        config=config, mesh=mesh, model_mode=model_mode, rngs=rngs
+    )


since we update the implementation from Linen to NNX, we also need to update its related unit test

mesakhcienet marked this pull request as ready for review July 16, 2025 03:07

mesakhcienet requested review from A9isha, RissyRan, SurbhiJainUSC, aireenmei, bvandermoon, gagika, gobbleturk, hengtaoguo, khatwanimohit, richjames0, shralex, vipannalla and yangyuwei as code owners July 16, 2025 03:07

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch from 38a123e to 89de7e2 Compare July 22, 2025 02:26

mesakhcienet changed the title ~~feat: convert SimpleDecodeLayer and SimpleMlpDecodeLayer into nnx~~ feat: convert SimpleDecoderLayer and SimpleMlpDecoderLayer into nnx Jul 22, 2025

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch 3 times, most recently from 4476b54 to eb99fd8 Compare July 28, 2025 01:39

bvandermoon reviewed Jul 28, 2025

View reviewed changes

Comment thread MaxText/layers/simple_layer.py Outdated

mesakhcienet requested a review from bvandermoon July 28, 2025 05:58

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch from 8a6291b to 260271c Compare July 28, 2025 08:27

mesakhcienet changed the title ~~feat: convert SimpleDecoderLayer and SimpleMlpDecoderLayer into nnx~~ feat: convert SimpleDecoderLayer and SimpleMlpDecoderLayer to use nnx Jul 28, 2025

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch from 1606b5a to ece793f Compare July 29, 2025 03:16

bvandermoon reviewed Jul 29, 2025

View reviewed changes

Comment thread MaxText/layers/nnx_wrappers.py Outdated

bvandermoon reviewed Jul 29, 2025

View reviewed changes

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch 2 times, most recently from ed83621 to ead2980 Compare July 29, 2025 07:50

mesakhcienet requested a review from bvandermoon July 29, 2025 08:11

cgarciae reviewed Aug 6, 2025

View reviewed changes

Comment thread MaxText/layers/nnx_wrappers.py Outdated

mesakhcienet requested a review from cgarciae August 6, 2025 09:25

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch 2 times, most recently from 6b913a3 to 5f3e877 Compare August 7, 2025 02:19

bvandermoon reviewed Aug 7, 2025

View reviewed changes

Comment thread MaxText/layers/simple_layer.py Outdated

mesakhcienet requested a review from bvandermoon August 7, 2025 03:54

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch 2 times, most recently from 86b4205 to e3d58ff Compare August 11, 2025 09:32

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch from e3d58ff to aa2cd98 Compare August 19, 2025 06:31

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch from aa2cd98 to 7da2050 Compare September 10, 2025 02:14

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch from 7da2050 to c0f964c Compare September 30, 2025 01:27

mesakhcienet requested review from jiangjy1982, shuningjin and suexu1025 as code owners September 30, 2025 01:27

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch 4 times, most recently from d579507 to 0f6b007 Compare October 2, 2025 04:13

bvandermoon reviewed Oct 3, 2025

View reviewed changes

NuojCheng approved these changes Oct 8, 2025

View reviewed changes

NuojCheng mentioned this pull request Oct 13, 2025

Onboarding explicit sharding to MaxText #2470

Merged

4 tasks

feat: migrate simple_layers from linen to nnx

d2bc8f4

mesakhcienet force-pushed the feat/simple-decoder-layer-nnx branch from 7562255 to d2bc8f4 Compare October 14, 2025 03:23

mesakhcienet commented Oct 14, 2025

View reviewed changes

mesakhcienet requested a review from bvandermoon October 15, 2025 01:22

bvandermoon approved these changes Oct 15, 2025

View reviewed changes

NuojCheng added the pull ready label Oct 15, 2025

copybara-service Bot merged commit 44d6528 into AI-Hypercomputer:main Oct 16, 2025
27 checks passed

Uh oh!

Conversation

mesakhcienet commented Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Tests

Environment

Changes

SimpleDecodeLayer

SimpleMlpDecodeLayer

Checklist

Uh oh!

Uh oh!

Uh oh!

bvandermoon left a comment

Choose a reason for hiding this comment

Uh oh!

mesakhcienet commented Jul 29, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bvandermoon Oct 3, 2025

Choose a reason for hiding this comment

Uh oh!

mesakhcienet Oct 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bvandermoon Oct 3, 2025

Choose a reason for hiding this comment

Uh oh!

mesakhcienet Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

NuojCheng left a comment

Choose a reason for hiding this comment

Uh oh!

mesakhcienet Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

mesakhcienet Oct 14, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mesakhcienet commented Jul 15, 2025 •

edited

Loading

`SimpleDecodeLayer`

`SimpleMlpDecodeLayer`

mesakhcienet Oct 14, 2025 •

edited

Loading