From d8b32cb705f5d9845d8c3e22747ff13f16802086 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 | 47 ++++++++++--- .../utils/tensor_handling.py | 68 +++++++++++++++++-- tests/unit/checkpointing_test.py | 14 +++- 3 files changed, 110 insertions(+), 19 deletions(-) diff --git a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py index 2032cf1c41..a07e5fbd9b 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 @@ -196,9 +197,13 @@ def transform_hf_state_to_mt_state(hf_state, target_tree, param_map_mt_to_hf, ho def tensor_getter(key): return hf_state.pop(key) - 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 +211,20 @@ 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 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 +235,18 @@ def tensor_getter(key): maxtext_config, ) - # Execute transformation and assign to flat_restored t_layer = time.time() - flat_restored[check_name] = load_fn() + flat_restored[target_path] = load_fn() max_logging.log(f"Transformed {check_name} from {hf_source} in {time.time() - t_layer:.4f}s") 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"] @@ -368,4 +373,24 @@ def load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_con 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..30ebb1d7ed 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,61 @@ def _binary_chunked_stack(tensors: List[np.ndarray], axis: int) -> np.ndarray: return np.concatenate([left, right], axis=axis) +def reshard_to_target(array, sharding): + """Reshards a local SingleDevice array cross-host to the target sharding.""" + if hasattr(array, "sharding"): + print("[reshard_to_target] Checking if target mesh is compatible...", flush=True) + array_mesh = getattr(array.sharding, "mesh", None) + target_mesh = getattr(sharding, "mesh", None) + is_compatible_named = ( + array_mesh is not None + and target_mesh is not None + and onp.array_equal(array_mesh.devices, target_mesh.devices) + ) + if is_compatible_named: + print("[reshard_to_target] Target mesh compatible. Directly returning JIT...", flush=True) + _reshard = jax.jit(lambda x: x, out_shardings=sharding) + return _reshard(array) + + host_id = jax.process_index() + print(f"[reshard_to_target] Target mesh NOT strictly compatible. Fallback path initiated. Process {host_id} identifying ownership...", flush=True) + + is_np = not hasattr(array, "addressable_shards") + 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"): + # Target owner device process index via device property to avoid vacuous length checks + owner_device = list(array.sharding.device_set)[0] + is_owner = (owner_device.process_index == host_id) + else: + is_owner = is_np + + print(f"[reshard_to_target] Process {host_id} is_owner: {is_owner}. Instantiating local buffer...", flush=True) + + if is_owner: + local_arr = onp.asarray(array) + else: + local_arr = onp.zeros(array.shape, dtype=array.dtype) + + print(f"[reshard_to_target] Process {host_id} instantiated shape {local_arr.shape}. Broadcasting via multihost_utils.broadcast_one_to_all...", flush=True) + + # Broadcast the array from the owner to all other hosts. + # Since safetensors_layout now processes resharding sequentially, this + # will not encounter cross-host transfer key collisions. + t_bcast = __import__('time').time() + global_replicated = jax.experimental.multihost_utils.broadcast_one_to_all( + local_arr, is_source=is_owner + ) + print(f"[reshard_to_target] Process {host_id} broadcast finished in {__import__('time').time() - t_bcast:.2f}s. Mapping to TPU array via jax.device_put...", flush=True) + + # global_replicated is a Host RAM Numpy Array present on all hosts. + # We use jax.device_put to slice this local RAM array onto the specific + # TPUs connected to each host governed by the global FSDP sharding. + res = jax.device_put(global_replicated, sharding) + print(f"[reshard_to_target] Process {host_id} device_put complete.", flush=True) + return res + + def _build_multi_axis_stacked_tensor( hf_source_keys: List[List[str]], tensor_getter_fn: Callable[[str], np.ndarray], @@ -89,17 +145,17 @@ def _build_multi_axis_stacked_tensor( 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) + processed_hf_tensor = reshard_to_target(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) + expert_tensor = reshard_to_target(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) + stacked_array = reshard_to_target(stacked_array, target_sharding) return stacked_array @@ -146,12 +202,12 @@ def _build_single_axis_stacked_tensor( 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) + processed_hf_tensor = reshard_to_target(processed_hf_tensor, slice_sharding) tensors_to_stack.append(processed_hf_tensor) 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) + stacked_array = reshard_to_target(stacked_array, target_sharding) return stacked_array @@ -162,7 +218,7 @@ def get_hf_loading_function(hf_source_keys_or_key, tensor_getter, hook_fn, mt_ta 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) + return reshard_to_target(array, leaf.sharding) else: return apply_hook_fns(getter(key), leaf, hook) 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)