  Transformers and LLMs have seen rapid adoption in all domains. Their sizes
have exploded to hundreds of billions of parameters and keep increasing. Under
these circumstances, the training of transformers is slow and often takes in
the order of weeks or months. Thanks to 3D model parallelism (data, pipeline,
and tensor-level parallelism), the training can scale to a large number of
GPUs, which reduces the duration of the training but dramatically increases the
cost. Even when a large number of GPUs are available, the aggregated GPU memory
is often not enough to hold the full training state (optimizer state, model
parameters, and gradients). To compensate, state-of-the-art approaches offload
the optimizer state at least partially to the host memory and perform hybrid
CPU-GPU computations. Such flexible solutions dramatically reduce the GPU
memory utilization, which makes it feasible to run the training on a smaller
number of GPUs at the cost of performance penalty. Unfortunately, the
challenges and bottlenecks of adopting this strategy are not sufficiently
studied by state-of-the-art, which results in poor management of the combined
host-GPU memory and poor overlapping between data movements and computations.
In this paper, we aim to fill this gap by characterizing the behavior of
offloaded training using the DeepSpeed runtime. Specifically, we study the GPU
memory utilization over time during each iteration, the activity on the PCIe
related to transfers between the host memory and the GPU memory, and the
relationship between resource utilization and the steps involved in each
iteration. Thanks to this study, we reveal opportunities for future
improvements of offloading solutions, which enable greater flexibility to
optimize the cost-performance trade-off in the context of transformer and LLM
training.
