Add PiSSA initialization for Megatron LoRA adapters - #1
Closed
kalectory wants to merge 11 commits into
Closed
Conversation
j316chuck
reviewed
Jun 30, 2026
j316chuck
approved these changes
Jun 30, 2026
…SkyRL's lock in E2E CI (NovaSky-AI#2022) ## What does this PR do? Fixes the nightly `gsm8k_tinker` / `gsm8k_tinker_fully_async` failures that started 2026-08-12: ``` ValueError: Server returned a JSON payload for self.model_cls=<class 'tinker.SampleResponse'> ..., which this SDK version only supports as proto. ``` **tinker 0.25.0 (released 2026-08-08) made its proto wire format mandatory**: it submits `forward_backward` request bodies as protobuf and rejects JSON `retrieve_future` results for `SampleResponse`/`ForwardBackwardOutput`. The SkyRL tinker server only speaks JSON today. The nightlies install the SDK unpinned into tinker-cookbook's own environment (which ships no lockfile), so every run floated to 0.25.0 and died at the first sample retrieval. Per review feedback, the version constraint lives in `pyproject.toml` rather than hardcoded in the CI scripts. Two commits: 1. **Cap the tinker SDK in the `tinker` extra** at `<=0.24.1` (the newest release that negotiates down to JSON: it sends `Accept: application/x-protobuf, application/json`, falls back to JSON responses, and its proto request path stays off because it is gated by a server-supplied `client_config` flag our server does not enable). The lock keeps main's tinker 0.24.0, which satisfies the cap; only the recorded requirement specifier changes. (Bumping the locked version to 0.24.1 requires a resolver run, which is currently blocked on main by an unrelated py3.14 megatron-split conflict between `vllm 0.26.0+cu129` and `mamba-ssm` over `apache-tvm-ffi` -- any full `uv lock` fails on main today. `uv lock --check` validates the updated lock.) 2. **Make the nightlies resolve the SDK from SkyRL's lock.** The scripts previously ran the cookbook client from `~/tinker-cookbook`'s own project with an unpinned `--with tinker`, so SkyRL's `pyproject.toml` never reached that environment. They now run the client from SkyRL's project, overlaying the cookbook checkout with `--with-editable "$COOKBOOK_DIR[math-rl,wandb]"`: the cookbook's `tinker>=0.9.0` requirement is already satisfied by the base environment's locked SDK, so uv does not fetch a newer one. The SDK version now has a single source of truth and moves only when `pyproject.toml`/`uv.lock` move. Once server-side proto support (NovaSky-AI#2021) lands, the cap can be lifted and the nightlies will follow the lock to newer SDKs automatically. ## Testing - Resolution verified locally at the pinned cookbook commit: the old invocation (`cd ~/tinker-cookbook && uv run --extra math-rl --with tinker ...`) resolves tinker 0.25.0; the new invocation resolves the locked 0.24.0 from SkyRL's lock, and the `math_rl.train` recipe plus wandb/datasets/torch all import cleanly in that environment. - `tests/tinker/` CPU suite (68 tests on the rebased branch, including the `test_api.py` integration tests that run the real SDK against a real server subprocess) passes with the locked SDK. - SDK 0.24.x's compatibility with the current JSON-only server was verified manually with a full `forward_backward` -> `forward` -> `optim_step` -> `save_weights_for_sampler` -> `sample` loop (JAX backend); the same loop on 0.25.0 reproduces the nightly failure exactly. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: Avi Basnet <avigyabb@stanford.edu> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Add custom wheels for aarch64/GB200. <img width="855" height="172" alt="image" src="https://github.com/user-attachments/assets/01e4e01b-de00-48e4-95c9-f272310ef49b" /> Tested gsm8k runs for both megatron and fsdp backends.
…rom Astral (NovaSky-AI#2039) # Upgrade Transformer Engine to 2.16 and source flash-attn + TE from Astral Follow-up to NovaSky-AI#2023. That PR moved `causal-conv1d` and `mamba-ssm` to Astral's GPU wheel index but had to leave `flash-attn` and `transformer-engine-torch` on hand-built fork wheels, because neither was usable with the Transformer Engine 2.11.0 we pinned. This upgrades the TE stack to 2.16 and moves both to Astral, so **no fork-built wheels remain** — on x86_64 or aarch64. ## Why TE 2.16 is the unlock Astral stamps every wheel with the (CUDA, torch) pair it was built against as a PEP 440 local version segment, e.g. `+cu.12.8.torch.2.11`. TE 2.11.0 could not tolerate that, in two independent ways: | Package | Why it failed on TE 2.11 | Why TE 2.16 fixes it | | --- | --- | --- | | `flash-attn` | PEP 440 makes `2.8.3+local > 2.8.3`, so TE's `get_attention_backend` gate (`max_version="2.8.3"`) rejected the wheel and **silently** fell back to the broken unfused attention path | TE strips the local segment before comparing (`PkgVersion(...).public`), added in 2.16 | | `transformer-engine-torch` | TE's meta package asserts **raw string equality** across `transformer-engine` / `-cu12` / `-torch` versions, which Astral's differing local segments break | Astral's meta wheel patches that check to compare base versions — a patch they ship only at 2.16 | The second point is why the TE trio has to move to Astral **as a set**: the patch lives in Astral's meta wheel, not upstream's. Mixing Astral's core/torch wheels with the PyPI meta package fails the assert at import. Verified in the built environment: ``` transformer-engine 2.16.0 transformer-engine-cu12 2.16.0+cu.12.8 transformer-engine-torch 2.16.0+cu.12.8.torch.2.11 flash-attn 2.8.3+cu.12.8.torch.2.11 fa.version=2.8.3 max_version=2.8.3 is_installed=True ``` `is_installed=True` is the key line — that is TE accepting flash-attn rather than silently disabling it. There are also zero `"Supported flash-attn versions are ..."` warnings (TE's rejection message) anywhere in the megatron suite output. ## Also in this PR - **Dead build config removed.** TE and flash-attn now always install as prebuilt wheels and are never compiled, so their `no-build-isolation-package`, `extra-build-dependencies`, and `extra-build-variables` entries are dead. Confirmed inert: removing them leaves `uv.lock` byte-identical, and both packages resolve to 8 wheels / 0 sdists. - **`transformer-engine-cu13` override no longer needed.** Astral's 2.16 meta package resolves its `core` extra to cu12 rather than cu13, so cu13 never enters the graph — it drops from 2 lock references to 0. - **aarch64 fork wheels retired.** NovaSky-AI#1936 added hand-built aarch64 artifacts for flash-attn and transformer-engine-torch for GB200. Astral publishes cp310–cp314 for both x86_64 and aarch64, so a single `sys_platform == 'linux'` marker covers both arches and those artifacts are no longer referenced. **See the caveat below.** ## Resolution delta Diffing resolved package versions against `main` — exactly 5 change, nothing else moves: ``` flash-attn 2.8.3 -> 2.8.3+cu.12.8.torch.2.11 transformer-engine 2.11.0 -> 2.16.0 transformer-engine-cu12 2.11.0 -> 2.16.0+cu.12.8 transformer-engine-torch 2.11.0 -> 2.16.0+cu.12.8.torch.2.11 nvdlfw-inspect (new) -> 0.2.2 # new TE 2.16 dependency ``` ## Testing All suites run on 4×H100: | Suite | Result | Time | | --- | --- | --- | | `-m megatron` | 130 passed, 2 skipped | 2:01 | | `-m megatron_models` | 3 passed, 4 skipped | 0:12 | | H100 CI (megatron half) | 5 passed | 0:25 | | H100 CI (fsdp half) | 2 passed | 0:11 | All skips are pre-existing static markers (`h100`-gated tests, and the existing `qwen3.5-moe` correctness skip), not anything this change introduces. Two things checked that turned out **not** to be regressions: megatron-core's `_SplitAlongDim` and `get_workspace` imports fail under TE 2.16 — but they fail identically under 2.11, so the `torch.split` fallbacks were already in use before this PR. ## Caveat for reviewers The aarch64/GB200 path is **untested**. This PR replaces the sm_100 aarch64 fork wheels from NovaSky-AI#1936 with Astral's aarch64 builds, and I could only test x86_64/H100. I have not confirmed Astral's aarch64 wheels include sm_100 support. If you'd rather not take that on here, the aarch64 fork entries can be kept and this PR scoped to x86_64 — say so and I'll split it. Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ovaSky-AI#2017) ## Problem `retrieve_future` returns `result_data`, which for sample requests is a large numeric blob — top-k logprobs for every prompt token run to a few MB. On the way from the engine to the caller, that blob was converted between Python objects and JSON text three extra times: 1. **Engine write** — `result.model_dump()` builds a `dict` of millions of floats, which `json.dumps` then walks again to produce the text SQLAlchemy stores. 2. **API read** — SQLAlchemy's JSON column calls `json.loads`, re-materializing every one of those floats as a Python object. 3. **API response** — FastAPI walks the whole structure with `jsonable_encoder`, then `json.dumps` it back to text. None of that work is used: the payload comes out of Pydantic and goes to the client unchanged. It's a few hundred ms of event-loop time per call, and because it's synchronous CPU work in the event loop, every other tenant queues behind it. ## Root cause Every one of those conversions is imposed by the column type. SQLAlchemy's `JSON` decodes on read and encodes on write by contract, so the application cannot opt out while `result_data` is declared as JSON. The fix is to stop declaring it as JSON. ## Changes - **`FutureDB.result_data` is now a `TEXT` column** holding pre-serialized JSON. - **Writers hand it a typed model.** The engine and both sample-forwarding clients serialize with `model_dump_json()`, which goes from model to JSON text inside pydantic-core without materializing the intermediate dict. The clients' failure paths build a `types.ErrorResponse` instead of an equivalent hand-rolled dict, so every write to the column goes through a typed model. Serialization sits at the single DB-write site in each client, which also means a payload for a row deleted mid-flight is never encoded at all. - **`retrieve_future` returns the stored text directly** in a `Response`, bypassing `jsonable_encoder` and a second `json.dumps`. The failure path still decodes, since it reports the error field — but as a `types.ErrorResponse` rather than by poking at a dict, which also turns malformed JSON there into the intended 500 instead of an unhandled `JSONDecodeError`. ## Performance Measured on a 2.4 MB result (4 sequences × 512 tokens, 4096 prompt tokens with top-20 prompt logprobs), best of 5: | | write | read + respond | total | |---|---|---|---| | before | 61.7 ms | 260.7 ms | **322.4 ms** | | after | 7.4 ms | 0.0 ms | **7.4 ms** | Worth noting for anyone who would reach for a faster JSON library first: `jsonable_encoder` alone was 168 ms of that 260 ms read path, and no encoder swap touches it. Dropping in orjson while keeping the JSON column lands at ~188 ms total. The win comes from not walking the payload at all. ## Compatibility Existing SQLite databases are unaffected — SQLAlchemy stores JSON columns as TEXT there. A Postgres deployment predating this change would need the column altered. There are no Alembic revisions in the tree today; the schema comes from `SQLModel.metadata.create_all` at startup. One narrow behavior change: a `FAILED` row whose payload has an `error` but no `status` now returns 500 rather than 400. No writer produces that shape. ## Testing `tests/tinker/` — 48 passed / 19 skipped; `tests/tinker/test_api.py` integration — 17 passed. The tests in `test_future_waiting.py` now store and assert on real result types instead of ad-hoc dicts, and the `retrieve_future` assertion is byte-exact so it actually catches a re-encode rather than passing on any equivalent JSON. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Changes the core async result path and DB column type (Postgres may need migration); behavior is low-risk for SQLite and correctness is covered by updated tests, but mis-serialized rows could break clients. > > **Overview** > Large sample results (e.g. multi‑MB prompt logprobs) no longer round-trip through Python dicts and FastAPI’s JSON encoder on every `retrieve_future` call. > > **`FutureDB.result_data`** is now a `TEXT` column holding pre-serialized JSON. The engine and inference forwarding clients write with `model_dump_json()`; failed paths use `types.ErrorResponse` instead of ad-hoc dicts. > > **`retrieve_future`** returns completed payloads via `raw_json_response()` so the stored bytes go straight to the client. `wait_for_future` / `poll_futures` keep `result_data` as JSON text end-to-end. Failed futures still decode with `ErrorResponse.model_validate_json` for 400 vs 500 handling. > > Tests assert real output types and byte-exact response bodies so a accidental re-encode would fail. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 1f1d77f. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
- Expose SkyRL's existing native GSPO implementation through the Tinker-compatible API. - Accept Tinker-style clipping thresholds and normalize them to SkyRL's epsilon configuration. ## Future Work - This path will fail on Jax, similar to the existing exposed DPPO and PPO_CRITIC paths. We generalize the PPO_CRITIC check to fail loudly on those paths too for now. ## Testing <img width="1870" height="854" alt="image" src="https://github.com/user-attachments/assets/9f0a526d-c0b6-4aaa-911d-4e5502429c1b" /> ``` uv run --extra dev pre-commit run --files skyrl/backends/skyrl_train_backend.py skyrl/tinker/api.py skyrl/tinker/types.py ``` Verified formatting, lint, and secret checks on all changed files. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Low Risk** > Changes are API surface and backend routing for a new loss_fn name; JAX fails fast on unsupported losses. No auth or data-path changes; training behavior on SkyRL-Train depends on existing native GSPO implementation. > > **Overview** > Adds **`gspo`** as a Tinker-compatible `loss_fn` on forward/backward, alongside existing PPO-style **`clip_low_threshold`** / **`clip_high_threshold`** config validation in the API layer. > > The SkyRL-Train backend now treats **`gspo`** like **`ppo`** for threshold normalization (Tinker ratios → SkyRL epsilon fields) but maps the internal loss name to **`gspo`** instead of **`regular`**. > > On JAX, unsupported losses are rejected with a single check against **`LOSS_TYPES`** (so **`gspo`**, **`dppo`**, and **`ppo_critic`** fail loudly). Tests cover that behavior via **`JaxBackendImpl._model_pass`**. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit efb4101. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Signed-off-by: Neil Kale <263453039+kalectory@users.noreply.github.com> Co-authored-by: Neil Kale <263453039+kalectory@users.noreply.github.com>
…I#2050) applying NVIDIA/TransformerEngine#3360 to TE 2.16 <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Changes attention backend selection at Megatron init via `exec` on TE internals; mitigated by SM guards, literal 2.16.0 source matching, and idempotent no-ops on unaffected GPUs, but wrong TE versions or arch-specific FA2 issues could still affect training memory or numerics. > > **Overview** > Backports **NVIDIA/TransformerEngine#3360** for the pinned **transformer-engine 2.16.0** so Megatron training can select **FlashAttention 2** for **head_dim 256** (e.g. Gemma 2/3) on GPUs outside TE’s SM allowlist (notably **sm103** B300/GB300, also sm86/sm89). > > A runtime patch **recompiles and rebinds** TE’s `get_attention_backend` to drop the erroneous `head_dim > 192` + compute-capability gate, keeping only FA2’s real limits (`<= 256`, `% 8 == 0`). It **no-ops** on sm80/90/100/120, when TE isn’t importable, or when TE source no longer matches 2.16.0, and clears TE’s memoized `_attention_backends` after apply. > > **Megatron only:** `patch_fa2_head_dim_allowlist()` runs in `make_megatron_module()` before `provide_distributed_model()` (policy and ref workers). A **verify script** and README document repro and numerics vs unfused attention. > > Remove this patch when the TE pin includes upstream #3360. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 1d9d8c9. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
… infra log on failure (NovaSky-AI#2051) # Fix colocated tinker E2E nightly: vLLM startup OOM on L4 The `SkyRL-GPU-E2E-CI-Tinker` nightly has failed on every run since 2026-08-14 (the first nightly after the vLLM 0.26.0 bump in NovaSky-AI#1854) with: ``` RuntimeError: Engine core initialization failed. See root cause above. Failed core proc(s): {} ``` raised from `VLLMServerActor.start()` at the first `save_weights_for_sampler`. (Failures on 08-08..08-13 were the separate unpinned tinker SDK break, fixed by NovaSky-AI#2022.) ## Root cause Reproduced on a 1x L4 workspace with a scaled-down (1-GPU) version of the CI backend config. The real error never reaches the job log (see "Log visibility" below); it is a CUDA OOM during vLLM engine startup, at the flashinfer sampler warmup that runs after CUDA graph capture: ``` torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 594.00 MiB. GPU 0 has a total capacity of 22.03 GiB of which 582.00 MiB is free. Process 29529 has 796.00 MiB memory in use. ... this process has 20.67 GiB in use. ``` Two factors combine: 1. **Engine startup ordering.** The main trainer entrypoint starts inference engines on empty GPUs before `build_models()`. The tinker `SkyRLTrainBackend` does the reverse: FSDP workers initialize first, and even after `offload_to_cpu()` each leaves ~800 MiB of unreclaimable CUDA context on the GPU that vLLM later profiles on. `gpu_memory_utilization` budgets a fraction of *total* (not free) memory, so vLLM still sizes its KV cache as if it had the whole card. 2. **vLLM 0.26.0's larger startup peak.** FULL_AND_PIECEWISE cudagraph capture plus the flashinfer sampler warmup transiently allocate ~3 GiB beyond the utilization budget. At 0.8 on a 22 GiB L4 with the FSDP context resident, this misses by ~12 MiB. This explains why only this nightly regressed: the fully-async tinker E2E is non-colocated (engines get empty GPUs), and the main colocated E2E starts engines before building models. ## Changes - `gsm8k_tinker.sh`: `gpu_memory_utilization` 0.8 → 0.7. Costs ~2 GiB of KV cache, irrelevant for 512-token GSM8K rollouts. - Both tinker E2E scripts: on any non-zero exit, the cleanup trap now dumps the tails of `server.log` **and** the newest `/tmp/skyrl-logs/infra-*.log`. ### Log visibility `VLLMServerActor` calls `redirect_actor_output_to_file()`, which sends actor output — including the vLLM engine's real traceback — to `/tmp/skyrl-logs/infra-*.log` on the cluster. The scripts previously only tailed `server.log`, and only when the server failed to boot, so a client-visible failure like this one showed nothing but the opaque re-raised `RayTaskError` (the empty `Failed core proc(s): {}` is a vLLM race where the engine-core proc dies before its exit code is captured). ## Verification On a 1x L4 with the 1-GPU equivalent of the CI backend config (`colocate_all=true`, 1 engine): - at `gpu_memory_utilization=0.8`, `save_weights_and_get_sampling_client()` fails with the exact CI error (OOM in flashinfer sampler warmup); - at `0.7`, the same call succeeds: engine starts, weights sync, KV cache wakes at 18.5/22 GiB. ## Note Any tinker + `colocate_all` user on ~24 GB GPUs will hit this same OOM at the default 0.8 utilization, since the lazy engine startup ordering is inherent to the backend. A follow-up could account for the training workers' CUDA context when sizing the KV budget, or start engines eagerly once the LoRA config is known. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Low Risk** > CI-only shell script changes; no production trainer or auth paths touched. > > **Overview** > Fixes the **colocated** tinker GSM8K nightly (`gsm8k_tinker.sh`) by lowering vLLM **`gpu_memory_utilization` from 0.8 to 0.7**, with comments explaining lazy engine startup after FSDP leaves CUDA context and vLLM 0.26’s larger startup peak OOMing 22 GiB L4s. The fully-async script keeps **0.8** (non-colocated engines). > > **Both** tinker E2E scripts replace the simple EXIT trap with a **`cleanup()`** that on any non-zero exit tails **`server.log`** and the newest **`/tmp/skyrl-logs/infra-*.log`** (Ray/vLLM errors) before tearing down the server process group—so CI logs show real OOM/tracebacks instead of opaque `RayTaskError`. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit 77ff96c. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> Signed-off-by: Avi Basnet <avigyabb@stanford.edu> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…ovaSky-AI#2021) ## What does this PR do? Fixes the nightly tinker cookbook e2e failures that started 2026-08-12: ``` ValueError: Server returned a JSON payload for self.model_cls=<class 'tinker.SampleResponse'> ..., which this SDK version only supports as proto. ``` The e2e scripts install the tinker SDK unpinned, and **tinker 0.25.0 (released 2026-08-08) made its proto wire paths mandatory**: 1. **Responses**: `retrieve_future` results for `sample` / `forward` / `forward_backward` must be proto-serialized when the client sends `Accept: application/x-protobuf`. Since 0.18.0 the SDK has sent `Accept: application/x-protobuf, application/json` and deserialized whichever format the server returned — JSON was a true client-side fallback, which is why our JSON-only server worked. 0.25.0 removed that fallback for these two response types. (SDKs <0.18.0 have no proto module at all and speak pure JSON.) 2. **Requests**: `forward_backward` bodies are submitted as proto (`Content-Type: application/x-protobuf`), and forward-only passes are routed to `/api/v1/forward_backward` via the proto `forward_only` flag instead of `/api/v1/forward`. Unlike the response path, this was never a fallback arrangement before 0.25.0: the proto write path (present since ~0.22) was gated by the server-supplied `client_config` flag `proto_write_fwdbwd` (default false), and our stub `client_config` never opted in, so JSON requests were simply the default. 0.25.0 removed the flag and made proto requests unconditional. ## Changes - **`skyrl/tinker/proto_serialization.py`** (new): mirrors the SDK's `tinker/proto/{request,response}_conv.py`. Serializers for `SampleResponse` / `ForwardBackwardOutput` (NaN encodes undefined prompt logprobs, `token_id=0 / logprob=-99999.0` sentinel fill for undefined top-k entries, batched int64 byte offsets for per-datum tensors), plus a parser for proto `ForwardBackwardRequest` bodies. The proto schema ships inside the `tinker` package, which the `tinker` extra already depends on -- no new dependency. - **`skyrl/tinker/api.py`**: `retrieve_future` content-negotiates on the `Accept` header (errors and all other result types stay JSON); `forward_backward` accepts both wire formats and maps `forward_only=true` to the `FORWARD` request type. `poll_futures`/`wait_for_future` now carry `request_type` alongside `(status, result_data)`. - Old JSON paths (including `/api/v1/forward`) are unchanged, so SDKs <=0.24.x keep working. ## Testing - New `tests/tinker/test_proto_serialization.py` round-trips the server encoding through the installed SDK's own proto codecs (12 tests), plus a proto-path test in `test_future_waiting.py`. - Manually verified a full loop (`forward_backward` -> `forward` -> `optim_step` -> `save_weights_for_sampler` -> `sample`) against a local JAX-backend server with **12 SDK versions spanning the supported range** (`tinker>=0.3.0`) -- all pass with this change; before it, 0.25.0 failed exactly like the nightly: | SDK versions | Wire behavior | Result | |---|---|---| | 0.3.0, 0.8.1, 0.13.1, 0.16.1 | pure JSON (no proto module) | ✅ | | 0.18.2, 0.20.0, 0.21.0, 0.22.0, 0.22.4, 0.23.4, 0.24.1 | proto responses (JSON-fallback), JSON requests | ✅ | | 0.25.0 | proto responses + proto requests, no fallback | ✅ | - `tests/tinker/` CPU suite passes (`skyrl_train/` GPU-cluster tests excluded -- they fail in my workspace with a pre-existing Ray version mismatch unrelated to this change). Notes for reviewers: - SDK >=0.18.0 sends `Accept: application/x-protobuf, application/json` on `retrieve_future` and understands proto responses, so the existing `test_api.py` integration tests (locked SDK 0.22.4) now exercise the proto response path end-to-end. - The SDK's zstd request compression (`proto_compress_fwdbwd`) is negotiated via `client_config`, which our server does not advertise, so compressed bodies never arrive. - I left the e2e scripts' `--with tinker` unpinned on purpose: catching SDK/server incompatibilities is exactly what those nightlies are for. Happy to pin to `tinker==0.25.0` instead if determinism is preferred. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Medium Risk** > Changes the Tinker API wire contract on hot training paths (forward/backward and sampling) with binary encoding; mistakes could break all 0.25.0 clients while JSON fallbacks must stay correct for older SDKs. > > **Overview** > Fixes nightly cookbook failures against **tinker SDK 0.25.0**, which requires protobuf for `sample` / `forward` / `forward_backward` results on `retrieve_future` and sends `forward_backward` bodies as protobuf (including forward-only via a `forward_only` flag). > > A new **`proto_serialization`** module mirrors the SDK’s proto codecs: it parses proto `ForwardBackwardRequest` bodies and serializes stored JSON results into `SampleResponse` and `ForwardBackwardOutput` wire bytes (tokens, logprobs, top-k sentinels, batched tensors). > > **`retrieve_future`** negotiates on `Accept: application/x-protobuf` for SAMPLE, EXTERNAL, FORWARD, and FORWARD_BACKWARD; errors and other futures stay JSON. **`forward_backward`** reads JSON or protobuf and maps `forward_only=true` to the FORWARD request type. Future polling now passes **`request_type`** through waiters. JSON **`/forward`** and legacy SDK paths are unchanged. > > Round-trip tests use the installed tinker SDK deserializers; **`EXTERNAL`** sample futures (external inference) serialize like **`SAMPLE`**. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit bd8043b. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> ## GPU e2e validation (2026-08-17) Both CI e2e tests were run as Anyscale jobs on this branch (4xL4, exact CI configs, cookbook client installing the **unpinned** tinker SDK -> 0.25.0, the strict proto-only client): - `gsm8k_tinker` (colocated, FSDP backend): **SUCCEEDED** (55m50s) - `gsm8k_tinker_fully_async` (non-colocated): **SUCCEEDED** (56m28s) The first attempt surfaced a real gap the local JAX-backend testing missed: `/asample` stores futures as `request_type=EXTERNAL` (not `SAMPLE`) whenever an external inference client is configured -- which is how the fsdp/megatron backends run sampling -- and those results bypassed the proto gate, reproducing the original failure. Fixed in `7c72fff0` (EXTERNAL results serialize as `SampleResponse`; failures still return JSON errors), with a regression unit test. Both jobs passed on the fixed commit, including the wandb reward assertions. ## Performance Microbenchmark of the response paths at CI-realistic payload shapes (single CPU core, round-tripped through the SDK's own deserializers; script: `serialize_result` vs the raw-JSON passthrough): | Payload | Wire size (JSON -> proto) | Server cost per result | Client parse (JSON -> proto) | |---|---|---|---| | sample, 4 seqs x 512 tokens (CI shape) | 51.5 KiB -> 16 KiB (3.2x) | +0.78 ms | 0.57 ms -> 0.02 ms | | sample + top-20 prompt logprobs, 2048-token prompt | 1.17 MiB -> 344 KiB (3.5x) | +34 ms | 21.9 ms -> 0.07 ms | | fwd_bwd, 512 datums x 600 tokens x 2 float fields | 11.4 MiB -> 2.4 MiB (4.8x) | +170 ms | 144.9 ms -> 2.8 ms | - **Client + wire: strictly faster.** Proto payloads are 3-5x smaller and client-side deserialization is `np.frombuffer` over binary instead of JSON-parsing decimal floats (~50x faster on large payloads; the JSON column above is `json.loads` only, a lower bound -- the SDK adds model construction on top). - **Server: a bounded regression vs the current JSON path, which NovaSky-AI#2017 made a zero-copy passthrough.** The proto path pays `json.loads(stored text) + validate + proto encode` per result: sub-millisecond for typical samples, ~170 ms worst-case for a full 512-datum fwd_bwd result. ~85% of that is the `json.loads` of the stored JSON text (a storage-format artifact, not a proto cost); the pre-NovaSky-AI#2017 baseline (`jsonable_encoder` + `json.dumps`) was ~300 ms on the same payload class, so this is still cheaper than the JSON path of two weeks ago. The SDK also chunks fwd_bwd requests (~5 MB estimated per chunk), so single results near the worst case are uncommon. - **End-to-end: indistinguishable.** The GPU e2e runs are GPU-bound; proto-client (0.25.0) and JSON-client (0.24.x) runs both complete in 55-56 min. Possible follow-ups (not in this PR): offload `json.loads + serialize` to a thread in `retrieve_future` so a worst-case result doesn't stall the event loop; longer term, store proto bytes in `FutureDB.result_data` for the four proto-serializable request types, which removes the parse entirely and makes every SDK >= 0.18.0 (all of which prefer proto via the Accept header) a zero-copy read. --------- Signed-off-by: Avi Basnet <avigyabb@stanford.edu> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
# What does this PR do? Fixes CI failures for the `gsm8k_fully_async_ci` fully async RL CI. Thresholds have been rebaselined after the async codepath changes: 1. stale KV cache reuse on weight sync (NovaSky-AI#1798), 2. policy_loss_type=rollout_is (NovaSky-AI#1850), 3. per-step cache salt (NovaSky-AI#1836). `loss/avg_final_rewards` in particular dropped from ~0.36 to ~0.24 once the fully async scripts moved to rollout_is, while eval accuracy stayed flat. <!-- CURSOR_SUMMARY --> --- > [!NOTE] > **Low Risk** > Test-only threshold and comment updates; no production training or inference logic changed. > > **Overview** > Updates **wandb metric gates** in `gsm8k_fully_async.sh` so the fully async GSM8K GPU E2E job passes after recent async training changes, without altering the training script invocation. > > **Threshold shifts** (5% allowance from CI runs since Jul 2026): eval min **0.56 → 0.50**, train reward min **0.32 → 0.198**, max tokens **283 → 285**, rollout/train logprob diff max **0.040 → 0.0193**. Comments now document the wider baseline window and tie the rebaseline to stale KV on weight sync (NovaSky-AI#1798), `rollout_is` policy loss (NovaSky-AI#1850), and per-step cache salt (NovaSky-AI#1836). > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit d589aac. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup> <!-- /CURSOR_SUMMARY --> --------- Signed-off-by: SumanthRH <sumanthrh@anyscale.com>
…g validation (NovaSky-AI#2029) ## What Injects the client-requested LoRA rank/alpha into the override dict passed to `from_cli_overrides` instead of assigning them onto the config afterwards (client wins over backend_config, matching the previous post-assignment semantics). ## Why `create_lora_training_client` requests were rejected whenever the backend config enables `trainer.policy.model.fake_int4_qat`: `TrainerConfig.__post_init__` validated `fake_int4_qat` against the default `lora.rank=0` and raised "fake_int4_qat requires LoRA" even for LoRA clients, because the client's rank/alpha arrived after validation. Part of the Kimi K2.x series (follow-up to NovaSky-AI#1862). Independent of the other PRs in the series. Made with [Cursor](https://cursor.com) Co-authored-by: Cursor <cursoragent@cursor.com>
Owner
Author
|
Superseded by NovaSky-AI#2069, which targets current upstream main. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Following Meng et al., we implement PiSSA, a reinitialization for low-rank adapters using singular value decomposition of the weight matrices.
Follow-ups
Testing
Numerical Error in PiSSA Initialization
We test the logprob numerical error in PiSSA initialization and find it to be on the same scale as LoRA initialization.
Training Qwen-3-0.6B with LoRA/PiSSA Results
We also test the algorithm via the SkyRL example provided in this PR, against the correspondign SkyRL example for Megatron LoRA on Qwen-3-0.6B.
Focused lint, compilation, shell syntax, and grouped-sharding tests pass locally. Dense and MoE B300 validation is in progress.