Feat/Decoder migration: DeepSeek/Gemma3/Gemma4/Llama4

[hsuan-lun-chiang]

Description

Implement and update the following models in NNX decoder that were are not supported in previous PR 2831:

• DeepSeek
• Gemma 3, 4
• Llama4

Tests

Test with different model and compare with Linen training. Details in the GDoc file

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.

[codecov]

Codecov Report

❌ Patch coverage is 31.34328% with 138 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	29.53%	119 Missing and 17 partials ⚠️
src/maxtext/layers/initializers.py	71.42%	1 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!

[RissyRan]

Thanks for the change! Before the review, could you ensure those are tested? Thank you!

cc @bvandermoon @entrpn

[hsuan-lun-chiang]

Thanks for the change! Before the review, could you ensure those are tested? Thank you!

cc @bvandermoon @entrpn

Hi @RissyRan, the Linen/NNX comparison logs for Llama4/Deepseek/Gemma4 are in the GDoc file. The PR is passing all unit tests with NNX flags (enable_nnx=True and pure_nnx_decoder=True), except for a few cases requiring further discussion, which we’ve documented here. Would you mind taking a look? Thank you!