diff --git a/src/maxtext/checkpoint_conversion/to_huggingface.py b/src/maxtext/checkpoint_conversion/to_huggingface.py index 80a30516f1..e843f81c56 100644 --- a/src/maxtext/checkpoint_conversion/to_huggingface.py +++ b/src/maxtext/checkpoint_conversion/to_huggingface.py @@ -246,8 +246,12 @@ def _validate_or_update_architecture(hf_config, max_config, override: bool): ("qk_rope_head_dim", "qk_rope_head_dim"), ("v_head_dim", "v_head_dim"), ("vocab_size", "vocab_size"), - ("global_head_dim", "global_head_dim"), - ("num_global_key_value_heads", "global_num_kv_heads"), + ("hc_mult", "mhc_expansion_rate"), + ("num_hash_layers", "first_num_hash_layers"), + ("index_n_heads", "indexer_n_heads"), + ("index_head_dim", "indexer_head_dim"), + ("o_lora_rank", "o_lora_rank"), + ("o_groups", "o_groups"), ] if max_config.attention_type == "mla": @@ -264,19 +268,17 @@ def _validate_or_update_architecture(hf_config, max_config, override: bool): attributes_to_check.append(("head_dim", "head_dim")) mismatches = [] - target_cfg = getattr(hf_config, "text_config", hf_config) or hf_config 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) # Handle None values @@ -311,7 +313,7 @@ def _validate_or_update_architecture(hf_config, max_config, override: bool): 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) else: mismatches.append(f"{hf_attr} (HF={hf_value} vs MaxText={mt_value})") @@ -534,13 +536,15 @@ def main(argv: Sequence[str]) -> None: mappings = _get_model_mappings(model_key, config.scan_layers, hf_config_obj.to_dict(), config) param_map = mappings["param_mapping"] shape_map = mappings["shape_mapping"] # HF target shapes + + hook_fn_map = mappings["hook_fn_mapping"] # 4. Extract and transform weights for Linen/NNX-SFT/NNX-RL checkpoints maxtext_state_dict = detect_and_extract_checkpoint(checkpoint_dict) # Validate that checkpoint keys match the parameter mapping - state_keys = {k.replace("_lora_a", "").replace("_lora_b", "") for k in maxtext_state_dict} + state_keys = {k.replace("_lora_a", "").replace("_lora_b", "") for k in maxtext_state_dict } filtered_map_keys = validate_and_filter_param_map_keys(param_map, state_keys) # When not converting a multimodal model, skip vision encoder weights even if diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index ce2eecb678..65da59cfa6 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -317,11 +317,9 @@ def get_maxtext_model_info(config): quant = quantizations.configure_quantization(config) maxtext_model_flax = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) - # Get abstract model structure (name, shape) without materializing the weights to save memory. - # Extract the 'params' collection from the abstract model state. This focuses checkpoint - # conversion on trainable model parameters; variables outside the 'params' collection - # (such as non-trainable state or optimizer buffers) are not included. - abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model_flax, config)["params"] + # Get abstract model structure (name, shape) without materializing the weights to save memory + # Keeps all collections (e.g. 'params', 'Tid2EidVar') in the tree structure + abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model_flax, config) abstract_params_flat, abstract_params_treedef = jax.tree_util.tree_flatten_with_path( abstract_params_tree, @@ -333,7 +331,7 @@ def get_maxtext_model_info(config): # preprocess state maxtext_abstract_dict = {} for mt_target_idx, (path_tuple, abstract_leaf_value) in enumerate(abstract_params_flat): - mt_param_key = "params-" + "-".join(param_key_parts_from_path(path_tuple)) + mt_param_key = "-".join(param_key_parts_from_path(path_tuple)) if isinstance(abstract_leaf_value, nn.LogicallyPartitioned): mt_target_shape = abstract_leaf_value.value.shape else: @@ -413,6 +411,7 @@ def _build_single_axis_stacked_tensor( hook_fns: Any, target_shape: tuple, config, + mt_key: str = "", ) -> np.ndarray: """Builds a MaxText tensor by stacking HF weights along a single axis. @@ -431,11 +430,16 @@ def _build_single_axis_stacked_tensor( """ tensors_to_stack = [] - if config.scan_layers: - # If it's a standard scanned layer, we use the configured param_scan_axis. + if config.scan_layers and "scanned_blocks" in mt_key: + if "MoEBiasVar" in mt_key or "Tid2EidVar" in mt_key: + axis_to_stack = 0 + else: + axis_to_stack = config.param_scan_axis + elif config.scan_layers and "MoeBlock" not in mt_key: axis_to_stack = config.param_scan_axis + elif config.scan_layers and "MoeBlock" in mt_key and "scanned_blocks" not in mt_key: + axis_to_stack = 0 else: - # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0). axis_to_stack = 0 # The hook function needs the shape of an individual slice, not the full stacked tensor. @@ -474,6 +478,8 @@ def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_t if not isinstance(hf_source_keys_or_key, list): # Case 1: Single hf key (str) def _loader(getter, key, shape, hook): + if key is None: + return apply_hook_fns(None, shape, hook) if isinstance(key, (list, tuple)): tensors = tuple(getter(k) for k in key) return apply_hook_fns(tensors, shape, hook) @@ -496,6 +502,7 @@ def _loader(getter, key, shape, hook): hook_fn, mt_target_shape_or_shapes, config, + mt_key, ) else: # isinstance(hf_source_keys_or_key[0], list) @@ -517,12 +524,9 @@ def _get_maxtext_indices_and_shapes(mt_param_key_or_keys, maxtext_abstract_dict) The index is the parameter's order in `maxtext_abstract_dict.keys()`. This function handles two forms of MaxText keys: - - `atomic_mt_key`: A single string representing one MaxText parameter that maps to HF parameter(s). - Example: "params-decoder-layers_0-self_attention-query-kernel" -> returns a single index and shape tuple. + - `atomic_mt_key`: A single string representing one MaxText parameter that map to HF parameter(s). - `composite_mt_key`: A tuple of strings representing multiple MaxText parameters derived from a single/bundled HF parameter source (e.g., HF gate_up_proj splitting into MT wi_0 and wi_1). - Example: ("params-decoder-layers_0-mlp-wi_0-kernel", "params-decoder-layers_0-mlp-wi_1-kernel") -> - returns lists of indices and shapes for each composite component. """ is_composite_mt_key = isinstance(mt_param_key_or_keys, tuple) # atomic_mt_key @@ -983,6 +987,8 @@ def main( } def _eager_getter(key): + if key is None: + return None if key not in hf_state_dict_numpy: raise ValueError(f"HuggingFace key {key} not found in state_dict.") v = hf_state_dict_numpy[key] @@ -1080,11 +1086,19 @@ def _eager_getter(key): max_logging.log(f"Elapse for transform: {(time.time() - start) / 60:.2f} min") print_ram_usage("Before creating full JAX tree") - # Create final MaxText parameters tree + max_logging.log(f"Length of final_mt_weights: {len(final_mt_weights)}") + max_logging.log(f"Treedef num_leaves: {abstract_params_treedef.num_leaves}") + # Create final MaxText parameters tree containing all collections jax_weights = jax.tree_util.tree_unflatten(abstract_params_treedef, final_mt_weights) del final_mt_weights, abstract_params_treedef + state_params = jax_weights + print_ram_usage("Before saving") + leaves = jax.tree_util.tree_leaves(state_params) + max_logging.log(f"Length of state_params leaves: {len(leaves)}") + if leaves: + max_logging.log(f"Type of first leaf before saving: {type(leaves[0])}") if lazy_load_tensors and not is_adapter_only: max_logging.log("Starting checkpoint save (loading weights just-in-time)...") else: @@ -1095,7 +1109,7 @@ def _eager_getter(key): # and sharded across virtual devices. save_weights_to_checkpoint( output_directory, - jax_weights, + state_params, simulated_cpu_devices_count, config.checkpoint_storage_use_ocdbt, config.checkpoint_storage_use_zarr3, diff --git a/src/maxtext/checkpoint_conversion/utils/hf_shape.py b/src/maxtext/checkpoint_conversion/utils/hf_shape.py index 85dd1d6ea0..152a49288c 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_shape.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_shape.py @@ -153,39 +153,6 @@ def GEMMA3_HF_WEIGHTS_TO_SHAPE(config): return shapes -def _get_gemma4_layer_attention_dims(text_cfg: dict, layer_idx: int, is_global: bool): - """Extracts (q_dim, kv_dim, norm_dim) for a Gemma 4 layer accounting for per_layer_config overrides.""" - num_attention_heads = text_cfg["num_attention_heads"] - num_key_value_heads = text_cfg["num_key_value_heads"] - head_dim = text_cfg["head_dim"] - global_head_dim = text_cfg.get("global_head_dim") or head_dim - num_global_key_value_heads = text_cfg.get("num_global_key_value_heads") or num_key_value_heads - - per_layer_config = text_cfg.get("per_layer_config") or {} - if isinstance(per_layer_config, list): - layer_override = ( - per_layer_config[layer_idx] - if layer_idx < len(per_layer_config) and isinstance(per_layer_config[layer_idx], dict) - else {} - ) - elif isinstance(per_layer_config, dict): - layer_override = ( - per_layer_config.get(layer_idx) - or per_layer_config.get(str(layer_idx)) - or per_layer_config.get(f"{layer_idx:02d}") - or {} - ) - else: - layer_override = {} - - l_heads = layer_override.get("num_attention_heads") or num_attention_heads - l_kv_heads = layer_override.get("num_key_value_heads") or ( - num_global_key_value_heads if is_global else num_key_value_heads - ) - l_head_dim = layer_override.get("head_dim") or (global_head_dim if is_global else head_dim) - return l_heads * l_head_dim, l_kv_heads * l_head_dim, l_head_dim - - def GEMMA4_HF_WEIGHTS_TO_SHAPE(config): """Generates shape mapping for Hugging Face Gemma4 parameters. @@ -212,22 +179,33 @@ def GEMMA4_HF_WEIGHTS_TO_SHAPE(config): hidden_size = text_cfg["hidden_size"] intermediate_size = text_cfg["intermediate_size"] num_hidden_layers = text_cfg["num_hidden_layers"] + num_attention_heads = text_cfg["num_attention_heads"] + num_key_value_heads = text_cfg["num_key_value_heads"] + num_global_key_value_heads = text_cfg.get("num_global_key_value_heads", num_key_value_heads) + head_dim = text_cfg["head_dim"] + global_head_dim = text_cfg.get("global_head_dim", head_dim) vocab_size = text_cfg["vocab_size"] num_experts = text_cfg.get("num_experts") num_experts = num_experts if num_experts is not None else 1 # "moe_intermediate_size" is the canonical key in Gemma4 config; fall back to "expert_intermediate_size" expert_intermediate_size = text_cfg.get("moe_intermediate_size") or text_cfg.get("expert_intermediate_size") - layer_types = text_cfg.get("layer_types", []) shapes[f"{text_base}.embed_tokens.weight"] = [vocab_size, hidden_size] shapes[f"{text_base}.norm.weight"] = [hidden_size] for i in range(num_hidden_layers): hf_prefix = f"{text_base}.layers.{i}" - is_global = (i < len(layer_types) and layer_types[i] == "full_attention") if layer_types else (i % 6) == 5 + is_global = (i % 6) == 5 - q_dim, kv_dim, norm_dim = _get_gemma4_layer_attention_dims(text_cfg, i, is_global) + 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 shapes[f"{hf_prefix}.self_attn.q_proj.weight"] = [q_dim, hidden_size] shapes[f"{hf_prefix}.self_attn.k_proj.weight"] = [kv_dim, hidden_size] @@ -313,7 +291,7 @@ def GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE(config): * derives global-vs-sliding from the per-model ``layer_types`` list (E2B has period-5, E4B has period-6), * emits the Per-Layer-Embedding parameters when ``hidden_size_per_layer_input`` > 0, - * omits k_proj/v_proj/k_norm/v_norm shapes on KV-shared layers, and + * omits k_proj/v_proj/k_norm/v_norm shapes on KV-shared layers, and * doubles ``intermediate_size`` on shared layers when ``use_double_wide_mlp`` is set (E2B). """ @@ -326,6 +304,11 @@ def GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE(config): hidden_size = text_cfg["hidden_size"] intermediate_size = text_cfg["intermediate_size"] num_hidden_layers = text_cfg["num_hidden_layers"] + num_attention_heads = text_cfg["num_attention_heads"] + num_key_value_heads = text_cfg["num_key_value_heads"] + num_global_key_value_heads = text_cfg.get("num_global_key_value_heads") or num_key_value_heads + head_dim = text_cfg["head_dim"] + global_head_dim = text_cfg.get("global_head_dim", head_dim) vocab_size = text_cfg["vocab_size"] layer_types = text_cfg.get("layer_types", []) @@ -347,7 +330,15 @@ def GEMMA4_SMALL_HF_WEIGHTS_TO_SHAPE(config): hf_prefix = f"{text_base}.layers.{i}" is_global = i < len(layer_types) and layer_types[i] == "full_attention" is_shared = num_kv_shared > 0 and i >= first_shared - q_dim, kv_dim, norm_dim = _get_gemma4_layer_attention_dims(text_cfg, i, is_global) + + 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 shapes[f"{hf_prefix}.self_attn.q_proj.weight"] = [q_dim, hidden_size] shapes[f"{hf_prefix}.self_attn.o_proj.weight"] = [hidden_size, q_dim] @@ -585,9 +576,9 @@ def DEEPSEEK_HF_WEIGHTS_TO_SHAPE(config): shared_intermediate_size = moe_intermediate_size * n_shared_experts layer_mapping.update( { - 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], + 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], } ) @@ -1174,105 +1165,89 @@ def QWEN3_VL_HF_WEIGHTS_TO_SHAPE(config): def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config): - """Returns a dictionary mapping HuggingFace weight names to shapes for DeepSeek V4.""" hidden_size = config["hidden_size"] vocab_size = config["vocab_size"] num_hidden_layers = config["num_hidden_layers"] q_lora_rank = config.get("q_lora_rank", 1024) kv_lora_rank = config.get("kv_lora_rank", 512) - # MaxText scales o_lora_rank by o_groups to get o_a_out_features (8192) o_lora_rank = config.get("o_lora_rank", 1024) * config.get("o_groups", 8) num_attention_heads = config["num_attention_heads"] head_dim = config["head_dim"] - moe_intermediate_size = config["moe_intermediate_size"] n_routed_experts = config["n_routed_experts"] - - # Dynamic hyperparameters replacing hardcoded literals: - hc_mult = config.get("hc_mult", 4) - hc_dim = (2 * hc_mult) + (hc_mult**2) # 2(4) + 16 = 24 - num_hash_layers = config.get("num_hash_layers", 3) # 3 - index_n_heads = config.get("index_n_heads", 64) # 64 - index_head_dim = config.get("index_head_dim", 128) # 128 + hc_mult = config.get("hc_mult", config.get("mhc_expansion_rate", 4)) + hc_dim = (2 * hc_mult) + (hc_mult**2) + num_hash_layers = config.get("num_hash_layers", config.get("first_num_hash_layers", 3)) + index_n_heads = config.get("index_n_heads", config.get("indexer_n_heads", 64)) + index_head_dim = config.get("index_head_dim", config.get("indexer_head_dim", 128)) mapping = { - "model.embed_tokens.weight": [vocab_size, hidden_size], - "model.norm.weight": [hidden_size], + "embed.weight": [vocab_size, hidden_size], + "norm.weight": [hidden_size], "head.weight": [vocab_size, hidden_size], - "model.hc_head.hc_fn": [hc_mult, hidden_size * hc_mult], - "model.hc_head.hc_base": [hc_mult], - "model.hc_head.hc_scale": [1], + "hc_head_fn": [hc_mult, hidden_size * hc_mult], + "hc_head_base": [hc_mult], + "hc_head_scale": [1], } for layer_idx in range(num_hidden_layers): - layer_prefix = f"model.layers.{layer_idx}" + layer_prefix = f"layers.{layer_idx}" layer_mapping = { - f"{layer_prefix}.input_layernorm.weight": [hidden_size], - f"{layer_prefix}.post_attention_layernorm.weight": [hidden_size], - f"{layer_prefix}.self_attn.q_a_proj.weight": [q_lora_rank, hidden_size], - f"{layer_prefix}.self_attn.q_a_norm.weight": [q_lora_rank], - f"{layer_prefix}.self_attn.q_b_proj.weight": [num_attention_heads * head_dim, q_lora_rank], - f"{layer_prefix}.self_attn.kv_proj.weight": [kv_lora_rank, hidden_size], - f"{layer_prefix}.self_attn.kv_norm.weight": [kv_lora_rank], - f"{layer_prefix}.self_attn.sinks": [num_attention_heads], - f"{layer_prefix}.self_attn.o_a_proj.weight": [o_lora_rank, hidden_size], - f"{layer_prefix}.self_attn.o_b_proj.weight": [hidden_size, o_lora_rank], - # MHC - f"{layer_prefix}.attn_hc.fn": [hc_dim, hidden_size * hc_mult], - f"{layer_prefix}.attn_hc.base": [hc_dim], - f"{layer_prefix}.attn_hc.scale": [num_hash_layers], - f"{layer_prefix}.ffn_hc.fn": [hc_dim, hidden_size * hc_mult], - f"{layer_prefix}.ffn_hc.base": [hc_dim], - f"{layer_prefix}.ffn_hc.scale": [num_hash_layers], - # MLP / MoE Block - f"{layer_prefix}.mlp.gate.weight": [n_routed_experts, hidden_size], - f"{layer_prefix}.mlp.gate.e_score_correction_bias": [n_routed_experts], - f"{layer_prefix}.mlp.shared_experts.gate_proj.weight": [moe_intermediate_size, hidden_size], - f"{layer_prefix}.mlp.shared_experts.up_proj.weight": [moe_intermediate_size, hidden_size], - f"{layer_prefix}.mlp.shared_experts.down_proj.weight": [hidden_size, moe_intermediate_size], + f"{layer_prefix}.attn_norm.weight": [hidden_size], + f"{layer_prefix}.ffn_norm.weight": [hidden_size], + f"{layer_prefix}.attn.wq_a.weight": [q_lora_rank, hidden_size], + f"{layer_prefix}.attn.q_norm.weight": [q_lora_rank], + f"{layer_prefix}.attn.wq_b.weight": [num_attention_heads * head_dim, q_lora_rank], + f"{layer_prefix}.attn.wkv.weight": [config.get("num_key_value_heads", 1)*config.get("head_dim", 64), hidden_size], + f"{layer_prefix}.attn.kv_norm.weight": [head_dim], + f"{layer_prefix}.attn.attn_sink": [num_attention_heads], + f"{layer_prefix}.attn.wo_a.weight": [o_lora_rank, (num_attention_heads * head_dim) // config.get("o_groups", 8)], + f"{layer_prefix}.attn.wo_b.weight": [hidden_size, o_lora_rank], + f"{layer_prefix}.hc_attn_fn": [hc_dim, hidden_size * hc_mult], + f"{layer_prefix}.hc_attn_base": [hc_dim], + f"{layer_prefix}.hc_attn_scale": [num_hash_layers], + f"{layer_prefix}.hc_ffn_fn": [hc_dim, hidden_size * hc_mult], + f"{layer_prefix}.hc_ffn_base": [hc_dim], + f"{layer_prefix}.hc_ffn_scale": [num_hash_layers], + f"{layer_prefix}.ffn.gate.weight": [n_routed_experts, hidden_size], + f"{layer_prefix}.ffn.gate.bias": [n_routed_experts], } - - # Experts - for e in range(n_routed_experts): - layer_mapping[f"{layer_prefix}.mlp.experts.{e}.w1.weight"] = [moe_intermediate_size, hidden_size] - layer_mapping[f"{layer_prefix}.mlp.experts.{e}.w3.weight"] = [moe_intermediate_size, hidden_size] - layer_mapping[f"{layer_prefix}.mlp.experts.{e}.w2.weight"] = [hidden_size, moe_intermediate_size] - - # Compressors - if layer_idx >= 2: - c_type = "csa" if (layer_idx % 2 == 0) else "hca" - if c_type == "csa": - layer_mapping.update( - { - f"{layer_prefix}.self_attn.compressor.kv_proj.weight": [1024, hidden_size], - f"{layer_prefix}.self_attn.compressor.gate_proj.weight": [1024, hidden_size], - f"{layer_prefix}.self_attn.compressor.position_bias": [4, 1024], - f"{layer_prefix}.self_attn.compressor.kv_norm.weight": [512], - } - ) - layer_mapping.update( - { - f"{layer_prefix}.self_attn.compressor.indexer.gate_proj.weight": [256, hidden_size], - f"{layer_prefix}.self_attn.compressor.indexer.kv_proj.weight": [256, hidden_size], - f"{layer_prefix}.self_attn.compressor.indexer.q_b_proj.weight": [ - index_n_heads * index_head_dim, - q_lora_rank, - ], - f"{layer_prefix}.self_attn.compressor.indexer.scorer.weights_proj.weight": [index_n_heads, hidden_size], - f"{layer_prefix}.self_attn.compressor.indexer.position_bias": [4, 256], - f"{layer_prefix}.self_attn.compressor.indexer.kv_norm.weight": [128], - } - ) - else: # hca - layer_mapping.update( - { - f"{layer_prefix}.self_attn.compressor.kv_proj.weight": [512, hidden_size], - f"{layer_prefix}.self_attn.compressor.gate_proj.weight": [512, hidden_size], - f"{layer_prefix}.self_attn.compressor.position_bias": [128, 512], - f"{layer_prefix}.self_attn.compressor.kv_norm.weight": [512], - } - ) - + if layer_idx < num_hash_layers: + 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 + if ratio > 0: + if ratio == 4: + layer_mapping[f"{layer_prefix}.attn.compressor.wgate.weight"] = [2 * head_dim, hidden_size] + layer_mapping[f"{layer_prefix}.attn.compressor.wkv.weight"] = [2 * head_dim, hidden_size] + layer_mapping[f"{layer_prefix}.attn.compressor.norm.weight"] = [head_dim] + layer_mapping[f"{layer_prefix}.attn.compressor.ape"] = [ratio, 2 * head_dim] + + # Indexer (csa only) + layer_mapping[f"{layer_prefix}.attn.indexer.compressor.wgate.weight"] = [2 * index_head_dim, hidden_size] + layer_mapping[f"{layer_prefix}.attn.indexer.compressor.wkv.weight"] = [2 * index_head_dim, hidden_size] + layer_mapping[f"{layer_prefix}.attn.indexer.compressor.norm.weight"] = [index_head_dim] + layer_mapping[f"{layer_prefix}.attn.indexer.compressor.ape"] = [ratio, 2 * index_head_dim] + layer_mapping[f"{layer_prefix}.attn.indexer.weights_proj.weight"] = [index_n_heads, hidden_size] + layer_mapping[f"{layer_prefix}.attn.indexer.wq_b.weight"] = [index_n_heads * index_head_dim, q_lora_rank] + else: + layer_mapping[f"{layer_prefix}.attn.compressor.wgate.weight"] = [head_dim, hidden_size] + layer_mapping[f"{layer_prefix}.attn.compressor.wkv.weight"] = [head_dim, hidden_size] + layer_mapping[f"{layer_prefix}.attn.compressor.norm.weight"] = [head_dim] + layer_mapping[f"{layer_prefix}.attn.compressor.ape"] = [ratio, head_dim] + + + for exp_idx in range(n_routed_experts): + layer_mapping[f"{layer_prefix}.ffn.experts.{exp_idx}.w1.weight"] = [moe_intermediate_size, hidden_size] + layer_mapping[f"{layer_prefix}.ffn.experts.{exp_idx}.w2.weight"] = [hidden_size, moe_intermediate_size] + layer_mapping[f"{layer_prefix}.ffn.experts.{exp_idx}.w3.weight"] = [moe_intermediate_size, hidden_size] + + layer_mapping[f"{layer_prefix}.ffn.shared_experts.w1.weight"] = [moe_intermediate_size, hidden_size] + layer_mapping[f"{layer_prefix}.ffn.shared_experts.w3.weight"] = [moe_intermediate_size, hidden_size] + layer_mapping[f"{layer_prefix}.ffn.shared_experts.w2.weight"] = [hidden_size, moe_intermediate_size] + + mapping.update(layer_mapping) return mapping diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 94e96173f1..aa407a6cae 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -514,7 +514,10 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def scale_rmsnorm_layer(input_tensor, target_shape): if saving_to_hf: @@ -782,7 +785,10 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def reshape_expert_kernel(input_tensor, target_shape=None): """Transposes expert weights. @@ -1143,7 +1149,10 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def permute_conv(input_tensor, target_shape=None): # MT: [K, 1, C] <-> HF: [C, 1, K] @@ -1284,7 +1293,7 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): hooks[f"{mlp_prefix}-shared_expert-wi_1-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose - # pyrefly: ignore[unsupported-operation] + hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = ( process_wi_0_wi_1 # pyrefly: ignore[unsupported-operation] ) @@ -1408,7 +1417,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F prefix = f"params-decoder-layers-layer_{block_idx}" # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] + mapping[f"{prefix}-input_layernorm-scale"] = [ f"model.layers.{i}.input_layernorm.weight" for i in hf_indices ] # pyrefly: ignore[bad-assignment] mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] @@ -1576,7 +1585,10 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def permute_conv(input_tensor, target_shape=None): # MT: [K, 1, C] <-> HF: [C, 1, K] @@ -1635,7 +1647,7 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal # Extract hf configuration parameters, without mtp num_main_layers = config["num_hidden_layers"] first_num_dense_layers = config["first_k_dense_replace"] - num_experts = config.get("n_routed_experts", 0) + num_experts = config.get("n_routed_experts", config.get("num_experts", maxtext_config.num_experts)) # Mapping for non-layer-specific weights mapping = { @@ -1731,13 +1743,39 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') num_main_layers = config["num_hidden_layers"] first_num_dense_layers = config["first_k_dense_replace"] + + + + + 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_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) + mapping = { - "params-decoder-logits_dense-kernel": reshape_kernel, + "params-token_embedder-embedding": unpad_hf_embedding_layer, + "params-decoder-logits_dense-kernel": unpad_logits_layer, + + } attention_need_reshape = { @@ -1919,7 +1957,10 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def reshape_bias(input_tensor, target_shape=None): """Reshapes biases between MaxText 2D (heads, dim) and HF 1D (hidden).""" @@ -1942,9 +1983,9 @@ def interleave(input_tensor, target_shape=None): """ if saving_to_hf: wi_0, wi_1 = input_tensor - wi_0_1 = jnp.empty(target_shape, dtype=wi_0.dtype) - wi_0_1 = wi_0_1.at[..., ::2].set(wi_0) - wi_0_1 = wi_0_1.at[..., 1::2].set(wi_1) + wi_0_1 = np.empty(target_shape, dtype=wi_0.dtype) # pyrefly: ignore[no-matching-overload] + wi_0_1[..., ::2] = wi_0 + wi_0_1[..., 1::2] = wi_1 return wi_0_1 else: wi_0_1 = input_tensor @@ -1969,7 +2010,6 @@ def interleave(input_tensor, target_shape=None): hooks[f"{prefix}-GptOssMlp-gate-kernel"] = transpose # `composite_mt_key`: A hook for combining multiple MaxText params. hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave # pyrefly: ignore[unsupported-operation] - # pyrefly: ignore[unsupported-operation] hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = ( interleave # pyrefly: ignore[unsupported-operation] ) @@ -2183,7 +2223,10 @@ def reshape_kernel_vision(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def reshape_conv3d_patch_embed(input_tensor, target_shape): """Reshape 3D conv patch embedding weight. @@ -2310,7 +2353,10 @@ def reshape_kernel_audio(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def reshape_conv2d_audio(input_tensor, target_shape): """Reshape Conv2D weight for audio. @@ -2891,10 +2937,10 @@ def _spec_active(gate): local_positions = list(range(attention_pattern_length - 1)) for subkey, suffix, gate in param_specs: if _spec_active(gate): - mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{local_prefix}-{subkey}"] = [ # pyrefly: ignore[no-matching-overload] [hf_layer(b * attention_pattern_length + l, suffix) for l in local_positions] for b in range(num_blocks) ] - mapping[f"{local_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{local_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[no-matching-overload] [hf_layer(b * attention_pattern_length + l, "self_attn.v_proj.weight") for l in local_positions] for b in range(num_blocks) ] @@ -2904,11 +2950,11 @@ def _spec_active(gate): global_position = attention_pattern_length - 1 for subkey, suffix, gate in param_specs: if _spec_active(gate): - mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{global_prefix}-{subkey}"] = [ # pyrefly: ignore[no-matching-overload] hf_layer(b * attention_pattern_length + global_position, suffix) for b in range(num_blocks) ] if not share_kv_projections: - mapping[f"{global_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment, no-matching-overload] + mapping[f"{global_prefix}-self_attention-value-kernel"] = [ # pyrefly: ignore[no-matching-overload] hf_layer(b * attention_pattern_length + global_position, "self_attn.v_proj.weight") for b in range(num_blocks) ] @@ -3303,7 +3349,10 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def scale_rmsnorm_layer(input_tensor, target_shape): # Shift of 1.0 is now folded into Gemma 4 text and vision checkpoint weights @@ -3579,7 +3628,10 @@ def reshape_kernel(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') # Identity mapping for Norms # Olmo3 checkpoints typically have weights ~1.0. @@ -3803,7 +3855,10 @@ def reshape_kernel_vision(input_tensor, target_shape): flipped_target_shape = np.flip(np.array(target_shape)) return input_tensor.reshape(flipped_target_shape).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') def reshape_conv3d_patch_embed(input_tensor, target_shape): """Reshape 3D conv patch embedding weight.""" @@ -3851,7 +3906,10 @@ def reshape_vision_attn_out(input_tensor, target_shape): if saving_to_hf: return input_tensor.reshape(hidden_size, hidden_size).T else: - return input_tensor.T.reshape(target_shape) + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') mapping["params-vision_encoder-Qwen3VLVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed @@ -3884,339 +3942,521 @@ def reshape_vision_attn_out(input_tensor, target_shape): mapping[f"{prefix}-mlp_0-kernel"] = reshape_kernel_vision mapping[f"{prefix}-mlp_2-kernel"] = reshape_kernel_vision - mapping["params-vision_encoder-Qwen3VLVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision - mapping["params-vision_encoder-Qwen3VLVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision - return mapping -# {maxtext model name: {maxtext weight name: hf weight name}} - - -def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): - """Maps MaxText parameter keys to HuggingFace parameter keys for DeepSeek V4.""" - n_layers = config["num_hidden_layers"] - num_experts = config.get("n_routed_experts", 8) +def DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): + def _get(cfg, key, default=None): + if isinstance(cfg, dict): + return cfg.get(key, default) + return getattr(cfg, key, default) + n_layers = _get(config, "num_hidden_layers", getattr(maxtext_config, "base_num_decoder_layers", 43)) + num_experts = maxtext_config.num_experts if getattr(maxtext_config, "num_experts", 0) > 0 else _get(config, "n_routed_experts", 256) + num_hash_layers = _get(config, "num_hash_layers", getattr(maxtext_config, "first_num_hash_layers", 3)) + mapping = { - "params-token_embedder-embedding": "model.embed_tokens.weight", - "params-decoder-decoder_norm-scale": "model.norm.weight", + "params-token_embedder-embedding": "embed.weight", + "params-decoder-decoder_norm-scale": "norm.weight", "params-decoder-logits_dense-kernel": "head.weight", - "params-decoder-hc_head-hc_fn": "model.hc_head.hc_fn", - "params-decoder-hc_head-hc_base": "model.hc_head.hc_base", - "params-decoder-hc_head-hc_scale": "model.hc_head.hc_scale", + "params-decoder-hc_head-hc_fn": "hc_head_fn", + "params-decoder-hc_head-hc_base": "hc_head_base", + "params-decoder-hc_head-hc_scale": "hc_head_scale", } - def add_layer_mapping(mt_layer_path, hf_layer_indices): - is_list = isinstance(hf_layer_indices, list) - - def get_hf_key(subpath): - if subpath is None: - return None - if is_list: - return [f"model.layers.{idx}.{subpath}" for idx in hf_layer_indices] - else: - return f"model.layers.{hf_layer_indices}.{subpath}" - - def get_hf_expert_keys(expert_subpath_template): - if is_list: - return [ - [f"model.layers.{idx}.mlp.experts.{e}.{expert_subpath_template}" for idx in hf_layer_indices] - for e in range(num_experts) - ] - else: - return [f"model.layers.{hf_layer_indices}.mlp.experts.{e}.{expert_subpath_template}" for e in range(num_experts)] - - layer_map = { - f"{mt_layer_path}-pre_self_attention_layer_norm-scale": get_hf_key("input_layernorm.weight"), - f"{mt_layer_path}-post_self_attention_layer_norm-scale": get_hf_key("post_attention_layernorm.weight"), - # Attention - f"{mt_layer_path}-self_attention-wq_a-kernel": get_hf_key("self_attn.q_a_proj.weight"), - f"{mt_layer_path}-self_attention-q_norm-scale": get_hf_key("self_attn.q_a_norm.weight"), - f"{mt_layer_path}-self_attention-wq_b-kernel": get_hf_key("self_attn.q_b_proj.weight"), - f"{mt_layer_path}-self_attention-wkv-kernel": get_hf_key("self_attn.kv_proj.weight"), - f"{mt_layer_path}-self_attention-kv_norm-scale": get_hf_key("self_attn.kv_norm.weight"), - f"{mt_layer_path}-self_attention-sinks": get_hf_key("self_attn.sinks"), - f"{mt_layer_path}-self_attention-o_a_proj-kernel": get_hf_key("self_attn.o_a_proj.weight"), - f"{mt_layer_path}-self_attention-o_b_proj-kernel": get_hf_key("self_attn.o_b_proj.weight"), - # mHC Attention - f"{mt_layer_path}-mhc_attention-mhc_norm-scale": None, - f"{mt_layer_path}-mhc_attention-pre_alpha": get_hf_key("attn_hc.fn"), - f"{mt_layer_path}-mhc_attention-post_alpha": get_hf_key("attn_hc.fn"), - f"{mt_layer_path}-mhc_attention-res_alpha": get_hf_key("attn_hc.fn"), - f"{mt_layer_path}-mhc_attention-pre_beta": get_hf_key("attn_hc.base"), - f"{mt_layer_path}-mhc_attention-post_beta": get_hf_key("attn_hc.base"), - f"{mt_layer_path}-mhc_attention-res_beta": get_hf_key("attn_hc.base"), - f"{mt_layer_path}-mhc_attention-pre_alpha_scale": get_hf_key("attn_hc.scale"), - f"{mt_layer_path}-mhc_attention-post_alpha_scale": get_hf_key("attn_hc.scale"), - f"{mt_layer_path}-mhc_attention-res_alpha_scale": get_hf_key("attn_hc.scale"), - # mHC MLP - f"{mt_layer_path}-mhc_mlp-mhc_norm-scale": None, - f"{mt_layer_path}-mhc_mlp-pre_alpha": get_hf_key("ffn_hc.fn"), - f"{mt_layer_path}-mhc_mlp-post_alpha": get_hf_key("ffn_hc.fn"), - f"{mt_layer_path}-mhc_mlp-res_alpha": get_hf_key("ffn_hc.fn"), - f"{mt_layer_path}-mhc_mlp-pre_beta": get_hf_key("ffn_hc.base"), - f"{mt_layer_path}-mhc_mlp-post_beta": get_hf_key("ffn_hc.base"), - f"{mt_layer_path}-mhc_mlp-res_beta": get_hf_key("ffn_hc.base"), - f"{mt_layer_path}-mhc_mlp-pre_alpha_scale": get_hf_key("ffn_hc.scale"), - f"{mt_layer_path}-mhc_mlp-post_alpha_scale": get_hf_key("ffn_hc.scale"), - f"{mt_layer_path}-mhc_mlp-res_alpha_scale": get_hf_key("ffn_hc.scale"), - # MoE Block - f"{mt_layer_path}-mlp-MoeBlock_0-gate-kernel": get_hf_key("mlp.gate.weight"), - # Shared Experts - f"{mt_layer_path}-mlp-shared_experts-wi_0-kernel": get_hf_key("mlp.shared_experts.gate_proj.weight"), - f"{mt_layer_path}-mlp-shared_experts-wi_1-kernel": get_hf_key("mlp.shared_experts.up_proj.weight"), - f"{mt_layer_path}-mlp-shared_experts-wo-kernel": get_hf_key("mlp.shared_experts.down_proj.weight"), - # Stacked Experts - f"{mt_layer_path}-mlp-MoeBlock_0-wi_0": get_hf_expert_keys("w1.weight"), - f"{mt_layer_path}-mlp-MoeBlock_0-wi_1": get_hf_expert_keys("w3.weight"), - f"{mt_layer_path}-mlp-MoeBlock_0-wo": get_hf_expert_keys("w2.weight"), - } - - if (is_list and hf_layer_indices[0] >= 3) or (not is_list and hf_layer_indices >= 3): - layer_map[f"{mt_layer_path}-mlp-MoeBlock_0-gate-bias"] = get_hf_key("mlp.gate.e_score_correction_bias") - - first_idx = hf_layer_indices[0] if is_list else hf_layer_indices - if first_idx >= 2: - if first_idx % 2 == 0: - layer_map.update( - { - f"{mt_layer_path}-self_attention-csa_compressor-kv_proj-kernel": get_hf_key( - "self_attn.compressor.kv_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-gate_proj-kernel": get_hf_key( - "self_attn.compressor.gate_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-position_bias": get_hf_key( - "self_attn.compressor.position_bias" - ), - f"{mt_layer_path}-self_attention-csa_compressor-kv_norm-scale": get_hf_key( - "self_attn.compressor.kv_norm.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-gate_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.gate_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-kv_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.kv_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-q_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.q_b_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-weights_proj-kernel": get_hf_key( - "self_attn.compressor.indexer.scorer.weights_proj.weight" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-position_bias": get_hf_key( - "self_attn.compressor.indexer.position_bias" - ), - f"{mt_layer_path}-self_attention-csa_compressor-indexer-kv_norm-scale": get_hf_key( - "self_attn.compressor.indexer.kv_norm.weight" - ), - } - ) - else: - layer_map.update( - { - f"{mt_layer_path}-self_attention-hca_compressor-kv_proj-kernel": get_hf_key( - "self_attn.compressor.kv_proj.weight" - ), - f"{mt_layer_path}-self_attention-hca_compressor-gate_proj-kernel": get_hf_key( - "self_attn.compressor.gate_proj.weight" - ), - f"{mt_layer_path}-self_attention-hca_compressor-position_bias": get_hf_key( - "self_attn.compressor.position_bias" - ), - f"{mt_layer_path}-self_attention-hca_compressor-kv_norm-scale": get_hf_key( - "self_attn.compressor.kv_norm.weight" - ), - } - ) - - mapping.update(layer_map) # pyrefly: ignore[no-matching-overload] - if not scan_layers: for i in range(n_layers): - add_layer_mapping(f"params-decoder-layers_{i}", i) - else: - for i in range(3): - add_layer_mapping(f"params-decoder-layers_{i}", i) - add_layer_mapping("params-decoder-scanned_blocks-layers_0", list(range(3, n_layers, 2))) - add_layer_mapping("params-decoder-scanned_blocks-layers_1", list(range(4, n_layers, 2))) + prefix = f"params-decoder-layers_{i}" + hf_prefix = f"layers.{i}" + + mapping[f"{prefix}-pre_self_attention_layer_norm-scale"] = f"{hf_prefix}.attn_norm.weight" + mapping[f"{prefix}-post_self_attention_layer_norm-scale"] = f"{hf_prefix}.ffn_norm.weight" + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = None + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = None + + # MHC Attention Alpha (Atomic) + mapping[f"{prefix}-mhc_attention-pre_alpha"] = f"{hf_prefix}.hc_attn_fn" + mapping[f"{prefix}-mhc_attention-post_alpha"] = f"{hf_prefix}.hc_attn_fn" + mapping[f"{prefix}-mhc_attention-res_alpha"] = f"{hf_prefix}.hc_attn_fn" + + # MHC Attention Beta (Atomic) + mapping[f"{prefix}-mhc_attention-pre_beta"] = f"{hf_prefix}.hc_attn_base" + mapping[f"{prefix}-mhc_attention-post_beta"] = f"{hf_prefix}.hc_attn_base" + mapping[f"{prefix}-mhc_attention-res_beta"] = f"{hf_prefix}.hc_attn_base" + + # MHC Attention Scale (Atomic) + mapping[f"{prefix}-mhc_attention-pre_alpha_scale"] = f"{hf_prefix}.hc_attn_scale" + mapping[f"{prefix}-mhc_attention-post_alpha_scale"] = f"{hf_prefix}.hc_attn_scale" + mapping[f"{prefix}-mhc_attention-res_alpha_scale"] = f"{hf_prefix}.hc_attn_scale" + + # MHC MLP Alpha (Atomic) + mapping[f"{prefix}-mhc_mlp-pre_alpha"] = f"{hf_prefix}.hc_ffn_fn" + mapping[f"{prefix}-mhc_mlp-post_alpha"] = f"{hf_prefix}.hc_ffn_fn" + mapping[f"{prefix}-mhc_mlp-res_alpha"] = f"{hf_prefix}.hc_ffn_fn" + + # MHC MLP Beta (Atomic) + mapping[f"{prefix}-mhc_mlp-pre_beta"] = f"{hf_prefix}.hc_ffn_base" + mapping[f"{prefix}-mhc_mlp-post_beta"] = f"{hf_prefix}.hc_ffn_base" + mapping[f"{prefix}-mhc_mlp-res_beta"] = f"{hf_prefix}.hc_ffn_base" + + # MHC MLP Scale (Atomic) + mapping[f"{prefix}-mhc_mlp-pre_alpha_scale"] = f"{hf_prefix}.hc_ffn_scale" + mapping[f"{prefix}-mhc_mlp-post_alpha_scale"] = f"{hf_prefix}.hc_ffn_scale" + mapping[f"{prefix}-mhc_mlp-res_alpha_scale"] = f"{hf_prefix}.hc_ffn_scale" + + # Attention Core + mapping[f"{prefix}-self_attention-q_norm-scale"] = f"{hf_prefix}.attn.q_norm.weight" + mapping[f"{prefix}-self_attention-kv_norm-scale"] = f"{hf_prefix}.attn.kv_norm.weight" + mapping[f"{prefix}-self_attention-wq_a-kernel"] = f"{hf_prefix}.attn.wq_a.weight" + mapping[f"{prefix}-self_attention-wq_b-kernel"] = f"{hf_prefix}.attn.wq_b.weight" + mapping[f"{prefix}-self_attention-wkv-kernel"] = f"{hf_prefix}.attn.wkv.weight" + mapping[f"{prefix}-self_attention-sinks"] = f"{hf_prefix}.attn.attn_sink" + + # Output projection + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = f"{hf_prefix}.attn.wo_a.weight" + mapping[f"{prefix}-self_attention-o_b_proj-kernel"] = f"{hf_prefix}.attn.wo_b.weight" + + # Compressors CSA and HCA + mapping[f"{prefix}-self_attention-csa_compressor-gate_proj-kernel"] = f"{hf_prefix}.attn.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-csa_compressor-kv_proj-kernel"] = f"{hf_prefix}.attn.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-csa_compressor-kv_norm-scale"] = f"{hf_prefix}.attn.compressor.norm.weight" + mapping[f"{prefix}-self_attention-csa_compressor-position_bias"] = f"{hf_prefix}.attn.compressor.ape" + + mapping[f"{prefix}-self_attention-hca_compressor-gate_proj-kernel"] = f"{hf_prefix}.attn.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-hca_compressor-kv_proj-kernel"] = f"{hf_prefix}.attn.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-hca_compressor-kv_norm-scale"] = f"{hf_prefix}.attn.compressor.norm.weight" + mapping[f"{prefix}-self_attention-hca_compressor-position_bias"] = f"{hf_prefix}.attn.compressor.ape" + + mapping[f"{prefix}-self_attention-csa_compressor-indexer-gate_proj-kernel"] = f"{hf_prefix}.attn.indexer.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_proj-kernel"] = f"{hf_prefix}.attn.indexer.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_norm-scale"] = f"{hf_prefix}.attn.indexer.compressor.norm.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-position_bias"] = f"{hf_prefix}.attn.indexer.compressor.ape" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-weights_proj-kernel"] = f"{hf_prefix}.attn.indexer.weights_proj.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-q_proj-kernel"] = f"{hf_prefix}.attn.indexer.wq_b.weight" + + # MoE + mapping[f"{prefix}-mlp-MoeBlock_0-gate-kernel"] = f"{hf_prefix}.ffn.gate.weight" + if i < num_hash_layers: + mapping[f"{prefix.replace('params-', 'Tid2EidVar-')}-mlp-MoeBlock_0-tid2eid"] = f"{hf_prefix}.ffn.gate.tid2eid" + else: + mapping[f"{prefix.replace('params-', 'MoEBiasVar-')}-mlp-MoeBlock_0-gate-bias"] = f"{hf_prefix}.ffn.gate.bias" + + mapping[f"{prefix}-mlp-shared_experts-wi_0-kernel"] = f"{hf_prefix}.ffn.shared_experts.w1.weight" + mapping[f"{prefix}-mlp-shared_experts-wi_1-kernel"] = f"{hf_prefix}.ffn.shared_experts.w3.weight" + mapping[f"{prefix}-mlp-shared_experts-wo-kernel"] = f"{hf_prefix}.ffn.shared_experts.w2.weight" + + mapping[f"{prefix}-mlp-MoeBlock_0-wi_0"] = [ + f"{hf_prefix}.ffn.experts.{e}.w1.weight" for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wi_1"] = [ + f"{hf_prefix}.ffn.experts.{e}.w3.weight" for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wo"] = [ + f"{hf_prefix}.ffn.experts.{e}.w2.weight" for e in range(num_experts) + ] - for i in range(3): - mapping[f"Tid2EidVar-decoder-layers_{i}-mlp-MoeBlock_0-tid2eid"] = f"model.layers.{i}.mlp.gate.tid2eid" + else: + # 1. Unrolled Prefix Layers + for i in range(num_hash_layers): + prefix = f"params-decoder-layers_{i}" + hf_prefix = f"layers.{i}" + + mapping[f"{prefix}-pre_self_attention_layer_norm-scale"] = f"{hf_prefix}.attn_norm.weight" + mapping[f"{prefix}-post_self_attention_layer_norm-scale"] = f"{hf_prefix}.ffn_norm.weight" + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = None + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = None + + mapping[f"{prefix}-mhc_attention-pre_alpha"] = f"{hf_prefix}.hc_attn_fn" + mapping[f"{prefix}-mhc_attention-post_alpha"] = f"{hf_prefix}.hc_attn_fn" + mapping[f"{prefix}-mhc_attention-res_alpha"] = f"{hf_prefix}.hc_attn_fn" + + mapping[f"{prefix}-mhc_attention-pre_beta"] = f"{hf_prefix}.hc_attn_base" + mapping[f"{prefix}-mhc_attention-post_beta"] = f"{hf_prefix}.hc_attn_base" + mapping[f"{prefix}-mhc_attention-res_beta"] = f"{hf_prefix}.hc_attn_base" + + mapping[f"{prefix}-mhc_attention-pre_alpha_scale"] = f"{hf_prefix}.hc_attn_scale" + mapping[f"{prefix}-mhc_attention-post_alpha_scale"] = f"{hf_prefix}.hc_attn_scale" + mapping[f"{prefix}-mhc_attention-res_alpha_scale"] = f"{hf_prefix}.hc_attn_scale" + + mapping[f"{prefix}-mhc_mlp-pre_alpha"] = f"{hf_prefix}.hc_ffn_fn" + mapping[f"{prefix}-mhc_mlp-post_alpha"] = f"{hf_prefix}.hc_ffn_fn" + mapping[f"{prefix}-mhc_mlp-res_alpha"] = f"{hf_prefix}.hc_ffn_fn" + + mapping[f"{prefix}-mhc_mlp-pre_beta"] = f"{hf_prefix}.hc_ffn_base" + mapping[f"{prefix}-mhc_mlp-post_beta"] = f"{hf_prefix}.hc_ffn_base" + mapping[f"{prefix}-mhc_mlp-res_beta"] = f"{hf_prefix}.hc_ffn_base" + + mapping[f"{prefix}-mhc_mlp-pre_alpha_scale"] = f"{hf_prefix}.hc_ffn_scale" + mapping[f"{prefix}-mhc_mlp-post_alpha_scale"] = f"{hf_prefix}.hc_ffn_scale" + mapping[f"{prefix}-mhc_mlp-res_alpha_scale"] = f"{hf_prefix}.hc_ffn_scale" + + mapping[f"{prefix}-self_attention-q_norm-scale"] = f"{hf_prefix}.attn.q_norm.weight" + mapping[f"{prefix}-self_attention-kv_norm-scale"] = f"{hf_prefix}.attn.kv_norm.weight" + mapping[f"{prefix}-self_attention-wq_a-kernel"] = f"{hf_prefix}.attn.wq_a.weight" + mapping[f"{prefix}-self_attention-wq_b-kernel"] = f"{hf_prefix}.attn.wq_b.weight" + mapping[f"{prefix}-self_attention-wkv-kernel"] = f"{hf_prefix}.attn.wkv.weight" + mapping[f"{prefix}-self_attention-sinks"] = f"{hf_prefix}.attn.attn_sink" + + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = f"{hf_prefix}.attn.wo_a.weight" + mapping[f"{prefix}-self_attention-o_b_proj-kernel"] = f"{hf_prefix}.attn.wo_b.weight" + + mapping[f"{prefix}-self_attention-csa_compressor-gate_proj-kernel"] = f"{hf_prefix}.attn.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-csa_compressor-kv_proj-kernel"] = f"{hf_prefix}.attn.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-csa_compressor-kv_norm-scale"] = f"{hf_prefix}.attn.compressor.norm.weight" + mapping[f"{prefix}-self_attention-csa_compressor-position_bias"] = f"{hf_prefix}.attn.compressor.ape" + + mapping[f"{prefix}-self_attention-hca_compressor-gate_proj-kernel"] = f"{hf_prefix}.attn.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-hca_compressor-kv_proj-kernel"] = f"{hf_prefix}.attn.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-hca_compressor-kv_norm-scale"] = f"{hf_prefix}.attn.compressor.norm.weight" + mapping[f"{prefix}-self_attention-hca_compressor-position_bias"] = f"{hf_prefix}.attn.compressor.ape" + + mapping[f"{prefix}-self_attention-csa_compressor-indexer-gate_proj-kernel"] = f"{hf_prefix}.attn.indexer.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_proj-kernel"] = f"{hf_prefix}.attn.indexer.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_norm-scale"] = f"{hf_prefix}.attn.indexer.compressor.norm.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-position_bias"] = f"{hf_prefix}.attn.indexer.compressor.ape" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-weights_proj-kernel"] = f"{hf_prefix}.attn.indexer.weights_proj.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-q_proj-kernel"] = f"{hf_prefix}.attn.indexer.wq_b.weight" + + mapping[f"{prefix}-mlp-MoeBlock_0-gate-kernel"] = f"{hf_prefix}.ffn.gate.weight" + mapping[f"{prefix.replace('params-', 'Tid2EidVar-')}-mlp-MoeBlock_0-tid2eid"] = f"{hf_prefix}.ffn.gate.tid2eid" + + mapping[f"{prefix}-mlp-shared_experts-wi_0-kernel"] = f"{hf_prefix}.ffn.shared_experts.w1.weight" + mapping[f"{prefix}-mlp-shared_experts-wi_1-kernel"] = f"{hf_prefix}.ffn.shared_experts.w3.weight" + mapping[f"{prefix}-mlp-shared_experts-wo-kernel"] = f"{hf_prefix}.ffn.shared_experts.w2.weight" + + mapping[f"{prefix}-mlp-MoeBlock_0-wi_0"] = [ + f"{hf_prefix}.ffn.experts.{e}.w1.weight" for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wi_1"] = [ + f"{hf_prefix}.ffn.experts.{e}.w3.weight" for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wo"] = [ + f"{hf_prefix}.ffn.experts.{e}.w2.weight" for e in range(num_experts) + ] - return mapping + # 2. Scanned Blocks + hca_layers = list(range(num_hash_layers, n_layers, 2)) + csa_layers = list(range(num_hash_layers + 1, n_layers, 2)) + + # Layer 0 in Scanned Blocks (HCA) + prefix = "params-decoder-scanned_blocks-layers_0" + mapping[f"{prefix}-pre_self_attention_layer_norm-scale"] = [f"layers.{i}.attn_norm.weight" for i in hca_layers] + mapping[f"{prefix}-post_self_attention_layer_norm-scale"] = [f"layers.{i}.ffn_norm.weight" for i in hca_layers] + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = None + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = None + + mapping[f"{prefix}-mhc_attention-pre_alpha"] = [f"layers.{i}.hc_attn_fn" for i in hca_layers] + mapping[f"{prefix}-mhc_attention-post_alpha"] = [f"layers.{i}.hc_attn_fn" for i in hca_layers] + mapping[f"{prefix}-mhc_attention-res_alpha"] = [f"layers.{i}.hc_attn_fn" for i in hca_layers] + + mapping[f"{prefix}-mhc_attention-pre_beta"] = [f"layers.{i}.hc_attn_base" for i in hca_layers] + mapping[f"{prefix}-mhc_attention-post_beta"] = [f"layers.{i}.hc_attn_base" for i in hca_layers] + mapping[f"{prefix}-mhc_attention-res_beta"] = [f"layers.{i}.hc_attn_base" for i in hca_layers] + + mapping[f"{prefix}-mhc_attention-pre_alpha_scale"] = [f"layers.{i}.hc_attn_scale" for i in hca_layers] + mapping[f"{prefix}-mhc_attention-post_alpha_scale"] = [f"layers.{i}.hc_attn_scale" for i in hca_layers] + mapping[f"{prefix}-mhc_attention-res_alpha_scale"] = [f"layers.{i}.hc_attn_scale" for i in hca_layers] + + mapping[f"{prefix}-mhc_mlp-pre_alpha"] = [f"layers.{i}.hc_ffn_fn" for i in hca_layers] + mapping[f"{prefix}-mhc_mlp-post_alpha"] = [f"layers.{i}.hc_ffn_fn" for i in hca_layers] + mapping[f"{prefix}-mhc_mlp-res_alpha"] = [f"layers.{i}.hc_ffn_fn" for i in hca_layers] + + mapping[f"{prefix}-mhc_mlp-pre_beta"] = [f"layers.{i}.hc_ffn_base" for i in hca_layers] + mapping[f"{prefix}-mhc_mlp-post_beta"] = [f"layers.{i}.hc_ffn_base" for i in hca_layers] + mapping[f"{prefix}-mhc_mlp-res_beta"] = [f"layers.{i}.hc_ffn_base" for i in hca_layers] + + mapping[f"{prefix}-mhc_mlp-pre_alpha_scale"] = [f"layers.{i}.hc_ffn_scale" for i in hca_layers] + mapping[f"{prefix}-mhc_mlp-post_alpha_scale"] = [f"layers.{i}.hc_ffn_scale" for i in hca_layers] + mapping[f"{prefix}-mhc_mlp-res_alpha_scale"] = [f"layers.{i}.hc_ffn_scale" for i in hca_layers] + + mapping[f"{prefix}-self_attention-q_norm-scale"] = [f"layers.{i}.attn.q_norm.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-kv_norm-scale"] = [f"layers.{i}.attn.kv_norm.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-wq_a-kernel"] = [f"layers.{i}.attn.wq_a.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-wq_b-kernel"] = [f"layers.{i}.attn.wq_b.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-wkv-kernel"] = [f"layers.{i}.attn.wkv.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-sinks"] = [f"layers.{i}.attn.attn_sink" for i in hca_layers] + + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = [f"layers.{i}.attn.wo_a.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-o_b_proj-kernel"] = [f"layers.{i}.attn.wo_b.weight" for i in hca_layers] + + mapping[f"{prefix}-self_attention-hca_compressor-gate_proj-kernel"] = [f"layers.{i}.attn.compressor.wgate.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-hca_compressor-kv_proj-kernel"] = [f"layers.{i}.attn.compressor.wkv.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-hca_compressor-kv_norm-scale"] = [f"layers.{i}.attn.compressor.norm.weight" for i in hca_layers] + mapping[f"{prefix}-self_attention-hca_compressor-position_bias"] = [f"layers.{i}.attn.compressor.ape" for i in hca_layers] + + mapping[f"{prefix}-mlp-MoeBlock_0-gate-kernel"] = [f"layers.{i}.ffn.gate.weight" for i in hca_layers] + mapping[f"{prefix.replace('params-', 'MoEBiasVar-')}-mlp-MoeBlock_0-gate-bias"] = [f"layers.{i}.ffn.gate.bias" for i in hca_layers] + + mapping[f"{prefix}-mlp-shared_experts-wi_0-kernel"] = [f"layers.{i}.ffn.shared_experts.w1.weight" for i in hca_layers] + mapping[f"{prefix}-mlp-shared_experts-wi_1-kernel"] = [f"layers.{i}.ffn.shared_experts.w3.weight" for i in hca_layers] + mapping[f"{prefix}-mlp-shared_experts-wo-kernel"] = [f"layers.{i}.ffn.shared_experts.w2.weight" for i in hca_layers] + + mapping[f"{prefix}-mlp-MoeBlock_0-wi_0"] = [ + [f"layers.{i}.ffn.experts.{e}.w1.weight" for i in hca_layers] for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wi_1"] = [ + [f"layers.{i}.ffn.experts.{e}.w3.weight" for i in hca_layers] for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wo"] = [ + [f"layers.{i}.ffn.experts.{e}.w2.weight" for i in hca_layers] for e in range(num_experts) + ] + # Layer 1 in Scanned Blocks (CSA) + prefix = "params-decoder-scanned_blocks-layers_1" + mapping[f"{prefix}-pre_self_attention_layer_norm-scale"] = [f"layers.{i}.attn_norm.weight" for i in csa_layers] + mapping[f"{prefix}-post_self_attention_layer_norm-scale"] = [f"layers.{i}.ffn_norm.weight" for i in csa_layers] + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = None + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = None + + mapping[f"{prefix}-mhc_attention-pre_alpha"] = [f"layers.{i}.hc_attn_fn" for i in csa_layers] + mapping[f"{prefix}-mhc_attention-post_alpha"] = [f"layers.{i}.hc_attn_fn" for i in csa_layers] + mapping[f"{prefix}-mhc_attention-res_alpha"] = [f"layers.{i}.hc_attn_fn" for i in csa_layers] + + mapping[f"{prefix}-mhc_attention-pre_beta"] = [f"layers.{i}.hc_attn_base" for i in csa_layers] + mapping[f"{prefix}-mhc_attention-post_beta"] = [f"layers.{i}.hc_attn_base" for i in csa_layers] + mapping[f"{prefix}-mhc_attention-res_beta"] = [f"layers.{i}.hc_attn_base" for i in csa_layers] + + mapping[f"{prefix}-mhc_attention-pre_alpha_scale"] = [f"layers.{i}.hc_attn_scale" for i in csa_layers] + mapping[f"{prefix}-mhc_attention-post_alpha_scale"] = [f"layers.{i}.hc_attn_scale" for i in csa_layers] + mapping[f"{prefix}-mhc_attention-res_alpha_scale"] = [f"layers.{i}.hc_attn_scale" for i in csa_layers] + + mapping[f"{prefix}-mhc_mlp-pre_alpha"] = [f"layers.{i}.hc_ffn_fn" for i in csa_layers] + mapping[f"{prefix}-mhc_mlp-post_alpha"] = [f"layers.{i}.hc_ffn_fn" for i in csa_layers] + mapping[f"{prefix}-mhc_mlp-res_alpha"] = [f"layers.{i}.hc_ffn_fn" for i in csa_layers] + + mapping[f"{prefix}-mhc_mlp-pre_beta"] = [f"layers.{i}.hc_ffn_base" for i in csa_layers] + mapping[f"{prefix}-mhc_mlp-post_beta"] = [f"layers.{i}.hc_ffn_base" for i in csa_layers] + mapping[f"{prefix}-mhc_mlp-res_beta"] = [f"layers.{i}.hc_ffn_base" for i in csa_layers] + + mapping[f"{prefix}-mhc_mlp-pre_alpha_scale"] = [f"layers.{i}.hc_ffn_scale" for i in csa_layers] + mapping[f"{prefix}-mhc_mlp-post_alpha_scale"] = [f"layers.{i}.hc_ffn_scale" for i in csa_layers] + mapping[f"{prefix}-mhc_mlp-res_alpha_scale"] = [f"layers.{i}.hc_ffn_scale" for i in csa_layers] + + mapping[f"{prefix}-self_attention-q_norm-scale"] = [f"layers.{i}.attn.q_norm.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-kv_norm-scale"] = [f"layers.{i}.attn.kv_norm.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-wq_a-kernel"] = [f"layers.{i}.attn.wq_a.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-wq_b-kernel"] = [f"layers.{i}.attn.wq_b.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-wkv-kernel"] = [f"layers.{i}.attn.wkv.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-sinks"] = [f"layers.{i}.attn.attn_sink" for i in csa_layers] + + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = [f"layers.{i}.attn.wo_a.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-o_b_proj-kernel"] = [f"layers.{i}.attn.wo_b.weight" for i in csa_layers] + + mapping[f"{prefix}-self_attention-csa_compressor-gate_proj-kernel"] = [f"layers.{i}.attn.compressor.wgate.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-kv_proj-kernel"] = [f"layers.{i}.attn.compressor.wkv.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-kv_norm-scale"] = [f"layers.{i}.attn.compressor.norm.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-position_bias"] = [f"layers.{i}.attn.compressor.ape" for i in csa_layers] + + mapping[f"{prefix}-self_attention-csa_compressor-indexer-gate_proj-kernel"] = [f"layers.{i}.attn.indexer.compressor.wgate.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_proj-kernel"] = [f"layers.{i}.attn.indexer.compressor.wkv.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_norm-scale"] = [f"layers.{i}.attn.indexer.compressor.norm.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-indexer-position_bias"] = [f"layers.{i}.attn.indexer.compressor.ape" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-indexer-weights_proj-kernel"] = [f"layers.{i}.attn.indexer.weights_proj.weight" for i in csa_layers] + mapping[f"{prefix}-self_attention-csa_compressor-indexer-q_proj-kernel"] = [f"layers.{i}.attn.indexer.wq_b.weight" for i in csa_layers] + + mapping[f"{prefix}-mlp-MoeBlock_0-gate-kernel"] = [f"layers.{i}.ffn.gate.weight" for i in csa_layers] + mapping[f"{prefix.replace('params-', 'MoEBiasVar-')}-mlp-MoeBlock_0-gate-bias"] = [f"layers.{i}.ffn.gate.bias" for i in csa_layers] + + mapping[f"{prefix}-mlp-shared_experts-wi_0-kernel"] = [f"layers.{i}.ffn.shared_experts.w1.weight" for i in csa_layers] + mapping[f"{prefix}-mlp-shared_experts-wi_1-kernel"] = [f"layers.{i}.ffn.shared_experts.w3.weight" for i in csa_layers] + mapping[f"{prefix}-mlp-shared_experts-wo-kernel"] = [f"layers.{i}.ffn.shared_experts.w2.weight" for i in csa_layers] + + mapping[f"{prefix}-mlp-MoeBlock_0-wi_0"] = [ + [f"layers.{i}.ffn.experts.{e}.w1.weight" for i in csa_layers] for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wi_1"] = [ + [f"layers.{i}.ffn.experts.{e}.w3.weight" for i in csa_layers] for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wo"] = [ + [f"layers.{i}.ffn.experts.{e}.w2.weight" for i in csa_layers] for e in range(num_experts) + ] -def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): - """Returns hook functions for transforming weights between MaxText and HuggingFace for DeepSeek V4.""" + return mapping - def transpose(input_tensor, target_shape=None): - return np.transpose(input_tensor) - def ones_norm(input_tensor, target_shape=None): - return np.ones(target_shape, dtype=np.float32) # pyrefly: ignore[no-matching-overload] +def DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): + def _get(cfg, key, default=None): + if isinstance(cfg, dict): + return cfg.get(key, default) + return getattr(cfg, key, default) - def identity(input_tensor, target_shape=None): - return input_tensor + n_layers = _get(config, "num_hidden_layers", getattr(maxtext_config, "base_num_decoder_layers", 43)) + num_hash_layers = _get(config, "num_hash_layers", getattr(maxtext_config, "first_num_hash_layers", 3)) - # Reshaping functions for wq_b, wkv, o_a_proj - def reshape_transpose_wq_b(input_tensor, target_shape=None): - # HF: [n_heads * q_head_dim, kv_lora_rank] - # MaxText: [kv_lora_rank, n_heads, q_head_dim] - if saving_to_hf: - tensor = input_tensor.reshape((input_tensor.shape[0], -1)) - return np.transpose(tensor) - tensor = np.transpose(input_tensor) # [kv_lora_rank, n_heads * q_head_dim] - return tensor.reshape(target_shape) - - def reshape_transpose_wkv(input_tensor, target_shape=None): - # HF: [n_kv_heads * (q_head_dim + v_head_dim), kv_lora_rank] - # MaxText: [kv_lora_rank, n_kv_heads, q_head_dim + v_head_dim] - if saving_to_hf: - tensor = input_tensor.reshape((input_tensor.shape[0], -1)) - return np.transpose(tensor) - tensor = np.transpose(input_tensor) - return tensor.reshape(target_shape) - - def reshape_transpose_o_a(input_tensor, target_shape=None): - # HF: [n_heads * v_head_dim, kv_lora_rank] (e.g. [8192, 4096]) - # MaxText: [n_heads, v_head_dim, kv_lora_rank] (e.g. [8, 4096, 1024]) - # We must reshape first and then permute (transpose) to get correct ordering. + def reshape_kernel(input_tensor, target_shape): + import numpy as np if saving_to_hf: - tensor = np.transpose(input_tensor, (0, 2, 1)) - return tensor.reshape(target_shape) - num_heads = target_shape[0] # pyrefly: ignore[unsupported-operation] - embed_dim = target_shape[1] # pyrefly: ignore[unsupported-operation] - kv_lora_rank = target_shape[2] # pyrefly: ignore[unsupported-operation] - tensor = input_tensor.reshape((num_heads, kv_lora_rank, embed_dim)) - return np.transpose(tensor, (0, 2, 1)) - - # Functions for mHC split - def mhc_split_fn_pre(input_tensor, target_shape=None): - return np.transpose(input_tensor[0:4, :]) - - def mhc_split_fn_post(input_tensor, target_shape=None): - return np.transpose(input_tensor[4:8, :]) - - def mhc_split_fn_res(input_tensor, target_shape=None): - return np.transpose(input_tensor[8:24, :]) - - def mhc_split_base_pre(input_tensor, target_shape=None): - return input_tensor[0:4] - - def mhc_split_base_post(input_tensor, target_shape=None): - return input_tensor[4:8] - - def mhc_split_base_res(input_tensor, target_shape=None): - return input_tensor[8:24].reshape(target_shape) + flipped_target_shape = np.flip(np.array(target_shape)) + return input_tensor.reshape(flipped_target_shape).T + else: + try: + return input_tensor.T.reshape(target_shape) + except Exception as e: + raise ValueError(f'RESHAPE FAILED: {str(e)} for target {target_shape} with input {input_tensor.shape}') - def mhc_split_scale_pre(input_tensor, target_shape=None): - return np.array([input_tensor[0]]).reshape(target_shape) + def reshape_o_a_proj(input_tensor, target_shape): + if saving_to_hf: + heads, head_dim, in_dim = input_tensor.shape[0], input_tensor.shape[1], input_tensor.shape[2] + if target_shape[0] == in_dim: + return input_tensor.transpose(2, 0, 1).reshape(target_shape) + else: + return input_tensor.transpose(0, 2, 1).reshape(target_shape) + else: + heads, head_dim, in_dim = target_shape[0], target_shape[1], target_shape[2] + if input_tensor.shape[0] == in_dim: + return input_tensor.reshape(in_dim, heads, head_dim).transpose(1, 2, 0) + else: + return input_tensor.reshape(heads, in_dim, head_dim).transpose(0, 2, 1) + + # Alpha Hooks + def reshape_mhc_pre(input_tensor, target_shape): + transposed = input_tensor.T + mix_hc = transposed.shape[-1] + k = int(np.sqrt(1 + mix_hc) - 1) + return transposed[..., :k] + + def reshape_mhc_post(input_tensor, target_shape): + transposed = input_tensor.T + mix_hc = transposed.shape[-1] + k = int(np.sqrt(1 + mix_hc) - 1) + return transposed[..., k:2*k] + + def reshape_mhc_res(input_tensor, target_shape): + transposed = input_tensor.T + mix_hc = transposed.shape[-1] + k = int(np.sqrt(1 + mix_hc) - 1) + return transposed[..., 2*k:] + + def composite_mhc_alpha(weights, target_shape=None): + pre_alpha, post_alpha, res_alpha = weights + return np.concatenate([pre_alpha, post_alpha, res_alpha], axis=-1).T + + # Beta Hooks + def reshape_mhc_beta_pre(input_tensor, target_shape): + mix_hc = input_tensor.shape[-1] + k = int(np.sqrt(1 + mix_hc) - 1) + return input_tensor[:k] + + def reshape_mhc_beta_post(input_tensor, target_shape): + mix_hc = input_tensor.shape[-1] + k = int(np.sqrt(1 + mix_hc) - 1) + return input_tensor[k:2*k] + + def reshape_mhc_beta_res(input_tensor, target_shape): + mix_hc = input_tensor.shape[-1] + k = int(np.sqrt(1 + mix_hc) - 1) + return input_tensor[2*k:].reshape(k, k) + + def composite_mhc_beta(weights, target_shape=None): + pre_beta, post_beta, res_beta = weights + return np.concatenate([pre_beta, post_beta, res_beta.flatten()], axis=-1) + + # Scale Hooks + def reshape_mhc_scale_pre(input_tensor, target_shape): + return input_tensor[0:1] + + def reshape_mhc_scale_post(input_tensor, target_shape): + return input_tensor[1:2] + + def reshape_mhc_scale_res(input_tensor, target_shape): + return input_tensor[2:3] + + def composite_mhc_scale(weights, target_shape=None): + pre_scale, post_scale, res_scale = weights + return np.concatenate([pre_scale, post_scale, res_scale], axis=-1) + + def mhc_dummy_norm(input_tensor, target_shape=None): + import numpy as np + dtype = input_tensor.dtype if input_tensor is not None else np.float32 + return np.ones(target_shape, dtype=dtype) + + 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 mhc_split_scale_post(input_tensor, target_shape=None): - return np.array([input_tensor[1]]).reshape(target_shape) + 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 mhc_split_scale_res(input_tensor, target_shape=None): - return np.array([input_tensor[2]]).reshape(target_shape) + mapping = { + "params-token_embedder-embedding": unpad_hf_embedding_layer, + "params-decoder-logits_dense-kernel": unpad_logits_layer, + "params-decoder-hc_head-hc_fn": reshape_kernel, + } - mapping = {} + def _attach_layer_hooks(prefix): + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = reshape_o_a_proj + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = mhc_dummy_norm + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = mhc_dummy_norm + + for key in [ + f"{prefix}-self_attention-wq_a-kernel", + f"{prefix}-self_attention-wq_b-kernel", + f"{prefix}-self_attention-wkv-kernel", + f"{prefix}-self_attention-o_b_proj-kernel", + f"{prefix}-mlp-MoeBlock_0-gate-kernel", + f"{prefix}-mlp-MoeBlock_0-wi_0", + f"{prefix}-mlp-MoeBlock_0-wi_1", + f"{prefix}-mlp-MoeBlock_0-wo", + f"{prefix}-mlp-shared_experts-wi_0-kernel", + f"{prefix}-mlp-shared_experts-wi_1-kernel", + f"{prefix}-mlp-shared_experts-wo-kernel", + f"{prefix}-self_attention-csa_compressor-gate_proj-kernel", + f"{prefix}-self_attention-csa_compressor-kv_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-gate_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-kv_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-weights_proj-kernel", + f"{prefix}-self_attention-csa_compressor-indexer-q_proj-kernel", + f"{prefix}-self_attention-hca_compressor-gate_proj-kernel", + f"{prefix}-self_attention-hca_compressor-kv_proj-kernel", + ]: + mapping[key] = reshape_kernel + + if saving_to_hf: + mapping[(f"{prefix}-mhc_attention-pre_alpha", f"{prefix}-mhc_attention-post_alpha", f"{prefix}-mhc_attention-res_alpha")] = composite_mhc_alpha + mapping[(f"{prefix}-mhc_mlp-pre_alpha", f"{prefix}-mhc_mlp-post_alpha", f"{prefix}-mhc_mlp-res_alpha")] = composite_mhc_alpha + mapping[(f"{prefix}-mhc_attention-pre_beta", f"{prefix}-mhc_attention-post_beta", f"{prefix}-mhc_attention-res_beta")] = composite_mhc_beta + mapping[(f"{prefix}-mhc_mlp-pre_beta", f"{prefix}-mhc_mlp-post_beta", f"{prefix}-mhc_mlp-res_beta")] = composite_mhc_beta + mapping[(f"{prefix}-mhc_attention-pre_alpha_scale", f"{prefix}-mhc_attention-post_alpha_scale", f"{prefix}-mhc_attention-res_alpha_scale")] = composite_mhc_scale + mapping[(f"{prefix}-mhc_mlp-pre_alpha_scale", f"{prefix}-mhc_mlp-post_alpha_scale", f"{prefix}-mhc_mlp-res_alpha_scale")] = composite_mhc_scale + else: + mapping[f"{prefix}-mhc_attention-pre_alpha"] = reshape_mhc_pre + mapping[f"{prefix}-mhc_attention-post_alpha"] = reshape_mhc_post + mapping[f"{prefix}-mhc_attention-res_alpha"] = reshape_mhc_res + mapping[f"{prefix}-mhc_mlp-pre_alpha"] = reshape_mhc_pre + mapping[f"{prefix}-mhc_mlp-post_alpha"] = reshape_mhc_post + mapping[f"{prefix}-mhc_mlp-res_alpha"] = reshape_mhc_res + mapping[f"{prefix}-mhc_attention-pre_beta"] = reshape_mhc_beta_pre + mapping[f"{prefix}-mhc_attention-post_beta"] = reshape_mhc_beta_post + mapping[f"{prefix}-mhc_attention-res_beta"] = reshape_mhc_beta_res + mapping[f"{prefix}-mhc_mlp-pre_beta"] = reshape_mhc_beta_pre + mapping[f"{prefix}-mhc_mlp-post_beta"] = reshape_mhc_beta_post + mapping[f"{prefix}-mhc_mlp-res_beta"] = reshape_mhc_beta_res + + mapping[f"{prefix}-mhc_attention-pre_alpha_scale"] = reshape_mhc_scale_pre + mapping[f"{prefix}-mhc_attention-post_alpha_scale"] = reshape_mhc_scale_post + mapping[f"{prefix}-mhc_attention-res_alpha_scale"] = reshape_mhc_scale_res + mapping[f"{prefix}-mhc_mlp-pre_alpha_scale"] = reshape_mhc_scale_pre + mapping[f"{prefix}-mhc_mlp-post_alpha_scale"] = reshape_mhc_scale_post + mapping[f"{prefix}-mhc_mlp-res_alpha_scale"] = reshape_mhc_scale_res - # Base mapping logic from original file - for key, hf_key in DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers).items(): - if hf_key is None: - mapping[key] = ones_norm - elif "token_embedder-embedding" in key: - mapping[key] = identity - elif "-wkv-kernel" in key: - mapping[key] = reshape_transpose_wkv - elif "-wq_b-kernel" in key: - mapping[key] = reshape_transpose_wq_b - elif "-o_a_proj-kernel" in key: - mapping[key] = reshape_transpose_o_a - elif "mhc" in key: - if "pre_alpha" in key and "scale" not in key: - mapping[key] = mhc_split_fn_pre - elif "post_alpha" in key and "scale" not in key: - mapping[key] = mhc_split_fn_post - elif "res_alpha" in key and "scale" not in key: - mapping[key] = mhc_split_fn_res - elif "pre_beta" in key: - mapping[key] = mhc_split_base_pre - elif "post_beta" in key: - mapping[key] = mhc_split_base_post - elif "res_beta" in key: - mapping[key] = mhc_split_base_res - elif "pre_alpha_scale" in key: - mapping[key] = mhc_split_scale_pre - elif "post_alpha_scale" in key: - mapping[key] = mhc_split_scale_post - elif "res_alpha_scale" in key: - mapping[key] = mhc_split_scale_res - elif "position_bias" in key: - mapping[key] = identity - elif "hc_head-hc_fn" in key: - mapping[key] = transpose - elif "hc_head-hc_base" in key or "hc_head-hc_scale" in key: - mapping[key] = identity - elif isinstance(hf_key, list): - mapping[key] = transpose - elif "-kernel" in key or "-embedding" in key or "-sinks" in key: - mapping[key] = transpose - - if saving_to_hf: - - def mhc_concat_fn(input_tensors, target_shape=None): - if len(input_tensors) != 3: - raise ValueError(f"mhc_concat_fn expected 3 tensors (pre, post, res), got {len(input_tensors)}") - tensors = [np.asarray(t) for t in input_tensors] - res = np.transpose(np.concatenate(tensors, axis=1)) - return res.reshape(target_shape) if target_shape is not None else res - - def mhc_concat_base(input_tensors, target_shape=None): - if len(input_tensors) != 3: - raise ValueError(f"mhc_concat_base expected 3 tensors (pre, post, res), got {len(input_tensors)}") - tensors = [np.asarray(t).ravel() for t in input_tensors] - res = np.concatenate(tensors, axis=0) - return res.reshape(target_shape) if target_shape is not None else res - - def mhc_concat_scale(input_tensors, target_shape=None): - if len(input_tensors) != 3: - raise ValueError(f"mhc_concat_scale expected 3 tensors (pre, post, res), got {len(input_tensors)}") - tensors = [np.asarray(t).ravel() for t in input_tensors] - res = np.concatenate(tensors, axis=0) - return res.reshape(target_shape) if target_shape is not None else res - - # Process composite mappings - keys_to_delete = [] - keys_to_add = {} - for key in list(mapping.keys()): - if "mhc" in key and "pre_alpha" in key and "scale" not in key: - post = key.replace("pre_alpha", "post_alpha") - res = key.replace("pre_alpha", "res_alpha") - keys_to_delete.extend([key, post, res]) - keys_to_add[(key, post, res)] = mhc_concat_fn - - if "mhc" in key and "pre_beta" in key: - post = key.replace("pre_beta", "post_beta") - res = key.replace("pre_beta", "res_beta") - keys_to_delete.extend([key, post, res]) - keys_to_add[(key, post, res)] = mhc_concat_base - - if "mhc" in key and "pre_alpha_scale" in key: - post = key.replace("pre_alpha_scale", "post_alpha_scale") - res = key.replace("pre_alpha_scale", "res_alpha_scale") - keys_to_delete.extend([key, post, res]) - keys_to_add[(key, post, res)] = mhc_concat_scale - - for k in set(keys_to_delete): - if k in mapping: - del mapping[k] - mapping.update(keys_to_add) + if not scan_layers: + for i in range(n_layers): + _attach_layer_hooks(f"params-decoder-layers_{i}") + else: + for i in range(num_hash_layers): + _attach_layer_hooks(f"params-decoder-layers_{i}") + _attach_layer_hooks("params-decoder-scanned_blocks-layers_0") + _attach_layer_hooks("params-decoder-scanned_blocks-layers_1") return mapping - PARAM_MAPPING = { "gemma2-2b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, "gemma2-9b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4256,7 +4496,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek2-16b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING, - "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING, + "deepseek4-284b": DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-20b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "gpt-oss-120b": GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4310,8 +4550,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "deepseek2-16b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek3-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, "deepseek3.2-671b": DEEPSEEK_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "deepseek4-tiny": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, - "deepseek4-284b": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "deepseek4-284b": DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN, "gpt-oss-20b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "gpt-oss-120b": GPT_OSS_TO_HF_PARAM_HOOK_FN, "qwen3-omni-30b-a3b": QWEN3_OMNI_MOE_MAXTEXT_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/configs/models/deepseek4-284b.yml b/src/maxtext/configs/models/deepseek4-284b.yml index 5689114145..6fe8a5a990 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/src/maxtext/configs/models/deepseek4-284b.yml @@ -22,6 +22,7 @@ base_mlp_dim: 2048 base_moe_mlp_dim: 2048 vocab_size: 129280 head_dim: 512 +qk_nope_head_dim: 512 qk_rope_head_dim: 64 # --- Standard Defaults --- @@ -59,6 +60,7 @@ routed_scaling_factor: 1.5 # --- Attention configuration --- +attention: "dot_product" attention_type: 'compressed' q_lora_rank: 1024 o_groups: 8 diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 81beb97020..7469c4047a 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -1021,14 +1021,14 @@ def generate_attention_mask( s_len = kv_seq_len - c_len def get_sliding_mask(s_len): - # Safely use segment_positions, or fall back to next_pos if None - if segment_positions is not None: - abs_q = segment_positions[:, :, None] - else: - local_next = next_pos[:, None] if isinstance(next_pos, jax.Array) else next_pos - abs_q = jnp.arange(q_seq_len)[None, :, None] + local_next - if model_mode == MODEL_MODE_AUTOREGRESSIVE and q_seq_len == 1: + # Safely use segment_positions, or fall back to next_pos if None + if segment_positions is not None: + abs_q = segment_positions[:, :, None] + else: + local_next = next_pos[:, None] if isinstance(next_pos, jax.Array) else next_pos + abs_q = jnp.arange(q_seq_len)[None, :, None] + local_next + if decoder_segment_ids is not None: is_valid = decoder_segment_ids[:, :s_len] == DECODING_ACTIVE_SEQUENCE_INDICATOR valid_indices = jnp.where(is_valid, jnp.arange(s_len)[None, :], -1) @@ -1042,13 +1042,18 @@ def get_sliding_mask(s_len): abs_k_prefill = jnp.broadcast_to(i, abs_k_ar.shape) abs_k = jnp.where(is_ar_cache, abs_k_ar, abs_k_prefill) - distance = abs_q - abs_k - in_window = (distance < self.sliding_window_size) if self.sliding_window_size is not None else True - return in_window & (distance >= 0) - - # For prefill and training phases (q_seq_len > 1) - abs_k = jnp.arange(s_len)[None, None, :] - distance = abs_q - abs_k + else: + abs_k = jnp.arange(s_len)[None, None, :] + distance = abs_q - abs_k + in_window = (distance < self.sliding_window_size) if self.sliding_window_size is not None else True + return in_window & (distance >= 0) + + # For prefill and training phases (q_seq_len > 1): + # Use global buffer coordinates so causal/sliding mask is valid across packed sequences + 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 + col_ids = jnp.arange(s_len)[None, None, :] + distance = row_ids - col_ids in_window = (distance < self.sliding_window_size) if self.sliding_window_size is not None else True return in_window & (distance >= 0) diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 17c9928831..52d64ed54e 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -4872,5 +4872,179 @@ def test_ulysses_divisible_accepted(self): self.attention._validate_kv_head_sharding(self._KV_KERNEL_AXES) # pylint: disable=protected-access +class DeepSeekV4AttentionMaskingTest(unittest.TestCase): + """Tests to validate AttentionOp masking logic for DeepSeek-V4 attention patterns.""" + + def setUp(self): + self.config = pyconfig.initialize([sys.argv[0], "src/maxtext/configs/base.yml"], run_name="test") + + def test_generate_attention_mask_local_sliding(self): + """Verifies AttentionType.LOCAL_SLIDING enforces both causal and sliding window constraints.""" + + # Test with multiple heads and different sequence lengths + for s_len in [1, 8, 128]: + op = AttentionOp( + config=self.config, + num_query_heads=4, + num_kv_heads=1, + max_target_length=256, + mesh=None, + attention_kernel="dot_product", + attention_type=AttentionType.LOCAL_SLIDING, + sliding_window_size=3, + ) + + batch_size = 1 + q_dummy = jnp.zeros((batch_size, s_len, 1, 128)) + k_dummy = jnp.zeros((batch_size, s_len, 1, 128)) + + mask = op.generate_attention_mask( + query=q_dummy, + key=k_dummy, + decoder_segment_ids=None, + model_mode="train", + ) + + self.assertEqual(mask.shape, (1, 1, 1, s_len, s_len)) + mask_np = np.array(mask)[0, 0, 0] + + # Expected float mask for window_size=3 + # Row 0: [0.0, INF, INF, INF, INF, ...] + # Row 1: [0.0, 0.0, INF, INF, INF, ...] + # Row 2: [0.0, 0.0, 0.0, INF, INF, ...] + # Row 3: [INF, 0.0, 0.0, 0.0, INF, ...] + if s_len > 1: + self.assertEqual(mask_np[0, 1], DEFAULT_MASK_VALUE) # strict causal + self.assertEqual(mask_np[0, 0], 0.0) + + if s_len >= 4: + self.assertEqual(mask_np[3, 0], DEFAULT_MASK_VALUE) # sliding window size=3 + self.assertEqual(mask_np[3, 1], 0.0) + + def test_generate_attention_mask_compressed(self): + """Verifies AttentionType.COMPRESSED stitches sliding window and float compressed_mask.""" + + batch_size = 1 + s_len = 8 + c_len = 2 + kv_len = s_len + c_len + + op = AttentionOp( + config=self.config, + num_query_heads=4, + num_kv_heads=1, + max_target_length=128, + mesh=None, + attention_kernel="dot_product", + attention_type=AttentionType.COMPRESSED, + sliding_window_size=3, + ) + + q_dummy = jnp.zeros((batch_size, s_len, 1, 128)) + k_dummy = jnp.zeros((batch_size, kv_len, 1, 128)) + + # Simulate a compressed float mask [batch, 1, s_len, c_len] + # In practice, this exactly mirrors what both HCA and CSA output: + # - HCA emits a simple mask blocking future blocks (batch, 1, seq_len, c_len) + # - CSA emits a sparse mask where only top-K blocks are 0.0, rest are -inf. + # We simulate this by making Block 0 invalid (-inf), and Block 1 valid (0.0). + compressed_mask = np.zeros((batch_size, 1, s_len, c_len), dtype=np.float32) + compressed_mask[:, :, :, 0] = DEFAULT_MASK_VALUE + compressed_mask = jnp.array(compressed_mask) + + mask = op.generate_attention_mask( + query=q_dummy, + key=k_dummy, + decoder_segment_ids=None, + model_mode="train", + compressed_mask=compressed_mask, + ) + + # Returned float mask should dynamically inherit the dimensionality of compressed_mask + # Because compressed_mask was 4D, the final mask should also be 4D: [batch, 1, s_len, kv_len] + self.assertEqual(mask.shape, (batch_size, 1, s_len, kv_len)) + mask_np = np.array(mask)[0, 0] + + # Uncompressed block (first s_len cols) follows sliding window float mask + self.assertEqual(mask_np[0, 1], DEFAULT_MASK_VALUE) + self.assertEqual(mask_np[0, 0], 0.0) + self.assertEqual(mask_np[3, 0], DEFAULT_MASK_VALUE) + self.assertEqual(mask_np[3, 1], 0.0) + + # Compressed block (last c_len cols) follows compressed_mask strictly + np.testing.assert_allclose(mask_np[:, s_len], DEFAULT_MASK_VALUE) + np.testing.assert_allclose(mask_np[:, s_len + 1], 0.0) + print("Mask logic for uncompressed & compressed attention passed perfectly.") + + def test_generate_attention_mask_compressed_all_modes(self): + """Verifies AttentionType.COMPRESSED across train, prefill, and autoregressive modes.""" + batch_size = 2 + s_len = 8 + c_len = 2 + kv_len = s_len + c_len + + op = AttentionOp( + config=self.config, + num_query_heads=4, + num_kv_heads=1, + max_target_length=128, + mesh=None, + attention_kernel="dot_product", + attention_type=AttentionType.COMPRESSED, + sliding_window_size=3, + ) + + # 1. Training mode (batch_size=2, 4D compressed_mask) + q_train = jnp.zeros((batch_size, s_len, 1, 128)) + k_train = jnp.zeros((batch_size, kv_len, 1, 128)) + c_mask_4d = jnp.zeros((batch_size, 1, s_len, c_len), dtype=jnp.float32) + mask_train = op.generate_attention_mask( + query=q_train, + key=k_train, + decoder_segment_ids=None, + model_mode="train", + compressed_mask=c_mask_4d, + ) + self.assertEqual(mask_train.shape, (batch_size, 1, s_len, kv_len)) + + # 2. Prefill mode (batch_size=2, 5D compressed_mask with segment_positions) + c_mask_5d = jnp.zeros((batch_size, 1, 1, s_len, c_len), dtype=jnp.float32) + seg_pos = jnp.arange(s_len)[None, :].repeat(batch_size, axis=0) + mask_prefill = op.generate_attention_mask( + query=q_train, + key=k_train, + decoder_segment_ids=None, + model_mode="prefill", + compressed_mask=c_mask_5d, + segment_positions=seg_pos, + ) + self.assertEqual(mask_prefill.shape, (batch_size, 1, 1, s_len, kv_len)) + + # 3. Autoregressive mode (q_seq_len=1, batch_size=2, decoder_segment_ids) + q_ar = jnp.zeros((batch_size, 1, 1, 128)) + k_ar = jnp.zeros((batch_size, kv_len, 1, 128)) + c_mask_ar = jnp.zeros((batch_size, 1, 1, c_len), dtype=jnp.float32) + seg_ids = jnp.ones((batch_size, 16), dtype=jnp.int32) + mask_ar = op.generate_attention_mask( + query=q_ar, + key=k_ar, + decoder_segment_ids=seg_ids, + model_mode="autoregressive", + compressed_mask=c_mask_ar, + ) + self.assertEqual(mask_ar.shape, (batch_size, 1, 1, 1, kv_len)) + + # 4. Compressed mask is None (fallback to uncompressed mask shape) + mask_none = op.generate_attention_mask( + query=q_train, + key=k_train, + decoder_segment_ids=None, + model_mode="train", + compressed_mask=None, + ) + self.assertEqual(mask_none.ndim, 4) + self.assertEqual(mask_none.shape[-1], kv_len) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/deepseek_v4_flash_vs_reference_test.py b/tests/unit/deepseek_v4_flash_vs_reference_test.py new file mode 100644 index 0000000000..34a0eb24ee --- /dev/null +++ b/tests/unit/deepseek_v4_flash_vs_reference_test.py @@ -0,0 +1,1581 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests validating DeepSeek-V4 MaxText components against official DeepSeek reference implementation. + +This test uses the official reference architecture from: +https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/blob/main/inference/model.py +and translates GPU TileLang kernels from: +https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash/blob/main/inference/kernel.py +into pure, unoptimized PyTorch CPU operations to verify numerical parity against MaxText implementations. +Parameter conversion is performed using `param_mapping.py`. +""" + +import dataclasses +import math +import sys +import unittest +from typing import Tuple, Optional, List + +import flax.linen as nn +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np +import torch +import torch.nn as nn_pt +import torch.nn.functional as F + +from maxtext.checkpoint_conversion.utils.param_mapping import ( + DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING, + DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN, +) +from maxtext.common.common_types import ( + Config, + DecoderBlockType, + AttentionType, + MODEL_MODE_TRAIN, + ShardMode, +) +from maxtext.configs import pyconfig +from maxtext.layers import initializers +from maxtext.layers.attention_compressed import CompressedAttention +from maxtext.layers.embeddings import Embed, DeepSeekV4RotaryEmbedding +from maxtext.layers.linears import DeepSeekV4GroupedLinear +from maxtext.layers.mhc import DeepSeek4HyperHead +from maxtext.layers.moe import RoutedMoE, RoutedAndSharedMoE +from maxtext.layers.normalizations import RMSNorm +from maxtext.models.deepseek4 import DeepSeek4DecoderLayer + + +# ============================================================================== +# Helper to create valid MaxText configuration for DeepSeek-V4 +# ============================================================================== + +def get_maxtext_config(**overrides) -> Config: + config_arguments = { + "model_name": "deepseek4-284b", + "override_model_config": True, + "per_device_batch_size": 1, + "matmul_precision": "highest", + "megablox": False, + "sparse_matmul": False, + "dtype": "float32", + "weight_dtype": "float32", + "base_num_decoder_layers": 5, + "num_experts": 4, + "num_experts_per_tok": 2, + "compress_ratios": [0, 0, 4, 128, 4], + "mtp_num_layers": 0, + "max_target_length": 32, + "max_prefill_predict_length": 32, + "skip_jax_distributed_system": True, + "scan_layers": False, + } + config_arguments.update(overrides) + argv = [sys.argv[0], "src/maxtext/configs/base.yml"] + return pyconfig.initialize(argv, **config_arguments) + + +# ============================================================================== +# 1. Official DeepSeek ModelArgs & Reference Config +# ============================================================================== + +@dataclasses.dataclass +class ModelArgs: + max_batch_size: int = 8 + max_seq_len: int = 4096 + dtype: str = "float32" + vocab_size: int = 129280 + dim: int = 4096 + inter_dim: int = 2048 + moe_inter_dim: int = 2048 + n_layers: int = 5 + n_dense_layers: int = 0 + n_heads: int = 64 + n_routed_experts: int = 4 + n_shared_experts: int = 1 + n_activated_experts: int = 2 + n_expert_groups: int = 1 + n_limited_groups: int = 1 + score_func: str = "sqrtsoftplus" + route_scale: float = 1.5 + q_lora_rank: int = 1024 + kv_lora_rank: int = 512 + o_lora_rank: int = 1024 + o_groups: int = 8 + head_dim: int = 512 + qk_nope_head_dim: int = 512 + qk_rope_head_dim: int = 64 + v_head_dim: int = 512 + sliding_window: int = 128 + index_topk: int = 512 + index_n_heads: int = 64 + index_head_dim: int = 128 + index_block_size: int = 1 + index_routing_dim: int = 32 + index_q_lora_rank: int = 16 + original_seq_len: int = 65536 + rope_theta: float = 10000.0 + rope_factor: float = 40.0 + beta_fast: int = 32 + beta_slow: int = 1 + mscale: float = 1.0 + hc_mult: int = 4 + hc_sinkhorn_iters: int = 20 + hc_eps: float = 1e-6 + norm_eps: float = 1e-6 + compress_rates: List[int] = dataclasses.field(default_factory=lambda: [0, 0, 4, 128, 4]) + n_hash_layers: int = 3 + n_mtp_layers: int = 0 + + +# ============================================================================== +# 2. CPU Reference Kernels (Pure PyTorch) +# ============================================================================== + +def precompute_freqs_cis(dim: int, end: int = 4096, theta: float = 10000.0) -> torch.Tensor: + """Computes RoPE frequency cis representation.""" + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + t = torch.arange(end, dtype=torch.float32) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + return freqs_cis + + +def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + """Applies interleaved rotary embedding to trailing dimensions.""" + dtype = x.dtype + x_c = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + while freqs_cis.ndim < x_c.ndim: + if freqs_cis.ndim == x_c.ndim - 1 and freqs_cis.shape[0] == x_c.shape[1]: + freqs_cis = freqs_cis.unsqueeze(0) + else: + freqs_cis = freqs_cis.unsqueeze(-2) + out = torch.view_as_real(x_c * freqs_cis).flatten(-2) + return out.to(dtype) + + +def apply_partial_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + """Applies interleaved rotary embedding to trailing channels of x.""" + rope_dim = freqs_cis.shape[-1] * 2 + if x.shape[-1] > rope_dim: + nope, rope = x[..., :-rope_dim], x[..., -rope_dim:] + rotated = apply_rotary_emb(rope, freqs_cis) + return torch.cat([nope, rotated], dim=-1) + else: + return apply_rotary_emb(x, freqs_cis) + + +def sparse_attn_cpu( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attn_sink: torch.Tensor, + topk_idxs: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Pure CPU unoptimized implementation of TileLang sparse_attn_kernel.""" + b, m, h, d = q.shape + topk = topk_idxs.shape[-1] + valid_mask = topk_idxs >= 0 + safe_idxs = torch.where(valid_mask, topk_idxs, torch.zeros_like(topk_idxs)) + + # k, v: [b, n, h, d] -> gather -> [b, m, topk, h, d] + k_expanded = k.unsqueeze(1).expand(b, m, -1, h, d) + v_expanded = v.unsqueeze(1).expand(b, m, -1, h, d) + gather_index = safe_idxs.unsqueeze(-1).unsqueeze(-1).expand(b, m, topk, h, d) + gathered_k = torch.gather(k_expanded, dim=2, index=gather_index) + gathered_v = torch.gather(v_expanded, dim=2, index=gather_index) + + # scores: [b, m, h, topk] + scores = torch.einsum("bmhd,bmthd->bmht", q.float(), gathered_k.float()) * softmax_scale + scores = torch.where(valid_mask.unsqueeze(2), scores, float("-inf")) + + sink = attn_sink.view(1, 1, h, 1).expand(b, m, h, 1).float() + all_scores = torch.cat([scores, sink], dim=-1) + all_weights = torch.softmax(all_scores, dim=-1) + attn_weights = all_weights[..., :topk] + + out = torch.einsum("bmht,bmthd->bmhd", attn_weights, gathered_v.float()) + return out.to(q.dtype) + + +def hc_split_sinkhorn_cpu( + mixes: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + hc_mult: int = 4, + sinkhorn_iters: int = 20, + eps: float = 1e-6, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pure CPU implementation of TileLang hc_split_sinkhorn.""" + pre = torch.sigmoid(mixes[..., :hc_mult] * hc_scale[0] + hc_base[:hc_mult]) + eps + post = 2.0 * torch.sigmoid(mixes[..., hc_mult : 2 * hc_mult] * hc_scale[1] + hc_base[hc_mult : 2 * hc_mult]) + comb = ( + mixes[..., 2 * hc_mult :].view(*mixes.shape[:-1], hc_mult, hc_mult) * hc_scale[2] + + hc_base[2 * hc_mult :].view(hc_mult, hc_mult) + ) + comb = torch.softmax(comb.float(), dim=-1) + eps + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + for _ in range(sinkhorn_iters - 1): + comb = comb / (comb.sum(dim=-1, keepdim=True) + eps) + comb = comb / (comb.sum(dim=-2, keepdim=True) + eps) + return pre.to(mixes.dtype), post.to(mixes.dtype), comb.to(mixes.dtype) + + +# ============================================================================== +# 3. Official DeepSeek PyTorch Reference Architecture (inference/model.py CPU version) +# ============================================================================== + +class ParallelEmbedding_PT(nn_pt.Module): + def __init__(self, vocab_size: int, dim: int): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + self.weight = nn_pt.Parameter(torch.randn(vocab_size, dim, dtype=torch.float32) * 0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.embedding(x, self.weight) + + +class RMSNorm_PT(nn_pt.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn_pt.Parameter(torch.ones(dim, dtype=torch.float32)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps) * self.weight).to(x.dtype) + + +class Linear_PT(nn_pt.Module): + def __init__(self, in_features: int, out_features: int, bias: bool = False): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.weight = nn_pt.Parameter(torch.randn(out_features, in_features, dtype=torch.float32) * 0.02) + if bias: + self.bias = nn_pt.Parameter(torch.zeros(out_features, dtype=torch.float32)) + else: + self.register_parameter("bias", None) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return F.linear(x, self.weight, self.bias) + + +class ColumnParallelLinear_PT(Linear_PT): + pass + + +class RowParallelLinear_PT(Linear_PT): + pass + + +class Compressor_PT(nn_pt.Module): + def __init__(self, rate: int, head_dim: int, args: ModelArgs): + super().__init__() + self.rate = rate + self.head_dim = head_dim + self.args = args + if rate > 0: + proj_dim = (2 if rate == 4 else 1) * head_dim + self.wkv = Linear_PT(args.dim, proj_dim) + self.wgate = Linear_PT(args.dim, proj_dim) + self.norm = RMSNorm_PT(head_dim, args.norm_eps) + self.ape = nn_pt.Parameter(torch.randn(rate, proj_dim, dtype=torch.float32) * 0.02) + self.register_buffer( + "freqs_cis", + precompute_freqs_cis( + args.qk_rope_head_dim, + args.max_seq_len, + args.rope_theta * (16.0 if rate > 0 else 1.0), + ), + persistent=False, + ) + + def overlap_transform(self, tensor: torch.Tensor, value=0.0): + b, s, _, _ = tensor.size() + ratio, d = self.rate, self.head_dim + new_tensor = tensor.new_full((b, s, 2 * ratio, d), value) + new_tensor[:, :, ratio:] = tensor[:, :, :, d:] + if s > 1: + new_tensor[:, 1:, :ratio] = tensor[:, :-1, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor, start_pos: int = 0) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + if self.rate == 0: + return None, None + bsz, seqlen, _ = x.size() + ratio, d = self.rate, self.head_dim + kv = self.wkv(x.float()) + score = self.wgate(x.float()) + remainder = seqlen % ratio + cutoff = seqlen - remainder + if cutoff == 0: + return torch.zeros(bsz, 0, d, dtype=x.dtype, device=x.device), None + if remainder > 0: + kv = kv[:, :cutoff] + score = score[:, :cutoff] + kv = kv.unflatten(1, (-1, ratio)) + score = score.unflatten(1, (-1, ratio)) + self.ape + if ratio == 4: + kv = self.overlap_transform(kv, 0.0) + score = self.overlap_transform(score, float("-inf")) + kv = (kv * score.softmax(dim=2)).sum(dim=2) + kv = self.norm(kv.to(x.dtype)) + freqs = self.freqs_cis[:cutoff:ratio] + kv = apply_partial_rotary_emb(kv, freqs) + return kv, freqs + + +class Indexer_PT(nn_pt.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.wq_b = Linear_PT(args.q_lora_rank, args.index_n_heads * args.index_head_dim) + self.weights_proj = Linear_PT(args.dim, args.index_n_heads) + self.compressor = Compressor_PT(4, args.index_head_dim, args) + self.register_buffer( + "freqs_cis", + precompute_freqs_cis(args.qk_rope_head_dim, args.max_seq_len, args.rope_theta), + persistent=False, + ) + + def forward(self, x: torch.Tensor, q_latent: torch.Tensor, start_pos: int = 0) -> Tuple[torch.Tensor, torch.Tensor]: + b, m, _ = x.shape + weights = self.weights_proj(x).float() * (self.args.index_n_heads ** -0.5) + q = self.wq_b(q_latent).view(b, m, self.args.index_n_heads, self.args.index_head_dim) + qr = self.freqs_cis[start_pos : start_pos + m] + q = apply_partial_rotary_emb(q, qr) + k, _ = self.compressor(x, start_pos) + n_win = k.size(1) if k is not None else 0 + if n_win == 0: + return torch.zeros(b, m, 0, dtype=torch.long, device=x.device), torch.zeros(m, 0, dtype=torch.bool, device=x.device) + # k has shape [b, n_win, index_head_dim] + scores = torch.einsum("bmhd,btd->bmth", q.float(), k.float()) * (self.args.index_head_dim ** -0.5) + scores = (scores.relu() * weights.unsqueeze(2)).sum(dim=-1) + + causal_mask = torch.arange(n_win, device=x.device).unsqueeze(0) >= (torch.arange(1, m + 1, device=x.device).unsqueeze(1) // 4) + scores = scores + torch.where(causal_mask.unsqueeze(0), float("-inf"), 0.0) + + topk_block_idxs = scores.topk(min(self.args.index_topk, n_win), dim=-1).indices + invalid = causal_mask.unsqueeze(0).expand(b, -1, -1).gather(dim=-1, index=topk_block_idxs) + final_topk_idxs = torch.where(invalid, torch.full_like(topk_block_idxs, -1), topk_block_idxs) + return final_topk_idxs, causal_mask + + +class Attention_PT(nn_pt.Module): + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.args = args + self.rate = args.compress_rates[layer_id] + self.n_heads = args.n_heads + self.head_dim = args.head_dim + self.v_head_dim = args.v_head_dim + + self.wq_a = Linear_PT(args.dim, args.q_lora_rank) + self.q_norm = RMSNorm_PT(args.q_lora_rank, args.norm_eps) + self.wq_b = Linear_PT(args.q_lora_rank, args.n_heads * args.head_dim) + self.wkv = Linear_PT(args.dim, args.head_dim) + self.kv_norm = RMSNorm_PT(args.qk_nope_head_dim, args.norm_eps) + + self.wo_a = ColumnParallelLinear_PT(args.n_heads * args.v_head_dim // args.o_groups, args.o_groups * args.o_lora_rank) + self.wo_b = RowParallelLinear_PT(args.o_groups * args.o_lora_rank, args.dim) + self.attn_sink = nn_pt.Parameter(torch.randn(args.n_heads, dtype=torch.float32) * 0.02) + + self.compressor = Compressor_PT(self.rate, args.v_head_dim, args) if self.rate > 0 else None + self.indexer = Indexer_PT(args) if self.rate == 4 else None + + self.register_buffer( + "freqs_cis", + precompute_freqs_cis(args.qk_rope_head_dim, args.max_seq_len, args.rope_theta), + persistent=False, + ) + + def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor: + b, m, _ = x.shape + q_latent = self.q_norm(self.wq_a(x)) + q = self.wq_b(q_latent).view(b, m, self.n_heads, self.head_dim) + qr = self.freqs_cis[start_pos : start_pos + m] + q = apply_partial_rotary_emb(q, qr) + kv = self.wkv(x) + if self.args.qk_nope_head_dim > 0: + kv = self.kv_norm(kv) + k_roped = apply_partial_rotary_emb(kv.unsqueeze(2), qr).squeeze(2) + v_unroped = kv + + # For standard/sliding attention, k has head dimension (v_head_dim) + k_proj = k_roped.unsqueeze(2).expand(-1, -1, self.n_heads, -1) + v_proj = v_unroped.unsqueeze(2).expand(-1, -1, self.n_heads, -1) + + scale = self.head_dim ** -0.5 + if self.rate == 0: + # Sliding window causal attention + scores = torch.einsum("bmhd,bnhd->bmhn", q.float(), k_proj.float()) * scale + pos_q = torch.arange(start_pos, start_pos + m, device=x.device).unsqueeze(1) + pos_k = torch.arange(0, m, device=x.device).unsqueeze(0) + causal_mask = pos_q < pos_k + sliding_mask = (pos_q - pos_k) >= self.args.sliding_window + mask = causal_mask | sliding_mask + scores = torch.where(mask.unsqueeze(1).unsqueeze(0), float("-inf"), scores) + sink = self.attn_sink.view(1, 1, self.n_heads, 1).expand(b, m, self.n_heads, 1).float() + all_scores = torch.cat([scores, sink], dim=-1) + weights = torch.softmax(all_scores, dim=-1)[..., :m] + out = torch.einsum("bmhn,bnhd->bmhd", weights, v_proj.float()).to(x.dtype) + elif self.rate == 4: + # CSA: compound attention with local sliding window + topk compressed blocks + comp_kv, _ = self.compressor(x, start_pos) + n_win = comp_kv.size(1) if comp_kv is not None else 0 + if n_win > 0: + topk_block_idxs, causal_block_mask = self.indexer(x, q_latent, start_pos) + + valid = topk_block_idxs >= 0 + entry_indices = torch.arange(n_win, device=x.device).view(1, 1, 1, n_win) + is_in_topk = topk_block_idxs.unsqueeze(-1) == entry_indices + is_valid_and_in_topk = is_in_topk & valid.unsqueeze(-1) + block_selected = is_valid_and_in_topk.any(dim=2) # [b, m, n_win] + + comp_kv_expanded = comp_kv.unsqueeze(2).expand(-1, -1, self.n_heads, -1) + all_k = torch.cat([k_proj, comp_kv_expanded], dim=1) + all_v = torch.cat([v_proj, comp_kv_expanded], dim=1) + + scores = torch.einsum("bmhd,bnhd->bmhn", q.float(), all_k.float()) * scale + + pos_q = torch.arange(start_pos, start_pos + m, device=x.device).unsqueeze(1) + pos_k = torch.arange(0, m, device=x.device).unsqueeze(0) + causal_mask = pos_q < pos_k + sliding_mask = (pos_q - pos_k) >= self.args.sliding_window + uncompressed_mask = causal_mask | sliding_mask + + scores[:, :, :, :m] = torch.where(uncompressed_mask.unsqueeze(0).unsqueeze(2), float("-inf"), scores[:, :, :, :m]) + scores[:, :, :, m:] = torch.where(block_selected.unsqueeze(2), scores[:, :, :, m:], float("-inf")) + + sink = self.attn_sink.view(1, 1, self.n_heads, 1).expand(b, m, self.n_heads, 1).float() + all_scores = torch.cat([scores, sink], dim=-1) + weights = torch.softmax(all_scores, dim=-1)[..., : m + n_win] + out = torch.einsum("bmhn,bnhd->bmhd", weights, all_v.float()).to(x.dtype) + else: + scores = torch.einsum("bmhd,bnhd->bmhn", q.float(), k_proj.float()) * scale + pos_q = torch.arange(start_pos, start_pos + m, device=x.device).unsqueeze(1) + pos_k = torch.arange(0, m, device=x.device).unsqueeze(0) + causal_mask = pos_q < pos_k + sliding_mask = (pos_q - pos_k) >= self.args.sliding_window + mask = causal_mask | sliding_mask + scores = torch.where(mask.unsqueeze(1).unsqueeze(0), float("-inf"), scores) + sink = self.attn_sink.view(1, 1, self.n_heads, 1).expand(b, m, self.n_heads, 1).float() + all_scores = torch.cat([scores, sink], dim=-1) + weights = torch.softmax(all_scores, dim=-1)[..., :m] + out = torch.einsum("bmhn,bnhd->bmhd", weights, v_proj.float()).to(x.dtype) + else: + # HCA: compressed attention (rate > 4) + comp_kv, _ = self.compressor(x, start_pos) + n_win = comp_kv.size(1) if comp_kv is not None else 0 + if n_win > 0: + comp_kv_expanded = comp_kv.unsqueeze(2).expand(-1, -1, self.n_heads, -1) + all_k = torch.cat([k_proj, comp_kv_expanded], dim=1) + all_v = torch.cat([v_proj, comp_kv_expanded], dim=1) + + scores = torch.einsum("bmhd,bnhd->bmhn", q.float(), all_k.float()) * scale + + pos_q = torch.arange(start_pos, start_pos + m, device=x.device).unsqueeze(1) + pos_k = torch.arange(0, m, device=x.device).unsqueeze(0) + causal_mask = pos_q < pos_k + sliding_mask = (pos_q - pos_k) >= self.args.sliding_window + uncompressed_mask = causal_mask | sliding_mask + + scores[:, :, :, :m] = torch.where(uncompressed_mask.unsqueeze(0).unsqueeze(2), float("-inf"), scores[:, :, :, :m]) + + pos_comp_k = torch.arange(0, n_win, device=x.device).unsqueeze(0) + causal_comp_mask = pos_comp_k >= ((pos_q + 1) // self.rate) + scores[:, :, :, m:] = torch.where((~causal_comp_mask).unsqueeze(0).unsqueeze(2), scores[:, :, :, m:], float("-inf")) + + sink = self.attn_sink.view(1, 1, self.n_heads, 1).expand(b, m, self.n_heads, 1).float() + all_scores = torch.cat([scores, sink], dim=-1) + weights = torch.softmax(all_scores, dim=-1)[..., : m + n_win] + out = torch.einsum("bmhn,bnhd->bmhd", weights, all_v.float()).to(x.dtype) + else: + scores = torch.einsum("bmhd,bnhd->bmhn", q.float(), k_proj.float()) * scale + pos_q = torch.arange(start_pos, start_pos + m, device=x.device).unsqueeze(1) + pos_k = torch.arange(0, m, device=x.device).unsqueeze(0) + causal_mask = pos_q < pos_k + sliding_mask = (pos_q - pos_k) >= self.args.sliding_window + mask = causal_mask | sliding_mask + scores = torch.where(mask.unsqueeze(1).unsqueeze(0), float("-inf"), scores) + sink = self.attn_sink.view(1, 1, self.n_heads, 1).expand(b, m, self.n_heads, 1).float() + all_scores = torch.cat([scores, sink], dim=-1) + weights = torch.softmax(all_scores, dim=-1)[..., :m] + out = torch.einsum("bmhn,bnhd->bmhd", weights, v_proj.float()).to(x.dtype) + + # Output projection: Grouped wo_a + wo_b + out = out.reshape(b, m, self.args.o_groups, -1) + # wo_a is applied per group + wo_a_w = self.wo_a.weight.reshape(self.args.o_groups, self.args.o_lora_rank, -1) + out = torch.einsum("bmgi,goi->bmgo", out.float(), wo_a_w.float()).reshape(b, m, -1) + out = self.wo_b(out.to(x.dtype)) + return out + + +class Gate_PT(nn_pt.Module): + def __init__(self, is_hash_layer: bool, args: ModelArgs): + super().__init__() + self.is_hash_layer = is_hash_layer + self.args = args + self.weight = nn_pt.Parameter(torch.randn(args.n_routed_experts, args.dim, dtype=torch.float32) * 0.02) + if is_hash_layer: + tids = torch.stack([ + torch.randperm(args.n_routed_experts)[: args.n_activated_experts] for _ in range(args.vocab_size) + ]).to(torch.int32) + self.tid2eid = nn_pt.Parameter(tids, requires_grad=False) + self.bias = None + else: + self.tid2eid = None + self.bias = nn_pt.Parameter(torch.randn(args.n_routed_experts, dtype=torch.float32) * 0.02) + + def forward(self, x: torch.Tensor, input_ids: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: + scores = F.linear(x.float(), self.weight.float()) + if self.args.score_func == "softmax": + scores = scores.softmax(dim=-1) + elif self.args.score_func == "sigmoid": + scores = scores.sigmoid() + elif self.args.score_func == "sqrtsoftplus": + scores = F.softplus(scores).sqrt() + original_scores = scores + if getattr(self, "bias", None) is not None: + scores = scores + self.bias + if self.is_hash_layer: + assert input_ids is not None + indices = self.tid2eid[input_ids].long() + else: + indices = scores.topk(self.args.n_activated_experts, dim=-1).indices + weights = torch.gather(original_scores, -1, indices) + if self.args.score_func in ("sigmoid", "sqrtsoftplus"): + weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) + weights = weights * self.args.route_scale + return weights, indices + + +class Expert_PT(nn_pt.Module): + def __init__(self, in_features: int, hidden_features: int, out_features: int): + super().__init__() + self.gate_proj = Linear_PT(in_features, hidden_features) + self.down_proj = Linear_PT(hidden_features, out_features) + self.up_proj = Linear_PT(in_features, hidden_features) + self.w1 = self.gate_proj + self.w2 = self.down_proj + self.w3 = self.up_proj + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = torch.clamp(self.gate_proj(x), max=10.0) + up = torch.clamp(self.up_proj(x), min=-10.0, max=10.0) + return self.down_proj(F.silu(gate) * up) + + +class MoE_PT(nn_pt.Module): + def __init__(self, is_hash_layer: bool, args: ModelArgs): + super().__init__() + self.args = args + self.gate = Gate_PT(is_hash_layer, args) + self.shared_experts = Expert_PT(args.dim, args.n_shared_experts * args.moe_inter_dim, args.dim) + self.experts = nn_pt.ModuleList([ + Expert_PT(args.dim, args.moe_inter_dim, args.dim) for _ in range(args.n_routed_experts) + ]) + + def forward(self, x: torch.Tensor, input_ids: Optional[torch.Tensor] = None) -> torch.Tensor: + weights, indices = self.gate(x, input_ids) + b, m, d = x.shape + out = torch.zeros_like(x) + for k in range(indices.shape[-1]): + idx_k = indices[..., k] + w_k = weights[..., k].unsqueeze(-1) + for e in range(self.args.n_routed_experts): + mask = (idx_k == e) + if mask.any(): + x_e = x[mask] + out_e = self.experts[e](x_e) + out[mask] += out_e * w_k[mask] + out += self.shared_experts(x) + return out + + +class Block_PT(nn_pt.Module): + def __init__(self, layer_id: int, args: ModelArgs): + super().__init__() + self.layer_id = layer_id + self.args = args + self.attn_norm = RMSNorm_PT(args.dim, args.norm_eps) + self.ffn_norm = RMSNorm_PT(args.dim, args.norm_eps) + self.attn = Attention_PT(layer_id, args) + self.ffn = MoE_PT(layer_id < args.n_hash_layers, args) + + self.hc_mult = args.hc_mult + hc_dim = self.hc_mult * args.dim + mix_hc = (2 + self.hc_mult) * self.hc_mult + self.hc_attn_fn = nn_pt.Parameter(torch.randn(mix_hc, hc_dim, dtype=torch.float32) * 0.02) + self.hc_ffn_fn = nn_pt.Parameter(torch.randn(mix_hc, hc_dim, dtype=torch.float32) * 0.02) + self.hc_attn_base = nn_pt.Parameter(torch.zeros(mix_hc, dtype=torch.float32)) + self.hc_ffn_base = nn_pt.Parameter(torch.zeros(mix_hc, dtype=torch.float32)) + self.hc_attn_scale = nn_pt.Parameter(torch.ones(3, dtype=torch.float32)) + self.hc_ffn_scale = nn_pt.Parameter(torch.ones(3, dtype=torch.float32)) + + def hc_pre(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor): + shape, dtype = x.size(), x.dtype + x_flat = x.flatten(2).float() + rsqrt = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + self.args.norm_eps) + mixes = F.linear(x_flat, hc_fn) * rsqrt + pre, post, comb = hc_split_sinkhorn_cpu( + mixes, hc_scale, hc_base, self.hc_mult, self.args.hc_sinkhorn_iters, self.args.hc_eps + ) + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2) + return y.to(dtype), post, comb + + def hc_post(self, x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb: torch.Tensor): + y = post.unsqueeze(-1) * x.unsqueeze(-2) + torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim=2) + return y.to(x.dtype) + + def forward(self, x: torch.Tensor, start_pos: int = 0, input_ids: Optional[torch.Tensor] = None) -> torch.Tensor: + x_attn, post, comb = self.hc_pre(x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base) + attn_out = self.attn(self.attn_norm(x_attn), start_pos) + x = self.hc_post(attn_out, x, post, comb) + + x_ffn, post, comb = self.hc_pre(x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base) + ffn_out = self.ffn(self.ffn_norm(x_ffn), input_ids) + x = self.hc_post(ffn_out, x, post, comb) + return x + + +class ParallelHead_PT(nn_pt.Module): + def __init__(self, vocab_size: int, dim: int, args: ModelArgs): + super().__init__() + self.vocab_size = vocab_size + self.dim = dim + self.args = args + self.weight = nn_pt.Parameter(torch.randn(vocab_size, dim, dtype=torch.float32) * 0.02) + + def hc_head(self, x: torch.Tensor, hc_fn: torch.Tensor, hc_scale: torch.Tensor, hc_base: torch.Tensor) -> torch.Tensor: + shape, dtype = x.size(), x.dtype + x_flat = x.flatten(2).float() + rsqrt = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + self.args.norm_eps) + mixes = F.linear(x_flat, hc_fn) * rsqrt + pre = torch.sigmoid(mixes * hc_scale + hc_base) + self.args.hc_eps + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=2) + return y.to(dtype) + + def forward( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + norm: RMSNorm_PT, + ) -> torch.Tensor: + x = self.hc_head(x, hc_fn, hc_scale, hc_base) + x = norm(x) + return F.linear(x, self.weight) + + +# ============================================================================== +# 4. Parameter Mapping Transfer Helpers +# ============================================================================== + +def _get_nested_pt_attr(obj, path: str): + """Fetches a nested attribute from a PyTorch module.""" + if path is None: + return None + parts = path.split(".") + curr = obj + for part in parts: + if part.isdigit(): + curr = curr[int(part)] + elif hasattr(curr, part): + curr = getattr(curr, part) + elif isinstance(curr, dict) and part in curr: + curr = curr[part] + else: + return None + return curr + + +def _apply_global_param_mapping(mt_model, pt_model, pt_config_dict: dict, mx_config: Config): + mapping = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING(pt_config_dict, mx_config, scan_layers=False) + hooks = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN(pt_config_dict, mx_config, scan_layers=False, saving_to_hf=False) + + for mt_key, hf_key in mapping.items(): + if "layers" in mt_key or (hf_key is not None and "layers" in str(hf_key)): + continue # handled by layer mapper + + if hf_key is None: + continue + + # Map directly + def set_value(obj, val): + import flax.nnx as nnx + if isinstance(obj, nnx.Variable): + obj.value = val + else: + obj.value = val + + if mt_key == "params-token_embedder-embedding": + val = pt_model.embed.weight.detach().numpy() + target_shape = mt_model.token_embedder.embedding.value.shape + if mt_key in hooks: + val = hooks[mt_key](val, target_shape=target_shape) + set_value(mt_model.token_embedder.embedding, val) + elif mt_key == "params-decoder-decoder_norm-scale": + val = pt_model.norm.weight.detach().numpy() + target_shape = mt_model.decoder.decoder_norm.scale.value.shape + if mt_key in hooks: + val = hooks[mt_key](val, target_shape=target_shape) + set_value(mt_model.decoder.decoder_norm.scale, val) + elif mt_key == "params-decoder-logits_dense-kernel": + val = pt_model.head.weight.detach().numpy() + target_shape = mt_model.decoder.logits_dense.kernel.value.shape + if mt_key in hooks: + val = hooks[mt_key](val, target_shape=target_shape) + set_value(mt_model.decoder.logits_dense.kernel, val) + elif mt_key == "params-decoder-hc_head-hc_fn": + if hasattr(pt_model, "hc_head_fn"): + val = pt_model.hc_head_fn.detach().numpy() + target_shape = mt_model.decoder.hc_head.hc_fn.value.shape + if mt_key in hooks: + val = hooks[mt_key](val, target_shape=target_shape) + set_value(mt_model.decoder.hc_head.hc_fn, val) + elif mt_key == "params-decoder-hc_head-hc_base": + if hasattr(pt_model, "hc_head_base"): + val = pt_model.hc_head_base.detach().numpy() + set_value(mt_model.decoder.hc_head.hc_base, val) + elif mt_key == "params-decoder-hc_head-hc_scale": + if hasattr(pt_model, "hc_head_scale"): + val = pt_model.hc_head_scale.detach().numpy() + set_value(mt_model.decoder.hc_head.hc_scale, val) + + +def _build_scanned_maxtext_params(pt_model, pt_config_dict: dict, mx_config: Config): + """Builds a complete Linen params dict for scanned DeepSeek-V4 using checkpoint conversion utilities.""" + from maxtext.checkpoint_conversion.to_maxtext import _get_hf_loading_function + from maxtext.checkpoint_conversion.utils.utils import param_key_parts_from_path + from maxtext.models import models + from maxtext.utils import maxtext_utils + + hf_tensors = {} + for name, param in pt_model.named_parameters(): + hf_tensors[name] = param.detach().cpu().numpy() + for name, buf in pt_model.named_buffers(): + hf_tensors[name] = buf.detach().cpu().numpy() + + def tensor_getter(key): + if key in hf_tensors: + return hf_tensors[key] + alt_key = ( + key.replace(".w1.weight", ".gate_proj.weight") + .replace(".w2.weight", ".down_proj.weight") + .replace(".w3.weight", ".up_proj.weight") + .replace(".w1.bias", ".gate_proj.bias") + .replace(".w2.bias", ".down_proj.bias") + .replace(".w3.bias", ".up_proj.bias") + ) + if alt_key in hf_tensors: + return hf_tensors[alt_key] + raise KeyError(f"Key {key} (and alt {alt_key}) not found in hf_tensors.") + + mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + maxtext_model = models.transformer_as_linen(mx_config, mesh, quant=None, model_mode=MODEL_MODE_TRAIN) + + abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model, mx_config) + abstract_params_flat, abstract_params_treedef = jax.tree_util.tree_flatten_with_path( + abstract_params_tree, + is_leaf=lambda x: isinstance(x, nn.LogicallyPartitioned), + ) + + maxtext_abstract_dict = {} + for mt_target_idx, (path_tuple, abstract_leaf_value) in enumerate(abstract_params_flat): + mt_param_key = "-".join(param_key_parts_from_path(path_tuple)) + if isinstance(abstract_leaf_value, nn.LogicallyPartitioned): + mt_target_shape = abstract_leaf_value.value.shape + else: + mt_target_shape = abstract_leaf_value.shape + maxtext_abstract_dict[mt_param_key] = (mt_target_idx, mt_target_shape) + + mapping = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING(pt_config_dict, mx_config, scan_layers=True) + hooks = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN(pt_config_dict, mx_config, scan_layers=True, saving_to_hf=False) + + final_mt_weights = [None] * len(abstract_params_flat) + + for mt_key, hf_key in mapping.items(): + if mt_key not in maxtext_abstract_dict: + continue + target_idx, target_shape = maxtext_abstract_dict[mt_key] + hook_fn = hooks.get(mt_key, None) + load_fn = _get_hf_loading_function(hf_key, tensor_getter, hook_fn, target_shape, mx_config, mt_key=mt_key) + val = load_fn() + final_mt_weights[target_idx] = val + + params = jax.tree_util.tree_unflatten(abstract_params_treedef, final_mt_weights) + return maxtext_model, params + + +def _apply_layer_param_mapping(mt_layer, pt_model, layer_idx: int, pt_config_dict: dict, mx_config: Config): + """Applies parameter mapping from PyTorch Block_PT to MaxText DeepSeek4DecoderLayer.""" + mapping = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_MAPPING(pt_config_dict, mx_config, scan_layers=False) + hooks = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN(pt_config_dict, mx_config, scan_layers=False, saving_to_hf=False) + + pt_prefix = f"layers.{layer_idx}." + for mt_key, hf_key in mapping.items(): + if f"layers_{layer_idx}" not in mt_key: + continue + + # Extract target NNX path + if "Tid2EidVar" in mt_key: + prefix = f"Tid2EidVar-decoder-layers_{layer_idx}-" + elif "MoEBiasVar" in mt_key: + prefix = f"MoEBiasVar-decoder-layers_{layer_idx}-" + else: + prefix = f"params-decoder-layers_{layer_idx}-" + + nnx_subpath = mt_key.replace(prefix, "").replace("-", ".") + parts = nnx_subpath.split(".") + obj = mt_layer + valid = True + for part in parts: + if hasattr(obj, part): + obj = getattr(obj, part) + else: + valid = False + break + if not valid or obj is None or not hasattr(obj, "value"): + continue + + target_shape = obj.value.shape + hook_fn = hooks.get(mt_key, lambda x, target_shape=None: x) + + if hf_key is None: + val = hook_fn(None, target_shape=target_shape) + elif isinstance(hf_key, list): + pt_vals = [_get_nested_pt_attr(pt_model, k.replace(pt_prefix, "")) for k in hf_key] + if any(v is None for v in pt_vals): + print(f"FAILED LIST LOOKUP for {mt_key}: {hf_key}", flush=True) + continue + pt_vals = [v.detach().numpy() for v in pt_vals] + slice_shape = target_shape[1:] + processed_vals = [hook_fn(v, target_shape=slice_shape) for v in pt_vals] + val = np.stack(processed_vals, axis=0) + print(f"HOOK CALL (LIST): mt_key={mt_key}, val_shape={val.shape}, target_shape={target_shape}", flush=True) + elif isinstance(hf_key, tuple): + pt_vals = [_get_nested_pt_attr(pt_model, k.replace(pt_prefix, "")) for k in hf_key] + if any(v is None for v in pt_vals): + print(f"FAILED TUPLE LOOKUP for {mt_key}: {hf_key}", flush=True) + continue + pt_vals = tuple(v.detach().numpy() for v in pt_vals) + val = hook_fn(pt_vals, target_shape=target_shape) + print(f"HOOK CALL (TUPLE): mt_key={mt_key}, val_shape={val.shape}, target_shape={target_shape}", flush=True) + else: + pt_attr = _get_nested_pt_attr(pt_model, hf_key.replace(pt_prefix, "")) + if pt_attr is None: + print(f"FAILED SINGLE LOOKUP for {mt_key}: {hf_key.replace(pt_prefix, '')}", flush=True) + continue + pt_val = pt_attr.detach().numpy() + # print(f"DEBUG: mt_key={mt_key}, hf_key={hf_key}, pt_val={pt_val.shape}, target_shape={target_shape}", flush=True) + print(f"HOOK CALL: mt_key={mt_key}, hf_key={hf_key}, pt_shape={pt_val.shape}, target_shape={target_shape}", flush=True) + val = hook_fn(pt_val, target_shape=target_shape) + + if val is not None: + setattr(obj, "value", jnp.array(val)) + + +# ============================================================================== +# 5. Unit Tests +# ============================================================================== + +class DeepSeekV4FlashEmbeddingTest(unittest.TestCase): + """Validates ParallelEmbedding_PT against MaxText Embed via parameter mapping.""" + + def setUp(self): + self.args = ModelArgs() + self.vocab_size = self.args.vocab_size + self.dim = self.args.dim + self.mx_config = get_maxtext_config() + self.mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + self.rngs = nnx.Rngs(0) + + def test_embedding_parity(self): + pt_embed = ParallelEmbedding_PT(self.vocab_size, self.dim) + mt_embed = Embed( + config=self.mx_config, + num_embeddings=self.vocab_size, + num_features=self.dim, + dtype=jnp.float32, + mesh=self.mesh, + rngs=self.rngs, + ) + + # Use mapping hook + hooks = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN({"num_hidden_layers": 5}, self.mx_config) + unpad_fn = hooks["params-token_embedder-embedding"] + mt_embed.embedding.value = jnp.array(unpad_fn(pt_embed.weight.detach().numpy(), mt_embed.embedding.value.shape)) + + input_ids_np = np.random.randint(0, self.vocab_size, size=(2, 16)) + pt_out = pt_embed(torch.tensor(input_ids_np, dtype=torch.long)).detach().numpy() + mt_out = np.array(mt_embed(jnp.array(input_ids_np))) + + np.testing.assert_allclose(mt_out, pt_out, rtol=1e-5, atol=1e-5) + + +class DeepSeekV4FlashRMSNormTest(unittest.TestCase): + """Validates RMSNorm_PT against MaxText RMSNorm.""" + + def setUp(self): + self.args = ModelArgs() + self.dim = self.args.dim + self.eps = self.args.norm_eps + self.rngs = nnx.Rngs(0) + + def test_rmsnorm_parity(self): + pt_norm = RMSNorm_PT(self.dim, self.eps) + mt_norm = RMSNorm(num_features=self.dim, epsilon=self.eps, dtype=jnp.float32, weight_dtype=jnp.float32, rngs=self.rngs) + + mt_norm.scale.value = jnp.array(pt_norm.weight.detach().numpy()) + + x_np = np.random.normal(size=(2, 16, self.dim)).astype(np.float32) + pt_out = pt_norm(torch.tensor(x_np)).detach().numpy() + mt_out = np.array(mt_norm(jnp.array(x_np))) + + np.testing.assert_allclose(mt_out, pt_out, rtol=1e-5, atol=1e-5) + + +class DeepSeekV4FlashRotaryEmbeddingTest(unittest.TestCase): + """Validates apply_rotary_emb / precompute_freqs_cis against DeepSeekV4RotaryEmbedding.""" + + def setUp(self): + self.args = ModelArgs() + self.head_dim = self.args.head_dim + self.qk_rope_head_dim = self.args.qk_rope_head_dim + self.seq_len = 32 + self.batch_size = 2 + self.num_heads = self.args.n_heads + + def test_main_rope(self): + self._run_rope_test(theta=10000.0) + + def test_compressed_rope(self): + self._run_rope_test(theta=160000.0) + + def _run_rope_test(self, theta: float): + freqs_cis = precompute_freqs_cis(self.qk_rope_head_dim, self.seq_len, theta) + mt_rope = DeepSeekV4RotaryEmbedding( + head_dim=self.head_dim, + partial_rotary_factor=self.qk_rope_head_dim / self.head_dim, + rope_theta=theta, + ) + + x_np = np.random.normal(size=(self.batch_size, self.seq_len, self.num_heads, self.head_dim)).astype(np.float32) + pos_np = np.arange(self.seq_len)[None, :].repeat(self.batch_size, axis=0) + + # PyTorch reference + x_pt = torch.tensor(x_np) + qr = freqs_cis[pos_np[0]].unsqueeze(1) # [seq_len, 1, qk_rope_head_dim//2] + pt_out = apply_partial_rotary_emb(x_pt, qr).detach().numpy() + + # MaxText + mt_out = np.array(mt_rope(jnp.array(x_np), jnp.array(pos_np), unsqueeze_dim=2)) + + np.testing.assert_allclose(mt_out, pt_out, rtol=1e-5, atol=1e-5) + + +class DeepSeekV4FlashGroupedLinearTest(unittest.TestCase): + """Validates ColumnParallelLinear_PT (wo_a) against MaxText DeepSeekV4GroupedLinear.""" + + def setUp(self): + self.args = ModelArgs() + self.in_features_per_group = (self.args.n_heads * self.args.v_head_dim) // self.args.o_groups + self.out_features_per_group = self.args.o_lora_rank + self.n_groups = self.args.o_groups + self.total_out_features = self.n_groups * self.out_features_per_group + self.rngs = nnx.Rngs(0) + self.mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + + def test_grouped_linear_parity(self): + pt_linear = ColumnParallelLinear_PT(self.in_features_per_group, self.total_out_features) + mt_linear = DeepSeekV4GroupedLinear( + in_features_per_group=self.in_features_per_group, + out_features=self.total_out_features, + n_groups=self.n_groups, + rngs=self.rngs, + ) + + # Use hook reshape_o_a_proj + hooks = DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN({"num_hidden_layers": 5}, None) + reshape_fn = hooks["params-decoder-layers_0-self_attention-o_a_proj-kernel"] + mt_linear.kernel.value = jnp.array(reshape_fn(pt_linear.weight.detach().numpy(), mt_linear.kernel.value.shape)) + + x_np = np.random.normal(size=(2, 16, self.n_groups, self.in_features_per_group)).astype(np.float32) + # PyTorch evaluation per group + w_groups = pt_linear.weight.view(self.n_groups, self.out_features_per_group, self.in_features_per_group) + pt_out = torch.einsum("bmgi,goi->bmgo", torch.tensor(x_np), w_groups).detach().numpy() + + mt_out = np.array(mt_linear(jnp.array(x_np))) + + np.testing.assert_allclose(mt_out, pt_out, rtol=1e-5, atol=1e-5) + + +class DeepSeekV4FlashMoEGateTest(unittest.TestCase): + """Validates Gate_PT (Hash and TopK) against MaxText RoutedMoE gate.""" + + def setUp(self): + self.args = ModelArgs() + self.dim = self.args.dim + self.vocab_size = self.args.vocab_size + self.num_experts = self.args.n_routed_experts + self.num_activated = self.args.n_activated_experts + self.mx_config = get_maxtext_config() + self.args.score_func = self.mx_config.routed_score_func + self.args.route_scale = self.mx_config.routed_scaling_factor + self.mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + self.rngs = nnx.Rngs(0) + + def test_hash_routing_gate_parity(self): + pt_gate = Gate_PT(is_hash_layer=True, args=self.args) + mt_moe = RoutedMoE( + config=self.mx_config, + num_experts=self.num_experts, + num_experts_per_tok=self.num_activated, + mesh=self.mesh, + kernel_init=initializers.nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed_moe", None), + is_hash_routing=True, + rngs=self.rngs, + ) + + # Copy weights & tid2eid + mt_moe.gate.kernel.value = jnp.array(pt_gate.weight.detach().numpy().T) + mt_moe.tid2eid.value = jnp.array(pt_gate.tid2eid.detach().numpy()) + + input_ids_np = np.random.randint(0, self.vocab_size, size=(2, 16)) + x_np = np.random.normal(size=(2, 16, self.dim)).astype(np.float32) + + pt_weights, pt_indices = pt_gate(torch.tensor(x_np), torch.tensor(input_ids_np, dtype=torch.long)) + + gate_logits, pre_bias_logits = mt_moe.gate(jnp.array(x_np)) + mt_weights, mt_indices = mt_moe.get_topk( + gate_logits, pre_bias_logits, rngs=self.rngs, input_ids=jnp.array(input_ids_np) + ) + + np.testing.assert_array_equal(np.array(mt_indices), pt_indices.detach().numpy()) + np.testing.assert_allclose(np.array(mt_weights), pt_weights.detach().numpy(), rtol=1e-5, atol=1e-5) + + +class DeepSeekV4FlashMoEBlockTest(unittest.TestCase): + """Validates full MoE_PT (Hash and TopK) against MaxText RoutedAndSharedMoE.""" + + def setUp(self): + self.args = ModelArgs() + self.dim = self.args.dim + self.inter_dim = self.args.inter_dim + self.vocab_size = self.args.vocab_size + self.num_experts = self.args.n_routed_experts + self.num_activated = self.args.n_activated_experts + self.mx_config = get_maxtext_config() + self.args.score_func = self.mx_config.routed_score_func + self.args.route_scale = self.mx_config.routed_scaling_factor + self.mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + self.rngs = nnx.Rngs(0) + + def test_hash_routing_moe_parity(self): + pt_moe = MoE_PT(is_hash_layer=True, args=self.args) + mt_moe = RoutedAndSharedMoE( + config=self.mx_config, + mesh=self.mesh, + kernel_init=initializers.nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("embed_moe", None), + rngs=self.rngs, + is_hash_routing=True, + ) + + class DummyLayer: + def __init__(self, moe): + self.mlp = moe + + class DummyPTLayer: + def __init__(self, moe): + self.ffn = moe + + pt_config_dict = { + "num_hidden_layers": 1, + "first_k_dense_replace": 0, + "n_routed_experts": self.args.n_routed_experts, + "num_experts_per_tok": self.args.n_activated_experts, + "first_num_hash_layers": 1, + } + _apply_layer_param_mapping(DummyLayer(mt_moe), DummyPTLayer(pt_moe), layer_idx=0, pt_config_dict=pt_config_dict, mx_config=self.mx_config) + + input_ids_np = np.random.randint(0, self.vocab_size, size=(2, 16)) + x_np = np.random.normal(size=(2, 16, self.dim)).astype(np.float32) + + pt_out = pt_moe(torch.tensor(x_np), torch.tensor(input_ids_np, dtype=torch.long)).detach().numpy() + mt_out, _, _ = mt_moe(jnp.array(x_np), input_ids=jnp.array(input_ids_np)) + + np.testing.assert_allclose(np.array(mt_out), pt_out, rtol=5e-2, atol=5e-2) + + +class DeepSeekV4FlashHyperHeadTest(unittest.TestCase): + """Validates ParallelHead_PT.hc_head against MaxText DeepSeek4HyperHead.""" + + def setUp(self): + self.args = ModelArgs() + self.dim = self.args.dim + self.hc_mult = self.args.hc_mult + self.vocab_size = self.args.vocab_size + self.mx_config = get_maxtext_config() + self.mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + self.rngs = nnx.Rngs(0) + + def test_hyperhead_parity(self): + pt_head = ParallelHead_PT(self.vocab_size, self.dim, self.args) + hc_fn = torch.randn(self.hc_mult, self.hc_mult * self.dim, dtype=torch.float32) * 0.02 + hc_base = torch.zeros(self.hc_mult, dtype=torch.float32) + hc_scale = torch.ones(1, dtype=torch.float32) + + mt_head = DeepSeek4HyperHead( + config=self.mx_config, + mesh=self.mesh, + rngs=self.rngs, + ) + + mt_head.hc_fn.value = jnp.array(hc_fn.detach().numpy().T) + mt_head.hc_base.value = jnp.array(hc_base.detach().numpy()) + mt_head.hc_scale.value = jnp.array(hc_scale.detach().numpy()) + + x_np = np.random.normal(size=(2, 16, self.hc_mult, self.dim)).astype(np.float32) + pt_out = pt_head.hc_head(torch.tensor(x_np), hc_fn, hc_scale, hc_base).detach().numpy() + mt_out = np.array(mt_head(jnp.array(x_np))) + + np.testing.assert_allclose(mt_out, pt_out, rtol=2e-5, atol=2e-5) + + +class DeepSeekV4FlashDecoderLayerTest(unittest.TestCase): + """Validates Block_PT vs MaxText DeepSeek4DecoderLayer across all layer types using param_mapping.py.""" + + def setUp(self): + self.batch_size = 2 + self.seq_len = 16 + self.args = ModelArgs() + self.dim = self.args.dim + self.vocab_size = self.args.vocab_size + self.mx_config = get_maxtext_config() + self.args.score_func = self.mx_config.routed_score_func + self.args.route_scale = self.mx_config.routed_scaling_factor + self.mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + self.rngs = nnx.Rngs(0) + + def _run_layer_test(self, layer_idx: int): + pt_block = Block_PT(layer_id=layer_idx, args=self.args) + mt_layer = DeepSeek4DecoderLayer( + config=self.mx_config, + model_mode="train", + mesh=self.mesh, + rngs=self.rngs, + layer_idx=layer_idx, + compress_ratio=self.args.compress_rates[layer_idx], + is_hash_routing=(layer_idx < self.args.n_hash_layers), + ) + + pt_config_dict = { + "num_hidden_layers": self.args.n_layers, + "first_k_dense_replace": 0, + "n_routed_experts": self.args.n_routed_experts, + "num_experts_per_tok": self.args.n_activated_experts, + "first_num_hash_layers": self.args.n_hash_layers, + } + _apply_layer_param_mapping(mt_layer, pt_block, layer_idx, pt_config_dict, self.mx_config) + + x_np = np.random.uniform(0.1, 1.0, size=(self.batch_size, self.seq_len, self.args.hc_mult, self.dim)).astype(np.float32) + pos_np = np.arange(self.seq_len)[None, :].repeat(self.batch_size, axis=0) + input_ids_np = np.random.randint(0, self.vocab_size, size=(self.batch_size, self.seq_len)) + + # PyTorch forward + pt_out = pt_block( + x=torch.tensor(x_np), + start_pos=0, + input_ids=torch.tensor(input_ids_np, dtype=torch.long), + ).detach().numpy() + + # MaxText forward + mt_out, _ = mt_layer( + inputs=jnp.array(x_np), + decoder_segment_ids=jnp.ones_like(pos_np, dtype=jnp.int32), + decoder_positions=jnp.array(pos_np), + deterministic=True, + model_mode="train", + decoder_input_tokens=jnp.array(input_ids_np), + ) + mt_out_np = np.array(mt_out) + + max_diff = np.max(np.abs(mt_out_np - pt_out)) + mean_diff = np.mean(np.abs(mt_out_np - pt_out)) + print(f"Layer {layer_idx} Parity -> max_diff: {max_diff:.6e}, mean_diff: {mean_diff:.6e}") + np.testing.assert_allclose(mt_out_np, pt_out, rtol=8e-2, atol=8e-2) + + def test_layer_0_sliding_hash(self): + self._run_layer_test(0) + + def test_layer_2_csa_hash(self): + self._run_layer_test(2) + + def test_layer_3_hca_topk(self): + self._run_layer_test(3) + + def test_layer_4_csa_topk(self): + self._run_layer_test(4) + + +# ============================================================================== +# Full Model Functional & Parity Test +# ============================================================================== +from maxtext.models.models import Transformer + +class Transformer_PT(nn_pt.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.vocab_size = args.vocab_size + self.n_layers = args.n_layers + self.embed = ParallelEmbedding_PT(args.vocab_size, args.dim) + self.layers = nn_pt.ModuleList([Block_PT(i, args) for i in range(args.n_layers)]) + self.norm = RMSNorm_PT(args.dim, args.norm_eps) + self.head = ParallelHead_PT(args.vocab_size, args.dim, args) + if args.hc_mult > 1: + self.hc_head_fn = nn_pt.Parameter(torch.randn(args.hc_mult, args.hc_mult * args.dim, dtype=torch.float32) * 0.02) + self.hc_head_base = nn_pt.Parameter(torch.zeros(args.hc_mult, dtype=torch.float32)) + self.hc_head_scale = nn_pt.Parameter(torch.ones(1, dtype=torch.float32)) + + def forward(self, input_ids: torch.Tensor, start_pos: int = 0): + h = self.embed(input_ids) + if self.args.hc_mult > 1: + h = h.unsqueeze(2).expand(-1, -1, self.args.hc_mult, -1) + for layer in self.layers: + h = layer(h, start_pos, input_ids) + + if self.args.hc_mult > 1: + logits = self.head(h, self.hc_head_fn, self.hc_head_scale, self.hc_head_base, self.norm) + else: + h_norm = self.norm(h) + logits = F.linear(h_norm, self.head.weight) + + return logits + + +class DeepSeekV4FlashFullModelTest(unittest.TestCase): + """Validates full PyTorch Transformer_PT against MaxText Transformer (DeepSeek4).""" + + def setUp(self): + self.batch_size = 2 + self.seq_len = 8 + self.vocab_size = 256 + self.args = ModelArgs(vocab_size=self.vocab_size) + self.dim = self.args.dim + self.mx_config = get_maxtext_config(vocab_size=self.vocab_size) + + def test_full_model_parity(self): + rng = jax.random.PRNGKey(0) + pt_config_dict = {"num_hidden_layers": 5, "num_hash_layers": 3} + + pt_model = Transformer_PT(self.args) + pt_model.eval() + + mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]), ("tensor",)) + mt_model = Transformer( + config=self.mx_config, + mesh=mesh, + quant=None, + model_mode="train", + rngs=nnx.Rngs(0), + ) + + input_ids = jnp.array([[1, 2, 3, 4, 5, 6, 7, 8], [8, 7, 6, 5, 4, 3, 2, 1]]) + + # Map Decoder Layers + for i in range(self.args.n_layers): + _apply_layer_param_mapping(getattr(mt_model.decoder, f'layers_{i}'), pt_model.layers[i], i, pt_config_dict, self.mx_config) + + _apply_global_param_mapping(mt_model, pt_model, pt_config_dict, self.mx_config) + + # Do forward pass + pt_model.eval() + import torch + with torch.no_grad(): + logits_pt = pt_model(torch.tensor(input_ids.tolist(), dtype=torch.long)) + + mt_out = mt_model( + decoder_input_tokens=input_ids, + decoder_positions=jnp.arange(self.seq_len)[None, :].repeat(self.batch_size, axis=0), + model_mode="train", + ) + # 1. Embedding check + pt_emb = pt_model.embed(torch.tensor(input_ids.tolist(), dtype=torch.long)).detach().numpy() + mt_emb = np.array(mt_model.token_embedder(input_ids)) + print(f"EMBEDDING DIFF: max={np.max(np.abs(pt_emb - mt_emb)):.6e}", flush=True) + + # Compare Layer 0 sub-blocks + pt_block = pt_model.layers[0] + mt_block = getattr(mt_model.decoder, "layers_0") + + pt_h = torch.tensor(pt_emb) + if self.args.hc_mult > 1: + pt_h = pt_h.unsqueeze(2).expand(-1, -1, self.args.hc_mult, -1) + + positions = jnp.broadcast_to(jnp.arange(self.seq_len, dtype=jnp.int32)[None, :], (self.batch_size, self.seq_len)) + segment_ids = jnp.zeros((self.batch_size, self.seq_len), dtype=jnp.int32) + + # 2a. Attention pre-mhc + pt_x_attn, pt_post, pt_comb = pt_block.hc_pre(pt_h, pt_block.hc_attn_fn, pt_block.hc_attn_scale, pt_block.hc_attn_base) + print(f"PT HC_PRE ATTN: {pt_x_attn.shape}", flush=True) + + # 2b. Compare Attention core + pt_attn_in = pt_block.attn_norm(pt_x_attn) + pt_attn_out = pt_block.attn(pt_attn_in, 0).detach().numpy() + + mt_attn_in = np.array(mt_block.pre_self_attention_layer_norm(jnp.array(pt_x_attn.detach().numpy()))) + print(f"ATTN NORM DIFF: max={np.max(np.abs(pt_attn_in.detach().numpy() - mt_attn_in)):.6e}", flush=True) + + # Q projection + pt_q_latent = pt_block.attn.q_norm(pt_block.attn.wq_a(pt_attn_in)) + mt_q_latent = mt_block.self_attention.q_norm(mt_block.self_attention.wq_a(jnp.array(mt_attn_in))) + print(f"ATTN Q_LATENT DIFF: max={np.max(np.abs(pt_q_latent.detach().numpy() - np.array(mt_q_latent))):.6e}", flush=True) + + pt_q = pt_block.attn.wq_b(pt_q_latent).view(self.batch_size, self.seq_len, self.args.n_heads, self.args.head_dim) + mt_q = mt_block.self_attention.wq_b(mt_q_latent) + print(f"ATTN Q_UP DIFF: max={np.max(np.abs(pt_q.detach().numpy() - np.array(mt_q))):.6e}", flush=True) + + # Q RoPE + qr_pt = pt_block.attn.freqs_cis[0 : self.seq_len] + pt_q_roped = apply_partial_rotary_emb(pt_q, qr_pt) + mt_q_roped = mt_block.self_attention._apply_rotary_embedding_v4(mt_q, positions, unsqueeze_dim=-2) + print(f"ATTN Q_ROPED DIFF: max={np.max(np.abs(pt_q_roped.detach().numpy() - np.array(mt_q_roped))):.6e}", flush=True) + + # KV projection + pt_kv = pt_block.attn.wkv(pt_attn_in) + if self.args.qk_nope_head_dim > 0: + pt_kv = pt_block.attn.kv_norm(pt_kv) + pt_k_roped = apply_partial_rotary_emb(pt_kv.unsqueeze(2), qr_pt).squeeze(2).unsqueeze(2).expand(-1, -1, self.args.n_heads, -1) + + mt_k_roped, _ = mt_block.self_attention.compressed_kv_projection(jnp.array(mt_attn_in), positions, "train") + print(f"ATTN KV_ROPED DIFF: max={np.max(np.abs(pt_k_roped.detach().numpy() - np.array(mt_k_roped))):.6e}", flush=True) + + # Test CompressedAttention directly + mt_attn_out, _ = mt_block.self_attention( + inputs_q=jnp.array(pt_attn_in.detach().numpy()), + inputs_kv=jnp.array(pt_attn_in.detach().numpy()), + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode="train", + ) + mt_attn_out = np.array(mt_attn_out) + print(f"ATTENTION CORE DIFF: max={np.max(np.abs(pt_attn_out - mt_attn_out)):.6e}", flush=True) + + # 2c. Compare MoE core + pt_x_ffn, _, _ = pt_block.hc_pre(pt_h, pt_block.hc_ffn_fn, pt_block.hc_ffn_scale, pt_block.hc_ffn_base) + pt_ffn_in = pt_block.ffn_norm(pt_x_ffn) + pt_ffn_out = pt_block.ffn(pt_ffn_in, torch.tensor(input_ids.tolist(), dtype=torch.long)).detach().numpy() + + mt_ffn_in = np.array(mt_block.post_self_attention_layer_norm(jnp.array(pt_x_ffn.detach().numpy()))) + print(f"FFN NORM DIFF: max={np.max(np.abs(pt_ffn_in.detach().numpy() - mt_ffn_in)):.6e}", flush=True) + + # Compare Shared Experts + pt_shared = pt_block.ffn.shared_experts(pt_ffn_in).detach().numpy() + mt_shared = np.array(mt_block.mlp.shared_experts(jnp.array(pt_ffn_in.detach().numpy()))) + print(f"MOE SHARED EXPERTS DIFF: max={np.max(np.abs(pt_shared - mt_shared)):.6e}", flush=True) + + # Compare Gate / Routing + pt_w, pt_idx = pt_block.ffn.gate(pt_ffn_in, torch.tensor(input_ids.tolist(), dtype=torch.long)) + gate_in = jnp.array(pt_ffn_in.detach().numpy()) + mt_gl, mt_pbl = mt_block.mlp.routed_moe.gate(gate_in) + mt_w, mt_idx = mt_block.mlp.routed_moe.get_topk(mt_gl, mt_pbl, input_ids=input_ids) + print(f"GATE INDICES MATCH: {np.array_equal(pt_idx.detach().numpy(), np.array(mt_idx))}", flush=True) + print(f"GATE WEIGHTS DIFF: max={np.max(np.abs(pt_w.detach().numpy() - np.array(mt_w))):.6e}", flush=True) + + # Compare Routed Experts + mt_routed, _, _ = mt_block.mlp.routed_moe(gate_in, input_ids=input_ids) + pt_routed = (pt_block.ffn(pt_ffn_in, torch.tensor(input_ids.tolist(), dtype=torch.long)) - pt_block.ffn.shared_experts(pt_ffn_in)).detach().numpy() + print(f"MOE ROUTED EXPERTS DIFF: max={np.max(np.abs(pt_routed - np.array(mt_routed))):.6e}", flush=True) + + mt_ffn_out, _, _ = mt_block.mlp( + jnp.array(pt_ffn_in.detach().numpy()), + input_ids=input_ids, + ) + print(f"MOE CORE DIFF: max={np.max(np.abs(pt_ffn_out - np.array(mt_ffn_out))):.6e}", flush=True) + + # Layer by Layer comparison + pt_curr_h = pt_h.clone() + mt_curr_h = jnp.array(pt_h.detach().numpy()) + for layer_i in range(self.args.n_layers): + pt_layer = pt_model.layers[layer_i] + mt_layer = getattr(mt_model.decoder, f"layers_{layer_i}") + + if layer_i == 2: + # Step by step layer 2 breakdown + l2_pt_x_attn, _, _ = pt_layer.hc_pre(pt_curr_h, pt_layer.hc_attn_fn, pt_layer.hc_attn_scale, pt_layer.hc_attn_base) + l2_pt_attn_in = pt_layer.attn_norm(l2_pt_x_attn) + l2_mt_attn_in = np.array(mt_layer.pre_self_attention_layer_norm(jnp.array(l2_pt_x_attn.detach().numpy()))) + print(f"L2 ATTN NORM DIFF: {np.max(np.abs(l2_pt_attn_in.detach().numpy() - l2_mt_attn_in)):.6e}", flush=True) + + l2_pt_q_latent = pt_layer.attn.q_norm(pt_layer.attn.wq_a(l2_pt_attn_in)) + l2_mt_q_latent = mt_layer.self_attention.q_norm(mt_layer.self_attention.wq_a(jnp.array(l2_mt_attn_in))) + print(f"L2 Q_LATENT DIFF: {np.max(np.abs(l2_pt_q_latent.detach().numpy() - np.array(l2_mt_q_latent))):.6e}", flush=True) + + l2_pt_comp_kv, _ = pt_layer.attn.compressor(l2_pt_attn_in, 0) + l2_mt_comp_kv, l2_mt_comp_mask = mt_layer.self_attention.csa_compressor( + jnp.array(l2_pt_attn_in.detach().numpy()), + jnp.array(l2_pt_q_latent.detach().numpy()), + positions, + None, + "train", + ) + print(f"L2 COMP KV DIFF: {np.max(np.abs(l2_pt_comp_kv.detach().numpy() - np.array(l2_mt_comp_kv[:, :, 0, :]))):.6e}", flush=True) + + l2_pt_topk, l2_pt_mask = pt_layer.attn.indexer(l2_pt_attn_in, l2_pt_q_latent, 0) + l2_mt_topk = mt_layer.self_attention.csa_compressor.indexer( + jnp.array(l2_pt_attn_in.detach().numpy()), + jnp.array(l2_pt_q_latent.detach().numpy()), + positions, + None, + "train", + ) + print(f"L2 INDEXER TOPK MATCH: {np.array_equal(l2_pt_topk.detach().numpy(), np.array(l2_mt_topk))}", flush=True) + print(f"L2 PT TOPK:\n{l2_pt_topk}", flush=True) + print(f"L2 MT TOPK:\n{l2_mt_topk}", flush=True) + + l2_pt_attn_out = pt_layer.attn(l2_pt_attn_in, 0).detach().numpy() + l2_mt_attn_out, _ = mt_layer.self_attention( + inputs_q=jnp.array(l2_pt_attn_in.detach().numpy()), + inputs_kv=jnp.array(l2_pt_attn_in.detach().numpy()), + decoder_segment_ids=segment_ids, + inputs_positions=positions, + deterministic=True, + model_mode="train", + ) + print(f"L2 ATTN OUT DIFF: {np.max(np.abs(l2_pt_attn_out - np.array(l2_mt_attn_out))):.6e}", flush=True) + + pt_curr_h = pt_layer(pt_curr_h, 0, torch.tensor(input_ids.tolist(), dtype=torch.long)) + mt_curr_h, _ = mt_layer( + mt_curr_h, + decoder_segment_ids=segment_ids, + decoder_positions=positions, + deterministic=True, + model_mode="train", + decoder_input_tokens=input_ids, + ) + diff = np.max(np.abs(pt_curr_h.detach().numpy() - np.array(mt_curr_h))) + print(f"LAYER {layer_i} [rate={self.args.compress_rates[layer_i]}]: OUTPUT DIFF: max={diff:.6e}", flush=True) + + # Head comparison + pt_head_out = pt_model.head.hc_head(pt_curr_h, pt_model.hc_head_fn, pt_model.hc_head_scale, pt_model.hc_head_base) + mt_head_out = mt_model.decoder.hc_head(mt_curr_h) + print(f"pt_head_out: shape={pt_head_out.shape}, min={pt_head_out.min():.4f}, max={pt_head_out.max():.4f}", flush=True) + print(f"mt_head_out: shape={mt_head_out.shape}, min={mt_head_out.min():.4f}, max={mt_head_out.max():.4f}", flush=True) + print(f"HC_HEAD DIFF: max={np.max(np.abs(pt_head_out.detach().numpy() - np.array(mt_head_out))):.6e}", flush=True) + + pt_norm_out = pt_model.norm(pt_head_out) + mt_norm_out = mt_model.decoder.decoder_norm(mt_head_out) + print(f"pt_norm_out: shape={pt_norm_out.shape}, min={pt_norm_out.min():.4f}, max={pt_norm_out.max():.4f}", flush=True) + print(f"mt_norm_out: shape={mt_norm_out.shape}, min={mt_norm_out.min():.4f}, max={mt_norm_out.max():.4f}", flush=True) + print(f"DECODER NORM DIFF: max={np.max(np.abs(pt_norm_out.detach().numpy() - np.array(mt_norm_out))):.6e}", flush=True) + + pt_logits = F.linear(pt_norm_out, pt_model.head.weight) + mt_logits = mt_model.decoder.logits_dense(mt_norm_out) + print(f"LOGITS DIFF: max={np.max(np.abs(pt_logits.detach().numpy() - np.array(mt_logits))):.6e}", flush=True) + + def assert_close(a, b, name): + max_diff = jnp.max(jnp.abs(a - b)) + mean_diff = jnp.mean(jnp.abs(a - b)) + print(f"{name} Parity -> max_diff: {max_diff:e}, mean_diff: {mean_diff:e}", flush=True) + np.testing.assert_allclose(a, b, atol=1e-2, rtol=1e-2) + + assert_close(np.array(mt_out), logits_pt.detach().numpy(), "Full Model Logits") + + def test_scanned_full_model_parity(self): + """Validates full PyTorch Transformer_PT against MaxText Transformer (DeepSeek4) with scan_layers=True.""" + # 7 layers: 3 prefix [0, 0, 4] + 4 scanned (2 blocks of [128, 4]) + n_layers = 7 + compress_rates = [0, 0, 4, 128, 4, 128, 4] + pt_config_dict = { + "num_hidden_layers": n_layers, + "num_hash_layers": 3, + } + + test_vocab_size = 256 + scanned_args = ModelArgs( + vocab_size=test_vocab_size, + n_layers=n_layers, + compress_rates=compress_rates, + n_hash_layers=3, + n_routed_experts=4, + n_activated_experts=2, + ) + + scanned_mx_config = get_maxtext_config( + vocab_size=test_vocab_size, + base_num_decoder_layers=n_layers, + compress_ratios=compress_rates, + first_num_hash_layers=3, + num_experts=4, + num_experts_per_tok=2, + scan_layers=True, + ) + + pt_model = Transformer_PT(scanned_args) + pt_model.eval() + + maxtext_model, params = _build_scanned_maxtext_params(pt_model, pt_config_dict, scanned_mx_config) + + input_ids = jnp.array([[1, 2, 3, 4, 5, 6, 7, 8], [8, 7, 6, 5, 4, 3, 2, 1]]) + + # PyTorch Forward Pass + with torch.no_grad(): + logits_pt = pt_model(torch.tensor(input_ids.tolist(), dtype=torch.long)) + + # MaxText Linen Forward Pass + positions = jnp.broadcast_to(jnp.arange(self.seq_len, dtype=jnp.int32)[None, :], (self.batch_size, self.seq_len)) + segment_ids = jnp.zeros((self.batch_size, self.seq_len), dtype=jnp.int32) + + @jax.jit + def run_mt(p, tokens, pos, seg): + return maxtext_model.apply( + p, + decoder_input_tokens=tokens, + decoder_positions=pos, + decoder_segment_ids=seg, + enable_dropout=False, + model_mode="train", + decoder_target_tokens=tokens, + decoder_target_mask=seg, + ) + + mt_out = run_mt(params, input_ids, positions, segment_ids) + + def assert_close(a, b, name): + max_diff = jnp.max(jnp.abs(a - b)) + mean_diff = jnp.mean(jnp.abs(a - b)) + print(f"{name} Parity -> max_diff: {max_diff:e}, mean_diff: {mean_diff:e}", flush=True) + np.testing.assert_allclose(a, b, atol=1e-2, rtol=1e-2) + + assert_close(np.array(mt_out), logits_pt.detach().numpy(), "Scanned Full Model Logits") + + +if __name__ == "__main__": + unittest.main() + + + diff --git a/tests/unit/deepseek_v4_vs_reference_test.py b/tests/unit/deepseek_v4_vs_reference_test.py index 66e052cdbb..3b4bb09816 100644 --- a/tests/unit/deepseek_v4_vs_reference_test.py +++ b/tests/unit/deepseek_v4_vs_reference_test.py @@ -1231,7 +1231,7 @@ def setUp(self): ) config_arguments = { - "model_name": "deepseek4-tiny", + "model_name": "deepseek4-284b", "override_model_config": True, "per_device_batch_size": 1, "matmul_precision": "highest", @@ -1471,7 +1471,7 @@ def setUp(self): self.rngs = nnx.Rngs(0) # Build MaxText config dictionary - argv = ["", "src/maxtext/configs/base.yml", "model_name=deepseek4-tiny"] + argv = ["", "src/maxtext/configs/base.yml", "model_name=deepseek4-284b"] config_arguments = { "attention": "dot_product", "dtype": "float32",