Add SkyRL-TX Qwen SFT and RL example

README.md

@@ -20,6 +20,7 @@ Use this repository to learn how to build distributed training jobs on Modal.
 - [**`lightning/`**](/lightning/) a simple lightning.ai Fabric example.
 - [**`nanoGPT/`**](/nanoGPT/) training Karpathy's nanoGPT reproduction of OpenAI's GPT-2.
 - [**`resnet50/`**](/resnet50/) training a ResNet50 model on the ImageNet dataset.
+- [**`skyrl-tx/`**](/skyrl-tx/) SkyRL-TX Tinker-compatible SFT and RL on Qwen3-8B.
 - [**`starcoder/`**](/starcoder) accelerated finetuning of Llama-2-7B on Rust and Go code, supporting either `torchrun` or `accelerate`.
 
 # Documentation

skyrl-tx/README.md

@@ -0,0 +1,164 @@
+# SkyRL-TX Qwen SFT and RL on Modal
+
+This example runs [SkyRL-TX](https://github.com/NovaSky-AI/SkyRL/tree/main/skyrl-tx)
+as a Tinker-compatible training server on a Modal multi-node GPU cluster.
+It uses `Qwen/Qwen3-8B` by default and exercises both:
+
+- supervised fine-tuning with Tinker's `cross_entropy` loss
+- policy-gradient RL with Tinker's `ppo` loss over sampled arithmetic rollouts
+
+The default topology is 2 nodes × 8 H100 GPUs. Within each node SkyRL-TX uses
+tensor parallelism; across nodes it uses JAX FSDP.
+
+## Tinker cookbook compatibility reports
+
+These notes assess 10 public
+[Tinker cookbook](https://github.com/thinking-machines-lab/tinker-cookbook)
+recipes against this example's pinned SkyRL-TX JAX backend and Modal launch
+shape:
+
+| Cookbook recipe | Status | Report |
+| --- | --- | --- |
+| `sl_loop` | Supported | [cookbook-sl-loop](cookbook-sl-loop/) |
+| `rl_loop` | Supported for default no-KL loop | [cookbook-rl-loop](cookbook-rl-loop/) |
+| `chat_sl` | Supported for text-only chat SFT | [cookbook-chat-sl](cookbook-chat-sl/) |
+| `math_rl` | Supported for default no-KL RL; partial for KL diagnostics | [cookbook-math-rl](cookbook-math-rl/) |
+| `code_rl` | Partial: sandbox infrastructure is external | [cookbook-code-rl](cookbook-code-rl/) |
+| `preference/dpo` | Partial: custom DPO smoke passed; full recipe still needs dataset adaptation | [cookbook-preference-dpo](cookbook-preference-dpo/) |
+| `preference/rlhf` | Partial: reward-policy smoke passed; full pipeline needs adaptation | [cookbook-preference-rlhf](cookbook-preference-rlhf/) |
+| `distillation/on_policy_distillation` | Partial: same-server teacher smoke passed; separate teacher models remain unvalidated | [cookbook-distillation-on-policy](cookbook-distillation-on-policy/) |
+| `search_tool` | Partial: retrieval stack is external | [cookbook-search-tool](cookbook-search-tool/) |
+| `vlm_classifier` | Partial API smoke only; real VLM support is not provided by this text-only example | [cookbook-vlm-classifier](cookbook-vlm-classifier/) |
+
+## Prerequisites
+
+- Modal CLI installed and authenticated
+- Modal environment selected: `export MODAL_ENVIRONMENT=<your-env>`
+- Modal secret `huggingface-secret` with `HF_TOKEN`
+- Access to multi-node GPU clusters
+
+Run all commands from the repo root.
+
+## Quickstart
+
+Download the Qwen checkpoint into the persistent Hugging Face cache volume:
+
+```bash
+modal run skyrl-tx/modal_train.py::download_model
+```
+
+Run the supervised fine-tuning smoke:
+
+```bash
+modal run --detach skyrl-tx/modal_train.py::run_sft
+```
+
+Run the RL smoke:
+
+```bash
+modal run --detach skyrl-tx/modal_train.py::run_rl
+```
+
+Run the cookbook compatibility smoke suite:
+
+```bash
+modal run --detach skyrl-tx/modal_train.py::run_cookbook --lora-rank 4
+```
+
+Run a single cookbook probe by name:
+
+```bash
+modal run --detach skyrl-tx/modal_train.py::run_cookbook --lora-rank 4 --example sl_loop
+```
+
+Use detached mode for the training jobs; the image build, model load, JAX
+initialization, and first compile can take several minutes.
+
+## Cluster sizing
+
+The launcher reads this sizing environment variable at import time:
+
+| Variable | Default | Purpose |
+| --- | --- | --- |
+| `SKYRL_TX_N_NODES` | `2` | Number of Modal containers in the JAX cluster |
+
+Each container always requests a full `H100:8` node. Partial H100 allocations do
+not work for multi-node SkyRL-TX runs.
+
+For the default `Qwen/Qwen3-8B` run:
+
+```text
+total GPUs = SKYRL_TX_N_NODES × 8 = 16 by default
+tensor_parallel_size = 8
+fully_sharded_data_parallel_size = SKYRL_TX_N_NODES = 2
+```
+
+## How it works
+
+`run_sft` and `run_rl` create a per-run ephemeral Modal Dict for coordination,
+then launch clustered `run_sft_cluster` or `run_rl_cluster` functions. Rank 0
+starts the SkyRL-TX Tinker API server:
+
+```bash
+uv run --extra gpu --extra tinker --extra jax -m skyrl.tinker.api \
+  --base-model Qwen/Qwen3-8B \
+  --backend jax \
+  --backend-config '{"tensor_parallel_size": 8, "fully_sharded_data_parallel_size": 2, ...}'
+```
+
+Ranks 1..N start SkyRL-TX JAX workers:
+
+```bash
+uv run --extra gpu --extra tinker --extra jax -m skyrl.backends.jax \
+  --coordinator-address <rank-0-ip>:7777 \
+  --num-processes <N> \
+  --process-id <rank>
+```
+
+After the API server reports healthy, rank 0 runs `sft_client.py`,
+`rl_client.py`, or `cookbook_smoke_client.py` against `http://localhost:8000`.
+
+## Volumes
+
+| Volume | Mount path | Purpose |
+| --- | --- | --- |
+| `skyrl-tx-hf-cache` | `/root/.cache/huggingface` | Qwen model cache |
+| `skyrl-tx-checkpoints` | `/checkpoints` | SkyRL-TX LoRA checkpoints |
+
+Each training job saves both a Tinker training-state checkpoint and a sampler
+checkpoint. The clients run a small evaluation pass after training, then load the
+saved sampler checkpoint for a sample/reward evaluation. Rank 0 lists the
+checkpoint files it found and commits the Modal volume. Successful runs print
+lines like:
+
+```text
+sft_state_checkpoint=file://...
+sft_eval_loss=...
+sft_sampler_checkpoint=file://...
+sft_sampler_eval_sample=...
+sft_checkpoint_artifact=... bytes=...
+sft_checkpoint_volume_committed=...
+rl_state_checkpoint=file://...
+rl_eval_loss_outputs=...
+rl_sampler_checkpoint=file://...
+rl_eval mean_reward=... trajectories=...
+rl_checkpoint_artifact=... bytes=...
+rl_checkpoint_volume_committed=...
+cookbook_result={"example": "...", "status": "PASS", ...}
+cookbook_summary={"passed": 10, "unexpected_failures": 0}
+cookbook_cookbook_results=.../cookbook_results.jsonl bytes=...
+```
+
+## Adjusting the smoke
+
+Both entrypoints expose small training-loop knobs:
+
+```bash
+modal run --detach skyrl-tx/modal_train.py::run_sft --steps 16 --lora-rank 8 --learning-rate 1e-6
+modal run --detach skyrl-tx/modal_train.py::run_rl --steps 8 --samples-per-prompt 4 --learning-rate 1e-6
+```
+
+The clients intentionally use tiny arithmetic datasets so the example validates
+the end-to-end SkyRL-TX path without requiring a full benchmark-scale run. The
+default `1e-6` learning rate keeps short LoRA smoke runs stable on the tiny
+batches used here.

skyrl-tx/cookbook-chat-sl/README.md

@@ -0,0 +1,66 @@
+# Tinker cookbook compatibility: `chat_sl`
+
+Status: **supported for text-only chat SFT**
+
+Source recipe: `tinker_cookbook.recipes.chat_sl.train`
+
+## What the cookbook example does
+
+`chat_sl` is the higher-level supervised chat-training recipe. It wraps the
+shared supervised trainer with dataset builders for public chat datasets such as
+Tulu3 and NoRobots, plus periodic evaluation and checkpointing.
+
+## Tinker API surface
+
+The recipe delegates to `tinker_cookbook.supervised.train`, so it uses:
+
+- `create_lora_training_client(...)` or checkpoint-based training-client restore
+- `forward_backward(..., loss_fn="cross_entropy")`
+- `forward(...)` through NLL/evaluation helpers
+- `optim_step(...)`
+- periodic and final state/sampler checkpoint saves
+
+## SkyRL-TX support
+
+The backend-facing training path is supported:
+
+- Cross-entropy chat SFT is the most direct SkyRL-TX use case.
+- Text-only chat renderers produce `EncodedTextChunk` inputs, which the JAX
+  SkyRL-TX backend consumes.
+- Periodic and final checkpoints are compatible with the checkpoint path used by
+  the Modal example.
+
+## Required adjustments
+
+- Pass `base_url=http://localhost:8000` and `model_name` matching the launched
+  SkyRL-TX base model.
+- Ensure the chosen cookbook renderer supports that model family. For the Modal
+  example's default, use a Qwen-compatible renderer.
+- If using the Modal launcher as-is, set the launcher/server `lora_rank` high
+  enough for the recipe's requested rank. Cookbook defaults often use rank 32,
+  while this repo's smoke defaults use rank 8.
+- Reduce dataset/batch sizes for a quick validation run.
+
+## Unsupported or risky pieces
+
+No text-SFT backend feature is missing. The risk is mostly scale: full public
+datasets plus frequent eval/checkpoint cadence will be much heavier than the
+small Modal smoke in this repo.
+
+## Executed SkyRL-TX smoke
+
+Smoke code: `skyrl-tx/cookbook_smoke_client.py::CookbookSmokeRunner.chat_sl`
+
+Validated with:
+
+```bash
+modal run skyrl-tx/modal_train.py::run_cookbook --lora-rank 4 --example chat_sl
+```
+
+Recorded result on 2 x `H100:8`: **PASS**. The smoke used the tokenizer chat
+template when available, trained only assistant-token labels with
+`cross_entropy`, and applied one optimizer step.
+
+```json
+{"example":"chat_sl","status":"PASS","loss_sum":-49.1875,"loss_values":32,"duration_seconds":0.771}
+```

skyrl-tx/cookbook-code-rl/README.md

@@ -0,0 +1,66 @@
+# Tinker cookbook compatibility: `code_rl`
+
+Status: **partial: model training path supported, sandbox infrastructure external**
+
+Source recipe: `tinker_cookbook.recipes.code_rl.train`
+
+## What the cookbook example does
+
+`code_rl` trains on programming tasks. The policy generates code, a sandbox runs
+tests, rewards come from test pass/fail and formatting signals, and the shared RL
+trainer updates the model from the resulting trajectories.
+
+## Tinker API surface
+
+- `save_weights_and_get_sampling_client()` for rollout sampling
+- `SamplingClient.sample(...)`
+- `forward_backward(..., loss_fn="importance_sampling")` by default
+- `optim_step(...)`
+- periodic/final checkpoints
+- optional async/off-policy training mode
+
+## SkyRL-TX support
+
+The Tinker backend requirements are mostly supported:
+
+- Text-only code prompts are ordinary `EncodedTextChunk` inputs.
+- Default `importance_sampling` training is supported.
+- Generated-token logprobs, optimizer steps, and checkpoints are supported.
+- Async/off-policy scheduling is client-side; it still submits standard sample
+  and forward/backward requests to the API.
+
+## Required adjustments
+
+- Pass `base_url=http://localhost:8000`, `model_name` matching the loaded server,
+  and a LoRA rank allowed by the server.
+- Provide a working sandbox backend. The cookbook supports SandboxFusion/Docker
+  or Modal sandboxing; SkyRL-TX does not provide that execution environment.
+- For Modal runs, do not rely on Modal NFS for checkpoint or sandbox state. Use
+  normal local files plus Modal Volumes, as the SkyRL-TX example does.
+- Keep `kl_penalty_coef=0.0` and `compute_post_kl=False` for the supported path.
+
+## Unsupported or risky pieces
+
+The code-execution sandbox is the main non-SkyRL dependency. Full-scale code RL
+also has high variance and long rollout latency, so it needs a dedicated Modal
+validation run before calling it production-ready. KL/post-KL settings remain
+unvalidated for this SkyRL-TX example.
+
+## Executed SkyRL-TX smoke
+
+Smoke code: `skyrl-tx/cookbook_smoke_client.py::CookbookSmokeRunner.code_rl`
+
+Validated with:
+
+```bash
+modal run skyrl-tx/modal_train.py::run_cookbook --lora-rank 4 --example code_rl
+```
+
+Recorded result on 2 x `H100:8`: **PASS** for the model-side path. The smoke
+sampled a code-shaped prompt, supplied a client-side reward, ran one
+`importance_sampling` forward/backward step, and applied one optimizer step. It
+did **not** run SandboxFusion/Docker/Modal sandbox tests.
+
+```json
+{"example":"code_rl","status":"PASS","loss_sum":-497.1251588719726,"loss_values":38,"duration_seconds":11.802,"external_dependency":"sandbox not exercised"}
+```

skyrl-tx/cookbook-distillation-on-policy/README.md

@@ -0,0 +1,67 @@
+# Tinker cookbook compatibility: `distillation/on_policy_distillation`
+
+Status: **partial: same-server teacher smoke validated; separate teacher models remain unvalidated**
+
+Source recipe: `tinker_cookbook.recipes.distillation.on_policy_distillation`
+
+## What the cookbook example does
+
+The on-policy distillation recipe has a student sample trajectories, then uses a
+teacher model's per-token logprobs to apply a reverse-KL training signal. Its
+default `kl_penalty_coef` is `1.0`, so teacher KL is the actual supervision.
+
+## Tinker API surface
+
+- student `create_lora_training_client(...)`
+- teacher `create_sampling_client(base_model=...)` or checkpoint-backed sampler
+- student rollout `sample(...)`
+- teacher `compute_logprobs_async(...)`
+- `forward_backward(..., loss_fn="importance_sampling")`
+- `optim_step(...)` and checkpoints
+
+## SkyRL-TX support
+
+Only the generic pieces are supported:
+
+- Text rollout sampling, generated-token logprobs, `importance_sampling`,
+  optimizer steps, and checkpointing are available.
+- A same-base teacher represented by a saved sampler checkpoint can use the
+  standard sampler path.
+
+## Required adjustments
+
+- Set `model_name` and `teacher_model` to the same model loaded by the SkyRL-TX
+  server, or provide teacher weights as a sampler checkpoint for that same base.
+- Pass `base_url=http://localhost:8000`.
+- Reduce `lora_rank`; cookbook defaults use rank 128, much higher than this
+  repo's smoke configuration.
+- For a backend-only smoke, set `kl_penalty_coef=0.0`; that converts it into a
+  standard no-reward/no-KL RL loop and no longer tests the recipe's core
+  distillation objective.
+
+## Unsupported or risky pieces
+
+As written, the recipe depends on teacher prompt logprobs through
+`compute_logprobs_async` and can use different teacher/student base models. The
+same-server teacher smoke passed, but the current SkyRL-TX Modal example still
+loads one text model. Treat separate teacher models, larger KL-heavy runs, and
+non-default student/teacher renderer combinations as unvalidated.
+
+## Executed SkyRL-TX smoke
+
+Smoke code: `skyrl-tx/cookbook_smoke_client.py::CookbookSmokeRunner.distillation_on_policy`
+
+Validated with:
+
+```bash
+modal run skyrl-tx/modal_train.py::run_cookbook --lora-rank 4 --example distillation_on_policy
+```
+
+Recorded result on 2 x `H100:8`: **PASS** for a same-server teacher. The smoke
+sampled from a checkpoint-backed teacher sampler, called `compute_logprobs` on
+the sampled trajectory, converted the teacher/student logprob difference into an
+advantage, ran one `importance_sampling` update, and applied one optimizer step.
+
+```json
+{"example":"distillation_on_policy","status":"PASS","loss_sum":-232.4375,"loss_values":26,"teacher_logprobs":12,"duration_seconds":48.351}
+```