[tinker] 8/n towards Kimi K2.6: protobuf forward_backward wire format (SDK >= 0.25)

skyrl/backends/skyrl_train_backend.py

@@ -945,7 +945,22 @@ def _forward_single_model_batch(
         )
         batch, pad_size = self._pad_batch(batch, micro_batch_size=micro_bs)
         model_id = prepared_batch.all_model_ids[0] if prepared_batch.all_model_ids else None
-        data = self._dispatch.forward(role, batch, model_id=model_id)
+        # Forward must run the SAME pipeline as forward_backward (loss path,
+        # forward_only): the loss-less worker forward takes Megatron's
+        # inference route, which bypasses training-only transforms like
+        # fake-INT4 QAT and the fused LM-head logprob, so its logprobs
+        # diverge from what gradients are computed against. The tinker
+        # FORWARD request always carries a loss_fn; use it (policy roles).
+        loss_fn = prepared_batch.all_loss_fns[0] if prepared_batch.all_loss_fns else None
+        loss_fn_config = next((c for c in prepared_batch.all_loss_fn_configs if c is not None), None)
+        if role != "critic" and loss_fn is not None:
+            self._validate_batch_role_and_loss(role, loss_fn)
+            loss_fn, loss_fn_config = self._normalize_policy_loss_request(role, loss_fn, loss_fn_config)
+            data = self._dispatch.forward(
+                role, batch, loss_fn=loss_fn, loss_fn_config=loss_fn_config, model_id=model_id
+            )
+        else:
+            data = self._dispatch.forward(role, batch, model_id=model_id)
 
         # Workers emit per-sample loss_fn_outputs directly. Trim padding entries.
         per_sample_outputs = data.loss_fn_outputs

skyrl/tinker/api.py

@@ -70,8 +70,8 @@
 
 async def wait_for_future(
     waiters: dict[int, set[asyncio.Future]], request_id: int, timeout: float
-) -> tuple[RequestStatus, dict | None] | None:
-    """Wait for ``request_id`` to finish, returning its ``(status, result_data)``.
+) -> tuple[RequestStatus, dict | None, types.RequestType] | None:
+    """Wait for ``request_id`` to finish, returning ``(status, result_data, request_type)``.
 
     Returns None if ``timeout`` elapses first, and raises KeyError if the request
     does not exist -- :func:`poll_futures` reports both.
@@ -118,15 +118,15 @@ async def poll_futures(
                     # The awaited ids go in as bound parameters, capped at 32766
                     # by SQLite (since 3.32) and 65535 by Postgres -- far above
                     # any plausible number of in-flight requests.
-                    statement = select(FutureDB.request_id, FutureDB.status, FutureDB.result_data).where(
-                        FutureDB.request_id.in_(awaited)
-                    )
+                    statement = select(
+                        FutureDB.request_id, FutureDB.status, FutureDB.result_data, FutureDB.request_type
+                    ).where(FutureDB.request_id.in_(awaited))
                     rows = (await session.exec(statement)).all()
 
                 # Pending requests are left alone to be picked up on a later tick.
-                outcomes: dict[int, tuple[RequestStatus, dict | None] | KeyError] = {
-                    request_id: (status, result_data)
-                    for request_id, status, result_data in rows
+                outcomes: dict[int, tuple[RequestStatus, dict | None, types.RequestType] | KeyError] = {
+                    request_id: (status, result_data, request_type)
+                    for request_id, status, result_data, request_type in rows
                     if status in TERMINAL_STATUSES
                 }
                 for request_id in set(awaited) - {request_id for request_id, *_ in rows}:
@@ -594,6 +594,257 @@ class ForwardRequest(BaseModel):
     forward_input: ForwardBackwardInput
 
 
+# ---------------------------------------------------------------------------
+# Protobuf wire-format support for /api/v1/forward_backward.
+#
+# tinker SDK >= 0.25 serializes ForwardBackwardRequest as protobuf
+# (Content-Type: application/x-protobuf, optionally zstd Content-Encoding) and
+# has no JSON fallback; older SDKs (<= 0.24) switch to the same proto path when
+# the server advertises `proto_write_fwdbwd` in /api/v1/client/config. The
+# decoder below converts the proto message into the existing pydantic request
+# models so the rest of the pipeline is wire-format agnostic. The proto schema
+# ships with the `tinker` package (a dependency of the `tinker` extra).
+# ---------------------------------------------------------------------------
+
+_PROTO_CONTENT_TYPE = "application/x-protobuf"
+
+# Cap for decompressed request bodies (zstd Content-Encoding).
+_MAX_FWDBWD_BODY_BYTES = 1 << 30
+
+
+def _proto_module():
+    try:
+        from tinker.proto import tinker_public_pb2 as public_pb
+    except ImportError as exc:  # pragma: no cover - tinker extra always ships it
+        raise HTTPException(
+            status_code=415,
+            detail="This server cannot decode application/x-protobuf bodies "
+            f"(tinker proto schema unavailable: {exc}). Send application/json.",
+        ) from exc
+    return public_pb
+
+
+def _decode_tensor_pb(tensor, public_pb) -> TensorData:
+    import numpy as np
+
+    encoding = tensor.WhichOneof("encoding")
+    if encoding != "dense":
+        raise HTTPException(status_code=422, detail=f"unsupported tensor encoding {encoding!r} (only dense)")
+    dtype_map = {
+        public_pb.DTYPE_FLOAT32: np.float32,
+        public_pb.DTYPE_INT64: np.int64,
+        public_pb.DTYPE_INT32: np.int32,
+    }
+    np_dtype = dtype_map.get(tensor.dtype)
+    if np_dtype is None:
+        raise HTTPException(status_code=422, detail=f"unsupported tensor dtype enum {tensor.dtype}")
+    try:
+        values = np.frombuffer(tensor.dense, dtype=np_dtype)
+    except ValueError as exc:
+        raise HTTPException(status_code=422, detail=f"malformed tensor data: {exc}") from exc
+    return TensorData(data=values.tolist())
+
+
+def _decode_chunk_pb(chunk) -> EncodedTextChunk:
+    import numpy as np
+
+    which = chunk.WhichOneof("chunk")
+    if which != "encoded_text":
+        # Image/audio chunks additionally require the server-side asset upload
+        # path the managed service uses; reject clearly rather than mangle.
+        raise HTTPException(status_code=422, detail=f"unsupported model-input chunk type {which!r} on the proto path")
+    try:
+        tokens = np.frombuffer(chunk.encoded_text.tokens, dtype=np.int32)
+    except ValueError as exc:
+        raise HTTPException(status_code=422, detail=f"malformed chunk tokens: {exc}") from exc
+    return EncodedTextChunk(tokens=tokens.tolist())
+
+
+def _decode_forward_backward_pb(body: bytes) -> tuple[ForwardBackwardRequest, bool]:
+    """Parse a proto ForwardBackwardRequest; returns (request, forward_only)."""
+    public_pb = _proto_module()
+    msg = public_pb.ForwardBackwardRequest()
+    try:
+        msg.ParseFromString(body)
+    except Exception as exc:
+        raise HTTPException(status_code=422, detail=f"invalid protobuf ForwardBackwardRequest: {exc}") from exc
+
+    data = []
+    for datum_pb in msg.data:
+        chunks = [_decode_chunk_pb(c) for c in datum_pb.model_input]
+        loss_fn_inputs = {key: _decode_tensor_pb(value, public_pb) for key, value in datum_pb.loss_fn_inputs.items()}
+        data.append(Datum(model_input=ModelInput(chunks=chunks), loss_fn_inputs=loss_fn_inputs))
+
+    try:
+        fb_input = ForwardBackwardInput(
+            data=data,
+            loss_fn=msg.loss_fn,
+            loss_fn_config=dict(msg.loss_fn_config) if msg.loss_fn_config else None,
+        )
+    except Exception as exc:
+        raise HTTPException(status_code=422, detail=f"invalid forward_backward payload: {exc}") from exc
+    return ForwardBackwardRequest(model_id=msg.model_id, forward_backward_input=fb_input), bool(msg.forward_only)
+
+
+_PROTO_MASK_LOGPROB = -99999.0  # tinker SDK sentinel for masked topk rows
+
+
+def _encode_forward_backward_output_pb(result: dict) -> bytes:
+    """Encode a stored ForwardBackwardOutput result dict as proto bytes.
+
+    Inverse of the SDK's ``deserialize_forward_backward_output``: one
+    ArrayRecord carrying every datum, each field a BatchedTensor of the
+    per-datum flat arrays with int64 byte offsets. Per-datum vectors are
+    flat, so ``trailing_shape`` stays empty.
+    """
+    import numpy as np
+
+    public_pb = _proto_module()
+    msg = public_pb.ForwardBackwardOutput()
+    msg.loss_fn_output_type = str(result.get("loss_fn_output_type") or "ArrayRecord")
+    for key, value in (result.get("metrics") or {}).items():
+        try:
+            msg.metrics[key] = float(value)
+        except (TypeError, ValueError):
+            continue
+
+    outputs = result.get("loss_fn_outputs") or []
+    if not outputs:
+        return msg.SerializeToString()
+
+    record = msg.loss_fn_outputs.add()
+    record.type_tag = msg.loss_fn_output_type
+    record.num_datums = len(outputs)
+
+    field_names: list[str] = []
+    for datum in outputs:
+        for name in datum:
+            if name not in field_names:
+                field_names.append(name)
+
+    for name in field_names:
+        dtype_label = "float32"
+        for datum in outputs:
+            td = datum.get(name)
+            if isinstance(td, dict) and td.get("dtype"):
+                dtype_label = td["dtype"]
+                break
+        np_dtype = np.int64 if dtype_label == "int64" else np.float32
+        proto_dtype = public_pb.DTYPE_INT64 if dtype_label == "int64" else public_pb.DTYPE_FLOAT32
+
+        flats = []
+        for datum in outputs:
+            td = datum.get(name)
+            data = (td.get("data") if isinstance(td, dict) else td) or []
+            flats.append(np.asarray(data, dtype=np_dtype).reshape(-1))
+        offsets = np.zeros(len(flats) + 1, dtype=np.int64)
+        for i, arr in enumerate(flats):
+            offsets[i + 1] = offsets[i] + arr.nbytes
+
+        bt = record.fields[name]
+        bt.data = b"".join(arr.tobytes() for arr in flats)
+        bt.offsets = offsets.tobytes()
+        bt.dtype = proto_dtype
+    return msg.SerializeToString()
+
+
+def _encode_sample_response_pb(result: dict) -> bytes:
+    """Encode a stored SampleOutput result dict as a proto SampleResponse."""
+    import math as _math
+
+    import numpy as np
+
+    public_pb = _proto_module()
+    msg = public_pb.SampleResponse()
+    for seq in result.get("sequences") or []:
+        seq_pb = msg.sequences.add()
+        seq_pb.stop_reason = (
+            public_pb.STOP_REASON_STOP if seq.get("stop_reason") == "stop" else public_pb.STOP_REASON_LENGTH
+        )
+        seq_pb.tokens = np.asarray(seq.get("tokens") or [], dtype=np.int32).tobytes()
+        logprobs = seq.get("logprobs")
+        if logprobs:
+            seq_pb.logprobs = np.asarray(logprobs, dtype=np.float32).tobytes()
+
+    prompt_logprobs = result.get("prompt_logprobs")
+    if prompt_logprobs:
+        arr = np.asarray(
+            [float("nan") if x is None else float(x) for x in prompt_logprobs],
+            dtype=np.float32,
+        )
+        msg.prompt_logprobs = arr.tobytes()
+
+    topk = result.get("topk_prompt_logprobs")
+    if topk:
+        k = max((len(row) for row in topk if row), default=0)
+        n = len(topk)
+        if n and k:
+            token_ids = np.zeros((n, k), dtype=np.int32)
+            logprobs = np.full((n, k), _PROTO_MASK_LOGPROB, dtype=np.float32)
+            for i, row in enumerate(topk):
+                if not row:
+                    continue  # all-masked row decodes to None client-side
+                for j, pair in enumerate(row[:k]):
+                    token_ids[i, j] = int(pair[0])
+                    logprobs[i, j] = float(pair[1])
+                    if not _math.isfinite(logprobs[i, j]):
+                        logprobs[i, j] = _PROTO_MASK_LOGPROB
+            msg.topk_prompt_logprobs.token_ids = token_ids.tobytes()
+            msg.topk_prompt_logprobs.logprobs = logprobs.tobytes()
+            msg.topk_prompt_logprobs.prompt_length = n
+            msg.topk_prompt_logprobs.k = k
+    return msg.SerializeToString()
+
+
+_PROTO_RESULT_ENCODERS = {
+    types.RequestType.FORWARD_BACKWARD: _encode_forward_backward_output_pb,
+    types.RequestType.FORWARD: _encode_forward_backward_output_pb,
+    types.RequestType.SAMPLE: _encode_sample_response_pb,
+    types.RequestType.EXTERNAL: _encode_sample_response_pb,
+}
+
+
+async def _parse_forward_backward_body(raw_request: Request) -> tuple[ForwardBackwardRequest, bool]:
+    """Decode a forward_backward body in either wire format.
+
+    JSON (SDK <= 0.24 default) and protobuf (SDK >= 0.25, or older SDKs when
+    `proto_write_fwdbwd` is advertised) are both accepted; anything else gets
+    a clean 415 instead of a 500 from FastAPI trying to utf-8 decode binary.
+    """
+    body = await raw_request.body()
+    if raw_request.headers.get("content-encoding", "").strip().lower() == "zstd":
+        import zstandard
+
+        try:
+            # Bounded like delta_payload.decompress_bytes: a small crafted body
+            # must not balloon into an arbitrarily large allocation. 1 GiB is
+            # far above any legitimate forward_backward payload.
+            body = zstandard.ZstdDecompressor().decompress(body, max_output_size=_MAX_FWDBWD_BODY_BYTES)
+        except Exception as exc:
+            raise HTTPException(status_code=422, detail=f"failed to zstd-decompress request body: {exc}") from exc
+
+    content_type = raw_request.headers.get("content-type", "").split(";")[0].strip().lower()
+    if content_type == _PROTO_CONTENT_TYPE:
+        return _decode_forward_backward_pb(body)
+    if content_type in ("application/json", ""):
+        try:
+            return ForwardBackwardRequest.model_validate_json(body), False
+        except Exception as exc:
+            import pydantic
+
+            if isinstance(exc, pydantic.ValidationError):
+                # exc.json() is fully JSON-safe (unlike FastAPI's default
+                # encoder, which crashes on bytes in the error context).
+                import json as _json
+
+                raise HTTPException(status_code=422, detail=_json.loads(exc.json())) from exc
+            raise HTTPException(status_code=422, detail=f"invalid JSON body: {exc}") from exc
+    raise HTTPException(
+        status_code=415,
+        detail=f"Unsupported content type {content_type!r}; send application/json or {_PROTO_CONTENT_TYPE}",
+    )
+
+
 class AdamParams(BaseModel):
     learning_rate: float = Field(default=1e-4, ge=0.0)
     beta1: float = Field(default=0.9, ge=0.0, lt=1.0)
@@ -824,6 +1075,12 @@ class WeightsInfoResponse(BaseModel):
 
 class ClientConfigResponse(BaseModel):
     pjwt_auth_enabled: bool = False
+    # Advertise the protobuf forward_backward path: SDK >= 0.25 requires it
+    # (its forward() asserts the flag and it has no JSON fallback), and older
+    # SDKs switch to proto when the flag is set. /api/v1/forward_backward
+    # decodes both wire formats either way.
+    proto_write_fwdbwd: bool = True
+    proto_compress_fwdbwd: bool = False
 
 
 @app.post("/api/v1/client/config", response_model=ClientConfigResponse)
@@ -1017,13 +1274,19 @@ async def get_training_run(model_id: str, session: AsyncSession = Depends(get_se
 
 
 @app.post("/api/v1/forward_backward", response_model=FutureResponse)
-async def forward_backward(request: ForwardBackwardRequest, session: AsyncSession = Depends(get_session)):
-    """Compute and accumulate gradients."""
+async def forward_backward(raw_request: Request, session: AsyncSession = Depends(get_session)):
+    """Compute and accumulate gradients.
+
+    Accepts JSON or protobuf bodies (see :func:`_parse_forward_backward_body`).
+    The proto schema carries a ``forward_only`` flag instead of using the
+    separate /forward route; honor it by dispatching the request type.
+    """
+    request, forward_only = await _parse_forward_backward_body(raw_request)
     await get_model(session, request.model_id)
 
     request_id = await create_future(
         session=session,
-        request_type=types.RequestType.FORWARD_BACKWARD,
+        request_type=types.RequestType.FORWARD if forward_only else types.RequestType.FORWARD_BACKWARD,
         model_id=request.model_id,
         request_data=request.forward_backward_input.to_types(),
     )
@@ -1266,8 +1529,23 @@ async def retrieve_future(request: RetrieveFutureRequest, req: Request):
     if row is None:
         raise HTTPException(status_code=408, detail="Timeout waiting for result")
 
-    status, result_data = row
+    status, result_data, request_type = row
     if status == RequestStatus.COMPLETED:
+        # SDK >= 0.25 only accepts forward_backward/sample results as
+        # proto (it sends Accept: application/x-protobuf and raises on
+        # a JSON body for those types). Encode on demand; fall back to
+        # JSON for clients that don't ask.
+        encoder = _PROTO_RESULT_ENCODERS.get(request_type)
+        if (
+            encoder is not None
+            and result_data is not None
+            and "error" not in result_data
+            and _PROTO_CONTENT_TYPE in req.headers.get("accept", "").lower()
+        ):
+            return fastapi.Response(
+                content=encoder(result_data),
+                media_type=_PROTO_CONTENT_TYPE,
+            )
         return result_data
 
     # Return 400 for handled errors (validation, etc.), 500 for unexpected failures