Shuwen check vma#3411
Open
Shuwen-Fang wants to merge 40 commits intomainfrom
Open
Conversation
Shuwen-Fang
requested review from
A9isha,
NicoGrande,
NuojCheng,
RissyRan,
SurbhiJainUSC,
aireenmei,
bvandermoon,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
igorts-git,
jesselu-google,
jiangjy1982,
khatwanimohit,
michelle-yooh,
parambole,
richjames0,
shralex,
shuningjin,
suexu1025 and
vipannalla
as code owners
Shuwen-Fang
requested review from
Obliviour,
SujeethJinesh,
gpolovets1,
jacoguzo,
mailvijayasingh,
mitalisi,
notabee and
shauryagup
as code owners
Imported from GitHub PR #2831 Migrate the Transformer decoder layer into NNX. Note: The following models are currently not supported: - DeepSeek - Gemma3 - Llama4 Support for these models will be added in a follow-up PR. Strategy: A `pure_nnx_decoder` flag is added to control whether NNX or Linen decoder shall be used. Initial migration doesn't include the pipeline NNX support. Conducted these tests. Details in the [GDoc file](https://docs.google.com/document/d/1NbUP3g5glgbC6bMyt44pwM_vQA1NR7U2rBUzfbTDwSs/edit?pli=1&resourcekey=0-9EUahtzL-hCycdu7l0grhQ&tab=t.htq5367h8au0) 1. Test with different model and compare with Linen training 2. Golden logits comparison 3. Inference 4. Checkpoint comparison (Including TreeStructure Comparison) 5. Sharding comparison TODOs: - NNX version of unit tests (future PRs) Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests tests and provided workload links above if applicable. - [x] 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](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files). Copybara import of the project: -- 073e916 by hsuan-lun-chiang <hsuan-lun.chiang@cienet.com>: Migrate Decoder to NNX Adding nnx_decoders.py in parallel with decoders.py 1. Dup and modifiy decoders.py on new file nnx_decoders.py 2. add new config pure_nnx_decoder to control if model will use NNXDecoder, default false for now 3. modify relative code to accomodate the change 4. add/modify unit test Merging this change closes #2831 COPYBARA_INTEGRATE_REVIEW=#2831 from CIeNET-International:feat/Migrate-Decoder-to-NNX 073e916 PiperOrigin-RevId: 884170982
…o add default tokenizer_path for default model
Imported from GitHub PR #3449 Move the `install_maxtext_extra_deps` deps directory to `dependencies` after `dependencies` was added to the PyPI package. This command still works/runs the expected installation: ``` uv pip install -e .[tpu] --resolution=lowest install_maxtext_tpu_github_deps install_maxtext_tpu_post_train_extra_deps ``` CI also builds this command Before submitting this PR, please make sure (put X in square brackets): - [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [x] I have necessary comments in my code, particularly in hard-to-understand areas. - [x] I have run end-to-end tests tests and provided workload links above if applicable. - [x] 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. Copybara import of the project: -- 277b08e by Branden Vandermoon <bvandermoon@google.com>: Move install_maxtext_extra_deps to dependencies directory Merging this change closes #3449 COPYBARA_INTEGRATE_REVIEW=#3449 from AI-Hypercomputer:bvandermoon-github-deps 277b08e PiperOrigin-RevId: 886463160
add qwen3-base to configs/types and checkpoint_conversion/param_mapping add qwen3-base configs to checkpoint_conversion/hf_model_configs pyink
PiperOrigin-RevId: 886627046
# Description FP8 path is still using tokamax internal backend APIs. The new `RaggedDotGroupSizes` was introduced ([pull3330](#3330)) for Tokamax public APIs in bf16 path, which broke FP8. # Tests Benchmarks were run internally. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [X] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [X] I have necessary comments in my code, particularly in hard-to-understand areas. - [X] I have run end-to-end tests tests and provided workload links above if applicable. - [X] 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](https://maxtext.readthedocs.io/en/latest/development.html#adding-new-documentation-files).
Shuwen-Fang
force-pushed
the
shuwen-check-vma
branch
from
8684153 to
8a2c5a3
Compare
|
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
12 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
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):
gemini-reviewlabel.