[feat] Fireworks AI Backend for EvalOnlyEntrypoint

.claude/docs/inference.md

@@ -20,6 +20,46 @@ Prefill-Decode disaggregation:
 - **Config**: `enable_pd=true` and `num_prefill` passed to `ServerGroup` constructor. Requires a `kv_connector`
 - **Server groups**: Separate prefill and decode `ServerGroup`s, one per engine.
 
+## Fireworks (external generation, eval-only)
+
+`generator.inference_engine.backend=fireworks` sends generation to the external Fireworks endpoint
+via `FireworksInferenceClient` (`inference_servers/fireworks_client.py`, built on the `fireworks-ai`
+SDK — install the `fireworks` extra). Token-in/token-out is preserved: prompts are sent as raw token
+ids and Fireworks returns the generated `token_ids` (`return_token_ids`), so the stock
+`SkyRLGymGenerator` works unchanged.
+
+- **Eval-only**: accepted only by `skyrl.train.entrypoints.main_generate` (`EvalOnlyEntrypoint`
+  overrides `get_inference_client`); `BasePPOExp.get_inference_client` raises for any non-vllm
+  backend, and the client raises on weight-sync methods.
+- **Config** (all under `generator.inference_engine.*`): `run_engines_locally=false`,
+  `served_model_name` (the Fireworks model id, e.g. `accounts/fireworks/models/gpt-oss-20b`),
+  `api_key`, and `hf_tokenizer_name` (the served model's HF tokenizer, e.g. `openai/gpt-oss-20b`;
+  only settable with this backend — `EvalOnlyEntrypoint.get_tokenizer` loads it instead of
+  `trainer.policy.model.path`, which this backend does not use) are required; `external_proxy_url`
+  is optional and is the server root **without** `/v1` (the SDK appends `/v1/completions`;
+  defaults to `https://api.fireworks.ai/inference`).
+- **Tokenizer pairing**: `hf_tokenizer_name` must be the served model's tokenizer — token ids are
+  consumed raw by the server, so a mismatch degrades generations silently instead of erroring.
+- **Sampling params**: converted by `get_fireworks_sampling_params` — vLLM-only keys
+  (`min_tokens`, `skip_special_tokens`, `include_stop_str_in_output`) are dropped with a warning.
+  Verified against the live endpoint (gpt-oss-20b): a matched `stop` string is excluded from
+  Fireworks' `text` field but its tokens are **included in `token_ids`**, and a natural stop ends
+  with the EOS token id. Since the client builds `response_ids` from `token_ids` and decodes
+  `responses` locally, the effective behavior matches the vLLM path
+  (`include_stop_str_in_output=True`, `skip_special_tokens=True`), including
+  `append_eos_token_after_stop_str_in_multi_turn` detection. `min_tokens` is the only true gap
+  (no Fireworks equivalent; zero-length completions are possible).
+- **Logprobs**: supported. Integer `logprobs` returns the legacy `LogProbs` shape
+  (`token_logprobs`, verified live to align 1:1 with `token_ids`); the client also parses the
+  OpenAI chat-style `NewLogProbs.content` shape. If logprobs are requested but the response
+  carries neither shape, `generate` raises immediately (a silent `None` would surface as a
+  confusing length-validation failure downstream). Note that eval uses
+  `generator.eval_sampling_params.logprobs` (not `sampling_params.logprobs`), and the stock
+  generator keys logprobs tracking off the sampling params actually passed per generation.
+- Example: `examples/eval/run_eval_fireworks.sh`. Tests:
+  `tests/backends/skyrl_train/inference_servers/test_fireworks_client.py` (offline, mocked HTTP
+  transport through the real SDK).
+
 ## Key Config Knobs
 
 All under `generator.inference_engine.*`:

.github/workflows/cpu_skyrl_train.yaml

@@ -70,7 +70,7 @@ jobs:
         with:
           activate-environment: true
       - name: Run cpu tests (without vllm)
-        run:  uv run --isolated --extra skyrl-train --extra dev pytest tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "not vllm"
+        run:  uv run --isolated --extra skyrl-train --extra fireworks --extra dev pytest tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "not vllm"
       - name: Run cpu tests (with vllm)
         run: uv run --isolated --extra fsdp --extra dev pytest tests/train/ tests/backends/skyrl_train/ --ignore=tests/backends/skyrl_train/gpu -m "vllm"
 

examples/eval/run_eval_fireworks.sh

@@ -0,0 +1,36 @@
+set -x
+
+# Evaluation-only generation for GSM8K against the external Fireworks endpoint
+# (generator.inference_engine.backend=fireworks). No local inference engines and
+# no vLLM: prompts are sent as token ids and Fireworks returns the generated
+# token ids (return_token_ids), so the stock generator works unchanged.
+#
+# hf_tokenizer_name must be the served model's tokenizer (token ids are sent
+# raw), and served_model_name is the Fireworks model id.
+# trainer.policy.model.path is not used by this backend.
+
+# uv run examples/train/gsm8k/gsm8k_dataset.py --output_dir $HOME/data/gsm8k
+# export FIREWORKS_AI_API_KEY=<your_key_here>
+# bash examples/eval/run_eval_fireworks.sh
+
+: "${FIREWORKS_AI_API_KEY:?export FIREWORKS_AI_API_KEY first}"
+
+DATA_DIR="$HOME/data/gsm8k"
+TOKENIZER="openai/gpt-oss-20b"
+FW_MODEL="accounts/fireworks/models/gpt-oss-20b"
+LOGGER="console"
+
+uv run --isolated --extra fireworks \
+  -m skyrl.train.entrypoints.main_generate \
+  data.val_data="['$DATA_DIR/validation.parquet']" \
+  trainer.logger="$LOGGER" \
+  trainer.placement.colocate_all=false \
+  generator.inference_engine.backend=fireworks \
+  generator.inference_engine.run_engines_locally=false \
+  generator.inference_engine.served_model_name="$FW_MODEL" \
+  generator.inference_engine.hf_tokenizer_name="$TOKENIZER" \
+  generator.inference_engine.api_key="$FIREWORKS_AI_API_KEY" \
+  generator.eval_sampling_params.max_generate_length=2048 \
+  generator.eval_sampling_params.temperature=0.7 \
+  environment.env_class=gsm8k \
+  "$@"

pyproject.toml

@@ -157,6 +157,12 @@ miniswe = [
 harbor = [
     "harbor[daytona,modal]; python_version >= '3.12'",
 ]
+# External Fireworks generation backend (eval-only, used via `main_generate`). Exact
+# prerelease pin: the fireworks-ai 1.x line supports token-in/token-out natively.
+fireworks = [
+    "skyrl[skyrl-train]",
+    "fireworks-ai==1.2.0a85",
+]
 
 dev = [
     "mkdocs",

skyrl/backends/skyrl_train/inference_servers/engine_utils.py

@@ -1,6 +1,7 @@
 import os
 from typing import Any, Dict, Optional, Union
 
+from loguru import logger
 from omegaconf import DictConfig, ListConfig
 
 from skyrl.train.config import SamplingParams
@@ -82,8 +83,63 @@ def get_vllm_sampling_params(sampling_params: Union[SamplingParams, DictConfig])
     return vllm_sampling_params
 
 
+# vLLM-only sampling keys the Fireworks /completions API rejects. `skip_special_tokens` and
+# `include_stop_str_in_output` are unnecessary: FireworksInferenceClient decodes responses locally
+# from token ids, and Fireworks includes a matched stop string's tokens in `token_ids` (only its
+# `text` field, which the client ignores, excludes it). `min_tokens` has no counterpart.
+_FIREWORKS_UNSUPPORTED_KEYS = ("min_tokens", "skip_special_tokens", "include_stop_str_in_output")
+
+
+def get_fireworks_sampling_params(sampling_params: Union[SamplingParams, DictConfig]) -> Dict[str, Any]:
+    """Convert sampling params to the subset Fireworks' OpenAI-schema ``/completions`` accepts.
+
+    All sources are merged first (typed fields, then DictConfig keys / ``additional_kwargs``),
+    then the result is sanitized: vLLM-only keys are dropped with a warning, out-of-range
+    ``top_k`` values are dropped with a warning (Fireworks accepts ``0..100``), and the
+    ``top_k=-1`` / ``min_p=0.0`` disable-sentinels and ``None`` values are dropped silently
+    (absence disables them; ``None`` would serialize as ``null`` via ``extra_body``).
+    """
+    stop_val = sampling_params.stop
+    params: Dict[str, Any] = {
+        "max_tokens": sampling_params.max_generate_length,
+        "temperature": sampling_params.temperature,
+        "top_p": sampling_params.top_p,
+        "top_k": sampling_params.top_k,
+        "min_p": sampling_params.min_p,
+        "logprobs": sampling_params.logprobs,
+        "stop": list(stop_val) if stop_val is not None else None,
+    }
+    if isinstance(sampling_params, DictConfig):
+        exclude_keys = ["max_generate_length"]  # renamed to max_tokens above
+        for key, value in sampling_params.items():
+            if key not in params and key not in exclude_keys:
+                if isinstance(value, ListConfig):
+                    value = list(value)
+                params[key] = value
+    else:
+        if sampling_params.additional_kwargs is not None:
+            for key, value in sampling_params.additional_kwargs.items():
+                if key not in params:
+                    params[key] = value
+
+    for key in _FIREWORKS_UNSUPPORTED_KEYS:
+        if key in params:
+            logger.warning(f"Dropping sampling param `{key}`: not supported by the Fireworks completions API.")
+            del params[key]
+    top_k = params.get("top_k")
+    if top_k is not None and not (isinstance(top_k, int) and 0 <= top_k <= 100):
+        if top_k != -1:  # -1 is the disable sentinel; absence disables top_k on Fireworks
+            logger.warning(f"Dropping sampling param `top_k={top_k}`: Fireworks accepts 0..100.")
+        del params["top_k"]
+    if params.get("min_p") is not None and params["min_p"] <= 0:
+        del params["min_p"]
+    return {key: value for key, value in params.items() if value is not None}
+
+
 def get_sampling_params_for_backend(backend: str, sampling_params: Union[SamplingParams, DictConfig]) -> Dict[str, Any]:
     if backend == "vllm":
         return get_vllm_sampling_params(sampling_params)
+    elif backend == "fireworks":
+        return get_fireworks_sampling_params(sampling_params)
     else:
         raise ValueError(f"Unsupported generation backend: {backend}")

skyrl/backends/skyrl_train/inference_servers/fireworks_client.py

@@ -0,0 +1,217 @@
+"""External Fireworks inference client (generation/eval only).
+
+Fireworks' OpenAI-compatible ``/completions`` accepts a pre-tokenized integer-array ``prompt``
+and, with ``return_token_ids=true``, returns the generated integer ``token_ids``. This gives
+token-in/token-out against an external endpoint with no re-tokenization drift, so the stock
+``SkyRLGymGenerator`` works unchanged.
+
+Built on the Fireworks v1 SDK (``fireworks-ai``, installed via the ``fireworks`` uv extra):
+``prompt`` accepts ``Iterable[Iterable[int]]``, ``return_token_ids`` is a first-class request
+param, and the response ``Choice`` declares ``token_ids``/``prompt_token_ids``. The SDK carries
+auth, retries with backoff, timeouts, and connection pooling. This module does not import vllm
+and has no control plane: wake/sleep/etc. are no-ops, weight sync raises. Only the eval-only
+entrypoint builds this client (training entrypoints reject ``backend='fireworks'``).
+"""
+
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import httpx
+from fireworks import AsyncFireworks
+
+from skyrl.backends.skyrl_train.inference_servers.base import (
+    InferenceEngineInput,
+    InferenceEngineInterface,
+    InferenceEngineOutput,
+)
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase
+
+_GEN_EVAL_ONLY = (
+    "Fireworks is a generation/eval-only backend (external hosted endpoint with no weight sync). "
+    "Use backend='vllm' for training."
+)
+
+# Server root of the Fireworks data plane. The SDK appends `/v1/completions` to an overridden
+# base_url, so this must NOT end in `/v1` (validate_inference_engine_cfg guards user-supplied
+# values).
+DEFAULT_FIREWORKS_BASE_URL = "https://api.fireworks.ai/inference"
+
+
+class FireworksInferenceClient(InferenceEngineInterface):
+    def __init__(
+        self,
+        model_name: str,
+        tokenizer: "PreTrainedTokenizerBase",
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        max_retries: int = 3,
+        request_timeout: float = 600.0,
+        *,
+        _http_client: Optional[httpx.AsyncClient] = None,
+    ):
+        """Args:
+        model_name: Fireworks model id used as the request ``model`` (e.g.
+            ``accounts/fireworks/models/gpt-oss-20b``).
+        tokenizer: The policy tokenizer; must be the served model's tokenizer since prompts
+            are sent as raw token ids.
+        base_url: Server root without ``/v1`` (defaults to the Fireworks data plane).
+        api_key: API key sent as ``Authorization: Bearer``. Always passed explicitly so the
+            SDK never falls back to the ``FIREWORKS_API_KEY`` env var; ``"EMPTY"`` placeholder
+            keeps keyless self-hosted endpoints constructible.
+        max_retries: SDK retry budget (backoff on 408/409/429/5xx and ``x-should-retry``).
+        request_timeout: Per-request timeout in seconds. Overrides the SDK's 60s default,
+            which is too short for long generations.
+        _http_client: Internal-reserved injectable httpx client, used by tests with
+            ``httpx.MockTransport``.
+        """
+        self._base_url = (base_url or DEFAULT_FIREWORKS_BASE_URL).rstrip("/")
+        self._model_name = model_name
+        self._tokenizer = tokenizer
+        self._client = AsyncFireworks(
+            base_url=self._base_url,
+            api_key=api_key or "EMPTY",
+            max_retries=max_retries,
+            timeout=request_timeout,
+            http_client=_http_client,
+        )
+
+    @property
+    def model_name(self) -> str:
+        return self._model_name
+
+    def get_endpoint_url(self) -> str:
+        return self._base_url
+
+    async def generate(
+        self,
+        input_batch: InferenceEngineInput,
+        model: Optional[str] = None,
+    ) -> InferenceEngineOutput:
+        prompt_token_ids = input_batch.get("prompt_token_ids")
+        if prompt_token_ids is None:
+            raise ValueError("FireworksInferenceClient only accepts `prompt_token_ids`, not `prompts`.")
+        if input_batch.get("mm_features"):
+            raise NotImplementedError("FireworksInferenceClient does not support multi-modal features.")
+
+        sampling_params = dict(input_batch.get("sampling_params") or {})
+        if sampling_params.get("n", 1) > 1:
+            raise ValueError("n > 1 is not supported. Use `config.generator.n_samples_per_prompt` instead.")
+        want_logprobs = sampling_params.get("logprobs") is not None
+
+        # model/prompt/return_token_ids are typed SDK params; the sampling dict rides extra_body
+        # (merged into the JSON request body) so additional_kwargs passthrough keeps working.
+        completion = await self._client.completions.create(
+            model=model or self._model_name,
+            prompt=prompt_token_ids,
+            return_token_ids=True,
+            extra_body=sampling_params,
+        )
+        if not completion.choices:
+            raise RuntimeError(f"Fireworks returned no choices: {completion!r}")
+        choices = sorted(completion.choices, key=lambda choice: choice.index)
+
+        response_ids: List[List[int]] = []
+        responses: List[str] = []
+        stop_reasons: List[str] = []
+        response_logprobs: List[Optional[List[float]]] = []
+        for choice in choices:
+            token_ids = choice.token_ids
+            # Re-encoding `choice.text` locally would silently reintroduce the retokenization
+            # drift this backend exists to avoid, so a missing field is a hard error.
+            assert token_ids is not None, (
+                f"Fireworks response missing `token_ids` for choice {choice.index} despite " "return_token_ids=true."
+            )
+            response_ids.append(list(token_ids))
+            # Decode locally to guarantee the InferenceEngineOutput invariant:
+            # tokenizer.decode(response_ids[i], skip_special_tokens=True) == responses[i].
+            responses.append(self._tokenizer.decode(token_ids, skip_special_tokens=True))
+            stop_reasons.append(choice.finish_reason or "stop")
+            if want_logprobs:
+                logprobs = self._extract_logprobs(choice)
+                # Silently emitting None here would surface far downstream as a confusing
+                # length-validation failure on GeneratorOutput["rollout_logprobs"].
+                if logprobs is None:
+                    raise RuntimeError(
+                        f"Sampling params requested logprobs but Fireworks returned none (or an "
+                        f"unrecognized shape) for choice {choice.index}. Set "
+                        f"generator.eval_sampling_params.logprobs=null (and "
+                        f"generator.sampling_params.logprobs=null) if logprobs are not needed."
+                    )
+                response_logprobs.append(logprobs)
+
+        return InferenceEngineOutput(
+            responses=responses,
+            response_ids=response_ids,
+            stop_reasons=stop_reasons,
+            response_logprobs=response_logprobs if want_logprobs else None,
+            prompt_logprobs=None,
+            rollout_expert_indices=None,
+        )
+
+    @staticmethod
+    def _extract_logprobs(choice: Any) -> Optional[List[float]]:
+        """Extract per-token logprobs from either Fireworks response shape.
+
+        ``choice.logprobs`` is a union of the classic completions shape (``LogProbs``, carries
+        ``token_logprobs``; what the live endpoint returns for integer ``logprobs``) and the
+        OpenAI chat-style shape (``NewLogProbs``, carries ``content`` items with ``.logprob``).
+        Null entries map to 0.0 rather than being dropped so the result stays aligned 1:1 with
+        the generated token ids (downstream validation asserts equal lengths).
+        """
+        logprobs = choice.logprobs
+        if logprobs is None:
+            return None
+        token_logprobs = getattr(logprobs, "token_logprobs", None)
+        if token_logprobs is not None:
+            return [logprob if logprob is not None else 0.0 for logprob in token_logprobs]
+        content = getattr(logprobs, "content", None)
+        if content is not None:
+            return [item.logprob if item.logprob is not None else 0.0 for item in content]
+        return None
+
+    async def completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
+        response = await self._client.completions.create(**request_payload.get("json", {}))
+        return response.model_dump()
+
+    async def chat_completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
+        response = await self._client.chat.completions.create(**request_payload.get("json", {}))
+        return response.model_dump()
+
+    async def render_chat_completion(self, request_payload: Dict[str, Any]) -> Dict[str, Any]:
+        raise NotImplementedError("render_chat_completion is not supported for the Fireworks backend.")
+
+    async def wake_up(self, *args: Any, **kwargs: Any):
+        # TODO: tokenizer handshake — probe with a text prompt + return_token_ids=true and
+        # compare the returned prompt_token_ids against a local tokenizer.encode() of the same
+        # string. A mispaired hf_tokenizer_name / served_model_name fails silently otherwise:
+        # the server consumes raw token ids, so a wrong tokenizer yields degraded generations,
+        # never an error.
+        return {}
+
+    async def sleep(self, *args: Any, **kwargs: Any):
+        return {}
+
+    async def reset_prefix_cache(self, reset_running_requests: bool = False):
+        return {}
+
+    async def pause_generation(self) -> None:
+        return
+
+    async def resume_generation(self) -> None:
+        return
+
+    async def finish_session(self, session_id: str) -> None:
+        return
+
+    async def teardown(self):
+        await self._client.close()
+
+    async def get_world_size(self) -> Tuple[int, int]:
+        raise NotImplementedError(_GEN_EVAL_ONLY)
+
+    async def init_weight_update_communicator(self, init_info):
+        raise NotImplementedError(_GEN_EVAL_ONLY)
+
+    async def update_named_weights(self, request):
+        raise NotImplementedError(_GEN_EVAL_ONLY)