Skip to content

Load balancing changes for Deepseek v4#4178

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

Load balancing changes for Deepseek v4#4178
dipakg-lang wants to merge 2 commits into
AI-Hypercomputer:dsv4_model_integratefrom
dipakg-lang:dsv4_model_integrate

Conversation

@dipakg-lang

@dipakg-lang dipakg-lang commented Jun 16, 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.

@parambole
feat(models): Integrate DeepSeek V4 architecture and routing
6deaacc 
This commit introduces full support for DeepSeek V4 by integrating its
compressed attention mechanisms, MoE routing, and architectural layers.

Key changes:
- Add `deepseek4.yml` configuration and `DeepSeek4DecoderLayer` implementation.
- Implement hybrid Hash Routing and Token Routing for MoE layers.
- Add prefix/suffix layer unrolling for non-uniform compression blocks.
- Fix Pydantic validation for base MLP dimensions.
- Bypass MLA instantiation in favor of native CompressedAttention (CSA/HCA).
@dipakg-lang dipakg-lang force-pushed the dsv4_model_integrate branch 4 times, most recently from 5b1fd33 to adc5670 Compare June 17, 2026 18:08
@dipakg-lang
Enabled auxilillary loss free load balancing and sequence wise load b…
6fdcb90 
…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_model_integrate branch from adc5670 to 6fdcb90 Compare June 17, 2026 18:09
@parambole parambole force-pushed the dsv4_model_integrate branch 18 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

@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

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

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

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

@darisoy darisoy Awaiting requested review from darisoy darisoy 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

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

@khatwanimohit khatwanimohit Awaiting requested review from khatwanimohit khatwanimohit 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

@hengtaoguo hengtaoguo Awaiting requested review from hengtaoguo hengtaoguo 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

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@dipakg-lang @parambole