Add Tinker API server for supervised fine-tuning with SkyRL backend

[tyler-griggs]

Summary

This PR adds the Tinker HTTP API server that enables external training orchestration using the SkyRL-Train backend. This allows tools like tinker-cookbook to run supervised fine-tuning (SFT) by making HTTP calls to control training, rather than embedding training logic directly.

Key capabilities:

• HTTP API for forward_backward() and optim_step() operations
• Support for LoRA training with dynamic model creation
• External control of mini-batching and training loops
• Per-sequence loss outputs for fine-grained control

Changes

New directories:

• skyrl_train/tinker/ - Tinker API server and engine

  • api.py - FastAPI server with training endpoints
  • engine.py - Background process handling training requests
  • backends/skyrl_train.py - SkyRL backend implementation
  • types.py - Type definitions for API contracts
  • loss_fns.py - Loss function implementations (JAX made optional)

• skyrl_train/tx_utils/ - Shared utilities

  • generator.py, log.py, models.py, etc.

Key modifications:

• Updated all imports from tx.* to skyrl_train.*
• Made JAX imports optional (only needed for JAX backend, not SkyRL)
• Removed JAX backend references (only SkyRL backend supported)
• Fixed engine subprocess path to use skyrl_train.tinker.engine

Architecture

tinker-cookbook (sl_loop.py)
  └─ HTTP requests to Tinker API
       └─ skyrl-train/skyrl_train/tinker/api.py (FastAPI server)
            └─ skyrl-train/skyrl_train/tinker/engine.py (background process)
                 └─ skyrl-train/skyrl_train/tinker/backends/skyrl_train.py
                      └─ WorkerDispatch → Ray actors → PolicyWorker

Usage

Start the Tinker API server:

cd ~/SkyRL/skyrl-train
uv run --extra vllm python -m skyrl_train.tinker.api \
    --base-model Qwen/Qwen3-0.6B \
    --backend skyrl_train \
    --port 8001

Run supervised fine-tuning from tinker-cookbook:

cd ~/tinker-cookbook
TINKER_API_KEY=test uv run python -m tinker_cookbook.recipes.sl_loop \
    base_url="http://localhost:8001" \
    model_name="Qwen/Qwen3-0.6B" \
    batch_size=4 \
    lora_rank=8

Test Results

Successfully tested with sl_loop.py from tinker-cookbook:

• Training runs at ~0.5-0.9s per step
• Loss decreases as expected (3.10 → 2.5-3.7 range)
• Forward/backward takes ~0.33-0.36s per mini-batch
• Gradient norms in expected range (678-1750)

Background

This integration was needed to enable external training orchestration for SFT workloads. The original Tinker code lived in skyrl-tx but required Flash Attention which had environment issues. Copying into skyrl-train allows it to use the working Flash Attention installation and existing SkyRL infrastructure.

Next Steps

• Add tests for Tinker API endpoints
• Add documentation for API usage
• Consider re-enabling checkpointing support in sl_loop
• Evaluate performance optimizations (ray.put() for batches)

🤖 Generated with Claude Code