fix deepspeed special test test_deepspeed_stage_3_save_warning

[awaelchli]

🐛 Bug

To Reproduce

Run on current master branch the following command:

SPECIAL_PATTERN=test_deepspeed_stage_3_save_warning bash tests/special_tests.sh

Error:

______________________________ test_deepspeed_stage_3_save_warning _______________________________

tmpdir = local('/tmp/pytest-of-adrian/pytest-1376/test_deepspeed_stage_3_save_wa0')

    @RunIf(min_gpus=2, deepspeed=True, special=True)
    def test_deepspeed_stage_3_save_warning(tmpdir):
        """
        Test to ensure that DeepSpeed Stage 3 gives a warning when saving.
        """
        model = BoringModel()
        trainer = Trainer(
            default_root_dir=tmpdir, plugins=[DeepSpeedPlugin(stage=3)], gpus=2, fast_dev_run=True, precision=16
        )
        trainer.fit(model)
        checkpoint_path = os.path.join(tmpdir, "model.pt")
        with pytest.warns(UserWarning, match="each worker will save a shard of the checkpoint within a directory."):
>           trainer.save_checkpoint(checkpoint_path)
E           Failed: DID NOT WARN. No warnings of type (<class 'UserWarning'>,) was emitted. The list of emitted warnings is: [].

tests/plugins/test_deepspeed_plugin.py:432: Failed

Expected behavior

Environment

• CUDA:
  - GPU:
  - GeForce RTX 3090
  - GeForce RTX 3090
  - GeForce RTX 3090
  - GeForce RTX 3090
  - GeForce RTX 3090
  - GeForce RTX 3090
  - GeForce RTX 3090
  - GeForce RTX 3090
  - available: True
  - version: 11.1
• Packages:
  - numpy: 1.20.2
  - pyTorch_debug: False
  - pyTorch_version: 1.9.0+cu111
  - pytorch-lightning: 1.5.0dev
  - tqdm: 4.60.0
• System:
  - OS: Linux
  - architecture:
  - 64bit
  - ELF
  - processor: x86_64
  - python: 3.8.8
  - version: added test model to do also #64-Ubuntu SMP Wed Dec 9 08:16:25 UTC 2020

Additional context

In our GPU CI as well, check section for "special tests"

[awaelchli]

TODO: git bisect

[carmocca]

TODO: git bisect

It's possible that the test itself failed in the original PR that added it (#8397). We don't know because the build log is not available.

The error code does not seem to get propagated, so special tests don't fail.