diff --git a/docs/recipes/reinforcement-learning.md b/docs/recipes/reinforcement-learning.md new file mode 100644 index 000000000..9dd7202ec --- /dev/null +++ b/docs/recipes/reinforcement-learning.md @@ -0,0 +1,243 @@ +--- +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. + +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 +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) + +## 🎯 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`** — 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` + +`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` | `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, 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 +the **DeepSpeed parity baseline** section). + +### 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 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) | +| `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 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 + +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 + 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. +- **`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. + +## 🆕 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. +- **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, from which the reward and rollout staleness are logged. + +## 📌 Open items + +- **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 new file mode 100644 index 000000000..0b97c5bdd --- /dev/null +++ b/examples/qwen_gspo_0.5b.yaml @@ -0,0 +1,76 @@ +# 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 +# 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: + # 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 + # 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 + # 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 + # 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..437c1e09c --- /dev/null +++ b/examples/qwen_gspo_7b.yaml @@ -0,0 +1,81 @@ +# 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 +# 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: + # 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: + 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 + # 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 + # 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/fast_llm/data/document/language_model.py b/fast_llm/data/document/language_model.py index b70814c88..6664f8de6 100644 --- a/fast_llm/data/document/language_model.py +++ b/fast_llm/data/document/language_model.py @@ -223,7 +223,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) @@ -231,7 +231,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/data/document/token.py b/fast_llm/data/document/token.py index 380dfd0ff..7a255dd90 100644 --- a/fast_llm/data/document/token.py +++ b/fast_llm/data/document/token.py @@ -117,12 +117,16 @@ def _get_model_input(self, begin: int, end: int, config: LengthPreprocessingConf global_cumulative_lengths = torch.from_numpy(padded_cumsum(self.lengths)).to( dtype=torch.int32, device=self.device ) + # 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`. 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`. model_input.global_document_index_q = torch.searchsorted( global_cumulative_lengths, torch.arange(begin, end, device=self.device), side="right", out_int32=True, - ) + ).clamp_(max=num_documents) model_input.num_documents_in_sequence = num_documents LengthModelInputPreprocessor( diff --git a/fast_llm/engine/schedule/config.py b/fast_llm/engine/schedule/config.py index 54ae03fcc..cf5a01ce9 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 58c641191..079cbd3a0 100644 --- a/fast_llm/engine/schedule/runner.py +++ b/fast_llm/engine/schedule/runner.py @@ -328,9 +328,10 @@ def _preprocess_data( # Time blocked on the data loader (input starvation), kept separate from the CPU # preprocessing below. Preprocessing runs interleaved with the schedule's compute, so its # clock is paused across each yield to exclude the compute happening between micro-batches. + n_micro_batches = context.schedule._eff_sequential_micro_batches if measure_time: wait_start = time.perf_counter() - model_inputs = [next(data_iterator) for _ in range(self._config.sequential_micro_batches)] + model_inputs = [next(data_iterator) for _ in range(n_micro_batches)] if measure_time: context.metrics["data_wait_time_ms"] += (time.perf_counter() - wait_start) * 1000 preprocess_start = time.perf_counter() @@ -351,7 +352,7 @@ def _preprocess_data( "grad_output": grad_output, "documents_seen": context.documents_seen, "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 e11307e1e..6c6094a78 100644 --- a/fast_llm/engine/training/trainer.py +++ b/fast_llm/engine/training/trainer.py @@ -116,10 +116,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, ) @@ -141,6 +143,46 @@ 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) -> 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) + 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, data_wait_time_ms + @abc.abstractmethod def _get_data(self) -> Data: pass @@ -221,13 +263,28 @@ 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, step_num_documents = self._runner.run_step( - train_iterator, - self._schedule, - iteration=self._completed_steps, - documents_seen=self._documents_seen, - return_metrics=is_logging, - ) + if self._config.schedule.docs_per_step > 0: + 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), + step_schedule, + iteration=self._completed_steps, + 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, + self._schedule, + iteration=self._completed_steps, + documents_seen=self._documents_seen, + return_metrics=is_logging, + ) self._documents_seen += step_num_documents diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index 298e691b9..a485ad2a8 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -133,6 +133,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 [ 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 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