Data Loading for Distributed Training
This documentation covers the management of datasets and dataloaders in PyTorch and TensorFlow, i.e., the entire pipeline that enables data to be passed in batches to the model during neural network training.
Practical examples in PyTorch and TensorFlow are provided at the bottom of the page.
This documentation is tailored to the Jean Zay environment:
- with large datasets (images, 3D images, videos, text, audio, etc.). Data loading is performed on CPUs, and the data is then transferred batch by batch to GPUs for model training or inference.
- with multi-GPU distribution using data parallelism. If you are not using this type of parallelism, you can simply ignore the distribution step.
The data loading and preprocessing pipeline for distributed training is illustrated below:
-
The Dataset is the collection of data, typically available as a structured set of files.
-
The Structure Discovery step involves identifying the structure of the dataset. At this stage, the data is not yet loaded into memory. For example, you retrieve the list of file paths or, more generally, the list of data pointers associated with their labels for loss function calculation in supervised learning.
-
Shuffle is the step where the dataset is randomly permuted. This is necessary for training and is performed at the beginning of each epoch.
The data pointers are stored in a memory buffer before being shuffled. Ideally, the buffer size should be large enough to store the entire dataset. However, memory constraints often make this impossible. A smaller buffer size, capable of holding ~1,000 or ~10,000 pointers, is considered acceptable and has no impact on training quality.
Note: For inference or validation, the Shuffle step is not performed. -
Distribution is the data parallelisation step that distributes the data across the various GPUs. Ideally, Distribution occurs after Shuffle. In practice, it is usually performed before Shuffle (or after Batch) to avoid inter-process communication. This is considered acceptable and has no perceptible impact on training quality.
-
Batch is the sampling step, i.e., splitting the dataset into fixed-size batches. An iterative loop is then run over all these batches.
Note: With data parallelism, you must distinguish between the batch size, which is defined during the input data preprocessing, and the mini-batch size, which corresponds to the portion of data sent to each GPU during training. -
Load and Transform is the step where the batch is loaded into memory, decoded, and transformed (preprocessed). This is the most CPU-intensive step and is therefore a clear bottleneck.
Some transformation operations are deterministic, meaning they generate the same transformed data in every epoch, while others are related to Data Augmentation and are random.
It is possible to perform deterministic operations upstream of the epoch loop to avoid recalculating them in each iteration. To do this, you can store the transformed data in a memory cache (in this case, the storage occurs during the first epoch, which will be longer than the subsequent ones). You can also create a custom input file, such as TFRecords, before training.
The random operations must be performed in each iteration. One possible optimisation is to offload these calculations to the GPU. Solutions such as DALI, for example, propose doing this for image Data Augmentation. GPUs are much faster than CPUs for this type of processing.
Note: Our experience on Jean Zay shows that, by optimising multiprocessing and Prefetch (see next point), CPUs are rarely congested for data preprocessing. To our knowledge, offloading operations to the GPU is therefore not essential on Jean Zay. -
Prefetch creates a FIFO (First In, First Out) queue for loading batches onto the GPU. This feature enables overlap between data transfers and computation.