def test_post_local_sgd_model_averaging(average_parameters_mock, tmp_path):
    """Test that when using DDP with post-localSGD, model averaging is called."""
    model = BoringModel()

    # test regular ddp does not call model averaging
    trainer = Trainer(
        fast_dev_run=True,
        accelerator="gpu",
        devices=2,
        strategy="ddp",
        default_root_dir=tmp_path,
        sync_batchnorm=True,
        enable_progress_bar=False,
        enable_model_summary=False,
    )

    trainer.fit(model)
    average_parameters_mock.assert_not_called()

    # test ddp with post-localSGD does call model averaging
    ddp_strategy = DDPStrategy(
        ddp_comm_state=post_localSGD.PostLocalSGDState(
            process_group=None,
            subgroup=None,
            start_localSGD_iter=8,
        ),
        ddp_comm_hook=post_localSGD.post_localSGD_hook,
        model_averaging_period=4,
    )

    trainer = Trainer(
        fast_dev_run=True,
        accelerator="gpu",
        devices=2,
        strategy=ddp_strategy,
        default_root_dir=tmp_path,
        sync_batchnorm=True,
    )

    trainer.fit(model)
    average_parameters_mock.assert_called()
Test that when using DDP with post-localSGD, model averaging is called.
tests/tests_pytorch/strategies/test_ddp_integration_comm_hook.py