Repository: Lightning-AI/pytorch-lightning
Pull request: Update Fabric CPU tests to work on GPU machines
## What does this PR do?

Fix CPU tests to work on GPU machines
Fixes #17371 

On testing on multi-GPU machines (the CI included) fails when running tests in `tests/tests_fabric` which were intended to run on the CPU due to the new `devices="auto"`  and `accelerator="auto"` features in versions `2.0+`
This PR changes the tests initialization of Fabric in these failing tests

```python
fabric = Fabric()
# changes to
fabric = Fabric(devices=1)
```
which fixes these tests and prevents the CI pipeline from hanging.

<details><summary>parity test</summary>
<p>
EDIT: This has been resolved by using `torch.use_deterministic_algorithms(True, warn_only=True)`


To fix the `tests/tests_fabric/parity/test_parity_simple.py::test_parity_single_device` I have marked it to run as a standalone. It seems that if the `make_deterministic` function
https://github.com/Lightning-AI/lightning/blob/97a61868fb5ff605a857b266a9b2a44b2330da71/tests/tests_fabric/parity/utils.py#L39-L44
is called after a previous model has been run without being deterministic the updated os.environ["CUBLAS_WORKSPACE_CONFIG"] variable is not read by torch / cuda.

```python
import os
import torch 
import torch.nn as nn

def make_deterministic():
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(1)
    torch.cuda.manual_seed(1)


layer = nn.Linear(32, 64).to("cuda")
data = torch.rand(32).to("cuda")
output = layer(data)

make_deterministic()  # setting deterministic algorithms after Torch operations causes error

output = layer(data)
```

</p>
</details> 

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- [x] Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
- Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


Problem statement:
Update Fabric tests with devices=1 where necessary
### Outline & Motivation

In 2.0, the defaults for accelerator, devices, strategy were all switched to "auto". This causes a problem for many tests that instantiate fabric like so

```py
fabric = Fabric()
```

and previously were assuming this would run on the CPU (devices=1). However, after 2.0, if you run this "cpu" test on a machine with one or multiple GPUs, lots are failing because of the auto setting switching to multi-gpu ddp. Another effect is that all tests that experinence this issue and are not marked as standalone will lead to the test suite hanging. 

It is annoying that we can't run CPU tests on a GPU machine.

### Pitch

Update all tests that are meant to run on the CPU with devices=1.

### Additional context

_No response_

cc @justusschock @awaelchli @carmocca @borda

Hints from issue discussion:
Hi @awaelchli, I can do this.
On my computer there is 1 failing in `tests/test_fabric`
```
FAILED tests/tests_fabric/test_fabric.py::test_backward_model_input_required - 
  RuntimeError: Expected all tensors to be on the same device, but found at least two devices, 
  cuda:0 and cpu! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
```

Can you clarify, is it to change all the tests that setup as `fabric = Fabric()` such that they always run on CPU or that they will make use of the GPU if it is there anyways?
eg. in the above failed test, the offending line is 
```python
  loss = model0(torch.randn(1, 1)).sum()
```
which can be fixed by passing the current device to the tensor, forcing the fabric device as CPU or marking the accelerator to run both if available?
```python
# run on current device
  loss = model0(torch.randn(1, 1, device=fabric.device)).sum()

# OR  run on CPU only
  fabric = Fabric(accelerator="cpu")

# OR  set the accelerator to run both if available
@pytest.mark.parametrize("accelerator", ["cpu", pytest.param("cuda", marks=RunIf(min_cuda_gpus=1))])
...
  loss = model0(torch.randn(1, 1, device=fabric.device)).sum()
```

pro-tip: If you comment this line https://github.com/Lightning-AI/lightning/blob/97a61868fb5ff605a857b266a9b2a44b2330da71/tests/tests_fabric/conftest.py#L137 and open a PR, GPU CI will run all tests (not just those marked with `@RunIF(min_cuda_gpus=...)`

This should uncover all the failing tests without having a GPU locally.
I do have a local GPU, so it's actually confusing me why those tests do not fail for me at all (they are not being skipped). Is it that they are failing on multi-GPU machines?
To see more tests fail, you need to run on a machine with two GPUs (is what I did). @carmocca 's trick will reveal that in our CI too, it runs on 2 Gpus

Commit message: standard tests all passing, set min gpus

Patch to review:
diff --git a/tests/tests_fabric/parity/test_parity_simple.py b/tests/tests_fabric/parity/test_parity_simple.py
index 1e262d98f8cab..be070d45c92ec 100644
--- a/tests/tests_fabric/parity/test_parity_simple.py
+++ b/tests/tests_fabric/parity/test_parity_simple.py
@@ -118,10 +118,10 @@ def train_fabric(fabric):
     "precision, accelerator",
     [
         (32, "cpu"),
-        pytest.param(32, "cuda", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param(32, "cuda", marks=RunIf(min_cuda_gpus=1, standalone=True)),
         # pytest.param(16, "cuda", marks=RunIf(min_cuda_gpus=1)),  # TODO: requires GradScaler
         pytest.param("bf16", "cpu", marks=RunIf(skip_windows=True)),
-        pytest.param("bf16", "cuda", marks=RunIf(min_cuda_gpus=1, bf16_cuda=True)),
+        pytest.param("bf16", "cuda", marks=RunIf(min_cuda_gpus=1, bf16_cuda=True, standalone=True)),
         pytest.param(32, "mps", marks=RunIf(mps=True)),
     ],
 )
diff --git a/tests/tests_fabric/plugins/precision/test_double_integration.py b/tests/tests_fabric/plugins/precision/test_double_integration.py
index a026cb89547a1..cfd0918af8ed0 100644
--- a/tests/tests_fabric/plugins/precision/test_double_integration.py
+++ b/tests/tests_fabric/plugins/precision/test_double_integration.py
@@ -51,5 +51,5 @@ def after_backward(self, model, optimizer):
 
 
 def test_double_precision():
-    fabric = DoublePrecisionBoringFabric(precision="64-true")
+    fabric = DoublePrecisionBoringFabric(devices=1, precision="64-true")
     fabric.run()
diff --git a/tests/tests_fabric/test_fabric.py b/tests/tests_fabric/test_fabric.py
index db9868fa35d1c..43cb90ff8128e 100644
--- a/tests/tests_fabric/test_fabric.py
+++ b/tests/tests_fabric/test_fabric.py
@@ -92,7 +92,7 @@ def test_setup_compiled_module(setup_method):
     """Test that an `OptimizedModule` can be passed to the setup method."""
     from torch._dynamo.eval_frame import OptimizedModule
 
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     model = nn.Linear(1, 2)
     compiled_model = torch.compile(model)
     assert isinstance(compiled_model, OptimizedModule)
@@ -170,7 +170,7 @@ def test_setup_module_parameters_on_different_devices(setup_method, move_to_devi
 
 def test_setup_module_and_optimizers():
     """Test that `setup()` can handle no optimizers, one optimizer, or multiple optimizers."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     model = nn.Linear(1, 2)
     optimizer0 = torch.optim.SGD(model.parameters(), lr=0.1)
     optimizer1 = torch.optim.Adam(model.parameters(), lr=0.1)
@@ -219,7 +219,7 @@ def test_setup_optimizers():
 
 def test_setup_twice_fails():
     """Test that calling `setup` with a model or optimizer that is already wrapped fails."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     model = nn.Linear(1, 2)
     optimizer = torch.optim.Adam(model.parameters())
 
@@ -234,7 +234,7 @@ def test_setup_twice_fails():
 
 def test_setup_module_twice_fails():
     """Test that calling `setup_module` with a model that is already wrapped fails."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     model = nn.Linear(1, 2)
 
     fabric_model = fabric.setup_module(model)
@@ -267,7 +267,7 @@ def test_setup_optimizers_not_supported(strategy_cls):
 
 def test_setup_tracks_num_models():
     """Test that setup() tracks how many times it has setup a model."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     model = nn.Linear(1, 2)
     optimizer = torch.optim.Adam(model.parameters())
 
@@ -293,7 +293,7 @@ def test_setup_dataloaders_unsupported_input():
 
 def test_setup_dataloaders_return_type():
     """Test that the setup method returns the dataloaders wrapped as FabricDataLoader and in the right order."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
 
     # single dataloader
     fabric_dataloader = fabric.setup_dataloaders(DataLoader(range(2)))
@@ -363,12 +363,12 @@ def test_setup_dataloaders_twice_fails():
 )
 def test_setup_dataloaders_move_to_device(fabric_device_mock):
     """Test that the setup configures FabricDataLoader to move the data to the device automatically."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     fabric_dataloaders = fabric.setup_dataloaders(DataLoader(Mock()), DataLoader(Mock()), move_to_device=False)
     assert all(dl.device is None for dl in fabric_dataloaders)
     fabric_device_mock.assert_not_called()
 
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     fabric_dataloaders = fabric.setup_dataloaders(DataLoader(Mock()), DataLoader(Mock()), move_to_device=True)
     assert all(dl.device == torch.device("cuda", 1) for dl in fabric_dataloaders)
     fabric_device_mock.assert_called()
@@ -380,7 +380,7 @@ def test_setup_dataloaders_distributed_sampler_not_needed():
     dataloader = DataLoader(Mock(), sampler=custom_sampler)
 
     # keep the custom sampler when not needed to replace
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     fabric_dataloader = fabric.setup_dataloaders(dataloader, use_distributed_sampler=True)
     assert fabric_dataloader.sampler is custom_sampler
 
@@ -455,7 +455,7 @@ def test_seed_everything():
     """Test that seed everything is static and sets the worker init function on the dataloader."""
     Fabric.seed_everything(3)
 
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     fabric_dataloader = fabric.setup_dataloaders(DataLoader(Mock()))
 
     assert fabric_dataloader.worker_init_fn.func is pl_worker_init_function
@@ -585,7 +585,7 @@ def test_backward_model_input_required():
     fabric.setup(model0, optimizer0)
     fabric.setup(model1, optimizer1)
 
-    loss = model0(torch.randn(1, 1)).sum()
+    loss = model0(torch.randn(1, 1, device=fabric.device)).sum()
 
     with pytest.raises(ValueError, match="please provide the model used to perform"):
         fabric.backward(loss)
@@ -604,7 +604,7 @@ def test_autocast():
 
 def test_no_backward_sync():
     """Test that `Fabric.no_backward_sync()` validates the strategy and model is compatible."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     model = nn.Linear(3, 3)
     with pytest.raises(TypeError, match="You need to set up the model first"):
         with fabric.no_backward_sync(model):
@@ -673,7 +673,7 @@ def fn_with_one_arg(arg):
 
 @mock.patch.dict(os.environ, {"LT_CLI_USED": "1"})  # pretend we are using the CLI
 def test_launch_and_cli_not_allowed():
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     with pytest.raises(RuntimeError, match=escape("Calling  `.launch()` again is not allowed")):
         fabric.launch()
 
@@ -829,7 +829,7 @@ def test_log_dict_input_parsing():
 @pytest.mark.parametrize("setup", [True, False])
 def test_save_wrapped_objects(setup, tmp_path):
     """Test that when modules and optimizers are in the state, they get unwrapped properly."""
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
     save_checkpoint_mock = Mock()
     fabric.strategy.save_checkpoint = save_checkpoint_mock
 
@@ -910,7 +910,7 @@ def test_all_reduce():
 
 @pytest.mark.parametrize("clip_val,max_norm", [(1e-3, None), (None, 1)])
 def test_grad_clipping(clip_val, max_norm):
-    fabric = Fabric()
+    fabric = Fabric(devices=1)
 
     fabric.strategy.clip_gradients_norm = Mock()
     fabric.strategy.clip_gradients_value = Mock()


Code review comments:
diff_hunk: @@ -118,10 +118,10 @@ def train_fabric(fabric):
     "precision, accelerator",
     [
         (32, "cpu"),
-        pytest.param(32, "cuda", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param(32, "cuda", marks=RunIf(min_cuda_gpus=1, standalone=True)),
path: tests/tests_fabric/parity/test_parity_simple.py
text: @user1:
these don't need to be marked as standalone. or do you see a good reason?
```suggestion
        pytest.param(32, "cuda", marks=RunIf(min_cuda_gpus=1)),
```

@author:
@user1 I put a reason for it at the end of the PR post. The tests fail because the `CUBLAS_WORKSPACE_CONFIG` needs to be set at the start of the python session.  I included a small test to show this behavior.

@user2:
In that case It would be better to address this environment variable leak. Standalone tests are considerably slower to run so it's better to avoid them if possible

@author:
I don't think it's really a leak, it happens because the variable isn't set at the beginning of the python run, and the error states it must be before any PyTorch applications.
```
RuntimeError: Deterministic behavior was enabled with either `torch.use_deterministic_algorithms(True)`
 or `at::Context::setDeterministicAlgorithms(true)`, but this operation is not deterministic 
 because it uses CuBLAS and you have CUDA >= 10.2. To enable deterministic behavior in this case, 
 you must set an environment variable before running your PyTorch application:
 CUBLAS_WORKSPACE_CONFIG=:4096:8 or CUBLAS_WORKSPACE_CONFIG=:16:8. 
 For more information, go to https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
``` 
I have reduced the reproducible error down to this
```python
import os
import torch 
import torch.nn as nn

def make_deterministic():
    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.benchmark = False
    torch.manual_seed(1)
    torch.cuda.manual_seed(1)


layer = nn.Linear(32, 64).to("cuda")
data = torch.rand(32).to("cuda")
output = layer(data)

make_deterministic()  # setting deterministic algorithms after Torch operations causes error

output = layer(data)
```
So the error occurs because there are previous tests which have run torch operations without setting the env variable.
Do you have suggestions for how I should solve this issue? Thanks @user1 @user2 

@author:
Actually if I use `torch.use_deterministic_algorithms(True, warn_only=True)` its all okay and the parity test still passes ~~so I guess this is a bug in Torch that it doesn't update its own context?~~ but CUDA does not register that the var was set (see below comment)

@author:
I think proper way to do it without setting the `warn_only=True` is to launch a subprocess

https://github.com/pytorch/pytorch/blob/e2923b521bd7fc67fa1e8f404397b20d1ec81668/torch/testing/_internal/common_utils.py#L1216-L1233
```
# Note that if a test requires CUDA to actually register the changed
# CUBLAS_WORKSPACE_CONFIG variable, a new subprocess must be created, because
# CUDA only checks the variable when the runtime initializes. Tests can be
# run inside a subprocess like so:
#
#   import subprocess, sys, os
#   script = '''
#   # Test code should go here
#   '''
#   try:
#       subprocess.check_output(
#           [sys.executable, '-c', script],
#           stderr=subprocess.STDOUT,
#           cwd=os.path.dirname(os.path.realpath(__file__)),
#           env=os.environ.copy())
#   except subprocess.CalledProcessError as e:
#       error_message = e.output.decode('utf-8')
#       # Handle exceptions raised by the subprocess here
```


@user1:
The current change in the latest commit LGTM too. 

> it happens because the variable isn't set at the beginning of the python run,

It would also be fine to move the make_deterministic() call to the beginning of `test_parity_single_device`. 

@author:
It would still have to be called twice anyways, because it resets the RNG seeds too?

diff_hunk: @@ -118,10 +118,10 @@ def train_fabric(fabric):
     "precision, accelerator",
     [
         (32, "cpu"),
-        pytest.param(32, "cuda", marks=RunIf(min_cuda_gpus=1)),
+        pytest.param(32, "cuda", marks=RunIf(min_cuda_gpus=1, standalone=True)),
         # pytest.param(16, "cuda", marks=RunIf(min_cuda_gpus=1)),  # TODO: requires GradScaler
         pytest.param("bf16", "cpu", marks=RunIf(skip_windows=True)),
-        pytest.param("bf16", "cuda", marks=RunIf(min_cuda_gpus=1, bf16_cuda=True)),
+        pytest.param("bf16", "cuda", marks=RunIf(min_cuda_gpus=1, bf16_cuda=True, standalone=True)),
path: tests/tests_fabric/parity/test_parity_simple.py
text: @user1:
```suggestion
        pytest.param("bf16", "cuda", marks=RunIf(min_cuda_gpus=1, bf16_cuda=True)),
```