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

[snehalv2002]

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

1. 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.
2. Scale Tensor Alias Resolution:
   • Added robust resolution in LazyHFLoader for common Hugging Face scale tensor naming conventions (.weight_scale, .scale, .weight_scale_inv, .scale_inv).
3. Selective Precision Preservation:
   • Unquantized layers (embeddings, norm scales) retain their high-precision dtypes even when --save_dtype="float8_e4m3fn" is specified.
4. 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]

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]

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]

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.