Add Fast-LLM as an alternative trainer (core integration) - #153
Open
jlamypoirier wants to merge 8 commits into
Open
Add Fast-LLM as an alternative trainer (core integration)#153jlamypoirier wants to merge 8 commits into
jlamypoirier wants to merge 8 commits into
Conversation
Wire Fast-LLM in as a drop-in replacement for the DeepSpeed trainer, selected by `use_fast_llm`. The DeepSpeed / HTTP path is unchanged when the flag is off. - launch.py: split run_finetune into a DeepSpeed and a Fast-LLM branch; launch `fast-llm train` via torchrun, build its config from the experiment config, pre-create the NCCL-broadcast TCPStore on rank 0, and thread use_fast_llm/weight_broadcast through process supervision. - vllm1.py: EngineManager owns the vLLM engine lifecycle and a Fast-LLM weight-update path that joins the trainer's persistent NCCL broadcast group, receives weights via a main-process Redis monitor thread, stamps the active model version on generated tokens, and tears the group down on training_finished. The HTTP weight-update endpoint is registered only in HTTP mode. - streams.py: RedisSharedStreamWriter for many-producer fan-in to a single Redis stream (orjson payloads), used by the preprocessor. - preprocess.py: convert preprocessed samples to Fast-LLM's streaming document format and write them to the shared stream. - state.py: listen to Fast-LLM's Redis event stream (weights_ready / training_finished) instead of the legacy trainer messages. - actor.py: stop the actor loop on the explicit training_finished event under Fast-LLM. - conf: use_fast_llm / weight_broadcast switches and the fast_llm trainer config scaffold; math.yaml enables the path. - Multi-node support (pod-IP exchange, whole-node GPU snapping, torchrun rendezvous) rides along and is inert on a single node. Not yet runtime-tested. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- conf/base.yaml: default use_fast_llm to false so DeepSpeed stays the default trainer (conf/math.yaml opts in); restore output_dir: ??? so a missing value fails with a clear mandatory-field error, not Path(None). - streams.py: drop RedisSharedStreamWriter's unused per-entry metadata (index/ts/writer_id) and mode="w" handling — the only caller writes the plain payload. README updated to match. - state.py / vllm1.py: share the fast_llm_events stream/payload constants and a fast_llm_event_version() helper between the two event consumers. - vllm1.py: bind manager/weight_update_mode before the try in create_engine so a failed engine assert isn't masked by a NameError in the finally; order @staticmethod outside @asynccontextmanager. - launch.py: keep the DeepSpeed finetune subprocess on the launcher's inherited stdout/stderr (leave the disabled path unchanged). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
C1: remove per-token model-version tagging producer from the core PR.
Its consumers (the async_llm/llm logprob parsers) live in a later PR, so
here it appended `:v<version>` to `token_id:<id>` tokens that the base
parsers `int(...split(":")[-1])` then dropped, corrupting logprobs. The
tagging feature (producer + consumer) belongs together in the metrics PR.
- F1: create_engine docstring — yields EngineManager, not a tuple.
- F2: drop redundant pre-raise logging + re-raise-only try/except in load_weights.
- F3: compute include_callbacks once (removes De Morgan-inverse guards).
- F4: drop dead inner imports of asyncio/time (already module-level).
- F5: unpack param_name directly (drop needless layer_name alias).
- F6: drop the REDIS_DATA_STREAM import alias.
- F7: finetune_frac -> finetune_fraction (match sibling locals).
- F8: dedup the torchrun cmd tail across multi/single-node branches.
- F9: inline single-use broadcast_port local.
- Fix inverted loss_masking_spans docstring (spans are masked out, label == -100).
- Downgrade poll_lag lag-check log from info to debug (fired every 0.5s).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
jlamypoirier
added a commit
that referenced
this pull request
Jul 16, 2026
The core-integration PR (#153) removed the version-tagging producer because its consumers — the async_llm/llm logprob parsers that strip `:v<version>` via parse_token_id_and_version — are only present here. Restore the producer so this PR carries the feature (producer + consumers) as a coherent whole. Keeps the vllm1.py fine-fixes propagated from #153 (load_weights logging, create_engine docstring, dead imports, param_name unpack). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Extract a shared read_fast_llm_events generator in state.py, consumed by both the trainer-state listener and the vLLM main-process monitor (single source of truth for the event wire format; transient read errors now retried in both consumers). - Remove the unused WorkerExtension/EngineManager close_communicator methods. - Collapse an unreachable branch in the process-monitor teardown loop. - Fix the broadcast-protocol docstring (broadcast_object / meta is None) and trim the create_engine docstring; drop restating comments. - Use world_map.dns_address_map directly instead of a getattr fallback so a broken invariant fails loudly. - Rename socket var s -> sock; note broadcast_store must stay bound to keep the TCPStore server socket alive. - Drop the never-overridden maxlen parameter from RedisSharedStreamWriter. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…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>
- 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>
_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>
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.
Adds Fast-LLM as an alternative trainer to PipelineRL, selected by
use_fast_llm. This is the core integration — the minimum needed to run the Fast-LLM trainer end-to-end. The DeepSpeed / HTTP path is unchanged when the flag is off.Second PR in a stack extracted from the
fast-llmbranch (#140), to make review tractable:main← #151 (trainer-agnostic fixes) ← this PR (core integration) ← #140 (optional: staleness metrics, experiment configs, tests). Intended to merge together.Description prepared with Claude Opus 4.8 (Claude Code).
What's here
launch.py— splitrun_finetuneinto a DeepSpeed and a Fast-LLM branch; launchfast-llm trainvia torchrun, build its config from the experiment config, pre-create the NCCL-broadcastTCPStoreon rank 0, and threaduse_fast_llm/weight_broadcastthrough process supervision.vllm1.py—EngineManagerowns the vLLM engine lifecycle and a Fast-LLM weight-update path: it joins the trainer's persistent NCCL broadcast group, receives weights via a main-process Redis monitor thread, stamps the active model version onto generated tokens, and tears the group down ontraining_finished. The HTTP weight-update endpoint is registered only in HTTP mode.streams.py—RedisSharedStreamWriterfor many-producer fan-in to a single Redis stream (orjson payloads), consumed by the preprocessor.preprocess.py— convert preprocessed samples to Fast-LLM's streaming document format and write them to the shared stream.state.py— listen to Fast-LLM's Redis event stream (weights_ready/training_finished) instead of the legacy trainer messages.actor.py— stop the actor loop on the explicittraining_finishedevent under Fast-LLM.conf/base.yaml,conf/math.yaml—use_fast_llm/weight_broadcastswitches and thefast_llmtrainer config scaffold;math.yamlenables the path.README.md— Fast-LLM trainer path, multi-node requirements, and install docs.Notes
Fast-LLMbranch name.🤖 Generated with Claude Code