Enabled auxilillary loss free load balancing and sequence wise load b…#4233
Enabled auxilillary loss free load balancing and sequence wise load b…#4233dipakg-lang wants to merge 1 commit into
Conversation
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
6d108ec to
bce41ae
Compare
bce41ae to
e25c23c
Compare
parambole
left a comment
There was a problem hiding this comment.
Hey @dipakg-lang, on a high level, the changes look good.
a) Can you add test cases for the dsv4-like structure?
b) Also, could you add the bug number for sequence load balancing loss?
|
e25c23c to
b026bf7
Compare
parambole
left a comment
There was a problem hiding this comment.
LGTM. Thank you for adding these changes.
a15c86a to
56f6962
Compare
|
Note: This is a fork repo, so unfortunately Gemini review doesn't work due to the API key thing. |
|
Thanks for the change! I have a few high level comments:
|
| adamw_mask: [".*gate.*bias.*"] | ||
|
|
||
| # --- Attention configuration --- | ||
| attention: 'dot_product' |
There was a problem hiding this comment.
We usually just put model configs in this yml file. The experimental config (like dot_product or flash, etc) usually is passed from cmd.
There was a problem hiding this comment.
okey, would you suggest we skip them and have the user pass it instead? as we only support the dot product right now, I felt it would be safer to have that value present here and later change when we remove it from there when we start supporting more options, thoughts?
There was a problem hiding this comment.
In this case, we commonly put assertion like this to indicate which case is not supported:
maxtext/src/maxtext/configs/types.py
Lines 2567 to 2572 in 7d8d169
There was a problem hiding this comment.
This already addressed in PR4153
- This
attention: 'dot_product'has been removed from deepseek4-284b.yml - types check is already added
maxtext/src/maxtext/configs/types.py
Lines 3135 to 3136 in 8e7ff2d
Please remove attention: 'dot_product' here as well.
There was a problem hiding this comment.
Removed dot produce here as well
|
|
||
| # Tiny model config for DeepSeek V4 for CPU execution and testing | ||
|
|
||
| base_emb_dim: 64 |
There was a problem hiding this comment.
Do you know why those model configs (if from v4) are not included in the 284B model? if using default value from function, it's risky - https://screenshot.googleplex.com/4na5Jwtha49zfZM
There was a problem hiding this comment.
adding @parambole to this conversation as well who has added the 284B model, so that we can resolve this.
There was a problem hiding this comment.
I think deepseek4-284b.yml from main branch is the up-to-date version. In particular, it has been updated in PR4153 based on review.
Should make this tiny version consistent with 284b when possible.
There was a problem hiding this comment.
Made consistent with 284b perfectly only scaled down version of it.
There was a problem hiding this comment.
Thanks! Didn't see qk_rope_head_dim in 284-version configs?
| pre_bias_logits = output | ||
|
|
||
| if self.use_bias: | ||
| # Architectural Note: Bias is an nnx.Param rather than nnx.Variable due to Linen/NNX state |
There was a problem hiding this comment.
I think we have turned on NNX by default. Shall we actually use nnx.variable for a cleaner solution? It seems nnx.param + stop_gradient is a shortcut/hack. Using nnx.variable, I would expect we don't need adamw_mask: [".*gate.*bias.*"]. Please let me know your thoughts!
cc @shuningjin involved discussion before
There was a problem hiding this comment.
Yes, I had the same suggestion earlier, but then we will have to manage the overhead of managing this variable ourselves that is why we have gone through this route, please let me know if you strongly feel about going through the variable route then I will have to create a plan on this variable management.
There was a problem hiding this comment.
we will have to manage the overhead of managing this variable
Oh I didn't realize that. I thought we were thinking of nnx.variable in this thread. Please let me know if I missed something.
Gemini says "Yes, the DeepSeek router bias is an absolute textbook use case for nnx.Variable"
There was a problem hiding this comment.
I would prefer nnx.variable to avoid manipulation of optimizer mask. Could you elaborate what is the "overhead of managing this variable"?
There was a problem hiding this comment.
Declaring the bias as an nnx.Variable instead of an nnx.Param introduces framework overhead because JAX requires custom state to be manually threaded through the train_step pipeline
(passed as inputs and returned as outputs) on every iteration. Additionally, because MaxText's sharding and checkpointing are built for the params tree, using a new state collection
would require writing custom PartitionSpec logic to distribute the bias across the mesh and custom serialization code to save and restore it. This is what motivated me to do design it in this way, otherwise if you remember I was the one pointing out need for this to be a variable. Please let me know if you would like to have discussion on this.
| lambda params: jax.tree_util.tree_map(lambda x: "frozen" if x else "trainable", freeze_mask_fn(params)), | ||
| ) | ||
|
|
||
| if getattr(config, "routed_bias", False) and getattr(config, "routed_bias_update_rate", 0.0) > 0.0: |
There was a problem hiding this comment.
This is needed if using nnx.variable?
There was a problem hiding this comment.
No, this part would be saved then. but we will have own variable management separately if not param
|
009d8da to
ae1fcba
Compare
|
There was a problem hiding this comment.
Thanks for the feature! I agree with @RissyRan in many aspects.
- More testing with Deepseek4 & Deepseek3
- loss-free lb - Routed bias implementation:
- it currently uses:
nnx.Param+ stop gradient + optimizer mask for adamw and muon - Can we simplify via
nnx.Variableto avoid optimizer mask?
- (additionally) loss-free lb - Routed bias update:
- in
train.pyuses complicated logic based on path. - Can we simplify?
For lb loss, perhaps the existing load_balance_loss is already sequence-wise. I added some notes in b/521990776
| pre_bias_logits = output | ||
|
|
||
| if self.use_bias: | ||
| # Architectural Note: Bias is an nnx.Param rather than nnx.Variable due to Linen/NNX state |
There was a problem hiding this comment.
I would prefer nnx.variable to avoid manipulation of optimizer mask. Could you elaborate what is the "overhead of managing this variable"?
|
|
||
| # Tiny model config for DeepSeek V4 for CPU execution and testing | ||
|
|
||
| base_emb_dim: 64 |
There was a problem hiding this comment.
I think deepseek4-284b.yml from main branch is the up-to-date version. In particular, it has been updated in PR4153 based on review.
Should make this tiny version consistent with 284b when possible.
| adamw_mask: [".*gate.*bias.*"] | ||
|
|
||
| # --- Attention configuration --- | ||
| attention: 'dot_product' |
There was a problem hiding this comment.
This already addressed in PR4153
- This
attention: 'dot_product'has been removed from deepseek4-284b.yml - types check is already added
maxtext/src/maxtext/configs/types.py
Lines 3135 to 3136 in 8e7ff2d
Please remove attention: 'dot_product' here as well.
| # Updates the shape to be aligned with state. | ||
| moe_bias_updates = jnp.array(moe_bias_updates[0]).transpose() | ||
| new_state = maxtext_utils.update_state_param(new_state, target_path, moe_bias_updates) | ||
| # Apply updates for Auxiliary-Loss-Free load balancing for the DeepSeek family. |
There was a problem hiding this comment.
The logic of routed bias update (in loss-free load balancing) seems really complicated. Is it possible to simplify?
There was a problem hiding this comment.
I think you are talking about this code which is about locating the bias in the tree, considering the nature of path of the bias, we have to have logic to locate it in such way that we don't miss it and the code continue to work for any other models as well is resulting in this code, if you have any ideas we can always incorporate them!
… balancing 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
ae1fcba to
cf88282
Compare
Description
Enabled auxilillary loss free load balancing and sequence wise load balancing for Deepseek.
The rest of the description includes relevant details and context, examples:
We require these changes to be compliant with the specification and also this will help us to have better use to hardware due to the load balancing between the experts.
FIXES: b/509933890
FIXES: b/521990776
Tests
Tested by running training loop with new tiny Deeepseek V4 model added as part of the commit,
here are the logs for testing and commands used for this :
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 :
=== DeepSeek V4 Load Balancing Variance Analysis (Step 0 vs Step 20) ===
Raw data collected for this analysis:
https://paste.googleplex.com/5060754624610304
https://paste.googleplex.com/5473518849490944
Checklist
Before submitting this PR, please make sure (put X in square brackets):