def test_deepspeed_summary(tmp_path):
    """Test to ensure that the summary contains the correct values when stage 3 is enabled and that the trainer enables
    the `DeepSpeedSummary` when DeepSpeed is used."""

    model = BoringModel()
    total_parameters = sum(x.numel() for x in model.parameters())

    class TestCallback(Callback):
        def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
            model_summary = DeepSpeedSummary(pl_module, max_depth=1)
            assert model_summary.total_parameters == total_parameters
            assert model_summary.trainable_parameters == total_parameters

            # check the additional params per device
            summary_data = model_summary._get_summary_data()
            params_per_device = summary_data[4][-1]
            assert int(params_per_device[0]) == (model_summary.total_parameters // 2)

    trainer = Trainer(
        strategy=DeepSpeedStrategy(stage=3),
        default_root_dir=tmp_path,
        accelerator="gpu",
        fast_dev_run=True,
        devices=2,
        precision="16-mixed",
        enable_model_summary=True,
        callbacks=[TestCallback()],
    )

    trainer.fit(model)
Test to ensure that the summary contains the correct values when stage 3 is enabled and that the trainer enables
    the `DeepSpeedSummary` when DeepSpeed is used.
tests/tests_pytorch/utilities/test_deepspeed_model_summary.py