[tinker] Support tinker SDK 0.32/0.33 (pin <0.34)

Widen the SDK pin to >=0.25.0,<0.34.0 and lock 0.33.1.

- proto_serialization: 0.28 renamed TopkPromptLogprobs -> TopkLogprobs
  (prompt_length -> length, same field numbers). Pick whichever the
  installed SDK schema has, like loss_fn_config_v2.
- JAX backend: emit loss:sum (sum of per-token losses) from
  forward_backward; SDK >= 0.31 custom-loss pops it. The skyrl-train
  backend got the same alias in #2455.
- API: accept 0.32's optimizer_config on create_model/load_weights and
  optimizer_params on optim_step, rejecting anything but Adam with a 422.
  Reject the new SampleRequest options the engines cannot produce
  (topk_sample_logprobs, target_prompt_logprobs, prompt_alt_tokens_k,
  0.33's prompt_logprobs_last_n) instead of returning None for them.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Gg8qLdfM4FSHLBjyQHZX2y
Signed-off-by: Avi Basnet <avigya.basnet@anyscale.com>

Author: avigyabb

.agents/docs/tinker.md

@@ -44,13 +44,15 @@ All endpoints are under `/api/v1/`. Requests are async -- submit via POST, get a
 
 ## SDK compatibility
 
-`pyproject.toml` pins `tinker>=0.25.0,<0.28.0`; `uv.lock` resolves 0.27.2. Verified against 0.25.0, 0.27.0 and 0.27.2. What each SDK line needs from the server:
+`pyproject.toml` pins `tinker>=0.25.0,<0.34.0`; `uv.lock` resolves 0.33.1. The server encodes/decodes with whichever SDK's `tinker_public_pb2` is installed, so `proto_serialization.py` feature-detects schema changes instead of assuming one line. What each SDK line needs from the server:
 
 - **0.25**: proto `forward_backward` bodies and proto `retrieve_future` results (`proto_serialization.py`).
-- **0.26/0.27**: `asample` must return `sample_sequence_ids`, one per `num_samples` (`SampleFutureResponse`); the SDK stamps them onto `SampledSequence.sequence_id` client-side, so the proto result does not carry them. `forward_backward` requests dual-write `loss_fn_config` into a legacy float map and `loss_fn_config_v2` (number/text union); the parser prefers v2 so string values reach the API model (which rejects them, since no SkyRL loss takes a string).
-- **0.28 is not supported yet**: it renames the proto `TopkPromptLogprobs` message to `TopkLogprobs` (`prompt_length` -> `length`), which breaks `serialize_sample_output` for `topk_prompt_logprobs > 0`, and adds `save_weights_external`.
+- **0.26/0.27**: `asample` must return `sample_sequence_ids`, one per `num_samples` (`SampleFutureResponse`); the SDK stamps them onto `SampledSequence.sequence_id` client-side. `forward_backward` requests dual-write `loss_fn_config` into a legacy float map and `loss_fn_config_v2` (number/text union); the parser prefers v2 so string values reach the API model (which rejects them, since no SkyRL loss takes a string).
+- **0.28**: proto `TopkPromptLogprobs` renamed to `TopkLogprobs` (`prompt_length` -> `length`); the top-k serializer picks whichever the installed SDK has.
+- **0.31+**: `forward_backward_custom` pops `loss:sum` from the result metrics (the Tinker service's name for the training loss). Both backends emit it: `skyrl_train_backend._extract_metrics` mirrors `total_loss:sum`, the JAX backend sums the per-token losses.
+- **0.32**: optimizer families. `create_model`/`load_weights` carry `optimizer_config` and `optim_step` sends non-Adam params under `optimizer_params` instead of `adam_params`; anything but Adam is a 422. New `SampleRequest` options the engines do not produce (`topk_sample_logprobs`, `target_prompt_logprobs`, `prompt_alt_tokens_k`, and 0.33's `prompt_logprobs_last_n`) are rejected rather than silently ignored.
 
-0.27 client features that are opt-in behind server-driven flags and that SkyRL leaves off (the stub `/api/v1/client/config` returns defaults): `sample_use_retrieve_futures` (`/api/v1/retrieve_futures` session poller), `sample_cancel_enabled` (`/api/v1/cancel_future`). Not implemented and ignored: `copy_weights`, checkpoint `user_metadata`, `Datum` provenance spans, `SampleRequest.record_stability_info`.
+Client features that are opt-in behind server-driven flags and that SkyRL leaves off (the stub `/api/v1/client/config` returns defaults): `sample_use_retrieve_futures` (`/api/v1/retrieve_futures` session poller), `sample_join_sampling_session` (`/api/v1/join_sampling_session`), dynamic-config `sample_cancel_enabled` (`/api/v1/cancel_future`). Not implemented: `copy_weights`, `save_weights_external` and external weights URLs, checkpoint `user_metadata`, `Datum` provenance spans, `SampleRequest.record_stability_info`, the checkpoint billing route.
 
 ## Concurrency and Batching
 

pyproject.toml

@@ -39,7 +39,7 @@ tpu = [
 ]
 
 tinker = [
-    "tinker>=0.25.0,<0.28.0",
+    "tinker>=0.25.0,<0.34.0",
     "fastapi[standard]",
     "sqlmodel",
     "sqlalchemy[asyncio]",

skyrl/backends/jax.py

@@ -757,11 +757,13 @@ def _model_pass(
         # Compute per-request results
         for request_id, _, start_idx, end_idx in request_batch_slices:
             loss_fn_outputs = []
+            loss_sum = 0.0
             # Compute per-example losses
             for i in range(start_idx, end_idx):
                 # Extract losses for this example's tokens
                 token_losses = token_losses_out[i]
                 token_logprobs = logprobs_out[i]
+                loss_sum += float(token_losses.sum())
                 loss_fn_outputs.append(
                     {
                         "elementwise_loss": {
@@ -780,7 +782,9 @@ def _model_pass(
             results[request_id] = types.ForwardBackwardOutput(
                 loss_fn_output_type="scalar",
                 loss_fn_outputs=loss_fn_outputs,
-                metrics={},
+                # The SDK's custom-loss path (>= 0.31) pops "loss:sum" from the
+                # result to report the surrogate loss.
+                metrics={"loss:sum": loss_sum},
             )
 
         return results

skyrl/tinker/api.py

@@ -25,6 +25,7 @@
     Field,
     Tag,
     ValidationError,
+    field_validator,
     model_validator,
 )
 from sqlalchemy.exc import IntegrityError
@@ -531,11 +532,23 @@ class LoRAConfig(BaseModel):
     )
 
 
+def _validate_adam_optimizer_config(optimizer_config: dict[str, Any] | None) -> dict[str, Any] | None:
+    """SDK >= 0.32 selects the optimizer family at model creation; SkyRL only has Adam."""
+    if optimizer_config is not None and optimizer_config.get("type", "adamw") != "adamw":
+        raise ValueError(
+            f"Unsupported optimizer_config type {optimizer_config.get('type')!r}; only 'adamw' is supported"
+        )
+    return optimizer_config
+
+
 class CreateModelRequest(BaseModel):
     session_id: str
     base_model: str
     lora_config: LoRAConfig
     model_role: str = "policy"
+    optimizer_config: dict[str, Any] | None = None
+
+    _check_optimizer = field_validator("optimizer_config")(_validate_adam_optimizer_config)
 
 
 class CreateModelResponse(BaseModel):
@@ -762,9 +775,21 @@ def to_types(self) -> types.AdamParams:
 
 class OptimStepRequest(BaseModel):
     model_id: str
-    adam_params: AdamParams
+    adam_params: AdamParams | None = None
+    # SDK >= 0.32 sends non-Adam families (e.g. Dimuon) under this key instead.
+    optimizer_params: dict[str, Any] | None = None
     seq_id: int | None = None
 
+    @model_validator(mode="after")
+    def validate_adam_only(self):
+        if self.optimizer_params is not None:
+            raise ValueError(
+                f"Unsupported optimizer_params type {self.optimizer_params.get('type')!r}; only Adam (adam_params) is supported"
+            )
+        if self.adam_params is None:
+            raise ValueError("adam_params is required")
+        return self
+
 
 class SaveWeightsForSamplerRequest(BaseModel):
     model_id: str
@@ -832,8 +857,31 @@ class SampleRequest(BaseModel):
     seq_id: int | None = None
     prompt_logprobs: bool | None = None
     topk_prompt_logprobs: int = Field(default=0, ge=0)
+    # SDK >= 0.32 / 0.33 request options SkyRL's engines do not produce. They
+    # are rejected rather than ignored, since the SDK would otherwise hand the
+    # caller None for data it asked for.
+    topk_sample_logprobs: int = 0
+    target_prompt_logprobs: Any | None = None
+    prompt_alt_tokens_k: int = 0
+    prompt_logprobs_last_n: int | None = None
     type: Literal["sample"] = "sample"
 
+    @model_validator(mode="after")
+    def validate_unsupported_options(self):
+        unsupported = [
+            name
+            for name, value in (
+                ("topk_sample_logprobs", self.topk_sample_logprobs),
+                ("target_prompt_logprobs", self.target_prompt_logprobs),
+                ("prompt_alt_tokens_k", self.prompt_alt_tokens_k),
+                ("prompt_logprobs_last_n", self.prompt_logprobs_last_n),
+            )
+            if value not in (0, None)
+        ]
+        if unsupported:
+            raise ValueError(f"Unsupported sampling options: {unsupported}")
+        return self
+
     @model_validator(mode="after")
     def validate_model_source(self):
         """Valid if:
@@ -862,8 +910,11 @@ class LoadWeightsRequest(BaseModel):
     path: str
     optimizer: bool = True
     seq_id: int | None = None
+    optimizer_config: dict[str, Any] | None = None
     type: Literal["load_weights"] | None = None
 
+    _check_optimizer = field_validator("optimizer_config")(_validate_adam_optimizer_config)
+
 
 class FutureResponse(BaseModel):
     future_id: str
@@ -1498,7 +1549,7 @@ async def optim_step(request: OptimStepRequest, session: AsyncSession = Depends(
         session=session,
         request_type=types.RequestType.OPTIM_STEP,
         model_id=request.model_id,
-        request_data=types.OptimStepInput(adam_params=request.adam_params.to_types()),
+        request_data=types.OptimStepInput(adam_params=request.adam_params.to_types()),  # validated non-None
         seq_id=request.seq_id,
     )
 

skyrl/tinker/proto_serialization.py

@@ -1,10 +1,11 @@
 """Proto wire-format conversion for the Tinker SDK's binary paths.
 
 Two directions, both required by tinker SDK >= 0.25.0 (older SDKs negotiate
-these via server-side client_config flags and fall back to JSON). Written
-against the SDK's ``tinker_public.proto`` as of 0.27.x; 0.28.0 renamed
-``TopkPromptLogprobs`` to ``TopkLogprobs`` (``prompt_length`` -> ``length``),
-which the top-k serializer below does not yet handle:
+these via server-side client_config flags and fall back to JSON). The server
+encodes with whichever SDK's ``tinker_public_pb2`` is installed, so schema
+changes across SDK lines (``loss_fn_config_v2`` in 0.27, the
+``TopkPromptLogprobs`` -> ``TopkLogprobs`` rename in 0.28) are feature-detected
+below rather than assumed:
 
 - **Responses**: the SDK retrieves ``sample``, ``forward``, and
   ``forward_backward`` results as protobuf (``Accept: application/x-protobuf``
@@ -65,6 +66,11 @@
 # against whichever SDK is installed, so feature-detect it.
 _HAS_LOSS_FN_CONFIG_V2 = "loss_fn_config_v2" in pb.ForwardBackwardRequest.DESCRIPTOR.fields_by_name
 
+# SDK >= 0.28 renamed the top-k matrix message (and its row-count field) when
+# it started reusing it for per-sequence sampled top-k.
+_TOPK_MESSAGE = getattr(pb, "TopkLogprobs", None) or pb.TopkPromptLogprobs
+_TOPK_LENGTH_FIELD = "length" if "length" in _TOPK_MESSAGE.DESCRIPTOR.fields_by_name else "prompt_length"
+
 
 def parse_forward_backward_request(body: bytes) -> tuple[dict, bool]:
     """Parse a proto ForwardBackwardRequest body into the JSON-equivalent
@@ -194,7 +200,7 @@ def serialize_sample_output(
     if topk_prompt_logprobs is not None:
         rows = topk_prompt_logprobs
         # k is not recorded in the result, so recover it from the widest row.
-        # With every row undefined, use k=1 so prompt_length stays encoded
+        # With every row undefined, use k=1 so the row count stays encoded
         # (the client maps fully-masked rows back to None).
         k = max((len(row) for row in rows if row), default=1)
         token_ids = np.full((len(rows), k), _TOPK_MASK_TOKEN_ID, dtype=np.int32)
@@ -204,11 +210,11 @@ def serialize_sample_output(
                 token_ids[i, j] = token_id
                 logprobs[i, j] = logprob
         proto.topk_prompt_logprobs.CopyFrom(
-            pb.TopkPromptLogprobs(
-                prompt_length=len(rows),
+            _TOPK_MESSAGE(
                 k=k,
                 token_ids=token_ids.tobytes(),
                 logprobs=logprobs.tobytes(),
+                **{_TOPK_LENGTH_FIELD: len(rows)},
             )
         )
 

tests/backends/test_jax_backend.py

@@ -292,7 +292,13 @@ def test_micro_batch_grad_accumulation():
     }
 
     # Run 1: micro-batching enabled
-    backend_micro.forward_backward(prepare_model_pass_batch(reqs))
+    outputs = backend_micro.forward_backward(prepare_model_pass_batch(reqs))
+
+    # The SDK's custom-loss path (>= 0.31) pops "loss:sum": the per-request sum of
+    # the elementwise losses it also receives.
+    for request_id, output in outputs.items():
+        expected = sum(sum(datum["elementwise_loss"]["data"]) for datum in output.loss_fn_outputs)
+        assert output.metrics["loss:sum"] == pytest.approx(expected, rel=1e-5), request_id
 
     adapter1_idx = backend_micro.models[adapter1_id].adapter_index
     adapter2_idx = backend_micro.models[adapter2_id].adapter_index

tests/tinker/test_api_validation.py

@@ -178,3 +178,73 @@ def test_nested_in_model_input(self):
         json_dict = types_input.model_dump(mode="json")
         recovered = types.ModelInput.model_validate(json_dict)
         assert recovered.chunks[0].data == _RAW_PNG
+
+
+# --- SDK >= 0.32 request fields SkyRL does not implement -------------------
+
+
+def _sample_request(**extra) -> api.SampleRequest:
+    return api.SampleRequest(
+        base_model="m",
+        prompt=api.ModelInput(chunks=[api.EncodedTextChunk(tokens=[1, 2])]),
+        sampling_params=api.SamplingParams(max_tokens=1),
+        **extra,
+    )
+
+
+def test_sample_request_accepts_default_new_options():
+    req = _sample_request(topk_sample_logprobs=0, prompt_alt_tokens_k=0, target_prompt_logprobs=None)
+    assert req.topk_sample_logprobs == 0
+
+
+@pytest.mark.parametrize(
+    "extra",
+    [
+        {"topk_sample_logprobs": 4},
+        {"prompt_alt_tokens_k": 2},
+        {"target_prompt_logprobs": {"data": [1], "dtype": "int64", "shape": [1, 1]}},
+        {"prompt_logprobs_last_n": 3},
+    ],
+)
+def test_sample_request_rejects_unsupported_options(extra):
+    """Rejecting beats ignoring: the SDK would otherwise return None for data the caller asked for."""
+    with pytest.raises(ValidationError, match="Unsupported sampling options"):
+        _sample_request(**extra)
+
+
+def test_create_model_accepts_adamw_optimizer_config():
+    req = api.CreateModelRequest(
+        session_id="s", base_model="m", lora_config=api.LoRAConfig(rank=1), optimizer_config={"type": "adamw"}
+    )
+    assert req.optimizer_config == {"type": "adamw"}
+
+
+def test_create_model_rejects_non_adam_optimizer_config():
+    with pytest.raises(ValidationError, match="only 'adamw' is supported"):
+        api.CreateModelRequest(
+            session_id="s",
+            base_model="m",
+            lora_config=api.LoRAConfig(rank=1),
+            optimizer_config={"type": "dimuon", "version": 1},
+        )
+
+
+def test_load_weights_rejects_non_adam_optimizer_config():
+    with pytest.raises(ValidationError, match="only 'adamw' is supported"):
+        api.LoadWeightsRequest(model_id="m", path="tinker://m/weights/c", optimizer_config={"type": "dimuon"})
+
+
+def test_optim_step_legacy_adam_params_wire_key():
+    req = api.OptimStepRequest(model_id="m", adam_params={"learning_rate": 1e-4})
+    assert req.adam_params is not None and req.adam_params.learning_rate == 1e-4
+
+
+def test_optim_step_rejects_non_adam_optimizer_params():
+    """SDK >= 0.32 sends non-Adam families under ``optimizer_params``."""
+    with pytest.raises(ValidationError, match="only Adam"):
+        api.OptimStepRequest(model_id="m", optimizer_params={"type": "dimuon", "learning_rate": 1e-3})
+
+
+def test_optim_step_requires_adam_params():
+    with pytest.raises(ValidationError, match="adam_params is required"):
+        api.OptimStepRequest(model_id="m")