From e9d6f601d0dd7d0b9ca6030adc16534c611477f6 Mon Sep 17 00:00:00 2001 From: Snehal Verma Date: Thu, 10 Sep 2026 21:02:20 +0000 Subject: [PATCH] feat(checkpoint_conversion): support direct FP8 and scale tensor ingestion 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. --- .../checkpoint_conversion/to_maxtext.py | 180 +++++++++++++++--- .../utils/tensor_handling.py | 7 +- .../checkpoint_conversion/utils/utils.py | 9 +- tests/unit/to_maxtext_test.py | 106 +++++++++++ 4 files changed, 270 insertions(+), 32 deletions(-) create mode 100644 tests/unit/to_maxtext_test.py diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index 17ced7f095..eb2b134466 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -83,10 +83,71 @@ except ImportError: torch = None - absl.logging.set_verbosity(absl.logging.INFO) # for max_logging.log +def resolve_scale_key(key: str, container) -> str: + """Resolves fallback keys from .weight_scale to .weight_scale_inv.""" + if key in container: + return key + if key.endswith(".weight_scale"): + alt_key = key + "_inv" + if alt_key in container: + return alt_key + return key + + +def _convert_tensor_to_numpy(tensor, save_dtype: str = "bfloat16") -> np.ndarray: + """Converts a PyTorch tensor or NumPy array to the target NumPy format. + + Args: + tensor: Input tensor from PyTorch/safetensors or NumPy array. + save_dtype: Target checkpoint storage dtype for weights (e.g. "bfloat16", "float8_e4m3fn"). + For FP8 checkpoints (save_dtype="float8_e4m3fn"), only quantized weights are stored as FP8. + Every other tensor keeps its source precision: companion scale tensors are stored in + whatever dtype the source checkpoint provides (e.g. bfloat16 for Qwen3.5-35B-A3B-FP8), + and unquantized modules (e.g. embeddings, router gates) remain in their native precision. + Source dtypes that are neither FP8, bfloat16, nor float32 are upcast to float32. + """ + if torch is not None and isinstance(tensor, torch.Tensor): + tensor = tensor.detach().cpu() + # If target is float32, convert all tensors to float32 + if save_dtype in ("float32", DType.FLOAT32): + return tensor.to(torch.float32).numpy() + + # Target is float8_e4m3fn (weight-only quantization mode): + # In an FP8 checkpoint, only quantized weights are stored as FP8. We preserve each + # tensor's source precision: companion scales keep whatever dtype the source provides + # (bfloat16 for Qwen3.5-35B-A3B-FP8), and unquantized modules (e.g. embeddings, router + # gates) remain in bfloat16. + if save_dtype in ("float8_e4m3fn", DType.FLOAT8_E4M3FN): + if tensor.dtype == torch.float8_e4m3fn: + return tensor.view(torch.uint8).numpy().view(ml_dtypes.float8_e4m3fn) + elif tensor.dtype == torch.bfloat16: + return tensor.view(torch.int16).numpy().view(ml_dtypes.bfloat16) + elif tensor.dtype == torch.float32: + return tensor.numpy() + else: + return tensor.to(torch.float32).numpy() + + # Target is bfloat16 + if save_dtype in ("bfloat16", DType.BFLOAT16): + if tensor.dtype == torch.bfloat16: + return tensor.view(torch.int16).numpy().view(ml_dtypes.bfloat16) + elif tensor.dtype == torch.float8_e4m3fn: + return tensor.to(torch.float32).numpy().astype(ml_dtypes.bfloat16) + else: + return tensor.to(torch.float32).numpy().astype(ml_dtypes.bfloat16) + + return tensor.numpy() + + if isinstance(tensor, np.ndarray): + if save_dtype in ("float32", DType.FLOAT32) and tensor.dtype != np.float32: + return tensor.astype(np.float32) + return tensor + return np.asarray(tensor) + + class LazyHFLoader: """ Loads Hugging Face weights on-demand to minimize RAM usage. @@ -108,10 +169,11 @@ class LazyHFLoader: can still occur in parallel. """ - def __init__(self, model_id, token, revision=None): + def __init__(self, model_id, token, revision=None, save_dtype: str = "bfloat16"): self.model_id = model_id self.token = token self.revision = revision + self.save_dtype = save_dtype # Whether loads from local directory self.is_local = os.path.isdir(self.model_id) self.shard_map = {} @@ -178,7 +240,9 @@ def get_tensor(self, key: str) -> np.ndarray: and reads only the required tensor's data from disk. """ # Handle single-file models (shard map key might be None or we just know the filename) - shard_name = self.shard_map.get(key) + resolved_key = resolve_scale_key(key, self.shard_map) + shard_name = self.shard_map.get(resolved_key) + if shard_name is None and None in self.shard_map: shard_name = self.shard_map[None] elif shard_name is None: @@ -205,8 +269,10 @@ def get_tensor(self, key: str) -> np.ndarray: # STEP 2: Lock ONLY the reading into RAM. # This prevents multiple threads from simultaneously allocating large chunks of RAM. with self._ram_lock: - with safe_open(local_path, framework="np", device="cpu") as f: - return f.get_tensor(key) + framework = "pt" if torch is not None else "np" + with safe_open(local_path, framework=framework, device="cpu") as f: + t = f.get_tensor(resolved_key) + return _convert_tensor_to_numpy(t, self.save_dtype) class LazyTensor: @@ -224,7 +290,16 @@ def __init__( ): self._load_fn = load_fn self.shape = shape - self.dtype = np.dtype(dtype) + try: + self.dtype = np.dtype(dtype) + except (TypeError, ValueError): + dtype_str = str(dtype) + if "float8_e4m3fn" in dtype_str: + self.dtype = np.dtype(ml_dtypes.float8_e4m3fn) + elif "bfloat16" in dtype_str: + self.dtype = np.dtype(ml_dtypes.bfloat16) + else: + self.dtype = np.dtype(np.float32) self.ndim = len(shape) self.name = name @@ -425,8 +500,11 @@ def _build_single_axis_stacked_tensor( tensors_to_stack = [] 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 + # If the target tensor rank exceeds param_scan_axis (e.g., multidimensional weights or 2D block scales), + # stack along param_scan_axis (typically axis 1 in scanned layers). For 1D tensors (e.g., per-layer scalar + # quantization scales like Llama 3.1 FP8 `kernel_scale` with shape (num_layers,)), stack along axis 0 + # because axis 1 does not exist on 1D tensors and stacking scalars along axis 0 produces the 1D vector. + axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0 else: # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0). axis_to_stack = 0 @@ -725,7 +803,7 @@ def convert_lora_to_maxtext_adapter( mt_adapter_tree = {} mapped_count = 0 - target_dtype = ml_dtypes.bfloat16 if save_dtype == "bfloat16" else np.float32 + target_dtype = ml_dtypes.bfloat16 if save_dtype in ("bfloat16", "float8_e4m3fn") else np.float32 collected_weights = {} @@ -850,6 +928,37 @@ def _merged_getter(key): return _merged_getter +def _extract_conversion_args(args: Sequence[str]) -> tuple[dict[str, Any], list[str]]: + """Extracts known checkpoint conversion args and returns remaining args.""" + conversion_arg_names = { + "save_dtype", + "hf_model_path", + "lazy_load_tensors", + "eager_load_method", + "revision", + "simulated_cpu_devices_count", + } + extracted = {} + cleaned_args = [] + for arg_item in args: + matched = False + for name in conversion_arg_names: + key_prefix = f"{name}=" + if arg_item.startswith(key_prefix) or arg_item.startswith(f"--{key_prefix}"): + val = arg_item.split("=", 1)[1] + if name == "lazy_load_tensors": + extracted[name] = str2bool(val) + elif name == "simulated_cpu_devices_count": + extracted[name] = int(val) + else: + extracted[name] = val + matched = True + break + if not matched: + cleaned_args.append(arg_item) + return extracted, cleaned_args + + def main( args: Sequence[str], lazy_load_tensors: bool = True, @@ -860,9 +969,17 @@ def main( simulated_cpu_devices_count: int = 16, ) -> None: overall_start = time.time() + extracted, cleaned_args = _extract_conversion_args(args) + save_dtype = extracted.get("save_dtype", save_dtype) + hf_model_path = extracted.get("hf_model_path", hf_model_path) + lazy_load_tensors = extracted.get("lazy_load_tensors", lazy_load_tensors) + eager_load_method = extracted.get("eager_load_method", eager_load_method) + revision = extracted.get("revision", revision) + simulated_cpu_devices_count = extracted.get("simulated_cpu_devices_count", simulated_cpu_devices_count) + args = cleaned_args # Check if the user is using an Instruct version. If so, use the base model architecture - for i, arg in enumerate(args): - if arg.startswith("model_name="): + for i, raw_param in enumerate(args): + if raw_param.startswith("model_name="): model_name_arg = args[i].split("=")[1] model_name_original = model_name_arg if "-Instruct" in model_name_arg: @@ -916,7 +1033,7 @@ def main( # Define the appropriate tensor getter based on mode if lazy_load_tensors: max_logging.log(f"Lazy loading ENABLED. Initializing LazyHFLoader for: {model_id}...") - hf_loader = LazyHFLoader(model_id, hf_token, revision=revision) + hf_loader = LazyHFLoader(model_id, hf_token, revision=revision, save_dtype=save_dtype) print_ram_usage("After LazyLoader init") tensor_getter = hf_loader.get_tensor @@ -976,22 +1093,12 @@ def main( } def _eager_getter(key): - if key not in hf_state_dict_numpy: + resolved_key = resolve_scale_key(key, hf_state_dict_numpy) + if resolved_key not in hf_state_dict_numpy: raise ValueError(f"HuggingFace key {key} not found in state_dict.") - v = hf_state_dict_numpy[key] - # target dtype is "float32" - if save_dtype == DType.FLOAT32: - return v.to(torch.float32).numpy() - # target dtype is "bfloat16" - elif save_dtype == DType.BFLOAT16: - # - torch.bfloat16 -> torch.float32 -> np.float32 -> ml_dtypes.bfloat16 - # As numpy doesn't accept bfloat16 directly, we convert to float32 first - # - torch.float16 -> np.float16 -> ml_dtypes.bfloat16 - # - torch.float32 -> np.float32 -> ml_dtypes.bfloat16 - if v.dtype == torch.bfloat16: - v = v.to(torch.float32) - return v.numpy().astype(ml_dtypes.bfloat16) - raise NotImplementedError(f"Save dtype {save_dtype} is not currently implemented.") + + v = hf_state_dict_numpy[resolved_key] + return _convert_tensor_to_numpy(v, save_dtype) tensor_getter = _eager_getter @@ -1153,7 +1260,7 @@ def _eager_getter(key): type=str, required=False, default="bfloat16", - choices=["float32", "bfloat16"], + choices=["float32", "bfloat16", "float8_e4m3fn"], help="Save MaxText weights in specified dtype", ) # Determines the logical sharding of the output checkpoint by partitioning @@ -1172,6 +1279,23 @@ def _eager_getter(key): parser.add_argument( "--simulated_cpu_devices_count", type=int, required=False, default=16, help="Sharding of checkpoint" ) + # Normalize key=value CLI arguments for local_args if passed without leading dashes + normalized_argv = [sys.argv[0]] + for raw_cli_arg in sys.argv[1:]: + for cli_prefix in ( + "save_dtype=", + "hf_model_path=", + "lazy_load_tensors=", + "eager_load_method=", + "revision=", + "simulated_cpu_devices_count=", + ): + if raw_cli_arg.startswith(cli_prefix): + raw_cli_arg = "--" + raw_cli_arg + break + normalized_argv.append(raw_cli_arg) + sys.argv = normalized_argv + # Parse local arguments # Parse known args returns the namespace AND the list of remaining arguments local_args, remaining_args = parser.parse_known_args() diff --git a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py index 039948b836..c4c064b82f 100644 --- a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py +++ b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py @@ -172,8 +172,11 @@ def _build_single_axis_stacked_tensor( target_dtype = target_leaf.dtype if hasattr(target_leaf, "dtype") else np.float32 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 + # If the target tensor rank exceeds param_scan_axis (e.g., multidimensional weights or 2D block scales), + # stack along param_scan_axis (typically axis 1 in scanned layers). For 1D tensors (e.g., per-layer scalar + # quantization scales like Llama 3.1 FP8 `kernel_scale` with shape (num_layers,)), stack along axis 0 + # because axis 1 does not exist on 1D tensors and stacking scalars along axis 0 produces the 1D vector. + axis_to_stack = config.param_scan_axis if len(target_shape) > config.param_scan_axis else 0 else: # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0). axis_to_stack = 0 diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index b3a78bb73c..7008461ade 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -286,8 +286,13 @@ def process_maxtext_param( # Case 2 or 3: The source tensor is stacked on a single axis. if maxtext_config.scan_layers: max_logging.log("\tscan") - # Case 2: Standard scanned layer. Stacked ONLY on the layer axis. - axis_to_slice = maxtext_config.param_scan_axis + # Case 2: Standard scanned layer. Stacked on layer axis (param_scan_axis) for multidimensional + # weights or 2D block scales, or axis 0 for 1D tensors (e.g. per-layer scalar scales like + # Llama 3.1 FP8 `kernel_scale` with shape (num_layers,)). + weight_sample = maxtext_param_weight[0] if isinstance(maxtext_param_weight, list) else maxtext_param_weight + axis_to_slice = ( + maxtext_config.param_scan_axis if getattr(weight_sample, "ndim", 0) > maxtext_config.param_scan_axis else 0 + ) else: max_logging.log("\tunscan moe") # Case 3: Unscanned MoE layer. Stacked ONLY on the expert axis. Assuming expert is axis 0. diff --git a/tests/unit/to_maxtext_test.py b/tests/unit/to_maxtext_test.py new file mode 100644 index 0000000000..b608c35026 --- /dev/null +++ b/tests/unit/to_maxtext_test.py @@ -0,0 +1,106 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for to_maxtext.py""" + +import unittest +import ml_dtypes +import numpy as np +import pytest + +try: + import torch +except ImportError: + torch = None + +pytestmark = [pytest.mark.decoupled_target] + +from maxtext.checkpoint_conversion.to_maxtext import ( + _convert_tensor_to_numpy, + _extract_conversion_args, + resolve_scale_key, +) + + +class ResolveScaleKeyTest(unittest.TestCase): + + def test_exact_match(self): + container = {"model.layers.0.mlp.gate.weight": 1} + self.assertEqual(resolve_scale_key("model.layers.0.mlp.gate.weight", container), "model.layers.0.mlp.gate.weight") + + def test_weight_scale_to_weight_scale_inv_fallback(self): + container = {"model.layers.0.mlp.experts.0.gate_proj.weight_scale_inv": 1} + resolved = resolve_scale_key("model.layers.0.mlp.experts.0.gate_proj.weight_scale", container) + self.assertEqual(resolved, "model.layers.0.mlp.experts.0.gate_proj.weight_scale_inv") + + def test_unmatched_key_returns_original(self): + container = {} + self.assertEqual( + resolve_scale_key("model.layers.0.mlp.experts.0.gate_proj.weight_scale", container), + "model.layers.0.mlp.experts.0.gate_proj.weight_scale", + ) + + +class ConvertTensorToNumpyTest(unittest.TestCase): + + def test_torch_float8_e4m3fn_to_numpy(self): + if torch is None: + self.skipTest("torch not available") + t = torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32).to(torch.float8_e4m3fn) + arr = _convert_tensor_to_numpy(t, save_dtype="float8_e4m3fn") + self.assertEqual(arr.dtype, ml_dtypes.float8_e4m3fn) + self.assertEqual(arr.shape, (3,)) + + def test_torch_float8_e4m3fn_to_float32(self): + if torch is None: + self.skipTest("torch not available") + t = torch.tensor([1.0, 2.0, -1.0], dtype=torch.float32).to(torch.float8_e4m3fn) + arr = _convert_tensor_to_numpy(t, save_dtype="float32") + self.assertEqual(arr.dtype, np.float32) + + def test_torch_bfloat16_to_numpy(self): + if torch is None: + self.skipTest("torch not available") + t = torch.tensor([1.0, 2.0], dtype=torch.bfloat16) + arr = _convert_tensor_to_numpy(t, save_dtype="bfloat16") + self.assertEqual(arr.dtype, ml_dtypes.bfloat16) + + def test_numpy_passthrough(self): + arr_in = np.array([1.0, 2.0], dtype=np.float32) + arr_out = _convert_tensor_to_numpy(arr_in, save_dtype="float8_e4m3fn") + self.assertTrue(np.array_equal(arr_in, arr_out)) + + +class ExtractConversionArgsTest(unittest.TestCase): + + def test_extract_mixed_flags(self): + args = [ + "base.yml", + "model_name=qwen3.5-35b-a3b-fp8", + "save_dtype=float8_e4m3fn", + "--lazy_load_tensors=false", + "eager_load_method=safetensors", + "--simulated_cpu_devices_count=4", + "steps=10", + ] + extracted, remaining = _extract_conversion_args(args) + self.assertEqual(extracted["save_dtype"], "float8_e4m3fn") + self.assertEqual(extracted["lazy_load_tensors"], False) + self.assertEqual(extracted["eager_load_method"], "safetensors") + self.assertEqual(extracted["simulated_cpu_devices_count"], 4) + self.assertEqual(remaining, ["base.yml", "model_name=qwen3.5-35b-a3b-fp8", "steps=10"]) + + +if __name__ == "__main__": + unittest.main()