Repository: pykeen/pykeen
Issue: Schlichtkrull sampler fails on large dataset due to upstream PyTorch size restrictions on multinomials
Root cause: External Dependency Breakage

### What is your question

Here is my current code

```python
import os
import torch
from pykeen.models import ConvE, RGCN
from pykeen.training import SLCWATrainingLoop, LCWATrainingLoop
from torch.optim import Adam
from pykeen.datasets import ConceptNet
from pykeen.evaluation import RankBasedEvaluator

device = torch.device("cuda:1")
LOG = "jotham_log.txt"
with open(LOG, "w") as log:
    print("Loading dataset...")
    print("Loading dataset...\n", file=log)
    dataset = ConceptNet()
    triples_factory = dataset.training
    print("Loading model...")
    print("Loading model...\n", file=log)
    model = RGCN(triples_factory=triples_factory,
                   embedding_dim=30,
                   random_seed=42
                  ).to(device)
    optimizer = Adam(params=model.get_grad_params())
    print("Instantiating Stochastic Local Closed World Assumption Training Loop...")
    print("Instantiating Stochastic Local Closed World Assumption Training Loop...\n", file=log)
    training_loop = LCWATrainingLoop(model=model,
                                      optimizer=optimizer,
                                      triples_factory=triples_factory,)
    print("Starting training...")
    print("Starting training...\n", file=log)
    losses = training_loop.train(
        triples_factory=triples_factory,
        num_epochs=200,
        batch_size=128,
        checkpoint_name="jotham/jotham_checkpoint.pt",
        checkpoint_frequency=5,
        checkpoint_on_failure=True,
        sampler="schlichtkrull",
        use_tqdm=True
    )
    PATH = "jotham_cc_kge.pth"
    print(f"Saving model to {PATH}    ...")
    print(f"Saving model to {PATH}    ...\n", file=log)
    torch.save(model.state_dict(), PATH)
    print("Evaluating model...")
    print("Evaluating model...\n", file=log)
    evaluator = RankBasedEvaluator()
    test_triples = dataset.testing.mapped_triples
    results = evaluator.evaluate(
        model=model,
        mapped_triples=test_triples,
        batch_size=128,
        additional_filter_triples=[
            dataset.training.mapped_triples,
            dataset.validation.mapped_triples
        ]
    )
    print("Writing results...")
    print("Writing results...\n", file=log)
    with open("results.txt", "w", encoding="utf-8") as f:
        print(results, file=f)
```

The issue that I get is number of categories cannot exceed 2^24 for the schlichtkrull_sampler. How am I supposed to train on ConceptNet then?

### Environment

 Key             | Value                    |
|-----------------|--------------------------|
| OS              | posix                    |
| Platform        | Linux                    |
| Release         | 4.15.0-55-generic        |
| Time            | Fri Dec 17 16:55:40 2021 |
| Python          | 3.8.8                    |
| PyKEEN          | 1.6.0                    |
| PyKEEN Hash     | UNHASHED                 |
| PyKEEN Branch   |                          |
| PyTorch         | 1.10.0+cu102             |
| CUDA Available? | true                     |
| CUDA Version    | 10.2                     |
| cuDNN Version   | 7605                     |


### 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

Hi @JothamWong ,

would you mind sharing the exact error message, too? If the problem comes from the Schlichtkrull sampler, the following code should also produce the error (while requiring less resources to execute)
```python
from pykeen.pipeline import pipeline
import logging

logging.basicConfig(level=logging.DEBUG)

result = pipeline(
    dataset="conceptnet",
    model="rgcn",
    model_kwargs=dict(
        embedding_dim=1,
    ),
    random_seed=42,
    training_loop="lcwa",
    training_kwargs=dict(
        num_epochs=1,
        batch_size=1,
        sampler="schlichtkrull",
    ),
)
```

Hello, 
here is the exact error message. I have used the code that you attached.

```python-traceback
Traceback (most recent call last):
  File "main2.py", line 6, in <module>
    result = pipeline(
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/pipeline/api.py", line 1068, in pipeline
    losses = training_loop_instance.train(
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/training_loop.py", line 319, in train
    result = self._train(
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/training_loop.py", line 494, in _train
    sub_batch_size, slice_size = self.sub_batch_and_slice(
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/training_loop.py", line 903, in sub_batch_and_slice
    sub_batch_size, finished_search, supports_sub_batching = self._sub_batch_size_search(
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/training_loop.py", line 994, in _sub_batch_size_search
    raise runtime_error
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/training_loop.py", line 982, in _sub_batch_size_search
    self._train(
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/training_loop.py", line 696, in _train
    raise e
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/training_loop.py", line 609, in _train
    for batch in batches:
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 560, in _next_data
    index = self._next_index()  # may raise StopIteration
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 512, in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/torch/utils/data/sampler.py", line 229, in __iter__
    for idx in self.sampler:
  File "/storage/jotham/anaconda3/lib/python3.8/site-packages/pykeen/training/schlichtkrull_sampler.py", line 105, in __iter__
    chosen_vertex = torch.multinomial(probabilities, num_samples=1)[0]
RuntimeError: number of categories cannot exceed 2^24
```

This seems to be https://github.com/pytorch/pytorch/issues/2576, which is still open on `torch` side.

Do you need to use `RGCN`, or are you just interested in training some model on `ConceptNet`?

I am interested in training some model on ConceptNet. I tried TransE and ConvE but I ran into some issues related to the model cannot support sub-batching or slicing. What model would be able to train on ConceptNet then? Thanks for the help

All models subclass from the new-style base class `ERModel` have automatic support for sub-batching / slicing. some of the more scalable models (all should be in `O((E + R)d)`) include
* ComplEx
* DistMA
* MuRE
* PairRE