diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index ce2eecb678..64b0c3e62b 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: @@ -517,12 +515,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 @@ -1041,9 +1036,9 @@ def _eager_getter(key): if not lazy_load_tensors: max_logging.log(f"maxtext param: {mt_param_key_or_keys}") - 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.get(mt_param_key_or_keys) + if hf_source_keys_or_key is None: + raise ValueError(f"MaxText parameter {mt_param_key_or_keys} not found in mapping.") hook_fn = hook_fn_map_mt.get(mt_param_key_or_keys) # Step 1: Resolves MaxText key(s) to target indices and shapes @@ -1080,10 +1075,12 @@ 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 + # 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") if lazy_load_tensors and not is_adapter_only: max_logging.log("Starting checkpoint save (loading weights just-in-time)...") @@ -1095,7 +1092,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/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 090c9879de..4fbabf94c5 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -3873,333 +3873,277 @@ def reshape_vision_attn_out(input_tensor, target_shape): 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.""" +def DEEPSEEK_V4_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-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) - - 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))) - - for i in range(3): - mapping[f"Tid2EidVar-decoder-layers_{i}-mlp-MoeBlock_0-tid2eid"] = f"model.layers.{i}.mlp.gate.tid2eid" + for i in range(n_layers): + if scan_layers: + prefix = f"params-decoder-scanned_blocks-layers_{i}" + else: + prefix = f"params-decoder-layers_{i}" + + # HC Norms (DeepSeek V4) mapped to standard MaxText Layer Norms + mapping[f"{prefix}-pre_self_attention_layer_norm-scale"] = f"layers.{i}.attn_norm.weight" + mapping[f"{prefix}-post_self_attention_layer_norm-scale"] = f"layers.{i}.ffn_norm.weight" + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = f"layers.{i}.attn_norm.weight" + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = f"layers.{i}.ffn_norm.weight" + + # MHC Attention Alpha (Atomic) + mapping[f"{prefix}-mhc_attention-pre_alpha"] = f"layers.{i}.hc_attn_fn" + mapping[f"{prefix}-mhc_attention-post_alpha"] = f"layers.{i}.hc_attn_fn" + mapping[f"{prefix}-mhc_attention-res_alpha"] = f"layers.{i}.hc_attn_fn" + + # MHC Attention Beta (Atomic) + mapping[f"{prefix}-mhc_attention-pre_beta"] = f"layers.{i}.hc_attn_base" + mapping[f"{prefix}-mhc_attention-post_beta"] = f"layers.{i}.hc_attn_base" + mapping[f"{prefix}-mhc_attention-res_beta"] = f"layers.{i}.hc_attn_base" + + # MHC Attention Scale (Atomic) + mapping[f"{prefix}-mhc_attention-pre_alpha_scale"] = f"layers.{i}.hc_attn_scale" + mapping[f"{prefix}-mhc_attention-post_alpha_scale"] = f"layers.{i}.hc_attn_scale" + mapping[f"{prefix}-mhc_attention-res_alpha_scale"] = f"layers.{i}.hc_attn_scale" + + # MHC MLP Alpha (Atomic) + mapping[f"{prefix}-mhc_mlp-pre_alpha"] = f"layers.{i}.hc_ffn_fn" + mapping[f"{prefix}-mhc_mlp-post_alpha"] = f"layers.{i}.hc_ffn_fn" + mapping[f"{prefix}-mhc_mlp-res_alpha"] = f"layers.{i}.hc_ffn_fn" + + # MHC MLP Beta (Atomic) + mapping[f"{prefix}-mhc_mlp-pre_beta"] = f"layers.{i}.hc_ffn_base" + mapping[f"{prefix}-mhc_mlp-post_beta"] = f"layers.{i}.hc_ffn_base" + mapping[f"{prefix}-mhc_mlp-res_beta"] = f"layers.{i}.hc_ffn_base" + + # MHC MLP Scale (Atomic) + mapping[f"{prefix}-mhc_mlp-pre_alpha_scale"] = f"layers.{i}.hc_ffn_scale" + mapping[f"{prefix}-mhc_mlp-post_alpha_scale"] = f"layers.{i}.hc_ffn_scale" + mapping[f"{prefix}-mhc_mlp-res_alpha_scale"] = f"layers.{i}.hc_ffn_scale" + + # Attention Core + mapping[f"{prefix}-self_attention-q_norm-scale"] = f"layers.{i}.attn.q_norm.weight" + mapping[f"{prefix}-self_attention-kv_norm-scale"] = f"layers.{i}.attn.kv_norm.weight" + mapping[f"{prefix}-self_attention-wq_a-kernel"] = f"layers.{i}.attn.wq_a.weight" + mapping[f"{prefix}-self_attention-wq_b-kernel"] = f"layers.{i}.attn.wq_b.weight" + mapping[f"{prefix}-self_attention-wkv-kernel"] = f"layers.{i}.attn.wkv.weight" + mapping[f"{prefix}-self_attention-sinks"] = f"layers.{i}.attn.attn_sink" + + # Output projection + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = f"layers.{i}.attn.wo_a.weight" + mapping[f"{prefix}-self_attention-o_b_proj-kernel"] = f"layers.{i}.attn.wo_b.weight" + + # Compressors CSA and HCA + mapping[f"{prefix}-self_attention-csa_compressor-gate_proj-kernel"] = f"layers.{i}.attn.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-csa_compressor-kv_proj-kernel"] = f"layers.{i}.attn.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-csa_compressor-kv_norm-scale"] = f"layers.{i}.attn.compressor.norm.weight" + mapping[f"{prefix}-self_attention-csa_compressor-position_bias"] = f"layers.{i}.attn.compressor.ape" + + mapping[f"{prefix}-self_attention-hca_compressor-gate_proj-kernel"] = f"layers.{i}.attn.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-hca_compressor-kv_proj-kernel"] = f"layers.{i}.attn.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-hca_compressor-kv_norm-scale"] = f"layers.{i}.attn.compressor.norm.weight" + mapping[f"{prefix}-self_attention-hca_compressor-position_bias"] = f"layers.{i}.attn.compressor.ape" + + mapping[f"{prefix}-self_attention-csa_compressor-indexer-gate_proj-kernel"] = f"layers.{i}.attn.indexer.compressor.wgate.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_proj-kernel"] = f"layers.{i}.attn.indexer.compressor.wkv.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-kv_norm-scale"] = f"layers.{i}.attn.indexer.compressor.norm.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-position_bias"] = f"layers.{i}.attn.indexer.compressor.ape" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-weights_proj-kernel"] = f"layers.{i}.attn.indexer.weights_proj.weight" + mapping[f"{prefix}-self_attention-csa_compressor-indexer-q_proj-kernel"] = f"layers.{i}.attn.indexer.wq_b.weight" + + # MoE + num_experts = config.get("n_routed_experts", 0) + if num_experts == 0: + # Failsafe if not found in config + num_experts = 8 + + mapping[f"{prefix}-mlp-MoeBlock_0-gate-kernel"] = f"layers.{i}.ffn.gate.weight" + num_hash_layers = config.get("first_num_hash_layers", 3) + if i < num_hash_layers: + mapping[f"Tid2EidVar-decoder-layers_{i}-mlp-MoeBlock_0-tid2eid"] = f"layers.{i}.ffn.gate.tid2eid" + else: + mapping[f"{prefix}-mlp-MoeBlock_0-gate-bias"] = f"layers.{i}.ffn.gate.bias" + + mapping[f"{prefix}-mlp-shared_experts-wi_0-kernel"] = f"layers.{i}.ffn.shared_experts.w1.weight" + mapping[f"{prefix}-mlp-shared_experts-wi_1-kernel"] = f"layers.{i}.ffn.shared_experts.w3.weight" + mapping[f"{prefix}-mlp-shared_experts-wo-kernel"] = f"layers.{i}.ffn.shared_experts.w2.weight" + + mapping[f"{prefix}-mlp-MoeBlock_0-wi_0"] = [ + f"layers.{i}.ffn.experts.{e}.w1.weight" for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wi_1"] = [ + f"layers.{i}.ffn.experts.{e}.w3.weight" for e in range(num_experts) + ] + mapping[f"{prefix}-mlp-MoeBlock_0-wo"] = [ + f"layers.{i}.ffn.experts.{e}.w2.weight" for e in range(num_experts) + ] return mapping -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.""" - - 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) - - def identity(input_tensor, target_shape=None): - return input_tensor +def DEEPSEEK_V4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False): + n_layers = config["num_hidden_layers"] - # 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): if saving_to_hf: - tensor = np.transpose(input_tensor, (0, 2, 1)) - return tensor.reshape(target_shape) - 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) + 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) - def mhc_split_scale_post(input_tensor, target_shape=None): - return np.array([input_tensor[1]]).reshape(target_shape) + def reshape_o_a_proj(input_tensor, target_shape): + heads, head_dim, in_dim = target_shape[0], target_shape[1], target_shape[2] + if saving_to_hf: + 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: + 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 + return np.ones(target_shape, dtype=input_tensor.dtype) - def mhc_split_scale_res(input_tensor, target_shape=None): - return np.array([input_tensor[2]]).reshape(target_shape) + mapping = { + "params-decoder-logits_dense-kernel": reshape_kernel, + "params-decoder-hc_head-hc_fn": reshape_kernel, + } - mapping = {} + for i in range(n_layers): + prefix = f"params-decoder-layers_{i}" + + # Transposed kernels + mapping[f"{prefix}-mhc_attention-mhc_norm-scale"] = mhc_dummy_norm + mapping[f"{prefix}-mhc_mlp-mhc_norm-scale"] = mhc_dummy_norm + + mapping[f"{prefix}-self_attention-o_a_proj-kernel"] = reshape_o_a_proj + + 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: + # Use composite tuple hooks for saving + 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 - # 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) + else: + # Use atomic slices for loading + 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 return mapping - PARAM_MAPPING = { "gemma2-2b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, "gemma2-9b": GEMMA2_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4239,7 +4183,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, @@ -4293,8 +4237,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/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index a6a02b37e1..94e9a2ea54 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -998,8 +998,9 @@ def extract_nnx_weights(weights_dict: dict) -> dict[str, np.ndarray]: def extract_linen_weights(weights_dict: dict) -> dict[str, np.ndarray]: """Extract weights from Linen checkpoint structure. - Linen checkpoints have structure: {'params': {'decoder': {'decoder_norm': {'scale': array}}}} - This function flattens it to: {'params-decoder-decoder_norm-scale': array} + Handles multi-collection structures like: + {'params': {'decoder': ...}, 'Tid2EidVar': {'decoder': ...}} + as well as single collection trees: {'decoder': ...} Args: weights_dict: Linen checkpoint weights dictionary @@ -1011,7 +1012,10 @@ def extract_linen_weights(weights_dict: dict) -> dict[str, np.ndarray]: leaves_with_paths = jax.tree_util.tree_leaves_with_path(weights_dict) for path_tuple, leaf_value in leaves_with_paths: path_keys = param_key_parts_from_path(path_tuple) - maxtext_param_key = "params-" + "-".join(path_keys) + if path_keys and (path_keys[0] == "params" or path_keys[0] == "Tid2EidVar" or path_keys[0].endswith("Var")): + maxtext_param_key = "-".join(path_keys) + else: + maxtext_param_key = "params-" + "-".join(path_keys) if not isinstance(leaf_value, (jax.Array, np.ndarray)): raise ValueError(f"Leaf value for {maxtext_param_key} is not an array. Type: {type(leaf_value)}.") result[maxtext_param_key] = leaf_value @@ -1051,13 +1055,13 @@ def detect_and_extract_checkpoint(checkpoint_dict: dict) -> dict[str, np.ndarray max_logging.log("Detected NNX-SFT checkpoint structure") return extract_nnx_weights(checkpoint_dict) else: - # Linen checkpoint: check if there's a nested 'params' key + # Linen checkpoint: pass multi-collection dictionary or single collection tree if isinstance(actual_weights_dict, dict) and "params" in actual_weights_dict: - actual_weights_dict = actual_weights_dict["params"] - max_logging.log("Detected Linen checkpoint structure") + max_logging.log("Detected Linen checkpoint structure (multi-collection)") + return extract_linen_weights(actual_weights_dict) else: max_logging.log("Detected Linen checkpoint structure (single params layer)") - return extract_linen_weights(actual_weights_dict) + return extract_linen_weights(actual_weights_dict) def load_hf_dict_from_transformers(model_id: str, token: str, revision: str | None = None, dtype: str = "auto"): @@ -1257,7 +1261,7 @@ def save_weights_to_checkpoint( raise RuntimeError("Failed to create Orbax checkpoint manager.") 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=jax_weights if ("params" in jax_weights.keys() if isinstance(jax_weights, dict) else False) else {"params": jax_weights}, tx=None, opt_state={} # type: ignore ) logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3))