Huge amount of CPU RAM needed during training

[davidecaroselli]

Hello team!

With the current version of fairseq we noticed that a huge amount of RAM (CPU RAM, not GPU RAM) is required in order to run the training. Moreover this is correlated to the number of GPU used on the same machine.

So my guess is that the binarized data used for the training is completely loaded in RAM for every GPU process, this will result in having the amount of CPU RAM is roughly:
RAM ~= (number of GPUs) * sizeof(binarized data).
If this is true, the amount of RAM needed for medium/large training sets is huge (hundreds of GB) wrt size of training set (less than 100 GB).

If this is the case, why can't we use a memory mapped training set? So that the amount of RAM depends exclusively on sizeof(binarized data) ?

I'm available to work on this if needed, can you please give me some code context, or a good starting point to begin with?

[myleott]

One option is to use --fix-batches-to-gpus, which will load only 1/N of the data in each process, assuming N workers.

mmapping the whole training set per node is reasonable as well and may have some interesting tradeoffs compare to the --fix-batches-to-gpus option.

[davidecaroselli]

From the documentation of --fix-batches-to-gpus I understand that there could be some problems with translation quality (am I correct?), but I don't really understand why.. If the whole cluster should work together to train on the training set without repetitions, why splitting batches per-GPU will result in problematic training?

In other words: which is the advantage of having duplicated batches per GPU? If a batch "k" is run on "GPU X", for sure won't be seen by "GPU Y" until next epoch, is that correct?
If this is the case, why preventing batch "k" to be replicated in all GPUs != "GPU X" should be a problem?

Thanks for the info @myleott ! I just want to understand better how fairseq work in order to be more confident into changing the training DataSet implementation.

[myleott]

When using --fix-batches-to-gpus then it will slightly reduce randomness, which can (very slightly) affect model quality if training on many GPUs.

This is because we will never shuffle batches across GPUs, only within a GPU.

Consider the normal case, training with 2 GPUs on a dataset with 4 batches. You might have something like:

             GPU 1     GPU 2
Epoch 0.0   batch 1   batch 2
Epoch 0.5   batch 3   batch 4
Epoch 1.0   batch 1   batch 3
Epoch 1.5   batch 4   batch 2

But with --fix-batches-to-gpus there will be less randomness, since batches 1/3 will always be on GPU 1 and batches 2/4 will always be on GPU 2:

             GPU 1     GPU 2
Epoch 0.0   batch 1   batch 2
Epoch 0.5   batch 3   batch 4
Epoch 1.0   batch 3   batch 4
Epoch 1.5   batch 1   batch 2

I'd suggest trying it out, you may find very similar or even identical results.

[davidecaroselli]

Hi @myleott

thanks for the explanation, this make totally sense!
Unfortunately this option does not seems to reduce CPU RAM consumption. I just tried to run my training on a 8GPU/128G RAM machine and the training failed again (side note: on a 8-GPU/256G RAM the training run smoothly).

So I would like to work on the MemoryMapped Dataset, because I think can definitely solve this issue. Can you point me to the parts of code that creates the current Dataset version and where it is used during training? So that I can start digging the code.

Thanks for your help!

PS: Can you re-open the issue? I think this could be useful to track the status of the development and the main place to share any feedbacks or request further info during the development!

[myleott]

Ah, you're right. If you're okay with slightly higher I/O from seeking, you can use the --lazy-load option [1].

This will switch to using an IndexedDataset instead of IndexedCachedDataset [2], which will seek instead of loading everything into memory.

If you want to make a memory mapped version, then you can extend IndexedDataset [3].

[1] https://github.com/pytorch/fairseq/blob/master/fairseq/tasks/translation.py#L56-L57
[2] https://github.com/pytorch/fairseq/blob/master/fairseq/tasks/translation.py#L141-L144
[3] https://github.com/pytorch/fairseq/blob/master/fairseq/data/indexed_dataset.py#L50

[davidecaroselli]

Hi @myleott !

I have an initial implementation of the Dataset, and it seems to work correctly. The problem is that I am having hard time trying to "plug in" the implementation in the flow of Fairseq. The problem I see is that Fairseq preprocess/train is entangled with the internal implementation of the Dataset (i.e. *.bin and *.idx files are directly accessed from outside the class itself!)

More problematic, here I see a bunch of public variables (https://github.com/pytorch/fairseq/blob/master/fairseq/data/indexed_dataset.py#L68) that are not compatible with the implementation I need. Even more relevant, I see that the IndexedRawTextDataset implementation does not expose them.

So basically I'm having troubles in understanding what is the "public" API every Dataset has to implement and what is instead a private implementation. More important the logic to load/delete/create a Dataset with the right implementation is spread and duplicated across multiple files (i.e. every Task seems to re-implement the loading depending on arguments).

So my question is: can I try to create a solid, consistent FairseqDataset interface that declares all an only public methods of every dataset? This first step should help me a lot in the creation of a new Dataset implementation.

[davidecaroselli]

wait.. I am confused.

There is already a FairseqDataset in fairseq_dataset.py but why no class in indexed_dataset.py extends it?

[myleott]

The "dataset" name is a bit overloaded. IndexedDataset is a standalone module that implements PyTorch's Dataset API.

There is also the FairseqDataset API, which is more closely connected to FairseqTasks. Currently the main two FairseqDatasets are LanguagePairDataset (for translation) and MonolingualDataset (for language modeling). You can extend this, but generally you'd be better off extending TranslationTask and returning a LanguagePairDataset that consists of PyTorch Datasets of your own (e.g., IndexedDataset, MemoryMappedIndexedDataset, etc.).

[davidecaroselli]

Thanks @myleott !

I have just completed the mmap implementation on my fork: https://github.com/davidecaroselli/fairseq/blob/features/mmap_dataset/fairseq/data/indexed_dataset.py#L264
I am now testing if the two implementations (IndexedDataset and MMapDataset) return the same lines after fairseq-preprocess.

In the meantime, I would like to replace the dataset flags --output-format and --lazy-load with a unique --dataset-type with possible values: raw, lazy, cached, mmap so that is much easier to select the correct implementation of the Dataset. Do you agree with such change? If not, what do you suggest?

[myleott]

Cool, that sounds reasonable. It’d be nice to reuse this logic in both preprocess and the Tasks (Translation and LanguageModel).

I’m also unsure if the name might be confusing, maybe dataloader-type instead? Although of course that’s overloaded with the PyTorch concept of a DataLoader :/

[davidecaroselli]

I’m also unsure if the name might be confusing, maybe dataloader-type instead? Although of course that’s overloaded with the PyTorch concept of a DataLoader :/

I see, for my understanding this option should be available in both preprocess.py and train.py, right? Something like --dataset-impl? I would like to keep the name --dataset-xxx because you're actually selecting the right implementation for a class called Dataset. IMHO if we use --dataloader that could lead to more confusion.

[myleott]

--dataset-impl sounds good. I wonder if we can also define this directly in the Task, and add necessary APIs to Task that can be called in preprocess. In general we’re trying to make Task become the primary interface for this kind of stuff, since we’d like to support a wider variety of tasks using the same CLI tools.