SkyRL is a scalable and flexible reinforcement learning framework designed to handle the complex memory and compute patterns required for large language model (LLM) alignment. It is optimized for distributed environments, enabling efficient training by decoupling data, model weights, and the training engine.
Image reference: https://docs.skyrl.ai/docs/tinker/architecture
SkyRL provides a backend implementation for the Tinker API using JAX, allowing you to scale LLM post-training on TPUs or GPUs . The integration contains three layers:
-
API Layer: FastAPI HTTP server that accepts Tinker API requests, stores them in a database, and returns future IDs for async polling.
-
Engine Layer: Background subprocess that polls the database, batches pending requests, and dispatches them to the backend.
-
Backend Layer: Translates Tinker operations into training and inference calls, managing Ray workers, JAX, FSDP2/Megatron training, and vLLM inference
See SkyRL architecture documentation for more details.
Why JAX and TPUs?
SkyRL utilizes JAX as its distributed backend for training and sampling during reinforcement learning. Because JAX is hardware-agnostic and offers primitives that scale seamlessly from local development to clusters of thousands of chips, SkyRL can extend model training across both GPUs and TPUs with little to no changes.
Furthermore, SkyRL leverages JAX capabilities such as JIT compilation to achieve high-throughput execution. This significantly boosts performance, particularly on TPUs, which are natively optimized for JAX-based workloads.
What We Are Building
In this tutorial, we will perform multi-host RL training (FSDP and tensor parallelism) on Qwen3-0.6B using SkyRL on GKE and Trillium TPUs.
Prerequisites
- GKE cluster with a 4x4 TPU v6e node pool. To get started, follow Deploy TPU workloads in GKE Standard
- kubectl installed , with the GKE authentication plugin.
Step 1: Configure access to the GKE cluster
Get cluster credentials to access the GKE cluster:
export PROJECT_ID=<YOUR_PROJECT_ID>
export CLUSTER_REGION=<YOUR_REGION>
export CLUSTER_NAME=<YOUR_CLUSTER_NAME>
gcloud container clusters get-credentials $CLUSTER_NAME --region $CLUSTER_REGION --project-id $PROJECT_ID
Step 2: Install the JobSet controller
In this tutorial, we will use the JobSet API to orchestrate parallel execution of SkyRL workers.
First, install the JobSet CRDs and controller:
VERSION=v0.11.1
kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/$VERSION/manifests.yaml
Verify installation by confirming that the JobSet controller is running:
kubectl -n jobset-system get pods
Step 3: Deploy JobSet
In this tutorial, we will create a JobSet resource that will orchestrate and deploy each SkyRL worker on a 4x4 TPU v6e slice. Worker 0 is configured to start the SkyRL Tinker API server and engine.
Create the JobSet manifest:
# jobset.yaml
apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: skyrl-tx-job
spec:
failurePolicy:
maxRestarts: 0
replicatedJobs:
- name: skyrl-workers
replicas: 1
template:
spec:
parallelism: 4
completions: 4
backoffLimit: 0
template:
spec:
affinity:
podAffinity:
requiredDuringSchedulingIgnoredDuringExecution:
- labelSelector:
matchExpressions:
- key: jobset.sigs.k8s.io/jobset-name
operator: In
values:
- skyrl-tx-job
topologyKey: cloud.google.com/gke-nodepool
subdomain: test
restartPolicy: Never
nodeSelector:
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
cloud.google.com/gke-tpu-topology: 4x4
containers:
- name: worker
image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
securityContext:
privileged: false
command:
- bash
- -c
- |
set -ex
git clone https://github.com/NovaSky-AI/SkyRL.git
cd SkyRL
if [ "$JOB_COMPLETION_INDEX" -eq "0" ]; then
uv run --extra tpu --extra tinker --extra jax -m skyrl.tinker.api \
--base-model "Qwen/Qwen3-0.6B" \
--session-timeout-sec 86400 \
--backend-config "{\"train_micro_batch_size\": 8, \"sample_max_num_sequences\": 256, \"tensor_parallel_size\": 4, \"fully_sharded_data_parallel_size\": 4, \"num_processes\": 4, \"coordinator_address\": \"skyrl-tx-job-skyrl-workers-0-0.skyrl-tx-job:7777\"}"
else
sleep 60
uv run --extra tpu --extra tinker --extra jax -m skyrl.backends.jax \
--coordinator-address skyrl-tx-job-skyrl-workers-0-0.skyrl-tx-job:7777 \
--num-processes 4 \
--process-id "$JOB_COMPLETION_INDEX"
fi
resources:
requests:
google.com/tpu: 4
limits:
google.com/tpu: 4
Apply the resource with kubectl:
kubectl apply -f jobset.yaml
Step 4: Verify the Tinker API server and engine
Retrieve logs from the first replicated job running the Tinker API server and engine:
kubectl logs -f -l batch.kubernetes.io/job-completion-index=0,job-name=skyrl-tx-job-skyrl-workers-0
Verify engine and workers have started from the output:
Installed 144 packages in 650ms
2026-03-26 17:18:44,231 - INFO - uvicorn.error: Started server process [401]
2026-03-26 17:18:44,233 - INFO - uvicorn.error: Waiting for application startup.
2026-03-26 17:18:45,877 - INFO - skyrl: Using internal engine for inference
2026-03-26 17:18:45,878 - DEBUG - skyrl: Detected API server uv run flags:
['--extra', 'tpu', '--extra', 'tinker', '--extra', 'jax']
2026-03-26 17:18:45,884 - INFO - skyrl: Started background engine with PID 404:
uv run --extra tpu --extra tinker --extra jax --extra tinker --extra jax -m
skyrl.tinker.engine --base-model Qwen/Qwen3-0.6B --backend jax --backend-config
{"train_micro_batch_size": 8, "sample_max_num_sequences": 256,
"tensor_parallel_size": 4, "fully_sharded_data_parallel_size": 4,
"num_processes": 4, "coordinator_address":
"skyrl-tx-job-skyrl-workers-0-0.skyrl-tx-job:7777"} --checkpoints-base
/tmp/skyrl_checkpoints --database-url
sqlite:////jax-ai-image/SkyRL/skyrl/tinker/tinker.db
--external-inference-api-key EMPTY --external-inference-lora-base
/tmp/lora_models --session-cleanup-interval-sec 60 --session-timeout-sec 300
2026-03-26 17:18:45,886 - INFO - uvicorn.error: Application startup complete.
2026-03-26 17:18:45,888 - INFO - uvicorn.error: Uvicorn running on
http://0.0.0.0:8000 (Press CTRL+C to quit)
warning: The `extra-build-dependencies` option is experimental and may change without warning. Pass `--preview-features extra-build-dependencies` to disable this warning.
WARNING:absl:Tensorflow library not found, tensorflow.io.gfile operations will use native shim calls. GCS paths (i.e. 'gs://...') cannot be accessed.
2026-03-26 17:20:32,331 - INFO - skyrl: JAX distributed initialized:
process_id=0 (4 total), local devices: 4, total devices: 16
Fetching 11 files: 100%|██████████| 11/11 [00:06<00:00, 1.83it/s]
Step 5: Run the RL training loop
In this example, we apply RL to enhance the model’s mathematical problem-solving capabilities. This walkthrough utilizes the math RL example from the tinker cookbook examples. Execute the following command to initiate the RL training process:
POD=$(kubectl get pod -l batch.kubernetes.io/job-completion-index=0,job-name=skyrl-tx-job-skyrl-workers-0 -o custom-columns=":metadata.name" --no-headers)
kubectl exec -ti $POD -- sh -c 'cd SkyRL && export TINKER_API_KEY=tml-dummy && uv run --with "tinker-cookbook[math-rl] @ git+https://github.com/thinking-machines-lab/tinker-cookbook.git@nightly" python -m tinker_cookbook.recipes.math_rl.train base_url=http://localhost:8000 model_name="Qwen/Qwen3-0.6B" group_size=4 groups_per_batch=100 learning_rate=1e-4 max_tokens=512 save_every=0'
You can observe the model’s progress across multiple rollouts, each assigned a reward based on its outcome. For instance, the following output represents a high-reward trajectory where the model successfully solved the math problem.
****** trajectory idx=2, reward=1 ******
Per-step metrics:
Step 0:
format: 1.0
correct: 1.0
---- datum ----
<|im_start|>user
What is 4 + 5?<|im_end|>
<|im_start|>assistant
9<|im_end|>
<|im_start|>user
What is 56 + 72?<|im_end|>
<|im_start|>assistant
<think>
Alright, the user is asking 56 plus 72. Let me add those numbers. 56 plus 70 is 126, and then add the remaining 2, so 128. Yep, that seems right. To make sure, I can break it down: 50 + 70 is 120, and 6 + 2 is 8, so 120 + 8 = 128. Final answer is 128.
</think>
128<|im_end|>
You can view the training metrics to view the model’s progress as it learns to solve Math problems. The metrics should be included in the output after each training step like the following:
Step 18
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Metric ┃ Value ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ env/all/ac_tokens_per_turn │ 484.094727 │
│ env/all/by_group/frac_all_bad │ 0.453125 │
│ env/all/by_group/frac_all_good │ 0.046875 │
│ env/all/by_group/frac_mixed │ 0.500000 │
│ env/all/correct │ 0.268555 │
│ env/all/format │ 0.212891 │
│ env/all/ob_tokens_per_turn │ 217.890625 │
│ env/all/reward/total │ 0.189844 │
│ env/all/total_ac_tokens │ 495713 │
│ env/all/total_episodes │ 1024 │
│ env/all/total_ob_tokens │ 223120 │
│ env/all/total_turns │ 1024 │
│ env/all/turns_per_episode │ 1.000000 │
│ optim/entropy │ 0.431664 │
│ optim/kl_sample_train_v1 │ 0.003004 │
│ optim/kl_sample_train_v2 │ 0.003044 │
│ optim/lr │ 0.000020 │
│ progress/batch │ 18 │
│ progress/done_frac │ 0.101064 │
│ skyrl.ai/grad_norm │ 0.003937 │
│ skyrl.ai/learning_rate │ 0.000020 │
│ time/assemble_training_data │ 0.725428 │
│ time/compute_full_batch_metrics_and_get_sampling_client │ 2.225906 │
│ time/compute_group_rewards:max │ 0.000005 │
│ time/compute_group_rewards:mean │ 0.000003 │
│ time/compute_kl_sample_train │ 0.125626 │
│ time/do_group_rollout_and_filter_constant_reward:max │ 98.730395 │
│ time/do_group_rollout_and_filter_constant_reward:mean │ 82.161999 │
│ time/do_train_step_and_get_sampling_client │ 66.614592 │
│ time/env_initial_observation:max │ 0.399646 │
│ time/env_initial_observation:mean │ 0.003795 │
│ time/env_step:max │ 0.008772 │
│ time/env_step:mean │ 0.000565 │
│ time/policy_sample:max │ 96.301080 │
│ time/policy_sample:mean │ 74.323960 │
│ time/prepare_minibatch │ 0.740415 │
│ time/sampling │ 98.864469 │
│ time/save_checkpoint │ 2.100218 │
│ time/save_checkpoint_and_get_sampling_client │ 2.100240 │
│ time/total │ 165.533009 │
│ time/train_step │ 63.602819 │
└─────────────────────────────────────────────────────────┴────────────┘