feat(models): add llama3.1-8b-fp8 model config and HuggingFace mapping - #5054
Draft
snehalv2002 wants to merge 4 commits into
Draft
feat(models): add llama3.1-8b-fp8 model config and HuggingFace mapping#5054snehalv2002 wants to merge 4 commits into
snehalv2002 wants to merge 4 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for FP8 weight-only storage with dynamic dequantization in MaxText, enabling more efficient model loading and inference. The changes include updates to checkpoint loading logic to handle FP8 scales, the addition of FP8-aware linear layers, and new model configurations. The review feedback suggests improving type compatibility by using 'Mapping' instead of 'dict' to better support 'FrozenDict', refactoring duplicated key resolution and FP8 check logic into reusable helper functions, and optimizing performance through set-based lookups.
Comment on lines
32
to
+33
| import jax | ||
| import jax.numpy as jnp |
There was a problem hiding this comment.
Import Mapping from collections.abc to support checking for Flax Linen's FrozenDict and other dictionary-like mapping types.
Suggested change
| import jax | |
| import jax.numpy as jnp | |
| import jax | |
| import jax.numpy as jnp | |
| from collections.abc import Mapping |
Comment on lines
+722
to
+847
| def _find_matching_meta_subtree(want_bare: Any, meta_tree: Any) -> Any: | ||
| """Unwraps metadata tree wrappers (e.g. 'params', 'model_params') to align with want_bare.""" | ||
| 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"] | ||
|
|
||
| return meta_tree | ||
|
|
||
|
|
||
| def _augment_target_with_scales(want_node: Any, meta_node: Any) -> Any: | ||
| """Augments `want_node` with `kernel_scale` from `meta_node` if present in checkpoint but not in want.""" | ||
| if not isinstance(want_node, dict) or not isinstance(meta_node, dict): | ||
| return want_node | ||
|
|
||
| augmented = {} | ||
| has_kernel = "kernel" in want_node | ||
| has_kernel_scale = "kernel_scale" in want_node | ||
| meta_has_kernel_scale = "kernel_scale" in meta_node | ||
|
|
||
| if has_kernel and not has_kernel_scale and meta_has_kernel_scale: | ||
| # Mode C: Target wants unquantized weights, but checkpoint has companion kernel_scale. | ||
| want_kernel = want_node["kernel"] | ||
| meta_scale = meta_node["kernel_scale"] | ||
| meta_kernel = meta_node.get("kernel") | ||
|
|
||
| # Restore kernel with stored dtype (e.g. float8) to avoid Orbax converting before dequantization | ||
| kernel_dtype = getattr(meta_kernel, "dtype", getattr(want_kernel, "dtype", jnp.bfloat16)) | ||
| augmented["kernel"] = jax.ShapeDtypeStruct( | ||
| shape=getattr(want_kernel, "shape", getattr(meta_kernel, "shape", ())), | ||
| dtype=kernel_dtype, | ||
| sharding=getattr(want_kernel, "sharding", None), | ||
| ) | ||
| augmented["kernel_scale"] = jax.ShapeDtypeStruct( | ||
| shape=getattr(meta_scale, "shape", ()), | ||
| dtype=getattr(meta_scale, "dtype", jnp.float32), | ||
| sharding=None, | ||
| ) | ||
| elif has_kernel: | ||
| augmented["kernel"] = want_node["kernel"] | ||
|
|
||
| for k, v in want_node.items(): | ||
| if k == "kernel": | ||
| continue | ||
| augmented[k] = _augment_target_with_scales( | ||
| v, | ||
| meta_node.get(k) if isinstance(meta_node, dict) else None, | ||
| ) | ||
|
|
||
| return augmented | ||
|
|
||
|
|
||
| def _augment_want_with_scales(want: Any, meta_tree: Any, is_nnx: bool, restore_key: str) -> Any: | ||
| """Augments the target params dictionary with scales from checkpoint metadata if needed.""" | ||
| if meta_tree is None: | ||
| return want | ||
|
|
||
| if is_nnx or restore_key in ("model_params", "model"): | ||
| meta_weights = _find_matching_meta_subtree(want, meta_tree) | ||
| return _augment_target_with_scales(want, meta_weights) | ||
| else: | ||
| # Linen: want is {"params": bare_weights} or bare_weights | ||
| if isinstance(want, dict) and "params" in want and len(want) == 1: | ||
| want_bare = want["params"] | ||
| meta_weights = _find_matching_meta_subtree(want_bare, meta_tree) | ||
| augmented_bare = _augment_target_with_scales(want_bare, meta_weights) | ||
| return {"params": augmented_bare} | ||
| else: | ||
| meta_weights = _find_matching_meta_subtree(want, meta_tree) | ||
| return _augment_target_with_scales(want, meta_weights) | ||
|
|
||
|
|
||
| def maybe_dequantize_restored_params(restored_weights: Any, want: Any) -> Any: | ||
| """Dequantizes restored weights if checkpoint contained companion kernel_scale but want did not. | ||
|
|
||
| For each parameter dictionary containing both 'kernel' and 'kernel_scale' where the target | ||
| model (`want`) does not expect 'kernel_scale', dynamically dequantizes 'kernel' to the | ||
| target compute dtype using `dequantize_weight` from `maxtext.layers.linears` and removes | ||
| 'kernel_scale' from the restored parameter dictionary. | ||
|
|
||
| Args: | ||
| restored_weights: The restored parameter PyTree from the checkpoint. | ||
| want: The expected parameter PyTree structure / ShapeDtypeStructs. | ||
|
|
||
| Returns: | ||
| The parameter PyTree with dequantized kernels and omitted kernel_scale where applicable. | ||
| """ | ||
| if not isinstance(restored_weights, dict): | ||
| return restored_weights | ||
|
|
||
| has_kernel = "kernel" in restored_weights | ||
| has_scale = "kernel_scale" in restored_weights | ||
| want_expects_scale = isinstance(want, dict) and "kernel_scale" in want | ||
|
|
||
| if has_kernel and has_scale and not want_expects_scale: | ||
| target_kernel = want.get("kernel") if isinstance(want, dict) else None | ||
| target_dtype = getattr(target_kernel, "dtype", jnp.bfloat16) | ||
|
|
||
| kernel = restored_weights["kernel"] | ||
| scale = restored_weights["kernel_scale"] | ||
| dequantized_kernel = linears.dequantize_weight(kernel, scale, compute_dtype=target_dtype) | ||
|
|
||
| out = {"kernel": dequantized_kernel} | ||
| for k, v in restored_weights.items(): | ||
| if k in ("kernel", "kernel_scale"): | ||
| continue | ||
| want_sub = want.get(k) if isinstance(want, dict) else None | ||
| out[k] = maybe_dequantize_restored_params(v, want_sub) | ||
| return out | ||
|
|
||
| out = {} | ||
| for k, v in restored_weights.items(): | ||
| want_sub = want.get(k) if isinstance(want, dict) else None | ||
| out[k] = maybe_dequantize_restored_params(v, want_sub) | ||
| return out |
There was a problem hiding this comment.
Using isinstance(..., dict) will fail for Flax Linen's FrozenDict (which does not inherit from dict in Python). This prevents any FP8 scaling or dequantization from being applied to Linen models. Replacing dict with Mapping ensures full compatibility with both NNX and Linen parameter structures.
def _find_matching_meta_subtree(want_bare: Any, meta_tree: Any) -> Any:
"""Unwraps metadata tree wrappers (e.g. 'params', 'model_params') to align with want_bare."""
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"]
return meta_tree
def _augment_target_with_scales(want_node: Any, meta_node: Any) -> Any:
"""Augments `want_node` with `kernel_scale` from `meta_node` if present in checkpoint but not in want."""
if not isinstance(want_node, Mapping) or not isinstance(meta_node, Mapping):
return want_node
augmented = {}
has_kernel = "kernel" in want_node
has_kernel_scale = "kernel_scale" in want_node
meta_has_kernel_scale = "kernel_scale" in meta_node
if has_kernel and not has_kernel_scale and meta_has_kernel_scale:
# Mode C: Target wants unquantized weights, but checkpoint has companion kernel_scale.
want_kernel = want_node["kernel"]
meta_scale = meta_node["kernel_scale"]
meta_kernel = meta_node.get("kernel")
# Restore kernel with stored dtype (e.g. float8) to avoid Orbax converting before dequantization
kernel_dtype = getattr(meta_kernel, "dtype", getattr(want_kernel, "dtype", jnp.bfloat16))
augmented["kernel"] = jax.ShapeDtypeStruct(
shape=getattr(want_kernel, "shape", getattr(meta_kernel, "shape", ())),
dtype=kernel_dtype,
sharding=getattr(want_kernel, "sharding", None),
)
augmented["kernel_scale"] = jax.ShapeDtypeStruct(
shape=getattr(meta_scale, "shape", ()),
dtype=getattr(meta_scale, "dtype", jnp.float32),
sharding=None,
)
elif has_kernel:
augmented["kernel"] = want_node["kernel"]
for k, v in want_node.items():
if k == "kernel":
continue
augmented[k] = _augment_target_with_scales(
v,
meta_node.get(k) if isinstance(meta_node, Mapping) else None,
)
return augmented
def _augment_want_with_scales(want: Any, meta_tree: Any, is_nnx: bool, restore_key: str) -> Any:
"""Augments the target params dictionary with scales from checkpoint metadata if needed."""
if meta_tree is None:
return want
if is_nnx or restore_key in ("model_params", "model"):
meta_weights = _find_matching_meta_subtree(want, meta_tree)
return _augment_target_with_scales(want, meta_weights)
else:
# Linen: want is {"params": bare_weights} or bare_weights
if isinstance(want, Mapping) and "params" in want and len(want) == 1:
want_bare = want["params"]
meta_weights = _find_matching_meta_subtree(want_bare, meta_tree)
augmented_bare = _augment_target_with_scales(want_bare, meta_weights)
return {"params": augmented_bare}
else:
meta_weights = _find_matching_meta_subtree(want, meta_tree)
return _augment_target_with_scales(want, meta_weights)
def maybe_dequantize_restored_params(restored_weights: Any, want: Any) -> Any:
"""Dequantizes restored weights if checkpoint contained companion kernel_scale but want did not.
For each parameter dictionary containing both 'kernel' and 'kernel_scale' where the target
model (`want`) does not expect 'kernel_scale', dynamically dequantizes 'kernel' to the
target compute dtype using `dequantize_weight` from `maxtext.layers.linears` and removes
'kernel_scale' from the restored parameter dictionary.
Args:
restored_weights: The restored parameter PyTree from the checkpoint.
want: The expected parameter PyTree structure / ShapeDtypeStructs.
Returns:
The parameter PyTree with dequantized kernels and omitted kernel_scale where applicable.
"""
if not isinstance(restored_weights, Mapping):
return restored_weights
has_kernel = "kernel" in restored_weights
has_scale = "kernel_scale" in restored_weights
want_expects_scale = isinstance(want, Mapping) and "kernel_scale" in want
if has_kernel and has_scale and not want_expects_scale:
target_kernel = want.get("kernel") if isinstance(want, Mapping) else None
target_dtype = getattr(target_kernel, "dtype", jnp.bfloat16)
kernel = restored_weights["kernel"]
scale = restored_weights["kernel_scale"]
dequantized_kernel = linears.dequantize_weight(kernel, scale, compute_dtype=target_dtype)
out = {"kernel": dequantized_kernel}
for k, v in restored_weights.items():
if k in ("kernel", "kernel_scale"):
continue
want_sub = want.get(k) if isinstance(want, Mapping) else None
out[k] = maybe_dequantize_restored_params(v, want_sub)
return out
out = {}
for k, v in restored_weights.items():
want_sub = want.get(k) if isinstance(want, Mapping) else None
out[k] = maybe_dequantize_restored_params(v, want_sub)
return out
Comment on lines
+181
to
+198
| resolved_key = key | ||
| shard_name = self.shard_map.get(resolved_key) | ||
| 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 |
There was a problem hiding this comment.
The fallback logic for resolving .weight_scale and .scale keys is duplicated across multiple places in this file (e.g., in get_tensor and _eager_getter). We can refactor this into a clean nested helper function resolve_key_fallback to improve maintainability and readability.
Suggested change
| resolved_key = key | |
| shard_name = self.shard_map.get(resolved_key) | |
| 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 | |
| def resolve_key_fallback(k, container): | |
| if k.endswith(".weight_scale"): | |
| base, suffixes = k[:-len(".weight_scale")], [".scale", ".weight_scale_inv", ".scale_inv"] | |
| elif k.endswith(".scale"): | |
| base, suffixes = k[:-len(".scale")], [".weight_scale", ".scale_inv", ".weight_scale_inv"] | |
| else: | |
| return k | |
| for suffix in suffixes: | |
| alt_key = base + suffix | |
| if alt_key in container: | |
| return alt_key | |
| return k | |
| resolved_key = key | |
| shard_name = self.shard_map.get(resolved_key) | |
| if shard_name is None: | |
| resolved_key = resolve_key_fallback(key, self.shard_map) | |
| shard_name = self.shard_map.get(resolved_key) |
Comment on lines
+228
to
+241
| final_key = resolved_key | ||
| 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 |
There was a problem hiding this comment.
We can reuse the resolve_key_fallback helper function here. Additionally, converting f.keys() to a set avoids repeated
f_keys = set(f.keys())
final_key = resolve_key_fallback(resolved_key, f_keys)
Comment on lines
+1042
to
+1055
| resolved_key = key | ||
| if resolved_key not in hf_state_dict_numpy: | ||
| 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 hf_state_dict_numpy: | ||
| resolved_key = alt_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 hf_state_dict_numpy: | ||
| resolved_key = alt_key | ||
| break |
There was a problem hiding this comment.
We can define and use the same resolve_key_fallback helper function here to keep the code clean and avoid deeply nested conditional blocks.
Suggested change
| resolved_key = key | |
| if resolved_key not in hf_state_dict_numpy: | |
| 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 hf_state_dict_numpy: | |
| resolved_key = alt_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 hf_state_dict_numpy: | |
| resolved_key = alt_key | |
| break | |
| def resolve_key_fallback(k, container): | |
| if k.endswith(".weight_scale"): | |
| base, suffixes = k[:-len(".weight_scale")], [".scale", ".weight_scale_inv", ".scale_inv"] | |
| elif k.endswith(".scale"): | |
| base, suffixes = k[:-len(".scale")], [".weight_scale", ".scale_inv", ".weight_scale_inv"] | |
| else: | |
| return k | |
| for suffix in suffixes: | |
| alt_key = base + suffix | |
| if alt_key in container: | |
| return alt_key | |
| return k | |
| resolved_key = resolve_key_fallback(key, hf_state_dict_numpy) |
| 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 |
There was a problem hiding this comment.
Import is_fp8_dtype from maxtext.common.common_types to avoid duplicating the FP8 dtype checking logic.
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 |
Comment on lines
+91
to
+104
| 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 |
There was a problem hiding this comment.
Remove the duplicate implementation of _is_fp8_dtype and alias it to the imported is_fp8_dtype helper from common_types.py.
_is_fp8_dtype = is_fp8_dtype
snehalv2002
changed the title
4 tasks
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
8 times, most recently
from
September 5, 2026 01:09
a3573f9 to
4b1be84
Compare
snehalv2002
force-pushed
the
pr/fp8-llama-onboarding
branch
from
September 8, 2026 20:26
d240f88 to
6117d8e
Compare
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
from
September 8, 2026 20:26
4b1be84 to
705fd0a
Compare
snehalv2002
force-pushed
the
pr/fp8-llama-onboarding
branch
from
September 8, 2026 20:43
6117d8e to
7bdd869
Compare
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
2 times, most recently
from
September 9, 2026 22:24
c427c2c to
2d1ad3e
Compare
snehalv2002
force-pushed
the
pr/fp8-llama-onboarding
branch
from
September 9, 2026 22:24
7bdd869 to
8007b8f
Compare
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
from
September 10, 2026 16:57
2d1ad3e to
feef8d6
Compare
snehalv2002
force-pushed
the
pr/fp8-llama-onboarding
branch
2 times, most recently
from
September 10, 2026 21:04
0700124 to
a9eaede
Compare
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
from
September 10, 2026 21:04
feef8d6 to
0f27eb6
Compare
… and types Implements lightweight FP8 weight-only storage with dynamic in-memory dequantization during the forward pass in MaxText. Key Changes: - common_types: Add is_fp8_dtype, get_weight_dtype with declarative unquantized module filtering. - quantizations: Add WeightQuantConfig dataclass and get_weight_quant_config helper for modular quantization configuration. - linears: Add dequantize_weight with scalar and 2D block-scale broadcasting, and integrate kernel_scale handling into DenseGeneral. - moe: Add FP8 weight storage and dynamic dequantization support for RoutedMoE (fused and unfused paths, per-expert and 3D block scaling). - decoders & nnx_decoders: Thread quantization configurations and preserve unquantized layer precisions. - configs: Define weight_dtype, weight_quant_type, weight_block_size, and unquantized_modules in base.yml and types.py. - unit tests: Add comprehensive unit test coverage for linear layers, MoE blocks, and NNX decoders under FP8.
…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 companion scale float32 precision and unquantized module bfloat16 precision. - 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.
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
from
September 10, 2026 21:10
0f27eb6 to
10a99d2
Compare
snehalv2002
force-pushed
the
pr/fp8-llama-onboarding
branch
from
September 10, 2026 21:10
a9eaede to
7230572
Compare
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
2 times, most recently
from
September 10, 2026 21:36
fffffd7 to
5e93173
Compare
snehalv2002
force-pushed
the
pr/fp8-orbax-restoration
branch
11 times, most recently
from
September 12, 2026 20:35
45f1098 to
e72dcf2
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Registers the dedicated
llama3.1-8b-fp8model configuration and parameter mapping, serving as the driver model and reference template for onboarding quantized architectures.Motivation & Context
With core dynamic dequantization and checkpointing infrastructure established (PRs 1-3), this PR adds the model definitions and parameter mapping required to run
llama3.1-8b-fp8directly from Hugging Face pre-quantized checkpoints (neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8).Key Changes
configs/models/llama3.1-8b-fp8.yml):weight_dtype: "float8_e4m3fn",dtype: "bfloat16",decoder_block: "llama2", andlogits_via_embedding: false.utils/globals.py):"llama3.1-8b-fp8": "neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8"toHF_IDS.param_mapping.py):LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPINGandHOOK_FNSto map self-attention and MLP scale tensors (kernel_scale).llama3.1-8b-fp8inhf_model_configs.pyandhf_shape.py.Part 4 of 5 in the FP8 Weight-Only Dynamic Dequantization series (depends on #5053, #5052, #5051).
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 model configuration test suite:
pytest tests/unit/configs_test.py -k "test_llama" -vResult:
18 passed, 0 failures.Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.