[MoE Refactor] Add sequence parallel tests to test_moe_layer.py

[bnellnm]

Purpose

Add sequence parallel tests to test_moe_layer.py. Unfortunately, they require 4 gpus so I'm not sure we can run them in automation. At least they can be run locally when validating changes to MoE.

cc @yzong-rh

Test Plan

Ran by hand with 4 gpus.

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gemini-code-assist]

Code Review

This pull request introduces support for sequence parallelism (SP) in MoE layer tests. Key changes include the addition of an sp_wrapper to handle sequence chunking and gathering, updates to test configurations and validation logic to support SP, and adjustments to weight chunking behavior when SP is enabled. Review feedback highlights a potential issue in the calculation of num_tokens_across_dp for sequence parallel configurations, which could lead to incorrect communication or out-of-bounds access. Additionally, there is a concern regarding the significant increase in FP8 quantization tolerance, which may reduce the effectiveness of correctness checks.

[mergify]

Hi @bnellnm, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[yzong-rh]

Also passed for me on 4xB200 on 0.20.2rc1.dev56+gf65376125.precompiled after merging with commit f6537612521df1156d8ac13f524e427aca908322 on main.

[yzong-rh]

Note: tolerance changed.

[mergify]

Hi @bnellnm, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

Hi @bnellnm, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10