feat(models): add AR value head and GAE track wiring - #1
Closed
yhl48 wants to merge 4 commits into
Closed
Conversation
Wire per-token critic values through Qwen3 replay and compute GAE advantages on RolloutTrack for the PPO critic path (issue Tencent-Hunyuan#86, part 2/3).
Keep transformers imports inside from_config / _packed_replay_supported as in upstream; only UniRL imports (ValueHead, ARReplayOutput) stay top-level.
Call _require_value_head_for_replay at the start of _replay_aware_forward so return_values=True fails fast before the transformer forward, not only from Qwen3ARStage.replay().
Extend ReplayResult with optional per-token values for PPO/GAE and replace ARReplayOutput in Qwen3 replay with the shared type.
Owner
Author
|
Closing in favor of upstream Tencent-Hunyuan#256 (rebased onto main after Tencent-Hunyuan#254 merged). |
2 tasks
leviking98z-rgb
added a commit
that referenced
this pull request
Aug 1, 2026
…cent-Hunyuan#214) * types: add Sample/Part endomorphism types with sample_id lineage (LIN-446) Squash merge of origin/LIN-446/main into LIN-444/main. Combines: - add Sample/Part endomorphism types - derive Sample/Part lineage from the sample_ids path * rollout: convert the rollout engine to Sample → Sample (LIN-454) Replace the RolloutReq/RolloutResp request/response pair with the unified recursive Sample/Part types across the dedicated rollout engines — generate: Sample → Sample, in place, no bridge. The model pipelines, the trainside engine, the trainers, reward, and the train stack stay on the old types for a later pass, so this branch is intentionally not runnable as a full training loop yet. Engines converted: - sglang (AR text/vlm), sglang_diffusion (image), vllm_omni (sd3 / qwen_image / hv15 single-DiT + HI3 multi-stage), composed (PE AR→diffusion chain). - base.py ABC flipped to generate(self, sample: Sample) -> Sample; engine/__init__ chunked_engine_generate_req -> chunked_engine_generate (split-by-root -> regroup -> concat, the Sample analogue of slice-by-index). Mechanics: - Pre-forked gen shells are filled by sampling_params type; positional lineage replaces parent_track; σ pinned onto DiffusionSamplingParams.sigmas; the x_T noise key is derived from the path-id lineage (OD-2); stage_config moved to Part.control. - Multi-input multimodal (chained input Parts) is gated with clear deferral raises — a deferred non-goal of this pass. P0 type fields: Part.control / conditions / fill, DiffusionSamplingParams.sigmas / init_noise_latent_shape, LatentSegment.initial_latents. Verification (rollout-only, no training): - scripts/check_sample_roundtrip.py structural oracle (5/5 contracts). - Real GPU rollout smokes pass for every converted engine and modality class: AR (sglang/Qwen3), diffusion (vllm_omni sd3, sglang_diffusion sd3), and the multi-stage composed PE chain (AR→diffusion). * rollout: support multi-input multimodal (IT2I + siblings) (LIN-454) Wire the previously-gated image+text / cot_text modalities (hi3_it2i, hi3_i2t, the sglang vlm path, and hi3_dit_recaption) onto the generate(Sample) -> Sample boundary. The type layer already supported it, so this is mostly un-gating: a second input rides as a chained input Part via the new Part.input_child(primitive) (branch-1 child, no sampling_params) so only the head stays a root, and Sample.conditioning() surfaces both primitives in turn order. Adapters locate the chained input by primitive type (image_input_part / new cot_text_from_sample); hi3_dit_recaption gets the ported cot_text build. Also fixes sglang text.py build_response to preserve intermediate input Parts (return [*parts[:-1], filled] instead of [parts[0], filled]) so a multi-input chain's gen Part keeps a valid parent. Verified engine-only with hand-built Samples: oracle check_sample_roundtrip 7/7 (incl. multi-input image-chain + cot_text fixtures), a real Qwen2.5-VL VLM run, and a full HunyuanImage-3 80B IT2I GPU run (image+text -> AR recaption -> edited image, chain filled). * LIN-480: migrate trainers/train-stack/logging onto the Sample/Part model Migrate the consumer side (all five trainers + train stack + wandb logging) from RolloutReq/RolloutResp/RolloutTrack to the Sample/Part endomorphism model (LIN-446) — the counterpart to LIN-454's rollout-engine conversion. Trainers now build request Samples (Part.input + .fork) and consume response Samples (frontier = parts[-1], stages located by sampling_params type); GRPO advantage / reward-propagation / token-balancing run through the methods already on Part/Sample. Scope: all trainers incl. it2i (multi-input via Part.input_child) and the HI3 unified two-engine path. Out of scope / unchanged: the reward service (score_and_attach migrates on a separate branch — its (*, req, track) signature does not yet match the trainers' score_and_attach(sample) call, so the train tail is gated on that branch), the trainside engine, and the rollout-engine internals. Highlights: - Sample/Part navigation on the type: root_group_ids, gen_parts/gen_part/ gen_part_index, with_parts, and tree-aware slice/select (a request Sample DP-scatters by prompt-tree, not by the parts list). - Per-trainer request builders + the HI3 two-engine _run_rollout_one stitching (AR recaption -> DiT image); the DiT x_T noise key is re-rooted from the globally-unique lineage so it stays unique across dp>1 replicas. - UnifiedModelTrainStack.train_track scatters the [input, ar, image] lineage as a unit so both stages shard at the same prompt boundaries. - wandb on Sample/parts: diffusion stage named "diffusion" consistently; num_samples reports the generated-sample count; zero-std groups bucket by the advantage grouping (advantage_group_ids). - disable_driver_xt flag on DiffusionSamplingParams (restores the unified DiT driver-x_T escape hatch). Verification: scripts/check_sample_roundtrip.py (14 contracts) + SD3 (vllm-omni), AR (sglang), and HI3 two-engine generation smokes pass on GPU. The dp>1 cluster is contract-tested + A2 engine-verified; its multinode end-to-end behavior is not yet verified, and the full train loop awaits the separate reward migration. * LIN-479: model pipelines + trainside engine onto the Sample/Part model Convert the model bundle's rollout boundary from RolloutReq/RolloutResp to the recursive Sample/Part types — the model-bundle half of the refactor, pairing with LIN-480 (trainers / train-stack): - SD3Pipeline.generate / Qwen3Pipeline.generate take a request Sample (pre-forked gen frontier) and fill it; a factored _conditions_for() is the single encode path shared by rollout and trainer-side replay (re-encode; Part.conditions left empty — no cache). - TrainsideRolloutEngine.generate -> Sample -> Sample: _ensure_sample_sigmas pins sigma onto the gen part's DiffusionSamplingParams; Part.slice/concat chunking. - NoiseRecipe.from_sample() — Sample-shaped x_T builder keyed on the gen part's path-id lineage. - check_sample_roundtrip.py +2 contracts (from_sample lineage, generate-fills- frontier); trainside_{sd3,ar}_smoke.py real-GPU rollout+replay smokes. Validated on H20 (pod unirl-gz-2): CPU oracle green; SD3 rollout->replay mean|dlogp|=0 (ratio~1); Qwen3 AR replay self-consistency 0 (the old_logp_source=replay contract — the in-process autoregress bf16 record vs replay fp32 differs ~3-5, harmless since trainside recipes use replay for old_logp). Stages (diffusion.py/ar.py) untouched. * reward: migrate score_and_attach to Sample → Sample (LIN-481) Convert the reward adapter from RolloutReq/RolloutTrack to the Sample/Part endomorphism model — the consumer-side counterpart LIN-480 left gated. The trainers already call score_and_attach(sample); this matches that contract. - RewardService.score_and_attach(*, req, track) -> RolloutTrack becomes score_and_attach(self, sample: Sample) -> Sample: score the frontier Part, build input primitives from Sample.conditioning() (nearest-ancestor caption), generated from frontier.primitive keyed by the backend preferred_input_kind, and root-sourced prompt metadata via the new Sample.root_metadata. - AR truncation/overlong shaping preserved verbatim, re-gated on the scored frontier being ARSamplingParams. - Delete the request/track expansion + metadata-normalization machinery (_KIND_TO_KEY, _normalize_prompt_metadata, _build_request_for_track): the Sample is row-aligned and DP-scatters by prompt-tree, nothing to reconcile. - Add primitive_modality_key + Sample.root_metadata; refresh reward docstrings/README to Sample/Part. - Add scripts/check_reward_roundtrip.py structural oracle (6 contracts; passes on-GPU-pod with real torch). * fix(sd3): carry encoded conditions on the trainside gen Part for train-stack replay SD3Pipeline.generate left Part.conditions empty (a "replay re-encodes" assumption), but the train stack's prepare_segment reads part.conditions directly with no re-encode wiring — so trainside SD3 training crashed in SD3Conditions.from_dict (text=None). Every sibling diffusion pipeline and the SD3 sglang_diffusion adapter already populate conditions via fill(conditions=...); bring trainside SD3 in line. Surfaced by the SD3 trainside e2e (trainside_sd3_smoke bypassed this by calling stage.replay directly, never going through train_track). * fix(vllm-omni): tile prompts to the gen fan-out for served diffusion rollout texts_from_sample asserted prompt count == gen sample count, but the request keeps the input Part un-fanned (one prompt per group) while the gen shell fans out samples_per_prompt via Sample.fork — so SD3/qwen-image served rollout crashed on rollout 1 with "prompt count 6 != gen sample count 96". vllm_omni runs num_outputs_per_prompt=1 (one output per gen sample), so tile each prompt across its group-by-parent-contiguous gen siblings. A 1:1 request (no fan-out) is unchanged. Surfaced by the SD3 sd3_vllmomni e2e (served LIN-454 path). * fix(sglang): allow unirl weight-sync classes through sglang 0.5.12 SafeUnpickler sglang 0.5.12's SafeUnpickler (CVE-2025-10164 guard) runs in the engine process and blocked the colocate tensor weight-sync: update_weights_from_tensor deserialization rejected unirl's vendored FlattenedTensorBucket / _rebuild_cuda_tensor_modified (both under the unirl. namespace) with "Blocked unsafe class loading". Register "unirl." on sglang's server-side allowlist in SGLangRolloutEngine.__init__ (the tp_size=1/use_ray=False scheduler runs in this process), mirroring what unirl's own SafeUnpickler already permits. Surfaced by the Qwen3 DRPO AR e2e (qwen3_drpo_4b_base_dapo_sglang): the full vertical (sglang gen -> MathVerify -> DRPO -> FSDP train) ran and crashed only at the first weight-sync. * fix(sglang): use native serializer for LLM weight-sync (sglang 0.5.12 compat) The prior allowlist patch (a1f33d1) registered "unirl." on the engine process's SafeUnpickler, but sglang's SRT scheduler deserializes the weight bucket in a SEPARATE spawned subprocess with a fresh allowlist, so the patch never took (same "Blocked unsafe class loading" recurred). Instead, serialize with sglang's NATIVE FlattenedTensorBucket / MultiprocessingSerializer / monkey_patch_torch_reductions (verified identical device-UUID IPC mapping) so the wire references sglang's own allowlisted classes — matching what the sglang_diffusion engine already does. Reverts the ineffective __init__ allowlist patch. Surfaced by the Qwen3 DRPO AR e2e: full vertical ran, crashed only at the initial colocate weight-sync into SRT. * LIN-495: migrate all model bundles onto the Sample → Sample rollout shape Convert every model pipeline's generate() from RolloutReq → RolloutResp to the recursive Sample/Part endomorphism, completing the model-bundle half of the rollout refactor (pairs with the earlier LIN-454/479/480/481 engine/trainer/ train-stack/reward migrations), then retire the legacy types. - Base contract: Pipeline.generate(self, sample: Sample) -> Sample; lift the request-side Sample readers out of the rollout layer into types/sample_ops.py so pipelines import them without a models→rollout inversion. - 7 Tier-1 diffusion pipelines (qwen_image, z_image, wan21, wan22, hunyuan_video, hunyuan_video15, ltx2): read prompt via sample.conditioning(), gen params off the frontier Part, fill the frontier shell; factor _conditions_for() shared by rollout + trainer-side replay. i2v image arrives via the input_child chain. - qwen_vl + flux2_klein (AR/edit with one extra input modality). - pe, bagel, hunyuan_image3 (multi-stage 2-gen-part flows): PEPipeline consumes the pre-forked [input, ar_shell, diff_shell] Sample and drives the child pipelines (mirrors the migrated ComposedRolloutEngine); bagel dispatcher reads task from parts[0].control and fills one or both gen Parts; HI3 dispatcher + 5 modes/ (incl. two-part t2ti). - Retire unirl/types/rollout_req.py + rollout_resp.py; relocate PrimitiveValue to types/primitives.py; drop the dead NoiseRecipe.from_rollout_req; sweep doc breadcrumbs across conditions/algorithms/sde/engine docs. All pipelines carry conditions=<conds>.to_dict() on frontier.fill — the train stack reads Part.conditions (GRPO/FlowGRPO re-type via conditions_cls.from_dict); it does NOT re-encode. (Caught by E2E pe-trainside on unirl-gz-2: an earlier revision left conditions empty and crashed at the first AR train step with "Qwen3ARConditions.from_dict: expected d['prompt'] ... got None".) Validated on TaiJi unirl-gz-2 (8×H20, GZ): compileall green; sd3-trainside reward grew 0.74→0.79 with ratio=1.0 / advantage_std≈0.49; pe-trainside (SD3+Qwen3-0.6B) runs the full 3-part chain with both AR + diffusion stacks training. Adds scripts/trainside_{qwen_image,qwen_vl,pe}_smoke.py (tier-representative GPU rollout→replay smokes). * LIN-503: agent-trajectory multi-conditioning — role-aware layer + multi-turn encode (both backends) Squash of LIN-503/main (8 commits) onto LIN-444/main. Takes the rollout Sample from single-turn-only to full multi-turn, multi-modal conditioning end-to-end. - Role-aware trajectory layer on Sample/Part (Phase 1/2): Turn, Part.role / resolved_role, turns(), text_conditioning() / vision_conditioning() fail-loud renderers, replace_frontier / with_filled_frontier write-back; diffusion/omni engines wired; sample_ops retired. - gap C — sglang AR engine encode consumes the trajectory: text.py / vlm.py build the chat conversation from text_/vision_conditioning() via a pure transpose + de-expand util (rollout/engine/sglang/utils/conversations.py); resolve_sampling fan-out fixed to the last-fork branch (parts[-2]). - Trainside conjugate — qwen3 / qwen_vl in-process encoders consume turns() via models/types/conversations.py (transpose, no de-expand, inline-PIL VLM fusion); stale "_conditions_for re-tokenize" docstrings + trainside smokes fixed. - Behavior-preserving on single-turn (byte-identical); parity-safe by construction (replay teacher-forces over the stored conditions["prompt"], not a re-encode). - Verification: CPU oracles (engine 8/8, trainside 6/6, layer 7/7) + 4 GPU multi-turn smokes (sglang/trainside × text/VLM) green on Qwen3-4B-Instruct and Qwen2.5-VL-3B-Instruct — the captured prompt carries user→assistant→tool (+ image), replay ratio=1. * feat(rollout): agentic tool-loop rollout — AgentLoop + ToolEnvironment (LIN-492) Add the agentic (multi-turn tool) rollout layer over the Sample/Part model, on top of LIN-503 multi-turn conditioning. Scope is ROLLOUT only — trainer wiring, reward, and a dataset for agentic GRPO are a follow-up (the trainer is still single-turn). - AgentLoop (unirl/rollout/loop/): environment-driven synchronous loop — fork -> generate -> env.step -> observe, bounded by max_turns. Structural Environment / RolloutEnginePort protocols; the existing SGLang engine satisfies the port unchanged. - ToolEnvironment + tools/ (Tool, CalculatorTool): parse <tool_call>, run a registered tool, feed the result back as a role="tool" observation, stop on a final answer (or max_turns). Safe ast-based calculator (no eval). - Sample.observe(role="tool"): append the world-response as a mask-0 input Part tagged for LIN-503 role-aware rendering (<tool_response>). Additive — no existing code path changed; the whole merge is 1648 insertions, 0 deletions. Verified: - CPU oracles: scripts/tool_env_smoke.py (9 checks), scripts/rollout_loop_smoke.py (6), scripts/check_sample_roundtrip.py (16, regression) — all green. - GPU (Qwen3-4B-Instruct): single-turn tool call (tool_env_ar_smoke); the closed loop — model calls the calculator, sees the exact result, and answers (tool_env_ar_loop_smoke); and a 72-trajectory reliability run at 97.2% correct with clean termination (tool_env_ar_reliability). Known follow-ups: trainer integration (multi-gen-Part train step + reward + dataset), per-sample termination for n>1 heterogeneous batches, more tools. * LIN-499: async per-group rollout-engine interface (squash-merged onto agentic) Squash-merge of LIN-499/main: the async per-group engine contract (agenerate core + sync generate facade + abort/pause/resume), all five engines migrated, Part.weight_version provenance, DevicePool.worker_max_concurrency knob, and the CPU async-contract test suite (11 tests pass on the merged tree). Conflict in types/sample.py resolved by keeping BOTH agentic's Part.role (LIN-503) and LIN-499's Part.weight_version — adjacent new fields, not mutually exclusive. * feat(rollout): agentic multi-turn rollout engine (rank-0 pull coordinator) AgenticRolloutEngine drives multi-turn (tool-use) rollout: a rank-0 coordinator (BROADCAST+RANK_ZERO, the NCCLWeightSync pattern) over a DP slab of per-worker persistent drain loops. Each worker runs one run_until_complete that pulls single-trajectory tasks from rank 0 and runs them as multi-turn agent loops on its inner engine's event loop (continuous-batched via the inner backend's semaphore). generate returns a flat List[Sample] of variable-depth trajectories; the GRPO n-group is recovered by bucketing on the shared prompt root id. - base: widen generate to RolloutOutput = Sample | List[Sample] (single-turn returns Sample, agentic returns List[Sample]); per-turn seams stay -> Sample. - engine/agentic: coordinator (set_workers/generate/next_task) + per-worker drain (run_drain/_drain/_run_one/_pull); _run_coro delegates to the inner so the drain and weight-sync verbs share one lock (the quiesce boundary); a per-worker trajectory cap distinct from the backend request semaphore; lifecycle/weight-sync verbs delegate to the inner engine. - loop: add async astep to the Environment protocol + ToolEnvironment (non-blocking tool boundary); make step re-entrant (turn derived from the sample, not a mutable counter) so one env serves concurrent trajectories; add agenerate to RolloutEnginePort. - tests: CPU contract tests (ragged List[Sample] + bucket-by-root, cap saturation, two distinct bounds, pull load-balancing, failure isolation, FIFO next_task) + astep parity / re-entrancy / slow-tool-yields. - scripts: multi-worker GPU smoke. Validated on a 2-worker Qwen3-8B H20 slab: ragged List[Sample] with correct tool-use answers, and cross-worker rank-0 aggregation proven safe (gen segments produced on 2 distinct workers, all spans plasma-backed in Ray's global object store, 48 TensorRefs hydrate from the driver) — risk #1 closed, no run_drain materialization needed. LIN-522 * feat(agentic): multi-turn agentic RL stack — AgenticTrainer, deep-research tools, ALFWorld Land the full multi-turn agentic RL capability on top of AgenticRolloutEngine. AgenticTrainer (unirl/trainer/agentic.py) GRPO over variable-depth trajectory lists: group-relative advantage across the n siblings of a prompt, every assistant turn concatenated into ONE on-policy train_track step (ratio ~ 1). The per-trajectory reward step (_rewards_and_groups) is overridable so tasks swap only the reward SOURCE. Deep-research task (answer-graded) M1 calculator + MathVerify and M2 search/visit tools + LLM-judge reward, with recipes and the train_deep_research entrypoint. ALFWorld baseline (AgenticEnvTrainer, env-sourced reward) Stateful per-trajectory ALFWorld adapter, an engine hook that attaches the environment's terminal-success return to the last turn, admissible-action snapping to avoid TextWorld PDDL crashes, and crash-exclusion (NaN reward -> neutral, zero advantage) so env-bug crashes cannot corrupt the GRPO gradient. Verified on-pod: task-success reward rises ~0.53 -> ~0.85 (peak 0.94), on-policy, on a fixed 8-game set (genuinely multi-turn: up to 10 env steps per trajectory, n=8 GRPO group). Testing: CPU contract tests cover the plumbing (reward attach, group advantages, env adapter, train_step assembly). The GPU end-to-end path is validated empirically (the rising curve), not yet in CI. * feat(rollout): stateful, decoupled tools for agentic RL (LIN-533) Add a StatefulTool seam so tools can hold per-trajectory state across turns, with guaranteed teardown and a first out-of-process tool. - StatefulTool(Tool): session_start -> execute_session -> session_end, keyed by a session id carried in the root control bag. ToolEnvironment dispatches on isinstance, so the stateless Tool path is unchanged. - ToolEnvironment: reset mints a uuid4 session id per stateful tool and stamps it into control["tool_sessions"] via _part_with_field (returns the request unchanged when there are no stateful tools); step/_run dispatch execute_session; async aclose ends sessions in the executor, swallowing errors. - AgenticRolloutEngine._run_one: finally-hook calls env.aclose on every path (success, crash, abort), duck-typed via getattr and wrapped so it can never re-raise into the drain. - AlfworldEnv.aclose: reclaim the episode + pooled template — fixes a leak when a trajectory dies in the engine between turns. - SandboxTool: persistent per-session Python REPL subprocess (lazy spawn in the executor, select-based timeout, killed on session_end); validates cross-turn variable reuse. - Tests: StatefulTool lifecycle + teardown on success and forced exception, SandboxTool REPL, and an ALFWorld leak regression. * feat(agentic): partial rollout — colocate + async trainers (LIN-531) Over-generate, interrupt generation at a turn boundary on weight sync, commit the fast complete GRPO groups, carry/drop the slow tail. - Engine: submit/poll/abort/drained over a background buffered drain; turn-boundary checkpoint-and-resume in _run_one. - Colocate driver: AgenticPartialTrainer / AgenticEnvPartialTrainer — keeps all GPUs on generation while cutting the straggler tail. - Disaggregated driver: AsyncAgenticTrainer / AsyncAgenticEnvTrainer. - Tail policy: carry (resumable tool envs) vs drop (stateful envs, ALFWorld). - Shared _advantage_train_and_log extracted from the barrier train_step (barrier path preserved); entrypoints + ALFWorld/deep-research recipes; per-rollout turn-histogram + committed/dropped instrumentation. ALFWorld speed study: colocate-partial beats the barrier only with group-level depth variance — Qwen3-8B ~18% faster, Qwen3-0.6B 27% slower. Integrates with LIN-533 teardown: the _run_one finally aclose now also fires on the abort/checkpoint path, releasing carried trajectories' episodes/sessions. * feat(deep-research): AReaL tongyi_deepresearch M2 (search + visit + judge) Reproduces AReaL's tongyi_deepresearch deep-research agent in UniRL's agentic stack (sync AgenticTrainer + AgenticRolloutEngine): - Verbatim Tongyi SYSTEM_PROMPT for the M2 recipe (deep_research_search_judge) - Hardened SearchTool: serper + serpapi providers, retries/backoff - Hardened VisitTool toward AReaL tool_visit.py: Jina read retries, structured evidence/summary extraction, content truncation - Robust LLM-judge verdict parsing (negative-first regex; fixes "not correct") - Cross-config invariant guards (AgenticTrainer + AgenticRolloutEngine) - prepare_asearcher: default ASearcherBase35k split, streaming load - Committed GPU training smoke (scripts/train_deep_research_smoke.py) Validated live: reward climbed 0.22 -> 0.50 over 19 rollouts (Qwen3-1.7B policy, Qwen2.5-72B judge, batch 128, n=8) before an infra-driven collapse. Follow-ups (not in this change): multi-turn token-recording reconciliation (Miles-style trim/accumulate) and a trajectory context-token cap for long-horizon runs. The shipped recipe defaults to max_turns=8 (safe). * test(deep-research): adversarial LLM-judge verdict-parse cases CPU unit test for _parse_verdict guarding the negative-first parse against the "incorrect" ⊃ "correct" substring trap and "not correct" / "wrong" phrasings the prior substring test misread. 20 cases; the coverage Phase C intended but the M2 squash omitted. Verified: 20 passed on agentic@ce13a48d. * refactor(types): grouping as id projections + is_gen kind signal Collapse the grouping/lineage duplication on the Sample/Part model: - sample_id: add ancestor_id(sid, depth) — grouping labels are id-prefix projections. root_group_ids and Sample.split project directly; the _root_groups_per_part walk (re-deriving what __post_init__ validated) is deleted. - compute_advantages: the group_ids label-list override becomes group_layer (None = immediate parent, 0 = root prompt; PE diffusion_group_scope="prompt" now passes group_layer=0 instead of threading root_group_ids labels). scope stays the normalization mode, with the historical global branch (unbiased std) kept verbatim for bit-parity with the shipped adv_normalization_scope: global recipes. - delete Part.advantage_group_ids (its only reader, the zero-std wandb metrics, buckets by sibling group_ids) and the dead Part.split. - kind signal: fork requires sampling_params (a paramless gen shell is unrepresentable); Part.is_gen is the single predicate behind gen_parts / resolved_role / the wandb gen filter. Ids stay pure addresses; role stays presentation-only. - contracts: check_group_layer_advantages pins per-layer GRPO values, the unbiased-std global formula, and fail-loud on paramless fork and negative group_layer. * refactor(rollout): make sync and async paths engine-native * refactor(rollout): sync-only engine contract; concurrent-safe sglang backends Drop agenerate/run_session/CoroutineFactory from the rollout-engine contract (BaseRolloutEngine, BaseSingleTurnRolloutEngine) and make the SGLang backends safely callable from concurrent threads, so the agentic drain can drive sync generate from one thread per trajectory: - native: replace the SessionRunner run_until_complete-under-drive-lock model with a serve/park LoopThread over SGLang's own engine.loop — callers submit coroutines threadsafe and stay in flight together; weight/memory verbs require quiesced generation and run with the loop parked (the Engine's sync wrappers still drive the idle loop themselves). - http: pure sync (urllib + threading.Semaphore + per-batch thread fan-out); the httpx client and the backend-owned event loop are gone. Controls are bounded 10s best-effort POSTs. - Backend protocol: generate is sync + thread-safe; async surface removed. New CPU tests cover the LoopThread lifecycle (concurrent overlap, semaphore bound, park/serve, quiesce guard, close-waits) and the sync HTTP backend against a stdlib ThreadingHTTPServer stub. * refactor(rollout): drop the async surface from the sync-core engines trainside, vllm_omni, sglang_diffusion, composed: delete agenerate (a to_thread wrapper over the locked sync path) and run_session + the LocalAsyncRuntime each held only to serve it; their sync generate paths are unchanged (the generate lock stays their concurrency story). Drop agenerate from the RolloutEnginePort protocol (its only consumer, AgentLoop, is sync). Delete unirl/rollout/engine/runtime.py and its test file; SessionRunner keeps a local transitional CoroutineFactory alias until the test fakes move off it. * refactor(agentic): thread-pool drain; sync env step/close; sync test suite AgenticRolloutEngine now drives trajectories on a per-drive thread pool (per_worker_concurrency threads = the trajectory cap; a thread holds its trajectory across tool-wait, preserving the two-bounds design) instead of coroutines on the inner engine's loop: _drain_worker pulls (blocking ray.get) and runs a fully-sync _run_one (inner.generate per turn + env.step + duck-typed env.close teardown). run_drain joins every drain thread and re-raises the first failure; a failed worker sets _stopping so siblings checkpoint. The barrier agenerate (test-only) is deleted; coordinator verbs and the turn-boundary checkpoint/carry semantics are unchanged, and the weight-sync quiesce is now 'abort joined the drain threads => decode-idle'. Environment protocol: astep/aclose (bridges that existed only for the deleted loop model) are removed; sync close(sample) joins reset/step as the guaranteed teardown hook (ToolEnvironment._end_sessions / AlfworldEnv._teardown_episode bodies, now public). SessionRunner is deleted (backends/base.py is protocol- only again). Tests: _fakes.py is a sync thread-safe backend/engine pair with a hold-gate for deterministic overlap; the engine tests assert the same contracts under threads (cap saturation, two distinct bounds, checkpoint/ resume, conservation, continuous batching); loop tests use step/close, with thread-based re-entrancy replacing the astep suite. * docs(rollout): sync-only contract in README + stale async docstring refs Rewrite the README's generation-interface and extending sections for the sync-only contract (threads, not asyncio; concurrent-caller requirement for agentic inners; the loop survives only inside the native SGLang backend). Update tool/env docstrings that still referenced the deleted astep/aclose/ SessionRunner, and guard the HTTP batch path against an empty wire (the old gather path returned [] there; a zero-worker pool would raise). * Fix trainside smoke replay conditions * Handle ragged AR evaluation batches * docs: document agentic integration and migration Add public Sample/Part and agent-loop guides, refresh trainer and recipe documentation, and remove stale retired-API references. Align formatting, executable modes, and pytest discovery so the upstream validation gate is reproducible; cover the merged SGLang top-k behavior. * chore: remove integration smoke artifacts * Fix Transformers checkpoint key remapping * Test model-aware checkpoint remapping * fix(sampling): require explicit generation frontiers * fix(agentic): finalize drained rollouts atomically * fix(agentic): apply safe partial-rollout tail policies * test: remove tests directory * fix(sample): resolve interior-stage conditioning at its own frontier Sample.conditioning is frontier-aligned by contract: turns() gathers one row per parts[-1] sample. UnifiedModelTrainer._build_request_sample forks hierarchically into [input, ar(P*N), image(P*N*M)], so HunyuanImage3 t2ti drove its AR pass at P*N*M width and then filled a P*N-row Part — a deterministic mismatch whenever the diffusion branch M > 1. Add Sample.conditioning_at(index), naming the parts[:index+1] replay idiom that the conditioning()/turns() docstrings already document, and rebuild t2ti on the shape BagelPipeline t2ti already uses: resolve AR prompts at the AR frontier, assert the width, fill the AR Part, then re-read conditioning() so the walk broadcasts prompt and CoT onto the image frontier. Collapse Bagel's hand-rolled equivalent onto the helper so the two t2ti implementations cannot drift. * fix(trainside): reject forward_batch_size for multi-stage Samples The chunk path assumes only parts[-1] is generated: it captures parts[:-1] once and keeps only chunk.parts[-1]. Given [input, ar, diffusion] a composed pipeline fills BOTH gen Parts per chunk, so the interior stage is regenerated and then discarded and the Sample returns carrying its original empty shell. For PE it is worse — PEPipeline.generate runs both levels inside one call and derives M = len(diff) // len(ar) from row counts, so slicing diff but not ar either raises a confusing arity error or silently mis-pairs prompts to images. Reject the combination instead. The guard keys on gen-Part count, not Part count (V2V is legitimately 3 Parts / 1 gen Part and sets forward_batch_size: 1), and fires whenever the knob is set rather than only when this batch would chunk, so a config cannot pass at one fan-out and corrupt at a larger one. The message points at per-stage chunking, which pe_sglang_* already does on the SGLang diffusion sub-engine. Swept all 113 recipes: 44 pair trainside with forward_batch_size, none of them multi-gen, so no existing recipe changes behavior. * fix(agentic): namespace fresh rollout roots per drive rollout_id alone cannot make root ids unique. _next_batch refills a short buffer by re-submitting under the SAME rollout_id — up to _MAX_REFILLS get_samples() calls per rollout step — and a data source may restart its ids on every call (DefaultDataSource numbers by batch position, not prompt identity). Two drives then namespace different source rows identically, so _GroupAssembler, which buckets purely by root id, merges siblings of unrelated prompts into one GRPO group and _build_tasks overwrites the buffered group's _gt_by_root answer. Add a monotonic _drive_seq applied to fresh roots only. Carried partials keep the ids they were submitted under and still rejoin their own siblings. * fix(agentic): mark infrastructure failures and exclude them from GRPO _run_one converted every exception into done=True with no failed status, so a backend outage, tool timeout, or context overflow entered GRPO as a legitimate low-scoring sibling and biased its whole group. _group_advantages already documents "NaN reward = crashed trajectory: excluded from the group's mean/std", but nothing ever produced that NaN, leaving the branch unreachable. Attach NaN on the failure path (dropping any partial env_reward, which does not describe a complete trajectory) while still returning done=True so the drain cannot stall. AgenticEnvTrainer now reports NaN for gen-less trajectories, and the answer-graded path overwrites graded-but-failed trajectories via a new _is_failed helper and logs the count. This intentionally reverses the previous "failed trajectory stays a legit group member" behaviour: scoring an infrastructure fault as a genuine miss manufactures a gradient for every sibling in the group. * fix(agentic): namespace partial rollout roots per drive Prevent refill batches from reusing root IDs and combining unrelated trajectories in partial-rollout GRPO groups. * Wire CUDA 13 compatibility into Taiji DRPO launcher * Source Taiji network environment before strict shell mode * Make Taiji launcher safe under inherited nounset * Deduplicate replicated DAPO source rows * Pin CUDA 13 compiler for SGLang JIT * Enable deterministic FSDP rollout parity for DRPO * Expose CUDA runtime to SGLang JIT linker * fix(veomni): select installed flash attention backend * fix(distributed): initialize transfer NCCL eagerly * fix(pe): use expandable CUDA allocator for long runs * fix(qwen3): bound cached decode growth * chore(examples): remove task-local launcher defaults * refactor(agentic): name entrypoints by capability, not benchmark The six agentic entrypoints were the only ones in the repo named after a dataset. Every other entrypoint is named for its trainer -- trainer/ar.py -> train_ar.py, whose docstring states outright that it serves both qwen_vl and qwen3 recipes. Apply that existing convention: train_deep_research.py -> train_agentic.py train_partial_deep_research.py -> train_agentic_partial.py train_async_deep_research.py -> train_agentic_async.py train_alfworld.py -> train_agentic_env.py train_partial_alfworld.py -> train_agentic_env_partial.py train_async_alfworld.py -> train_agentic_env_async.py Each name now maps 1:1 onto its trainer module, and both axes are legible: _env is the reward source (environment return vs graded terminal answer), _partial/_async is the execution topology. Topology is a suffix so the whole family groups under train_agentic*; train_async_ar.py keeps its prefix rather than widen the diff. Recipes stay benchmark-named -- the entrypoint names the capability, the recipe names the task. Each docstring gains a train_ar.py-style line naming the recipe family it serves, and the README tables become reward-source x topology matrices instead of benchmark lists. Bodies are unchanged: same trainer classes, same constructor arguments, same config_name defaults. No behavioral change. Note: python -m unirl.train_alfworld and train_deep_research no longer exist; six user-facing module paths moved. * chore: apply ruff-format to qwen3-omni sample-native files Three files added by adf70e8 were never run through ruff-format, so `pre-commit run --all-files` -- which the Lint workflow runs on every PR to main -- fails on this branch. Formatting only; no logic touched. * fix(agentic): score env-reward faults NaN on every path, not just the barrier 7ad5b34 established that an infrastructure fault must be excluded from GRPO rather than scored as a genuine miss -- "scoring an infrastructure fault as a genuine miss manufactures a gradient for every sibling in the group" -- but it only reached AgenticEnvTrainer. The colocate-partial and fully-async env variants each carried their own copy of _rewards_and_groups and kept the pre-fix 0.0, agentic_env_async.py even retaining the comment "gen-less / failed trajectory stays a legit group member" that the fix says it reverses. Concretely: when a trajectory faults before its first turn the engine attaches no reward (_attach_env_reward returns early on an empty gen list), so the trainer's missing-reward branch decides the sentinel. NaN is dropped from the group's mean/std and given zero advantage; 0.0 instead drags the mean, shifts the std, and hands every sibling a gradient the model never earned. A fault AFTER turn one was already fine on all paths -- the engine attaches NaN there and it passes straight through. Rather than patch two literals, hoist the method into a shared _EnvRewardSource mixin that all three env trainers mix in ahead of their trainer base. One implementation, so the next sentinel change cannot land on some paths and miss others. Net -83/+30 in the trainers; the reward SOURCE becomes a named concept matching how the docstrings already describe it. Tests, restoring coverage this branch lost: - tests/rollout/test_agentic_failure_marking.py, deleted by c566479 "test: remove tests directory" while the fix it guarded stayed. Restored verbatim; it still passes against current code. - tests/trainer/test_agentic_env_reward_source.py, new. The gap that let this bug through: nothing ever asserted the env TRAINERS' sentinel. Parametrized over all three env paths, plus a structural test that they resolve to one shared function. Verified against the pre-fix tree: exactly 3 failures, naming the two buggy paths and the divergence itself. 18 pass after the fix. * test: remove tests directory Removes the 13 remaining test files (1117 lines), including the two added in 57edcca. No workflow in .github/ runs pytest and the root pyproject declares no testpaths, so nothing in CI changes. * refactor(examples): drop the M1 calculator scaffold, port M2 to all topologies deep_research_calc_mathverify{,_partial,_async}.yaml were bring-up scaffolding: M1 existed to prove the agentic loop "WITHOUT external services (the existing CalculatorTool + a rule-based MathVerify reward) ... so M2 only swaps in the search/visit tools and the LLM-judge reward (config-only)". M2 landed and the scaffold has no remaining purpose. The catch: all three were the hydra config_name DEFAULTS for the three answer-graded entrypoints, and search_judge existed in barrier flavour only. Deleting them would have left train_agentic_partial and train_agentic_async with no recipe. Repointing both at the barrier recipe is not an option either: it declares TensorWeightSync, so the async entrypoint would compose and then fail at runtime, since disaggregated training needs NCCLWeightSync to cross the slab boundary. So port M2 to the two missing topologies, grafting each one's topology block off the calc recipe it replaces and leaving tools/reward/system-prompt identical: deep_research_search_judge_partial.yaml colocate, TensorWeightSync, oversample_batch_size 6, partial_rollout, worker_max_concurrency 24 deep_research_search_judge_async.yaml disaggregated, NCCLWeightSync, train_fraction 0.5, oversample_batch_size 8, mem_fraction_static 0.8 tail_policy: carry is correct for both -- SearchTool and VisitTool subclass the stateless Tool, not StatefulTool, so a carried Sample holds all resume state (the same reason the calc recipes carried). Verified: all six agentic entrypoints compose against their defaults with --cfg job --resolve, and the resolved configs carry the right weight-sync class, mem_fraction_static, and prompts_per_rollout == oversample_batch_size per topology. check-recipe-targets resolves 2135 paths; hooks and bash -n clean. Neither new recipe has been train-verified -- they are compose-checked ports. CalculatorTool itself is left in place; no recipe uses it now, but it stays a supported library tool. --------- Co-authored-by: Jianghai <72591262+CjhHa1@users.noreply.github.com> Co-authored-by: CjhHa1 <cjh18671720497@outlook.com> Co-authored-by: leviking98z-rgb <leviking98z@gmail.com>
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
1 participant
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
Part 2/3 for Tencent-Hunyuan#86 (PPO value critic + GAE on AR models). Stacks on Tencent-Hunyuan#254 (
feat/gae-advantage).ValueHeadandARReplayOutput; optionaluse_value_headon Qwen3 attaches a per-token critic during replay.TextSegmentwith packedvalues,returns, andtoken_advantages.scatter_terminal_rewards(scalar outcome reward → last response token) andRolloutTrack.compute_gae_advantages(per-sample GAE viacu_seqlens, no cross-trajectory leak).return_values=Truewithout a value head.Not wired into trainer or PPO loss yet — that is PR 3.
Related Issue
Part of Tencent-Hunyuan#86 (roadmap Tencent-Hunyuan#25). Stacks on Tencent-Hunyuan#254.
Test Plan
CPU (Python 3.12+ env):
Locally verified on Python 3.9 for the track/GAE tests above; full project pytest needs 3.12+.
GPU end-to-end (Qwen3 +
use_value_head=True+ replay values): not run in CI here; follow-up in PR 3 recipe.Compatibility / Risk
use_value_head: bool = Falseon Qwen3 — default off, no checkpoint/API break for existing GRPO runs.Reviewer Notes
main, rebase this branch ontomainbefore upstream review.Checklist