Add support for SPMD PP for deepseek decoder block #1687
Conversation
gobbleturk
requested review from
khatwanimohit,
bvandermoon,
vipannalla,
RissyRan,
richjames0,
rni418,
gagika,
shralex,
yangyuwei,
SurbhiJainUSC,
hengtaoguo,
A9isha,
wang2yn84,
wyzhang,
mitalisi,
gpolovets1 and
mailvijayasingh
as code owners
gobbleturk
force-pushed
the
mattdavidow-pipeline-deepseek
branch
from
23fd9a9 to
03e1524
Compare
gobbleturk
force-pushed
the
mattdavidow-pipeline-deepseek
branch
from
03e1524 to
ee35916
Compare
we can pipeline 58 but the PP rank must divide the number of layers - 58 isn't a particularly divisor friendly number, I updated this in PR description |
gobbleturk
force-pushed
the
mattdavidow-pipeline-deepseek
branch
from
ee35916 to
46bbf6d
Compare
| y, _ = self.scan_decoder_layers(cfg, dense_layer, cfg.first_num_dense_layers, "dense_layers", mesh)( | ||
| y, | ||
| decoder_segment_ids, | ||
| decoder_positions, | ||
| deterministic, | ||
| model_mode, | ||
| ) | ||
| if num_moe_layers_outside_pp > 0: | ||
| y, _ = self.scan_decoder_layers(cfg, moe_layer, num_moe_layers_outside_pp, "moe_layers", mesh)( | ||
| y, | ||
| decoder_segment_ids, | ||
| decoder_positions, | ||
| deterministic, | ||
| model_mode, | ||
| ) | ||
| y = self.pipeline_module(y, decoder_segment_ids, decoder_positions, deterministic, model_mode, partition_spec=partition_spec) |
There was a problem hiding this comment.
Just for my understanding, how are the 56 layers that need to be pipelined being specified here? Is it just because they aren't being used in self.scan_decoder_layers?
There was a problem hiding this comment.
56 was an example, it can be set via pipeline_parallel_layers which was added in a previous PR
gobbleturk
force-pushed
the
mattdavidow-pipeline-deepseek
branch
from
46bbf6d to
14300d9
Compare
There was a problem hiding this comment.
Thanks Matt! One more comment, can we add a unit test in moe_test.py like [this](
maxtext/MaxText/tests/moe_test.py
Line 423 in 0add45d
| ) | ||
| remaining_layers = self.config.num_decoder_layers - self.config.pipeline_parallel_layers | ||
| if remaining_layers > 0: | ||
| logical_axis_rules_pp_as_dp = maxtext_utils.logical_axis_rules_pp_act_as_dp(self.config.logical_axis_rules) |
There was a problem hiding this comment.
Do you think we could add some flexibility here? Instead of act as dp, could we act based on passed in config (not sure if doable). DP may not be helpful.
For instance, for DS-v3 (61 layers with 58 MoE layers) with pipeline_parallel_layers=56, we could run FSDP or EP from config bellow?
pipeline_parallel_layers=56 && ici_fsdp=-1/ici_ep=-1?
There was a problem hiding this comment.
with pipeline_parallel_layers=56 there are only 2 moe layers that will be treated with PP replaced with DP. These two layers will still be sharded by the other sharding strategies - e.g. if config was EP_ICI=16 FSDP_ICI=16 PP_DCN=16, these two layers will be sharded as EP_ICI=16, FSDP_ICI=16 DP_DCN=16, e.g. still the weights are sharded 256 ways
I do like the idea of flexibility, but perhaps we can save this as potential future work?
There was a problem hiding this comment.
these two layers will be sharded as EP_ICI=16, FSDP_ICI=16 DP_DCN=16, e.g. still the weights are sharded 256 ways
SG! I thought it will only have DP sharding.
Description
Add support for using PP with deepseek, including with the new feature
pipeline_parallel_layerswhich only pipelines a subset of layers. This change can help out with SPMD pipelining since PP must divide the number of layers. E.g. for deepseek there are 58 sparse layers, which does not have many friendly divisors.With this PR we can pipeline just a subset of the sparse layers, e.g. 56 of them and set PP=8. Other layers will be sharded like DP.
Tests
Ran some locally of smaller v2-16B model and added AOT test - I will paste XPROFS in a bit
Checklist
Before submitting this PR, please make sure (put X in square brackets):