feat(models): add qwen3.5-35b-a3b-fp8 model config and HuggingFace mapping

[snehalv2002]

Description

Adds FP8 weight-only model configuration and Hugging Face checkpoint conversion mappings for Qwen 3.5 35B (Qwen/Qwen3.5-35B-A3B-FP8).

Summary of Changes:

1. Model Configuration:
   • Added src/maxtext/configs/models/qwen3.5-35b-a3b-fp8.yml configuring weight_dtype: "float8_e4m3fn" and compute dtype: "bfloat16" with full Qwen 3.5 MoE architecture parameters (40 layers, cycle interval 4, 256 routed experts, 1 shared expert, 8 experts per token, GatedDeltaNet linear attention).
2. Hugging Face Model Registry:
   • Registered "qwen3.5-35b-a3b-fp8" and "qwen3.5-35b-fp8" pointing to "Qwen/Qwen3.5-35B-A3B-FP8" in HF_IDS (src/maxtext/utils/globals.py).
3. Parameter Mapping & Scale Hooks:
   • Updated QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING in src/maxtext/checkpoint_conversion/utils/param_mapping.py to map companion kernel_scale parameters for self-attention (Q, K, V, Out), linear attention (in_proj_qkvz, in_proj_ba, out_proj), MLP shared expert, and MoE routed experts across both scanned and unscanned modes.
   • Updated QWEN3_5_MAXTEXT_TO_HF_PARAM_HOOK_FN with reshape_scale and transposition hooks for scale tensors.
   • Registered "qwen3.5-35b-a3b-fp8" and "qwen3.5-35b-fp8" in PARAM_MAPPING and HOOK_FNS.
4. Model Config & Shape Registries:
   • Registered "qwen3.5-35b-a3b-fp8" and "qwen3.5-35b-fp8" in HF_MODEL_CONFIGS and HF_SHAPE.
5. Unit Tests:
   • Registered qwen3.5-35b-a3b.yml and qwen3.5-35b-a3b-fp8.yml in tests/unit/configs_test.py (QWEN_CONFIGS).

Tests

• Unit tests for all Qwen model configurations:

  PYTHONPATH=src pytest tests/unit/configs_test.py -k "qwen" -v

  Result: 13 passed, 0 failed.
• Full configuration suite:

  PYTHONPATH=src pytest tests/unit/configs_test.py -v

  Result: 78 passed, 0 failed.
• Pyconfig test suite:

  PYTHONPATH=src pytest tests/unit/pyconfig_test.py -v

  Result: 29 passed, 0 failed.

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 introduces support for the FP8 version of the Qwen 3.5 35B model (qwen3.5-35b-a3b-fp8 and qwen3.5-35b-fp8), adding model configurations, global variables, and parameter mapping hooks for the FP8 weight scales (kernel_scale). The review feedback highlights critical issues in the checkpoint conversion logic: first, the shape definitions for the new weight_scale parameters must be added to QWEN3_5_HF_WEIGHTS_TO_SHAPE to prevent KeyError during conversion; second, the generic reshape_scale hook is insufficient for composite scale keys (such as in_proj_qkvz-kernel_scale, in_proj_ba-kernel_scale, and routed experts scales), which require custom helper functions to correctly handle splitting and concatenating tuples of scale tensors.

[gemini-code-assist]

The QWEN3_5_HF_WEIGHTS_TO_SHAPE function in hf_shape.py needs to be updated to define the shapes of the newly mapped weight_scale parameters (such as self_attn.q_proj.weight_scale, linear_attn.in_proj_qkv.weight_scale, etc.). Without these shape definitions, any attempt to convert checkpoints for the FP8 models (e.g., MaxText -> HF) will fail with a KeyError when looking up the expected shapes of the scale tensors.

[gemini-code-assist]

To support the composite scale keys (such as in_proj_qkvz-kernel_scale, in_proj_ba-kernel_scale, and the routed experts scale tuple), we need custom helper functions to handle splitting and concatenating these scale tensors. Using reshape_scale directly on composite keys will cause AttributeError or incorrect outputs because composite keys expect tuples of tensors when saving to HF, and receive tuples of tensors when loading from HF.

  def reshape_scale(input_tensor, target_shape=None):
    if target_shape is None:
      return input_tensor
    if input_tensor.ndim == 2:
      return input_tensor.transpose().reshape(target_shape)
    return input_tensor.reshape(target_shape)

  def process_wi_0_wi_1_scale(input_tensor, target_shape=None):
    if saving_to_hf:
      wi_0, wi_1 = input_tensor
      return np.concatenate([wi_0, wi_1], axis=-1)
    else:
      return np.split(input_tensor, 2, axis=-1)

  def split_qkvz_scale(input_tensor, target_shape=None):
    if saving_to_hf:
      conv_dim = 2 * H_k * D_k + H_v * D_v
      return input_tensor[:conv_dim], input_tensor[conv_dim:]
    else:
      qkv_scale, z_scale = input_tensor
      return np.concatenate([qkv_scale, z_scale], axis=0)

  def split_ba_scale(input_tensor, target_shape=None):
    if saving_to_hf:
      return input_tensor[:H_v], input_tensor[H_v:]
    else:
      b_scale, a_scale = input_tensor
      return np.concatenate([b_scale, a_scale], axis=0)

[gemini-code-assist]

Use the newly defined split_qkvz_scale and split_ba_scale helper functions to correctly handle the composite scale keys.

Suggested change

	hooks[f"{prefix}-attention-in_proj_qkvz-kernel_scale"] = reshape_scale
	hooks[f"{prefix}-attention-in_proj_ba-kernel_scale"] = reshape_scale
	hooks[f"{prefix}-attention-in_proj_qkvz-kernel_scale"] = split_qkvz_scale
	hooks[f"{prefix}-attention-in_proj_ba-kernel_scale"] = split_ba_scale

[gemini-code-assist]

Use the newly defined process_wi_0_wi_1_scale helper function to correctly handle the composite routed experts scale key.

Suggested change

	hooks[(f"{mlp_prefix}-routed_experts-wi_0-kernel_scale", f"{mlp_prefix}-routed_experts-wi_1-kernel_scale")] = (
	reshape_scale
	)
	hooks[(f"{mlp_prefix}-routed_experts-wi_0-kernel_scale", f"{mlp_prefix}-routed_experts-wi_1-kernel_scale")] = (
	process_wi_0_wi_1_scale
	)