test(fp8): add PyTorch reference parity suite and logit checker metrics

[snehalv2002]

Description

Adds an end-to-end PyTorch reference equivalence test suite and enhances tests/utils/forward_pass_logit_checker.py with numerical error and KL divergence metrics for FP8 validation.

Motivation & Context

To guarantee numerical correctness and bit-level compatibility between MaxText's dynamic dequantization and reference PyTorch implementations, this PR adds granular module-level unit tests and side-by-side logits comparison tools.

Key Changes

1. PyTorch Reference Unit Suite (tests/unit/llama_fp8_vs_reference_test.py):
   • Implemented 5 test suites comparing PyTorch FP8 modules against MaxText modules with identical weights and inputs:
     • DenseGeneral scalar, per-channel, and block-wise scaling parity.
     • MlpBlock (SwiGLU) FP8 parity.
     • Attention (RoPE + GQA) FP8 parity.
     • LlamaDecoderLayer unscanned and scanned parity.
     • Full NNXDecoder pipeline parity.
2. Enhanced Logit Checker (tests/utils/forward_pass_logit_checker.py):
   • Added maximum absolute difference (atol), relative difference (rtol), top-k rank agreement, and KL divergence ($D_{KL}$) evaluation when running with --run_hf_model=True.

Part 5 of 5 in the FP8 Weight-Only Dynamic Dequantization series (depends on #5053, #5052, #5051, #5054).

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Ran the complete reference parity test suite:

pytest tests/unit/llama_fp8_vs_reference_test.py -v

Result: 5 passed, 0 failures.

Ran side-by-side forward pass logit verification against Hugging Face reference model (neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8):

python -m tests.utils.forward_pass_logit_checker \
  src/maxtext/configs/models/llama3.1-8b-fp8.yml \
  model_name=llama3.1-8b-fp8 \
  load_parameters_path=/tmp/maxtext_llama3.1_8b_fp8/0/items \
  --run_hf_model=True \
  --max_kl_div=0.05

Result:

• Average KL Divergence ($D_{KL}$): 4.7988e-03 (< 0.05 limit).
• Top-10 Rank Agreement: 96.7% (29/30 tokens across prompts).
• Top-1 Agreement: 100%.

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 with dynamic on-the-fly dequantization, specifically adding configurations and mappings for the llama3.1-8b-fp8 model. Key changes include adding fallback mechanisms for scale keys during checkpoint conversion, implementing dynamic dequantization of restored parameters when loading checkpoints, and updating linear layers to support kernel_scale parameters and FP8 weight types. The review feedback focuses on optimizing checkpoint loading by caching keys as a set, replacing dict type checks with collections.abc.Mapping to support FrozenDict, eliminating code duplication by using the centralized is_fp8_dtype helper, and ensuring that bias is initialized in the activation compute dtype when weight_dtype is FP8 to prevent precision loss.

[gemini-code-assist]

When weight_dtype is an FP8 type, bias should be initialized and stored in the compute/activation dtype (e.g., bfloat16) rather than FP8, to avoid precision loss and compatibility issues.

Suggested change

	if self.use_bias:
	bias_axes = self.kernel_axes[-len(self.out_features_shape) :]
	bias_shape = kernel_shape[-len(self.out_features_shape) :]
	try:
	bias_val = default_bias_init(rngs.params(), bias_shape, self.weight_dtype)
	except (TypeError, ValueError):
	bias_val = default_bias_init(rngs.params(), bias_shape, self.dtype).astype(self.weight_dtype)
	self.bias = nnx.Param(
	default_bias_init(rngs.params(), bias_shape, self.weight_dtype),
	bias_val,
	sharding=bias_axes,
	)
	if self.use_bias:
	bias_axes = self.kernel_axes[-len(self.out_features_shape) :]
	bias_shape = kernel_shape[-len(self.out_features_shape) :]
	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)
	self.bias = nnx.Param(
	bias_val,
	sharding=bias_axes,
	)

[gemini-code-assist]

Import Mapping from collections.abc to robustly check for dictionary-like objects (such as Flax's FrozenDict), which do not inherit from the standard dict class.

Suggested change

	import jax
	import jax.numpy as jnp
	import jax
	from collections.abc import Mapping
	import jax.numpy as jnp

[gemini-code-assist]

Use Mapping instead of dict to support Flax's FrozenDict and other dictionary-like containers.

Suggested change

	if not isinstance(want_bare, dict) or not isinstance(meta_tree, dict):
	return meta_tree
	want_keys = set(want_bare.keys())
	if want_keys and want_keys.issubset(meta_tree.keys()):
	return meta_tree
	for wrapper in ("params", "model_params", "model", "items"):
	if wrapper in meta_tree and isinstance(meta_tree[wrapper], dict):
	sub = meta_tree[wrapper]
	if want_keys and want_keys.issubset(sub.keys()):
	return sub
	if wrapper == "params" and "params" in sub and isinstance(sub["params"], dict):
	if want_keys and want_keys.issubset(sub["params"].keys()):
	return sub["params"]
	if not isinstance(want_bare, Mapping) or not isinstance(meta_tree, Mapping):
	return meta_tree
	want_keys = set(want_bare.keys())
	if want_keys and want_keys.issubset(meta_tree.keys()):
	return meta_tree
	for wrapper in ("params", "model_params", "model", "items"):
	if wrapper in meta_tree and isinstance(meta_tree[wrapper], Mapping):
	sub = meta_tree[wrapper]
	if want_keys and want_keys.issubset(sub.keys()):
	return sub
	if wrapper == "params" and "params" in sub and isinstance(sub["params"], Mapping):
	if want_keys and want_keys.issubset(sub["params"].keys()):
	return sub["params"]

[gemini-code-assist]

Use Mapping instead of dict to support Flax's FrozenDict and other dictionary-like containers.

Suggested change

	if not isinstance(want_node, dict) or not isinstance(meta_node, dict):
	return want_node
	if not isinstance(want_node, Mapping) or not isinstance(meta_node, Mapping):
	return want_node

[gemini-code-assist]

Use Mapping instead of dict to support Flax's FrozenDict and other dictionary-like containers.

Suggested change

	augmented[k] = _augment_target_with_scales(
	v,
	meta_node.get(k) if isinstance(meta_node, dict) else None,
	)
	augmented[k] = _augment_target_with_scales(
	v,
	meta_node.get(k) if isinstance(meta_node, Mapping) else None,
	)

[gemini-code-assist]

Instead of duplicating the FP8 check logic locally, import is_fp8_dtype from maxtext.common.common_types to maintain consistency and reduce code duplication.

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]

Use the imported is_fp8_dtype instead of the removed local helper.

Suggested change

	if has_scale is None:
	should_have_scale = (
	_is_fp8_dtype(self.weight_dtype)
	or (kernel_scale_init is not None)
	or (scale_shape is not None)
	)
	if has_scale is None:
	should_have_scale = (
	is_fp8_dtype(self.weight_dtype)
	or (kernel_scale_init is not None)
	or (scale_shape is not None)
	)

[gemini-code-assist]

Use the imported is_fp8_dtype instead of the removed local helper.

Suggested change

	if not _is_fp8_dtype(kernel.dtype) and kernel_scale is None:
	kernel = jnp.asarray(kernel, self.dtype)
	if not is_fp8_dtype(kernel.dtype) and kernel_scale is None:
	kernel = jnp.asarray(kernel, self.dtype)

[gemini-code-assist]

Use the imported is_fp8_dtype instead of the removed local helper.

Suggested change

	if has_scale is None:
	should_have_scale = (
	_is_fp8_dtype(self.weight_dtype)
	or (kernel_scale_init is not None)
	or (scale_shape is not None)
	)
	else:
	should_have_scale = has_scale
	if has_scale is None:
	should_have_scale = (
	is_fp8_dtype(self.weight_dtype)
	or (kernel_scale_init is not None)
	or (scale_shape is not None)
	)
	else:
	should_have_scale = has_scale

[gemini-code-assist]

Use the imported is_fp8_dtype instead of the removed local helper.

Suggested change

	if _is_fp8_dtype(kernel.dtype) or kernel_scale is not None:
	kernel = dequantize_weight(kernel, kernel_scale, compute_dtype=self.dtype)
	else:
	kernel = jnp.asarray(kernel, self.dtype)
	if is_fp8_dtype(kernel.dtype) or kernel_scale is not None:
	kernel = dequantize_weight(kernel, kernel_scale, compute_dtype=self.dtype)
	else:
	kernel = jnp.asarray(kernel, self.dtype)