Add cli for frozenlake gemma4-e2b-it recipe

examples/frozenlake/configs/gemma4_e2b.yaml

@@ -0,0 +1,152 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+model_config:
+  model_name: "gemma-4-e2b"
+  model_id: "google/gemma-4-E2B-it"
+  model_source: "huggingface"
+  model_path: ""
+  mesh:
+    shape: "(8,1)"
+    axis_names: "('fsdp','tp')"
+  remat_config: "DECODER"
+  use_flash_attention: True
+  flash_attention_block_size: 256
+  use_sliding_window_kv_cache: False
+  dtype: "bfloat16"
+  rng_seed: 42
+actor_model_config:
+  load_dtype: "float32"
+  mesh:
+    shape: "(8,1)"
+    axis_names: "('fsdp','tp')"
+reference_model_config:
+  mesh: null
+  same_mesh_as: "actor"
+rollout_model_config:
+  mesh: null
+  same_mesh_as: "actor"
+tokenizer_config:
+  tokenizer_type: "huggingface"
+  tokenizer_path: "google/gemma-4-E2B-it"
+  add_bos: True
+
+# Agentic GRPO configuration
+training_mode: "agentic_grpo"
+agent_class_path: "examples.frozenlake.agent.FrozenLakeAgent"
+agent_kwargs:
+  use_multistep_prompt: True
+env_class_path: "examples.frozenlake.env.FrozenLakeEnv"
+env_kwargs:
+  max_steps: 8
+
+
+data_module: "examples.frozenlake.data"
+data_config:
+  split: "train"
+  shuffle: True
+  seed: 42
+  data_dir: "/tmp/data/frozenlake"
+
+chat_parser_config:
+  type: "gemma4"
+  enable_thinking: False
+
+batch_size: 64
+num_batches: 5
+num_test_batches: 2
+num_train_epochs: 1
+
+rl_training_config:
+  mini_batch_size: 64
+  train_micro_batch_size: 2
+  actor_optimizer_config:
+    opt_type: "adamw"
+    peak_value: 1e-6
+    schedule_type: "warmup_cosine_decay_schedule"
+    init_value: 0.0
+    end_value: 0.0
+    warmup_ratio: 0.1
+    warmup_steps: 0
+    decay_steps: 9
+    b1: 0.9
+    b2: 0.95
+    weight_decay: 0.0
+    max_grad_norm: 100.0
+  eval_every_n_steps: 10
+  max_steps: 5
+  metrics_logging_options:
+    log_dir: "/tmp/logging"
+    flush_every_n_steps: 1
+  checkpointing_options:
+    save_interval_steps: 50
+    max_to_keep: 4
+  profiler_options: {}
+  compute_logps_chunk_size: 2048
+
+rollout_config:
+  total_generation_steps: 2048
+  max_prompt_length: 2048
+  temperature: 0.7
+  top_p: 1.0
+  top_k: 0
+  return_logprobs: True
+  rollout_vllm_init_with_random_weights: True
+  rollout_vllm_sampling_kwargs:
+    skip_special_tokens: False
+
+rollout_engine: "vllm"
+
+vllm_config:
+  model_version: "google/gemma-4-E2B-it"
+  hbm_utilization: 0.2
+  tpu_backend_type: "jax"
+  server_mode: True
+  async_scheduling: False
+  max_num_seqs: 64
+  max_num_batched_tokens: 16384
+  data_parallel_size: 1
+  tensor_parallel_size: 8
+  kwargs:
+    kv_cache_metrics: True
+    disable_log_stats: False
+    enable_prefix_caching: False
+    dtype: "bfloat16"
+    limit_mm_per_prompt:
+      image: 0
+      video: 0
+      audio: 0
+    hf_overrides:
+      final_logit_softcapping: 30.0
+      text_config:
+        final_logit_softcapping: 30.0
+
+offload_to_cpu: False
+
+reward_functions: []
+
+agentic_grpo_config:
+  num_generations: 8
+  num_iterations: 1
+  beta: 0.0
+  epsilon: 0.003
+  epsilon_high: 0.005
+  loss_algo: "gspo-token"
+  loss_agg_mode: "sequence-mean-token-mean"
+  kl_loss_mode: "low_var_kl"
+  advantage_estimator: "rloo"
+  sampler_is: "token"
+  sampler_is_threshold: 2.0
+  degenerate_group_masking: False
+  max_concurrency: 512

examples/frozenlake/data.py

@@ -21,6 +21,8 @@
 
 import numpy as np
 import pandas as pd
+import datasets as datasets_lib
+import grain
 
 
 DEFAULT_DIR = os.getcwd()
@@ -84,6 +86,43 @@ def save_dataset(data: list[dict], filepath: str) -> None:
   print(f"Saved {len(data)} entries to {filepath}")
 
 
+def create_dataset(
+    split: str = "train",
+    data_dir: str = "/tmp/data/frozenlake",
+    seed: int = 42,
+    train_size: int = 10000,
+    test_size: int = 100,
+    **kwargs
+) -> grain.MapDataset:
+  """Lively generates the dataset, saves it to a local directory, and returns a MapDataset.
+
+  This function acts as the entry point for the Tunix GRPO pipeline.
+  """
+  os.makedirs(data_dir, exist_ok=True)
+  filepath = os.path.join(data_dir, f"{split}.parquet")
+
+  if not os.path.exists(filepath):
+    size = train_size if split == "train" else test_size
+    print(f"Dynamically generating {size} instances for '{split}' split...")
+    seeds, sizes, ps = generate_dataset_parameters(size, random_seed=seed)
+    data = [
+        get_frozenlake_dict(s, sizes[idx], ps[idx])
+        for idx, s in enumerate(seeds)
+    ]
+    save_dataset(data, filepath)
+  else:
+    print(f"Loading existing dataset from {filepath}")
+
+  df = pd.read_parquet(filepath)
+  hf_ds = datasets_lib.Dataset.from_pandas(df)
+
+  def process_item(item):
+    item["prompts"] = ""
+    return item
+
+  return grain.MapDataset.source(hf_ds).map(process_item)
+
+
 def main():
   """Main function to generate and save FrozenLake datasets.
 

examples/frozenlake/run_gemma4_e2b.sh

@@ -0,0 +1,33 @@
+#!/bin/bash
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Install dependencies:
+pip install gymnasium
+
+set -x # Enable xtrace
+
+batch_size=${batch_size:-64}
+num_batches=${num_batches:-5}
+
+echo "Using parameters:"
+echo "  Batch Size: $batch_size"
+echo "  Num Batches: $num_batches"
+
+python3 -m tunix.cli.grpo_main \
+  tunix/cli/base_agentic_config.yaml \
+  override_config_file=examples/frozenlake/configs/gemma4_e2b.yaml \
+  batch_size=$batch_size \
+  num_batches=$num_batches \
+  "$@"

examples/frozenlake/train_frozenlake.py

@@ -1,6 +1,6 @@
 """Agentic FrozenLake GRPO recipe for Gemma4-2B on a single TPU host.
 
-Designed for v5p-8 / v6e-4 -class hosts where actor, reference, and rollout
+Designed for v5p-4 / v6e-8 -class hosts where actor, reference, and rollout
 share a single mesh. Configuration is env-driven so the same image runs
 unchanged on any spot VM:
 
@@ -191,7 +191,7 @@
 # some headroom without provisioning a huge unused KV-cache pool — on a
 # shared trainer+rollout mesh that KV-cache pool consumes HBM that the
 # trainer needs at peak (logits + activations + optimizer state).
-VLLM_MAX_NUM_SEQS = 64
+VLLM_MAX_NUM_SEQS = 32
 VLLM_MAX_BATCHED_TOKENS = VLLM_MAX_NUM_SEQS * 2 * 1024 // 8
 
 NUM_ITERATIONS = 1
@@ -369,7 +369,7 @@ def process_item(item):
 
 config = model_lib.ModelConfig.gemma4_e2b()
 if ENABLE_REMAT:
-  config.remat_config = model_lib.RematConfig.BLOCK
+  config.remat_config = model_lib.RematConfig.DECODER
 if ENABLE_FLASH_ATTENTION:
   config.use_flash_attention = True
   config.flash_attention_block_size = 256
@@ -453,7 +453,7 @@ def process_item(item):
     # max_seq_len rather than the vLLM default. Once vLLM-TPU gains support
     # for sleep/wake_up, this can be relaxed since the KV pool can be
     # offloaded to host RAM during train_step.
-    "rollout_vllm_hbm_utilization": 0.28,
+    "rollout_vllm_hbm_utilization": 0.2,
     "rollout_vllm_tpu_backend_type": "jax",
     "rollout_vllm_server_mode": True,
     # Async scheduling adds an extra in-flight step that can race weight sync;
@@ -529,11 +529,12 @@ def process_item(item):
         # invokes the trainer ``mini_batch_size // train_micro_batch_size``
         # times, so the optimizer still sees a ``mini_batch_size`` gradient
         # per update.
-        train_micro_batch_size=1,
-        compute_logps_micro_batch_size=1,
+        train_micro_batch_size=2,
+        compute_logps_micro_batch_size=2,
         metrics_logging_options=metrics_logging_options,
         checkpoint_root_directory=CKPT_DIR,
         checkpointing_options=checkpointing_options,
+        compute_logps_chunk_size=2048,
     ),
     rollout_config=rollout_engine_config,
 )

tunix/cli/config.py

@@ -279,6 +279,8 @@ def _config_mapping(self, key: str) -> dict[str, Any]:
           f"Expected config section {key!r} to be a mapping, got"
           f" {type(value).__name__}."
       )
+    if isinstance(value, omegaconf.DictConfig):
+      return omegaconf.OmegaConf.to_container(value, resolve=True)
     return dict(value)
 
   def _mutable_config_mapping(self, key: str) -> MutableMapping[str, Any]:
@@ -475,36 +477,25 @@ def check_supported_workflow(self) -> None:
     supported_sources["gemma"] = ["kaggle", "internal", "maxtext"]
     supported_sources["gemma2"] = ["kaggle", "internal", "maxtext"]
     supported_sources["gemma3"] = ["gcs", "internal", "maxtext"]
-
-    if model_name.startswith("gemma3") or model_name.startswith("gemma-3"):
-      expected_sources = supported_sources["gemma3"]
-      if model_source not in expected_sources:
-        raise ValueError(
-            f"Model '{model_name}' must use source(s) {expected_sources}, but"
-            f" got '{model_source}'."
-        )
-    elif model_name.startswith("gemma2") or model_name.startswith("gemma-2"):
-      expected_sources = supported_sources["gemma2"]
-      if model_source not in expected_sources:
-        raise ValueError(
-            f"Model '{model_name}' must use source(s) {expected_sources}, but"
-            f" got '{model_source}'."
-        )
+    supported_sources["gemma4"] = ["kaggle", "huggingface"]
+
+    if model_name.startswith(("gemma4", "gemma-4")):
+      model_family = "gemma4"
+    elif model_name.startswith(("gemma3", "gemma-3")):
+      model_family = "gemma3"
+    elif model_name.startswith(("gemma2", "gemma-2")):
+      model_family = "gemma2"
     elif model_name.startswith("gemma"):
-      expected_sources = supported_sources["gemma"]
-      if model_source not in expected_sources:
-        raise ValueError(
-            f"Model '{model_name}' must use source(s) {expected_sources}, but"
-            f" got '{model_source}'."
-        )
+      model_family = "gemma"
     else:
-      # Default case for other models
-      expected_sources = supported_sources["other"]
-      if model_source not in expected_sources:
-        raise ValueError(
-            f"Model '{model_name}' must use source(s) {expected_sources}, but"
-            f" got '{model_source}'."
-        )
+      model_family = "other"
+
+    expected_sources = supported_sources[model_family]
+    if model_source not in expected_sources:
+      raise ValueError(
+          f"Model '{model_name}' must use source(s) {expected_sources}, but"
+          f" got '{model_source}'."
+      )
 
   def _get_nested_config(self, keys: Sequence[str]) -> Any:
     """Helper to retrieve a value from a nested dictionary."""