test(fp8): add PyTorch reference parity suite and logit checker metrics by snehalv2002 · Pull Request #5055 · AI-Hypercomputer/maxtext

snehalv2002 · 2026-08-28T22:29:18Z

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

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.
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 · 2026-08-28T22:32:23Z

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


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 · 2026-08-28T22:32:23Z

 import jax
+import jax.numpy as jnp


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 · 2026-08-28T22:32:23Z

+  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"]


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 · 2026-08-28T22:32:23Z

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


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 · 2026-08-28T22:32:23Z

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


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 · 2026-08-28T22:32:24Z

 import flax.linen as nn

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


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 · 2026-08-28T22:32:24Z

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


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 · 2026-08-28T22:32:24Z

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


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 · 2026-08-28T22:32:24Z

+    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


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 · 2026-08-28T22:32:24Z

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


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)

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

b4c07de

gemini-code-assist Bot reviewed Aug 28, 2026

View reviewed changes

snehalv2002 changed the title ~~Pr/fp8 reference parity~~ test(fp8): add PyTorch reference parity suite and logit checker metrics Aug 28, 2026

snehalv2002 changed the base branch from main to pr/fp8-llama-onboarding August 28, 2026 22:54

Navigation Menu

Uh oh!

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

test(fp8): add PyTorch reference parity suite and logit checker metrics#5055
snehalv2002 wants to merge 1 commit into
pr/fp8-llama-onboardingfrom
pr/fp8-reference-parity

snehalv2002 commented Aug 28, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

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

	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

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

Uh oh!

Conversation

snehalv2002 commented Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation & Context

Key Changes

Tests

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

snehalv2002 commented Aug 28, 2026 •

edited

Loading