Repository: pykeen/pykeen
Issue: How to interpret the state of my GPU cluster when I get a CUDA OOM error?
Root cause: Configuration Error

I run into a out-of-memory issue when I try a simple training on the sample Nations dataset on our on-premise GPU cluster.

 

I have carefully checked the GPU memory usage before I start the training. I choose cuda:2, which in principle has a GPU Memory of 11019MiB.

 

The code I attach can run, but will throw the out-of-memory error during the training:

![image](https://user-images.githubusercontent.com/16422009/158821205-345488da-9e07-4d07-9949-573d7e50aca6.png)

![image](https://user-images.githubusercontent.com/16422009/158821271-4403cd71-bec8-4c7e-bb3c-d42b1bcc6a71.png)

 



 

 

I have tried to change the embedding_dim to a smaller number, it will postpone the error, for example Epoch 49. But since this is a really small problem, I wonder why there can be a memory issue ?

 

Is there something wrong with my configuration ?

### How to reproduce
```python
from pykeen.triples import TriplesFactory
from pykeen.pipeline import pipeline
import torch

result = pipeline(
    dataset='Nations',
    model='TransE',
    model_kwargs=dict(embedding_dim=100),
    stopper='early',
    evaluator='RankBasedEvaluator',
    training_kwargs=dict(num_epochs=50, batch_size=32),
    optimizer_kwargs=dict(lr=1.0e-2),
    result_tracker='csv',
    result_tracker_kwargs=dict(
        name='transE.csv',
    ),
    device=torch.device("cuda:2" if torch.cuda.is_available() else "cpu"),
)
```

### Environment

| Key             | Value                    |
|-----------------|--------------------------|
| OS              | posix                    |
| Platform        | Linux                    |
| Release         | 5.11.0-37-generic        |
| Time            | Thu Mar 17 10:16:20 2022 |
| Python          | 3.8.12                   |
| PyKEEN          | 1.7.0                    |
| PyKEEN Hash     | UNHASHED                 |
| PyKEEN Branch   |                          |
| PyTorch         | 1.10.1                   |
| CUDA Available? | true                     |
| CUDA Version    | 11.3                     |
| cuDNN Version   | 8200                     |
 

Some information of GPU cluster
```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.63.01    Driver Version: 470.63.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:04:00.0 Off |                  N/A |
| 83%   70C    P2   243W / 250W |  10230MiB / 11019MiB |     91%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce ...  Off  | 00000000:05:00.0 Off |                  N/A |
|121%   84C    P2   228W / 250W |   8465MiB / 11019MiB |     89%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  NVIDIA GeForce ...  Off  | 00000000:08:00.0 Off |                  N/A |
| 39%   37C    P8     6W / 250W |      8MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  NVIDIA GeForce ...  Off  | 00000000:09:00.0 Off |                  N/A |
| 37%   34C    P8     3W / 250W |   4901MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   4  NVIDIA GeForce ...  Off  | 00000000:85:00.0 Off |                  N/A |
|101%   88C    P2   194W / 250W |  10823MiB / 11019MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   5  NVIDIA GeForce ...  Off  | 00000000:86:00.0 Off |                  N/A |
| 56%   75C    P2   177W / 250W |  10823MiB / 11019MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   6  NVIDIA GeForce ...  Off  | 00000000:89:00.0 Off |                  N/A |
| 41%   34C    P8    18W / 250W |   1783MiB / 11019MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   7  NVIDIA GeForce ...  Off  | 00000000:8A:00.0 Off |                  N/A |
|135%   83C    P2   248W / 250W |  10823MiB / 11019MiB |    100%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
```

### Additional information

_No response_

For some reason a small amount of memory is reserved on cuda:0 as well. In your nvidia-smi output most other devices are already in heavy use, so I suspect that is why it fails.

> 

@Koenkalle yes, you are totally right. I can confirm that if memory of CUDA:0 is enough. The error is gone. 

Actually I notice that pykeen creates two processes, one on CUDA:2, one CUDA:0 (seems only created for evaluation), which is the cause of the out-of-memory error.


> Actually I notice that pykeen creates two processes, one on CUDA:2, one CUDA:0 (seems only created for evaluation), which is the cause of the out-of-memory error.

The latter sounds like a bug in the device resolution during evaluation.

As a hotfix, you can set the environment variable `CUDA_VISIBLE_DEVICES` to the GPU you want to use, e.g., 

```console
CUDA_VISIBLE_DEVICES=2 python3 my_script.py
```

In this case, you would also need to adapt your code to use `cuda:0`, or simply `cuda`, cf. https://pytorch.org/docs/stable/tensor_attributes.html#torch.torch.device.

You can also pass `device=None` (which is also the default), in which case PyKEEN will select `cuda` if a GPU is available, and `cpu` otherwise.