Training with Multiple GPUs Using Model Parallelism¶
Training deep learning models can be resource intensive. Even with a powerful GPU, some models can take days or weeks to train. Large long short-term memory (LSTM) recurrent neural networks can be especially slow to train, with each layer, at each time step, requiring eight matrix multiplications. Fortunately, given cloud services like AWS, machine learning practitioners often have access to multiple machines and multiple GPUs. One key strength of MXNet is its ability to leverage powerful heterogeneous hardware environments to achieve significant speedups.
There are two primary ways that we can spread a workload across multiple devices. In a previous document, we addressed data parallelism, an approach in which examples within a batch are divvied among the available devices. With data parallelism, each device stores a complete copy of the model. Here, we explore model parallelism, a different approach. Instead of splitting the batch among the devices, we partition the model itself. Most commonly, we achieve model parallelism by assigning the parameters (and computation) of different layers of the network to different devices.
In particular, we will focus on LSTM recurrent networks. LSTMS are powerful sequence models, that have proven especially useful for natural language translation, speech recognition, and working with time series data. For a general high-level introduction to LSTMs, see the excellent tutorial by Christopher Olah. For a working example of LSTM training with model parallelism, see example/model-parallelism-lstm/.
Model Parallelism: Using Multiple GPUs As a Pipeline¶
Model parallelism in deep learning was first proposed for the extraordinarily large convolutional layer in GoogleNet. From this implementation, we take the idea of placing each layer on a separate GPU. Using model parallelism in such a layer-wise fashion provides the benefit that no GPU has to maintain all of the model parameters in memory.
In the preceding figure, each LSTM layer is assigned to a different GPU. After GPU 1 finishes computing layer 1 for the first sentence, it passes its output to GPU 2. At the same time, GPU 1 fetches the next sentence and starts training. This differs significantly from data parallelism. Here, there is no contention to update the shared model at the end of each iteration, and most of the communication happens when passing intermediate results between GPUs.
In the current implementation, the layers are defined in lstm_unroll().
Workload Partitioning¶
Implementing model parallelism requires knowledge of the training task. Here are some general heuristics that we find useful:
- To minimize communication time, place neighboring layers on the same GPUs.
- Be careful to balance the workload between GPUs.
- Remember that different kinds of layers have different computation-memory properties.
Let’s take a quick look at the two pipelines in the preceding diagram. They both have eight layers with a decoder and an encoder layer. Based on our first principle, it’s unwise to place all neighboring layers on separate GPUs. We also want to balance the workload across GPUs. Although the LSTM layers consume less memory than the decoder/encoder layers, they consume more computation time because of the dependency of the unrolled LSTM. Thus, the partition on the left will be faster than the one on the right because the workload is more evenly distributed.
Currently, the layer partition is implemented in lstm.py and configured in lstm_ptb.py using the group2ctx option.
Apply Bucketing to Model Parallelism¶
To achieve model parallelism while using bucketing, you need to unroll an LSTM model for each bucket to obtain an executor for each. For details about how the model is bound, see lstm.py.
On the other hand, because model parallelism partitions the model/layers, the input data has to be transformed/transposed to the agreed shape. For more details, see bucket_io.
