Skip to content

DO NOT MERGE Clean check vma#3501

Open
Shuwen-Fang wants to merge 3 commits intomainfrom
clean-check-vma
Open

DO NOT MERGE Clean check vma#3501
Shuwen-Fang wants to merge 3 commits intomainfrom
clean-check-vma

Conversation

@Shuwen-Fang
Copy link
Copy Markdown
Collaborator

@Shuwen-Fang Shuwen-Fang commented Mar 25, 2026

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, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@Shuwen-Fang Shuwen-Fang changed the title Clean check vma DO NOT MERGE Clean check vma Mar 25, 2026
@codecov
Copy link
Copy Markdown

codecov Bot commented Mar 25, 2026 •
edited
Loading

Codecov Report

❌ Patch coverage is 63.63636% with 8 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/layers/moe.py 66.66% 5 Missing ⚠️
src/maxtext/kernels/megablox/ops.py 66.66% 0 Missing and 2 partials ⚠️
src/maxtext/utils/gradient_accumulation.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

Shuwen-Fang added a commit that referenced this pull request Mar 31, 2026
@Shuwen-Fang
PR #3501: DO NOT MERGE Clean check vma
ce40301 
Imported from GitHub PR #3501

# 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, 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:

--
ec1b98b by Shuwen Fang <shuwenf@google.com>:

moe formatting changes

--
7618024 by Shuwen Fang <shuwenf@google.com>:

check vma changes

--
41b53c9 by Shuwen Fang <shuwenf@google.com>:

check vma changes

Merging this change closes #3501

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3501 from AI-Hypercomputer:clean-check-vma 41b53c9
PiperOrigin-RevId: 889398062
@Shuwen-Fang Shuwen-Fang mentioned this pull request Mar 31, 2026
Closed
4 tasks
Shuwen-Fang added a commit that referenced this pull request Mar 31, 2026
@Shuwen-Fang
PR #3501: DO NOT MERGE Clean check vma
b6dc0fc 
Imported from GitHub PR #3501

# 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, 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:

--
ec1b98b by Shuwen Fang <shuwenf@google.com>:

moe formatting changes

--
7618024 by Shuwen Fang <shuwenf@google.com>:

check vma changes

--
41b53c9 by Shuwen Fang <shuwenf@google.com>:

check vma changes

Merging this change closes #3501

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3501 from AI-Hypercomputer:clean-check-vma 41b53c9
PiperOrigin-RevId: 889398062
Shuwen-Fang added a commit that referenced this pull request Mar 31, 2026
@Shuwen-Fang
PR #3501: DO NOT MERGE Clean check vma
fbb59f6 
Imported from GitHub PR #3501

# 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, 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:

--
ec1b98b by Shuwen Fang <shuwenf@google.com>:

moe formatting changes

--
7618024 by Shuwen Fang <shuwenf@google.com>:

check vma changes

--
41b53c9 by Shuwen Fang <shuwenf@google.com>:

check vma changes

Merging this change closes #3501

FUTURE_COPYBARA_INTEGRATE_REVIEW=#3501 from AI-Hypercomputer:clean-check-vma 41b53c9
PiperOrigin-RevId: 889398062
@github-actions
Copy link
Copy Markdown

github-actions Bot commented Apr 25, 2026

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 github-actions Bot added the stale Automatically applied to stale PRs. label Apr 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@michelle-yooh michelle-yooh Awaiting requested review from michelle-yooh michelle-yooh is a code owner

@gagika gagika Awaiting requested review from gagika gagika 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

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@jesselu-google jesselu-google Awaiting requested review from jesselu-google jesselu-google is a code owner

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla 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

@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner

@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

pull ready stale Automatically applied to stale PRs.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@Shuwen-Fang