Skip to content

[JAX] Remove shard_map from MoEBlock to support quant before FSDP AG using Grouped quant+GEMM custom partitioning rules#3131

Draft
jberchtold-nvidia wants to merge 38 commits into
NVIDIA:mainfrom
jberchtold-nvidia:jberchtold/quant-before-fsdp-ag-gmm
Draft

[JAX] Remove shard_map from MoEBlock to support quant before FSDP AG using Grouped quant+GEMM custom partitioning rules#3131
jberchtold-nvidia wants to merge 38 commits into
NVIDIA:mainfrom
jberchtold-nvidia:jberchtold/quant-before-fsdp-ag-gmm

Use 2D output from grouped quant and support GSPMD partition rules

766c36a
Select commit
Loading
Failed to load commit list.
DCO / DCO required action Jun 15, 2026 in 1s

DCO

There are 7 commits incorrectly signed off. This means that the author(s) of these commits failed to include a Signed-off-by line in their commit message.

To avoid having PRs blocked in the future, always include Signed-off-by: Author Name <authoremail@example.com> in every commit message. You can also do this automatically by using the -s flag (i.e., git commit -s).

Here is how to fix the problem so that this code can be merged.


Rebase the branch

If you have a local git environment and meet the criteria below, one option is to rebase the branch and add your Signed-off-by lines in the new commits. Please note that if others have already begun work based upon the commits in this branch, this solution will rewrite history and may cause serious issues for collaborators (described in the git documentation under "The Perils of Rebasing").

You should only do this if:

  • You are the only author of the commits in this branch
  • You are absolutely certain nobody else is doing any work based upon this branch
  • There are no empty commits in the branch (for example, a DCO Remediation Commit which was added using --allow-empty)

To add your Signed-off-by line to every commit in this branch:

  1. Ensure you have a local copy of your branch by checking out the pull request locally via command line.
  2. In your local branch, run: git rebase HEAD~38 --signoff
  3. Force push your changes to overwrite the branch: git push --force-with-lease origin jberchtold/quant-before-fsdp-ag-gmm

Summary

Commit sha: b14df81, Author: Jeremy Berchtold, Committer: Jeremy Berchtold; The sign-off is missing.
Commit sha: 60a0b50, Author: Jeremy Berchtold, Committer: Jeremy Berchtold; The sign-off is missing.
Commit sha: 786fa1d, Author: Jeremy Berchtold, Committer: Jeremy Berchtold; The sign-off is missing.
Commit sha: 0eee8b8, Author: tdophung, Committer: tdophung; The sign-off is missing.
Commit sha: 776c5ef, Author: tdophung, Committer: tdophung; The sign-off is missing.
Commit sha: acb610f, Author: tdophung, Committer: tdophung; The sign-off is missing.
Commit sha: 458d1c4, Author: tdophung, Committer: tdophung; The sign-off is missing.