Fast-LLM trainer integration with vLLM v1 weight broadcast by bigximik · Pull Request #140 · ServiceNow/PipelineRL

bigximik · 2026-05-06T08:09:54Z

Status: WIP — handover from Denis (2026-05-06)

This branch is not ready to merge. It's the in-progress integration of Fast-LLM as an alternative trainer to DeepSpeed, with weight broadcast to vLLM v1 over a persistent NCCL group instead of HTTP. I'm leaving the integration project — this PR captures everything needed to pick it up.

Read this first: docs/FAST_LLM_INTEGRATION.md — canonical handover (architecture, per-file changes, glossary, all known issues with file:line citations, testing guide, operations notes, open questions).

Stats: 79 commits ahead of main, ~8 400 insertions / 195 deletions across 35 files (mostly new tests + integration plumbing + handover docs).

What works today

Fast-LLM (gspo branch) trainer launches under torchrun, joins a persistent NCCL broadcast group, and pushes weights to vLLM v1 workers in place. No HTTP weight upload.
Coordinated NCCL teardown (training_finished event over redis → vLLM destroys process group → both sides hit the collective barrier together) — dist.destroy_process_group() no longer hangs.
4-node multi-node smoke verified end-to-end on both fast-llm GSPO and DeepSpeed PPO (see "Smoke result" below).
GSPO loss math matches DeepSpeed exactly: grad_norm parity, grpo_new_logprobs matches step-by-step over a 400-step run (see chart below).

Companion Fast-LLM PR

This PipelineRL branch pins to the gspo branch in Fast-LLM (PR #502). The Fast-LLM PR contains:

GSPO loss kernel (sequence-level geometric-mean IS-ratio clipping)
Decoupled loss/gradient divisors (loss /num_documents, grad /num_documents²) + SDP loss correction — exact match to DeepSpeed's 1/batch_size dual-factor math
fp32_lm_head flag matching vLLM's bf16_last_layer_fp32 precision (otherwise IS ratios drift)
metrics: GRPOMetricsLevel enum (none/basic/with_entropy) — merged from PR #494 (Joel's metrics refactor)

Once that PR merges to Fast-LLM main, the README install step here should be revved from git checkout gspo → git checkout main and this PR rebased onto a fresh main.

What's NOT done yet

Fix actor _prefetch_to_doc_target overshoot (pipelinerl/actor.py:613). Causes premature run end on long runs (50+ steps). Workaround: bump max_train_steps ~20%. Real fix: trainer signals "done" instead of actor inferring.
Address rollout retry exhaustion under bursts (pipelinerl/async_llm.py:137-146). Two consecutive aborts can drop a rollout permanently. Allow more retries or evict stuck rollouts.
Investigate reward lag vs DS (~2-point gap at step 400 in actor/reward_mean — see chart below). Root cause unknown; newlp parity is confirmed so the gap is upstream of the trainer.
Resolve commented-out pyproject.toml overrides (pyproject.toml:81-87). The [tool.uv] block force-overrides transformers>=4.51.0 / accelerate>=1.7.0 because tapeagents==0.1.16 pins them lower; [tapeagents] extra is broken at runtime. Either bump tapeagents or drop the extra on this branch.
Close fast-llm finetune metric gaps, e.g. rl/ess (effective sample size — diagnostic for data/policy drift).
Bump base image + vLLM version. Currently pinned to interactive-toolkit:25.12-py3-vllm014rc1redis (PyTorch 25.12, vLLM 0.14.0rc1). Move to the latest base PyTorch + vLLM that both Fast-LLM and PipelineRL support; re-run smoke after.

Known issues (with code references)

Issue	Symptom	Site	Memory ref
Actor overshoot ends runs early	`TimeoutError: No document received after 600s` near final step	`pipelinerl/actor.py:158, 613-614`	`project_actor_samples_target_overshoot_bug.md`
Rollout retry exhaustion	Rollout stuck in actor's `in_progress` after `attempt=2/2 abort`	`pipelinerl/async_llm.py:137-146`	`project_stall_investigation.md`
Reward lag vs DS	`actor/reward_mean` ~2 points below DS at step 400	unknown (upstream of trainer)	`project_fastllm_reward_lag_after_gspo_fix.md`

Current limitation (not a bug): streams=files is not implemented for use_fast_llm=true — Fast-LLM only ships RedisStreamingDataset. Use streams=redis. See project_streams_files_not_supported_fast_llm.md.

Training curves (400-step run): fast-llm GSPO vs DeepSpeed GSPO

Compared runs:

fast-llm: math_7b_4node_fastllm_gspo_20260505_122944 (divisor² + SDP fix)
DS: math_7b_ds_fastllm_4node_20260428_135427 (matching GSPO config: policy_loss=gspo, epsilon_low=3e-3, 400 steps)

new_logprobs — fast-llm matches DS step-by-step (the GSPO loss math fix is correct):

actor/reward_mean — fast-llm lags DS by ~2 points at step 400 (open issue):

How to verify locally

See examples/interactive/fast_llm_4node.sh and examples/interactive/ds_4node.sh — both follow the README install.

# inside an interactive 4-node EAI session, after the README install:
bash examples/interactive/fast_llm_4node.sh   # fast-llm + vLLM v1 + GSPO
bash examples/interactive/ds_4node.sh         # DeepSpeed + vLLM v1 + PPO (reference)

Both run a 2-step smoke and finish in ~10 minutes. Override MAX_TRAIN_STEPS=N for longer runs.

Smoke result (last verified 2026-05-06)

Smoke	EAI Job	Step 1 grad_norm	Step 2 grad_norm	Step 1 newlp	Step 2 newlp	NaN
fast-llm GSPO	`59f3b62f`	0.166	0.173	-0.171	-0.162	0
DeepSpeed PPO	`084ef7d8`	0.201	0.247	-0.162	-0.146	0

Per-step wall time ~80–120 s for both — fast-llm and DS run at comparable speed at this scale.

Code change summary

See docs/FAST_LLM_INTEGRATION.md §5 "Per-file changes" for the file-by-file table. Highlights:

pipelinerl/launch.py: TCPStore pre-creation for broadcast rendezvous (workaround for torchrun client-only TORCHELASTIC_USE_AGENT_STORE=True); fast_llm.callbacks.streaming.broadcast.* injection.
pipelinerl/state.py: fast-llm event-stream listener thread; samples_processed=0 initialization to avoid startup deadlock.
pipelinerl/vllm1.py: init_actor_update_group/destroy_actor_update_group with WEIGHTS_BROADCAST_PG_NAME; training_finished handler for coordinated NCCL teardown.
pipelinerl/async_llm.py: rollout retry on vLLM aborted request (weight-update collision).
tests/: weight-broadcast tests (test_vllm1_fast_llm_broadcast.py), full vLLM v1 integration (test_vllm1_integration.py), multi-node topology (test_world_multinode.py), actor error handling.

Reviewer checklist

This is a draft PR for handover, not for merge. Reviewer should:

Read docs/FAST_LLM_INTEGRATION.md end-to-end.
Skim README §"Install FastLLM+PipelineRL" — it should reproduce on a fresh interactive job.
Run bash examples/interactive/fast_llm_4node.sh and confirm step 1-2 metrics in finetune/stdout_node0.log.
Pick up the TODO list above; create separate issues/PRs for each item.

…stop traning was not received

… better abab pattern detection in generations results to test weight bradcast correctnes, some refactoring

…fig for testing

[WIP] Adding tests to vllm actor for Fast-LLM integration

The two vLLM output-path hooks run per generated token / per response. Guard the annotation blocks so an unexpected shape (or a future vLLM change) disables per-token tagging and falls back to the per-rollout version instead of raising into the output processor or response builder. Self-disables after the first failure (logged once) so it degrades cheaply rather than throwing per token. Verified the patched seams exist and match on vLLM 0.16.0, 0.18.1, and 0.24.0: - OpenAIServingChat._create_chat_logprobs in chat_completion/serving.py, kwargs call convention, token string is `token_id:<id>` (via _get_decoded_token / format_token_id_placeholder). - LogprobsProcessor.update_from_output in v1/engine/logprobs.py, self.logprobs a list[dict[token_id, Logprob]] (FlatLogprobs alternative is guarded). - Logprob is a mutable dataclass; OpenAIBaseModel is extra="allow" and not frozen. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…hook - Log the raw trainer step (`completed_step`) in the actor's published stats alongside `trainer_model_version`, so it is no longer write-only. - Make the `update_from_output` hook forward `*args/**kwargs` to the original, matching the `_create_chat_logprobs` seam, so a vLLM signature drift degrades to the per-rollout fallback instead of raising into the output processor. - Rename `before` -> `previous_length`; trim two comments (testing directive and an overstated guarantee). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

- tests/test_model_version.py: cover parse_token_id_and_version (with/without the :v<n> suffix, v0, multi-digit) and convert_to_fast_llm_format's model_version padding/fallback (per-token left-pad, pad-with-first when no scalar, scalar broadcast, absent, full-completion). Logic previously verified standalone. - Reword the serving-class import comment: the chat_completion package predates 0.18.1, so drop the specific-version attribution. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Forward raw reward, document-count model version, and per-token model version

Brings 28 commits of PipelineRL main into the fast-llm integration branch (PR #140). One content conflict, in pipelinerl/async_llm.py: main added tool-call parsing + content-or-empty for the TIR domain in the same region where fast-llm raises asyncio.TimeoutError on finish_reason=="abort" (vLLM weight-update-collision retry). Resolved as a union: keep the abort retry, then build LLMOutput with tool_calls. All fast-llm-specific plumbing (weight broadcast in vllm1.py, use_fast_llm dispatch in launch.py, the redis event listener in state.py) is untouched by main and preserved. New from main: tir domain, memory_debug, and finetune_loop/preprocess/rollouts updates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

_pause_generation detected the installed vLLM API at runtime and, on older builds lacking pause_generation(mode=), fell back to wait_for_inflight_requests=False — which aborts in-flight requests on every weight-update pause (the abort-storm cause). PipelineRL now pins vLLM 0.18.1, which has mode=, so the fallback branch is dead. Collapse to the mode="keep" path and drop the now-unused inspect import. The finish_reason=="abort" handler in async_llm.py is left in place as a defensive path in case aborts are reintroduced. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The bf16_last_layer_fp32 quantization forces a bf16 body with an fp32 lm_head, overriding --dtype. For a fp16/fp32 inference run (e.g. matching a fp16/fp32 trainer to keep the importance ratios exact), it silently kept the body at bf16. Gate it on vllm_config.vllm_kwargs.dtype: skip the quant for non-bf16 dtypes so vLLM honors --dtype; bf16/auto/unset are unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Single source of truth for the run's hyperparameters (both the PipelineRL/DS side and the fast_llm subtree), so the launch script carries no config. Compose with: python -m pipelinerl.launch --config-name math_qwen05_gspo. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…atch pairs) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Replace the single math_qwen05_gspo.yaml with a shared base (math_qwen05_gspo_common: sampling/topology/preprocessor RL prep/orchestration) plus one file per trainer backend: math_qwen05_gspo_fllm (Fast-LLM) and math_qwen05_gspo_ds (DeepSpeed). Must-match trainer values are interpolated from the shared base (micro_batch_size<-seq_length, train_iters<-max_train_steps, FL loss epsilon<-finetune.rl.epsilon) so they can't drift. Launch with --config-name math_qwen05_gspo_fllm | _ds. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…config base.yaml's fast_llm node inlined a full slate of training hyperparameters (train_iters, micro_batch_size, schedule, the loss, zero_stage, compute_dtype) — asymmetric with the DeepSpeed side (whose hyperparameters live in the finetune: base group, not inline). Reduce base's fast_llm node to contract + scaffold only (launcher-filled nulls, streaming dataset, weight-broadcast callback, pretrained loader, RL-streaming invariants). Move the hyperparameters into the experiment config, which already set most; math_qwen05_gspo_fllm now also sets compute_dtype and export.interval and defines the loss under a clean 'gspo' key. launch.py only fills nulls in kept sections + reads data.micro_batch_size (supplied by the experiment config), so this is safe; counting.yaml already carries its own fast_llm node. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

These shell scripts hardcoded deployment-specific infrastructure details (container image, storage paths, run/identity values) that don't belong in a public repository. The reproducible recipe lives in the committed hydra configs (conf/math_qwen05_gspo_*.yaml); launching is environment glue kept out of the repo. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

math_qwen7b_gspo_fllm and math_qwen7b_gspo_ds inherit the 0.5B per-backend configs and override only the model and the memory knobs the larger model needs (ZeRO-2, MLP recompute), keeping the shared recipe single-sourced. Node count is chosen by the launcher (env WORLD_SIZE), not the config. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

get_optimizer threads betas into the adamw_torch / cpuadam optimizers, sourced from finetune.adam_beta1 / adam_beta2 (defaults 0.9 / 0.999 = torch's, so existing runs are unchanged). The GSPO DeepSpeed configs set the sqrt-rule betas directly, replacing the ADAM_BETA1/2 environment hack. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Wrap load_dataset so the existing call sites keep passing trust_remote_code: it is forwarded when the installed datasets accepts it and dropped otherwise, instead of raising a TypeError on newer datasets. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Same load_dataset wrapper as the math domain, so chartqa's load_dataset calls don't raise on newer datasets. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The Fast-LLM streaming callback broadcasts the cumulative document count under key `documents_seen`, but both weights_ready consumers read `document_count`, which is never present. `event.get("document_count")` was always None, so the model version silently fell back to `step`. The per-token model_version (vllm1.py tags each generated token) and the propagated weight version (state.py) then rode the step clock (~1/step) while the trainer's staleness metric subtracts them from documents_seen (~256/step), so gspo_staleness grew unbounded (≈ documents_seen) instead of reporting a bounded lag. DeepSpeed uses the legacy version field and was unaffected. Metric-only impact on the current runs (max_lag: null, pop_old_data: true -> staleness never gates filtering); corrects the logged staleness for future runs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The wandb_name config field existed but was ignored: init_wandb always set the display name from the run-dir path, so runs showed up as the full `pipelinerl_runs/pipelinerl_<name>_<ts>/<component>` path — unreadable in a legend. Use cfg.wandb.wandb_name (+ the component suffix, so per-service runs stay distinguishable) as the display name when set; fall back to the run-dir path otherwise. The run id stays derived from the unique run-dir path, so a short name can repeat across experiments without colliding. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

- Delete the stale FAST_LLM_INTEGRATION.md handover doc and its chart images - Remove the orphaned RedisSharedStreamReader (no caller; Fast-LLM reads the stream via its own dataset) and revert read_stream to its single-arg form - Drop the commented-out old DeepSpeed run_finetune block in launch.py - Remove unused vllm1.py imports and the test-only is_extension_loaded RPC - Strip the PIPELINERL_DEBUG logger reconfig and verbose logger.debug tracing from both weight-update receive paths - Prune stale README bits (dead handover-doc links, WIP status banner, outdated read_stream/shared doc) - Drop TestWorkerExtension.test_extension_loaded (exercised the removed RPC) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Bring the extracted trainer-agnostic fixes (PR #151) in as the base of the Fast-LLM integration so #140 stacks on top and its diff shows only the integration-specific delta. Conflicts (the review improvements #151 gained after extraction) resolved in favor of #151: dedup the datasets>=4 load_dataset shim into pipelinerl/domains/_datasets_compat.py (drop the now-dead `import inspect` from both loaders), drop the invalid "bf16" vLLM dtype token, and label the unpacked preprocess log correctly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Strip the debug.log_data_pipeline instrumentation from actor.py (the Process-A and Process-B jsonl loggers, inline json/pathlib imports, and per-cycle timing probes) and drop the config flag. This was ad-hoc debug scaffolding; removing it de-weaves the trainer-agnostic rollout-robustness changes so they can be extracted separately. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

The previous commit removed the actor.py jsonl diagnostics that consumed this flag, but it is also read by the fast-llm data-pipeline wiring in launch.py and preprocess.py (which enable fast-llm's built-in input-starvation logging), so the key must stay. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

- preprocess.py: _check_group_sizes accepts 1..group_size, matching the actor's relaxed group-size assert (a 4xx-skipped or empty rollout is filtered out, so a group may publish fewer than `attempts` rollouts); the strict == check crashed the preprocessor on those. - async_llm.py: remove the unreachable `finish_reason == "abort"` block (aborts are already handled by the retry loop and make_training_text) and its now-unused `import asyncio`. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Brings in the lag min/mean/max metric change (#152).

rafapi and others added 30 commits December 12, 2025 14:09

basic launcher

d33420b

single channel streams

acad1e7

update

9a7eab0

Merge branch 'main' into fast-llm

b526d8d

enable fast-llm for basic streaming

9543265

enable fast-llm for basic streaming

1d9ac8d

set cwd

6d42ade

convert samples for fast-llm format

e8674e1

listen to trainer events from redis

e78844e

put all the samples in a single stream

d16a829

fix accessing non complete dists

67b3f2b

change to calling fast-llm without conda

dc04911

fix imports for vllm .14.1rc1

30a5346

tmp changes to not install dependencies which are in base image already

9df2feb

2 gpu and 1 gpu integration tests

49707aa

changes for tests

a893686

added instruction to install PipelineRL+FastLLM

641e457

added fast-llm bradcast functionality and test

7d14e80

refactoring of traning helper

dad6242

added stop traning support and changed engine clean up logic to warn …

402f29a

…stop traning was not received

added 3 and 4 gpu tests (not run yet), added traning end event, added…

3848fbc

… better abab pattern detection in generations results to test weight bradcast correctnes, some refactoring

fix setting current device for tp and pp cases

e170185

fix multi actor generation abab patter consistency check

3541dfe

fast-llm weight update bug fix some other changes

5560ada

note update

f4cc7a4

fix no weight bradcast case with fast-llm

f78be64

added data for grpo loss to send to fast-llm, added base fast-llm con…

807f47c

…fig for testing

fast-llm weights broadcast integration

407be1d

removed duplicate option

b7e2109

Merge pull request #128 from ServiceNow/fast-llm

a3f6ed2

[WIP] Adding tests to vllm actor for Fast-LLM integration

jlamypoirier and others added 4 commits July 7, 2026 17:04

Merge pull request #150 from ServiceNow/jlp_rl_metrics

c77bdda

Forward raw reward, document-count model version, and per-token model version

jlamypoirier changed the title ~~[WIP] Fast-LLM trainer integration with vLLM v1 weight broadcast — handover~~ Fast-LLM trainer integration with vLLM v1 weight broadcast Jul 10, 2026

jlamypoirier and others added 16 commits July 10, 2026 09:02

Config comment: describe variants abstractly (no dangling filenames)

69fc60c

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Annotate finetune vs fast_llm config sections (backend split + must-m…

c9bbb06

…atch pairs) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Apply the datasets>=4 trust_remote_code shim to chartqa too

75d5892

Same load_dataset wrapper as the math domain, so chartqa's load_dataset calls don't raise on newer datasets. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

jlamypoirier mentioned this pull request Jul 14, 2026

Trainer-agnostic fixes extracted from the fast-llm branch #151

Open

jlamypoirier changed the base branch from main to jlp_pipelinerl_unrelated_fixes July 14, 2026 22:48

jlamypoirier and others added 6 commits July 14, 2026 19:24

Merge branch 'jlp_pipelinerl_unrelated_fixes' into fast-llm

fdb36d0

Merge branch 'jlp_pipelinerl_unrelated_fixes' into fast-llm

3f2bbb5

Merge jlp_pipelinerl_unrelated_fixes into fast-llm

7923420

Brings in the lag min/mean/max metric change (#152).

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Fast-LLM trainer integration with vLLM v1 weight broadcast#140

Fast-LLM trainer integration with vLLM v1 weight broadcast#140
bigximik wants to merge 129 commits into
jlp_pipelinerl_unrelated_fixesfrom
fast-llm

bigximik commented May 6, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

bigximik commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Status: WIP — handover from Denis (2026-05-06)

What works today

Companion Fast-LLM PR

What's NOT done yet

Known issues (with code references)

Training curves (400-step run): fast-llm GSPO vs DeepSpeed GSPO

How to verify locally

Smoke result (last verified 2026-05-06)

Code change summary

Reviewer checklist

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

bigximik commented May 6, 2026 •

edited

Loading