[albertvillanova]

This PR writes the FP8 checkpoint of the tiny Mistral Large 4 model in the layout of mistralai/Mistral-Large-4-1T-A52B, instead of quantizing it at load.

Part of #7137.

Motivation

When the model is quantized at load and saved by transformers, the file stores fused expert scales named weight_scale_inv next to per-expert weights. The script also cast the router and the RoPE alpha to F32. The reference stores per-expert qscale_weight scales and keeps both tensors in BF16.

Changes

• Quantize the linears the reference stores in FP8 in 32x32 blocks (scaled down from 128x128). Use power-of-two scales, like the reference's UE8M0 scales, stored per expert as qscale_weight (F32).
• Keep gate.weight and lerope.alpha in BF16, like the reference.
• Mirror the reference quantization_config, with only weight_block_size scaled down.
• Reload the written checkpoint before the smoke test, and push it as written.

Dry run at transformers 5.22.0: the checkpoint matches the reference names and dtypes, it reloads with the per-expert scales merged, and the smoke test passes. The dtype check in the script compares against the consolidated reference once #8238 is merged.

---

Note

Low Risk
Changes are confined to the tiny-model generation script and checkpoint export path, not runtime inference or core library logic.

Overview
The tiny Mistral Large 4 generator no longer relies on FineGrainedFP8Config and load-time quantization. It saves a BF16 checkpoint, rewrites selected linear weights in model.safetensors to float8_e4m3fn in 32×32 blocks with per-block power-of-two qscale_weight tensors (matching the reference consolidated layout), restores quantization_config with shrunk weight_block_size, reloads from disk for the smoke test, and uploads via push_dir_to_hub instead of hub-push on the in-memory quantized model.

Router gate.weight and RoPE alpha are no longer upcast to FP32 during export, aligning dtypes with the reference checkpoint naming and scale conventions.

^{Reviewed by Cursor Bugbot for commit faff858. Bugbot is set up for automated code reviews on this repo. Configure here.}

[albertvillanova]

The CI model has been regenerated and the corresponding Hub PR is open. It needs to be merged once this PR is approved:

• https://huggingface.co/trl-internal-testing/tiny-MistralLarge4ForConditionalGeneration/discussions/2

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[mangeshraut712]

The in-script FP8 conversion is a reasonable way to avoid needing an FP8 GPU at load time, but a couple of edges will corrupt or change the checkpoint:

1. All-zero blocks become NaN. For a block whose amax is 0,

   torch.log2(blocks.abs().amax(dim=(1, 3)) / torch.finfo(torch.float8_e4m3fn).max)

   is -inf, exp2 is 0, and blocks / scale is 0/0 → NaN before the float8_e4m3fn cast. Tiny or structured-init weights can easily have a zero tile. I’d clamp the scale to a minimum of 1 (or skip / write a zero scale) when amax == 0.

2. The fp32 restore for the router and RoPE scale was dropped. The previous code forced layer.mlp.gate.weight and rotary_emb.alpha to float32 because “the reference keeps [them] in fp32”. Those names are not in fp8_weights, so they now stay bfloat16 from model.to(dtype=torch.bfloat16). If check_dtype_pattern does not cover them, the published tiny checkpoint will silently disagree with the reference dtypes.

3. view(rows // 32, 32, cols // 32, 32) throws if a matched weight is not divisible by 32. hidden_size / intermediate_size / moe_intermediate_size are, but if a wq / indexer.wq_b ever comes out 32×16 (e.g. from kv_lora_rank or index_head_dim), the script crashes instead of skipping or padding. A divisibility check with a clear error would help.

[albertvillanova]

Tested with MODEL_REVISIONS pointing at the regenerated model, trl-internal-testing/tiny-MistralLarge4ForConditionalGeneration at refs/pr/2, on top of current main. The suite is green: https://github.com/huggingface/trl/actions/runs/42846518953

[qgallouedec]

Thanks, confirmed on my side