Repository: pykeen/pykeen
Issue: Automatic Memory Optimization can spawn batches of elements > INT_MAX
Root cause: Resource Mishandling

**Describe the bug**
Setup: FB15k237 (with inverse triples), Tesla V100, LCWA training loop

New GPUs have lots of VRAM - I'm testing on a V100 32GB.

Automatic memory optimization loop here
https://github.com/pykeen/pykeen/blob/dfc407eecf1f1960de4670e8f30e526779da5a5c/src/pykeen/evaluation/evaluator.py#L331-L335

quickly inflates the batch size to 1024 / 2048 / 4096. 
Now here is the problem: in the `create_sparse_positive_filter_()` there is a `torch.nonzero()` operator which can't process tensors with more elements than `INT_MAX` (2 billion something):
https://github.com/pykeen/pykeen/blob/dfc407eecf1f1960de4670e8f30e526779da5a5c/src/pykeen/evaluation/evaluator.py#L396-L398

Note that `entity_filter_test` has the shape `[batch_size, num_triples]`, and in our FB15k237 scenario the probed tensor of size `[4096, 544K]` fits onto GPU, but has >INT_MAX elements, so torch throws an error:

```
RuntimeError: nonzero is not supported for tensors with more than INT_MAX elements,   file a support request
```

I believe the same will appear when people will start putting larger datasets like YAGO3-10 on large GPUs with 1-N scoring.



**To Reproduce**
1. Have a 32GB+ GPU
2. Start a vanilla script with DistMult on FB15k 237 with automatic memory optimization
3. LCWA loop
4. [See error]

**Expected behavior**
Suggestions:
1. Hardcode a new ending constraint like `values_dict['batch_size'] * maximum_triples < INT_MAX`
2. Change the procedure with `torch.nonzero()` to something that can handle larger tensors

**Environment:**

```
pykeen              1.0.5
torch               1.7.1
```

**Additional information**
I actually stumbled upon this bug when found out that `RankBasedEvaluator` does not accept `batch_size` in the constructor, and hence in the early stopper:
https://github.com/pykeen/pykeen/blob/dfc407eecf1f1960de4670e8f30e526779da5a5c/src/pykeen/stoppers/early_stopping.py#L130-L131

The first eval iteration works fine (with the `EarlyStopper` defined batch size, not `Evaluator`), but then this line forces the early_stopper to have the `None` value, which, in turn, hardcodes evaluation batch_size to 1 in pykeen 1.0.5 and 32 in 1.1.0:

https://github.com/pykeen/pykeen/blob/dfc407eecf1f1960de4670e8f30e526779da5a5c/src/pykeen/evaluation/evaluator.py#L546-L550

I can't try it on 1.1.0 due to the blocking issue #261 
Still, it might deserve its own issue 🤔 




Hi @migalkin ,

thanks for reporting this issue! This is interesting, since this behaviour is not documented in PyTorch's documentation ([here](https://pytorch.org/docs/stable/generated/torch.nonzero.html)), and since the return type is a long tensor, also unexpected.

I am also tagging @lvermue who is most familiar with this code.

That limitation was recently acknowledged by pytorch here: https://discuss.pytorch.org/t/torch-geometric/106975 
So I presume it indeed exists 

I opened an issue on PyTorch: https://github.com/pytorch/pytorch/issues/51871 since I couldn't find any for it.

> I can't try it on 1.1.0 due to the blocking issue #261

I re-open this issue: When #261 is fixed, @migalkin may test this fix again on his GPU :slightly_smiling_face: 

@migalkin since #261 is fixed (with #270), would you please re-run your example with the current `master` on the V100 and let us know whether the problem was resolved?

Update: I'm in the `add-compat` branch , after solving the compatibility issue I'm running the same script with automatic memory optimization, and now it crashes during evaluation in a model forward pass (TransE in my case).

Trace:
```
File "...pykeen/src/pykeen/training/training_loop.py", line 308, in train
    drop_last=drop_last,
  File "...pykeen/src/pykeen/training/training_loop.py", line 608, in _train
    raise e
  File "...pykeen/src/pykeen/training/training_loop.py", line 595, in _train
    if stopper is not None and stopper.should_evaluate(epoch) and stopper.should_stop(epoch):
  File "...pykeen/src/pykeen/stoppers/early_stopping.py", line 128, in should_stop
    do_time_consuming_checks=self.evaluation_batch_size is None,
  File "...pykeen/src/pykeen/evaluation/evaluator.py", line 178, in evaluate
    do_time_consuming_checks=do_time_consuming_checks,
  File "...pykeen/src/pykeen/evaluation/evaluator.py", line 595, in evaluate
    filtering_necessary=filtering_necessary,
  File "...pykeen/src/pykeen/evaluation/evaluator.py", line 671, in _evaluate_batch
    batch_scores_of_corrupted = model.predict_h(batch[:, 1:3], slice_size=slice_size)
  File "...pykeen/src/pykeen/models/base.py", line 332, in predict_h
    scores = self.score_h_inverse(rt_batch=rt_batch, slice_size=slice_size)
  File "...pykeen/src/pykeen/models/base.py", line 542, in score_h_inverse
    return self.score_t(hr_batch=t_r_inv)
  File "...pykeen/src/pykeen/models/unimodal/trans_e.py", line 110, in score_t
    return -torch.norm(h[:, None, :] + r[:, None, :] - t[None, :, :], dim=-1, p=self.scoring_fct_norm)
RuntimeError: CUDA out of memory. Tried to allocate 22.13 GiB (GPU 0; 47.46 GiB total capacity; 37.59 GiB already allocated; 2.20 GiB free; 44.32 GiB reserved in total by PyTorch)
```

So apparently the batch size grew too big for a 48 GB GPU 🤔 


@migalkin Can you provide your script?

@lvermue  Sure, here it is

```python
from pykeen.models import TransE #DistMult, TransE
from pykeen.datasets import FB15k237, Kinships, WN18RR
from torch.optim import Adam
from pykeen.training import SLCWATrainingLoop, LCWATrainingLoop
from pykeen.evaluation import RankBasedEvaluator
from pykeen.trackers import WANDBResultTracker
from pykeen.stoppers import EarlyStopper
from pykeen.losses import MarginRankingLoss, BCEWithLogitsLoss, SoftplusLoss
import torch

def main():

    dataset = FB15k237(create_inverse_triples=True)
    training_triples_factory = dataset.training

    loss = BCEWithLogitsLoss()
    
    model = TransE(triples_factory=training_triples_factory, embedding_dim=100, loss=loss)

    optimizer = Adam(params=model.get_grad_params(), lr=0.001)
    print(f"Number of params: {sum(p.numel() for p in model.parameters())}")

    training_loop = LCWATrainingLoop(model=model, optimizer=optimizer)

    valid_evaluator = RankBasedEvaluator()

    early_stopper = EarlyStopper(
            model=model,
            relative_delta=0.0005,
            evaluation_triples_factory=dataset.validation,
            frequency=1,
            patience=1000,
            result_tracker=None,
            evaluation_batch_size=256,
            evaluator=valid_evaluator,
        )

    training_loop.train(num_epochs=1000,
                        batch_size=256,
                        result_tracker=None,
                        stopper=early_stopper,
                        label_smoothing=0.1)

    # Pick an evaluator


    # Get triples to test
    mapped_triples = dataset.testing.mapped_triples

    # Evaluate
    test_evaluator = RankBasedEvaluator()
    results = test_evaluator.evaluate(model, mapped_triples, batch_size=128)
    print(results)

if __name__ == '__main__':
    main()
```

@migalkin I could consistently replicate your error using your script and un-setting all batch sizes so AMO tries to derive the right sizes on an A100-40GB. I have a proposed a fix in #433. I have also tested the fix on your add-compat branch and it consistently fixed the bug 😊 Does this fix the issue for you as well?