Repository: pykeen/pykeen
Issue: Question on the OOM issue and accelerating the model
Root cause: Architectural Constraint

### What is your question

Hello, 

I'm trying to use Pykeen for training a large graph that contains about 85m entities/relations. 
My dimension is 10 and ideally, it should takes about 4 * 85m * 10 = 3,400,000,000, nearly 3G of storage. 
My RAM is 250G and I have 20G in GPU.
I've tried to reduce the dimension and the batch size. 
However, I always got this return after it loads about 200G in RAM: 
```python-traceback
Traceback (most recent call last):
  File "Pykeen.py", line 16, in <module>
    embedding_dim=10,
  File "/dataZ/pykeen-master/src/pykeen/pipeline/api.py", line 1232, in pipeline
    **trainils: not found
  File "/dataZ/pykeen-master/src/pykeen/training/training_loop.py", line 375, in train
    pin_memory=pin_memory,
  File "/dataZ/pykeen-master/src/pykeen/training/training_loop.py", line 501, in _train
    training_instances=training_instances,
  File "/dataZ/pykeen-master/src/pykeen/training/training_loop.py", line 930, in sub_batch_and_slice
     supports_sub_batching=supports_sub_batching,
  File "/dataZ/pykeen-master/src/pykeen/training/slcwa.py", line 123, in _slice_size_search
     raise MemoryError("The current model can't be trained on this hardware with these parameters.")
MmoryError: The current model can't be trained on this hardware with these parameters.
```
 My code is like this:
 
```python
from pykeen.triples import TriplesFactory
from pykeen.pipeline import pipeline
from pykeen.datasets.nations import NATIONS_TRAIN_PATH, NATIONS_TEST_PATH

training = '/dataZ/wikidata/trainT.txt'
testing = '/dataZ/wikidata/testT.txt'

result = pipeline(
    model='TransE',
    training=training,      # TRAIN_PATH
    testing=testing,        # TEST_PATH
    validation=validation,
    training_kwargs=dict(batch_size=16),
    model_kwargs=dict(
        embedding_dim=10,
    ),
)
```

Am I making some mistakes on the code or it's just impossible to run on such a large-scaled dataset?

Thanks in advanced for your answer :)

 

### Environment

| Key             | Value                       |
|-----------------|-----------------------------|
| OS              | posix                       |
| Platform        | Linux                       |
| Release         | 3.10.0-1160.36.2.el7.x86_64 |
| Time            | Mon Feb 28 05:14:43 2022    |
| Python          | 3.7.11                      |
| PyKEEN          | 1.7.1-dev                   |
| PyKEEN Hash     | UNHASHED                    |
| PyKEEN Branch   |                             |
| PyTorch         | 1.10.2+cu113                |
| CUDA Available? | true                        |
| CUDA Version    | 11.3                        |
| cuDNN Version   | 8200                        |

### Issue Template Checks

- [X] This is not a bug report (use a different issue template if it is)
- [X] This is not a feature request (use a different issue template if it is)
- [X] I've read the text explaining why including environment information is important and understand if I omit this information that my issue will be dismissed

Update:

I've re-configurate the dimension to 5 and the batch size to 8.
Now it works. 
However, it will need more than 100 days to finish the training.
Is there some configurations to accelerate it?

```
Training epochs on cuda:0:   0%|                        | 0/5 [00:00<?, ?epoch/s]
Training batches on cuda:0:   0%|                         | 62458/61276953 [2:29:58<2431:28:00,  6.99batch/s]
```

Hi @Yansera ,

would you mind to share the full traceback of the `MemoryError` you received, as well as the size of your triples factories?

By default, `pykeen` uses the filtered rank-based evaluation, in which the scores of true triples known from training are excluded in evaluation. To this end, it moves all known true triples to GPU here
https://github.com/pykeen/pykeen/blob/ef456cde3f02c52ad46d9d6deede2b61167d37e3/src/pykeen/evaluation/evaluator.py#L638-L642

This should consume around `3 * num_training_triples + num_validation_triples * 8` additional bytes (`*8` since torch's default index type is `torch.long`).


>  Is there some configurations to accelerate it?

You should be able to significantly increase the batch-size. Since you are using the SLCWA training loop, each batch will compute `batch_size * (1 + num_negatives)` scores. The additional memory requirements for this should be negligible compared to the huge amount of memory required to store the entity embeddings.

Since you are dealing with a large graph, the [NodePiece](https://github.com/pykeen/pykeen/blob/master/src/pykeen/models/unimodal/node_piece.py) model may be of interest to you, which can greatly reduce the number of trainable parameters. It is also used by the top entries of the [OGB WikiKG2 leaderboard](https://ogb.stanford.edu/docs/leader_linkprop/#ogbl-wikikg2) - a dataset which was extracted from Wikidata, too. You can find the paper [here](https://arxiv.org/abs/2106.12144).

Thanks for your feedback!
I will look at that model!

My task is to compare different models, hence, I also need to training more models, that's the problem :)

And for the error traceback, I've got this:

```
# python3 Pykeen.py
No random seed is specified. Setting to 843812282.
INFO:pykeen.training.training_loop:Starting sub_batch_size search for training now...
INFO:pykeen.training.training_loop:Even sub_batch_size=1 does not fit in memory with these parameters
WARNING:pykeen.training.slcwa:This model supports sub-batching, but it also requires slicing, which is not possible for sLCWA
Traceback (most recent call last):
  File "Pykeen.py", line 16, in <module>
    embedding_dim=10,
  File "/dataZ/pykeen-master/src/pykeen/pipeline/api.py", line 1232, in pipeline
    **trainils: not found
  File "/dataZ/pykeen-master/src/pykeen/training/training_loop.py", line 375, in train
    pin_memory=pin_memory,
  File "/dataZ/pykeen-master/src/pykeen/training/training_loop.py", line 501, in _train
    training_instances=training_instances,
  File "/dataZ/pykeen-master/src/pykeen/training/training_loop.py", line 930, in sub_batch_and_slice
     supports_sub_batching=supports_sub_batching,
  File "/dataZ/pykeen-master/src/pykeen/training/slcwa.py", line 123, in _slice_size_search
     raise MemoryError("The current model can't be trained on this hardware with these parameters.")
MmoryError: The current model can't be trained on this hardware with these parameters.
```

@Yansera could you try the following snippet? 

```python
from pykeen.pipeline import pipeline

if __name__ == "__main__":
    training = "/dataZ/wikidata/trainT.txt"
    testing = "/dataZ/wikidata/testT.txt"

    result = pipeline(
        interaction="TransE",
        interaction_kwargs=dict(p=2),
        dimensions=dict(d=10),
        training=training,  # TRAIN_PATH
        testing=testing,  # TEST_PATH
        training_kwargs=dict(batch_size=16),
    )
```

TL;DR: These kind of models should always support slicing. Since you are already using SLCWA, I am not sure whether this will actually help you, but it's worth a try.

_Long Version_:
When training with the given batch-size fails due to memory reasons, PyKEEN tries the following two techniques to perform an equivalent training in terms of the computed gradients, and parameter updates, but trying to trade an increased runtime against smaller memory footprint (somewhat similar to what, e.g., huggingface does with the explicitly set [Gradient Accumulation](https://huggingface.co/docs/transformers/master/en/main_classes/deepspeed#gradient-accumulation).).
The following two tricks are performed:

1. _Sub-Batching_: this means that we divide the batch into smaller sub-batches along the batch-dimension. Then each sub-batch is treated as a "normal" batch, doing a forward pass to compute a loss, and a backward pass to compute the gradient. However, we do _not_ call the optimizer's `step` method, i.e., the parameter update, until we have accumulated the gradients of each sub-batch corresponding to a single batch.

2. _Slicing_: when sub-batching fails, i.e., we still get out-of-memory errors when training with one batch-element at a time, we continue with _slicing_. Note that typically we compute multiple scores for each (positive) batch element, e.g., in LCWA we have `num_entities_ many tail scores, and for SLCWA, we have `1 + num_negatives` many scores. In slicing, we apply the same trick as sub-batching, just to another dimension of the tensor. Not all old-style models, i.e., those subclassing from [_OldAbstractModel](https://pykeen.readthedocs.io/en/stable/api/pykeen.models._OldAbstractModel.html#pykeen.models._OldAbstractModel), do support this, e.g., TransE doesn't seem to do. With new-style models, i.e,, those subclassing from [ERModel](https://pykeen.readthedocs.io/en/stable/api/pykeen.models.ERModel.html#pykeen.models.ERModel), we have added support of slicing to _any_ interaction function. For each interaction function, where the default model has not been ported to new-style yet, e.g., TransE, you can still explicitly construct an equivalent new-style model by using the `interaction` parameter of the `pipeline` instead of `model`. Typically, you may also have to provide some additional shape information baked into the old-style models, cf. `dimensions` in the example above.

Thanks for the entire explanation and the proposition! :) 
I will take the time to read your recommendations.  
I've tried the new script and it always does not work with the same error traceback. Sorry for that