def test_ddp_fp16_compress_wrap_sgd_comm_hook(tmp_path):
    """Test for DDP FP16 compress wrapper for SGD hook."""
    model = BoringModel()
    strategy = TestDDPStrategy(
        expected_ddp_comm_hook_name=default.fp16_compress_wrapper(powerSGD.powerSGD_hook).__qualname__,
        ddp_comm_state=powerSGD.PowerSGDState(process_group=None),
        ddp_comm_hook=powerSGD.powerSGD_hook,
        ddp_comm_wrapper=default.fp16_compress_wrapper,
    )
    trainer = Trainer(
        max_epochs=1,
        accelerator="gpu",
        devices=2,
        strategy=strategy,
        default_root_dir=tmp_path,
        sync_batchnorm=True,
        fast_dev_run=True,
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    trainer.fit(model)
    assert trainer.state.finished, f"Training failed with {trainer.state}"
Test for DDP FP16 compress wrapper for SGD hook.
tests/tests_pytorch/strategies/test_ddp_integration_comm_hook.py