[EPLB] Fix balancedness metric computation and add verbose reporting

[arpera]

Purpose

This PR fixes the balancedness metric computation in EPLB to correctly reflect the load imbalance per MoE layer.

Previously, balancedness was calculated by computing the average and max load across layers (dim=0) instead of across ranks (dim=1), and then summing those values.

This PR changes the balancedness metric to:

1. Compute the max/mean ratio of tokens per rank independently for each MoE layer.
2. Average these per-layer ratios to produce the final balancedness metric.

This change ensures the logged metric accurately represents the average severity of the bottleneck at each MoE layer.

Additional changes in this PR:

• Added a log_balancedness_verbose configuration flag (default: False). When enabled, EPLB logs a detailed multi-line report per logging interval, which includes a per-layer / per-rank token table to help debug expert routing.
• Documented the log_balancedness_verbose flag in docs/serving/expert_parallel_deployment.md.
• Updated the docs/serving/expert_parallel_deployment.md documentation to include previously added EPLB configuration fields that were missing from the docs:
  • log_balancedness_interval (originally introduced in #29499)
  • communicator (introduced in #33176)

Test Plan

Manual vLLM local launches to verify the logs.

Validation Result

Should not affect production performance since balancedness is calculated and printed only when the log_balancedness option is set in the vLLM config.

Example of the new verbose logging output:

---

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.

[mergify]

Hi @arpera, 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

[gemini-code-assist]

Code Review

This pull request refactors the balancedness calculation in eplb_state.py to use an average of per-layer ratios instead of a global ratio. The feedback suggests using float64 for these calculations to prevent potential precision loss as model depth or batch sizes increase.

[arpera]

Valid points. Would it be ok if I heavily refactor this logging function and instead avg_tokens, max_tokens, and balancedness introduce more information about each layer? For example, I'll print hottest and coldest experts on current step, tokens' distribution among ranks on each layer on one step, distribution of experts among ranks, etc. I mean there won't be any one aggregated metric (like balancedness now), but instead I'll output detailed info about step. @ilmarkov What do you think about this idea?

[tlrmchlsmth]

Valid points. Would it be ok if I heavily refactor this logging function and instead avg_tokens, max_tokens, and balancedness introduce more information about each layer? For example, I'll print hottest and coldest experts on current step, tokens' distribution among ranks on each layer on one step, distribution of experts among ranks, etc. I mean there won't be any one aggregated metric (like balancedness now), but instead I'll output detailed info about step. @ilmarkov What do you think about this idea?

I think that would be great to expose. Maybe guard it behind a verbosity that we set in the eplb config?

[arpera]

Maybe guard it behind a verbosity that we set in the eplb config?

Yes, great! Could you tell me please one more detail what is the name of this flag in the eplb config?

[tlrmchlsmth]

Could you tell me please one more detail what is the name of this flag in the eplb config?

OH, sorry for the confusion - I was suggesting you add a flag to the eplb config (verbose_logging?)

[arpera]

Yes, I was also thinking about this idea. Now work in progress, stay tuned.

[mergify]

Documentation preview: https://vllm--39178.org.readthedocs.build/en/39178/

[ilmarkov]

I see this is still in progress — apologies if some of my comments are on things you're already planning to change. Happy to re-review once you're ready.

Nice touch with the ANSI heat coloring. Note though that in practice most users won't see it: in production, vLLM logs are typically collected from containers (Docker, Kubernetes). I would suggest to consider saving verbose stats in some structured format (maybe not in stdout/stderr) to ease an analysis.

[vadiklyutiy]

Friendly ping

[arpera]

@ilmarkov @tlrmchlsmth, PR is ready for a final review. Please have a look

[ilmarkov]

Thanks for the update! I'd still insist on saving the verbose log into a file for easier analysis. Also added some comments on improving the logged info.

[mergify]

Documentation preview: https://vllm--39178.org.readthedocs.build/en/39178/

[mergify]

Hi @arpera, 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

[arpera]

@ilmarkov I presume if everything else is ok then let us approve this PR. Issue from mergify is wrong, see 41377

[ilmarkov]

Thank you for the update!

The non-verbose part changes look good to me. In the verbose, I'd need some changes if we still want this piece after expert_load_dump_dir introduction. Although, I don't see the case we manually check stderr given that we have machine-parsable dumping with the same frequency. Maybe, for single node debugging only.

@tlrmchlsmth what do you think?

[arpera]

@vadiklyutiy, can we merge this? I think I fixed all the issues and change is ready for merge now.