Add gradient accumulation for training (single-GPU/DDP/FSDP2)

[eldarkurtic]

Summary

Adds 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 (per-step batch × N) without extra per-microbatch memory. Matches the convention used by peer speculative-decoding frameworks (DeepSpec, TorchSpec, SpecForge): loss scaled by 1/N, optimizer + LR scheduler stepped once per accumulation window, counting in optimizer (macro) steps.

Behavior

• Macro-step counting: global_step, --log-freq, --checkpoint-freq (when < 1), and the LR scheduler total all advance once per accumulation window (an accum=N run does 1/N the scheduler steps per epoch).
• Drop remainder: 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 (used by the scheduler math, the __init__ guard, and the loop boundary).
• _maybe_no_sync: on DDP, skips the gradient all-reduce on non-boundary microbatches and syncs the accumulated gradient on the boundary. Single-GPU (raw module) and FSDP2 (fully_shard, DTensor params) fall back to a no-op context and stay correct because gradients accumulate additively into .grad.
• Guardrail: Trainer raises early if N exceeds the batches-per-epoch (which would otherwise run zero optimizer steps).
• accum=1 is behavior-identical to the previous one-step-per-batch loop (regression-safe).

Verification

Numerically confirmed equivalent to the mean of per-microbatch gradients (accumulated grad vs. reference) on real hardware:

Path	Launch	cadence (accum=1 / 4)	max abs grad diff
Single-GPU	python	13 / 3	3.7e-9
DDP	torchrun ×8	13 / 3	9.3e-9
FSDP2 (fully_shard)	torchrun ×4	13 / 3	2.3e-9

All three exercise the real Trainer.train_epoch → _accumulate_and_step → _maybe_no_sync path. DDP/FSDP2 use different data per rank, so cross-rank averaging composes with accumulation.

Unit tests added: scheduler-step math (accum, remainder, accum=1 regression), numerical grad-equivalence on a toy model, _maybe_no_sync behavior (single-process DDP), the accum-too-large guard, and CLI parsing/rejection.

Test plan

• python -m pytest tests/unit/train/test_gradient_accumulation.py tests/unit/train/test_trainer_scheduler.py tests/unit/train/test_cli_args.py — passing.
• make quality (ruff + format + mypy) clean on changed files.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: a8fcf34e-3330-47ca-bd36-d184d4aca451

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Added a validated gradient-accumulation CLI option, wired it into TrainerConfig, and updated training to count complete accumulation windows for optimizer, scheduler, logging, and checkpoint steps. Added DDP synchronization, remainder handling, validation, documentation, and unit tests.

Changes

Gradient accumulation training

Layer / File(s)	Summary
CLI configuration and documentation scripts/train.py, src/speculators/train/trainer.py, docs/cli/train.md, tests/unit/train/test_cli_args.py	Adds --gradient-accumulation-steps, validates positive values, defaults to 1, passes the value into TrainerConfig, and documents its training semantics.
Optimizer-step resolution and validation src/speculators/train/trainer.py, tests/unit/train/test_trainer_scheduler.py, tests/unit/train/test_gradient_accumulation.py	Computes optimizer and scheduler steps from complete accumulation windows, drops trailing remainders, and rejects configurations that produce no optimizer step per epoch.
Accumulation execution and synchronization src/speculators/train/trainer.py, tests/unit/train/test_gradient_accumulation.py	Accumulates scaled gradients, skips DDP synchronization on non-boundary microsteps, performs optimizer and scheduler operations at boundaries, and gates global-step and checkpoint updates accordingly.

Possibly related PRs

• vllm-project/speculators#651: Updates the same scheduler step-resolution logic and its unit tests.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly matches the main change: adding gradient accumulation support across training backends.
Description check	✅ Passed	The description is directly related to the changeset and accurately summarizes the new gradient accumulation behavior.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews

	Protection	Waiting on
🔴	Require approval from approved reviewers list	👀 reviews

🔴 Require approval from approved reviewers list

Waiting for any of

• approved-reviews-by = dsikka
• approved-reviews-by = fynnsu
• approved-reviews-by = orestis-z
• approved-reviews-by = rahul-tuli
• approved-reviews-by = shanjiaz

This rule is failing.

All pull requests must have at least one approving review from a member of the approved reviewers list before merging.

• any of:
  • approved-reviews-by = dsikka
  • approved-reviews-by = fynnsu
  • approved-reviews-by = orestis-z
  • approved-reviews-by = rahul-tuli
  • approved-reviews-by = shanjiaz

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @eldarkurtic.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/train/trainer.py`:
- Around line 191-197: Update the training-state persistence and resume
validation around the trainer’s accumulation setup: store
gradient_accumulation_steps in training_state.json, read it when resuming, and
reject checkpoints whose stored value differs from the configured value. Treat
missing accumulation metadata in legacy state as 1, while preserving the
existing zero-optimizer-step validation.
- Around line 506-519: Update checkpoint cadence calculations in the training
loop around _optimizer_steps_per_epoch, deriving step_interval from
optimizer-step counts rather than raw microbatch counts. Compare the
optimizer-step position using local_step // accum in the checkpoint modulo
logic, while preserving boundary alignment and the existing end-of-epoch guard.
- Around line 125-132: Update _optimizer_steps_per_epoch to validate that accum
is at least 1 before performing floor division, raising the trainer API’s
appropriate validation error for non-positive values. Add a direct TrainerConfig
test covering gradient_accumulation_steps=0 and a negative value, while
preserving the existing drop-remainder behavior for valid accumulation values.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: d25c7b4e-c115-47c0-a052-1c0c4805ac84

📥 Commits

Reviewing files that changed from the base of the PR and between 77045b3 and cab6621.

📒 Files selected for processing (6)

• docs/cli/train.md
• scripts/train.py
• src/speculators/train/trainer.py
• tests/unit/train/test_cli_args.py
• tests/unit/train/test_gradient_accumulation.py
• tests/unit/train/test_trainer_scheduler.py

[mergify]

The quality checks have failed. Please run make style and make quality under
the root directory to address the lint failures. You will need to install the
dev optional install to get the required linting packages:
https://github.com/vllm-project/speculators/blob/main/CONTRIBUTING.md

[rahul-tuli]

Hi @eldarkurtic #841 just landed, quick heads up

The new config CLI (#841) is in main, so parse_args() in scripts/train.py is gone every flag is now a pydantic Field in src/speculators/train/config/schema.py. Your --gradient-accumulation-steps just needs to move there. Only the CLI wiring changes; the Trainer logic, tests, and docs all rebase cleanly.

The gist:

• Delete the _positive_int helper + the add_argument(...) block (no more parse_args).
• Add the flag as a Field on TrainerArgs, with a >= 1 validator:

# src/speculators/train/config/schema.py — in TrainerArgs
gradient_accumulation_steps: int = Field(
    default=1,
    description="Accumulate gradients over N microbatches before each optimizer step ...",
)

@field_validator("gradient_accumulation_steps")
@classmethod
def _validate_gradient_accumulation_steps(cls, v: int) -> int:
    if v < 1:
        raise ValueError("--gradient-accumulation-steps must be a positive integer (>= 1)")
    return v

• Everything else stays — main() still gets args.gradient_accumulation_steps through the flatten() bridge, so your TrainerConfig(...) line is untouched. In test_cli_args.py, just drop the from scripts.train import parse_args line (the _parse helper is already on the new resolver).

Easiest path: I already rebased + migrated your commit (authorship kept). Grab the patch below and:

git fetch origin
git checkout -B add-gradient-accumulation origin/main
git am pr859-new-cli.patch
git push --force-with-lease origin add-gradient-accumulation

Tested on the rebased result CLI, scheduler, grad-accum, resume, and schema tests all pass. 👍 Shout if you'd rather I just push it up for you.

Full patch — pr859-new-cli.patch

From 116463c82c0c7d8befcc14b64df8d35d830f29a4 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) <noreply@anthropic.com>
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        |  17 ++
 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, 417 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 bb0c6f9..3b21327 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 cedf2d6..62216f8 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 1de8d83..207a2c9 100644
--- a/src/speculators/train/config/schema.py
+++ b/src/speculators/train/config/schema.py
@@ -400,6 +400,15 @@ 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,
+        description="Accumulate gradients over N microbatches before each optimizer "
+        "step, giving an effective batch size of `per-step batch * N` without extra "
+        "per-microbatch memory. 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. Must be >= 1 and must not change across a resume.",
+    )
 
     @field_validator("checkpoint_freq")
     @classmethod
@@ -413,6 +422,14 @@ class TrainerArgs(_Group):
             )
         return v
 
+    @field_validator("gradient_accumulation_steps")
+    @classmethod
+    def _validate_gradient_accumulation_steps(cls, v: int) -> int:
+        if v < 1:
+            raise ValueError("--gradient-accumulation-steps must be a positive "
+                             "integer (>= 1)")
+        return v
+
 
 class LoggingArgs(_Group):
     logger: str = Field(
diff --git a/src/speculators/train/trainer.py b/src/speculators/train/trainer.py
index e5f8572..1af4e96 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 @@ class Trainer:
         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 @@ class Trainer:
                 "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 @@ class Trainer:
                 # 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 @@ class Trainer:
         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 @@ class Trainer:
         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 @@ class Trainer:
                 )
 
             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 @@ class Trainer:
                     },
                     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 @@ class Trainer:
             ):
                 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 dbd47fc..bd352aa 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 0000000..f67ceb5
--- /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 8421f94..b7ee188 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 2aa47a5..cda9937 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).
-- 
2.47.3

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @eldarkurtic.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @eldarkurtic.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork