[skyrl-train] Add SFT support via forward_backward(loss_fn="cross_entropy")

Enables supervised fine-tuning using the Tinker-compatible API.

Changes:
- ppo_utils.py: Add CROSS_ENTROPY loss type and cross_entropy_loss() function
- worker.py: Add SFT code path that returns per-token logprobs and elementwise_loss
- worker_dispatch.py: Add loss_fn and loss_fn_config params to forward_backward()
- dispatch.py: Update MeshDispatch to pass through kwargs (loss_fn, loss_fn_config)
- replay_buffer.py: Make action_log_probs optional in Experience
- worker_utils.py: Use .get() for optional fields; handle non-scalar metrics

New:
- examples/sft/: Minimal SFT example demonstrating the API

This enables PR #871 (SkyRL-train backend for Tinker) to return proper
per-token values instead of placeholder data.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: tyler-griggs

skyrl-train/examples/sft/README.md

@@ -0,0 +1,30 @@
+# SFT (Supervised Fine-Tuning) Example
+
+This example demonstrates how to use SkyRL's training infrastructure for supervised fine-tuning (SFT).
+
+## Usage
+
+```bash
+uv run --isolated --extra vllm python examples/sft/sft_trainer.py
+```
+
+## How It Works
+
+1. **Load Dataset**: Uses a small subset of the Alpaca dataset
+2. **Tokenize**: Converts instruction/output pairs into token sequences
+3. **Create Batch**: Builds a `TrainingInputBatch` with:
+   - `sequences`: Token IDs (left-padded)
+   - `attention_mask`: 1 for real tokens, 0 for padding
+   - `loss_mask`: 1 for response tokens to compute loss on
+4. **Train**: Calls `forward_backward(loss_fn="cross_entropy")` for SFT
+
+## Loss Functions
+
+The `loss_fn` parameter supports:
+
+| Loss Function | Use Case |
+|--------------|----------|
+| `cross_entropy` | Supervised fine-tuning |
+| `regular` / `ppo` | PPO with clipping |
+| `gspo` | Group Sequence Policy Optimization |
+| ... | See `PolicyLossRegistry` for all options |

skyrl-train/examples/sft/sft_trainer.py

@@ -0,0 +1,192 @@
+"""
+Minimal SFT (Supervised Fine-Tuning) trainer using WorkerDispatch.
+
+This script demonstrates SFT using the same forward_backward interface as RL training,
+but with loss_fn="cross_entropy" to compute simple negative log-likelihood loss.
+
+Usage:
+    # First, make sure you have Ray installed and a GPU available
+    uv run --isolated --extra vllm python examples/sft/sft_trainer.py
+
+This example:
+1. Loads a small subset of the Alpaca dataset
+2. Tokenizes examples into prompt + completion format
+3. Uses WorkerDispatch.forward_backward(loss_fn="cross_entropy") for SFT
+4. Demonstrates the Tinker-compatible API for supervised fine-tuning
+"""
+
+import ray
+import hydra
+import torch
+from datasets import load_dataset
+from loguru import logger
+from omegaconf import DictConfig
+from transformers import AutoTokenizer
+from tqdm import tqdm
+
+from ray.util.placement_group import placement_group
+
+from skyrl_train.training_batch import TrainingInputBatch
+from skyrl_train.entrypoints.main_base import config_dir
+from skyrl_train.workers.worker_dispatch import WorkerDispatch
+from skyrl_train.workers.worker import PPORayActorGroup
+from skyrl_train.workers.fsdp.fsdp_worker import PolicyWorker
+from skyrl_train.utils.utils import initialize_ray, validate_cfg
+from skyrl_train.utils import get_ray_pg_ready_with_timeout
+
+
+def get_sft_config() -> DictConfig:
+    """Get config with SFT-specific overrides."""
+    with hydra.initialize_config_dir(config_dir=config_dir):
+        cfg = hydra.compose(config_name="ppo_base_config")
+
+    # Use a small model for testing
+    cfg.trainer.policy.model.path = "Qwen/Qwen2.5-0.5B-Instruct"
+    cfg.trainer.placement.policy_num_gpus_per_node = 1
+    cfg.generator.inference_engine_tensor_parallel_size = 1
+    cfg.trainer.logger = "console"
+    cfg.trainer.micro_train_batch_size_per_gpu = 2
+
+    validate_cfg(cfg)
+    return cfg
+
+
+def tokenize_sft_example(example: dict, tokenizer, max_length: int = 512) -> dict | None:
+    """Tokenize a single SFT example (instruction + output).
+
+    Returns dict with input_ids, attention_mask, num_actions (response length),
+    or None if the example was fully truncated.
+    """
+    instruction = example.get("instruction", "")
+    input_text = example.get("input", "")
+    output = example.get("output", "")
+
+    # Combine instruction and input
+    if input_text:
+        prompt = f"{instruction}\n\n{input_text}"
+    else:
+        prompt = instruction
+
+    # Tokenize prompt and full sequence separately to find boundary
+    prompt_tokens = tokenizer(prompt, add_special_tokens=True, truncation=True, max_length=max_length)
+    full_text = f"{prompt}\n\n{output}"
+    full_tokens = tokenizer(full_text, add_special_tokens=True, truncation=True, max_length=max_length)
+
+    prompt_len = len(prompt_tokens["input_ids"])
+    full_len = len(full_tokens["input_ids"])
+    num_actions = full_len - prompt_len
+
+    # Skip examples where response was fully truncated
+    if num_actions <= 0:
+        return None
+
+    return {
+        "input_ids": full_tokens["input_ids"],
+        "attention_mask": full_tokens["attention_mask"],
+        "num_actions": num_actions,
+    }
+
+
+def collate_sft_batch(examples: list, tokenizer) -> TrainingInputBatch:
+    """Collate tokenized examples into a TrainingInputBatch.
+
+    Creates the batch format expected by forward_backward with cross_entropy loss:
+    - sequences: [batch_size, seq_len] - token IDs (left-padded)
+    - attention_mask: [batch_size, seq_len] - 1 for real tokens, 0 for padding
+    - loss_mask: [batch_size, num_actions] - 1 for tokens to compute loss on
+    """
+    max_len = max(len(ex["input_ids"]) for ex in examples)
+    max_num_actions = max(ex["num_actions"] for ex in examples)
+
+    sequences = []
+    attention_masks = []
+    loss_masks = []
+
+    for ex in examples:
+        pad_len = max_len - len(ex["input_ids"])
+        # Left-pad sequences (SkyRL convention)
+        sequences.append([tokenizer.pad_token_id] * pad_len + ex["input_ids"])
+        attention_masks.append([0] * pad_len + ex["attention_mask"])
+        # Per-example loss_mask: 0s for padding, 1s only for this example's response tokens
+        action_pad = max_num_actions - ex["num_actions"]
+        loss_masks.append([0] * action_pad + [1] * ex["num_actions"])
+
+    batch = TrainingInputBatch(
+        {
+            "sequences": torch.tensor(sequences, dtype=torch.long),
+            "attention_mask": torch.tensor(attention_masks, dtype=torch.long),
+            "loss_mask": torch.tensor(loss_masks, dtype=torch.long),
+        }
+    )
+    batch.metadata = {"response_length": max_num_actions}
+    return batch
+
+
+def main():
+    """Run a minimal SFT training loop."""
+    cfg = get_sft_config()
+    initialize_ray(cfg)
+
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(cfg.trainer.policy.model.path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    logger.info("Loading dataset...")
+    # Use a small subset for demonstration
+    dataset = load_dataset("yahma/alpaca-cleaned", split="train[:100]")
+
+    logger.info("Tokenizing dataset...")
+    tokenized = [tokenize_sft_example(ex, tokenizer) for ex in dataset]
+    tokenized = [ex for ex in tokenized if ex is not None]  # Filter truncated
+    logger.info(f"Kept {len(tokenized)} examples after filtering truncated")
+
+    logger.info("Initializing policy worker...")
+    num_gpus = cfg.trainer.placement.policy_num_gpus_per_node
+    pg = placement_group([{"GPU": num_gpus, "CPU": num_gpus}], strategy="PACK")
+    get_ray_pg_ready_with_timeout(pg, timeout=30)
+
+    actor_group = PPORayActorGroup(
+        cfg,
+        num_nodes=1,
+        num_gpus_per_node=num_gpus,
+        ray_actor_type=PolicyWorker,
+        pg=pg,
+        num_gpus_per_actor=0.75,
+        colocate_all=False,
+        sequence_parallel_size=cfg.trainer.policy.sequence_parallel_size,
+    )
+    ray.get(actor_group.async_init_model(cfg.trainer.policy.model.path))
+
+    dispatch = WorkerDispatch(cfg, policy_actor_group=actor_group)
+
+    # Training loop
+    batch_size = 4
+    num_steps = 10
+    logger.info(f"Starting SFT training for {num_steps} steps...")
+
+    for step in tqdm(range(num_steps)):
+        # Create batch from tokenized examples
+        start_idx = (step * batch_size) % len(tokenized)
+        batch_examples = tokenized[start_idx : start_idx + batch_size]
+        if len(batch_examples) < batch_size:
+            batch_examples = tokenized[:batch_size]  # Wrap around
+
+        batch = collate_sft_batch(batch_examples, tokenizer)
+
+        # Forward-backward with cross-entropy loss (Tinker API style)
+        metrics = dispatch.forward_backward("policy", batch, loss_fn="cross_entropy")
+
+        # Optimizer step
+        grad_norm = dispatch.optim_step("policy")
+
+        if step % 5 == 0:
+            loss_val = metrics.get("final_loss", metrics.get("loss", "N/A"))
+            logger.info(f"Step {step}: loss={loss_val:.4f}, grad_norm={grad_norm}")
+
+    logger.info("SFT training complete!")
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    main()

skyrl-train/skyrl_train/dataset/replay_buffer.py

@@ -57,7 +57,7 @@ class Experience:
     """
 
     sequences: Integer[torch.Tensor, "batch seq_len"]
-    action_log_probs: Float[torch.Tensor, "batch response_len"]
+    action_log_probs: Optional[Float[torch.Tensor, "batch response_len"]]
     base_action_log_probs: Optional[Float[torch.Tensor, "batch response_len"]]
     values: Optional[Float[torch.Tensor, "batch response_len"]]
     returns: Optional[Float[torch.Tensor, "batch response_len"]]
@@ -74,7 +74,8 @@ class Experience:
     @torch.no_grad()
     def to_device(self, device: torch.device) -> None:
         self.sequences = to(self.sequences, device)
-        self.action_log_probs = to(self.action_log_probs, device)
+        if self.action_log_probs is not None:
+            self.action_log_probs = to(self.action_log_probs, device)
         if self.base_action_log_probs is not None:
             self.base_action_log_probs = to(self.base_action_log_probs, device)
         if self.values is not None:
@@ -94,7 +95,8 @@ def to_device(self, device: torch.device) -> None:
 
     def pin_memory(self):
         self.sequences = pin_memory(self.sequences)
-        self.action_log_probs = pin_memory(self.action_log_probs)
+        if self.action_log_probs is not None:
+            self.action_log_probs = pin_memory(self.action_log_probs)
         if self.base_action_log_probs is not None:
             self.base_action_log_probs = pin_memory(self.base_action_log_probs)
         if self.values is not None:

skyrl-train/skyrl_train/distributed/dispatch.py

@@ -121,7 +121,9 @@ class MeshDispatch(Dispatch):
     """
 
     @classmethod
-    def dispatch(cls, actor_infos: List[ActorInfo], method: str, data: TrainingInputBatch) -> List[ObjectRef]:
+    def dispatch(
+        cls, actor_infos: List[ActorInfo], method: str, data: TrainingInputBatch, **kwargs
+    ) -> List[ObjectRef]:
         assert len(actor_infos) > 0, "actor_infos must be a non-empty list"
         object_refs = []
         dp_size = actor_infos[0].rank.dp_size
@@ -134,7 +136,7 @@ def dispatch(cls, actor_infos: List[ActorInfo], method: str, data: TrainingInput
         for actor_info in actor_infos:
             # index into tensordict to get the correct data to send
             data_to_send = data_chunks[actor_info.rank.dp]
-            object_refs.append(getattr(actor_info.handle, method).remote(data_to_send))
+            object_refs.append(getattr(actor_info.handle, method).remote(data_to_send, **kwargs))
         return object_refs
 
     @classmethod
@@ -159,24 +161,14 @@ def sync_collect(cls, actor_infos: List[ActorInfo], object_refs: List[ObjectRef]
 
     @classmethod
     def validate_dispatch_args(cls, *args, **kwargs) -> Tuple[Tuple, Dict[str, Any]]:
-        sig = inspect.signature(cls.dispatch)
-        # pass dummy actor_infos and method_name
-        bound_args = sig.bind([], "dummy", *args, **kwargs)
-        bound_args.apply_defaults()
-
-        # Check if there are any extra arguments
-        if len(bound_args.arguments) > 3:  #  data, actor_infos, method_name
-            # remove actor_infos and method_name - not added by user
-            bound_args.arguments.pop("actor_infos")
-            bound_args.arguments.pop("method")
-            raise ValueError(f"MeshDispatch only accepts 'data' as an argument, got extra args: {bound_args.arguments}")
-
-        data = bound_args.arguments.get("data")
+        # First positional arg must be data (TrainingInputBatch)
+        if not args:
+            raise ValueError("MeshDispatch requires 'data' as first positional argument")
+        data = args[0]
         if not isinstance(data, TrainingInputBatch):
-            raise ValueError(f"For MeshDispatch, `data` entry should be a `TrainingInput`, got {data}")
-        args = (data,)
-        kwargs = {}
-        return args, kwargs
+            raise ValueError(f"For MeshDispatch, `data` entry should be a `TrainingInputBatch`, got {type(data)}")
+        # Pass through data as positional arg, and any kwargs (e.g., loss_fn, loss_fn_config)
+        return (data,), kwargs
 
 
 class PassThroughDispatch(Dispatch):

skyrl-train/skyrl_train/utils/ppo_utils.py

@@ -471,6 +471,7 @@ class PolicyLossType(StrEnum):
     CLIP_COV = "clip_cov"
     KL_COV = "kl_cov"
     SAPO = "sapo"
+    CROSS_ENTROPY = "cross_entropy"
 
 
 class PolicyLossRegistry(BaseFunctionRegistry):
@@ -500,6 +501,7 @@ def repopulate_registry(cls):
             "clip_cov": [PolicyLossType.CLIP_COV, compute_policy_loss_clip_cov],
             "kl_cov": [PolicyLossType.KL_COV, compute_policy_loss_kl_cov],
             "sapo": [PolicyLossType.SAPO, sapo_policy_loss],
+            "cross_entropy": [PolicyLossType.CROSS_ENTROPY, cross_entropy_loss],
         }
 
         for pl_name, (pl_type, pl_func) in pl_types.items():
@@ -878,6 +880,48 @@ def compute_policy_loss_kl_cov(
     return pg_loss, 0.0
 
 
+@register_policy_loss(PolicyLossType.CROSS_ENTROPY)
+def cross_entropy_loss(
+    log_probs: torch.Tensor,
+    old_log_probs: torch.Tensor,
+    advantages: torch.Tensor,
+    config: DictConfig,
+    loss_mask: Optional[torch.Tensor] = None,
+    rollout_logprobs: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, float]:
+    """
+    Cross-entropy loss for supervised fine-tuning (SFT).
+
+    This loss function computes the negative log-likelihood of the target tokens,
+    ignoring the old_log_probs and advantages which are only used for RL.
+
+    The loss is computed as: -log_probs * loss_mask, summed over all tokens.
+    This matches Tinker's cross_entropy semantics where the loss is a simple sum.
+
+    Args:
+        log_probs: Log probabilities from the model for each token
+        old_log_probs: Ignored (only used for RL losses)
+        advantages: Ignored (only used for RL losses)
+        config: Algorithm configuration
+        loss_mask: Mask indicating which tokens to include in loss (1=include, 0=ignore)
+        rollout_logprobs: Ignored (only used for RL losses)
+
+    Returns:
+        Tuple of (loss, clip_ratio) where clip_ratio is always 0.0 for SFT
+    """
+    # Simple negative log-likelihood: -log p(token)
+    elementwise_loss = -log_probs
+
+    # Apply loss mask and sum (matching Tinker's SUM reduction semantics)
+    if loss_mask is not None:
+        loss = (elementwise_loss * loss_mask).sum()
+    else:
+        loss = elementwise_loss.sum()
+
+    # No clipping in cross-entropy loss
+    return loss, 0.0
+
+
 def reduce_loss(
     loss: torch.Tensor,
     loss_mask: Optional[torch.Tensor],