From ea39e9e32af1ea9e4ce3d65e340271e7bc5b6a7c Mon Sep 17 00:00:00 2001 From: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> Date: Fri, 24 Jul 2026 03:17:05 -0400 Subject: [PATCH] Add gradient accumulation for training (single-GPU/DDP/FSDP2) Introduce a `--gradient-accumulation-steps` / `TrainerConfig.gradient_accumulation_steps` knob (default 1) that accumulates gradients over N microbatches before each optimizer step, growing the effective batch size without extra per-microbatch memory. - Loss is scaled by 1/N; the optimizer, gradient clip, and LR scheduler run once per accumulation window (macro-step counting for global_step, log_freq, checkpoint cadence, and scheduler total steps). - Trailing microbatches that don't fill a full window are dropped each epoch; `_optimizer_steps_per_epoch` is the single source of truth for this policy. - `_maybe_no_sync` skips the DDP all-reduce on non-boundary microbatches; the single-GPU and FSDP2 paths fall back to a no-op context and stay correct because gradients accumulate additively into `.grad`. - Trainer raises early when N exceeds the batches-per-epoch (would never step). - accum=1 is behavior-identical to the previous one-step-per-batch loop. Verified numerically equivalent to the mean of per-microbatch gradients on single-GPU, 8-GPU DDP, and 4-GPU FSDP2 (fully_shard), plus unit tests for the scheduler math, no_sync behavior, cadence, and CLI wiring. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Eldar Kurtic <8884008+eldarkurtic@users.noreply.github.com> --- docs/cli/train.md | 2 + scripts/train.py | 1 + src/speculators/train/config/schema.py | 8 + src/speculators/train/trainer.py | 157 +++++++++++--- tests/unit/train/test_cli_args.py | 22 ++ .../unit/train/test_gradient_accumulation.py | 195 ++++++++++++++++++ tests/unit/train/test_mid_epoch_resume.py | 18 ++ tests/unit/train/test_trainer_scheduler.py | 29 +++ 8 files changed, 408 insertions(+), 24 deletions(-) create mode 100644 tests/unit/train/test_gradient_accumulation.py diff --git a/docs/cli/train.md b/docs/cli/train.md index bb0c6f9f2..3b21327c3 100644 --- a/docs/cli/train.md +++ b/docs/cli/train.md @@ -114,6 +114,8 @@ torchrun --standalone --nproc_per_node=4 scripts/train.py \ - **`--lr`** (float, default: `1e-4`) Learning rate. +- **`--gradient-accumulation-steps`** (int, default: `1`) Accumulate gradients over N microbatches before each optimizer step, giving an effective batch size of `per-step batch × N` without extra per-microbatch memory. Must be ≥ 1. Notes: step-based counters (`global_step`, `--scheduler-total-steps`, `--log-freq`, and `--checkpoint-freq` when `< 1`) are counted in optimizer steps, so an accumulated run performs `1/N` as many scheduler steps per epoch; trailing microbatches that don't fill a full window are dropped each epoch; and the value must not change across a resume. On multi-GPU DDP, the gradient all-reduce is skipped on non-boundary microbatches (`no_sync`). + - **`--train-data-ratio`** (float, default: `0.9`) Ratio of data to use for training, the rest of the provided data will be used for validation. - **`--no-resume-from-checkpoint`** (flag) Disable automatic checkpoint resumption. Without this flag, this script will automatically load the latest checkpoint in `{save-path}` if one exists. diff --git a/scripts/train.py b/scripts/train.py index cedf2d6de..62216f8fe 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -675,6 +675,7 @@ def main(cfg: TrainConfig): # noqa: C901 hidden_states_dtype=hidden_states_dtype, log_freq=args.log_freq, fsdp_shard=args.fsdp_shard, + gradient_accumulation_steps=args.gradient_accumulation_steps, ) trainer = Trainer(draft_model, trainer_config, train_loader, val_loader) diff --git a/src/speculators/train/config/schema.py b/src/speculators/train/config/schema.py index 1de8d832e..fc34a4a7f 100644 --- a/src/speculators/train/config/schema.py +++ b/src/speculators/train/config/schema.py @@ -400,6 +400,14 @@ class TrainerArgs(_Group): "parameters are fully replicated (DDP-like). Enable when the model does not " "fit in a single GPU's memory.", ) + gradient_accumulation_steps: int = Field( + default=1, + ge=1, + description="Accumulate gradients over N microbatches before each optimizer " + "step (effective batch = per-step batch * N). The trailing partial window is " + "dropped each epoch; step-based counters (global_step, scheduler total, " + "log/checkpoint cadence) are measured in optimizer steps.", + ) @field_validator("checkpoint_freq") @classmethod diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py index e5f8572d8..1af4e9691 100644 --- a/src/speculators/train/trainer.py +++ b/src/speculators/train/trainer.py @@ -1,3 +1,4 @@ +import contextlib import json import logging import time @@ -98,6 +99,7 @@ class TrainerConfig(NamedTuple): lr: float num_epochs: int save_path: str + gradient_accumulation_steps: int = 1 resume_from_checkpoint: bool = False train_call_kwargs: dict | None = None val_call_kwargs: dict | None = None @@ -121,6 +123,16 @@ class TrainerConfig(NamedTuple): max_steps: int | None = None +def _optimizer_steps_per_epoch(num_batches: int, accum: int) -> int: + """Number of optimizer steps taken in one epoch under gradient accumulation. + + Each optimizer step consumes ``accum`` microbatches; the trailing partial + window (``num_batches % accum`` microbatches) is dropped. Single source of + truth for the drop-remainder policy. + """ + return num_batches // accum + + def _resolve_scheduler_steps( config: TrainerConfig, train_loader_len: int, @@ -132,7 +144,10 @@ def _resolve_scheduler_steps( default of 1% of the resolved total steps. ``scheduler_total_steps`` defaults to ``num_epochs * train_loader_len``. """ - default_total_steps = config.num_epochs * train_loader_len + steps_per_epoch = _optimizer_steps_per_epoch( + train_loader_len, config.gradient_accumulation_steps + ) + default_total_steps = config.num_epochs * steps_per_epoch scheduler_total_steps = ( config.scheduler_total_steps if config.scheduler_total_steps is not None @@ -173,6 +188,16 @@ def __init__( self.rank = get_rank() self.train_loader = train_loader self.val_loader = val_loader + + accum = config.gradient_accumulation_steps + if accum < 1: + raise ValueError(f"gradient_accumulation_steps must be >= 1, got {accum}.") + if _optimizer_steps_per_epoch(len(train_loader), accum) == 0: + raise ValueError( + f"gradient_accumulation_steps={accum} exceeds the number of " + f"batches per epoch ({len(train_loader)}); no optimizer step would " + "ever run. Lower gradient_accumulation_steps or add more data." + ) self.is_distributed = is_distributed() self.resume_from_checkpoint = config.resume_from_checkpoint acc = torch.accelerator.current_accelerator() @@ -197,6 +222,7 @@ def _save_training_state(self, epoch: int, local_step: int) -> None: "epoch": epoch, "local_step": local_step, "global_step": self.global_step, + "gradient_accumulation_steps": self.config.gradient_accumulation_steps, } p = self._training_state_path(epoch) p.parent.mkdir(parents=True, exist_ok=True) @@ -222,6 +248,16 @@ def setup_trainer(self): # Check if this was a mid-epoch checkpoint — if so, resume # from within that epoch rather than jumping to the next one. state = self._load_training_state() + # Accumulation windows are aligned to the saved microbatch index, so + # resuming with a different accum would misalign them; reject it. + saved_accum = state.get("gradient_accumulation_steps", 1) + if state and saved_accum != self.config.gradient_accumulation_steps: + raise ValueError( + "Cannot resume: checkpoint was trained with " + f"gradient_accumulation_steps={saved_accum}, but this run " + f"uses {self.config.gradient_accumulation_steps}. " + "They must match to resume." + ) is_mid_epoch = ( state and state.get("epoch") == self.checkpointer.previous_epoch @@ -371,6 +407,76 @@ def make_scheduler(opt: torch.optim.Optimizer): if self.resume_from_checkpoint and self.checkpointer.previous_epoch != -1: self.checkpointer.load_scheduler_state_dict(self.schedulers) + def _maybe_no_sync(self, is_boundary: bool): + """Skip DDP gradient all-reduce on non-boundary accumulation micro-steps. + + Returns ``model.no_sync()`` only for a real ``DistributedDataParallel`` + model on a non-boundary micro-step; otherwise a no-op context. Single-GPU + (raw module) and FSDP2 (``fully_shard``) fall through to ``nullcontext`` and + remain correct because gradients accumulate additively into ``.grad``. + """ + if isinstance(self.model, DistributedDataParallel) and not is_boundary: + return self.model.no_sync() + return contextlib.nullcontext() + + def _accumulate_and_step( + self, + loss: torch.Tensor, + accum: int, + is_window_start: bool, + is_boundary: bool, + timer: "_StepTimer", + ) -> dict[str, float]: + """Run one accumulation micro-step; step the optimizer only at a boundary. + + Gradients are zeroed at the start of each accumulation window and the loss is + scaled by ``accum`` so the accumulated gradient equals the mean over the + effective (window) batch. On a boundary micro-step the gradient is clipped, + the optimizer and schedulers step. LRs are captured before the scheduler step + to match the pre-step logging semantics. Returns those LRs by optimizer name. + """ + if is_window_start: + self._optimizers_zero_grad() + with self._maybe_no_sync(is_boundary): + (loss / accum).backward() + + current_lrs = { + type(opt).__name__: opt.param_groups[0]["lr"] for opt in self.optimizers + } + if is_boundary: + torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) + timer.mark("bwd") + self._optimizers_step() + self._schedulers_step() + timer.mark("opt") + return current_lrs + + def _maybe_save_step_checkpoint( + self, + epoch: int, + local_step: int, + accum: int, + opt_steps_per_epoch: int, + step_interval: int | None, + is_boundary: bool, + ) -> None: + """Save a mid-epoch checkpoint at the configured optimizer-step cadence. + + Cadence is measured in optimizer steps so it aligns with accumulation-window + boundaries; saves too close to the end of the epoch are skipped. ``local_step`` + (a microbatch index on a window boundary) is what gets persisted, so resume + lands on a fresh window start. + """ + if not is_boundary or step_interval is None or self.config.save_best: + return + opt_step = local_step // accum + if ( + opt_step % step_interval == 0 + and opt_steps_per_epoch - opt_step >= step_interval * MIN_STEP_PCT + # Avoid saving back to back at the end of each epoch + ): + self.maybe_save_checkpoint(epoch, local_step=local_step) + def _optimizers_zero_grad(self): for opt in self.optimizers: opt.zero_grad() @@ -432,17 +538,28 @@ def train_epoch(self, epoch: int): if self.rank == 0: train_loader = tqdm(train_loader, desc=f"Epoch {epoch}") # type: ignore[assignment] + accum = self.config.gradient_accumulation_steps + opt_steps_per_epoch = _optimizer_steps_per_epoch(num_steps, accum) + # Sub-epoch checkpoint cadence is measured in optimizer steps so it aligns + # with accumulation-window boundaries (microbatch counts would rarely match). step_interval = ( - max(1, round(num_steps * self.config.checkpoint_freq)) + max(1, round(opt_steps_per_epoch * self.config.checkpoint_freq)) if self.config.checkpoint_freq < 1 else None ) + # Last microbatch index that completes a full accumulation window; any + # trailing microbatches that can't fill a window are dropped. + last_boundary = opt_steps_per_epoch * accum t_before_fetch = time.perf_counter() timer = _StepTimer() for local_step_rel, batch in enumerate(train_loader, 1): # local_step is 1-based index into the *full* epoch (not the slice). local_step = local_step_rel + skip_steps - timer.reset(self.global_step % self.config.log_freq == 0) + if local_step > last_boundary: + break + is_window_start = (local_step - 1) % accum == 0 + is_boundary = local_step % accum == 0 + timer.reset(is_boundary and self.global_step % self.config.log_freq == 0) timer.mark_value("start", t_before_fetch) gpu_batch = { @@ -461,18 +578,9 @@ def train_epoch(self, epoch: int): ) timer.mark("fwd") - self._optimizers_zero_grad() - loss.backward() - torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0) - - timer.mark("bwd") - self._optimizers_step() - - current_lrs = { - type(opt).__name__: opt.param_groups[0]["lr"] for opt in self.optimizers - } - self._schedulers_step() - timer.mark("opt") + current_lrs = self._accumulate_and_step( + loss, accum, is_window_start, is_boundary, timer + ) t_before_fetch = timer.now() or time.perf_counter() profile = None @@ -501,7 +609,8 @@ def train_epoch(self, epoch: int): }, extra={"step": self.global_step}, ) - self.global_step += 1 + if is_boundary: + self.global_step += 1 if ( self.config.max_steps is not None @@ -509,14 +618,14 @@ def train_epoch(self, epoch: int): ): break - if ( - step_interval is not None - and not self.config.save_best - and local_step % step_interval == 0 - and num_steps - local_step >= step_interval * MIN_STEP_PCT - # Avoid saving back to back ay the end of each epoch - ): - self.maybe_save_checkpoint(epoch, local_step=local_step) + self._maybe_save_step_checkpoint( + epoch, + local_step, + accum, + opt_steps_per_epoch, + step_interval, + is_boundary, + ) @torch.no_grad() def val_epoch(self, epoch: int) -> dict[str, float] | None: diff --git a/tests/unit/train/test_cli_args.py b/tests/unit/train/test_cli_args.py index dbd47fce5..bd352aa93 100644 --- a/tests/unit/train/test_cli_args.py +++ b/tests/unit/train/test_cli_args.py @@ -2,6 +2,8 @@ import argparse +import pytest + from speculators.models.dflash.core import DFlashDraftModel from speculators.models.dspark.core import DSparkDraftModel from speculators.models.eagle3.core import Eagle3DraftModel @@ -169,3 +171,23 @@ def test_no_norm_before_fc_flag(monkeypatch): def test_no_norm_output_flag(monkeypatch): args = _parse(monkeypatch, ["--no-norm-output"]) assert args.norm_output is False + + +# --------------------------------------------------------------------------- +# Gradient accumulation +# --------------------------------------------------------------------------- + + +def test_gradient_accumulation_steps_default(monkeypatch): + args = _parse(monkeypatch, []) + assert args.gradient_accumulation_steps == 1 + + +def test_gradient_accumulation_steps_explicit(monkeypatch): + args = _parse(monkeypatch, ["--gradient-accumulation-steps", "8"]) + assert args.gradient_accumulation_steps == 8 + + +def test_gradient_accumulation_steps_rejects_zero(monkeypatch): + with pytest.raises(SystemExit): + _parse(monkeypatch, ["--gradient-accumulation-steps", "0"]) diff --git a/tests/unit/train/test_gradient_accumulation.py b/tests/unit/train/test_gradient_accumulation.py new file mode 100644 index 000000000..f67ceb5e9 --- /dev/null +++ b/tests/unit/train/test_gradient_accumulation.py @@ -0,0 +1,195 @@ +"""Unit tests for gradient accumulation. + +Covers the two correctness-critical pieces of the feature without spinning up a +full ``Trainer.train_epoch`` run: + +- The numerical identity that N accumulated ``(loss / N).backward()`` micro-steps + produce the same gradient as one backward over the concatenated batch (this is + why loss is scaled by ``1 / gradient_accumulation_steps`` in the loop). +- ``Trainer._maybe_no_sync`` returns a real ``no_sync`` context only for a DDP model + on a non-boundary micro-step, and a no-op context otherwise. +""" + +import contextlib +import os +from types import SimpleNamespace +from typing import cast + +import pytest +import torch +import torch.distributed as dist +from torch import nn +from torch.nn.parallel import DistributedDataParallel +from torch.utils.data import DataLoader, TensorDataset + +from speculators.model import SpeculatorModel +from speculators.train.trainer import Trainer, TrainerConfig + +# --------------------------------------------------------------------------- +# Numerical equivalence: accumulate == single large batch +# --------------------------------------------------------------------------- + + +def _grad_vector(model: nn.Module) -> torch.Tensor: + return torch.cat( + [p.grad.reshape(-1) for p in model.parameters() if p.grad is not None] + ) + + +@pytest.mark.parametrize( + ("accum", "batch"), + [(1, 4), (4, 4), (8, 3)], +) +def test_accumulated_grad_matches_single_large_batch(accum: int, batch: int) -> None: + """N micro-steps of ``(loss / N).backward()`` == one backward over the whole + batch with mean-reduction loss (grads accumulate additively into ``.grad``).""" + torch.manual_seed(0) + in_dim, out_dim = 5, 3 + model = nn.Linear(in_dim, out_dim).double() + loss_fn = nn.MSELoss() # mean reduction + + # Equal-size micro-batches so mean(concat) == (1/N) * sum(mean(chunk_i)). + xs = [torch.randn(batch, in_dim, dtype=torch.float64) for _ in range(accum)] + ys = [torch.randn(batch, out_dim, dtype=torch.float64) for _ in range(accum)] + + # Accumulation path (mirrors train_epoch: zero at window start, loss / accum). + model.zero_grad(set_to_none=True) + for x, y in zip(xs, ys, strict=True): + (loss_fn(model(x), y) / accum).backward() + acc_grad = _grad_vector(model) + + # Reference: single backward over the concatenated batch. + model.zero_grad(set_to_none=True) + loss_fn(model(torch.cat(xs)), torch.cat(ys)).backward() + ref_grad = _grad_vector(model) + + assert torch.allclose(acc_grad, ref_grad, atol=1e-10, rtol=1e-8) + + +def test_accum_one_is_exact_single_step() -> None: + """accum=1 must be bit-for-bit identical to a plain single backward.""" + torch.manual_seed(1) + model = nn.Linear(4, 2).double() + x = torch.randn(6, 4, dtype=torch.float64) + y = torch.randn(6, 2, dtype=torch.float64) + loss_fn = nn.MSELoss() + + model.zero_grad(set_to_none=True) + (loss_fn(model(x), y) / 1).backward() + acc_grad = _grad_vector(model) + + model.zero_grad(set_to_none=True) + loss_fn(model(x), y).backward() + plain_grad = _grad_vector(model) + + assert torch.equal(acc_grad, plain_grad) + + +# --------------------------------------------------------------------------- +# _maybe_no_sync +# --------------------------------------------------------------------------- + + +def _call_maybe_no_sync(model: nn.Module, is_boundary: bool): + # Bind the unbound method to a minimal stand-in carrying only `.model`. + stand_in = cast("Trainer", SimpleNamespace(model=model)) + return Trainer._maybe_no_sync(stand_in, is_boundary) + + +def test_maybe_no_sync_plain_module_is_nullcontext() -> None: + model = nn.Linear(2, 2) + for is_boundary in (True, False): + ctx = _call_maybe_no_sync(model, is_boundary) + assert isinstance(ctx, contextlib.nullcontext) + + +@pytest.fixture +def single_process_group(): + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29591") + dist.init_process_group(backend="gloo", rank=0, world_size=1) + try: + yield + finally: + dist.destroy_process_group() + + +def test_maybe_no_sync_ddp_skips_sync_off_boundary(single_process_group) -> None: + ddp = DistributedDataParallel(nn.Linear(2, 2)) + + # Non-boundary micro-step: a real no_sync context (not the no-op). + off_boundary = _call_maybe_no_sync(ddp, is_boundary=False) + assert not isinstance(off_boundary, contextlib.nullcontext) + + # Boundary micro-step: sync must happen, so a no-op context. + on_boundary = _call_maybe_no_sync(ddp, is_boundary=True) + assert isinstance(on_boundary, contextlib.nullcontext) + + +# --------------------------------------------------------------------------- +# Validation: accum larger than an epoch is a hard error (never steps) +# --------------------------------------------------------------------------- + + +def test_trainer_rejects_accum_larger_than_epoch() -> None: + """accum > batches-per-epoch would run zero optimizer steps; reject it early.""" + loader = DataLoader(TensorDataset(torch.arange(3)), batch_size=1) # 3 batches + cfg = TrainerConfig( + lr=1e-4, + num_epochs=1, + save_path="unused", # error raises before any checkpoint I/O + gradient_accumulation_steps=5, + ) + with pytest.raises(ValueError, match="exceeds the number of"): + Trainer(cast("SpeculatorModel", nn.Identity()), cfg, loader) + + +@pytest.mark.parametrize("accum", [0, -1]) +def test_trainer_rejects_non_positive_accum(accum: int) -> None: + """The Trainer API guards accum >= 1 even when constructed directly (bypassing + the CLI validator): 0 would ZeroDivisionError and negatives negate gradients.""" + loader = DataLoader(TensorDataset(torch.arange(3)), batch_size=1) + cfg = TrainerConfig( + lr=1e-4, + num_epochs=1, + save_path="unused", + gradient_accumulation_steps=accum, + ) + with pytest.raises(ValueError, match="must be >= 1"): + Trainer(cast("SpeculatorModel", nn.Identity()), cfg, loader) + + +# --------------------------------------------------------------------------- +# Sub-epoch checkpoint cadence is measured in optimizer steps +# --------------------------------------------------------------------------- + + +def test_step_checkpoint_cadence_uses_optimizer_steps() -> None: + """Mid-epoch checkpoints fire on the optimizer-step cadence and stay aligned to + accumulation boundaries. With 100 microbatches, accum=4, checkpoint_freq=0.5 the + interval is 12 optimizer steps, so the only qualifying save is at optimizer step + 12 (microbatch 48); step 24 is dropped by the end-of-epoch guard.""" + num_steps, accum = 100, 4 + opt_steps_per_epoch = num_steps // accum # 25 + step_interval = max(1, round(opt_steps_per_epoch * 0.5)) # 12 + + saved: list[int] = [] + stand_in = cast( + "Trainer", + SimpleNamespace( + config=SimpleNamespace(save_best=False), + maybe_save_checkpoint=lambda epoch, local_step: saved.append(local_step), + ), + ) + for local_step in range(1, num_steps + 1): + Trainer._maybe_save_step_checkpoint( + stand_in, + 0, + local_step, + accum, + opt_steps_per_epoch, + step_interval, + local_step % accum == 0, + ) + + assert saved == [48] diff --git a/tests/unit/train/test_mid_epoch_resume.py b/tests/unit/train/test_mid_epoch_resume.py index 8421f9410..b7ee188b1 100644 --- a/tests/unit/train/test_mid_epoch_resume.py +++ b/tests/unit/train/test_mid_epoch_resume.py @@ -127,6 +127,7 @@ def _make_trainer( trained_steps: list[tuple[int, int, int]], resume: bool = False, epochs: int = 1, + accum: int = 1, ) -> _MockTrainer: cfg = TrainerConfig( save_path=save_path, @@ -136,6 +137,7 @@ def _make_trainer( checkpoint_freq=0.3, log_freq=1, scheduler_type="none", + gradient_accumulation_steps=accum, ) trainer = _MockTrainer(_dummy_model(), cfg, _make_loader()) trainer._trained_steps = trained_steps @@ -169,6 +171,7 @@ def test_mid_epoch_checkpoint_saves_training_state( "epoch": 0, "local_step": step_interval, "global_step": step_interval, + "gradient_accumulation_steps": 1, } assert state == expected @@ -403,3 +406,18 @@ def test_fast_skip_sampler_slice_avoids_skipped_getitem( assert sampler.generated_for_epoch == 0 assert sampler._cached_generated_batches == (0, sampler.all_batches[3:]) assert dataset.seen_indices == list(range(3, 10)) + + +def test_resume_rejects_gradient_accumulation_mismatch( + trained_steps: list[tuple[int, int, int]], +) -> None: + """Resuming with a different gradient_accumulation_steps misaligns accumulation + windows, so it is rejected rather than silently corrupting the run.""" + with tempfile.TemporaryDirectory() as tmpdir: + # Run 1: persist a checkpoint trained with accum=2. + t1 = _make_trainer(tmpdir, trained_steps=trained_steps, accum=2) + t1.maybe_save_checkpoint(0, local_step=2) # boundary for accum=2 + + # Run 2: resuming with a different accum must fail loudly. + with pytest.raises(ValueError, match="Cannot resume"): + _make_trainer(tmpdir, trained_steps=[], resume=True, accum=3) diff --git a/tests/unit/train/test_trainer_scheduler.py b/tests/unit/train/test_trainer_scheduler.py index 2aa47a541..cda993752 100644 --- a/tests/unit/train/test_trainer_scheduler.py +++ b/tests/unit/train/test_trainer_scheduler.py @@ -23,6 +23,35 @@ def test_scheduler_steps_default_to_one_percent_of_training_steps(): assert warmup_steps == 1 +def test_scheduler_steps_scale_down_with_gradient_accumulation(): + # 5 epochs * (20 // 4) = 25 optimizer steps. + warmup_steps, total_steps = _resolve_scheduler_steps( + make_config(gradient_accumulation_steps=4), 20 + ) + + assert total_steps == 25 + assert warmup_steps == 0 # 25 // 100 + + +def test_scheduler_steps_drop_remainder_with_gradient_accumulation(): + # 21 // 4 == 5 (the trailing microbatch is dropped), so 5 epochs * 5 == 25. + _, total_steps = _resolve_scheduler_steps( + make_config(gradient_accumulation_steps=4), 21 + ) + + assert total_steps == 25 + + +def test_scheduler_steps_accum_one_is_unchanged(): + # Regression guard: accum=1 must match the pre-accumulation behavior. + warmup_steps, total_steps = _resolve_scheduler_steps( + make_config(gradient_accumulation_steps=1), 20 + ) + + assert total_steps == 100 + assert warmup_steps == 1 + + def test_scheduler_total_steps_only_defaults_warmup_to_one_percent_of_total(): # default_total_steps is num_epochs * loader_len = 100, but the explicit # scheduler_total_steps override must drive the 1% warmup fallback (10, not 1).