diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index c2ff040979..22ea98009c 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -87,6 +87,67 @@ absl.logging.set_verbosity(absl.logging.INFO) # for max_logging.log +def get_key_mapper(model_name): + if model_name.startswith("deepseek4"): + def map_hf_key(hf_key): + # Rule 11: Global keys + if hf_key == "model.embed_tokens.weight": + return "embed.weight" + if hf_key == "model.norm.weight": + return "norm.weight" + if hf_key.startswith("model.hc_head.hc_"): + return hf_key.replace("model.hc_head.hc_", "hc_head_") + + k = hf_key + if k.startswith("model."): + k = k[6:] + + k = k.replace("experts..", "experts.") + k = k.replace(".self_attn.", ".attn.") + + k = k.replace(".attn_hc.fn", ".hc_attn_fn") + k = k.replace(".attn_hc.base", ".hc_attn_base") + k = k.replace(".attn_hc.scale", ".hc_attn_scale") + k = k.replace(".ffn_hc.fn", ".hc_ffn_fn") + k = k.replace(".ffn_hc.base", ".hc_ffn_base") + k = k.replace(".ffn_hc.scale", ".hc_ffn_scale") + + k = k.replace(".input_layernorm.weight", ".attn_norm.weight") + k = k.replace(".post_attention_layernorm.weight", ".ffn_norm.weight") + k = k.replace(".mlp.", ".ffn.") + k = k.replace("e_score_correction_bias", "bias") + + k = k.replace(".shared_experts.gate_proj.weight", ".shared_experts.w1.weight") + k = k.replace(".shared_experts.up_proj.weight", ".shared_experts.w3.weight") + k = k.replace(".shared_experts.down_proj.weight", ".shared_experts.w2.weight") + + k = k.replace(".compressor.indexer.gate_proj.weight", ".indexer.compressor.wgate.weight") + k = k.replace(".compressor.indexer.kv_proj.weight", ".indexer.compressor.wkv.weight") + k = k.replace(".compressor.indexer.q_b_proj.weight", ".indexer.wq_b.weight") + k = k.replace(".compressor.indexer.scorer.weights_proj.weight", ".indexer.weights_proj.weight") + k = k.replace(".compressor.indexer.position_bias", ".indexer.compressor.ape") + k = k.replace(".compressor.indexer.kv_norm.weight", ".indexer.compressor.norm.weight") + + k = k.replace(".compressor.gate_proj.weight", ".compressor.wgate.weight") + k = k.replace(".compressor.kv_proj.weight", ".compressor.wkv.weight") + k = k.replace(".compressor.position_bias", ".compressor.ape") + k = k.replace(".compressor.kv_norm.weight", ".compressor.norm.weight") + + k = k.replace(".attn.q_a_proj.weight", ".attn.wq_a.weight") + k = k.replace(".attn.q_a_norm.weight", ".attn.q_norm.weight") + k = k.replace(".attn.q_b_proj.weight", ".attn.wq_b.weight") + k = k.replace(".attn.kv_proj.weight", ".attn.wkv.weight") + k = k.replace(".attn.sinks", ".attn.attn_sink") + k = k.replace(".attn.o_a_proj.weight", ".attn.wo_a.weight") + k = k.replace(".attn.o_b_proj.weight", ".attn.wo_b.weight") + + return k + return map_hf_key + else: + return lambda k: k + + + class LazyHFLoader: """ Loads Hugging Face weights on-demand to minimize RAM usage. @@ -108,10 +169,11 @@ class LazyHFLoader: can still occur in parallel. """ - def __init__(self, model_id, token, revision=None): + def __init__(self, model_id, token, revision=None, map_fn=None): self.model_id = model_id self.token = token self.revision = revision + self.map_fn = map_fn or (lambda k: k) # Whether loads from local directory self.is_local = os.path.isdir(self.model_id) self.shard_map = {} @@ -178,13 +240,17 @@ def get_tensor(self, key: str) -> np.ndarray: and reads only the required tensor's data from disk. """ # Handle single-file models (shard map key might be None or we just know the filename) - shard_name = self.shard_map.get(key) + mapped_key = self.map_fn(key) + shard_name = self.shard_map.get(mapped_key) if shard_name is None and None in self.shard_map: shard_name = self.shard_map[None] elif shard_name is None: - # Fallback: sometimes keys in index don't perfectly match requested keys if there are prefix mismatches. - # You might need advanced fuzzy matching here if you encounter errors. - raise ValueError(f"Key {key} not found in HF checkpoint index.") + # Fallback to unmapped key if mapped is not found + shard_name = self.shard_map.get(key) + if shard_name is not None: + mapped_key = key + else: + raise ValueError(f"Key {key} (mapped to {mapped_key}) not found in HF checkpoint index.") if shard_name in self._local_shard_paths: local_path = self._local_shard_paths[shard_name] @@ -206,7 +272,7 @@ def get_tensor(self, key: str) -> np.ndarray: # This prevents multiple threads from simultaneously allocating large chunks of RAM. with self._ram_lock: with safe_open(local_path, framework="np", device="cpu") as f: - return f.get_tensor(key) + return f.get_tensor(mapped_key) class LazyTensor: @@ -318,7 +384,7 @@ def get_maxtext_model_info(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 - abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model_flax, config)["params"] + 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, @@ -389,29 +455,34 @@ def _build_single_axis_stacked_tensor( hook_fns: Any, target_shape: tuple, config, + mt_param_name: str = "", ) -> np.ndarray: - """Builds a MaxText tensor by stacking HF weights along a single axis. + """Builds a MaxText tensor by stacking HF weights along one axis. - This function handles both standard scanned layers (e.g., attention) and - unscanned MoE layers (which are stacked along the expert axis). + This function handles two cases: + 1. Stacking layers for standard scanned blocks (e.g. attention/norm layers). + In this case, axis_to_stack is config.param_scan_axis. + 2. Stacking experts for unscanned MoeBlock layers. + In this case, axis_to_stack is 0. Args: - hf_source_keys: A 1D list of Hugging Face parameter names. - tensor_getter_fn: A callable that takes a HF key and returns the tensor (as numpy array). + hf_source_keys: A list of Hugging Face parameter names. + tensor_getter_fn: A callable that takes a HF key and returns the tensor. hook_fns: The hook function(s) to apply to each individual weight. target_shape: The final shape of the target MaxText tensor. config: The MaxText pyconfig object. + mt_param_name: The name of the MaxText parameter. Returns: The final, assembled NumPy array for the MaxText parameter. """ tensors_to_stack = [] - if config.scan_layers: + if config.scan_layers and "scanned_blocks" in mt_param_name: # If it's a standard scanned layer, we use the configured param_scan_axis. axis_to_stack = config.param_scan_axis else: - # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0). + # Otherwise, if an unscanned MoE layer (or scan_layers is False), 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. @@ -432,7 +503,14 @@ def _build_single_axis_stacked_tensor( return np.stack(tensors_to_stack, axis=axis_to_stack) -def _get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_target_shape_or_shapes, config): +def _get_hf_loading_function( + hf_source_keys_or_key, + tensor_getter, + hook_fn, + mt_target_shape_or_shapes, + config, + mt_param_name: str = "", +): """Determine the loading function for HF keys. This function natively supports `composite_hf_key` mapping (where multiple HF keys @@ -450,6 +528,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) @@ -472,6 +552,7 @@ def _loader(getter, key, shape, hook): hook_fn, mt_target_shape_or_shapes, config, + mt_param_name=mt_param_name, ) else: # isinstance(hf_source_keys_or_key[0], list) @@ -891,11 +972,13 @@ def main( hf_state_dict_numpy = None hf_loader = None + model_name_for_path = model_name_original or config.model_name + map_fn = get_key_mapper(model_name_for_path) # Define the appropriate tensor getter based on mode if lazy_load_tensors: max_logging.log(f"Lazy loading ENABLED. Initializing LazyHFLoader for: {model_id}...") - hf_loader = LazyHFLoader(model_id, hf_token, revision=revision) + hf_loader = LazyHFLoader(model_id, hf_token, revision=revision, map_fn=map_fn) print_ram_usage("After LazyLoader init") tensor_getter = hf_loader.get_tensor @@ -955,9 +1038,13 @@ def main( } def _eager_getter(key): - 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] + mapped_key = map_fn(key) + if mapped_key not in hf_state_dict_numpy: + if key in hf_state_dict_numpy: + mapped_key = key + else: + raise ValueError(f"HuggingFace key {key} (mapped to {mapped_key}) not found in state_dict.") + v = hf_state_dict_numpy[mapped_key] # target dtype is "float32" if save_dtype == DType.FLOAT32: return v.to(torch.float32).numpy() @@ -981,7 +1068,7 @@ def _eager_getter(key): model_key = config.model_name # load config hf_config_obj = HF_MODEL_CONFIGS[model_key] - hf_config_dict = hf_config_obj.to_dict() + hf_config_dict = hf_config_obj.to_dict() if hasattr(hf_config_obj, "to_dict") else hf_config_obj # example of param mapping (gemma2, maxtext:huggingface): # "params-decoder-layers_{maxtext_layer_idx}-pre_self_attention_norm_global-scale": # f"model.layers.{global_layer_idx}.input_layernorm.weight", @@ -1013,9 +1100,9 @@ def _eager_getter(key): if not lazy_load_tensors: max_logging.log(f"maxtext param: {mt_param_key_or_keys}") - hf_source_keys_or_key = param_map_mt_to_hf.get(mt_param_key_or_keys) - if hf_source_keys_or_key is None: + if mt_param_key_or_keys not in param_map_mt_to_hf: raise ValueError(f"MaxText parameter {mt_param_key_or_keys} not found in mapping.") + hf_source_keys_or_key = param_map_mt_to_hf[mt_param_key_or_keys] hook_fn = hook_fn_map_mt.get(mt_param_key_or_keys) # Step 1: Resolves MaxText key(s) to target indices and shapes @@ -1032,6 +1119,7 @@ def _eager_getter(key): hook_fn, mt_target_shape_or_shapes, config, + mt_param_name=mt_param_key_or_keys, ) # Step 3: Load hf keys and convert to maxtext keys diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 8bb262958c..5d4a2182b9 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1701,6 +1701,124 @@ def __init__(self, **kwargs): # {maxtext model name: hf model config} HF_MODEL_CONFIGS = { + "deepseek4-tiny": { + "architectures": [ + "DeepseekV4ForCausalLM" + ], + "attention_bias": False, + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_deepseek_v4.DeepseekV4Config", + "AutoModel": "modeling_deepseek_v4.DeepseekV4Model", + "AutoModelForCausalLM": "modeling_deepseek_v4.DeepseekV4ForCausalLM" + }, + "bos_token_id": 100000, + "compress_rates": [1, 1, 1], + "csa_heads": 8, + "eos_token_id": 100001, + "first_k_dense_replace": 3, + "hca_heads": 8, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 2048, + "kv_lora_rank": 512, + "max_position_embeddings": 163840, + "model_type": "deepseek_v4", + "moe_intermediate_size": 1024, + "moe_layer_freq": 1, + "n_group": 4, + "n_routed_experts": 8, + "n_shared_experts": 1, + "norm_topk_prob": False, + "num_attention_heads": 64, + "num_experts_per_tok": 3, + "num_hidden_layers": 7, + "num_key_value_heads": 64, + "num_nextn_predict_layers": 1, + "pretraining_tp": 1, + "q_lora_rank": 1024, + "qk_nope_head_dim": 64, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 40, + "mscale": 1.0, + "mscale_all_dim": 1.0, + "original_max_position_embeddings": 4096, + "type": "yarn" + }, + "rope_theta": 10000.0, + "routed_scaling_factor": 16.0, + "scoring_func": "sigmoid", + "seq_aux": True, + "tie_word_embeddings": False, + "topk_group": 2, + "topk_method": "group_limited_greedy", + "use_cache": True, + "v_head_dim": 64, + "vocab_size": 102400 + }, + "deepseek4-284b": { + "architectures": [ + "DeepseekV4ForCausalLM" + ], + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 0, + "eos_token_id": 1, + "expert_dtype": "fp4", + "hc_eps": 1e-06, + "hc_mult": 4, + "hc_sinkhorn_iters": 20, + "head_dim": 512, + "hidden_act": "silu", + "hidden_size": 4096, + "index_head_dim": 128, + "index_n_heads": 64, + "index_topk": 512, + "initializer_range": 0.02, + "max_position_embeddings": 1048576, + "model_type": "deepseek_v4", + "moe_intermediate_size": 2048, + "n_routed_experts": 256, + "n_shared_experts": 1, + "norm_topk_prob": True, + "num_attention_heads": 64, + "num_experts_per_tok": 6, + "num_hidden_layers": 43, + "num_hash_layers": 3, + "num_key_value_heads": 1, + "num_nextn_predict_layers": 1, + "o_groups": 8, + "o_lora_rank": 1024, + "q_lora_rank": 1024, + "qk_rope_head_dim": 64, + "rms_norm_eps": 1e-06, + "rope_scaling": { + "beta_fast": 32, + "beta_slow": 1, + "factor": 16, + "original_max_position_embeddings": 65536, + "type": "yarn" + }, + "rope_theta": 10000.0, + "routed_scaling_factor": 1.5, + "scoring_func": "sqrtsoftplus", + "sliding_window": 128, + "swiglu_limit": 10.0, + "tie_word_embeddings": False, + "topk_method": "noaux_tc", + "torch_dtype": "bfloat16", + "transformers_version": "4.57.1", + "use_cache": True, + "vocab_size": 129280, + "compress_rope_theta": 160000, + "compress_ratios": [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, 0] + }, + "gemma2-2b": gemma2_2b_config, "gemma2-9b": gemma2_9b_config, "gemma2-27b": gemma2_27b_config, diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 80707a2abd..61bb195715 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -3858,6 +3858,232 @@ def reshape_vision_attn_out(input_tensor, target_shape): # {maxtext model name: {maxtext weight name: hf weight name}} + +def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): + n_layers = config["num_hidden_layers"] + num_experts = config.get("n_routed_experts", 8) + + mapping = { + "params-token_embedder-embedding": "model.embed_tokens.weight", + "params-decoder-decoder_norm-scale": "model.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", + } + + 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 or first_idx == 2: + 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) + + 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", [3, 5]) + add_layer_mapping("params-decoder-scanned_blocks-layers_1", [4, 6]) + + for i in range(3): + mapping[f"Tid2EidVar-decoder-layers_{i}-mlp-MoeBlock_0-tid2eid"] = f"model.layers.{i}.mlp.gate.tid2eid" + + return mapping + +def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): + def transpose(input_tensor, target_shape=None): + return np.transpose(input_tensor) + + def transpose_stack(input_tensor, target_shape=None): + # input_tensor is a list of tensors + stacked = np.stack(input_tensor, axis=0) # [E, out, in] + return np.transpose(stacked, (0, 2, 1)) # [E, in, out] + + def ones_norm(input_tensor, target_shape=None): + return np.ones(target_shape, dtype=np.float32) + + def identity(input_tensor, target_shape=None): + return input_tensor + + # 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] + tensor = np.transpose(input_tensor) # [kv_lora_rank, n_heads * q_head_dim] + n_heads = config["num_attention_heads"] + 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] + 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. + num_heads = target_shape[0] + embed_dim = target_shape[1] + kv_lora_rank = target_shape[2] + 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) + + def mhc_split_scale_pre(input_tensor, target_shape=None): + return np.array([input_tensor[0]]).reshape(target_shape) + def mhc_split_scale_post(input_tensor, target_shape=None): + return np.array([input_tensor[1]]).reshape(target_shape) + def mhc_split_scale_res(input_tensor, target_shape=None): + return np.array([input_tensor[2]]).reshape(target_shape) + + mapping = {} + + # 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 type(hf_key) == list: + mapping[key] = transpose if not saving_to_hf else transpose_stack + elif "-kernel" in key or "-embedding" in key or "-sinks" in key: + mapping[key] = transpose + + return mapping + PARAM_MAPPING = { "gemma2-2b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, "gemma2-9b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -3896,6 +4122,8 @@ def reshape_vision_attn_out(input_tensor, target_shape): "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-tiny": DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING, + "deepseek4-284b": DEEPSEEKV4_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, @@ -3948,6 +4176,8 @@ def reshape_vision_attn_out(input_tensor, target_shape): "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, "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/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 75ee3dffcc..b92ea936fa 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -1236,8 +1236,13 @@ def save_weights_to_checkpoint( if checkpoint_manager is None: raise RuntimeError("Failed to create Orbax checkpoint manager.") + if isinstance(jax_weights, dict) and "params" in jax_weights: + state_params = jax_weights + else: + state_params = {"params": jax_weights} + state_new = train_state.TrainState( - step=step_number_to_save_new_ckpt, apply_fn=None, params={"params": jax_weights}, tx=None, opt_state={} # type: ignore + step=step_number_to_save_new_ckpt, apply_fn=None, params=state_params, tx=None, opt_state={} # type: ignore ) logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))