From e26ba9f9ef2a584ae9897fc98204de97622953a3 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 27 May 2026 15:33:26 -0400 Subject: [PATCH 01/11] Add fp32_lm_head flag for vLLM precision parity 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) --- fast_llm/layers/language_model/config.py | 7 +++++ fast_llm/layers/language_model/head.py | 34 +++++++++++++++++++----- 2 files changed, 35 insertions(+), 6 deletions(-) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index bde33f297..6a0bfcfd6 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -131,6 +131,13 @@ class LanguageModelHeadConfig(BlockConfig): hint=FieldHint.architecture, valid=skip_valid_if_none(check_field(Assert.gt, 0)), ) + fp32_lm_head: bool = Field( + default=False, + desc="Upcast input and weight to float32 before the lm_head linear. " + "Matches vLLM's bf16_last_layer_fp32 quantization so new_logprobs and old_logprobs " + "are computed at the same numerical precision, keeping the IS ratio near 1 at init.", + hint=FieldHint.feature, + ) prediction_heads: int = Field( default=1, desc="Prediction heads.", diff --git a/fast_llm/layers/language_model/head.py b/fast_llm/layers/language_model/head.py index 22c750082..eb67cd553 100644 --- a/fast_llm/layers/language_model/head.py +++ b/fast_llm/layers/language_model/head.py @@ -22,7 +22,7 @@ ) from fast_llm.layers.language_model.loss.config import LanguageModelLabelEntropyLossConfig from fast_llm.layers.language_model.loss.loss import LanguageModelLoss -from fast_llm.tensor import TensorMeta +from fast_llm.tensor import TensorMeta, accumulate_gradient from fast_llm.utils import Assert, safe_merge_dicts logger = logging.getLogger(__name__) @@ -252,9 +252,17 @@ def _logits_loss_forward_backward_partial( split_index: int = 0, return_logits: bool = False, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if self._config.fp32_lm_head: + input_dtype = input_.dtype + input_ = input_.to(torch.float32) + # detach → requires_grad=False → output_parallel_linear_backward skips weight grad + weight = self.output_weights.detach().to(torch.float32) + else: + weight = self.output_weights + logits, context = output_parallel_linear_forward( input_=input_, - weight=self.output_weights, + weight=weight, bias=None, group=self._parallel_dim.group if self._vocab_parallel else None, sequence_parallel=self._sequence_parallel and self._vocab_parallel, @@ -285,12 +293,26 @@ def _logits_loss_forward_backward_partial( if loss_value is not None: losses_.append(loss_value.detach()) - if grad is not None and self._config.final_logit_softcap is not None: + if not self.training or grad is None: + return sum(losses_) if losses_ else None, None + + if self._config.final_logit_softcap is not None: grad = _softcap_backward(grad, logits, self._config.final_logit_softcap) - return sum(losses_) if losses_ else None, ( - output_parallel_linear_backward(grad, context) if self.training else None - ) + input_grad = output_parallel_linear_backward(grad, context) + if self._config.fp32_lm_head: + # Weight grad was skipped because weight.requires_grad=False; accumulate manually. + # context: (input_, weight, bias, group, sequence_parallel, ...) + saved_input = context[0] + if context[4]: # sequence_parallel + from fast_llm.core.ops import gather_op + + saved_input = gather_op(saved_input, context[3], dim=0) + grad_weight = grad.flatten(0, -2).t().mm(saved_input.flatten(0, -2)) + accumulate_gradient(self.output_weights, grad_weight.to(self.output_weights.dtype)) + input_grad = input_grad.to(input_dtype) + + return sum(losses_) if losses_ else None, input_grad def get_loss_definitions(self) -> list[LossDef]: return [ From 5afa8c7078cc121e0f19b29cd53603ec0ab8ce91 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 27 May 2026 15:37:19 -0400 Subject: [PATCH 02/11] Add docs_per_step for dynamic microbatch accumulation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- fast_llm/data/document/language_model.py | 3 +- fast_llm/engine/schedule/config.py | 10 ++ fast_llm/engine/schedule/runner.py | 5 +- fast_llm/engine/schedule/schedule.py | 38 +++-- fast_llm/engine/training/trainer.py | 61 ++++++- tests/layers/test_docs_per_step.py | 204 +++++++++++++++++++++++ 6 files changed, 297 insertions(+), 24 deletions(-) create mode 100644 tests/layers/test_docs_per_step.py diff --git a/fast_llm/data/document/language_model.py b/fast_llm/data/document/language_model.py index 16114cb80..96ab2b7b9 100644 --- a/fast_llm/data/document/language_model.py +++ b/fast_llm/data/document/language_model.py @@ -207,7 +207,7 @@ def _set_target_inputs( model_input.targets.append(target_input) - def _get_label_counts(self, mask: torch.Tensor): + def _get_label_counts(self, mask: torch.Tensor) -> torch.Tensor: # Count the number of non-masked labels in each document through cumulative sums. mask_cumsum = torch.cat([mask.new_zeros(1), mask.cumsum(0)]) length_cumsum = torch.tensor([0] + self.lengths, device=self.device).cumsum(0) @@ -215,7 +215,6 @@ def _get_label_counts(self, mask: torch.Tensor): labels_per_document = label_count_cumsum[1:] - label_count_cumsum[:-1] # Expand to one entry per token: find each token's document index via the sorted # length cumsum, then look up that document's label count. - # TODO: Document index already computed in `LengthModelInputPreprocessor`. document_index = torch.searchsorted( length_cumsum[1:], torch.arange(len(mask), device=self.device), side="right" ) diff --git a/fast_llm/engine/schedule/config.py b/fast_llm/engine/schedule/config.py index 29720b90b..2920c1334 100644 --- a/fast_llm/engine/schedule/config.py +++ b/fast_llm/engine/schedule/config.py @@ -21,6 +21,16 @@ class ScheduleConfig(Config): hint=FieldHint.core, valid=check_field(Assert.gt, 0), ) + docs_per_step: int = Field( + default=0, + desc="Target number of documents (rollouts) per optimizer step, globally across all data-parallel ranks. " + "When >0, each training step dynamically accumulates microbatches until the globally all-reduced " + "document count reaches this value, then triggers the optimizer step. " + "depth_first_micro_batches is ignored when this is set. " + "0 = use depth_first_micro_batches as-is (fixed microbatch count per step).", + hint=FieldHint.feature, + valid=check_field(Assert.geq, 0), + ) breadth_first_micro_batches: int = Field( default=1, desc="Number of micro-batches processed breadth-first, i.e., interleaved across model stages.", diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py index b2e212946..128b95e8e 100644 --- a/fast_llm/engine/schedule/runner.py +++ b/fast_llm/engine/schedule/runner.py @@ -320,7 +320,8 @@ def _preprocess_data( if context.schedule.phase.is_training else None ) - model_inputs = [next(data_iterator) for _ in range(self._config.sequential_micro_batches)] + n_micro_batches = context.schedule._eff_sequential_micro_batches + model_inputs = [next(data_iterator) for _ in range(n_micro_batches)] model_inputs[0][0].share_batch_data( [model_input for model_inputs_ in model_inputs for model_input in model_inputs_], self._distributed ) @@ -336,7 +337,7 @@ def _preprocess_data( extra_kwargs={ "grad_output": grad_output, "micro_batch": micro_batch, - "num_micro_batches": self._config.sequential_micro_batches, + "num_micro_batches": n_micro_batches, "micro_batch_splits": self._config.micro_batch_splits, }, ) diff --git a/fast_llm/engine/schedule/schedule.py b/fast_llm/engine/schedule/schedule.py index 6f7bf1d95..845b5df82 100644 --- a/fast_llm/engine/schedule/schedule.py +++ b/fast_llm/engine/schedule/schedule.py @@ -115,15 +115,17 @@ def __init__( batch_meta: list[ModelInput], distributed_config: DistributedConfig, phase: PhaseType, + _depth_first_override: int | None = None, ): super().__init__(config) + self._depth_first_override = _depth_first_override self._multi_stage = multi_stage self._distributed_config = distributed_config self._num_stages = len(self._multi_stage.stages) self._phase = phase self._is_training = self._phase.is_training - if self._config.num_inputs < self._distributed_config.pipeline_parallel: + if self._eff_num_inputs < self._distributed_config.pipeline_parallel: warnings.warn("Not enough input to achieve true pipeline parallelism.") # Setup the activation metas. @@ -155,9 +157,25 @@ def __init__( def phase(self) -> PhaseType: return self._phase + @property + def _eff_depth_first(self) -> int: + return ( + self._depth_first_override + if self._depth_first_override is not None + else self._config.depth_first_micro_batches + ) + + @property + def _eff_sequential_micro_batches(self) -> int: + return self._eff_depth_first * self._config.breadth_first_micro_batches + + @property + def _eff_num_inputs(self) -> int: + return self._eff_sequential_micro_batches * self._config.micro_batch_splits + @property def samples_per_batch(self) -> int: - return self._config.sequential_micro_batches * self._distributed_config.batch_data_parallel + return self._eff_sequential_micro_batches * self._distributed_config.batch_data_parallel def iterate(self, pipeline_rank: int | None = None) -> typing.Iterator[Step]: return iter(self._steps if pipeline_rank is None else self._device_steps[pipeline_rank]) @@ -189,7 +207,7 @@ def _create_index(self) -> None: Assert.in_range( step.index, 0, - self._config.num_inputs, + self._eff_num_inputs, ) Assert.incl(step.type_, (StepType.forward, StepType.backward)) step.global_index = i @@ -205,7 +223,7 @@ def _create_index(self) -> None: Assert.custom(all, self._device_steps) # Consistency checks step_map = self._step_map.copy() - for data_index in range(self._config.num_inputs): + for data_index in range(self._eff_num_inputs): for type_ in (StepType.forward, StepType.backward): for stage in range(0 if type_ == StepType.forward else self._first_grad_stage, self._num_stages): assert ( @@ -470,14 +488,11 @@ def _create_steps(self) -> tuple[list[Step], int]: first_grad_stage += 1 else: first_grad_stage = self._num_stages - for depth_first_micro_batch in range(self._config.depth_first_micro_batches): + for depth_first_micro_batch in range(self._eff_depth_first): for stage in range(self._num_stages): for breadth_first_micro_batch in range(self._config.breadth_first_micro_batches): for micro_batch_split in range(self._config.micro_batch_splits): - micro_batch = ( - breadth_first_micro_batch * self._config.depth_first_micro_batches - + depth_first_micro_batch - ) + micro_batch = breadth_first_micro_batch * self._eff_depth_first + depth_first_micro_batch steps.append( Step( stage=stage, @@ -492,10 +507,7 @@ def _create_steps(self) -> tuple[list[Step], int]: for stage in reversed(range(first_grad_stage, self._num_stages)): for breadth_first_micro_batch in range(self._config.breadth_first_micro_batches): for micro_batch_split in reversed(range(self._config.micro_batch_splits)): - micro_batch = ( - breadth_first_micro_batch * self._config.depth_first_micro_batches - + depth_first_micro_batch - ) + micro_batch = breadth_first_micro_batch * self._eff_depth_first + depth_first_micro_batch steps.append( Step( stage=stage, diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py index 1ed18c449..77a88377e 100644 --- a/fast_llm/engine/training/trainer.py +++ b/fast_llm/engine/training/trainer.py @@ -115,10 +115,12 @@ def setup(self, distributed: Distributed, run: Run) -> None: preprocessing_config = self._multi_stage.get_preprocessing_config( PhaseType.training, self._config.schedule.micro_batch_splits ) + self._single_mb_meta = preprocessing_config.get_input_meta(self._data.config.micro_batch_size) + self._schedule_cache: dict[int, Schedule] = {} self._schedule = Schedule( config=self._config.schedule, multi_stage=self._multi_stage, - batch_meta=preprocessing_config.get_input_meta(self._data.config.micro_batch_size), + batch_meta=self._single_mb_meta, distributed_config=self._config.model.distributed, phase=PhaseType.training, ) @@ -140,6 +142,41 @@ def setup(self, distributed: Distributed, run: Run) -> None: self._is_setup = True + def _get_or_build_schedule(self, n_microbatches: int) -> Schedule: + if n_microbatches not in self._schedule_cache: + bfmb = self._config.schedule.breadth_first_micro_batches + depth_first = n_microbatches // bfmb + self._schedule_cache[n_microbatches] = Schedule( + config=self._config.schedule, + multi_stage=self._multi_stage, + batch_meta=self._single_mb_meta, + distributed_config=self._config.model.distributed, + phase=PhaseType.training, + _depth_first_override=depth_first, + ) + return self._schedule_cache[n_microbatches] + + def _prefetch_to_doc_target(self, data_iterator) -> list: + target = self._config.schedule.docs_per_step + bfmb = self._config.schedule.breadth_first_micro_batches + buffer = [] + total_docs = 0 + while total_docs < target: + mb = next(data_iterator) + mb[0].share_batch_data(mb, self._distributed) + total_docs += mb[0].num_documents_in_batch + buffer.append(mb) + Assert.eq( + len(buffer) % bfmb, + 0, + msg=f"Fetched {len(buffer)} microbatches not divisible by breadth_first_micro_batches={bfmb}", + ) + # Reset num_documents_in_batch to the step total on all microbatches + for mb in buffer: + for mi in mb: + mi.num_documents_in_batch = total_docs + return buffer + @abc.abstractmethod def _get_data(self) -> Data: pass @@ -220,12 +257,22 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: # TODO: Data loader hates getting all micro-batches at once. # (Also preprocessing adds overhead) - reduced_losses, update_successful, train_metrics = self._runner.run_step( - train_iterator, - self._schedule, - iteration=self._completed_steps, - return_metrics=is_logging, - ) + if self._config.schedule.docs_per_step > 0: + buffer = self._prefetch_to_doc_target(train_iterator) + step_schedule = self._get_or_build_schedule(len(buffer)) + reduced_losses, update_successful, train_metrics = self._runner.run_step( + iter(buffer), + step_schedule, + iteration=self._completed_steps, + return_metrics=is_logging, + ) + else: + reduced_losses, update_successful, train_metrics = self._runner.run_step( + train_iterator, + self._schedule, + iteration=self._completed_steps, + return_metrics=is_logging, + ) # Advanced, skipped, and Nan iterations. if update_successful: diff --git a/tests/layers/test_docs_per_step.py b/tests/layers/test_docs_per_step.py new file mode 100644 index 000000000..b288a934f --- /dev/null +++ b/tests/layers/test_docs_per_step.py @@ -0,0 +1,204 @@ +""" +Unit tests for docs_per_step. + +Covers: + 1. Divisor scaling in fused_grpo_loss_forward_backward + 2. Schedule._eff_depth_first / _eff_sequential_micro_batches / _eff_num_inputs properties + 3. Trainer._prefetch_to_doc_target accumulation logic +""" + +import dataclasses +import types + +import pytest +import torch + +from fast_llm.engine.schedule.config import ScheduleConfig +from fast_llm.engine.schedule.schedule import Schedule +from fast_llm.layers.language_model.loss.policy_gradient import fused_grpo_loss_forward_backward + +device = "cuda" if torch.cuda.is_available() else "cpu" +_atol = 1e-4 if device == "cuda" else 1e-5 + + +# --------------------------------------------------------------------------- +# 1. Divisor-scaling correctness in raw kernels +# --------------------------------------------------------------------------- + + +def test_grpo_divisor_scales_loss(): + """Halving the divisor should double the loss.""" + torch.manual_seed(10) + n_tok, vocab = 16, 32 + logits = torch.randn(n_tok, vocab, device=device) + target = torch.randint(0, vocab, (n_tok,), device=device) + advantages = torch.randn(n_tok, device=device) + old_lp = torch.randn(n_tok, device=device) - 2.0 + + d1 = float(n_tok) + d2 = float(n_tok) * 2 + + loss1, _, _ = fused_grpo_loss_forward_backward(logits, target, advantages, old_lp, divisor=d1) + loss2, _, _ = fused_grpo_loss_forward_backward(logits, target, advantages, old_lp, divisor=d2) + + assert ( + abs(loss1.item() - 2.0 * loss2.item()) < _atol * 10 + ), f"Expected loss(d1) ≈ 2*loss(d2), got {loss1.item():.6f} vs {2*loss2.item():.6f}" + + +# --------------------------------------------------------------------------- +# 2. Schedule._eff_* properties +# --------------------------------------------------------------------------- + + +def _make_bare_schedule(depth_first: int, breadth_first: int, splits: int, override: int | None) -> Schedule: + """Create a Schedule with __init__ bypassed to test the _eff_* properties only.""" + config = ScheduleConfig( + depth_first_micro_batches=depth_first, + breadth_first_micro_batches=breadth_first, + micro_batch_splits=splits, + ) + sched = object.__new__(Schedule) + # Minimal attributes used by the three _eff_* properties. + object.__setattr__(sched, "_config", config) + object.__setattr__(sched, "_depth_first_override", override) + # samples_per_batch also needs _distributed_config.batch_data_parallel + fake_distributed = types.SimpleNamespace(batch_data_parallel=1) + object.__setattr__(sched, "_distributed_config", fake_distributed) + return sched + + +def test_schedule_eff_properties_no_override(): + sched = _make_bare_schedule(depth_first=4, breadth_first=2, splits=3, override=None) + assert sched._eff_depth_first == 4 + assert sched._eff_sequential_micro_batches == 8 # 4 * 2 + assert sched._eff_num_inputs == 24 # 8 * 3 + assert sched.samples_per_batch == 8 # 8 * dp=1 + + +def test_schedule_eff_properties_with_override(): + sched = _make_bare_schedule(depth_first=4, breadth_first=2, splits=3, override=7) + assert sched._eff_depth_first == 7 # override wins + assert sched._eff_sequential_micro_batches == 14 # 7 * 2 + assert sched._eff_num_inputs == 42 # 14 * 3 + assert sched.samples_per_batch == 14 # 14 * dp=1 + + +def test_schedule_eff_properties_override_equals_config(): + """Override equal to config value → same result as no override.""" + sched_no = _make_bare_schedule(depth_first=3, breadth_first=2, splits=1, override=None) + sched_yes = _make_bare_schedule(depth_first=3, breadth_first=2, splits=1, override=3) + assert sched_no._eff_depth_first == sched_yes._eff_depth_first + assert sched_no._eff_sequential_micro_batches == sched_yes._eff_sequential_micro_batches + assert sched_no._eff_num_inputs == sched_yes._eff_num_inputs + + +def test_schedule_samples_per_batch_uses_eff(): + """samples_per_batch should scale with _eff_sequential, not config.sequential.""" + sched = _make_bare_schedule(depth_first=2, breadth_first=2, splits=1, override=5) + # Config says depth_first=2 → sequential=4; override=5 → eff_sequential=10 + assert sched._eff_sequential_micro_batches == 10 + assert sched.samples_per_batch == 10 # dp=1 + + +# --------------------------------------------------------------------------- +# 3. _prefetch_to_doc_target accumulation logic +# --------------------------------------------------------------------------- + + +@dataclasses.dataclass +class _FakeMicrobatch: + """Stub for a single split of one microbatch.""" + + num_documents: int + num_documents_in_batch: int | None = None + + @classmethod + def share_batch_data(cls, inputs, distributed): + """Mimic TokenModelInput.share_batch_data with group=None (single process).""" + if inputs[0].num_documents_in_batch is None: + total = sum(inp.num_documents for inp in inputs) + for inp in inputs: + inp.num_documents_in_batch = total + + +def _fake_iterator(doc_counts: list[int]): + """Yield [_FakeMicrobatch(n)] for each n in doc_counts.""" + for n in doc_counts: + yield [_FakeMicrobatch(num_documents=n)] + + +class _StubTrainer: + """Concrete stub that exposes only the interface _prefetch_to_doc_target needs.""" + + # Borrow the method directly so it runs against this stub's attributes. + from fast_llm.engine.training.trainer import Trainer as _Trainer + + _prefetch_to_doc_target = _Trainer._prefetch_to_doc_target + + +def _make_fake_trainer(docs_per_step: int, bfmb: int = 1): + """Create a _StubTrainer with the attributes _prefetch_to_doc_target reads.""" + schedule_cfg = types.SimpleNamespace( + docs_per_step=docs_per_step, + breadth_first_micro_batches=bfmb, + ) + config = types.SimpleNamespace(schedule=schedule_cfg) + distributed = types.SimpleNamespace(batch_data_group=None) + + trainer = _StubTrainer() + trainer._config = config + trainer._distributed = distributed + return trainer + + +def test_prefetch_stops_at_target(): + """Buffer should stop growing once cumulative docs ≥ docs_per_step.""" + trainer = _make_fake_trainer(docs_per_step=6, bfmb=1) + # Each microbatch has 2 docs; need ≥6 → expect 3 microbatches + it = _fake_iterator([2, 2, 2, 2, 2]) + buffer = trainer._prefetch_to_doc_target(it) + + assert len(buffer) == 3, f"Expected 3 microbatches, got {len(buffer)}" + + +def test_prefetch_resets_num_documents_in_batch(): + """After the call, every microbatch input has num_documents_in_batch = step total.""" + trainer = _make_fake_trainer(docs_per_step=5, bfmb=1) + # 3 docs, 3 docs → total=6 (overshoots 5, stops after 2nd) + it = _fake_iterator([3, 3, 3]) + buffer = trainer._prefetch_to_doc_target(it) + + step_total = sum(mb[0].num_documents for mb in buffer) + for mb in buffer: + for mi in mb: + assert ( + mi.num_documents_in_batch == step_total + ), f"Expected num_documents_in_batch={step_total}, got {mi.num_documents_in_batch}" + + +def test_prefetch_overshoot_is_included(): + """A microbatch that pushes the total over the target IS included (not dropped).""" + trainer = _make_fake_trainer(docs_per_step=5, bfmb=1) + it = _fake_iterator([4, 4]) # 4 < 5, then 8 ≥ 5 → 2 microbatches + buffer = trainer._prefetch_to_doc_target(it) + assert len(buffer) == 2 + assert buffer[-1][0].num_documents_in_batch == 8 # step total = 4+4 + + +def test_prefetch_divisibility_check(): + """Raises when fetched count is not divisible by breadth_first_micro_batches.""" + trainer = _make_fake_trainer(docs_per_step=4, bfmb=2) + # Each microbatch has 5 docs → only 1 mb needed, but 1 % 2 != 0 + it = _fake_iterator([5, 5, 5]) + with pytest.raises(Exception): + trainer._prefetch_to_doc_target(it) + + +def test_prefetch_exact_divisibility(): + """No error when fetched count is exactly divisible by breadth_first_micro_batches.""" + trainer = _make_fake_trainer(docs_per_step=4, bfmb=2) + # 2 docs each → need ≥4 → fetch 2 microbatches → 2 % 2 == 0 + it = _fake_iterator([2, 2, 2, 2]) + buffer = trainer._prefetch_to_doc_target(it) + assert len(buffer) == 2 From cc8768e11265c8b7086ab9894fe6cbe6b5b299f5 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Thu, 25 Jun 2026 12:40:17 -0400 Subject: [PATCH 03/11] Log num_documents and documents_seen training metrics 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 --- fast_llm/engine/training/trainer.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py index 77a88377e..069f6b44b 100644 --- a/fast_llm/engine/training/trainer.py +++ b/fast_llm/engine/training/trainer.py @@ -221,6 +221,7 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: skipped_iters = 0 nan_iters = 0 total_losses = {loss_def.name: 0.0 for loss_def in self._loss_definitions} + total_documents_seen = 0 # Profiling profiler = self._config.profiling.get_profiler( @@ -259,6 +260,9 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: # (Also preprocessing adds overhead) if self._config.schedule.docs_per_step > 0: buffer = self._prefetch_to_doc_target(train_iterator) + # `_prefetch_to_doc_target` broadcasts the step document total onto every microbatch. + step_num_documents = buffer[0][0].num_documents_in_batch + total_documents_seen += step_num_documents step_schedule = self._get_or_build_schedule(len(buffer)) reduced_losses, update_successful, train_metrics = self._runner.run_step( iter(buffer), @@ -267,6 +271,7 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: return_metrics=is_logging, ) else: + step_num_documents = None reduced_losses, update_successful, train_metrics = self._runner.run_step( train_iterator, self._schedule, @@ -304,6 +309,11 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: metrics_key = PhaseType.training metrics[metrics_key] = { "batch_size": self._batch_size, + **( + {"num_documents": step_num_documents, "documents_seen": total_documents_seen} + if step_num_documents is not None + else {} + ), **{ name: (value / advanced_iters if advanced_iters > 0 else float("nan")) for name, value in total_losses.items() From b57b72919f058b66e0bcd75406d9427ea64c4abe Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 22 Jun 2026 14:47:38 -0400 Subject: [PATCH 04/11] Fix GSPO segment index out-of-bounds on padded sequences 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 --- fast_llm/data/document/token.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/fast_llm/data/document/token.py b/fast_llm/data/document/token.py index 4d1af453d..add8f268a 100644 --- a/fast_llm/data/document/token.py +++ b/fast_llm/data/document/token.py @@ -118,16 +118,19 @@ def _get_model_input(self, begin: int, end: int, config: TokenPreprocessingConfi global_cumulative_lengths = torch.from_numpy(padded_cumsum(self.lengths)).to( dtype=torch.int32, device=self.device ) + # Exclude the trailing padding "length" from the count. + num_documents_in_sequence = len(self.lengths) - (1 if self.unpadded_length < len(self.tokens) else 0) + # Padding tokens fall past the last real document, so `searchsorted` would assign them the + # phantom (num_documents + 1)-th index — one past the per-segment buffer sized by + # `num_documents_in_sequence`. Clamp them onto the last real document; their target is masked + # so the contribution is zero, and the 1-based index stays within `num_documents_in_sequence`. model_input.global_document_index_q = torch.searchsorted( global_cumulative_lengths, torch.arange(begin, end, device=self.device), side="right", out_int32=True, - ) - # Exclude the padding "length" from the count. - model_input.num_documents_in_sequence = len(self.lengths) - ( - 1 if self.unpadded_length < len(self.tokens) else 0 - ) + ).clamp_(max=num_documents_in_sequence) + model_input.num_documents_in_sequence = num_documents_in_sequence LengthModelInputPreprocessor( lengths=lengths, From 6c7c6890638a0daf6b806db6c7471a243a27396c Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 14:02:42 -0400 Subject: [PATCH 05/11] Document GSPO/GRPO RL recipe and reference Qwen configs 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 --- docs/recipes/reinforcement-learning.md | 188 +++++++++++++++++++++++++ examples/qwen_gspo_0.5b.yaml | 67 +++++++++ examples/qwen_gspo_7b.yaml | 71 ++++++++++ mkdocs.yaml | 1 + 4 files changed, 327 insertions(+) create mode 100644 docs/recipes/reinforcement-learning.md create mode 100644 examples/qwen_gspo_0.5b.yaml create mode 100644 examples/qwen_gspo_7b.yaml diff --git a/docs/recipes/reinforcement-learning.md b/docs/recipes/reinforcement-learning.md new file mode 100644 index 000000000..0011b8c3a --- /dev/null +++ b/docs/recipes/reinforcement-learning.md @@ -0,0 +1,188 @@ +--- +title: Reinforcement Learning with GSPO / GRPO +--- + +Fast-LLM can act as the **trainer** in an asynchronous reinforcement-learning loop. An +external RL orchestrator generates rollouts with [vLLM](https://github.com/vllm-project/vllm) +samplers, scores them into per-token advantages, and streams the packed rollouts to the +Fast-LLM trainer; after each optimizer step the updated weights are broadcast back to the +samplers over NCCL. Fast-LLM itself only sees packed sequences that carry, per document, +the **advantage**, the **sampling log-probabilities** (the behaviour policy `π_old`), and +the **document lengths**. The policy-gradient objective (GSPO or GRPO) is computed in the +language-model head. + +This recipe documents the training-side configuration used for a set of reference +Qwen2.5 experiments, and the reasoning behind each setting. The orchestration layer, the +rollout data pipeline, and the reward model are out of scope — this page is about how to +configure the Fast-LLM trainer. + +Two ready-to-adapt override fragments accompany this page: + +- [`examples/qwen_gspo_0.5b.yaml`](https://github.com/ServiceNow/Fast-LLM/blob/main/examples/qwen_gspo_0.5b.yaml) +- [`examples/qwen_gspo_7b.yaml`](https://github.com/ServiceNow/Fast-LLM/blob/main/examples/qwen_gspo_7b.yaml) + +They contain only the RL-relevant overrides; the base model config (architecture, vocab, +tokenizer, RoPE) is imported from the Qwen2.5 checkpoint by the orchestrator, and the +data stream is supplied at run time. + +## 🎯 The policy-gradient loss + +The RL objective is one entry in the head's `losses` dictionary. The dictionary key is a +free label; the `type` field selects the loss: + +```yaml +model: + base_model: + head: + fp32_lm_head: true + losses: + gspo: # free label; use `grpo` for GRPO + type: gspo # `gspo` or `grpo` + epsilon_low: 0.003 # lower clip on the log-prob ratio + epsilon_high: 0.004 # upper clip on the log-prob ratio + logits_scale_factor: 1.4285714 +``` + +- **`type`** — `gspo` (sequence-level importance ratio, geometric-mean per segment) or + `grpo` (token-level ratio). Both derive from `LanguageModelPolicyGradientLossConfig` + and share the fields below. +- **`epsilon_low` / `epsilon_high`** — the asymmetric clip on the importance ratio. GSPO + operates on a per-segment geometric-mean ratio, so its useful epsilon range is far + tighter than GRPO's token-level clip; the reference runs use `0.003 / 0.004`. +- **`logits_scale_factor`** — set to `1 / sampling_temperature` so the trainer's + log-probabilities are computed at the same temperature the sampler used. With sampling + temperature `0.7` this is `1.4285714`. Leaving it at the default `1.0` while sampling at + `T < 1` inflates the importance ratio and destabilizes training. +- **`metrics`** — *(GRPO only.)* Opts into diagnostic metrics (clip fraction, entropy, + ratio statistics). This field lives on the GRPO config and is **rejected by a `gspo` + loss**; do not set it on a GSPO run. See [Open items](#-open-items). + +### `fp32_lm_head` + +`fp32_lm_head: true` keeps the language-model head in fp32 regardless of the body compute +dtype. The samplers compute their logits through an fp32 path, and matching the trainer +head to it removes a source of cross-engine numerical drift. This is independent of the +body precision — keep it on for bf16, fp16, and fp32 runs alike. + +## ⚙️ Reference recipe + +The reference experiments train Qwen2.5-0.5B and Qwen2.5-7B on a math-reasoning task with +GSPO. The shared recipe: + +| Setting | Value | Notes | +| --- | --- | --- | +| Loss | GSPO, `epsilon_low=0.003`, `epsilon_high=0.004` | | +| `logits_scale_factor` | `1.4285714` | `= 1 / 0.7` | +| `fp32_lm_head` | `true` | | +| `docs_per_step` | `256` | rollouts accumulated per optimizer step | +| `depth_first_micro_batches` | `48` | | +| `micro_batch_size` | `10000` tokens | tracks the packed sequence length | +| Learning rate | `5e-7`, constant, `0` warmup | | +| `beta_1` / `beta_2` | `0.974004` / `0.999750` | sqrt-rule, `m=4` | +| `gradient_norm_clipping` | `76.8` | `= 0.3 × docs_per_step` | +| Seed | `43` | | +| Sampling temperature | `0.7` | | + +Per-size differences: + +| | Qwen2.5-0.5B | Qwen2.5-7B | +| --- | --- | --- | +| Head | tied | untied | +| `zero_stage` | `2` | `3` | +| `mlp.recompute_level` | — | `full` | +| Trainer sharding | 1 replica (FSDP size 1) | 4 replicas (FSDP size 4) | + +Each size was run in three precision arms — **bf16** (default), **fp16-matched**, and +**fp32-matched** — alongside a **DeepSpeed** parity baseline (see +[DeepSpeed parity](#-deepspeed-parity-baseline)). + +### Derived quantities + +- **`gradient_norm_clipping = 0.3 × docs_per_step`.** The per-step gradient is normalized + by the document count, so the raw gradient norm scales with `docs_per_step`. To hold the + *effective* clip fixed at `0.3`, the threshold must scale the same way: `0.3 × 256 = 76.8`. + If you change `docs_per_step`, rescale the clip. +- **`docs_per_step` should equal the number of rollouts per step.** The loss divides by the + document count seen across the accumulated micro-batches; setting `docs_per_step` to the + rollout count makes that divisor equal to the batch size, matching the reference-engine + normalization. +- **Adam betas follow the sqrt-rule for an effective batch multiplier `m`:** + `beta = beta_default ** (1/m)`. With `m = 4`, `0.9 → 0.974004` and `0.999 → 0.999750`. + +## 🔬 Precision matching + +The single most important correctness constraint: **the trainer compute dtype must match +the sampler dtype.** The trainer recomputes `π` on the rollouts to form the importance +ratio `exp(new_logprob − old_logprob)`; `old_logprob` comes from the sampler. If the two +engines run at different precisions, the ratio is systematically biased even at step 0, +which corrupts the clip decisions and collapses reward. + +- **bf16 (default).** Trainer body bf16, sampler bf16, head fp32 on both. This is the + baseline and needs no precision overrides. +- **fp16-matched / fp32-matched.** Set `model.distributed.compute_dtype` to `float16` / + `float32` **and** set the sampler dtype to the same value. Do not set only one — a naive + `compute_dtype=float16` with a bf16 sampler is the mismatch failure above (reward drops + to noise and throughput falls sharply from the diverging clip mask). + +Precision *level* is not the driver of the trainer↔sampler gap; precision *mismatch* is. +In the reference runs the matched fp16 and fp32 arms track (and slightly beat) bf16, which +is consistent with bf16's 8-bit mantissa sitting near the noise floor of the very tight +GSPO epsilon. + +## 🧮 DeepSpeed parity baseline + +For A/B comparison against a DeepSpeed ZeRO-3 trainer, match these on the DeepSpeed side: + +| Fast-LLM | DeepSpeed equivalent | +| --- | --- | +| `optimizer.learning_rate.base: 5e-7` | `learning_rate: 5e-7` | +| `learning_rate.warmup_iterations: 0` | `num_warmup_steps: 0` | +| `learning_rate.decay_style: constant` | `lr_scheduler_type: constant` | +| `schedule.docs_per_step: 256` | `gradient_accumulation_passes: 256` | +| `optimizer.beta_1/beta_2` | Adam `beta1/beta2` (same values) | +| `fp32_lm_head: true` | fp32 LM head (matched) | + +DeepSpeed runs ZeRO-3 bf16 with gradient checkpointing. The samplers, temperature, +`epsilon`, seed, and sequence shape are identical across both trainers. + +## 🧷 Configuration gotchas + +- **`micro_batch_size` is in tokens and must be ≥ the packed sequence length.** One + sequence per micro-batch, so it tracks the sequence length (10000 here) and cannot be + set smaller. Lowering trainer memory means lowering sequence length and + `micro_batch_size` together. +- **Layering over an orchestrator-generated config.** If the orchestrator already emits a + Fast-LLM config and you override via the command line, use add-or-override (`++`) for + keys it already sets (e.g. `multi_stage.zero_stage`) and add (`+`) only for keys it does + not. +- **GSPO rejects `metrics`.** The `metrics` field exists on the GRPO config only; a `gspo` + loss raises a validation error if it is set. Leave it unset for GSPO. +- **`normalize_by_documents` no longer exists.** Document normalization is applied + unconditionally; there is no config knob for it. + +## 🆕 Relevant trainer behavior + +- **Gradient-accumulation fix (single-replica trainers).** A prior bug suppressed the + decoder gradient under multi-micro-batch accumulation when the trainer used a single + FSDP shard with shared gradient buffers. This is fixed. It only affected FSDP-size-1 + trainers (the 0.5B reference run); sharded trainers (the 7B run, FSDP size 4) were never + affected. +- **`docs_per_step` accumulation.** The scheduler accumulates micro-batches until + `docs_per_step` documents are seen, then takes one optimizer step, so the loss divisor is + the full step's document count rather than a single micro-batch's. +- **Document / reward / model-version metrics.** The trainer tracks `documents_seen` + (checkpointed and offered as a W&B x-axis) and carries the reward and a per-token model + version through the batch, logging reward and version statistics. + +## 📌 Open items + +Known limitations and proposed follow-ups for the reference experiments: + +- **Diagnostic metrics for GSPO.** The `metrics` field is GRPO-only today; exposing it on + the shared policy-gradient base would let GSPO runs log clip fraction, entropy, and ratio + statistics. This is a pending code change, not a config option. +- **7B rollout-staleness instability.** The 7B runs show reward spikes traced to rollout + staleness and an abort storm when in-flight generations are cancelled on weight sync — + *not* to numerical noise (fp16 and bf16 spike at the same rate). Proposed mitigations + (not yet adopted): bound rollout lag, raise the weight-update interval, avoid aborting + in-flight generations on sync, and double-buffer weights. diff --git a/examples/qwen_gspo_0.5b.yaml b/examples/qwen_gspo_0.5b.yaml new file mode 100644 index 000000000..7cbe4f402 --- /dev/null +++ b/examples/qwen_gspo_0.5b.yaml @@ -0,0 +1,67 @@ +# Fast-LLM training overrides for GSPO reinforcement learning on Qwen2.5-0.5B. +# +# This is NOT a standalone training config. Reinforcement learning is driven by an +# external async-RL orchestrator that generates rollouts with vLLM samplers and +# broadcasts the trainer weights back to the samplers over NCCL each step. The +# orchestrator builds the base model config by importing the Qwen2.5-0.5B checkpoint +# (architecture, vocab, RoPE, tokenizer) and supplies the rollout data stream +# (packed sequences carrying advantages, sampling log-probabilities and per-document +# lengths). The keys below are the RL-relevant overrides layered on top of that base +# config; only these differ from a standard fine-tune. See +# `docs/recipes/reinforcement-learning.md` for the full recipe and rationale. +# +# Reference run: Qwen2.5-0.5B, seed 43, single trainer replica (FSDP size 1). + +model: + distributed: + # bfloat16 is the default. For a fp16 or fp32 run, set this AND the sampler dtype + # to the same value — a trainer/sampler precision mismatch corrupts the importance + # ratios and collapses reward. See the "Precision matching" section in the recipe. + compute_dtype: bfloat16 + sequence_data_parallel: 1 + timeout: 3600 + multi_stage: + zero_stage: 2 + base_model: + head: + # Keep the LM head in fp32 to match the sampler's fp32-logit path (cross-engine + # numerical matching, independent of the body compute dtype). + fp32_lm_head: true + losses: + # The dict key is a free label; `type` selects the loss. Use `grpo` for GRPO. + gspo: + type: gspo + epsilon_low: 0.003 + epsilon_high: 0.004 + # = 1 / sampling_temperature (temperature 0.7). Rescales logits so the + # trainer's log-probabilities match the sampler's. + logits_scale_factor: 1.4285714 + +data: + # Token budget per micro-batch. Must be >= the longest packed sequence, so it tracks + # the orchestrator's sequence length (10000 here); it cannot be smaller. + micro_batch_size: 10000 + +schedule: + depth_first_micro_batches: 48 + # Accumulate gradients until this many documents (rollouts) are seen, then step. + # Set equal to the number of rollouts per step so the loss divisor matches the batch. + docs_per_step: 256 + +optimizer: + learning_rate: + base: 5.0e-7 + warmup_iterations: 0 + decay_style: constant + decay_iterations: 400 + # Adjusted for the effective batch multiplier (sqrt-rule, m=4): beta = beta_default**(1/m). + beta_1: 0.974004 + beta_2: 0.999750 + # 0.3 * docs_per_step. The per-step gradient is divided by the document count, so the + # clip threshold scales with docs_per_step to keep the effective clip fixed at 0.3. + gradient_norm_clipping: 76.8 + +training: + train_iters: 2000 + checkpoint: + interval: 100 diff --git a/examples/qwen_gspo_7b.yaml b/examples/qwen_gspo_7b.yaml new file mode 100644 index 000000000..7578bf773 --- /dev/null +++ b/examples/qwen_gspo_7b.yaml @@ -0,0 +1,71 @@ +# Fast-LLM training overrides for GSPO reinforcement learning on Qwen2.5-7B. +# +# This is NOT a standalone training config. Reinforcement learning is driven by an +# external async-RL orchestrator that generates rollouts with vLLM samplers and +# broadcasts the trainer weights back to the samplers over NCCL each step. The +# orchestrator builds the base model config by importing the Qwen2.5-7B checkpoint +# (architecture, vocab, RoPE, tokenizer) and supplies the rollout data stream +# (packed sequences carrying advantages, sampling log-probabilities and per-document +# lengths). The keys below are the RL-relevant overrides layered on top of that base +# config. See `docs/recipes/reinforcement-learning.md` for the full recipe and rationale. +# +# Reference run: Qwen2.5-7B (untied head), seed 43, trainer sharded across 4 replicas +# (ZeRO-3, FSDP size 4). Same recipe as the 0.5B config; the differences are the +# ZeRO stage and MLP recompute, both to fit the 7B trainer in memory. + +model: + distributed: + # bfloat16 is the default. For a fp16 or fp32 run, set this AND the sampler dtype + # to the same value — a trainer/sampler precision mismatch corrupts the importance + # ratios and collapses reward. See the "Precision matching" section in the recipe. + compute_dtype: bfloat16 + sequence_data_parallel: 1 + timeout: 3600 + multi_stage: + # Shard weights, gradients and optimizer state across the trainer replicas. + zero_stage: 3 + base_model: + decoder: + block: + mlp: + # Recompute MLP activations in the backward pass to save trainer memory. + recompute_level: full + head: + # Qwen2.5-7B has an untied head; the sampler's fp32-logit path binds here, so + # keep the trainer head in fp32 to match it. + fp32_lm_head: true + losses: + # The dict key is a free label; `type` selects the loss. Use `grpo` for GRPO. + gspo: + type: gspo + epsilon_low: 0.003 + epsilon_high: 0.004 + # = 1 / sampling_temperature (temperature 0.7). + logits_scale_factor: 1.4285714 + +data: + # Token budget per micro-batch. Must be >= the longest packed sequence, so it tracks + # the orchestrator's sequence length (10000 here); it cannot be smaller. + micro_batch_size: 10000 + +schedule: + depth_first_micro_batches: 48 + # Accumulate gradients until this many documents (rollouts) are seen, then step. + docs_per_step: 256 + +optimizer: + learning_rate: + base: 5.0e-7 + warmup_iterations: 0 + decay_style: constant + decay_iterations: 400 + # Adjusted for the effective batch multiplier (sqrt-rule, m=4): beta = beta_default**(1/m). + beta_1: 0.974004 + beta_2: 0.999750 + # 0.3 * docs_per_step. + gradient_norm_clipping: 76.8 + +training: + train_iters: 2000 + checkpoint: + interval: 100 diff --git a/mkdocs.yaml b/mkdocs.yaml index df5d874bf..a86313691 100644 --- a/mkdocs.yaml +++ b/mkdocs.yaml @@ -179,6 +179,7 @@ nav: - Upcycle Llama 3B to MoE: recipes/upcycle-llama-3b-to-moe.md - Instruction Finetuning: recipes/instruction-finetuning.md - Generate: recipes/generate.md + - Reinforcement Learning: recipes/reinforcement-learning.md - Reference: - User Guide: - Configuration: user_guide/configuration.md From d2f4f239dbbcd6072d36535597f5543ed73127c4 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 14:58:22 -0400 Subject: [PATCH 06/11] RL recipe: GSPO metrics available (#555), add lag/instability measurement #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 --- docs/recipes/reinforcement-learning.md | 69 +++++++++++++++++++------- examples/qwen_gspo_0.5b.yaml | 4 ++ examples/qwen_gspo_7b.yaml | 4 ++ 3 files changed, 60 insertions(+), 17 deletions(-) diff --git a/docs/recipes/reinforcement-learning.md b/docs/recipes/reinforcement-learning.md index 0011b8c3a..340bb926d 100644 --- a/docs/recipes/reinforcement-learning.md +++ b/docs/recipes/reinforcement-learning.md @@ -53,9 +53,13 @@ model: log-probabilities are computed at the same temperature the sampler used. With sampling temperature `0.7` this is `1.4285714`. Leaving it at the default `1.0` while sampling at `T < 1` inflates the importance ratio and destabilizes training. -- **`metrics`** — *(GRPO only.)* Opts into diagnostic metrics (clip fraction, entropy, - ratio statistics). This field lives on the GRPO config and is **rejected by a `gspo` - loss**; do not set it on a GSPO run. See [Open items](#-open-items). +- **`metrics`** — opts into diagnostic metrics; shared by both GSPO and GRPO. `none` + (default) logs nothing extra; `basic` logs importance-ratio, KL and advantage statistics + plus the rollout reward and staleness; `with_entropy` also logs the policy entropy. The + extra statistics need a second full-vocab softmax, so `basic`/`with_entropy` require + `pipeline_parallel == 1` and add a compute cost that is skipped entirely at `none`. This + is the switch to turn on when measuring lag and instability — see + the **Measuring lag and instability** section. ### `fp32_lm_head` @@ -94,7 +98,7 @@ Per-size differences: Each size was run in three precision arms — **bf16** (default), **fp16-matched**, and **fp32-matched** — alongside a **DeepSpeed** parity baseline (see -[DeepSpeed parity](#-deepspeed-parity-baseline)). +the **DeepSpeed parity baseline** section). ### Derived quantities @@ -145,6 +149,33 @@ For A/B comparison against a DeepSpeed ZeRO-3 trainer, match these on the DeepSp DeepSpeed runs ZeRO-3 bf16 with gradient checkpointing. The samplers, temperature, `epsilon`, seed, and sequence shape are identical across both trainers. +## 📊 Measuring lag and instability + +Set `metrics: basic` on the policy-gradient loss to log the diagnostics needed to quantify +rollout staleness and training instability (`with_entropy` adds the policy entropy on top). +The statistics are gated behind `metrics != none`, so `none` logs none of them. + +Enabled by `metrics: basic` (each prefixed with the loss name, e.g. `gspo_`), logged as +mean plus max/min where applicable: + +| Metric | Meaning | +| --- | --- | +| `staleness` | `documents_seen − model_version`: documents trained since each token's generating policy was synced — the rollout **lag**. `max_staleness` is the worst-case lag. | +| `train_samples_reward` | Mean rollout reward. Averaged over the sample-filtered batch it is biased, so treat it as a diagnostic, not a policy-performance metric. | +| `clipped_ratio_fraction` | Fraction of segments (GSPO) / tokens (GRPO) whose importance ratio was clipped — the primary instability indicator. | +| `ratio_new_old`, `ratio_new_old_sum`, `ratio_new_old_squared_sum` | Importance-ratio mean and the sums to derive its variance. | +| `kl_new_old` | KL(new ‖ old) between the trainer and sampler policies. | +| `advantage`, `max_advantage`, `min_advantage` | Advantage statistics. | +| `num_tokens` | Labeled-token count, for per-token averages. | +| `entropy` | Policy entropy (only with `metrics: with_entropy`). | + +The streaming trainer additionally logs, unconditionally, `weight_sync_time_ms` (wall-clock +cost of broadcasting weights to the samplers each step) and `documents_seen` (cumulative +document count, also selectable as a W&B x-axis). + +GSPO's ratio/clip statistics are segment-level and GRPO's are token-level, so the suffixes +match where the meaning lines up but the units differ. + ## 🧷 Configuration gotchas - **`micro_batch_size` is in tokens and must be ≥ the packed sequence length.** One @@ -155,8 +186,9 @@ DeepSpeed runs ZeRO-3 bf16 with gradient checkpointing. The samplers, temperatur Fast-LLM config and you override via the command line, use add-or-override (`++`) for keys it already sets (e.g. `multi_stage.zero_stage`) and add (`+`) only for keys it does not. -- **GSPO rejects `metrics`.** The `metrics` field exists on the GRPO config only; a `gspo` - loss raises a validation error if it is set. Leave it unset for GSPO. +- **`metrics` requires `pipeline_parallel == 1`.** The diagnostic metrics compute a second + full-vocab softmax, which pipeline parallelism would split; setting `basic`/`with_entropy` + under `pipeline_parallel > 1` is rejected. The reference runs use `pipeline_parallel == 1`. - **`normalize_by_documents` no longer exists.** Document normalization is applied unconditionally; there is no config knob for it. @@ -170,19 +202,22 @@ DeepSpeed runs ZeRO-3 bf16 with gradient checkpointing. The samplers, temperatur - **`docs_per_step` accumulation.** The scheduler accumulates micro-batches until `docs_per_step` documents are seen, then takes one optimizer step, so the loss divisor is the full step's document count rather than a single micro-batch's. +- **Diagnostic metrics for GSPO and GRPO.** The `metrics` field is shared by both losses, + so GSPO runs can log importance-ratio, KL, advantage, clip-fraction and (optionally) + entropy statistics — see the **Measuring lag and instability** section. - **Document / reward / model-version metrics.** The trainer tracks `documents_seen` (checkpointed and offered as a W&B x-axis) and carries the reward and a per-token model - version through the batch, logging reward and version statistics. + version through the batch, from which the reward and rollout staleness are logged. ## 📌 Open items -Known limitations and proposed follow-ups for the reference experiments: - -- **Diagnostic metrics for GSPO.** The `metrics` field is GRPO-only today; exposing it on - the shared policy-gradient base would let GSPO runs log clip fraction, entropy, and ratio - statistics. This is a pending code change, not a config option. -- **7B rollout-staleness instability.** The 7B runs show reward spikes traced to rollout - staleness and an abort storm when in-flight generations are cancelled on weight sync — - *not* to numerical noise (fp16 and bf16 spike at the same rate). Proposed mitigations - (not yet adopted): bound rollout lag, raise the weight-update interval, avoid aborting - in-flight generations on sync, and double-buffer weights. +- **Current focus: newer sampler + measure lag.** The 7B runs show reward spikes traced to + rollout staleness and an abort storm when in-flight generations are cancelled on a weight + sync — *not* to numerical noise (fp16 and bf16 spike at the same rate). The immediate step + is to re-run on a sampler build that keeps in-flight generations across a weight update + (instead of aborting them), with `metrics: basic` enabled, and quantify the remaining + staleness and clip fraction. See + the **Measuring lag and instability** section. +- **Later, once the metrics are in hand.** Candidate mitigations — to be judged against the + measured staleness rather than adopted blind — bound the rollout lag, raise the + weight-update interval, and double-buffer weights. Deferred until the lag measurements exist. diff --git a/examples/qwen_gspo_0.5b.yaml b/examples/qwen_gspo_0.5b.yaml index 7cbe4f402..bf6140d4c 100644 --- a/examples/qwen_gspo_0.5b.yaml +++ b/examples/qwen_gspo_0.5b.yaml @@ -36,6 +36,10 @@ model: # = 1 / sampling_temperature (temperature 0.7). Rescales logits so the # trainer's log-probabilities match the sampler's. logits_scale_factor: 1.4285714 + # Log lag/instability diagnostics (staleness, reward, clip fraction, ratio + # variance, KL, advantage). `with_entropy` also logs entropy; `none` disables + # them and skips the extra softmax. Requires pipeline_parallel == 1. + metrics: basic data: # Token budget per micro-batch. Must be >= the longest packed sequence, so it tracks diff --git a/examples/qwen_gspo_7b.yaml b/examples/qwen_gspo_7b.yaml index 7578bf773..3bcc2c7c1 100644 --- a/examples/qwen_gspo_7b.yaml +++ b/examples/qwen_gspo_7b.yaml @@ -42,6 +42,10 @@ model: epsilon_high: 0.004 # = 1 / sampling_temperature (temperature 0.7). logits_scale_factor: 1.4285714 + # Log lag/instability diagnostics (staleness, reward, clip fraction, ratio + # variance, KL, advantage). `with_entropy` also logs entropy; `none` disables + # them and skips the extra softmax. Requires pipeline_parallel == 1. + metrics: basic data: # Token budget per micro-batch. Must be >= the longest packed sequence, so it tracks From 2557e71f2bced4f1f3d538955e0994fa8012565d Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 13 Jul 2026 12:13:51 -0400 Subject: [PATCH 07/11] =?UTF-8?q?RL=20recipe:=20lower=20ZeRO=20stage=20(0.?= =?UTF-8?q?5B=E2=86=921,=207B=E2=86=922),=20document=20DeepSpeed-side=20st?= =?UTF-8?q?age?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- docs/recipes/reinforcement-learning.md | 21 ++++++++++++++++----- examples/qwen_gspo_0.5b.yaml | 4 +++- examples/qwen_gspo_7b.yaml | 6 ++++-- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/docs/recipes/reinforcement-learning.md b/docs/recipes/reinforcement-learning.md index 340bb926d..cc5664a09 100644 --- a/docs/recipes/reinforcement-learning.md +++ b/docs/recipes/reinforcement-learning.md @@ -92,10 +92,15 @@ Per-size differences: | | Qwen2.5-0.5B | Qwen2.5-7B | | --- | --- | --- | | Head | tied | untied | -| `zero_stage` | `2` | `3` | +| `zero_stage` | `1` | `2` | | `mlp.recompute_level` | — | `full` | | Trainer sharding | 1 replica (FSDP size 1) | 4 replicas (FSDP size 4) | +At FSDP size 1 (the 0.5B trainer) ZeRO shards across nothing, so the stage is a no-op — +`1` is simply the minimal setting. On the 4-GPU 7B trainer, ZeRO-2 shards optimizer state +and gradients (enough) and avoids ZeRO-3's per-step parameter all-gather; go to `3` only if +it OOMs (the fp32 lm-head logits are the large transient). + Each size was run in three precision arms — **bf16** (default), **fp16-matched**, and **fp32-matched** — alongside a **DeepSpeed** parity baseline (see the **DeepSpeed parity baseline** section). @@ -135,7 +140,7 @@ GSPO epsilon. ## 🧮 DeepSpeed parity baseline -For A/B comparison against a DeepSpeed ZeRO-3 trainer, match these on the DeepSpeed side: +For A/B comparison against a DeepSpeed trainer, match these on the DeepSpeed side: | Fast-LLM | DeepSpeed equivalent | | --- | --- | @@ -145,9 +150,15 @@ For A/B comparison against a DeepSpeed ZeRO-3 trainer, match these on the DeepSp | `schedule.docs_per_step: 256` | `gradient_accumulation_passes: 256` | | `optimizer.beta_1/beta_2` | Adam `beta1/beta2` (same values) | | `fp32_lm_head: true` | fp32 LM head (matched) | - -DeepSpeed runs ZeRO-3 bf16 with gradient checkpointing. The samplers, temperature, -`epsilon`, seed, and sequence shape are identical across both trainers. +| `multi_stage.zero_stage: 1 / 2` | `deepspeed_config: deepspeed_stage{1,2}_bf16` | + +The orchestrator's DeepSpeed trainer picks its ZeRO stage from its own DeepSpeed config +(a `deepspeed_config` selecting a stage JSON), independently of the Fast-LLM `zero_stage`. +The ZeRO stage is a memory/sharding choice and is mathematically equivalent across stages, +so it does **not** affect the reward comparison and the two trainers need not match — pick +the lowest stage that fits (stage 1 for a single-GPU trainer, stage 2 for the 4-GPU 7B +trainer). The samplers, temperature, `epsilon`, seed, and sequence shape are identical +across both trainers. ## 📊 Measuring lag and instability diff --git a/examples/qwen_gspo_0.5b.yaml b/examples/qwen_gspo_0.5b.yaml index bf6140d4c..75007a7a9 100644 --- a/examples/qwen_gspo_0.5b.yaml +++ b/examples/qwen_gspo_0.5b.yaml @@ -21,7 +21,9 @@ model: sequence_data_parallel: 1 timeout: 3600 multi_stage: - zero_stage: 2 + # Single trainer GPU (FSDP size 1) — ZeRO shards across nothing, so the stage is a + # no-op; 1 is the minimal choice. Raise only if the trainer is sharded over >1 GPU. + zero_stage: 1 base_model: head: # Keep the LM head in fp32 to match the sampler's fp32-logit path (cross-engine diff --git a/examples/qwen_gspo_7b.yaml b/examples/qwen_gspo_7b.yaml index 3bcc2c7c1..3ecb0cb09 100644 --- a/examples/qwen_gspo_7b.yaml +++ b/examples/qwen_gspo_7b.yaml @@ -22,8 +22,10 @@ model: sequence_data_parallel: 1 timeout: 3600 multi_stage: - # Shard weights, gradients and optimizer state across the trainer replicas. - zero_stage: 3 + # ZeRO-2 shards optimizer state + gradients across the 4-GPU trainer (parameters stay + # replicated). Sufficient here and avoids ZeRO-3's per-step parameter all-gather. Bump + # to 3 only if the trainer OOMs — the fp32 lm-head logits are the large transient. + zero_stage: 2 base_model: decoder: block: From d3e31861db848e2ce6a0507cbb2020030e91c3dc Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 13 Jul 2026 12:52:57 -0400 Subject: [PATCH 08/11] RL recipe: mark example YAMLs as illustrative, point to the operative 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 --- docs/recipes/reinforcement-learning.md | 16 +++++++++++----- examples/qwen_gspo_0.5b.yaml | 5 ++++- examples/qwen_gspo_7b.yaml | 5 ++++- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/docs/recipes/reinforcement-learning.md b/docs/recipes/reinforcement-learning.md index cc5664a09..446abae61 100644 --- a/docs/recipes/reinforcement-learning.md +++ b/docs/recipes/reinforcement-learning.md @@ -16,15 +16,21 @@ Qwen2.5 experiments, and the reasoning behind each setting. The orchestration la rollout data pipeline, and the reward model are out of scope — this page is about how to configure the Fast-LLM trainer. -Two ready-to-adapt override fragments accompany this page: +The **operative config is a single committed file on the orchestrator side** — one hydra +config that pins every hyperparameter for the run (both the orchestrator/DeepSpeed settings +and the Fast-LLM `fast_llm:` subtree) and is what the launcher composes; the launch command +itself carries no hyperparameters, only environment (paths, W&B identity). That file is the +source of truth for a run. + +Two illustrative excerpts accompany this page as a readable reference for the Fast-LLM half +of that config — they are **not** the operative source and **not** standalone-runnable (the +model architecture is imported from the Qwen2.5 checkpoint and the data comes from the +stream); treat them as documentation of the knobs, and edit the orchestrator hydra config to +actually change a run: - [`examples/qwen_gspo_0.5b.yaml`](https://github.com/ServiceNow/Fast-LLM/blob/main/examples/qwen_gspo_0.5b.yaml) - [`examples/qwen_gspo_7b.yaml`](https://github.com/ServiceNow/Fast-LLM/blob/main/examples/qwen_gspo_7b.yaml) -They contain only the RL-relevant overrides; the base model config (architecture, vocab, -tokenizer, RoPE) is imported from the Qwen2.5 checkpoint by the orchestrator, and the -data stream is supplied at run time. - ## 🎯 The policy-gradient loss The RL objective is one entry in the head's `losses` dictionary. The dictionary key is a diff --git a/examples/qwen_gspo_0.5b.yaml b/examples/qwen_gspo_0.5b.yaml index 75007a7a9..0b97c5bdd 100644 --- a/examples/qwen_gspo_0.5b.yaml +++ b/examples/qwen_gspo_0.5b.yaml @@ -1,4 +1,7 @@ -# Fast-LLM training overrides for GSPO reinforcement learning on Qwen2.5-0.5B. +# ILLUSTRATIVE EXCERPT of the Fast-LLM knobs for GSPO RL on Qwen2.5-0.5B — NOT the operative +# config and NOT standalone-runnable. A run is driven by a single committed hydra config in +# the RL orchestrator; this file is a readable reference for the fast_llm.* values only. To +# change a run, edit that hydra config, not this file. # # This is NOT a standalone training config. Reinforcement learning is driven by an # external async-RL orchestrator that generates rollouts with vLLM samplers and diff --git a/examples/qwen_gspo_7b.yaml b/examples/qwen_gspo_7b.yaml index 3ecb0cb09..604af5af2 100644 --- a/examples/qwen_gspo_7b.yaml +++ b/examples/qwen_gspo_7b.yaml @@ -1,4 +1,7 @@ -# Fast-LLM training overrides for GSPO reinforcement learning on Qwen2.5-7B. +# ILLUSTRATIVE EXCERPT of the Fast-LLM knobs for GSPO RL on Qwen2.5-7B — NOT the operative +# config and NOT standalone-runnable. A run is driven by a single committed hydra config in +# the RL orchestrator; this file is a readable reference for the fast_llm.* values only. To +# change a run, edit that hydra config, not this file. # # This is NOT a standalone training config. Reinforcement learning is driven by an # external async-RL orchestrator that generates rollouts with vLLM samplers and From 20de9cc54ac1e378c2e81fbadf9da10ea71a0640 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 13 Jul 2026 14:40:49 -0400 Subject: [PATCH 09/11] RL recipe: operative config is per-backend (shared base + fllm/ds), not one file Co-Authored-By: Claude Opus 4.8 --- docs/recipes/reinforcement-learning.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/recipes/reinforcement-learning.md b/docs/recipes/reinforcement-learning.md index 446abae61..76c899b7f 100644 --- a/docs/recipes/reinforcement-learning.md +++ b/docs/recipes/reinforcement-learning.md @@ -16,10 +16,11 @@ Qwen2.5 experiments, and the reasoning behind each setting. The orchestration la rollout data pipeline, and the reward model are out of scope — this page is about how to configure the Fast-LLM trainer. -The **operative config is a single committed file on the orchestrator side** — one hydra -config that pins every hyperparameter for the run (both the orchestrator/DeepSpeed settings -and the Fast-LLM `fast_llm:` subtree) and is what the launcher composes; the launch command -itself carries no hyperparameters, only environment (paths, W&B identity). That file is the +The **operative config is committed on the orchestrator side** — a shared hydra base plus one +config per trainer backend (Fast-LLM / DeepSpeed) that together pin every hyperparameter for +the run (the orchestrator/preprocessor settings, and the trainer's — the Fast-LLM `fast_llm:` +subtree or the DeepSpeed `finetune:` optimizer). The launcher composes those; the launch command +itself carries no hyperparameters, only environment (paths, W&B identity). Those configs are the source of truth for a run. Two illustrative excerpts accompany this page as a readable reference for the Fast-LLM half From 36dd34a2885397390d6de8f1b14ea5e3eb9097c1 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Mon, 13 Jul 2026 17:13:58 -0400 Subject: [PATCH 10/11] docs(rl): 7B ZeRO stage 3 (recompute leaves memory tight for stage 2) 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 --- docs/recipes/reinforcement-learning.md | 16 +++++++++------- examples/qwen_gspo_7b.yaml | 9 +++++---- 2 files changed, 14 insertions(+), 11 deletions(-) diff --git a/docs/recipes/reinforcement-learning.md b/docs/recipes/reinforcement-learning.md index 76c899b7f..9dd7202ec 100644 --- a/docs/recipes/reinforcement-learning.md +++ b/docs/recipes/reinforcement-learning.md @@ -99,14 +99,16 @@ Per-size differences: | | Qwen2.5-0.5B | Qwen2.5-7B | | --- | --- | --- | | Head | tied | untied | -| `zero_stage` | `1` | `2` | +| `zero_stage` | `1` | `3` | | `mlp.recompute_level` | — | `full` | | Trainer sharding | 1 replica (FSDP size 1) | 4 replicas (FSDP size 4) | At FSDP size 1 (the 0.5B trainer) ZeRO shards across nothing, so the stage is a no-op — -`1` is simply the minimal setting. On the 4-GPU 7B trainer, ZeRO-2 shards optimizer state -and gradients (enough) and avoids ZeRO-3's per-step parameter all-gather; go to `3` only if -it OOMs (the fp32 lm-head logits are the large transient). +`1` is simply the minimal setting. On the 4-GPU 7B trainer, MLP recompute is already on, so +memory is tight: ZeRO-2 (optimizer state + gradients sharded, parameters replicated) may not +fit, so `3` — which also shards parameters — is the safe default. Drop to `2` only if there +is headroom (it avoids ZeRO-3's per-step parameter all-gather; the fp32 lm-head logits are +the large transient). Each size was run in three precision arms — **bf16** (default), **fp16-matched**, and **fp32-matched** — alongside a **DeepSpeed** parity baseline (see @@ -157,14 +159,14 @@ For A/B comparison against a DeepSpeed trainer, match these on the DeepSpeed sid | `schedule.docs_per_step: 256` | `gradient_accumulation_passes: 256` | | `optimizer.beta_1/beta_2` | Adam `beta1/beta2` (same values) | | `fp32_lm_head: true` | fp32 LM head (matched) | -| `multi_stage.zero_stage: 1 / 2` | `deepspeed_config: deepspeed_stage{1,2}_bf16` | +| `multi_stage.zero_stage: 1 / 3` | `deepspeed_config: deepspeed_stage{1,3}_bf16` | The orchestrator's DeepSpeed trainer picks its ZeRO stage from its own DeepSpeed config (a `deepspeed_config` selecting a stage JSON), independently of the Fast-LLM `zero_stage`. The ZeRO stage is a memory/sharding choice and is mathematically equivalent across stages, so it does **not** affect the reward comparison and the two trainers need not match — pick -the lowest stage that fits (stage 1 for a single-GPU trainer, stage 2 for the 4-GPU 7B -trainer). The samplers, temperature, `epsilon`, seed, and sequence shape are identical +the lowest stage that fits (stage 1 for a single-GPU trainer, stage 3 for the 4-GPU 7B +trainer with recompute on). The samplers, temperature, `epsilon`, seed, and sequence shape are identical across both trainers. ## 📊 Measuring lag and instability diff --git a/examples/qwen_gspo_7b.yaml b/examples/qwen_gspo_7b.yaml index 604af5af2..437c1e09c 100644 --- a/examples/qwen_gspo_7b.yaml +++ b/examples/qwen_gspo_7b.yaml @@ -25,10 +25,11 @@ model: sequence_data_parallel: 1 timeout: 3600 multi_stage: - # ZeRO-2 shards optimizer state + gradients across the 4-GPU trainer (parameters stay - # replicated). Sufficient here and avoids ZeRO-3's per-step parameter all-gather. Bump - # to 3 only if the trainer OOMs — the fp32 lm-head logits are the large transient. - zero_stage: 2 + # ZeRO-3 shards parameters too (on top of optimizer state + gradients). With MLP + # recompute already on, the 4-GPU trainer runs near its memory limit, so ZeRO-2's + # replicated parameters may not fit — stage 3 is the safe default. Drop to 2 only if + # there is headroom (2 avoids the per-step parameter all-gather). + zero_stage: 3 base_model: decoder: block: From d3d20226a2395281eca4a779e85b8688621248d2 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 15 Jul 2026 17:50:00 -0400 Subject: [PATCH 11/11] Measure RL data-wait in the prefetch loop, not the schedule The docs_per_step path drains the data loader into a buffer before run_step, so the schedule runner's data_wait_time_ms reads ~0 -- it pops an in-memory buffer. Time the next(data_iterator) pulls in _prefetch_to_doc_target (the real input-starvation wait) and overwrite train_metrics["data_wait_time_ms"] on logging steps. Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/training/trainer.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py index 4300756ff..6c6094a78 100644 --- a/fast_llm/engine/training/trainer.py +++ b/fast_llm/engine/training/trainer.py @@ -157,13 +157,18 @@ def _get_or_build_schedule(self, n_microbatches: int) -> Schedule: ) return self._schedule_cache[n_microbatches] - def _prefetch_to_doc_target(self, data_iterator) -> list: + def _prefetch_to_doc_target(self, data_iterator) -> tuple[list, float]: target = self._config.schedule.docs_per_step bfmb = self._config.schedule.breadth_first_micro_batches buffer = [] total_docs = 0 + # Time blocked pulling micro-batches from the stream — the real input-starvation wait, which + # the schedule runner can't measure because it receives an already-prefetched buffer. + data_wait_time_ms = 0.0 while total_docs < target: + wait_start = time.perf_counter() mb = next(data_iterator) + data_wait_time_ms += (time.perf_counter() - wait_start) * 1000 mb[0].share_batch_data(mb, self._distributed) total_docs += mb[0].num_documents_in_batch buffer.append(mb) @@ -176,7 +181,7 @@ def _prefetch_to_doc_target(self, data_iterator) -> list: for mb in buffer: for mi in mb: mi.num_documents_in_batch = total_docs - return buffer + return buffer, data_wait_time_ms @abc.abstractmethod def _get_data(self) -> Data: @@ -259,7 +264,7 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: # TODO: Data loader hates getting all micro-batches at once. # (Also preprocessing adds overhead) if self._config.schedule.docs_per_step > 0: - buffer = self._prefetch_to_doc_target(train_iterator) + buffer, data_wait_time_ms = self._prefetch_to_doc_target(train_iterator) step_schedule = self._get_or_build_schedule(len(buffer)) reduced_losses, update_successful, train_metrics, step_num_documents = self._runner.run_step( iter(buffer), @@ -268,6 +273,10 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: documents_seen=self._documents_seen, return_metrics=is_logging, ) + if is_logging: + # run_step sees the prefetched buffer, so its own data-wait timer reads ~0; + # replace it with the real stream-pull wait measured during prefetch. + train_metrics["data_wait_time_ms"] = data_wait_time_ms else: reduced_losses, update_successful, train_metrics, step_num_documents = self._runner.run_step( train_iterator,