Gradient accumulation drops final partial window and advances LR scheduler per microbatch

Describe the bug

The supervised single-device, multi-device (both optimizer modes), and DDP epoch loops discard the final partial gradient-accumulation window when the number of training batches is not divisible by the accumulation interval.

The same loops call learner.lr_step() per microbatch rather than per optimizer update. With an accumulation interval of k, the first k - 1 scheduler values are unused and non-constant schedules advance approximately k times too quickly. DDP additionally advances the scheduler peer_count times for each local microbatch.

Consequences:

• batch_count % accumulation_interval batches do not affect the model each epoch.
• If accumulation_interval > batch_count, an epoch can process and report every batch without performing any optimizer update.
• LR schedules are compressed whenever accumulation is enabled.
• grads_accumulation(0) is accepted and silently behaves like an update every batch.

Affected code:

• crates/burn-train/src/learner/supervised/strategies/single/epoch.rs
• crates/burn-train/src/learner/supervised/strategies/multi/epoch.rs
• crates/burn-train/src/learner/supervised/strategies/ddp/epoch.rs

Confirmed on upstream main at 52c101f455e7228f35051f11f544d672373938be.

To reproduce

This integration test uses the existing crates/burn-train/tests/common toy model and dataloaders. The loader produces two training batches, while the accumulation interval is three.

mod common;

use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};

use burn_core::tensor::Device;
use burn_optim::lr_scheduler::{LrScheduler, LrSchedulerRecord};
use burn_train::{SupervisedTraining, logger::InMemoryMetricLogger};
use common::*;

#[derive(Clone)]
struct CountingScheduler {
    calls: Arc<AtomicUsize>,
}

impl LrScheduler for CountingScheduler {
    fn step(&mut self) -> f64 {
        self.calls.fetch_add(1, Ordering::SeqCst);
        0.1
    }

    fn to_record(&self) -> LrSchedulerRecord {
        LrSchedulerRecord::new()
    }

    fn load_record(&mut self, _record: LrSchedulerRecord) {}
}

#[test]
fn partial_accumulation_window_is_dropped() {
    let device = Device::flex().autodiff();
    let model = ToyModel::new(&device);
    let before = model.weight.val().try_into_vec_as::<f32>().unwrap();

    let optim = burn_optim::SgdConfig::new().init();
    let calls = Arc::new(AtomicUsize::new(0));
    let scheduler = CountingScheduler {
        calls: calls.clone(),
    };
    let learner = burn_train::Learner::new(model, optim, scheduler);

    // 4 items / batch size 2 = 2 training batches.
    let (dl_train, dl_valid) = make_dataloaders();
    let dir = tempfile::tempdir().unwrap();

    let result = SupervisedTraining::new(dir.path(), dl_train, dl_valid)
        .num_epochs(1)
        .grads_accumulation(3)
        .with_metric_logger(InMemoryMetricLogger::new())
        .with_application_logger(None)
        .launch(learner);

    let after = result
        .model
        .weight
        .val()
        .try_into_vec_as::<f32>()
        .unwrap();

    // Both assertions currently pass:
    assert_eq!(before, after, "both batches' gradients were discarded");
    assert_eq!(calls.load(Ordering::SeqCst), 2);
}

Run with:

cargo test -p burn-train --test grad_accum_repro --no-default-features -- --nocapture

Observed: the test passes. Both batches are reported as processed, the model is unchanged, and the scheduler advances twice despite there being zero optimizer updates.

Expected behavior

At natural epoch completion, a non-empty partial accumulation window should produce one final optimizer update. The scheduler should advance exactly once per optimizer update, including that final update.

For the reproduction above:

• The model weights should change.
• There should be one optimizer update.
• The scheduler should advance once.

Interrupted epochs and dataloader errors should not necessarily flush; the final flush can be limited to natural exhaustion.

Minimal fix

• Validate that the accumulation interval is greater than zero.
• Move lr_step() immediately before actual optimizer calls.
• After natural dataloader exhaustion, drain and apply accumulated gradients when accumulation_current > 0.
• Apply the same behavior to single-device, both multi-device optimizer modes, and DDP.