fix and test EPLB balancedness calculation

[jdebache]

Purpose

Fix EPLB balancedness logging to aggregate rank load within each MoE layer. The
current reduction uses the layer axis, so it can report perfect balance when one
EP rank receives all tokens in every layer.

Test Plan

pytest -q tests/distributed/test_eplb_algo.py

Test Result

The new regression case has equal token totals per layer with all tokens routed
to one rank. Before the fix, it produced avg_tokens=100 and max_tokens=100;
after the fix, it produces the expected avg_tokens=50 and max_tokens=200.

15 passed, 1 skipped

---

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
• Documentation updates are not required for this metrics-only fix.

[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.

[ilmarkov]

LGTM. Basically revival of #39178 without additional changes

[tlrmchlsmth]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83351 for commit 43d62a1d20d5.

[jdebache]

It is unclear to me how the failure in CI is related to my changes. It looks infra related, and this step seems to be intermittently failing according to the CI dashboard.
@tlrmchlsmth: if you agree, would it be okay for you to merge it?