Te/main DIFF#2

Closed

phu0ngng wants to merge 80 commits into

mainfrom

Owner

phu0ngng commented Oct 8, 2025

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

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.

jberchtold-nvidia and others added 30 commits

August 12, 2025 12:29


          Quantizations backup

f196660

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>


          Quantizations dot_general integration

44b413e


          Support TE checkpointing policies

1413c26


          Add support for TE fused layernorm_mlp

213a639


          Documentation for TransformerEngineQuantization

f629134


          Fix dummy initializer shapes to support divisibility by 32 for MXFP8

b102e65


          Avoid fp8_autocast in integration and instead specify recipe directly…

7c90533

… when creating quantizer sets


          Cast gamma to input dtype to fix issue with TE norm primitive

33efed9


          Update recipe arg name to fp8_recipe

816e025


          Support block size function on quantization utility to support MXFP8 …

f31d488

…init


          Remove test-maxtext.sh

411effb


          Lock TE pre-attention norm behind an env var

3dc3de7


          Merge pull request AI-Hypercomputer#5 from jberchtold-nvidia/jberchto…

4447f68

…ld/te_integration

Reworked MaxText+TE integration


          add run_single_node script

8eceb43

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          print percentage for normalized

22449a4

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          update test-maxtext-te.sh

410d797

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          clean up traces in test-maxtext-te.sh

12c875a

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          Merge pull request AI-Hypercomputer#7 from phu0ngng/patch/te_integrat…

cb6d584

…ion_v2

Add test scripts to v2 integration


          fix normalize baseline to 0%

8a39a36

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          Fix normalize baseline to 0%

6cc84e8

Fix normalize baseline to 0%


          Name TEWrapper module in integration with suffix for specific op

030d547


          Merge pull request AI-Hypercomputer#9 from jberchtold-nvidia/dev/jber…

…chtold/specific-te-wrapper-names

More specific te wrapper names


          if profile, run 1 decoder_layer

ec12df0

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          add trace to output-tag

89317ff

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          Plot loss curves helper script

15a1061


          modify config file for tracing

614a9f5

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          Merge pull request AI-Hypercomputer#11 from jberchtold-nvidia/dev/jbe…

a621aa0

…rchtold/plot-loss-curves-script

Plot loss curves script


          clean up

e378cd7

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          Raise error when dropout used in TE MLP

36dbcb5

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>


          Merge pull request AI-Hypercomputer#13 from jberchtold-nvidia/dev/jbe…

b8da497

…rchtold/error-when-dropout-used-in-mlp

Raise error when dropout used in TE mlp

phu0ngng and others added 26 commits

October 8, 2025 10:58


          minor fixes for fused_mlp

394d056

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          added cs to the test script

10a6cc9

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          cleanup

a924639

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          merge with upstream

27fb574

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          move te to gpu/te

9c9256a

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          TE Dense op integration

0b85c44

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Jeremy Berchtold <jberchtold@nvidia.com>


          Merge pull request AI-Hypercomputer#24 from phu0ngng/te/main

cbb4dfa

Merge with upstream + minor changes


          added tpsp to mesh_resource

e03152a

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          added missing license

5a6aab0

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          added nvfp4

4d19d6e

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          added more training tests

57ae109

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          Merge remote-tracking branch 'official-upstream/main' into te/main

00eb792


          Fix use_te in gemma3

05fb74c

Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>


          Merge pull request AI-Hypercomputer#27 from jberchtold-nvidia/te/main

0f1ef19

Merge in upstream MaxText


          fetch changes from the TE Dense PR

f066a94

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          TE Dense op integration

18ed99f

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
Co-authored-by: Jeremy Berchtold <jberchtold@nvidia.com>


          added tpsp to mesh_resource

88659a8

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          added missing license

d9810d3

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          added nvfp4

3b440fe

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          added more training tests

f50fad6

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          rebase

a7c02c0

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          cleanup

d819e67

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          minor fix

7fd658e

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          cleanup

d8f6a49

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          minor fix

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


          add comment

1560b4d

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

github-actions Bot commented Dec 19, 2025

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 Bot added the stale label

github-actions Bot commented Dec 27, 2025

This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it.

github-actions Bot closed this

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

Labels