From 3ef2d41f0cf20c403379f9d909bfb441424b8054 Mon Sep 17 00:00:00 2001 From: Lance Wang Date: Sun, 23 Aug 2026 22:51:29 +0000 Subject: [PATCH] Register llama3-8b for HuggingFace conversion llama3-8b appeared in HF_IDS but in none of the tables that actually convert weights, so it looked supported and failed on a missing key. Its MaxText config is byte-identical to llama3.1-8b's, so it reuses that family's mapping, shapes and hooks; only the source repo differs. The shared HF config carries 3.1's rope_scaling and 131072 context, which reach a consumer reading the written config.json rather than the weight mapping the conversion uses. mistral-7b is left unregistered, and the note beside its config records why. Wiring it to the Llama family converts cleanly -- 291 of 291 arrays, every name and shape matching the repo -- and produces numerically wrong weights: forward_pass_logit_checker measured KL 1.32e-01 mean and 3.03e-01 max against the 3e-3 the repo's own conversion tests use. Setting rope_max_timescale to v0.1's 10000 rather than the config's 1e6 made it worse, so the rope timescale is not what differs. Until mistral has a mapping of its own, the KeyError from converting it is the honest outcome: a structurally perfect checkpoint full of wrong numbers is the worse failure. The config sits here so the next attempt starts from the measurement instead of repeating it. --- .../utils/hf_model_configs.py | 39 +++++++++++++++++++ .../checkpoint_conversion/utils/hf_shape.py | 2 + .../utils/param_mapping.py | 3 ++ src/maxtext/utils/globals.py | 2 + 4 files changed, 46 insertions(+) diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 89abd56d4c..076c853b46 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -1640,6 +1640,40 @@ def __init__(self, **kwargs): mixtral_8x7b_config = transformers.MixtralConfig(**mixtral_8x7b_dict) +# from https://huggingface.co/mistralai/Mistral-7B-v0.1/blob/main/config.json +# +# Not registered in HF_MODEL_CONFIGS. Wiring mistral-7b to the Llama family's mapping, hooks and +# shapes converts cleanly -- 291/291 arrays, every name and shape matching the repo -- and produces +# numerically wrong weights: forward_pass_logit_checker measured KL 1.32e-01 mean / 3.03e-01 max +# against the 3e-3 the repo's own conversion tests use. Setting rope_max_timescale to v0.1's 10000 +# rather than the config's 1e6 made it worse (1.58e-01 / 3.65e-01), so the rope timescale is not +# what differs. Mistral needs its own mapping, and until it has one, the KeyError from converting +# it is the honest outcome: a structurally perfect checkpoint full of wrong numbers is worse. +mistral_7b_dict = { + "architectures": ["MistralForCausalLM"], + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 14336, + "max_position_embeddings": 32768, + "model_type": "mistral", + "num_attention_heads": 32, + "num_hidden_layers": 32, + "num_key_value_heads": 8, + "rms_norm_eps": 1e-05, + "rope_theta": 10000.0, + "sliding_window": 4096, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "transformers_version": "4.34.0.dev0", + "use_cache": True, + "vocab_size": 32000, +} +mistral_7b_config = transformers.MistralConfig(**mistral_7b_dict) + + # from https://huggingface.co/mistralai/Mixtral-8x22B-Instruct-v0.1/blob/main/config.json mixtral_8x22b_dict = { "architectures": ["MixtralForCausalLM"], @@ -1913,6 +1947,11 @@ def __init__(self, **kwargs): "qwen3-vl-2b": qwen3_vl_2b_config, "qwen3-vl-4b": qwen3_vl_4b_config, "qwen3-vl-30b-a3b": qwen3_vl_30b_a3b_config, + # Llama 3 and 3.1 8B are the same architecture, and their MaxText configs are identical, so + # the same HF config describes both for conversion purposes. It does carry 3.1's + # rope_scaling and 131072 context, which only matter to a consumer reading the written + # config.json rather than to the weight mapping itself. + "llama3-8b": llama31_8b_config, "llama3.1-8b": llama31_8b_config, "llama3.1-8b-Instruct": llama31_8b_config, "llama3.1-70b": llama31_70b_config, diff --git a/src/maxtext/checkpoint_conversion/utils/hf_shape.py b/src/maxtext/checkpoint_conversion/utils/hf_shape.py index 85dd1d6ea0..ff1c7a3d92 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_shape.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_shape.py @@ -1300,6 +1300,8 @@ def DEEPSEEKV4_HF_WEIGHTS_TO_SHAPE(config): "qwen3-vl-2b": QWEN3_VL_HF_WEIGHTS_TO_SHAPE, "qwen3-vl-4b": QWEN3_VL_HF_WEIGHTS_TO_SHAPE, "qwen3-vl-30b-a3b": QWEN3_VL_HF_WEIGHTS_TO_SHAPE, + # Same weights and shapes as 3.1; see the note beside llama3-8b in hf_model_configs.py. + "llama3-8b": LLAMA31_HF_WEIGHTS_TO_SHAPE, "llama3.1-8b": LLAMA31_HF_WEIGHTS_TO_SHAPE, "llama3.1-8b-Instruct": LLAMA31_HF_WEIGHTS_TO_SHAPE, "llama3.1-70b": LLAMA31_HF_WEIGHTS_TO_SHAPE, diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 94e96173f1..e3ccac2cc1 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -4245,6 +4245,8 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_MAPPING, + # Same architecture as 3.1, so the same mapping applies. + "llama3-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING, @@ -4299,6 +4301,7 @@ def mhc_concat_scale(input_tensors, target_shape=None): "qwen3-vl-2b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-4b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, "qwen3-vl-30b-a3b": QWEN3_VL_MAXTEXT_TO_HF_PARAM_HOOK_FN, + "llama3-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-8b-Instruct": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, "llama3.1-70b": LLAMA31_MAXTEXT_TO_HF_PARAM_HOOK_FN, diff --git a/src/maxtext/utils/globals.py b/src/maxtext/utils/globals.py index 30f6e65124..70cf6330f1 100644 --- a/src/maxtext/utils/globals.py +++ b/src/maxtext/utils/globals.py @@ -66,6 +66,8 @@ "qwen3-vl-2b": "Qwen/Qwen3-VL-2B-Instruct", "qwen3-vl-4b": "Qwen/Qwen3-VL-4B-Instruct", "qwen3-vl-30b-a3b": "Qwen/Qwen3-VL-30B-A3B-Instruct", + # Llama 3 and 3.1 8B have byte-identical MaxText configs; only the source repo differs. + "llama3-8b": "meta-llama/Meta-Llama-3-8B", "llama3.1-8b": "meta-llama/Llama-3.1-8B", "llama3.1-8b-Instruct": "meta-llama/Llama-3.1-8B-Instruct", "llama3.1-70b-Instruct": "meta-llama/Llama-3.1-70B-Instruct",