Dsv4 param mapping

[snehalv2002]

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:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

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):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request adds support for DeepSeek-V4 checkpoint conversion, refactors attention masking logic, and introduces comprehensive reference tests. However, several critical issues were identified in the review: direct usage of hf_config instead of target_cfg breaks validation and updates for multimodal models; removing Gemma 4 attention dimension helpers breaks per_layer_config overrides; modifying DeepSeek-V2/V3 shape mappings breaks backward compatibility; incorrect dimension indexing and lack of unpadding in embedding and logit layers cause shape mismatch crashes during Hugging Face to MaxText conversion; and hardcoding compress ratios in the DeepSeek-V4 shape mapping limits custom configurations.

[gemini-code-assist]

The removal of target_cfg and checking hf_config directly will break architecture validation for multimodal models (like Qwen2-VL) where text-specific attributes reside in hf_config.text_config. This causes hasattr(hf_config, hf_attr) to return False and silently skip all validation checks. We should restore target_cfg to ensure nested configurations are validated correctly.

Suggested change

	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)

[gemini-code-assist]

When updating the configuration, we should write to target_cfg instead of hf_config directly, to ensure nested configurations (like text_config in multimodal models) are updated correctly.

Suggested change

	if not is_match:
	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)
	if not is_match:
	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)

[gemini-code-assist]

The removal of _get_gemma4_layer_attention_dims and inline calculation of attention dimensions completely breaks support for per_layer_config overrides in Gemma 4 models. Gemma 4 models with heterogeneous layers or different head dimensions per layer rely on per_layer_config to specify these parameters. Hardcoding the global num_attention_heads, num_key_value_heads, and head_dim will cause shape mismatches and loading failures for these checkpoints. We should restore the per_layer_config handling.

[gemini-code-assist]

In DEEPSEEK_HF_WEIGHTS_TO_SHAPE (which is the shape mapping for DeepSeek-V2/V3), the keys were changed from mlp.shared_experts.gate_proj.weight to ffn.shared_experts.w1.weight. This will break checkpoint conversion for DeepSeek-V2 and DeepSeek-V3 because their Hugging Face checkpoints use the mlp prefix and standard Llama-like projection names. DeepSeek-V4 has its own separate shape mapping function DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE, so we should revert this change in DEEPSEEK_HF_WEIGHTS_TO_SHAPE to preserve compatibility with V2/V3.

Suggested change

	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],

[gemini-code-assist]

In unpad_hf_embedding_layer, target_vocab_size is set to target_shape[0]. When saving_to_hf is False (converting HF to MaxText), target_shape is the MaxText shape (emb_dim, vocab_size). Thus, target_shape[0] is emb_dim (e.g., 4096) instead of vocab_size (e.g., 128000). This causes the function to slice the wrong dimension and return a tensor with incorrect shape, leading to a shape mismatch crash. We should check saving_to_hf and use the correct dimension index for vocab_size.

Suggested change

	def unpad_hf_embedding_layer(input_tensor, target_shape):
	target_vocab_size = target_shape[0]
	if input_tensor.shape[0] == target_vocab_size:
	return input_tensor[:target_vocab_size, :]
	else:
	# MaxText is (emb_dim, vocab_size)
	return input_tensor[:, :target_vocab_size].T
	def unpad_hf_embedding_layer(input_tensor, target_shape):
	if saving_to_hf:
	target_vocab_size = target_shape[0]
	return input_tensor[:, :target_vocab_size].T
	else:
	target_vocab_size = target_shape[1]
	return input_tensor[:target_vocab_size, :].T

[gemini-code-assist]

In unpad_logits_layer, when saving_to_hf is False, it calls reshape_kernel(input_tensor, target_shape). However, reshape_kernel expects the total number of elements to match. Since the HF vocab size is padded (e.g., 129280) and MaxText vocab size is unpadded (e.g., 128000), the total number of elements is different, and reshape_kernel will raise a ValueError: RESHAPE FAILED. We must unpad the vocab dimension just like we do for the embedding layer.

Suggested change

	def unpad_logits_layer(input_tensor, target_shape):
	if saving_to_hf:
	return input_tensor[:, :target_shape[0]].T
	else:
	return reshape_kernel(input_tensor, target_shape)
	def unpad_logits_layer(input_tensor, target_shape):
	if saving_to_hf:
	return input_tensor[:, :target_shape[0]].T
	else:
	target_vocab_size = target_shape[1]
	return input_tensor[:target_vocab_size, :].T

[gemini-code-assist]

The list of compress ratios is hardcoded to a specific pattern of length 43. This makes it impossible to convert any model with a different number of layers or different compress ratios (like smaller test models or custom configurations). We should retrieve the compress ratios dynamically from the configuration (e.g., config.get("compress_rates") or config.get("compress_ratios")) to ensure robustness and maintainability.