diff --git a/examples/frozenlake/configs/gemma4_e2b.yaml b/examples/frozenlake/configs/gemma4_e2b.yaml new file mode 100644 index 000000000..9c6129ac7 --- /dev/null +++ b/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 diff --git a/examples/frozenlake/data.py b/examples/frozenlake/data.py index 3997c2343..9e41b18c1 100644 --- a/examples/frozenlake/data.py +++ b/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. diff --git a/examples/frozenlake/run_gemma4_e2b.sh b/examples/frozenlake/run_gemma4_e2b.sh new file mode 100644 index 000000000..62046f497 --- /dev/null +++ b/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 \ + "$@" diff --git a/examples/frozenlake/train_frozenlake.py b/examples/frozenlake/train_frozenlake.py index 4f4282f03..813b78a5b 100644 --- a/examples/frozenlake/train_frozenlake.py +++ b/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, ) diff --git a/tunix/cli/config.py b/tunix/cli/config.py index 5b0aeb631..ef57d053e 100644 --- a/tunix/cli/config.py +++ b/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.""" diff --git a/tunix/cli/grpo_main.py b/tunix/cli/grpo_main.py index 3e9a50e01..f95c6d4b0 100644 --- a/tunix/cli/grpo_main.py +++ b/tunix/cli/grpo_main.py @@ -187,6 +187,7 @@ def resolve_owner( else: role_to_owner[role] = resolve_owner(role, set()) continue + role_to_owner[role] = resolve_owner(role, set()) return role_to_owner @@ -251,6 +252,7 @@ def create_role_to_mesh(self): ) return {role: owner_to_mesh[owner] for role, owner in role_to_owner.items()} + # ------------------------------------------------------------------ # Rollout config # ------------------------------------------------------------------ @@ -395,10 +397,11 @@ def _rollout_engine_extra( rollout_vllm_server_mode_submission_threshold=submission_threshold, rollout_vllm_server_mode_submission_timeout_s=submission_timeout_s, rollout_vllm_async_scheduling=vllm.get("async_scheduling", True), - tensor_parallel_size=( - rollout_shape[1] if len(rollout_shape) > 1 else 1 + tensor_parallel_size=vllm.get( + "tensor_parallel_size", + rollout_shape[1] if len(rollout_shape) > 1 else 1, ), - data_parallel_size=rollout_shape[0], + data_parallel_size=vllm.get("data_parallel_size", rollout_shape[0]), rollout_vllm_max_num_seqs=max_num_seqs, rollout_vllm_max_num_batched_tokens=max_batched_tokens, rollout_vllm_kwargs=vllm.get( @@ -521,10 +524,13 @@ def create_rl_cluster(self, tokenizer): ) else: graph_def, params = nnx.split(reference_model) - actor_model = nnx.merge( - graph_def, - jax.tree.map(jnp.copy, params), - ) + actor_dtype_str = actor_model_config.get("load_dtype") + if actor_dtype_str: + actor_dtype = getattr(jnp, actor_dtype_str) + params = jax.tree.map(lambda x: x.astype(actor_dtype), params) + else: + params = jax.tree.map(jnp.copy, params) + actor_model = nnx.merge(graph_def, params) cluster_config = self.create_cluster_config( role_to_mesh=role_to_mesh, @@ -707,11 +713,15 @@ def _create_chat_parser(self, tokenizer: Any) -> Any: """Instantiate a chat parser based on chat_parser_config.type.""" from tunix.rl.agentic.parser.chat_template_parser import parser as chat_parser_lib # pylint: disable=g-import-not-at-top - parser_type = self._config_mapping("chat_parser_config").get( - "type", "default" - ) + chat_cfg = self._config_mapping("chat_parser_config") + parser_type = chat_cfg.get("type", "default") if parser_type == "qwen": return chat_parser_lib.QwenChatTemplateParser(tokenizer) + if parser_type == "gemma4": + return chat_parser_lib.Gemma4ChatTemplateParser( + tokenizer, + enable_thinking=chat_cfg.get("enable_thinking", False) + ) return chat_parser_lib.DefaultChatTemplateParser(tokenizer) def _load_class_from_path(self, dotted_path: str) -> type[Any]: @@ -888,4 +898,4 @@ def main(argv, **kwargs): if __name__ == "__main__": - app.run(main) + app.run(main) \ No newline at end of file diff --git a/tunix/cli/utils/model.py b/tunix/cli/utils/model.py index 38749257c..bc55f4c5f 100644 --- a/tunix/cli/utils/model.py +++ b/tunix/cli/utils/model.py @@ -196,6 +196,9 @@ def create_model( 'flash_attention_block_size', 1024 ), remat_config=model_config.get('remat_config', 1), + dtype=model_config.get('dtype'), + load_dtype=model_config.get('load_dtype'), + use_sliding_window_kv_cache=model_config.get('use_sliding_window_kv_cache'), ) if model_config.get('lora_config'): diff --git a/tunix/models/automodel.py b/tunix/models/automodel.py index 11c8cf405..7c9541716 100644 --- a/tunix/models/automodel.py +++ b/tunix/models/automodel.py @@ -568,11 +568,22 @@ def from_pretrained( # pick corresponding config based on model version model_params = call_model_config(naming_info.model_name) + # Get load_dtype explicitly from kwargs + load_dtype_str = kwargs.get('load_dtype') + load_dtype = getattr(jnp, load_dtype_str) if load_dtype_str and isinstance(load_dtype_str, str) else load_dtype_str + # Apply any model config field overrides passed via kwargs (e.g. # use_flash_attention, flash_attention_block_size). if dataclasses.is_dataclass(model_params): valid_fields = {f.name for f in dataclasses.fields(model_params)} - overrides = {k: v for k, v in kwargs.items() if k in valid_fields} + overrides = {k: v for k, v in kwargs.items() if k in valid_fields and v is not None} + if 'remat_config' in overrides and isinstance(overrides['remat_config'], str): + model_module = get_model_module(naming_info.model_name, ModelModule.MODEL) + if hasattr(model_module, 'RematConfig'): + overrides['remat_config'] = getattr(model_module.RematConfig, overrides['remat_config']) + if 'dtype' in overrides and isinstance(overrides['dtype'], str): + overrides['dtype'] = getattr(jnp, overrides['dtype']) + if overrides: logging.info('Applying model config overrides: %s', overrides) model_params = dataclasses.replace(model_params, **overrides) @@ -583,6 +594,7 @@ def from_pretrained( resolved_model_path, model_params, mesh, + dtype=load_dtype, ) return model, resolved_model_path diff --git a/tunix/models/gemma4/model.py b/tunix/models/gemma4/model.py index d37d1e770..203708c96 100644 --- a/tunix/models/gemma4/model.py +++ b/tunix/models/gemma4/model.py @@ -207,6 +207,13 @@ def gemma4_e2b( vision_encoder=vision.VisionEncoderConfig(use_clipped_linears=True), ) + @classmethod + def gemma4_e2b_it( + cls, + sharding_config: ShardingConfig = ShardingConfig.get_default_sharding(), + ) -> 'ModelConfig': + return cls.gemma4_e2b(sharding_config=sharding_config) + @classmethod def gemma4_e4b( cls, @@ -1716,4 +1723,4 @@ def get_model_input(self): @property def num_embed(self) -> int: - return self.config.num_embed + return self.config.num_embed \ No newline at end of file