Thanks for the answer!



 iffiX:

DistributedSampler in this case could also be a major overhead contributor, since internally the DataLoader will use _MultiProcessingDataLoaderIter , which uses a inter process queue to get data from sub-processes


I was just testing some bottlenecks in my code an it indeed seems like the DistributedSampler is a major culprit, when used with shuffle=True. This might be because the dataset I was using in the tests has a length of 320 million samples. And it might also be the reason why I wasnt using bigger batch sizes because 10 GB of RAM seems to be occupied by the DistributedSampler.
So I decided to use the sampler on a single GPU and look at the effects.
code:
import torch
import os
from transformers import AdamW, BertConfig, TrainingArguments, Trainer
from datetime import datetime

from prediction_module import path_vocab, path_storage
from prediction_module.protein_datasets import ProteinBertMaskedLMDataset
from prediction_module.protein_models import ProteinBertForMaskedLM


def train_single(dist_sampler=False, n_steps=100, shuffle=False):
    rank = 0
    dataset = ProteinBertMaskedLMDataset(
        path_vocab, os.path.join(path_storage, "data", "uniparc", "uniparc_train_sorted.h5"),
    )
    config = BertConfig(
        vocab_size=dataset.tokenizer.vocab_size,
        max_position_embeddings=dataset.input_size
    )
    model = ProteinBertForMaskedLM(config)
    model.cuda(rank)
    model.train()

    optimizer = AdamW(model.parameters(), lr=1e-5)  # create optimizer
    # Data loading code

    sampler = DistributedSampler(dataset, 1, rank, shuffle=shuffle) if dist_sampler else None
    loader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=False, num_workers=0, pin_memory=False,
                                         sampler=sampler, collate_fn=dataset.collate_fn)
    print("start trainig")
    start = datetime.now()
    for epoch in range(1):
        for i, inputs in enumerate(loader):
            for k, v in inputs.items():
                if isinstance(v, torch.Tensor):
                    inputs[k] = v.cuda(rank, non_blocking=True)
            optimizer.zero_grad()
            outputs = model(**inputs)
            loss = outputs[0]
            loss = loss.mean()
            loss.backward()
            optimizer.step()
            if i >= n_steps:
                break
    print("Training complete in:", str(datetime.now() - start))


if __name__ == "__main__":
    train_single()
    train_single(True)
    train_single(True, shuffle=True)


The results:

start trainig
Training complete in: 0:00:07.233262
start trainig
Training complete in: 0:00:19.662416
start trainig
Training complete in: 0:03:29.437496

Additionally, the RAM on the GPU required for the first two function calls is about 3,5 GB and for the last one about 14 GB. I am not sure what is going on here, but that seems very weird.