[Fix] Classify SessionServer endpoints and fail closed on trace errors by matrix72c · Pull Request #2061 · InternLM/xtuner

matrix72c · 2026-09-01T09:07:20Z

Summary

SessionServer is a catch-all HTTP proxy, but only two requests represent model generations: POST /v1/messages and POST /v1/chat/completions.

This change makes endpoint policy explicit before request/response processing and makes trace-enabled generations fail closed when the response cannot be recorded. Auxiliary API calls remain transparent, while the client-visible result stays consistent with the training trace stored by SessionServer.

The response-reliability changes previously proposed in #2059 are consolidated here so endpoint and trace policy are reviewed in one PR.

Problem

The proxy cannot safely infer the endpoint from the JSON body. For example, POST /v1/messages/count_tokens contains messages but is not a model generation. Treating it as one can invoke on_request, inject generation-only fields, filter tools, or invoke on_response unexpectedly.

For trace-enabled generations, malformed JSON, non-object response bodies, incomplete SSE streams, downstream disconnects, or a failing on_response hook can leave the trace store without a complete assistant turn. Returning a normal success response in that situation allows the caller to observe a turn that cannot be used for training.

Implementation

Add a small private _is_generation_endpoint(method, path) predicate. It uses the normalized method and path, ignores query strings, and recognizes only POST /v1/messages and POST /v1/chat/completions as generation endpoints.
Keep format detection independent from generation classification, so /v1/messages/count_tokens, /v1/messages/batches, and future /v1/messages/... paths still receive the correct Anthropic headers.
Keep endpoint policy explicit and local to _handle_request; no ContextVar, wrapper handler, shared enum, or cross-project abstraction is introduced.
For non-generation JSON object requests, remove only the top-level SessionServer-owned session_id. Non-object bodies and objects without session_id remain byte-for-byte unchanged whenever possible.
Derive response tracing from both the endpoint policy and the upstream status code. Non-2xx responses never enter the cleaner or trace hooks, and explicit upstream error envelopes are passed through without being wrapped again.
For trace-enabled streams, retain the original upstream chunks and delay [DONE] until parsing and on_response complete. If parsing or the response hook fails, return a native 500 error or SSE error event without [DONE].
Continue draining the upstream stream after a downstream connection reset so a complete trace can still be collected.
Handle stream preparation and connection-reset behavior directly in the request handler instead of adding a one-use helper.

Successful generation behavior, the public SessionServer API, and the existing /v1/responses 501 rejection remain unchanged. No LMDeploy code or management endpoint implementation is changed, and fix/chunk-loss-detached-head-memory is untouched.

Tests

Added focused aiohttp fake-upstream coverage for:

generation and non-generation method/path classification;
transparent count_tokens and auxiliary endpoint forwarding with no generation hooks;
successful training responses and delayed stream completion;
evaluation requests that skip on_response;
upstream JSON/SSE errors and all non-2xx responses remaining unchanged;
malformed, non-object, incomplete, and untraceable 2xx responses failing closed;
response-hook failures producing native errors without [DONE]; and
downstream disconnects before headers and during streaming while the upstream is still drained and traced.

Validation performed:

Focused endpoint/trace suite: 25 passed
ruff check: All checks passed
python -m py_compile: passed
git diff --check: passed

A PPO critic shares the actor's backbone but needs a scalar head emitting one value per token instead of vocabulary logits. Add a generic derivation so no architecture needs a hand-written value variant. - Extract a `build_head` hook in `Dense`/`MoE` so the output head is overridable, defaulting to the existing vocabulary head. - Add `xtuner/v1/model/value.py`: - `as_value_config()` derives a critic config from any transformer or compose config. For compose (VLM) configs only `text_config` is converted; the vision tower and projector are shared verbatim. - `ValueModelMixin` maps the head to its own `value_head.weight` checkpoint key, so every backbone key still matches the actor's and a critic can initialize straight from an actor checkpoint. When that key is absent the head is initialized with Normal(0, 1/(hidden_size+1)) following Open-Reasoner-Zero; a default 0.02-std head would emit large arbitrary values that GAE propagates into every advantage. - Derived configs force `tie_word_embeddings=False` (a `[1, hidden_size]` head cannot alias the embedding matrix), `mtp_config=None`, `z_loss_cfg=None`, and `mesh_prefix="critic"`. `balancing_loss_cfg` is preserved since MoE routers still need load balancing in a critic. - Add `TrainEngine.load_dcp(strict=)` so a critic checkpoint can require an exact match instead of silently missing its value head. Default `None` keeps the previous frozen-params-derived behavior.

Building a critic from an actor checkpoint crashed in `from_hf` with "a leaf Variable that requires grad is being used in an in-place operation". `init_params` writes the value head in place. For a sharded parameter it does so via `param.copy_(...)`, which autograd forbids on a leaf that requires grad -- precisely what a freshly built, trainable value head is. Every other initializer in the codebase runs under `no_grad` for this reason. The unit tests missed it because a plain CPU tensor takes the other branch in `init_params`, so the regression test drives the DTensor path directly. Found by an end-to-end run; the trainer died during worker construction.

Add `ValueLossConfig`/`Kwargs`/`Context` plus `value_loss`, implementing both plain MSE and the clipped PPO value objective. Clipping confines the prediction to a trust region around the pre-update value and takes the larger of the clipped and unclipped errors, so a critic cannot move too far while rollouts are reused across passes. The context follows the existing loss protocol, so it plugs into the unmodified model forward path (`model -> LMHead -> loss_ctx.forward`): `build_batches` reduces one denominator over every micro-batch, data-parallel rank and gradient-accumulation step, making the loss invariant to how a batch is sharded. Verified against a single-rank reference under 2-rank gloo, with uneven per-rank token counts, and under a sequence-parallel split. `mode="chunk"` is rejected: chunking exists to avoid materializing `[tokens, vocab_size]` logits, which a `[tokens, 1]` value head never does, so accepting it would only add an untested path. Also add `explained_variance`, computed from cross-rank reducible sums. It is the primary signal for whether a critic is learning at all -- 0 means it predicts no better than the mean return -- and returns `None` when the return variance is too small for the ratio to be meaningful.

`ValueLossConfig.build` materializes its kwargs on the accelerator, so on a GPU host the loss context lives on cuda while the tests' reference tensors were built on the CPU. Every comparison and every mixed-tensor call then failed -- five tests that passed on a CPU-only box. Create reference tensors on the same device as the context, and reduce on the CPU in the distributed tests since their process group is gloo.

`kl_penalty` computed the per-token KL and immediately reduced it to a weighted scalar, so the unreduced estimate was unreachable. PPO with KL-in-reward needs exactly that unreduced tensor: the penalty enters the token reward, flows through the GAE recursion, and shapes the value targets, which is not equivalent to adding KL to the loss. Split out `kl_divergence_per_token` and make `kl_penalty` its weighted sum. Behavior is unchanged -- verified bitwise identical to the previous implementation across every estimator name in float32 and bfloat16.

The RL worker built an optimizer but never a scheduler, so `lr_cfg` was accepted and silently ignored: warmup and decay had no effect on RL runs. PPO needs this, since actor and critic each require their own schedule. - Move the scheduler construction into `LRConfig.build(optimizer, total_steps)` so it is parameterized by optimizer rather than bound to one engine, and refactor the SFT trainer to delegate to it. Verified to reproduce the previous SFT schedule to floating-point precision for every lr type. - Build it as a single `LambdaLR` instead of a `SequentialLR` over warmup and decay. `SequentialLR.load_state_dict` restores counters but never reapplies the learning rate, and its next steps then diverge from the original schedule; this silently affected SFT resume too. The schedule now also holds at `lr_min` past `total_steps` rather than running off the end of the curve. - `LRScheduler.load_state_dict` never writes the restored rate back to the optimizer, so the first step after any resume used a stale rate -- 0.0 when warmup is configured. `_ResumableLambdaLR` reapplies it on load. - Advance the schedule only when the optimizer update was actually applied. `step_optimizer` skips non-finite or over-threshold gradients, and decaying the rate on updates that never happened would drift the schedule; add `TrainEngine.optimizer_step_will_apply` to expose that decision. - Persist and restore scheduler state in the worker checkpoint, honoring the existing `load_scheduler` flag, and warn rather than fail on checkpoints written before this existed. Log the current `lr` per training step.

`AdvantageEstimator.compute(rewards, group)` maps a group of scalar rewards to one scalar advantage per completion, which cannot express GAE: GAE needs per-token value predictions and recurses along the sequence. - Add `TokenLevelAdvantageEstimator`, a parallel ABC for estimators that require a critic. The trainer dispatches on it to decide whether to build a value model and where advantages are computed. Existing estimators are untouched. - Add `GAEEstimator` plus `GAEAdvantageConfig`, the single source of truth for gamma/lambda, and `terminal_token_rewards`, which places each trajectory's score on its final action token. - The recursion walks action tokens only, so interleaved observation tokens in agentic rollouts neither receive an advantage nor discount the chain. Packed boundaries reset it; each trajectory's last action bootstraps from zero. - The scan is vectorized across trajectories: actions are gathered into a `[num_trajectories, max_actions]` matrix and reduced column by column, so the Python loop length is the longest action count rather than the token count. This is 3-30x faster than a per-token loop on 32k packs while agreeing to float32 precision, verified by differential tests against a textbook reference implementation. A closed-form `(gamma*lambda)**t` weighting was rejected: it underflows to zero past ~8k tokens. - Add `normalize_advantages`, reducing moments over the whole process group so the effective step size does not depend on data sharding. Sequence parallelism needs no special case since SP ranks hold replicas, scaling both the statistics and the count. - Thread `token_rewards` from `ColateItem` through packing to `WorkerInputItem`. Per-token tensors are declared once in `TOKEN_TENSOR_KEYS` so adding another signal does not touch the packing and DP-padding paths separately. Also correct the stale `advantage: float` annotation, which has been a per-token list.

`fit` handled input preparation, reference/old logprobs, importance-sampling metrics, loss calibration and the optimizer loop in one 311-line method. PPO needs the last two parts verbatim but reaches them differently: advantages come from a critic forward and GAE rather than a group reward baseline. Extract the loss-calibration and optimizer-step loop into `_train_actor`, which both paths can call. This is a pure move -- the only change to the relocated code is collecting log items into a returned list instead of mutating `worker_log_item` -- so PPO can reuse it without duplicating entropy, importance-sampling, metric and profiling handling. `fit` drops to 200 lines with no behavior change.

Wire the value model, GAE and value loss into the RL worker so PPO runs end to end. The critic lives in the actor's worker process and shares its placement group, avoiding a second Ray actor and a second copy of the rollout plumbing. Both are full models, so a colocated run cannot hold them on the accelerator at once: `PPOPhase` sequences the swap explicitly, and `offload_model`, `onload_model` and `save_hf` refuse transitions that would fault in the wrong model mid-step. One step runs the critic forward under `no_grad`, computes GAE from the frozen values, trains the critic on the resulting returns, then swaps in the actor and reuses `_train_actor` unchanged. Sequence parallelism is supported: GAE is a sequential recursion over a whole trajectory and cannot be evaluated on a shard, so value shards are all-gathered and every rank computes the same targets redundantly -- far cheaper than the communication a distributed scan would need. This follows the existing gather/compute/split pattern in `compute_rollout_is`. `pack_max_length` must be divisible by `sp_size` so the gather returns exactly `pack_max_length` and no shard padding has to be trimmed. Trainer-side changes are deliberately small: the existing data preparation gains a terminal `token_rewards` tensor instead of a second 300-line PPO copy, and the group-baseline estimator is skipped since PPO has a learned baseline. Because of that baseline, uniform-reward groups still carry signal and are no longer wasted work. A validator keeps the critic and a token-level estimator enabled together and pushes gamma/lambda down to the worker, so they are configured in exactly one place. Critic health is reported as explained variance alongside value/return means and clip fraction, reduced across ranks. Explained variance is the metric that says whether the critic is learning at all: 0 means it predicts no better than the mean return, which is the usual first sign that PPO is misconfigured. Checkpoints round-trip critic weights, optimizer and scheduler; `save_hf` exports the critic under `critic/`. Loss-side KL and SFT interleaving raise `NotImplementedError` rather than silently misbehaving, since both would need a third resident model.

The KL reward phase faulted the actor in for a behavior-logprob forward via `_onload_actor`, which only accepted `ALL_OFFLOADED`. A warmup step skips the actor update and leaves the phase at `ACTOR_READY`, so the very next rollout crashed with "Cannot onload the actor from phase actor_ready" -- the first warmup step worked and the second one died. `ACTOR_READY` already means the actor is resident, so the transition needs a phase stamp but no device transfer. Accepting it makes `_onload_actor` idempotent for that state and keeps the critic from ever sharing the accelerator with the actor. Add a state-machine test that pins the transition and asserts no redundant host-to-device transfer happens. (cherry picked from commit 8e3b3dee441f63d46b9ee076bd711dcfc5e1623d)

The critic's `reduced_critic_*` sums (valid count, value/return sums and squares, squared error, clip count) were silently dropped: the model's `ModelForwardExtraLogInfo.get()` only emits keys from its hardcoded reduction whitelist, and the critic sums were not on it. `_finalize_critic_metrics` therefore received nothing, and the tracker never recorded `critic/explained_variance`, the primary health signal for PPO. Add the keys to the shared whitelist, following the same pattern the policy metrics use. All are sums, which matches the reduce op applied per micro-batch.

Two additions that make PPO stable in practice. KL-in-reward: classic RLHF PPO subtracts `beta * KL(policy || reference)` from the token reward before advantage estimation, so the penalty is discounted and bootstrapped like any other reward and the critic learns to predict it. The loss-side KL used by the group-baseline algorithms cannot express this -- it reaches the policy gradient but leaves the value function unaware -- so the two are mutually exclusive rather than additive. Verified: with beta=0.2 and a per-token KL of 0.5 over four tokens, the value target at t=0 shifts by exactly -0.4, which a loss-side penalty would not do. The penalty runs before the critic phase, transiently faulting in the actor for behavior log probabilities and the reference model for its own, then evicting both before the critic needs the accelerator. `behavior_logprobs="rollout"` reuses the inference engine's values and skips the extra actor forward, at the cost of measuring divergence from the sampling policy including any train/inference numerical mismatch. Critic warmup: a freshly initialized value head predicts noise, so the first advantages are noise too. `warmup_steps` runs those steps with the critic only -- no actor forward, no optimizer step, and no scheduler step, so the actor's schedule stays aligned with the updates it actually receives. The reference model is now built when either KL path is enabled, and mean KL is logged so the penalty can be tuned against a target divergence.

Adds the entry point a reader needs to see how the PPO pieces fit together: selecting `GAEAdvantageConfig` is what turns an RL run into PPO, `critic_cfg` supplies the value function it requires, and `as_value_config` derives that critic from the actor so it initializes straight from the actor checkpoint. The comments record the decisions that are not obvious from the API: why `cpu_memory_per_worker` is 48 GiB rather than the 16 GiB the GRPO examples use (two models plus their Adam state live in host memory between phases), why the critic runs a higher learning rate and more passes than the policy, why `prompt_repeat_k` can be small once a learned baseline replaces the group baseline, and why `gamma=1.0` is standard for RLHF. Add a config smoke test asserting the invariants the file is meant to demonstrate: critic and GAE enabled together, gamma/lambda reaching the worker, distinct device meshes, MTP and weight tying disabled on the critic, and KL applied to the reward rather than the loss.

Expose the run-size knobs (steps, batch, repeat-k, response length, rollout TP, evaluation) as environment variables with the previous values as defaults. An end-to-end smoke run can then exercise the exact config that ships as the reference, instead of a divergent copy that could drift from it.

Add `EXP_TRACKER` so a smoke run can select the jsonl tracker and have its metrics asserted on programmatically rather than eyeballed in TensorBoard. Defaults to tensorboard, unchanged for normal runs.

The trainer asserts `num_cpus_per_worker * num_workers + 10` against what Ray sees in the pod, so the rjob `--cpu` request must match the config or startup aborts. Expose it as `CPUS_PER_WORKER`, defaulting to the previous value.

…ependencies

…collector

…nstrumentation

…experts-uint16

…ched-head-memory

matrix72c and others added 30 commits August 18, 2026 10:15

feat(rl): support sandbox dependency groups

fbf6ade

style(rl): format sandbox pool test

a4533e2

Merge branch 'main' into feat/sandbox-pool-dependencies

8715365

feat(trace): support external OTLP collectors

4a032dd

fix(trace): normalize external collector paths

0d1920d

fix(trace): narrow external endpoint types

7681fde

feat(trace): instrument sandbox agents and SessionServer

482a44b

fix(trace): defer external JSONL loading

d296013

test(trace): avoid integration import conflict

9da4e53

fix(trace): normalize external JSONL path for mypy

1578a77

fix(model): initialize missing value head on every rank

30a382f

docs(rl): let the PPO example emit machine-readable metrics

32683a1

Add `EXP_TRACKER` so a smoke run can select the jsonl tracker and have its metrics asserted on programmatically rather than eyeballed in TensorBoard. Defaults to tensorboard, unchanged for normal runs.

fix(rl): keep critic metric collectives rank-consistent

fe16af6

fix(rl): store routed experts as uint16

7909d9c

Merge branch 'main' into feat/sandbox-otel-instrumentation

9eb590f

matrix72c force-pushed the fix/session-server-endpoint-classification branch from d80dd84 to 6165c7b Compare September 1, 2026 13:01

matrix72c added 5 commits September 2, 2026 09:07

Merge remote-tracking branch 'upstream/main' into feat/sandbox-pool-d…

fcb3b8d

…ependencies

Merge remote-tracking branch 'upstream/main' into feat/external-otel-…

4c80db2

…collector

Merge remote-tracking branch 'upstream/main' into feat/sandbox-otel-i…

1ebdb7c

…nstrumentation

Merge remote-tracking branch 'upstream/main' into fix/compact-routed-…

d2a84df

…experts-uint16

Merge remote-tracking branch 'upstream/main' into fix/chunk-loss-deta…

c67df69

…ched-head-memory

matrix72c force-pushed the fix/session-server-endpoint-classification branch from dd8954e to cd15002 Compare September 2, 2026 03:36

matrix72c changed the title ~~Classify SessionServer endpoints before proxying~~ [Fix] Classify SessionServer endpoints and fail closed on trace errors Sep 2, 2026

matrix72c changed the title ~~[Fix] Classify SessionServer endpoints and fail closed on trace errors~~ Fix SessionServer endpoint classification and trace failures Sep 2, 2026

matrix72c changed the title ~~Fix SessionServer endpoint classification and trace failures~~ [Fix] Classify SessionServer endpoints and fail closed on trace errors Sep 2, 2026

matrix72c added 19 commits September 4, 2026 17:06

Merge branch 'feat/sandbox-pool-dependencies' into integration

25307b5

Merge branch 'feat/external-otel-collector' into integration

baac803

Merge branch 'feat/sandbox-otel-instrumentation' into integration

ea7bbfe

Merge integration into feat/ppo

b2d170e

Merge branch 'feat/sandbox-pool-dependencies' into integration

9d1a459

Merge branch 'feat/external-otel-collector' into integration

701140c

Merge branch 'feat/sandbox-otel-instrumentation' into integration

3b909f1

Merge branch 'feat/ppo' into integration

0b6922d

Merge integration into compact routed experts topic

e817562

Merge branch 'feat/sandbox-pool-dependencies' into integration

46515a9

Merge branch 'feat/external-otel-collector' into integration

fa9ffea

Merge branch 'feat/sandbox-otel-instrumentation' into integration

d1d89e8

Merge branch 'feat/ppo' into integration

cb2af63

Merge branch 'fix/compact-routed-experts-uint16' into integration

5416270

Merge branch 'fix/chunk-loss-detached-head-memory' into integration

c055f11

fix: restrict SessionServer tracing to generation endpoints

d431dc1

refactor: classify session server endpoints locally

09b2c1e

fix: fail closed on untraceable generation responses

4ba8ad0

style: apply ruff formatting to SessionServer changes

76eaaba

matrix72c force-pushed the fix/session-server-endpoint-classification branch from 5a8433a to 76eaaba Compare September 4, 2026 09:12

Navigation Menu

Uh oh!

[Fix] Classify SessionServer endpoints and fail closed on trace errors - #2061

[Fix] Classify SessionServer endpoints and fail closed on trace errors#2061
matrix72c wants to merge 58 commits into
InternLM:mainfrom
matrix72c:fix/session-server-endpoint-classification

matrix72c commented Sep 1, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

matrix72c commented Sep 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Problem

Implementation

Tests

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

matrix72c commented Sep 1, 2026 •

edited

Loading