Fix sft_llama3_demo_tpu.ipynb

Author: hsuan-lun-chiang

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -895,7 +895,7 @@ def extract_nnx_weights(weights_dict: dict) -> dict[str, np.ndarray]:
         path_keys.append(val_str)
 
     # Skip NNX RNG state variables (not model weights)
-    if "to_nnx__rngs" in path_keys or any(k.endswith("_rngs") for k in path_keys):
+    if "to_nnx__rngs" in path_keys or any(k == "rngs" or k.endswith("_rngs") for k in path_keys):
       continue
     # Skip if this is the "value" key itself - we want the parent path
     if path_keys[-1] == "value":

src/maxtext/trainers/diloco/diloco.py

@@ -28,7 +28,6 @@
 import drjax
 from flax import nnx
 from flax import struct
-from flax.training import train_state
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Int32, Key, PyTree, UInt32

src/maxtext/utils/maxtext_utils.py

@@ -1667,10 +1667,15 @@ def _make_named_sharding(v):
         rules = composite_rules(context_rules, local_rules)
         raw_sharding = from_sharding_rules(out_sharding, rules)
         mesh_axis_names = mesh.axis_names if mesh is not None else ()
-        sanitized_sharding = [
-            x if (x is None or (isinstance(x, str) and x in mesh_axis_names) or isinstance(x, tuple)) else None
-            for x in raw_sharding
-        ]
+
+        def _sanitize(x):
+          if isinstance(x, list):
+            x = tuple(x)
+          if x is None or (isinstance(x, str) and x in mesh_axis_names) or isinstance(x, tuple):
+            return x
+          return None
+
+        sanitized_sharding = [_sanitize(x) for x in raw_sharding]
         pspec = PartitionSpec(*sanitized_sharding)
       else:
         pspec = PartitionSpec(*out_sharding)