[Performance][Model Loader] Skip non-local expert weights during EP model loading

[esmeetu]

Purpose

In DP+EP deployments, every rank currently reads all expert weights from disk via safe_open().get_tensor(), only for FusedMoE.weight_loader to discard non-local experts afterward. For large MoE models (e.g. Kimi-K2.5 with 384 experts at 591GB), each rank reads the full 591GB but only keeps ~144GB (dense + 1/8 experts).

This PR moves the filtering before the disk read by checking the tensor name against the local expert set in the weight iterator, so f.get_tensor() is never called for non-local experts.

• Skip reading non-local expert weights from disk during model loading when expert parallelism (EP) is enabled
• Each rank only reads its own expert shard + shared dense weights, avoiding ~87% of storage I/O for typical MoE models
• No change to non-EP or non-MoE loading paths

Test Plan

• Verify on MoE model with EP enabled: only local expert weights loaded, model output unchanged
• Verify on 3D MoE model: gpt-oss-120b
• Verify non-MoE model loading is unaffected (local_expert_ids remains None)
• Verify EP=1 (no filtering) produces identical results to baseline
• Test with different EP sizes (8, 16, 24) to confirm correct expert distribution

Benchmark Results

Kimi-K2.5-NVFP4 (591GB, 384 experts):

DP/EP=4 (1 node × 4 GPUs)

Metric	main	PR	Speedup
Loading Weights	96.58s	67.59s	1.4x
Model Loading Total	103.53s	72.89s	1.4x

DP/EP=8 (2 nodes × 4 GPUs)

Metric	main	PR	Speedup
Loading Weights	58.45s	25.35s	2.3x
Model Loading Total	63.42s	32.04s	2.0x

DP/EP=16 (4 nodes × 4 GPUs)

Metric	main	PR	Speedup
Loading Weights	54.84s	21.71s	2.5x
Model Loading Total	83.46s	39.71s	2.2x

Note on warm vs cold cache: The numbers above were measured with OS page cache already populated (repeated runs). On cold start (first load after boot, deployment, or scale-up), the speedup is expected to be
larger because the dominant cost shifts from CPU-side mmap fault handling to network filesystem I/O, where reducing read volume from 591GB to ~144GB per rank (EP=8) has a proportionally greater effect. Cold
cache testing was not performed but can be reproduced by running sync && echo 3 > /proc/sys/vm/drop_caches (requires root) before the first load.

Speedup scales with EP size: higher EP → more experts skipped → greater I/O reduction.

EP size	Experts/rank	Per-rank I/O	Expert I/O skipped
EP=4	96	~208GB	75%
EP=8	48	~144GB	87.5%
EP=16	24	~112GB	93.75%
EP=24	16	~101GB	95.8%

---

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.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request introduces a significant performance optimization for loading Mixture-of-Experts (MoE) models with expert parallelism (EP) enabled. By filtering out non-local expert weights before they are read from disk, it effectively reduces I/O, which is particularly beneficial for large models. The implementation is well-structured, with the core filtering logic encapsulated in a new ep_weight_filter.py module and accompanied by a comprehensive test suite. The changes are correctly integrated into the model loading pipeline. I've found one critical issue regarding the calculation of expert parallelism rank and size, which omits the prefill context parallelism dimension. Addressing this will ensure the correctness of expert filtering in all configurations.

[mergify]

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

uv pip install pre-commit
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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 90784738f0

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[mergify]

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

uv pip install pre-commit
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