Use tuple keys in flax flatten/unflatten for checkpoint conversion.

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,29 +197,36 @@ 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...")
 
   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 +237,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,11 +353,21 @@ 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)
@@ -368,4 +382,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

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,14 +49,96 @@ 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):
+  """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."""
+  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(array.shape, dtype=array.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)
+  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 _build_multi_axis_stacked_tensor(
     hf_source_keys: List[List[str]],
     tensor_getter_fn: Callable[[str], np.ndarray],
     hook_fns: Any,
     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 two axes directly in place on device."""
   if hasattr(target_leaf, "sharding"):
     target_shape = target_leaf.shape
     target_sharding = target_leaf.sharding
@@ -65,41 +148,20 @@ def _build_multi_axis_stacked_tensor(
     target_sharding = None
     target_dtype = target_leaf.dtype if hasattr(target_leaf, "dtype") else np.float32
 
-  mt_slice_shape = target_shape[2:]
-
-  # Pre-derive the compatible sharding specs to avoid rank mismatches
-  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))
-  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)
+      # Retrieve array natively on TPU (SingleDeviceSharding on Host 0)
+      layer_tensors_for_expert.append(tensor_getter_fn(hf_key_single))
 
     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)
+    stacked_array = reshard_to_target(stacked_array, target_sharding, hook_fns=hook_fns, target_shape=target_shape)
   return stacked_array
 
 
@@ -120,38 +182,16 @@ 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
-
-  # 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)
-
-  if target_sharding is not None and hasattr(target_sharding, "spec"):
-    spec_list = list(target_sharding.spec)
-    del spec_list[axis_to_stack]
-    slice_sharding = jax.sharding.NamedSharding(target_sharding.mesh, jax.sharding.PartitionSpec(*spec_list))
-  else:
-    slice_sharding = target_sharding
+  axis_to_stack = config.param_scan_axis if config.scan_layers else 0
 
   tensors_to_stack = []
   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)
+    tensors_to_stack.append(tensor_getter_fn(hf_key_single))
 
   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, hook_fns=hook_fns, target_shape=target_shape)
   return stacked_array
 
 
@@ -161,10 +201,11 @@ 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)
+        return reshard_to_target(getter(key), leaf.sharding, hook_fns=hook, target_shape=leaf.shape)
       else:
-        return apply_hook_fns(getter(key), leaf, hook)
+        # Fallback Local Path
+        local_arr, _ = get_safe_local_array(getter(key))
+        return apply_hook_fns(local_arr, leaf.shape, hook)
 
     return partial(
         _loader,

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)