Register llama3-8b for HuggingFace conversion

[ecnal-cienet]

Description

llama3-8b appeared in HF_IDS but in none of the tables that actually convert weights — HF_MODEL_CONFIGS, PARAM_MAPPING, HF_SHAPE, HOOK_FNS — 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 config.json this writes, not the weight mapping the conversion itself uses.

This also adds mistral-7b's config without registering it, and records why. Wiring mistral 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 (1.58e-01 / 3.65e-01), so the rope timescale is not what differs.

Mistral needs a mapping of its own. Until it has one, the KeyError from converting it is the honest outcome: a structurally perfect checkpoint full of wrong numbers is the worse failure, because nothing downstream notices. The config sits beside the note so the next attempt starts from the measurement rather than repeating it.

Tests

The four tables load and agree:

llama3-8b    HF_IDS=True  HF_MODEL_CONFIGS=True  PARAM_MAPPING=True  HF_SHAPE=True
mistral-7b   HF_IDS=True  HF_MODEL_CONFIGS=False PARAM_MAPPING=False HF_SHAPE=False

e2e log

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request adds support for the Llama 3 8B model by mapping it to the existing Llama 3.1 8B configurations, shapes, and parameter mappings, and introduces a configuration definition for Mistral 7B. The reviewer points out that reusing the Llama 3.1 8B configuration directly for Llama 3 8B is problematic because it outputs incorrect RoPE scaling and context length (8192 vs 131072) in the generated config.json. They suggest defining a dedicated llama3_8b_config with the correct parameters to prevent downstream issues.

[gemini-code-assist]

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,
)

[gemini-code-assist]

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,

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!