Expose native GSPO through Tinker API

skyrl/backends/jax.py

@@ -634,8 +634,9 @@ def _model_pass(
         """
         if not prepared_batch.all_model_inputs:
             return {}
-        if "ppo_critic" in prepared_batch.all_loss_fns:
-            raise ValueError("ppo_critic is only supported by the SkyRL-Train backend")
+        unsupported_loss_fns = set(prepared_batch.all_loss_fns) - LOSS_TYPES.keys()
+        if unsupported_loss_fns:
+            raise ValueError(f"Loss functions {sorted(unsupported_loss_fns)} are not supported by the JAX backend")
 
         results = {}
 

skyrl/backends/skyrl_train_backend.py

@@ -817,7 +817,7 @@ def _normalize_policy_loss_request(
                 normalized_config["dppo"] = dppo_overrides
             return loss_fn, normalized_config or None
 
-        if loss_fn != "ppo":
+        if loss_fn not in {"ppo", "gspo"}:
             return loss_fn, loss_fn_config
 
         normalized_config = dict(loss_fn_config or {})
@@ -827,7 +827,7 @@ def _normalize_policy_loss_request(
             normalized_config["eps_clip_low"] = 1.0 - clip_low_threshold
         if clip_high_threshold is not None:
             normalized_config["eps_clip_high"] = clip_high_threshold - 1.0
-        return "regular", normalized_config or None
+        return ("regular" if loss_fn == "ppo" else "gspo"), normalized_config or None
 
     def forward_backward(
         self,

skyrl/tinker/api.py

@@ -602,13 +602,14 @@ class ForwardBackwardInput(BaseModel):
         "cross_entropy": set(),
         "importance_sampling": set(),
         "ppo": {"clip_low_threshold", "clip_high_threshold", "value_clip"},
+        "gspo": {"clip_low_threshold", "clip_high_threshold"},
         "cispo": {"clip_low_threshold", "clip_high_threshold"},
         "ppo_critic": {"value_clip"},
         "dppo": {"delta_low", "delta_high"},
     }
 
     data: list[Datum]
-    loss_fn: Literal["cross_entropy", "importance_sampling", "ppo", "cispo", "ppo_critic", "dppo"]
+    loss_fn: Literal["cross_entropy", "importance_sampling", "ppo", "gspo", "cispo", "ppo_critic", "dppo"]
     loss_fn_config: dict[str, float] | None = None
 
     @model_validator(mode="after")

skyrl/tinker/types.py

@@ -160,7 +160,7 @@ class Datum(BaseModel):
 
 class ForwardBackwardInput(BaseModel):
     data: list[Datum]
-    loss_fn: Literal["cross_entropy", "importance_sampling", "ppo", "cispo", "ppo_critic", "dppo"]
+    loss_fn: Literal["cross_entropy", "importance_sampling", "ppo", "gspo", "cispo", "ppo_critic", "dppo"]
     loss_fn_config: dict[str, float] | None = None
 
 
@@ -341,6 +341,7 @@ class PreparedSampleBatch(BaseModel):
     "cross_entropy",
     "importance_sampling",
     "ppo",
+    "gspo",
     "cispo",
     "ppo_critic",
     "dppo",

tests/backends/test_jax_backend.py

@@ -6,7 +6,7 @@
 import optax
 import pytest
 
-from skyrl.backends.jax import JaxBackend, JaxBackendConfig
+from skyrl.backends.jax import JaxBackend, JaxBackendConfig, JaxBackendImpl
 from skyrl.tinker import api, types
 from skyrl.tinker.engine import prepare_model_pass_batch, prepare_sample_batch
 from skyrl.tinker.types import LoraConfig, OptimStepInput
@@ -172,6 +172,17 @@ def make_fwd_bwd_input(token_lists: list[list[int]]) -> types.ForwardBackwardInp
     return types.ForwardBackwardInput(data=samples, loss_fn="cross_entropy")
 
 
+@pytest.mark.parametrize("loss_fn", ["gspo", "dppo", "ppo_critic"])
+def test_model_pass_rejects_unsupported_loss(loss_fn):
+    model_pass_input = make_fwd_bwd_input([[1, 2, 3]])
+    model_pass_input.loss_fn = loss_fn
+    prepared_batch = prepare_model_pass_batch({"request": ("model", model_pass_input)})
+
+    backend = object.__new__(JaxBackendImpl)
+    with pytest.raises(ValueError, match=f"{loss_fn}.*not supported by the JAX backend"):
+        backend._model_pass(prepared_batch, None)
+
+
 def _assert_tree_allclose(t1, t2, rtol=1e-3, atol=1e-3, min_match_pct=99.0):
     """Assert that at least min_match_pct% of elements in two trees are close."""
     leaves1 = jax.tree.leaves(t1)