feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext by snehalv2002 · Pull Request #5052 · AI-Hypercomputer/maxtext

snehalv2002 · 2026-08-28T22:28:43Z

Description

Enhances MaxText's standalone checkpoint conversion tool (src/maxtext/checkpoint_conversion/to_maxtext.py) to directly ingest, serialize, and store native 8-bit float weights (float8_e4m3fn, float8_e5m2) and companion scale tensors into Orbax format.

Motivation & Context

Previously, to_maxtext.py only supported save_dtype="bfloat16" and save_dtype="float32". Converting pre-quantized multi-terabyte FP8 Hugging Face checkpoints forced CPU upcasting to BF16, blowing up disk and RAM requirements by 2x.

Key Changes

Zero-Copy PyTorch $\to$ JAX FP8 Bridging:
- Reinterprets PyTorch torch.float8_e4m3fn / torch.float8_e5m2 tensors as uint8 views before bridging to ml_dtypes.float8_e4m3fn in NumPy memory, avoiding unsupported standard NumPy float8 casts.
Scale Tensor Alias & Multi-Dimensional Scale Resolution:
- Added robust resolution in LazyHFLoader for common Hugging Face scale tensor naming conventions (.weight_scale, .scale, .weight_scale_inv, .scale_inv).
- Supported scalar, 1D per-channel, and 2D block-scaled grids (e.g. 128x128) during tensor loading and conversion.
Selective Precision Preservation & Declarative Filtering:
- Unquantized layers (embeddings, norm scales, router gates) retain their high-precision dtypes even when --save_dtype="float8_e4m3fn" is specified, adhering to declarative unquantized_modules.
Scan Axis Bounds Safety (tensor_handling.py, utils.py):
- Added dimension guards (if tensor.ndim > param_scan_axis) when stacking parameters across scanned decoder layers to safely handle lower-rank scale tensors.
CLI Support:
- Expanded --save_dtype arguments to accept float8_e4m3fn and float8_e5m2.

Part 2 of 5 in the FP8 Weight-Only Dynamic Dequantization series (depends on #5053).

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

Tested end-to-end checkpoint conversion of Hugging Face FP8 models (neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8):

python -m maxtext.checkpoint_conversion.to_maxtext \
  model_name=llama3.1-8b-fp8 \
  hf_model_path=neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8 \
  base_output_directory=/tmp/maxtext_llama3.1_8b_fp8 \
  save_dtype=float8_e4m3fn \
  --lazy_load_tensors=True

Verified that linear weights are stored in float8_e4m3fn and scale tensors in float32.

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 in MaxText, including changes to checkpoint conversion, linear layers, and decoder layers to support FP8 weights alongside scalar, per-channel, and block-wise scales. Feedback on these changes includes addressing a potential lossy double conversion of float16/bfloat16 tensors during checkpoint loading, extracting duplicated scale key resolution logic in to_maxtext.py into a helper function, and removing the redundant local _is_fp8_dtype helper in linears.py in favor of the existing utility in common_types.py.

gemini-code-assist · 2026-08-28T22:31:38Z

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


Converting torch.float16 and torch.bfloat16 unconditionally to ml_dtypes.bfloat16 inside get_tensor can lead to a lossy double conversion when save_dtype is configured as float32 (i.e., float16 -> bfloat16 -> float32).

Since LazyTensor.__array__ already handles casting the retrieved numpy array to the target save_dtype (via arr.astype(dtype)), get_tensor should simply return the closest native numpy representation (e.g., float32 for bfloat16 to avoid PyTorch conversion issues, and native float16 for float16).

Suggested change

elif t.dtype == torch.bfloat16:

return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)

elif t.dtype == torch.float16:

return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)

elif t.dtype == torch.bfloat16:

return t.to(torch.float32).numpy()

elif t.dtype == torch.float16:

return t.numpy()

gemini-code-assist · 2026-08-28T22:31:38Z

+    if shard_name is None:
+      # Check fallback for .weight_scale vs .scale and inverse scales
+      if resolved_key.endswith(".weight_scale"):
+        for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]:
+          alt_key = resolved_key[:-len(".weight_scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break
+      elif resolved_key.endswith(".scale"):
+        for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]:
+          alt_key = resolved_key[:-len(".scale")] + suffix
+          if alt_key in self.shard_map:
+            resolved_key = alt_key
+            shard_name = self.shard_map[resolved_key]
+            break


The fallback logic for resolving .weight_scale vs .scale and inverse scales is duplicated three times in this file (here, in get_tensor's second block, and in _eager_getter).

To improve maintainability and reduce redundancy, consider extracting this logic into a helper function at the module level:

def resolve_scale_key(key: str, container) -> str: """Resolves fallback keys for weight scales and inverse scales.""" if key in container: return key if key.endswith(".weight_scale"): for suffix in [".scale", ".weight_scale_inv", ".scale_inv"]: alt_key = key[:-len(".weight_scale")] + suffix if alt_key in container: return alt_key elif key.endswith(".scale"): for suffix in [".weight_scale", ".scale_inv", ".weight_scale_inv"]: alt_key = key[:-len(".scale")] + suffix if alt_key in container: return alt_key return key

You can then simplify this block to:

resolved_key = resolve_scale_key(key, self.shard_map) shard_name = self.shard_map.get(resolved_key)

gemini-code-assist · 2026-08-28T22:31:38Z

+def _is_fp8_dtype(dtype: Any) -> bool:
+  """Checks whether a dtype is an FP8 data type."""
+  if dtype is None:
+    return False
+  try:
+    canon_dtype = _canonicalize_dtype(dtype)
+  except (TypeError, ValueError):
+    return False
+
+  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))
+  return canon_dtype in fp8_types


The _is_fp8_dtype helper function is duplicated here. MaxText already has a standard is_fp8_dtype utility defined in maxtext.common.common_types.

You can import is_fp8_dtype from maxtext.common.common_types and use it directly instead of redefining _is_fp8_dtype locally.

snehalv2002 · 2026-09-09T20:14:06Z

+
+if "HF_HOME" not in os.environ and os.path.exists("/dev/shm"):
+  os.environ["HF_HOME"] = "/dev/shm/hf_cache"
+


remove this it was only for local development

snehalv2002 · 2026-09-09T20:45:34Z

  checkpointing.wait_until_finished(checkpoint_manager)

  max_logging.log(f"Elapse for checkpoint save: {(time.time() - start) / 60:.2f} min")
+


delete this new line addition

snehalv2002 · 2026-09-09T20:47:10Z

+      with safe_open(local_path, framework=framework, device="cpu") as f:
+        final_key = resolve_scale_key(resolved_key, f.keys())
+        t = f.get_tensor(final_key)
+        if torch is not None and isinstance(t, torch.Tensor):
+          if hasattr(torch, "float8_e4m3fn") and t.dtype == torch.float8_e4m3fn:
+            return t.view(torch.uint8).numpy().view(ml_dtypes.float8_e4m3fn)
+          elif hasattr(torch, "float8_e5m2") and t.dtype == torch.float8_e5m2:
+            return t.view(torch.uint8).numpy().view(ml_dtypes.float8_e5m2)
+          elif t.dtype == torch.bfloat16:
+            if self.save_dtype in ("float32", DType.FLOAT32):
+              return t.to(torch.float32).numpy()
+            return t.view(torch.int16).numpy().view(ml_dtypes.bfloat16)
+          elif t.dtype == torch.float16:
+            if self.save_dtype in ("float32", DType.FLOAT32):
+              return t.to(torch.float32).numpy()
+            elif self.save_dtype in ("bfloat16", DType.BFLOAT16):
+              return t.to(torch.float32).numpy().astype(ml_dtypes.bfloat16)
+            return t.numpy()
+          else:
+            return t.numpy()
+        return t


can we simplify this somehow?

gobbleturk · 2026-09-10T17:37:50Z

  if config.scan_layers:
    # If it's a standard scanned layer, we use the configured param_scan_axis.
-    axis_to_stack = config.param_scan_axis
+    axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0


can you include a comment here?

My guess is this is used for 1D scales?

Yes for 1D scales. Added a comment to specify this

codecov · 2026-09-10T21:27:55Z

Codecov Report

❌ Patch coverage is 5.10204% with 93 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/checkpoint_conversion/to_maxtext.py	4.21%	91 Missing ⚠️
src/maxtext/checkpoint_conversion/utils/utils.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

…stion in to_maxtext Extends the MaxText HuggingFace checkpoint conversion pipeline to directly ingest FP8 weight tensors and companion scale tensors. Key Changes: - to_maxtext: Support float8_e4m3fn and float8_e5m2 save dtypes in _convert_tensor_to_numpy and _eager_getter, preserving each tensor's source precision (companion scales keep their source dtype; unquantized modules stay bfloat16). - to_maxtext: Support scale key suffix resolution (.weight_scale, .scale, .scale_inv, .weight_scale_inv) in LazyHFLoader and _eager_getter. - tensor_handling & to_maxtext: Handle 1D per-layer scalar scale stacking along axis 0 when scan_layers=True. - utils: Support slicing 1D MaxText parameters along axis 0 during checkpoint export. - tests: Add unit tests verifying FP8 tensor casting and scale suffix resolution in hf_checkpoint_conversion_test.py.

gemini-code-assist Bot reviewed Aug 28, 2026

View reviewed changes

snehalv2002 changed the title ~~Pr/fp8 to maxtext~~ feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext Aug 28, 2026

This was referenced Aug 28, 2026

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

Open

feat(models): add llama3.1-8b-fp8 model config and HuggingFace mapping #5054

Draft

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

Draft

snehalv2002 changed the base branch from main to pr/fp8-dequant-engine August 28, 2026 22:52

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 56098db to 3d13d99 Compare September 2, 2026 17:56

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 3d13d99 to 4e38337 Compare September 2, 2026 21:27

snehalv2002 force-pushed the pr/fp8-to-maxtext branch 6 times, most recently from cb0ebb5 to f54bfa9 Compare September 8, 2026 20:43

snehalv2002 commented Sep 9, 2026

View reviewed changes

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from f54bfa9 to 44cc3a7 Compare September 9, 2026 20:43

snehalv2002 commented Sep 9, 2026

View reviewed changes

snehalv2002 marked this pull request as ready for review September 9, 2026 22:09

snehalv2002 requested review from RissyRan, bvandermoon, gobbleturk, hengtaoguo, jiangjy1982, parambole and shuningjin as code owners September 9, 2026 22:09

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from ffb5bec to a1edf65 Compare September 10, 2026 16:56

gobbleturk reviewed Sep 10, 2026

View reviewed changes

gobbleturk approved these changes Sep 10, 2026

View reviewed changes

Shuwen-Fang self-requested a review September 10, 2026 20:46

Shuwen-Fang approved these changes Sep 10, 2026

View reviewed changes

snehalv2002 force-pushed the pr/fp8-to-maxtext branch 2 times, most recently from dbf4b02 to fde4737 Compare September 10, 2026 21:10

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from fde4737 to ed14346 Compare September 11, 2026 00:39

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from ed14346 to 205de7a Compare September 11, 2026 17:47

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 205de7a to 41f7756 Compare September 11, 2026 18:22

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 41f7756 to ebc3f73 Compare September 11, 2026 21:24

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from ebc3f73 to cb11b27 Compare September 11, 2026 21:27

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from cb11b27 to 5e61891 Compare September 12, 2026 00:14

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 5e61891 to f3a466f Compare September 12, 2026 00:16

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from f3a466f to e4f3ad5 Compare September 12, 2026 04:14

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from e4f3ad5 to 69bc314 Compare September 12, 2026 11:00

snehalv2002 force-pushed the pr/fp8-to-maxtext branch from 69bc314 to c84e000 Compare September 12, 2026 20:35

Navigation Menu

Uh oh!

feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext - #5052

feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion in to_maxtext#5052
snehalv2002 wants to merge 1 commit into
pr/fp8-dequant-enginefrom
pr/fp8-to-maxtext

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!

snehalv2002 Sep 9, 2026

Uh oh!

gemini-code-assist Bot Aug 28, 2026

Uh oh!

snehalv2002 Sep 9, 2026 •

edited

Loading

Uh oh!

snehalv2002 Sep 9, 2026

Uh oh!

snehalv2002 Sep 9, 2026

Uh oh!

gobbleturk Sep 10, 2026

Uh oh!

snehalv2002 Sep 10, 2026

Uh oh!

codecov Bot commented Sep 10, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants


		if "HF_HOME" not in os.environ and os.path.exists("/dev/shm"):
		os.environ["HF_HOME"] = "/dev/shm/hf_cache"

		checkpointing.wait_until_finished(checkpoint_manager)

		max_logging.log(f"Elapse for checkpoint save: {(time.time() - start) / 60:.2f} min")

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!

snehalv2002 Sep 9, 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!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

gobbleturk Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

snehalv2002 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

snehalv2002 commented Aug 28, 2026 •

edited

Loading

snehalv2002 Sep 9, 2026 •

edited

Loading

codecov Bot commented Sep 10, 2026 •

edited

Loading