From 96c72f725ea34dee634a1c2e2a37c3f70732cf73 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Fri, 3 Jul 2026 20:18:29 -0700 Subject: [PATCH] Use tuple keys in flax flatten/unflatten for checkpoint conversion. PiperOrigin-RevId: 942334915 --- .../utils/load_dynamic.py | 90 ++++++-- .../utils/tensor_handling.py | 211 ++++++++++++++---- tests/unit/checkpointing_test.py | 14 +- 3 files changed, 240 insertions(+), 75 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py index 2032cf1c41..0cbe2a6e00 100644 --- a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py +++ b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py @@ -169,6 +169,7 @@ def load_sharded_hf_state(path): """ t0 = time.time() context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.SAFETENSORS) + context.safetensors.ignore_load_sharding = True with context: metadata = ocp_v1.pytree_metadata(path) simple_abstract_state = metadata.metadata @@ -186,19 +187,27 @@ def combine_sharding(sds, single_sharding): max_logging.log("Reading raw Safetensors into memory (Distributed Sharded GCS Download)...") hf_state = ocp_v1.load_pytree(path, sharded_abstract_state) max_logging.log(f"load_sharded_hf_state took {time.time() - t0:.2f}s") - return hf_state + return hf_state, simple_abstract_state -def transform_hf_state_to_mt_state(hf_state, target_tree, param_map_mt_to_hf, hook_fn_map_mt, maxtext_config): +def transform_hf_state_to_mt_state(hf_state, target_tree, param_map_mt_to_hf, hook_fn_map_mt, maxtext_config, simple_abstract_state): """Transforms HF state into MaxText state by applying param mappings and mathematical hooks.""" t0 = time.time() def tensor_getter(key): - return hf_state.pop(key) + info = simple_abstract_state.get(key) + arr = hf_state.pop(key, None) + if info is not None: + return arr, info.shape, info.dtype + return arr, None, None - flat_target = flax.traverse_util.flatten_dict(target_tree, sep=".") + flat_target = flax.traverse_util.flatten_dict(target_tree, sep=None) flat_restored = flat_target.copy() + # Create a lookup mapping from stringified/joined path to the original tuple path + # Make sure we stringify any integer indices to support mixed arrays in paths securely. + path_str_to_tuple = {".".join(map(str, path)): path for path in flat_target} + mapped_count = 0 keys_missed = [] max_logging.log("Starting fast in-memory Distributed Transformations...") @@ -206,19 +215,22 @@ def tensor_getter(key): for mt_key, hf_source in param_map_mt_to_hf.items(): mt_name = mt_key.replace("params-", "").replace("-", ".") - # Determine the correct key in flat_target + # Determine the correct key in path_str_to_tuple check_name = mt_name - if check_name not in flat_target: - if f"params.{mt_name}" in flat_target: + if check_name not in path_str_to_tuple: + if f"params.{mt_name}" in path_str_to_tuple: check_name = f"params.{mt_name}" - elif mt_key.replace("-", ".") in flat_target: + elif f"model.{mt_name}" in path_str_to_tuple: + check_name = f"model.{mt_name}" + elif mt_key.replace("-", ".") in path_str_to_tuple: check_name = mt_key.replace("-", ".") - if check_name not in flat_target: + if check_name not in path_str_to_tuple: keys_missed.append(mt_name) continue - target_leaf = flat_target[check_name] + target_path = path_str_to_tuple[check_name] + target_leaf = flat_target[target_path] hook_fn = hook_fn_map_mt.get(mt_key) load_fn = get_hf_loading_function( @@ -229,19 +241,15 @@ def tensor_getter(key): maxtext_config, ) - # Execute transformation and assign to flat_restored - t_layer = time.time() - flat_restored[check_name] = load_fn() - - max_logging.log(f"Transformed {check_name} from {hf_source} in {time.time() - t_layer:.4f}s") + flat_restored[target_path] = load_fn() mapped_count += 1 if mapped_count == 0: max_logging.log(f"All transformations missed! Sample missed mt_names: {keys_missed[:5]}") - max_logging.log(f"Sample flat_target keys: {list(flat_target.keys())[:5]}") + max_logging.log(f"Sample flat_target keys: {list(path_str_to_tuple.keys())[:5]}") max_logging.log(f"Successfully mapped {mapped_count} parameters.") - restored_params = flax.traverse_util.unflatten_dict(flat_restored, sep=".") + restored_params = flax.traverse_util.unflatten_dict(flat_restored, sep=None) if "params" in restored_params: restored_params = restored_params["params"] @@ -349,23 +357,57 @@ def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_con param_map_mt_to_hf, hook_fn_map_mt = get_hf_config_and_mappings(maxtext_config) max_logging.log(f"[1/3] Mappings derived in {time.time() - t_total:.2f}s") - target_tree = ( - abstract_unboxed_pre_state.to_pure_dict() - if isinstance(abstract_unboxed_pre_state, nnx.State) - else abstract_unboxed_pre_state.params - ) + if isinstance(abstract_unboxed_pre_state, nnx.State): + # In NNX, abstract_unboxed_pre_state contains both model and optimizer states. + # We only want to target and transform the model variables. + model_state = getattr(abstract_unboxed_pre_state, "model", None) + if model_state is not None: + target_tree = ( + model_state.to_pure_dict() + if hasattr(model_state, "to_pure_dict") + else model_state + ) + else: + target_tree = abstract_unboxed_pre_state.to_pure_dict() + else: + # In Linen, params is only the model parameters. + target_tree = abstract_unboxed_pre_state.params t1 = time.time() - hf_state = load_sharded_hf_state(path) + hf_state_raw = load_sharded_hf_state(path) + if isinstance(hf_state_raw, tuple): + hf_state, simple_abstract_state = hf_state_raw + else: + hf_state, simple_abstract_state = hf_state_raw, {} max_logging.log(f"[2/3] Distributed Sharded GCS load completed in {time.time() - t1:.2f}s") t2 = time.time() # Transform Hugging Face weight tensors on-the-fly into MaxText format # in-memory. This is done in-memory on each host, sharded across the mesh. restored_params = transform_hf_state_to_mt_state( - hf_state, target_tree, param_map_mt_to_hf, hook_fn_map_mt, maxtext_config + hf_state, target_tree, param_map_mt_to_hf, hook_fn_map_mt, maxtext_config, simple_abstract_state ) max_logging.log(f"[3/3] CPU Transformations completed in {time.time() - t2:.2f}s") max_logging.log(f"Total safetensors_dynamic duration: {time.time() - t_total:.2f}s") + if restored_params and "params" in restored_params: + restored_params = restored_params["params"] + + def _filter_shape_dtype_structs(d): + if not isinstance(d, dict): + return d + res = {} + for k, v in d.items(): + if isinstance(v, dict): + sub = _filter_shape_dtype_structs(v) + if sub: + res[k] = sub + elif not hasattr(v, "sharding") or not isinstance(v, jax.ShapeDtypeStruct): + # Only keep things that are not ShapeDtypeStruct, + # meaning real JAX arrays that we restored + res[k] = v + return res + + restored_params = _filter_shape_dtype_structs(restored_params) + return None, restored_params diff --git a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py index 508697624b..d82b7f89d6 100644 --- a/src/maxtext/checkpoint_conversion/utils/tensor_handling.py +++ b/src/maxtext/checkpoint_conversion/utils/tensor_handling.py @@ -18,6 +18,7 @@ from typing import Any, Callable, List import jax import jax.numpy as np +import numpy as onp def apply_hook_fns(weight, target_shape, hook_fns): @@ -48,6 +49,104 @@ def _binary_chunked_stack(tensors: List[np.ndarray], axis: int) -> np.ndarray: return np.concatenate([left, right], axis=axis) +def get_safe_local_array(array, shape, dtype): + """Safely extracts a SingleDevice sharded tensor into a purely local CPU Numpy Array + (returning Zeros on nodes that do not own the payload) to enable math processing.""" + if array is None: + return onp.zeros(shape, dtype=dtype), False + + is_np = not hasattr(array, "addressable_shards") + host_id = jax.process_index() + + if not is_np and len(array.addressable_shards) > 0: + is_owner = True + elif not is_np and hasattr(array, "sharding") and hasattr(array.sharding, "device_set"): + owner_device = list(array.sharding.device_set)[0] + is_owner = (owner_device.process_index == host_id) + else: + is_owner = is_np + + if is_owner: + return onp.asarray(array), is_owner + return onp.zeros(shape, dtype=dtype), is_owner + + +def reshard_to_target(array, sharding, hook_fns=None, target_shape=None, is_owner=None): + """Reshards a local CPU array cross-host explicitly to the target sharding.""" + if hasattr(array, "sharding") and isinstance(array.sharding, jax.sharding.NamedSharding): + # For fully formed NamedSharding inputs, we process the hooks normally here. + if hook_fns is not None and target_shape is not None: + array = apply_hook_fns(array, target_shape, hook_fns) + _reshard = jax.jit(lambda x: x, out_shardings=sharding) + return _reshard(array) + + # Secure purely CPU memory + if is_owner is None: + local_arr, is_owner = get_safe_local_array(array, target_shape if target_shape is not None else array.shape, array.dtype if array is not None else np.float32) + else: + local_arr = onp.asarray(array) + + if hook_fns is not None and target_shape is not None: + local_arr = apply_hook_fns(local_arr, target_shape, hook_fns) + local_arr = onp.asarray(local_arr) + + import math + bytes_per_slice = math.prod(local_arr.shape[1:]) * local_arr.itemsize if len(local_arr.shape) > 1 else local_arr.itemsize + max_bytes_per_chunk = 128 * 1024 * 1024 # 128 MB max buffer per transmission + + if bytes_per_slice >= max_bytes_per_chunk: + chunk_size = 1 + else: + chunk_size = max(1, max_bytes_per_chunk // bytes_per_slice) + + total_len = local_arr.shape[0] + + # Broadcast over network sequentially in mathematically stable chunks + # to rigorously prevent Host OOM Resource Exhausted Limits on the FSDP Grid. + if total_len <= chunk_size or total_len == 0: + global_replicated = jax.experimental.multihost_utils.broadcast_one_to_all( + local_arr, is_source=is_owner + ) + cpu_replicated = onp.asarray(global_replicated) + del global_replicated + else: + replicated_chunks = [] + for i in range(0, total_len, chunk_size): + chunk_arr_cpu = local_arr[i : i + chunk_size] + global_replicated_chunk = jax.experimental.multihost_utils.broadcast_one_to_all( + chunk_arr_cpu, is_source=is_owner + ) + replicated_chunks.append(onp.asarray(global_replicated_chunk)) + del global_replicated_chunk + + cpu_replicated = onp.concatenate(replicated_chunks, axis=0) + del replicated_chunks + + del local_arr + + # jax.make_array_from_callback directly pulls the memory slices mathematically matched + # by the FSDP Sharding layout and instantiates them purely into device memory! + res = jax.make_array_from_callback( + cpu_replicated.shape, + sharding, + lambda index: cpu_replicated[index] + ) + return res + + +def unpack_getter_output(val): + """Safely unpacks the (array, shape, dtype) tuple or defaults to None for mock arrays.""" + if isinstance(val, tuple) and len(val) == 3: + return val + return val, None, None + +def unpack_getter_output(val): + """Safely unpacks the (array, shape, dtype) tuple or defaults to None for mock arrays.""" + if isinstance(val, tuple) and len(val) == 3: + return val + return val, None, None + + def _build_multi_axis_stacked_tensor( hf_source_keys: List[List[str]], tensor_getter_fn: Callable[[str], np.ndarray], @@ -55,7 +154,7 @@ def _build_multi_axis_stacked_tensor( target_leaf: Any, config, ) -> np.ndarray: - """Builds a MaxText tensor by stacking HF weights along two axes (experts and layers) directly in place on device.""" + """Builds a MaxText tensor by stacking HF weights along experts and layers using CPU-TPU hybrid phases.""" if hasattr(target_leaf, "sharding"): target_shape = target_leaf.shape target_sharding = target_leaf.sharding @@ -65,41 +164,52 @@ def _build_multi_axis_stacked_tensor( target_sharding = None target_dtype = target_leaf.dtype if hasattr(target_leaf, "dtype") else np.float32 + # Slice shape is (7168, 2048) or similar, from index 2 onwards mt_slice_shape = target_shape[2:] - # Pre-derive the compatible sharding specs to avoid rank mismatches + # pre-derive layer sharding spec (removes layer dimension at axis 1) if target_sharding is not None and hasattr(target_sharding, "spec"): - # Target shape is (experts, layers, ...) -> slice is from index 2 onwards - spec_list = list(target_sharding.spec)[2:] - slice_sharding = jax.sharding.NamedSharding(target_sharding.mesh, jax.sharding.PartitionSpec(*spec_list)) - # Stacking layer shards - layer_spec_list = list(target_sharding.spec)[1:] - layer_sharding = jax.sharding.NamedSharding(target_sharding.mesh, jax.sharding.PartitionSpec(*layer_spec_list)) + spec_list = list(target_sharding.spec) + del spec_list[1] # Delete layers axis + layer_sharding = jax.sharding.NamedSharding(target_sharding.mesh, jax.sharding.PartitionSpec(*spec_list)) else: - slice_sharding = target_sharding layer_sharding = target_sharding - all_expert_tensors = [] - # Outer loop iterates through experts - for layer_keys_for_expert in hf_source_keys: - layer_tensors_for_expert = [] - # Inner loop iterates through layers for the current expert - for hf_key_single in layer_keys_for_expert: - hf_tensor_numpy = tensor_getter_fn(hf_key_single) - processed_hf_tensor = apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns) - - if target_sharding is not None: - processed_hf_tensor = jax.device_put(processed_hf_tensor, slice_sharding) - layer_tensors_for_expert.append(processed_hf_tensor) - - expert_tensor = _binary_chunked_stack(layer_tensors_for_expert, axis=0) - if target_sharding is not None: - expert_tensor = jax.device_put(expert_tensor, layer_sharding) - all_expert_tensors.append(expert_tensor) - - stacked_array = _binary_chunked_stack(all_expert_tensors, axis=0).astype(target_dtype) - if target_sharding is not None: - stacked_array = jax.device_put(stacked_array, target_sharding) + # Transpose key matrix from [256 experts][61 layers] to [61 layers][256 experts] + # This allows us to CPU-stack all experts for a single layer together (takes only ~5.2 GB CPU RAM) + keys_by_layer = list(zip(*hf_source_keys)) + + expert_tensors_across_layers = [] + + for j, expert_keys_for_layer in enumerate(keys_by_layer): + expert_tensors_numpy = [] + + # Establish ownership dynamically based on the first expert in this layer + sample_tensor, sample_shape, sample_dtype = unpack_getter_output(tensor_getter_fn(expert_keys_for_layer[0])) + _, is_owner = get_safe_local_array(sample_tensor, sample_shape if sample_shape is not None else mt_slice_shape, sample_dtype if sample_dtype is not None else target_dtype) + + # Fetch and format all 256 experts for this single layer on CPU + for hf_key_single in expert_keys_for_layer: + hf_tensor_raw, raw_shape, raw_dtype = unpack_getter_output(tensor_getter_fn(hf_key_single)) + local_arr, _ = get_safe_local_array(hf_tensor_raw, raw_shape if raw_shape is not None else mt_slice_shape, raw_dtype if raw_dtype is not None else target_dtype) + processed = apply_hook_fns(local_arr, mt_slice_shape, hook_fns) + expert_tensors_numpy.append(onp.asarray(processed)) + + # Stack experts locally on Host CPU (Shape: 256, 7168, 2048) + stacked_layer = onp.stack(expert_tensors_numpy, axis=0) + + # Reshard this layer to TPU. Peak HBM overhead per node: (5.2GB / 256) = 20 Megabytes! + if layer_sharding is not None: + sharded_layer = reshard_to_target(stacked_layer, layer_sharding, is_owner=is_owner) + else: + sharded_layer = jax.device_put(stacked_layer, jax.local_devices()[0]) + + # Block CPU thread here to force the TPU sequencer queue to flush + sharded_layer.block_until_ready() + expert_tensors_across_layers.append(sharded_layer) + + # Stack the 61 FSDP-sharded layers on TPU (takes < 0.1s!) + stacked_array = _binary_chunked_stack(expert_tensors_across_layers, axis=1).astype(target_dtype) return stacked_array @@ -120,15 +230,8 @@ def _build_single_axis_stacked_tensor( target_sharding = None 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 - else: - # Otherwise, if an unscanned MoE layer, and we stack along the expert axis (0). - axis_to_stack = 0 + axis_to_stack = config.param_scan_axis if config.scan_layers else 0 - # The hook function needs the shape of an individual slice, not the full stacked tensor. - # We calculate it by removing the stacking dimension from the final target shape. mt_slice_shape_list = list(target_shape) del mt_slice_shape_list[axis_to_stack] mt_slice_shape = tuple(mt_slice_shape_list) @@ -141,17 +244,21 @@ def _build_single_axis_stacked_tensor( slice_sharding = target_sharding tensors_to_stack = [] + + # Reshard each layer slice sequentially. Since there are only 61 layers, this is extremely fast. for hf_key_single in hf_source_keys: - hf_tensor_numpy = tensor_getter_fn(hf_key_single) - processed_hf_tensor = apply_hook_fns(hf_tensor_numpy, mt_slice_shape, hook_fns) - - if target_sharding is not None: - processed_hf_tensor = jax.device_put(processed_hf_tensor, slice_sharding) - tensors_to_stack.append(processed_hf_tensor) + hf_tensor_raw, raw_shape, raw_dtype = unpack_getter_output(tensor_getter_fn(hf_key_single)) + local_arr, is_owner = get_safe_local_array(hf_tensor_raw, raw_shape if raw_shape is not None else mt_slice_shape, raw_dtype if raw_dtype is not None else target_dtype) + if slice_sharding is not None: + processed = reshard_to_target(local_arr, slice_sharding, hook_fns=hook_fns, target_shape=mt_slice_shape, is_owner=is_owner) + else: + processed = apply_hook_fns(local_arr, mt_slice_shape, hook_fns) + + # Block CPU thread here to force the TPU sequencer queue to flush + processed.block_until_ready() + tensors_to_stack.append(processed) stacked_array = _binary_chunked_stack(tensors_to_stack, axis=axis_to_stack).astype(target_dtype) - if target_sharding is not None: - stacked_array = jax.device_put(stacked_array, target_sharding) return stacked_array @@ -161,10 +268,16 @@ def get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_ta # Case 1: Single hf key (str) def _loader(getter, key, leaf, hook): if hasattr(leaf, "sharding"): - array = apply_hook_fns(getter(key), leaf.shape, hook) - return jax.device_put(array, device=leaf.sharding) + t_jax, raw_shape, raw_dtype = unpack_getter_output(getter(key)) + local_arr, is_owner = get_safe_local_array(t_jax, raw_shape if raw_shape is not None else leaf.shape, raw_dtype if raw_dtype is not None else (leaf.dtype if hasattr(leaf, "dtype") else np.float32)) + res = reshard_to_target(local_arr, leaf.sharding, hook_fns=hook, target_shape=leaf.shape, is_owner=is_owner) + res.block_until_ready() + return res else: - return apply_hook_fns(getter(key), leaf, hook) + # Fallback Local Path + t_jax, raw_shape, raw_dtype = unpack_getter_output(getter(key)) + local_arr, _ = get_safe_local_array(t_jax, raw_shape if raw_shape is not None else leaf.shape, raw_dtype if raw_dtype is not None else (leaf.dtype if hasattr(leaf, "dtype") else np.float32)) + return apply_hook_fns(local_arr, leaf.shape, hook) return partial( _loader, diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index 2907f4f314..e99dd855a9 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -140,6 +140,16 @@ def getter_fn(key): np.testing.assert_allclose(result[1, 0], tensors["expert_1.layer_0.weight"]) np.testing.assert_allclose(result[1, 1], tensors["expert_1.layer_1.weight"]) + def test_reshard_to_target_single_device_sharding(self): + from maxtext.checkpoint_conversion.utils.tensor_handling import reshard_to_target + arr = jax.device_put(np.ones((4, 4), dtype=np.float32), jax.devices()[0]) + # Target sharding + target_sharding = jax.sharding.NamedSharding(self.mesh, jax.sharding.PartitionSpec("x", None)) + # Reshard to target + result = reshard_to_target(arr, target_sharding) + self.assertEqual(result.shape, (4, 4)) + np.testing.assert_allclose(result, 1.0) + class LoadDynamicTest(parameterized.TestCase): """Tests for cache downloads and dynamic loading of safetensors.""" @@ -235,7 +245,7 @@ def __init__(self): dummy_ret_val, loaded_vars = load_dynamic.load_safetensors_dynamic_state(path, abstract_state, config) self.assertIsNone(dummy_ret_val) - self.assertEqual(loaded_vars, {"params": {}}) + self.assertEqual(loaded_vars, {}) mock_hf_fs.assert_called_once_with(token="dummy_token") mock_sync.assert_called_once_with("dynamic_hf_download_complete") @@ -301,7 +311,7 @@ def __init__(self): self.assertIsNotNone(loaded_vars) # Assert values match - loaded_weight = loaded_vars["params"]["token_embedder"]["embedding"] + loaded_weight = loaded_vars["token_embedder"]["embedding"] np.testing.assert_allclose(loaded_weight, dummy_weight)