def test_deepspeed_multigpu_stage_2_accumulated_grad_batches(tmp_path, offload_optimizer):
    """Test to ensure with Stage 2 and multiple GPUs, accumulated grad batches works."""

    class VerificationCallback(Callback):
        def __init__(self):
            self.on_train_batch_start_called = False

        def on_train_batch_start(self, trainer, pl_module: LightningModule, batch: Any, batch_idx: int) -> None:
            deepspeed_engine = trainer.strategy.model
            assert trainer.global_step == deepspeed_engine.global_steps
            self.on_train_batch_start_called = True

    model = ModelParallelClassificationModel()
    dm = ClassifDataModule()
    verification_callback = VerificationCallback()
    strategy = DeepSpeedStrategy(stage=2, offload_optimizer=offload_optimizer)
    strategy.config["zero_force_ds_cpu_optimizer"] = False
    trainer = Trainer(
        default_root_dir=tmp_path,
        # TODO: this test fails with max_epochs >1 as there are leftover batches per epoch.
        # there's divergence in how Lightning handles the last batch of the epoch with how DeepSpeed does it.
        # we step the optimizers on the last batch but DeepSpeed keeps the accumulation for the next epoch
        max_epochs=1,
        strategy=strategy,
        accelerator="gpu",
        devices=2,
        limit_train_batches=5,
        limit_val_batches=2,
        precision="16-mixed",
        accumulate_grad_batches=2,
        callbacks=[verification_callback],
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    assert trainer.limit_train_batches % trainer.accumulate_grad_batches != 0, "leftover batches should be tested"
    trainer.fit(model, datamodule=dm)
    assert verification_callback.on_train_batch_start_called
Test to ensure with Stage 2 and multiple GPUs, accumulated grad batches works.
tests/tests_pytorch/strategies/test_deepspeed.py