[WIP] add deepseek-v3

[bzantium]

What does this PR do?

This PR adds the codes for the DeepSeekV3.
code relies heavily on original remote code.

resolved: #35425

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case: DeepSeek V3 Support #35425
• [] Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

to: @ArthurZucker

[Rocketknight1]

Hi @bzantium, this looks great so far! We'll need added tests for the model + a green CI, and then feel free to ping me to assign a reviewer, or if you have any problems with the port.

[ArthurZucker]

Ultra kudos! It's super nice
Mostly missing tests, here you can use a similar approach to the gemma2 tests, which use inheritance!

[cuichenx]

@bzantium Thanks for the amazing work! I was wondering if you were able to train V3 with FSDP? If so how many GPUs did you need? Thanks!

[ArthurZucker]

One big thing would be TP support, the base_tp_plan would probably need to be updated to make sure each mlp's gat up down have the correct order, unless the direct usage of dist remove this need

[casper-hansen]

This is great work and I'm looking forward to try it out. For multi-token prediction, is this planned to be implemented in this PR via the num_nextn_predict_layers attribute in the config?

[bzantium]

Thanks for the comments in detail; following your comments, I revised code quite a lot and fixed some mismatch between original code and this PR. I checked the outputs from both are the same. I think now I can add test codes. For TP support, I think they can be applied only for mlp layer but not for self_attn because they have functions like split on the hidden_dim. I added as following:

    base_model_tp_plan = {
        "layers.*.gate_proj": "colwise",
        "layers.*.up_proj": "colwise",
        "layers.*.down_proj": "rowwise",
    }

to: @ArthurZucker

[bzantium]

@bzantium Thanks for the amazing work! I was wondering if you were able to train V3 with FSDP? If so how many GPUs did you need? Thanks!

I did not try training yet, since this PR only supports for inference currently. I plan to add trainable code afterwards.

[bzantium]

This is great work and I'm looking forward to try it out. For multi-token prediction, is this planned to be implemented in this PR via the num_nextn_predict_layers attribute in the config?

thanks for attention, I don't have plan to implement multi-token prediction for this time since there is no additional parameters for additional layer which is required for multi-token prediction.

[casper-hansen]

This is great work and I'm looking forward to try it out. For multi-token prediction, is this planned to be implemented in this PR via the num_nextn_predict_layers attribute in the config?

thanks for attention, I don't have plan to implement multi-token prediction for this time since there is no additional parameters for additional layer which is required for multi-token prediction.

The weights are provided. DeepSeek provided a neat description of this. Since num_nextn_predict_layers=1 and num_hidden_layers=61, the MTP module is layer 61.

[bzantium]

This is great work and I'm looking forward to try it out. For multi-token prediction, is this planned to be implemented in this PR via the num_nextn_predict_layers attribute in the config?

thanks for attention, I don't have plan to implement multi-token prediction for this time since there is no additional parameters for additional layer which is required for multi-token prediction.

The weights are provided. DeepSeek provided a neat description of this. Since num_nextn_predict_layers=1 and num_hidden_layers=61, the MTP module is layer 61.

I missed this! thanks for the reference, let me check how to deal with it.

[ArthurZucker]

Yep a lot better! For training you can also have a look at switch transformers and granite MOE where we do have losses and etc

[ArthurZucker]

I think query keys and values also need to appear here no? 🤗

[bzantium]

I'm not sure how TP is implemented here but multi latent attention implementation makes it hard to apply TP because they use lora_a -> RMSNorm -> lora_b -> split and so forth.

q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))).view(hidden_shape).transpose(1, 2)
q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
k_pe = k_pe.view(*input_shape, 1, self.qk_rope_head_dim).transpose(1, 2)
kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(hidden_shape).transpose(1, 2)

does colwise_rep do all_gather after mm operation?

[ArthurZucker]

nice, IDK how many experts are actually there, but given how important the model is, we might want to use more optimized version of MoE implementation. The best we have in transformers is for SwitchTransformers ! Were we skip experts that won't have any tokens.

One thing is also that this is not compile compatible "yet"?

[bzantium]

I will check this out! thanks :)

[ArthurZucker]

Do you want me to jump on the PR and help you merge this faster @bzantium ? 🤗

[bzantium]

Do you want me to jump on the PR and help you merge this faster @bzantium ? 🤗

Of course! As you commented, it looks like there's still a lot to work left. (code optimization, training code, testing code and multi-token prediction if possible)
to: @ArthurZucker