import importlib.util, os, shutil as _sh
if importlib.util.find_spec("tunix") is None:
   print("Installing Tunix + JAX ecosystem — this takes ~5-8 min…")
   %pip install -q ipywidgets tensorboardX transformers grain nest_asyncio
   %pip install -q datasets huggingface_hub "numpy>2"
   %pip install -q tensorflow tensorflow_datasets
   %pip install -q git+https://github.com/jax-ml/jax
   %pip install -q git+https://github.com/google/tunix
   %pip install -q git+https://github.com/google/qwix
   %pip uninstall -q flax -y
   %pip install -q git+https://github.com/google/flax
   %pip uninstall -q wandb -y
   print("\n\n✅ Install done. The runtime will RESTART now.")
   print("👉  After it restarts, RUN THIS CELL AGAIN to start training.\n")
   os.kill(os.getpid(), 9)
import getpass, functools, json, re
import numpy as np
os.environ["WANDB_MODE"] = "disabled"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
   try:
       from google.colab import userdata
       HF_TOKEN = userdata.get("HF_TOKEN")
   except Exception:
       HF_TOKEN = getpass.getpass("Hugging Face token (needs Gemma license access): ")
os.environ["HF_TOKEN"] = HF_TOKEN or ""
import nest_asyncio; nest_asyncio.apply()
import tensorflow as tf; tf.config.set_visible_devices([], "GPU")
import jax, jax.numpy as jnp, optax, grain, qwix
from flax import nnx
from orbax import checkpoint as ocp
from huggingface_hub import snapshot_download, login
from datasets import load_dataset
from tunix.generate import sampler as sampler_lib
from tunix.generate import tokenizer_adapter as tokenizer_lib
from tunix.models.gemma3 import model as gemma_lib
from tunix.models.gemma3 import params_safetensors as params_safetensors_lib
from tunix.models.gemma3 import params as gemma_params
from tunix.rl import rl_cluster as rl_cluster_lib
from tunix.rl.grpo.grpo_learner import GRPOConfig, GRPOLearner
from tunix.rl.rollout import base_rollout
from tunix.sft import metrics_logger
if HF_TOKEN:
   login(token=HF_TOKEN)
devices = jax.devices()
print("JAX backend :", jax.default_backend(), "|", len(devices), "device(s):", devices)
IS_GPU = _sh.which("nvidia-smi") is not None
if IS_GPU and jax.default_backend() == "cpu":
   print("⚠️  GPU runtime but JAX sees CPU only. Re-run `pip install -U \"jax[cuda12]\"` "
         "and restart, or switch to a TPU runtime.")
MODEL_ID = "google/gemma-3-1b-it"
RANK, ALPHA = 32, 32.0
MAX_PROMPT_LENGTH      = 256
TOTAL_GENERATION_STEPS = 512
NUM_GENERATIONS        = 2
NUM_ITERATIONS         = 1
BETA                   = 0.08
EPSILON                = 0.2
TEMPERATURE, TOP_P, TOP_K = 0.9, 1.0, 50
MAX_STEPS   = 100
TRAIN_LIMIT = MAX_STEPS
NUM_TEST    = 16
LEARNING_RATE, B1, B2, WEIGHT_DECAY = 3e-6, 0.9, 0.99, 0.1
WARMUP_STEPS, MAX_GRAD_NORM = int(0.1 * MAX_STEPS), 0.1
CKPT_DIR, TB_DIR = "/content/ckpts/", "/content/tb/grpo"
N = jax.device_count()
MESH = [(N, 1), ("fsdp", "tp")]
mesh = jax.make_mesh(*MESH, axis_types=(jax.sharding.AxisType.Auto,) * 2)