[tx] DeepseekV3 implementation #889
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces the JAX implementation for the DeepseekV3 model. The implementation is comprehensive and covers the model's unique features like Multi-Head Latent Attention and Mixture of Experts with shared experts. The code is well-structured.
My review focuses on a critical bug that will prevent the model from running, along with some suggestions to improve maintainability by reducing code duplication and avoiding magic numbers. Addressing these points will make the implementation more robust and easier to maintain.
| # Precompute RoPE frequencies | ||
| # qk_rope_head_dim = config.qk_rope_head_dim | ||
| # original_seq_len = getattr(config, "original_seq_len", config.max_position_embeddings) | ||
| # rope_factor = getattr(config, "rope_factor", 1.0) | ||
| # beta_fast = getattr(config, "beta_fast", 32) | ||
| # beta_slow = getattr(config, "beta_slow", 1) | ||
|
|
||
| # TODO: Swap out like llama's rope? | ||
| # self.freqs_cis = precompute_freqs_cis( | ||
| # dim=qk_rope_head_dim, | ||
| # max_seq_len=config.max_position_embeddings, | ||
| # rope_theta=config.rope_theta, | ||
| # original_seq_len=original_seq_len, | ||
| # rope_factor=rope_factor, | ||
| # beta_fast=beta_fast, | ||
| # beta_slow=beta_slow, | ||
| # ) |
There was a problem hiding this comment.
This block for precomputing RoPE frequencies is commented out, but self.freqs_cis is used in DeepseekV3Model.__call__ at line 571. This will raise an AttributeError at runtime.
Looking at the DeepseekV3MLA implementation, the freqs_cis parameter is not used. Instead, apply_rope is called, which computes the frequencies on the fly.
To fix this, you should remove the freqs_cis parameter from the entire call chain, as it appears to be unused. This involves:
- Removing
freqs_cis: jax.Arrayfrom the signature ofDeepseekV3MLA.__call__. - Removing
freqs_cis: jax.Arrayfrom the signature ofDeepseekV3DecoderLayer.__call__. - Removing the
freqs_cis=self.freqs_cisargument from thelayer()call withinDeepseekV3Model.__call__.
This will resolve the crash and align the code with the current apply_rope implementation. You can then address the TODO about swapping the RoPE implementation in a separate change.
skyrl-tx/tx/models/deepseekv3.py
Outdated
| ) | ||
|
|
||
| # Bias only for specific model sizes (7168 hidden_size in original) | ||
| self.use_bias = config.hidden_size == 7168 |
There was a problem hiding this comment.
Hardcoding the magic number 7168 to determine self.use_bias makes the code brittle and less maintainable. If a new model variant is introduced that also requires this bias, this line would need to be updated. A better approach would be to introduce a dedicated boolean flag in the DeepseekV3Config, such as use_router_bias, to control this behavior explicitly.
skyrl-tx/tx/models/deepseekv3.py
Outdated
| # Bias only for specific model sizes (7168 hidden_size in original) | ||
| self.use_bias = config.hidden_size == 7168 | ||
| if self.use_bias: | ||
| from tx.layers.util import Param |
There was a problem hiding this comment.
This import statement is located inside a conditional block within __init__. According to PEP 8, all imports should be at the top of the file. This improves code readability and avoids potential circular import issues or unexpected behavior. Please move from tx.layers.util import Param to the top of the file with the other imports.
skyrl-tx/tx/models/deepseekv3.py
Outdated
| class DeepseekV3SharedMLP(nnx.Module): | ||
| """Always active shared experts.""" | ||
|
|
There was a problem hiding this comment.
The DeepseekV3SharedMLP class is nearly identical to DeepseekV3MLP, with the only significant difference being the intermediate_size. This creates code duplication, which can make maintenance harder.
To improve this, consider refactoring them into a single, more generic MLP class (e.g., SwiGLU) that accepts intermediate_size as a parameter in its __init__ method. You could then instantiate this class with config.intermediate_size for the standard MLP and with the calculated shared_inter_dim for the shared MLP part.
2 participants
[WIP]
Addresses #865