feat(checkpointing): add dequantize-on-load parameter restoration

[snehalv2002]

Description

Upgrades src/maxtext/common/checkpointing.py to support restoring FP8 Orbax checkpoints and introduces automatic in-memory dequantize-on-load.

Motivation & Context

Full pre-training and fine-tuning with optimizers like AdamW require unquantized BF16 master weights, as optimizer states and gradient accumulation are not supported on raw 8-bit floats. This feature enables loading a quantized FP8 checkpoint directly into an unquantized model (weight_dtype="bfloat16"), saving ~50% disk storage and transfer bandwidth while automatically reconstructing bfloat16 weights in memory upon restore without requiring an offline conversion step.

How to Use Dequantize-on-Load

To use dequantize-on-load, specify the unquantized model name (or set weight_dtype=bfloat16), but point load_parameters_path directly at an FP8 quantized Orbax checkpoint:

# Example: Restore an FP8 checkpoint into an unquantized BF16 model for decode / eval:
python3 -m maxtext.inference.decode \
  src/maxtext/configs/base.yml \
  model_name=qwen3.5-35b-a3b \
  load_parameters_path=/path/to/fp8_checkpoint/0/items \
  per_device_batch_size=1

# Example: Restore an FP8 checkpoint into an unquantized BF16 model for fine-tuning / training:
python3 -m maxtext.trainers.pre_train.train \
  src/maxtext/configs/base.yml \
  model_name=qwen3.5-35b-a3b \
  load_parameters_path=/path/to/fp8_checkpoint/0/items

What happens:

1. MaxText instantiates the model architecture in memory using bfloat16 (the target specification want expects BF16 weights and no scale parameters).
2. load_params_from_path inspects the checkpoint metadata, discovers the companion scale tensors (kernel_scale, wi_0_scale, etc.), and temporarily augments the restore schema to retrieve both the FP8 weights and scales from disk.
3. Restored FP8 weights are dynamically dequantized in memory to bfloat16 using quantizations.dequantize_weight, and companion scale tensors are dropped.
4. The model receives pure bfloat16 weights and runs without any runtime quantization overhead.

Key Changes

1. Target State Augmentation (_augment_want_with_scales):
   • Discovers companion scale tensors ({weight}_scale) present in checkpoint metadata but omitted from the unquantized target model (want).
   • Augments the target abstract schema so Orbax restores both the weights and their scale parameters without schema mismatches (supporting NNX and Linen layouts).
2. In-Memory Dequantization (maybe_dequantize_restored_params):
   • Single-pass traversal over restored parameters: when a weight and its companion scale are present but the target model expects an unquantized weight, dequantizes the weight to the target dtype and strips the companion scale.
   • Preserves unquantized companion parameters (e.g. bias) intact.
3. Pre-load Architectural Validation:
   • Compares checkpoint metadata against the target model before loading, raising a descriptive ValueError on shape mismatches rather than failing mid-restore.
4. Direct FP8 Loading Preserved:
   • If the target model expects weight_dtype="float8_e4m3fn" (e.g. model_name=qwen3.5-35b-a3b-fp8), weights and scales are loaded directly without dequantization.

Part 3 in the FP8 Weight-Only Dynamic Dequantization series (depends on #5053, #5052).

Tests

• Unit tests in tests/unit/checkpointing_test.py (FP8DequantizeOnLoadTest):
  • test_load_fp8_checkpoint_into_bf16_nnx_model: Verifies NNX dequantize-on-load to BF16, scale removal, and unquantized bias preservation.
  • test_load_fp8_checkpoint_into_fp8_nnx_model: Verifies direct FP8 restoration preserves weights and scale tensors.
  • test_load_fp8_checkpoint_into_bf16_linen_dict: Verifies legacy Linen dictionary layout dequantization.
  • test_load_fp8_checkpoint_shape_mismatch_raises: Verifies pre-load shape mismatch detection.

PYTHONPATH=src JAX_PLATFORMS=cpu pytest tests/unit/checkpointing_test.py -k "FP8DequantizeOnLoadTest" -v

Result: 4 passed, 0 failures (full test suite: 24 passed).

• Linters passed cleanly: codespell, pylint (10.00/10), pyink.

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 FP8 weight-only storage and dynamic dequantization during checkpoint loading. It adds fallback logic for weight scales, implements dynamic dequantization of restored parameters when the target model expects unquantized weights, and updates dense linear layers to support block-wise quantization scales. The review feedback identifies several critical issues, including a bug in checkpoint loading for NNX models where dequantized weights are not written back, a performance bottleneck from repeated O(N) lookups on safetensors keys, incorrect unconditional conversion of float16 to bfloat16, potential AttributeErrors on older JAX versions, and numerical instability risks from initializing bias as FP8.

[gemini-code-assist]

In load_params_from_path, when is_nnx is True, restored_weights is retrieved from restored_collection["params"]. After dequantizing, restored_weights is reassigned to the new dequantized dictionary, but restored_collection["params"] is never updated. This means the returned restored_collection still contains the old, non-dequantized weights, completely breaking the feature for NNX models. We should update restored_collection["params"] with the new restored_weights when is_nnx is True.

Suggested change

	restored_weights = maybe_dequantize_restored_params(restored_weights, want)
	if not is_nnx:
	restored_collection = restored_weights
	restored_weights = maybe_dequantize_restored_params(restored_weights, want)
	if is_nnx and restore_key not in ("model_params", "model"):
	restored_collection["params"] = restored_weights
	else:
	restored_collection = restored_weights

[gemini-code-assist]

Calling f.keys() repeatedly in a loop or checking final_key not in f.keys() can be extremely slow because f.keys() returns a list in safetensors, making each lookup an O(N) operation. Converting f.keys() to a set once before the lookups will make them O(1) and significantly improve performance, especially for models with many keys.

Suggested change

	if final_key not in f.keys():
	if final_key.endswith(".weight_scale"):
	for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
	alt_key = final_key[:-len(".weight_scale")] + suffix
	if alt_key in f.keys():
	final_key = alt_key
	break
	elif final_key.endswith(".scale"):
	for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
	alt_key = final_key[:-len(".scale")] + suffix
	if alt_key in f.keys():
	final_key = alt_key
	break
	f_keys = set(f.keys())
	if final_key not in f_keys:
	if final_key.endswith(".weight_scale"):
	for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
	alt_key = final_key[:-len(".weight_scale")] + suffix
	if alt_key in f_keys:
	final_key = alt_key
	break
	elif final_key.endswith(".scale"):
	for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
	alt_key = final_key[:-len(".scale")] + suffix
	if alt_key in f_keys:
	final_key = alt_key
	break

[gemini-code-assist]

In the generic get_tensor reader, unconditionally converting float16 tensors to bfloat16 is incorrect and can cause unexpected behavior or precision loss if the caller actually expects float16. Since PyTorch CPU tensors support .numpy() directly for float16, we should return the tensor in its original dtype.

Suggested change

	elif t.dtype == torch.float16:
	return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)
	elif t.dtype == torch.float16:
	return t.numpy()

[gemini-code-assist]

Directly referencing jnp.float8_e4m3fn and jnp.float8_e5m2 will raise an AttributeError on older JAX versions that do not support these FP8 dtypes. We should safely check for their existence using hasattr or getattr before adding them to fp8_types.

Suggested change

	fp8_types = [jnp.float8_e4m3fn, jnp.float8_e5m2]
	for attr in ("float8_e4m3fnuz", "float8_e5m2fnuz", "float8_e4m3b11fnuz"):
	if hasattr(jnp, attr):
	fp8_types.append(getattr(jnp, attr))
	fp8_types = []
	for attr in ("float8_e4m3fn", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "float8_e4m3b11fnuz"):
	if hasattr(jnp, attr):
	fp8_types.append(getattr(jnp, attr))

[gemini-code-assist]

Import is_fp8_dtype from common_types to avoid duplicating the FP8 dtype checking logic in linears.py.

Suggested change

	from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config, Shape
	from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config, Shape, is_fp8_dtype

[gemini-code-assist]

To avoid code duplication and improve maintainability, we should reuse the is_fp8_dtype function imported from common_types instead of duplicating the complex checking logic here.

def _is_fp8_dtype(dtype: Any) -> bool:
  """Checks whether a dtype is an FP8 data type."""
  return is_fp8_dtype(dtype)

[gemini-code-assist]

If self.weight_dtype is FP8, the bias is initialized as FP8. Quantizing bias to FP8 is highly discouraged as it can cause severe numerical instability and loss of model quality. Bias should remain in the compute dtype (self.dtype or bfloat16).

      bias_dtype = self.dtype if _is_fp8_dtype(self.weight_dtype) else self.weight_dtype
      try:
        bias_val = default_bias_init(rngs.params(), bias_shape, bias_dtype)
      except (TypeError, ValueError):
        bias_val = default_bias_init(rngs.params(), bias_shape, self.dtype).astype(bias_dtype)

[codecov]

Codecov Report

❌ Patch coverage is 84.31373% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/checkpointing.py	84.31%	4 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!

[Shuwen-Fang]

Would this work for per tensor, per-channel, & block-wise quantization?

[snehalv2002]

in terms of naming and functionality yes, the maxtext scales are always defined as f"{k}_scale" (we introduced this naming scheme in the first PR in the stack) and the dequantization method works for all three types of quantization you mentioned above using the quantizations.dequantize_weight() module which already existed in maxtext.