Skip to content

Dsv4 load balancing#4190

Open
dipakg-lang wants to merge 2 commits into
AI-Hypercomputer:dsv4_model_integratefrom
dipakg-lang:dsv4_load_balancing
Open

Dsv4 load balancing#4190
dipakg-lang wants to merge 2 commits into
AI-Hypercomputer:dsv4_model_integratefrom
dipakg-lang:dsv4_load_balancing

Conversation

@dipakg-lang

@dipakg-lang dipakg-lang commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

  • why is this change being made,
  • the problem being solved and any relevant context,
  • why this is a good solution,
  • some information about the specific implementation,
  • shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@dipakg-lang dipakg-lang force-pushed the dsv4_load_balancing branch from 663c5f6 to 8b75018 Compare June 17, 2026 22:23
@codecov

codecov Bot commented Jun 17, 2026 •
edited
Loading

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 69.38776% with 15 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/maxtext/trainers/pre_train/train.py 62.16% 13 Missing and 1 partial ⚠️
src/maxtext/common/metric_logger.py 0.00% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@dipakg-lang dipakg-lang force-pushed the dsv4_load_balancing branch from 8b75018 to 128daba Compare June 17, 2026 22:43
@parambole parambole force-pushed the dsv4_model_integrate branch from 631941c to 9a571d5 Compare June 17, 2026 23:19
@dipakg-lang
Enabled auxilillary loss free load balancing and sequence wise load b…
cd0a54e 
…alancing for Deepseek.

Tested by running training loop with new tiny Deeepseek V4 model added as part of the commit,
here are the logs for testing

Without load balancing active logs : https://paste.googleplex.com/6421399878107136

with load balancing logs : https://paste.googleplex.com/6551357300539392

Here are the results actived for reducing the varience :

1 === DeepSeek V4 Load Balancing Variance Analysis (Step 0 vs Step 20) ===
    2
    3 | Layer Index | Routing Type | Step 0 Var (Baseline) | Step 20 Var (Run A) | Step 20 Var (Run B) | Improvement (A vs B) |
    4 |-------------|--------------|-----------------------|---------------------|---------------------|----------------------|
    5 |           0 | Hash Routed  |            3932160.00 |          3932160.00 |          3932160.00 |                0.00% |
    6 |           1 | Hash Routed  |            3932160.00 |          3932160.00 |          3932160.00 |                0.00% |
    7 |           2 | Hash Routed  |            3932160.00 |          3932160.00 |          3932160.00 |                0.00% |
    8 |           3 | Top-K Routed |               7409.38 |             7509.25 |             3672.12 |               51.10% |
    9 |           4 | Top-K Routed |               3158.38 |             3230.12 |             1216.00 |               62.35% |
   10 |           5 | Top-K Routed |               5713.38 |             5772.75 |             2359.38 |               59.13% |
   11 |           6 | Top-K Routed |               8295.25 |             8082.50 |             3674.12 |               54.54% |
   12 |           7 | Top-K Routed |               4765.62 |             4614.62 |             1212.75 |               73.72% |
   13 |           8 | Top-K Routed |               4960.75 |             4923.12 |             1663.50 |               66.21% |
   14 |           9 | Top-K Routed |               3905.50 |             3816.25 |             1316.88 |               65.49% |
   15 |          10 | Top-K Routed |               5057.00 |             4981.12 |             2257.75 |               54.67% |
   16 |          11 | Top-K Routed |              10446.62 |            10381.62 |             5565.75 |               46.39% |
   17 |          12 | Top-K Routed |               9538.50 |             9529.25 |             5319.12 |               44.18% |
   18 |          13 | Top-K Routed |               7031.38 |             7131.25 |             3270.25 |               54.14% |
   19 |          14 | Top-K Routed |               4852.00 |             4900.12 |             1906.88 |               61.09% |
   20 |          15 | Top-K Routed |               9306.12 |             9342.88 |             4733.75 |               49.33% |
   21 |          16 | Top-K Routed |               5811.25 |             5749.50 |             2110.88 |               63.29% |
   22 |          17 | Top-K Routed |               6715.62 |             6874.25 |             2664.12 |               61.24% |
   23 |          18 | Top-K Routed |               8145.50 |             7869.25 |             3383.75 |               57.00% |
   24 |          19 | Top-K Routed |               6042.12 |             5908.62 |             2353.00 |               60.18% |
   25 |          20 | Top-K Routed |               8559.88 |             8158.25 |             4333.38 |               46.88% |
   26 |          21 | Top-K Routed |              11742.25 |            11943.62 |             7563.50 |               36.67% |
   27 |          22 | Top-K Routed |               4959.62 |             5014.88 |             1998.62 |               60.15% |
   28 |          23 | Top-K Routed |               7717.12 |             7751.88 |             3879.88 |               49.95% |
   29 |          24 | Top-K Routed |               9017.75 |             9307.88 |             4702.75 |               49.48% |
   30 |          25 | Top-K Routed |              14127.12 |            14111.25 |             8079.25 |               42.75% |
   31 |          26 | Top-K Routed |               5074.25 |             5194.12 |             1675.50 |               67.74% |
   32 |          27 | Top-K Routed |              11919.50 |            11204.38 |             6470.75 |               42.25% |
   33 |          28 | Top-K Routed |              12241.75 |            12998.62 |             7624.12 |               41.35% |
   34 |          29 | Top-K Routed |               9384.50 |             9005.00 |             5052.00 |               43.90% |
   35 |          30 | Top-K Routed |               9698.62 |             9678.25 |             5231.75 |               45.94% |
   36 |          31 | Top-K Routed |              12244.25 |            12392.75 |             7249.25 |               41.50% |
   37 |          32 | Top-K Routed |              10030.00 |             9972.62 |             4755.50 |               52.31% |
   38 |          33 | Top-K Routed |               7265.00 |             6973.62 |             3271.75 |               53.08% |
   39 |          34 | Top-K Routed |              11945.50 |            11940.62 |             6076.88 |               49.11% |
   40 |          35 | Top-K Routed |              12917.50 |            13740.00 |             7210.62 |               47.52% |
   41 |          36 | Top-K Routed |              15011.62 |            15083.00 |             8870.62 |               41.19% |
   42 |          37 | Top-K Routed |              10294.12 |            10176.25 |             5907.50 |               41.95% |
   43 |          38 | Top-K Routed |               8928.62 |             9236.00 |             5136.62 |               44.38% |
   44 |          39 | Top-K Routed |              15633.62 |            15171.00 |             9684.75 |               36.16% |
   45 |          40 | Top-K Routed |               7687.75 |             7658.12 |             4521.25 |               40.96% |
   46 |          41 | Top-K Routed |              12485.12 |            12270.38 |             6933.25 |               43.50% |
   47 |          42 | Top-K Routed |              17641.25 |            17163.50 |            10974.12 |               36.06% |
   48 |-------------|--------------|-----------------------|---------------------|---------------------|----------------------|
   49 | TOTAL/AVG   | Top-K Only   |             357681.12 |           356762.50 |           185883.62 |               47.90% |

Raw data collected for this analysis:

https://paste.googleplex.com/5060754624610304

https://paste.googleplex.com/5473518849490944
@dipakg-lang dipakg-lang force-pushed the dsv4_load_balancing branch from 128daba to cd0a54e Compare June 18, 2026 00:00
@parambole parambole force-pushed the dsv4_model_integrate branch 6 times, most recently from 228323c to 48bdd7a Compare June 18, 2026 05:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit is a code owner

@jshin1394 jshin1394 Awaiting requested review from jshin1394 jshin1394 is a code owner

@liudangyi liudangyi Awaiting requested review from liudangyi liudangyi is a code owner

@richjames0 richjames0 Awaiting requested review from richjames0 richjames0 is a code owner

@shralex shralex Awaiting requested review from shralex shralex is a code owner

@darisoy darisoy Awaiting requested review from darisoy darisoy is a code owner

@RissyRan RissyRan Awaiting requested review from RissyRan RissyRan is a code owner

@michelle-yooh michelle-yooh Awaiting requested review from michelle-yooh michelle-yooh is a code owner

@gagika gagika Awaiting requested review from gagika gagika is a code owner

@suexu1025 suexu1025 Awaiting requested review from suexu1025 suexu1025 is a code owner

@parambole parambole Awaiting requested review from parambole parambole is a code owner

@shuningjin shuningjin Awaiting requested review from shuningjin shuningjin is a code owner

@jiangjy1982 jiangjy1982 Awaiting requested review from jiangjy1982 jiangjy1982 is a code owner

@gobbleturk gobbleturk Awaiting requested review from gobbleturk gobbleturk is a code owner

@bvandermoon bvandermoon Awaiting requested review from bvandermoon bvandermoon is a code owner

@NicoGrande NicoGrande Awaiting requested review from NicoGrande NicoGrande is a code owner

@NuojCheng NuojCheng Awaiting requested review from NuojCheng NuojCheng is a code owner

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo is a code owner

@abhinavclemson abhinavclemson Awaiting requested review from abhinavclemson abhinavclemson is a code owner

@vipannalla vipannalla Awaiting requested review from vipannalla vipannalla is a code owner

@SurbhiJainUSC SurbhiJainUSC Awaiting requested review from SurbhiJainUSC SurbhiJainUSC is a code owner

@A9isha A9isha Awaiting requested review from A9isha A9isha is a code owner

@aireenmei aireenmei Awaiting requested review from aireenmei aireenmei is a code owner

@dipannita08 dipannita08 Awaiting requested review from dipannita08 dipannita08 is a code owner

@igorts-git igorts-git Awaiting requested review from igorts-git igorts-git is a code owner

Assignees

No one assigned

Labels

gemini-review

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@dipakg-lang @parambole