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

# Conflicts: # pipelinerl/state.py # pipelinerl/vllm1.py

…ck, drop dead code - Coerce samples_processed to int in the preprocessor progress log, avoiding a TypeError before the first SamplesProcessed arrives on the non-Fast-LLM path. - Consolidate the trainer-finished check into TrainerState.is_finished(), replacing three near-identical copies across actor.py and preprocess.py. - Extract a _node_suffix() helper shared by both finetune launchers. - Remove the dead create_engine(cleanup=...) parameter and teardown block (the sole caller always passes cleanup=False) and its now-unused import. - Inline the single-use write_sample_for_fast_llm wrapper. - Type-hint read_fast_llm_events; rename r/gname to full names; regroup vllm1 imports; drop a forward-ref quote, restating comments, and the _inspect alias. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

# Conflicts: # pipelinerl/preprocess.py

- Drop the internal container-registry FQDN from the 'use prebuilt image' step; keep the image tag and reference the internal registry generically. - Genericize the EAI_PROFILE datacenter code to a placeholder. - Fast-LLM's `gspo` branch was deleted after its functionality merged to main; drop the now-broken `git checkout gspo` step so the install instructions work against main. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

`_run_finetune_fast_llm` derived `experiment_name` from the save_dir path and set it unconditionally, so the Fast-LLM trainer's wandb run always logged under the run-dir path (e.g. `pipelinerl_runs/.../finetune`) instead of grouping with the actor/preprocess runs, which `init_wandb` names `{wandb_name}/{component}`. Prefer `{wandb_name}/finetune` when a run name is set; fall back to the old path-derived name otherwise. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…configs - vLLM per-token model-version tagging: producer patches in vllm1.py plus consumers in async_llm.py, llm.py, and rollouts.py. - Preprocessor/state/actor RL metrics: reward and model-version pass-through, completed-step back-pressure, and the data-pipeline logging flag. - GSPO experiment configs (Qwen 0.5B / 7B, DeepSpeed and Fast-LLM variants), the counting task config, and the conf/base.yaml logging flag. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Name the Fast-LLM finetune wandb run from wandb_name

- conf/counting.yaml: the fast_llm GRPO loss entry had no `type:`, so it fails to dispatch on the Fast-LLM backend (base defines no default loss). Add `type: grpo`. - preprocess.py: the debug data-pipeline log file was opened outside the stream `with` block and never closed. Register it on a contextlib.ExitStack entered by that block so the handle closes on exit. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

`_run_finetune_fast_llm` derived `experiment_name` from the save_dir path and set it unconditionally, so the Fast-LLM trainer's wandb run always logged under the run-dir path (e.g. `pipelinerl_runs/.../finetune`) instead of grouping with the actor/preprocess runs, which `init_wandb` names `{wandb_name}/{component}`. Prefer `{wandb_name}/finetune` when a run name is set; fall back to the old path-derived name otherwise. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

- vllm1.py: _create_chat_logprobs captures *args/**kwargs like update_from_output, so a vLLM signature change disables tagging instead of crashing at bind time; bounds-guard token_ids[index] to match top_logprobs[index]; log {error!r} on import failure. - preprocess.py: hoist `import json` to module top, use the existing Path import, attribute-access cfg.debug.log_data_pipeline, and use an identity check for pipeline_log_file. - math_qwen05_gspo_fllm.yaml: drop inert decay_iterations under a constant schedule; reword the losses comment. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Test layer of the stacked breakdown of PR #140, extracted on top of the config/metrics PR (#155). Covers the vLLM v1 weight-broadcast path, world setup, actor error handling, launch-process monitoring, and model-version tagging. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

_run_finetune_fast_llm referenced finetune_rank in its multi-node torchrun args, but the name is only bound inside _node_suffix. Define it in the multi-node branch where it is used. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Apply the coarse + fine review findings: - Delete test_actor_error_handling.py — every test reimplemented the actor error-handling logic inline and asserted on the copy, so it could not catch a regression in the real actor. - Drop the generic-path topology matrix (TestWeightUpdateTP2, TestWeightUpdateMultiActor); equivalent multi-topology coverage lives on the Fast-LLM path in test_vllm1_fast_llm_broadcast.py. - Move the EngineManager import out of conftest top level into the factory fixture so the CPU-only unit tests collect without vLLM/torch. - Consolidate the duplicated stream_process_output / kill_process_tree helpers into server_weight_update_utils.py and drop the now-pointless injected stream_process_output_fn parameter. - Remove dead helpers (force_kill_process, check_pattern_detected), unused fixtures/params/imports, and the dead timeout parameter. - Minor typing and style fixes (builtin generics, strict zip, f-string and implicit-Optional cleanups). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

# Conflicts: # tests/conftest.py # tests/distributed_trainer_helper.py # tests/fast_llm_trainer_helper.py # tests/server_weight_update_utils.py # tests/test_vllm1_fast_llm_broadcast.py # tests/test_vllm1_integration.py # tests/test_world_multinode.py # tests/trainer_test_utils.py # tests/weight_update_utils.py

…r exit code - Delete TestPodIPExchange, TestHostfileCreation, TestRedisHostMultiNode: each asserted on a hand-copy of launch.py logic (_simulate_pod_ip_exchange, _compute_streams_host, the hostfile host-list) rather than the real code path, so a regression in launch.py would leave them green. The helper _simulate_pod_ip_exchange stays as setup for TestDeepSpeedCommand, which drives the real _run_finetune_deepspeed. - Assert trainer_proc.returncode in (0, None) after the wait loop in the server / broadcast pattern tests, so a trainer that crashes during final cleanup fails instead of passing on the already-captured pattern. - Fine cleanups: drop unused `import tempfile`, redundant local re-imports (argparse-as-ap, omegaconf, per-function pathlib/sys.path), unused thread handles and loop counter; hoist the broadcast import out of the inner closure; return the dtype string directly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…suite run Running the suite for the first time exposed 11 failures in test_world_multinode.py, all test-setup gaps (not product bugs): - 8× `Missing key wandb_name`: the fast-llm test cfgs (_capture_cmd, _capture_fast_llm_files) omitted wandb.wandb_name, which _run_finetune_fast_llm reads. Add it (None) so the code reaches the torchrun/naming assertions. - 3× `WorldMap has no dns_address_map`: multinode finetune always runs after _exchange_pod_ips (which sets dns_address_map), but _make_world_map produced a pre-exchange map. Set dns_address_map for world_size > 1 in the helper to mirror production. Delete test_deepspeed_multinode_uses_dns_names_without_exchange: its "without exchange" premise is unreachable for multinode, and the after-exchange sibling already covers the DNS-in-filter assertion (plus the no-pod-IP-leak check). Verified on GPU: test_world_multinode + test_launch_process_monitoring now 29 passed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

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 8 commits July 16, 2026 12:45

Merge branch 'jlp_fast_llm_core' into fast-llm

0c810be

# Conflicts: # pipelinerl/state.py # pipelinerl/vllm1.py

Merge branch 'jlp_fast_llm_core' into fast-llm

bccc00a

# Conflicts: # pipelinerl/preprocess.py

Merge branch 'jlp_fast_llm_core' into fast-llm

33442e9

Merge branch 'jlp_fast_llm_metrics_configs' into fast-llm

2f5526a

jlamypoirier changed the base branch from jlp_fast_llm_core to jlp_fast_llm_metrics_configs July 16, 2026 21:47

jlamypoirier mentioned this pull request Jul 16, 2026

Fast-LLM RL: per-token model-version tagging, metrics, and GSPO configs #155

Open

jlamypoirier and others added 9 commits July 17, 2026 12:27

Merge pull request #154 from ServiceNow/jlp_fix_fl_finetune_wandb_name

e5c8a1d

Name the Fast-LLM finetune wandb run from wandb_name

Merge branch 'jlp_fast_llm_metrics_configs' into fast-llm

dfb0f64

Merge remote-tracking branch 'origin/fast-llm' into fast-llm

ad93e3e

Merge branch 'jlp_fast_llm_core' into jlp_fast_llm_metrics_configs

3a18787

Merge branch 'jlp_fast_llm_metrics_configs' into fast-llm

d948d52

Merge branch 'jlp_fast_llm_metrics_configs' into fast-llm

77c6262

jlamypoirier changed the base branch from jlp_fast_llm_metrics_configs to main July 17, 2026 18:25

jlamypoirier mentioned this pull request Jul 17, 2026

Fast-LLM integration: test suite #156

Open

jlamypoirier and others added 8 commits July 17, 2026 16:15

Fix NameError in multi-node Fast-LLM launch

89df400

_run_finetune_fast_llm referenced finetune_rank in its multi-node torchrun args, but the name is only bound inside _node_suffix. Define it in the multi-node branch where it is used. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Merge branch 'jlp_fast_llm_core' into jlp_fast_llm_metrics_configs

aa96f38

Merge branch 'jlp_fast_llm_metrics_configs' into jlp_fast_llm_tests

8689150

Merge branch 'jlp_fast_llm_metrics_configs' into fast-llm

53033d2

Navigation Menu

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 175 commits into
mainfrom
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