fix: resolve weight sync and checkpointing issues in RL training

Author: mesakhcienet

src/maxtext/integration/vllm/maxtext_vllm_rollout.py

@@ -51,21 +51,38 @@ def _create_model_converter(model_name: str, config: Any, mesh: jax.sharding.Mes
   return None
 
 
+def _find_scanned_layer_idx(key_tuple, container_names=("layers", "scanned_blocks", "layers_remainder")):
+  """Returns (container_idx, container_name) if a scanned layer structure is found, else (-1, None)."""
+  for name in container_names:
+    for i in range(len(key_tuple) - 1):
+      if key_tuple[i] == name and isinstance(key_tuple[i + 1], str) and key_tuple[i + 1].startswith("layers_"):
+        return i, name
+  return -1, None
+
+
 def unroll_gemma_scanned_weights(weights):
   """Workaround for tunix unstacking bug with Gemma 3/4 scanned blocks.
 
-  tunix fails to map nested layers like `layers.layers_0` to `layers_X`.
-  We manually unroll them here if we detect the structure.
+  tunix fails to map nested layers like `layers.layers_0` to `layers_X`
+  if the target expects integer keys (as in nnx.List).
+  We manually unroll them here, keeping the keys as tuples with integers.
   """
-  if not hasattr(weights, "to_pure_dict"):
+  if hasattr(weights, "to_pure_dict"):
+    pure_dict = weights.to_pure_dict()
+  elif hasattr(weights, "to_dict"):
+    pure_dict = weights.to_dict()
+  elif isinstance(weights, dict):
+    pure_dict = weights
+  else:
     return weights
 
-  flat_w = flatten_dict(weights.to_pure_dict(), sep="/")
+  flat_w = flatten_dict(pure_dict)
   new_flat_w = {}
 
+  logging.debug("MaxTextVllmSampler: First 5 keys in flat_w: %s", list(flat_w.keys())[:5])
+
   # Check if this is actually a scanned Gemma 3/4 checkpoint
-  # by looking for the scanned nested structure.
-  is_gemma_scanned = any("decoder/layers/layers_0/" in k or "decoder/scanned_blocks/layers_0/" in k for k in flat_w)
+  is_gemma_scanned = any(_find_scanned_layer_idx(k)[0] != -1 for k in flat_w)
 
   if not is_gemma_scanned:
     return weights
@@ -76,27 +93,28 @@ def unroll_gemma_scanned_weights(weights):
   pattern_keys = set()
   scan_length = 0
   for k, v in flat_w.items():
-    if "decoder/layers/layers_" in k or "decoder/scanned_blocks/layers_" in k:
-      layer_sub_idx = k.split("layers_")[-1].split("/")[0]
-      pattern_keys.add(int(layer_sub_idx))
-      # In MaxText, Gemma uses param_scan_axis=1, so the scan dimension is at axis 1
-      if hasattr(v, "shape") and len(v.shape) > 1:
+    container_idx, name = _find_scanned_layer_idx(k)
+    if container_idx != -1 and name != "layers_remainder":
+      layer_sub_idx = int(k[container_idx + 1].split("layers_")[1])
+      pattern_keys.add(layer_sub_idx)
+      if hasattr(v, "shape") and len(v.shape) >= 2:
         scan_length = max(scan_length, v.shape[1])
 
   pattern_length = max(pattern_keys) + 1 if pattern_keys else 0
   logging.info("MaxTextVllmSampler: Discovered scan_length=%d, pattern_length=%d", scan_length, pattern_length)
 
   unrolled_count = 0
   for k, v in flat_w.items():
-    if "decoder/layers/layers_" in k or "decoder/scanned_blocks/layers_" in k:
-      # Unstack the array along the 1st axis
-      if "decoder/scanned_blocks/layers_" in k:
-        parts = k.split("decoder/scanned_blocks/layers_")
-      else:
-        parts = k.split("decoder/layers/layers_")
+    if "dropout" in k or "rngs" in k:
+      new_flat_w[k] = v
+      continue
+
+    container_idx, container_name = _find_scanned_layer_idx(k)
 
-      layer_sub_idx = int(parts[1].split("/")[0])
-      suffix = "/" + "/".join(parts[1].split("/")[1:])
+    if container_idx != -1 and container_name in ("layers", "scanned_blocks"):
+      layer_sub_idx = int(k[container_idx + 1].split("layers_")[1])
+      prefix = k[:container_idx] + ("layers",)
+      suffix = k[container_idx + 2 :]
 
       if hasattr(v, "shape") and len(v.shape) > 1:
         v_swapped = jnp.swapaxes(v, 1, 0)
@@ -106,25 +124,29 @@ def unroll_gemma_scanned_weights(weights):
 
       for i in range(scan_length):
         global_idx = i * pattern_length + layer_sub_idx
-        # Map back to nnx.List format which uses layers/X/ instead of layers_X
-        new_flat_w[f"decoder/layers/{global_idx}{suffix}"] = unstacked[i]
+        new_k = prefix + (global_idx,) + suffix
+        new_flat_w[new_k] = unstacked[i]
         unrolled_count += 1
 
-    elif "decoder/layers_remainder/layers_" in k:
-      layer_sub_idx = int(k.split("decoder/layers_remainder/layers_")[1].split("/")[0])
-      suffix = "/" + "/".join(k.split("decoder/layers_remainder/layers_")[1].split("/")[1:])
+    elif container_idx != -1 and container_name == "layers_remainder":
+      layer_sub_idx = int(k[container_idx + 1].split("layers_")[1])
+      prefix = k[:container_idx] + ("layers",)
+      suffix = k[container_idx + 2 :]
 
       global_idx = scan_length * pattern_length + layer_sub_idx
-      new_flat_w[f"decoder/layers/{global_idx}{suffix}"] = v
+      new_k = prefix + (global_idx,) + suffix
+      new_flat_w[new_k] = v
       unrolled_count += 1
     else:
       new_flat_w[k] = v
 
+  assert unrolled_count > 0, "MaxTextVllmSampler: Detected scanned structure, but failed to unroll any layers!"
+
   logging.info(
       "MaxTextVllmSampler: Successfully unrolled %d scanned tensor components into vLLM-compatible nnx.List format.",
       unrolled_count,
   )
-  return unflatten_dict(new_flat_w, sep="/")
+  return unflatten_dict(new_flat_w)
 
 
 class MaxTextVllmSampler(VllmSampler):
@@ -152,12 +174,7 @@ def update_params(
   ):
     """Update the vLLM runner weights from a MaxText state tree."""
     if self._converter is None:
-      # --- Workaround for tunix unstacking bug with Gemma 3/4 scanned blocks ---
-      # tunix fails to map nested layers like `layers.layers_0` to `layers_X`.
-      # We manually unroll them here if we detect the structure.
       updated_weights = unroll_gemma_scanned_weights(updated_weights)
-      # --- End Workaround ---
-
       super().update_params(updated_weights, filter_types)
       return None
 

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -381,7 +381,7 @@ def build_reward_fns(trainer_config: Any, make_reward_fn: Callable) -> list:
   ]
 
 
-def create_rl_components(
+def create_rl_components(  # pylint: disable=too-many-positional-arguments
     trainer_config,
     sampler_config,
     sampler_devices,

tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh

@@ -43,6 +43,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
     model_name=${MODEL_NAME} enable_single_controller=${use_pathways} \
     checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \
     rollout_tensor_parallelism=1 \
+    use_standalone_converter=true \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     vllm_additional_config='{"maxtext_config": {"model_name": "gemma3-4b", "log_config": "false"}}'
 

tests/post_training/unit/vllm_rollout_unroll_test.py

@@ -99,22 +99,24 @@ def test_correctly_unrolls_gemma_scanned_weights(self):
     decoder_dict = unrolled["decoder"]
 
     # Should contain keys 0 to 6 under layers
-    self.assertIn("0", decoder_dict["layers"])
-    self.assertIn("1", decoder_dict["layers"])
-    self.assertIn("2", decoder_dict["layers"])
-    self.assertIn("3", decoder_dict["layers"])
-    self.assertIn("4", decoder_dict["layers"])
-    self.assertIn("5", decoder_dict["layers"])
-    self.assertIn("6", decoder_dict["layers"])
+    self.assertIn(0, decoder_dict["layers"])
+    self.assertIn(1, decoder_dict["layers"])
+    self.assertIn(2, decoder_dict["layers"])
+    self.assertIn(3, decoder_dict["layers"])
+    self.assertIn(4, decoder_dict["layers"])
+    self.assertIn(5, decoder_dict["layers"])
+    self.assertIn(6, decoder_dict["layers"])
+
+    self.assertIsInstance(list(decoder_dict["layers"].keys())[0], int)
 
     # Check that values are correctly sliced
-    np.testing.assert_array_equal(decoder_dict["layers"]["0"]["attn"]["wq"], np.array([[0], [0]]))
-    np.testing.assert_array_equal(decoder_dict["layers"]["1"]["attn"]["wq"], np.array([[1], [1]]))
-    np.testing.assert_array_equal(decoder_dict["layers"]["2"]["attn"]["wq"], np.array([[2], [2]]))
-    np.testing.assert_array_equal(decoder_dict["layers"]["3"]["attn"]["wq"], np.array([[3], [3]]))
-    np.testing.assert_array_equal(decoder_dict["layers"]["4"]["attn"]["wq"], np.array([[4], [4]]))
-    np.testing.assert_array_equal(decoder_dict["layers"]["5"]["attn"]["wq"], np.array([[5], [5]]))
-    np.testing.assert_array_equal(decoder_dict["layers"]["6"]["attn"]["wq"], np.array([[6], [6]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][0]["attn"]["wq"], np.array([[0], [0]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][1]["attn"]["wq"], np.array([[1], [1]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][2]["attn"]["wq"], np.array([[2], [2]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][3]["attn"]["wq"], np.array([[3], [3]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][4]["attn"]["wq"], np.array([[4], [4]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][5]["attn"]["wq"], np.array([[5], [5]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][6]["attn"]["wq"], np.array([[6], [6]]))
 
   @pytest.mark.cpu_only
   def test_correctly_unrolls_gemma3_gemma4_scanned_blocks(self):
@@ -150,7 +152,8 @@ def test_correctly_unrolls_gemma3_gemma4_scanned_blocks(self):
     unrolled = unroll_gemma_scanned_weights(weights)
 
     decoder_dict = unrolled["decoder"]
-    self.assertIn("0", decoder_dict["layers"])
-    self.assertIn("6", decoder_dict["layers"])
-    np.testing.assert_array_equal(decoder_dict["layers"]["0"]["attn"]["wq"], np.array([[0], [0]]))
-    np.testing.assert_array_equal(decoder_dict["layers"]["6"]["attn"]["wq"], np.array([[6], [6]]))
+    self.assertIn(0, decoder_dict["layers"])
+    self.assertIn(6, decoder_dict["layers"])
+    self.assertIsInstance(list(decoder_dict["layers"].keys())[0], int)
+    np.testing.assert_array_equal(decoder_dict["layers"][0]["attn"]["wq"], np.array([[0], [0]]))
+    np.testing.assert_array_equal(decoder_dict["layers"][6]["attn"]["wq"], np.array([[6], [6]]))