def test_deepspeed_custom_activation_checkpointing_params_forwarded(tmp_path):
    """Ensure if we modify the activation checkpointing parameters, we pass these to deepspeed.checkpointing.configure
    correctly."""
    import deepspeed

    ds = DeepSpeedStrategy(
        partition_activations=True,
        cpu_checkpointing=True,
        contiguous_memory_optimization=True,
        synchronize_checkpoint_boundary=True,
    )

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        fast_dev_run=1,
        strategy=ds,
        precision="16-mixed",
        accelerator="gpu",
        devices=1,
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    with mock.patch(
        "deepspeed.checkpointing.configure", wraps=deepspeed.checkpointing.configure
    ) as deepspeed_checkpointing_configure:
        trainer.fit(model)

    deepspeed_checkpointing_configure.assert_called_with(
        mpu_=None, partition_activations=True, contiguous_checkpointing=True, checkpoint_in_cpu=True, profile=None
    )
Ensure if we modify the activation checkpointing parameters, we pass these to deepspeed.checkpointing.configure
    correctly.
tests/tests_pytorch/strategies/test_deepspeed.py