Dsv4 mtp support - #5002
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 adds support for DeepSeek-V4, including MTP blocks and updated parameter mapping. The reviewer identified several critical issues: the removal of target_cfg in to_huggingface.py breaks multimodal model support; the deletion of _get_gemma4_layer_attention_dims breaks Gemma 4 per-layer configurations; renaming weight keys in hf_shape.py breaks compatibility with official HF checkpoints; a potential TypeError exists in attention_op.py when next_pos is None; and hardcoding compression ratios in hf_shape.py reduces configuration flexibility.
| for hf_attr, mt_attr in attributes_to_check: | ||
| # Skip checks if MaxText config doesn't have the attribute (shouldn't happen for valid configs) | ||
| if not hasattr(max_config, mt_attr): | ||
| # Skip checks if the HF config doesn't have this attribute (e.g. layer_norm_eps vs rms_norm_eps) | ||
| if not hasattr(hf_config, hf_attr): | ||
| continue | ||
|
|
||
| # Skip checks if the HF config doesn't have this attribute or raises AmbiguousGlobalPerLayerAttributeError | ||
| try: | ||
| hf_value = getattr(target_cfg, hf_attr) | ||
| except (AttributeError, ValueError, RuntimeError): | ||
| # Skip checks if MaxText config doesn't have the attribute (shouldn't happen for valid configs) | ||
| if not hasattr(max_config, mt_attr): | ||
| continue | ||
|
|
||
| hf_value = getattr(hf_config, hf_attr) | ||
| mt_value = getattr(max_config, mt_attr) |
There was a problem hiding this comment.
Removing target_cfg and directly using hf_config breaks checkpoint conversion for multimodal models (such as Qwen-VL or Gemma-3) where text-related configuration parameters are nested under hf_config.text_config.
Please restore the target_cfg fallback logic to maintain backward compatibility.
| for hf_attr, mt_attr in attributes_to_check: | |
| # Skip checks if MaxText config doesn't have the attribute (shouldn't happen for valid configs) | |
| if not hasattr(max_config, mt_attr): | |
| # Skip checks if the HF config doesn't have this attribute (e.g. layer_norm_eps vs rms_norm_eps) | |
| if not hasattr(hf_config, hf_attr): | |
| continue | |
| # Skip checks if the HF config doesn't have this attribute or raises AmbiguousGlobalPerLayerAttributeError | |
| try: | |
| hf_value = getattr(target_cfg, hf_attr) | |
| except (AttributeError, ValueError, RuntimeError): | |
| # Skip checks if MaxText config doesn't have the attribute (shouldn't happen for valid configs) | |
| if not hasattr(max_config, mt_attr): | |
| continue | |
| hf_value = getattr(hf_config, hf_attr) | |
| mt_value = getattr(max_config, mt_attr) | |
| target_cfg = getattr(hf_config, "text_config", hf_config) or hf_config | |
| for hf_attr, mt_attr in attributes_to_check: | |
| # Skip checks if the HF config doesn't have this attribute (e.g. layer_norm_eps vs rms_norm_eps) | |
| if not hasattr(target_cfg, hf_attr): | |
| continue | |
| # Skip checks if MaxText config doesn't have the attribute (shouldn't happen for valid configs) | |
| if not hasattr(max_config, mt_attr): | |
| continue | |
| hf_value = getattr(target_cfg, hf_attr) | |
| mt_value = getattr(max_config, mt_attr) |
| if override: | ||
| max_logging.log(f"⚠️ Overwriting HF Config '{hf_attr}': {hf_value} -> {mt_value} (from MaxText '{mt_attr}')") | ||
| setattr(target_cfg, hf_attr, mt_value) | ||
| setattr(hf_config, hf_attr, mt_value) |
There was a problem hiding this comment.
When override is enabled, setting the attribute directly on hf_config instead of target_cfg will fail to update the nested text configuration for multimodal models.
| setattr(hf_config, hf_attr, mt_value) | |
| setattr(target_cfg, hf_attr, mt_value) |
| if is_global: | ||
| q_dim = num_attention_heads * global_head_dim | ||
| kv_dim = num_global_key_value_heads * global_head_dim | ||
| norm_dim = global_head_dim | ||
| else: | ||
| q_dim = num_attention_heads * head_dim | ||
| kv_dim = num_key_value_heads * head_dim | ||
| norm_dim = head_dim |
There was a problem hiding this comment.
Deleting _get_gemma4_layer_attention_dims and hardcoding global config values breaks Gemma 4 models that use per_layer_config to define heterogeneous attention configurations across different layers.
Please restore _get_gemma4_layer_attention_dims and use it to determine the dimensions per layer.
| f"{layer_prefix}.ffn.shared_experts.w1.weight": [shared_intermediate_size, hidden_size], | ||
| f"{layer_prefix}.ffn.shared_experts.w3.weight": [shared_intermediate_size, hidden_size], | ||
| f"{layer_prefix}.ffn.shared_experts.w2.weight": [hidden_size, shared_intermediate_size], |
There was a problem hiding this comment.
Changing the Hugging Face weight names in DEEPSEEK_HF_WEIGHTS_TO_SHAPE (which is used for DeepSeek V2/V3) from mlp.shared_experts.gate_proj/up_proj/down_proj to ffn.shared_experts.w1/w3/w2 breaks compatibility with official Hugging Face DeepSeek-V2/V3 checkpoints.
| f"{layer_prefix}.ffn.shared_experts.w1.weight": [shared_intermediate_size, hidden_size], | |
| f"{layer_prefix}.ffn.shared_experts.w3.weight": [shared_intermediate_size, hidden_size], | |
| f"{layer_prefix}.ffn.shared_experts.w2.weight": [hidden_size, shared_intermediate_size], | |
| f"{layer_prefix}.mlp.shared_experts.gate_proj.weight": [shared_intermediate_size, hidden_size], | |
| f"{layer_prefix}.mlp.shared_experts.up_proj.weight": [shared_intermediate_size, hidden_size], | |
| f"{layer_prefix}.mlp.shared_experts.down_proj.weight": [hidden_size, shared_intermediate_size], |
| local_next = next_pos[:, None] if isinstance(next_pos, jax.Array) else next_pos | ||
| row_ids = jnp.arange(q_seq_len)[None, :, None] + local_next |
There was a problem hiding this comment.
If next_pos is None (which is common during training when simple/unpacked sequences are used), local_next will be None. Adding None to jnp.arange(...) will raise a TypeError and crash the training run.
Please default local_next to 0 if next_pos is None.
| local_next = next_pos[:, None] if isinstance(next_pos, jax.Array) else next_pos | |
| row_ids = jnp.arange(q_seq_len)[None, :, None] + local_next | |
| local_next = 0 if next_pos is None else (next_pos[:, None] if isinstance(next_pos, jax.Array) else next_pos) | |
| row_ids = jnp.arange(q_seq_len)[None, :, None] + local_next |
| layer_mapping[f"{layer_prefix}.ffn.gate.tid2eid"] = [vocab_size, config.get("num_experts_per_tok", 2)] | ||
|
|
||
| # Compressor logic fixing the clash! | ||
| ratio = [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4][layer_idx] if layer_idx < 43 else 4 |
There was a problem hiding this comment.
Hardcoding a 43-element list of compress ratios directly in DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE makes the code fragile and prevents supporting custom DeepSeek-V4 configurations with different layer counts or ratios.
Please retrieve compress_ratios or compress_rates from the config with a fallback.
compress_ratios = config.get("compress_ratios", config.get("compress_rates", [0, 0, 4, 128, 4]))
ratio = compress_ratios[layer_idx] if layer_idx < len(compress_ratios) else 41 participant
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:
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):
gemini-reviewlabel.