Skip to content

Register llama3-8b for HuggingFace conversion #4974

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ecnal-cienet wants to merge 1 commit into
base: feat/convert-post-training-tunix-onload
Choose a base branch
from
Draft

Register llama3-8b for HuggingFace conversion #4974

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions src/maxtext/checkpoint_conversion/utils/hf_model_configs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

@gemini-code-assist gemini-code-assist Bot Aug 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

To ensure that the generated config.json for Llama 3 8B is correct, we should define a dedicated llama3_8b_config that does not include the Llama 3.1 rope_scaling and has the correct max_position_embeddings of 8192. Reusing the Llama 3.1 config directly will cause downstream Hugging Face consumers to load the model with incorrect RoPE scaling and context length, leading to wrong outputs.

mistral_7b_config = transformers.MistralConfig(**mistral_7b_dict)


llama3_8b_config = transformers.LlamaConfig(
    vocab_size=128256,
    hidden_size=4096,
    intermediate_size=14336,
    num_hidden_layers=32,
    num_attention_heads=32,
    num_key_value_heads=8,
    max_position_embeddings=8192,
    head_dim=128,
    rms_norm_eps=1e-5,
    bos_token_id=128000,
    eos_token_id=128001,
    attention_bias=False,
    attention_dropout=0.0,
    hidden_act="silu",
    initializer_range=0.02,
    mlp_bias=False,
    model_type="llama",
    pretraining_tp=1,
    rope_scaling=None,
    rope_theta=500000.0,
    tie_word_embeddings=False,
    use_cache=True,
)



# from https://huggingface.co/mistralai/Mixtral-8x22B-Instruct-v0.1/blob/main/config.json
mixtral_8x22b_dict = {
"architectures": ["MixtralForCausalLM"],
Expand Down Expand Up @@ -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,
Comment on lines +1950 to +1954

@gemini-code-assist gemini-code-assist Bot Aug 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Use the newly defined llama3_8b_config for "llama3-8b" to ensure the correct context length and RoPE scaling are written to config.json.

Suggested change
# 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,
# Llama 3 and 3.1 8B are the same architecture, and their MaxText configs are identical, so
# they share the same weight mapping, shapes, and hooks. However, Llama 3 8B has a different
# context length (8192) and no rope_scaling, so it uses its own HF config to ensure the
# written config.json is correct.
"llama3-8b": llama3_8b_config,

"llama3.1-8b": llama31_8b_config,
"llama3.1-8b-Instruct": llama31_8b_config,
"llama3.1-70b": llama31_70b_config,
Expand Down
2 changes: 2 additions & 0 deletions src/maxtext/checkpoint_conversion/utils/hf_shape.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
3 changes: 3 additions & 0 deletions src/maxtext/checkpoint_conversion/utils/param_mapping.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
2 changes: 2 additions & 0 deletions src/maxtext/utils/globals.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading