Add fp32 LM head, docs_per_step accumulation, and document-count metrics#520
Open
jlamypoirier wants to merge 16 commits into
Open
Add fp32 LM head, docs_per_step accumulation, and document-count metrics#520jlamypoirier wants to merge 16 commits into
jlamypoirier wants to merge 16 commits into
Conversation
jlamypoirier
changed the title
When True, upcasts the LM head linear's input and weight to FP32 before the matmul, matching vLLM's bf16_last_layer_fp32 quantization. This lets the trainer compute log-probabilities at the same numerical precision as the actor's sampling, so the importance-sampling ratio starts near 1.0 instead of being inflated by trainer/actor precision mismatch. The detached FP32 weight has requires_grad=False, which makes output_parallel_linear_backward skip the weight-grad path. The FSDP gradient contract is restored by computing grad_weight explicitly and accumulating into the original BF16 param's grad_buffer via accumulate_gradient. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
A schedule config field that replaces the static microbatch count with a
runtime document-count target. Matches DeepSpeed's
gradient_accumulation_passes semantics for RL: each microbatch holds one
rollout and the step boundary is set by total rollouts rather than a
fixed microbatch count.
- ScheduleConfig.docs_per_step — when >0, Trainer._prefetch_to_doc_target
fetches microbatches one at a time, all-reduces the per-microbatch doc
count, and stops once the global total reaches the target. The final
step total is broadcast to every microbatch so the loss normalization
stays consistent.
- Trainer._get_or_build_schedule(N) builds and caches a per-N Schedule
with _depth_first_override = N // breadth_first_micro_batches, reusing
the schedule machinery without touching the runner.
- Schedule._eff_{depth_first,sequential_micro_batches,num_inputs} expose
the effective values under an override.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
jlamypoirier
force-pushed
the
jlp_rl_features
branch
from
debc3dd to
5afa8c7
Compare
jlamypoirier
changed the title
Surface the per-step document count produced by `_prefetch_to_doc_target` (the loss-normalization denominator) and the cumulative document total as training metrics. Lets the dynamic `docs_per_step` accumulation be verified in production and gives documents-seen as a cross-run x-axis. Gated on `docs_per_step > 0`; no effect on the static-schedule path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 task
jlamypoirier
changed the title
Padding tokens fall past the last real document, so searchsorted assigned them a phantom (num_documents+1)-th index, one past the per-segment buffer sized by num_documents_in_sequence -> CUDA device-side assert in the GSPO index_add_. Clamp the 1-based document index onto the last real document; padding targets are masked so the contribution is zero. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Resolved two conflicts where PR #520 and main (#559) both added document-count tracking with different designs: - fast_llm/data/document/token.py: kept PR 520's GSPO padded-sequence index clamp, retargeted onto main's consolidated `num_documents` local (dropped the duplicate `num_documents_in_sequence`). - fast_llm/engine/training/trainer.py: kept PR 520's `docs_per_step` prefetch branch but adopted main's 4-value `run_step` contract, `documents_seen` param, and checkpointed `self._documents_seen`; dropped the now-dead local `total_documents_seen` accumulator. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Add a single source of truth for the reinforcement-learning trainer setup: a recipe page (docs/recipes/reinforcement-learning.md) covering the GSPO/GRPO loss config, the reference Qwen2.5-0.5B/7B recipe, precision matching, the DeepSpeed parity baseline, config gotchas, and open items; plus two override fragments (examples/qwen_gspo_0.5b.yaml, examples/qwen_gspo_7b.yaml) with the RL-relevant fast_llm.* keys. Register the page in the mkdocs nav. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ment #555 (merged) lifted the `metrics` field to the shared policy-gradient base, so GSPO now carries it. Correct the recipe (it no longer "rejects metrics" / is no longer a pending change), document the `none`/`basic`/`with_entropy` levels and the pipeline_parallel==1 constraint, and add a "Measuring lag and instability" section listing the staleness/reward/clip/ratio/KL/advantage/entropy metrics (gated behind metrics != none) plus the streaming-trainer weight-sync/documents metrics. Reframe open items around the newer-sampler re-run. Enable `metrics: basic` in both example configs. Reference sections by name (emoji headings slug differently under markdownlint vs mkdocs, so avoid fragments). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…tage
zero_stage was carrying PipelineRL's default (0.5B) / a conservative memory
choice (7B). The 0.5B trainer is a single GPU (FSDP size 1) so the stage is a
no-op — use 1. On the 4-GPU 7B trainer ZeRO-2 shards optimizer state + gradients
and avoids ZeRO-3's per-step parameter all-gather; 3 only if it OOMs. Add the
DeepSpeed-side mapping (deepspeed_config: deepspeed_stage{1,2}_bf16) and note the
ZeRO stage is math-equivalent, so it doesn't affect the FL-vs-DS comparison.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… config The runnable config is a single committed hydra config on the orchestrator side; the examples/*.yaml are a readable excerpt of the Fast-LLM knobs, not the source and not standalone-runnable. Say so explicitly to avoid drift/confusion. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ot one file Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The 7B trainer runs MLP recompute, so it sits near its memory limit and ZeRO-2's replicated parameters may not fit. Default the 7B example and the recipe doc to zero_stage 3 / deepspeed_stage3_bf16, matching the committed experiment config and the original runs. Stage 2 stays documented as the drop-down option when there is headroom. 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.
Summary
Training-side changes for RL fine-tuning, consolidated into a single PR to
main(previously stacked PRs #526 + #520; #526 is closed in favor of this one).
fp32_lm_head— FP32 LM head logitsNew
fp32_lm_headflag onLanguageModelHeadConfig(defaultFalse). When enabled, the LM headupcasts input and weight to FP32 for the logits projection and casts back, matching vLLM's
bf16_last_layer_fp32, so the trainer computes log-probabilities at the same precision the actorsampled with. Includes the gradient-flow fix for the detached FP32 weight copy — gradients are
accumulated back into the BF16 parameter's buffer.
Dynamic
docs_per_stepaccumulationNew
ScheduleConfig.docs_per_stepfield. When>0, each step accumulates microbatches one at atime, all-reduces the per-microbatch document count, and stops once the global total reaches the
target, instead of using a fixed microbatch count. The final step total is broadcast to every
microbatch so the loss-normalization denominator stays consistent. Off by default
(
docs_per_step=0keeps the original static-schedule path).num_documents/documents_seenmetricsLogs the per-step document count (the divisor
docs_per_stepproduces) and the cumulative documenttotal as training metrics — lets the dynamic accumulation be verified, and gives documents-seen as a
cross-run x-axis. Gated on
docs_per_step>0; no effect on the static path.GSPO segment-index fix for padded sequences
Clamp
global_document_index_qtonum_documents_in_sequenceinfast_llm/data/document/token.py.Padding tokens fall past the last real document, so
searchsortedassigned them a phantomout-of-range segment index, causing a CUDA device-side assert in the GSPO
index_add_. Paddingtargets are masked, so clamping them onto the last real document contributes zero.
Test plan
pytest tests/layers/test_docs_per_step.pySplit from #502.