practical tips for accelerating neural network training

We continue to study how to speed up the training of neural networks. IN last article we dived into the theoretical aspects of this problem. Today let's move on to practice.

We will look at several interesting studies that demonstrate the effectiveness of various approaches to accelerating neural networks on a variety of tasks and datasets. Then we will discuss practical recommendations for choosing and combining optimization methods and tell you which tools are best to use for profiling and monitoring the learning process. To top it off, we’ll look at useful libraries for fast and efficient development.

Research and its results

To comprehensively evaluate the effectiveness of neural network optimization methods, it is not enough to limit ourselves to theoretical analysis alone. It is also necessary to test these methods on real application problems and generally accepted data sets. Fortunately, many researchers have already conducted similar experiments. Of course, it will not be possible to consider all of them, but I took a sample of five studies as an example.

Comparison of various acceleration techniques on real problems and datasets

One promising direction is gradient optimization. Yuzhong Yun and colleagues at work “Z-score normalization of gradients to speed up neural network training“proposed the ZNorm method, which equalizes the scale of gradients between network layers. Experiments conducted on standard image classification datasets (CIFAR-10 and ImageNet) showed that ZNorm outperforms methods such as clipping and gradient centralization, providing faster convergence and better accuracy. The key advantage of ZNorm is that it requires virtually no additional computation. This makes it a good acceleration method.

Another rapidly developing area is optimizing the training of graph neural networks (GNNs). Hesham Mostafa and co-authors in the article “Accelerating distributed training of graph neural networks for billion-scale graphs» presented the FastSample method, aimed at working with extremely large graphs. It combines a new graph partitioning algorithm that minimizes inter-node communications and an optimized sampling kernel that reduces the amount of data sent. Experiments on large-scale graphs have demonstrated a twofold acceleration of training of modern GNN architectures (GraphSAGE, GAT) without loss of accuracy. Moreover, FastSample has demonstrated excellent scalability in a distributed environment, which makes it possible to process graphs of record sizes (hundreds of billions of edges) with a moderate increase in computing resources.

Below is a graph of the speedup in sample time and total training time when training on a single node on the ogbn-papers100M dataset. The highly optimized sampling kernels in DGL were used as a baseline. The graph shows the speedups for different mini-batch sizes (1024, 2048, 4096, 8192, and 10240) and different fanout values ​​for each of the three GNN layers in the model.

Many researchers rely on specialized hardware accelerators. A striking example is the MaxK-GNN system, designed by Hongwu Peng and his team to train graph neural networks on GPUs. MaxK-GNN effectively combines hardware and algorithmic optimization. At the hardware level, specialized kernels are implemented for forward and reverse pass operations. Algorithmically, a new nonlinearity MaxK is proposed, which is theoretically justified as a universal approximator. Experiments on numerous datasets have shown that MaxK-GNN achieves the quality of state-of-the-art models with 3-4 times acceleration.

Hybrid quantum-classical approaches, which are still being explored, deserve special mention. In work “Hybrid Quantum-Classical Scheduling to Accelerate Training of Newtonian Gradient Descent Neural Networks» Pingzhi Li and colleagues proposed the Q-Newton scheduler to speed up the training of neural networks with the second order (Newton's method). The key idea is the distribution of subproblems between quantum and classical solvers of linear systems based on a heuristic estimate of the condition number. Tests on synthetic data showed a manifold reduction in training time compared to traditional first-order (SGD) approaches. And although the practical use of Q-Newton will require fairly powerful quantum computers, such hybrid circuits have a good chance of overcoming the computational limitations of classical systems.

Another important trend is the adaptation of modern neural network architectures to low-power devices. Matteo Presciutto in the article “Compress and accelerate neural networks on resource-constrained hardware for real-time inference» used quantization of weights and activations to port neural networks to DSP processors widely used in embedded systems. The proposed technique made it possible to reduce the size of models by 4 times with minimal loss of accuracy (less than 1% on test data). This opens the way for neural networks to real-time systems – from smartphones to autopilots.

Analysis of the impact of optimizations on model performance and quality

So, what did these optimizations ultimately achieve?

1. ZNorm gradient normalization improves the speed and robustness of training, allowing you to produce better models in less time. At the same time, ZNorm requires virtually no computational effort, which makes it attractive for accelerating modern deep architectures. As a result, ZNorm gives a multiple increase in the speed of convergence compared to basic methods.

2. FastSample's graph sampling optimization significantly reduces GNN training time on very large sparse graphs. This occurs by minimizing inter-node communications and the volume of data sent. A separate bonus is excellent scalability in a distributed environment, allowing you to process graphs of record size (hundreds of billions of edges) with a moderate increase in computing resources. It is important that at the same time it is possible to maintain high accuracy of the models.

3. Hardware accelerators like MaxK-GNN provide a multiple increase in the performance of training graph neural networks while maintaining the quality of the models. They clearly demonstrate the synergistic effect of joint optimization of the hardware and algorithm levels. Dedicated forward and backward pass kernels efficiently utilize GPU resources, and algorithmic innovations such as MaxK non-linearity improve the representativeness of models. The results of competent level synergy are obvious.

4. Hybrid quantum-classical approaches, as in the Q-Newton scheduler, have the potential to significantly speed up the training of second-order neural networks compared to classical methods. And in the experiments conducted, the training time was reduced by up to 4 times. However, further progress in this technology requires quantum computers of sufficient power, which do not yet exist.

5. Model compression algorithms, such as quantization, can significantly speed up inference and reduce memory requirements with minimal degradation in accuracy. For example, moving from 32-bit to 8-bit quantization reduces the model size by 4 times with virtually no loss of accuracy. This makes it possible to run modern neural networks on low-power embedded systems, such as DSP processors, in real-time scenarios.

This is just a series of studies on this topic, but they well illustrate the general trend: the leap is to be. Although some of the approaches considered, especially from the field of quantum computing, still remain at the level of fundamental research, the general vector towards specialization and hybridization is obvious.

Recommendations and tips

Now I want to give some practical recommendations that will be useful to you in experiments with acceleration.

Selecting the optimal combination of methods for a specific task

With all the variety of modern methods for optimizing neural networks, there is no universal solution suitable for any task. Each approach has its own strengths and weaknesses that must be taken into account when choosing an optimization strategy.

The first important factor is the characteristics of the problem itself and the data. For example, if we are working with sequential data such as text or time series, then recurrent architectures like LSTM or GRU may be a good choice. For them, it is critically important to correctly organize the transfer of state between sequence steps, using, for example, effective implementations of CUDA kernels or frameworks such as cuDNN.

If the task is related to graph processing, then you should pay attention to methods for optimizing graph neural networks. For large sparse graphs, methods like FastSample are well suited, minimizing communications between cluster nodes. For dense graphs, you can use hardware accelerators like MaxK-GNN, which effectively utilize GPU resources.

Another important aspect is the available computing resources. If the model is trained on a single GPU, then the focus should be on algorithmic optimizations like ZNorm or quantization. If you have a cluster with many nodes, you can use distributed learning methods such as Horovod or PyTorch DDP.

Finally, you need to take into account the features of the model itself. For deep convolutional networks, optimizing convolution operations, such as using Winograd or FFT methods, is critical. For transformers, the key role is played by optimizing the attention operation, which accounts for the main computational load. Here you can apply approximate attention calculation methods, like Linformer or Reformer.

In general, choosing the optimal combination of methods is always a compromise between speed, accuracy and resource costs. Therefore, in practice, a thorough analysis of the problem and experimental selection of the best combination of architecture, learning algorithms and hardware platform are necessary.

Profiling and monitoring the learning process

To identify bottlenecks and optimize the neural network training process, it is necessary to carry out its profiling and monitoring. Modern frameworks such as PyTorch and TensorFlow provide rich capabilities for collecting and visualizing various performance metrics.

For example, PyTorch has a built-in profiler, torch.profiler, which allows you to measure execution time and memory consumption at the level of individual statements. With its help, you can find the most costly operations and focus optimization efforts on them. PyTorch also provides convenient hooks for collecting intermediate tensor values, gradients, and other quantities during forward and backward passes. This is useful for debugging and monitoring the training process.

In TensorFlow, you can use the TensorBoard tool for profiling, which allows you to visualize the computational graph, analyze resource utilization, and find bottlenecks. The tf.profiler profiler can also be useful, providing detailed information about execution time and memory consumption at the level of individual operations.

Tools like Weights and Biases (wandb) or TensorBoard are good for monitoring the learning process in real time. They allow you to log various metrics, visualize their dynamics, compare experiments, and track resource utilization. This helps to quickly diagnose problems with convergence, overfitting, and suboptimal use of resources.

When distributing training on a cluster, it is important to monitor network activity and load balancing between nodes. To do this, you can use both standard tools like Ganglia or Prometheus, and specialized solutions for deep learning clusters, such as Horovod Timeline or TensorFlow's tf.contrib.timeline.

In addition to high-level metrics, profiling performance at the level of individual CUDA cores and operations is useful for deep optimization. To do this, you can use tools like Nsight Compute and Nsight Systems from NVIDIA, which provide detailed information about GPU performance and resource utilization. They help identify bottlenecks, optimize memory use, and effectively parallelize calculations.

Regular profiling and monitoring is the key to understanding model behavior and effectively optimizing the training process. This is especially important when working with large and complex models, where the cost of error is high and resources are limited. Smart use of profiling and visualization tools helps you quickly find and fix problems, thereby increasing the efficiency of research and development.

Data pipeline optimization

Another important aspect of effective neural network training is data pipeline optimization. This includes the steps of extracting, transforming and loading data, as well as feeding it efficiently into the model.

One of the main problems is the input/output bottleneck (I/O bottleneck), when the data feed speed does not keep up with the speed of calculations on the GPU. To solve this problem, you need to parallelize and optimize I/O and data preprocessing operations as much as possible.

The first step is to use fast data storage formats such as TFRecord, Feather, Parquet. They allow you to efficiently serialize and deserialize tensors, as well as store them in compressed form, saving disk space and read time.

Then you need to organize a data preprocessing pipeline that will perform all the necessary transformations on the fly, without saving intermediate results to disk. This allows you to avoid filling up memory and parallelize calculations. A good example of such a pipeline is tf.data in TensorFlow or torch.utils.data in PyTorch.

A key technique for speeding up the pipeline is prefetching, that is, loading the next batch of data in the background while the model processes the current batch. This allows I/O time to overlap with computation time and almost completely eliminate the I/O bottleneck. Most deep learning frameworks support prefetching out of the box, you just need to properly configure the buffer size and number of threads.

Another useful trick is to cache your most frequently accessed data in RAM or on an SSD. This avoids reading the same data over and over again from slow network storage or a remote database. Caching is especially effective when iterating over the same data set multiple times, such as when training with a large number of epochs.

Finally, for distributed learning on a cluster, it is important to minimize the amount of data sent between nodes. To do this, you can use data compression and quantization techniques, as well as algorithms for their efficient distribution between nodes, such as Ring AllReduce or Parameter Server.

The pipeline optimization stage is often underestimated as an aspect of effective neural network training. Although a properly constructed ETL pipeline can speed up data processing by orders of magnitude and save a significant amount of computing resources. Conversely, suboptimal I/O organization can negate all efforts to optimize the model and infrastructure.

Useful libraries and frameworks

Finally, I would like to note several useful libraries and frameworks that can significantly facilitate and speed up the development and training of neural networks.

First up is PyTorch Lightning, a high-level add-on to PyTorch that provides convenient abstractions for organizing code, distributed training, logging, and visualizing metrics. Lightning allows you to focus on the model itself and quickly experiment with different architectures and hyperparameters without getting distracted by low-level details.

A similar role for TensorFlow is played by Keras, a high-level API for building and training neural networks. Keras provides a simple and intuitive interface for defining the model architecture, selecting the optimizer and loss function, and setting up the training process. It is compatible with various backends (TensorFlow, Theano, CNTK) and can be used either independently or as part of more complex pipelines.

For working with large language models (BERT, GPT, T5, etc.), the Transformers library from Hugging Face is an indispensable tool. It provides ready-made pretrained models and convenient interfaces for fine-tuning and inference of transformers for various natural language processing (NLP) tasks. The library is optimized for effective training on GPU and TPU, supports various frameworks and is constantly updated with new state-of-the-art models.

If you need to quickly deploy a service based on a neural network, then the FastAPI framework may be a good choice. It allows you to create a REST API for your model with minimal effort and provides high performance and automatic documentation generation. FastAPI is built on a modern asynchronous Python stack (Starlette, Pydantic, uvicorn) and is well suited for microservice architectures.

Finally, for monitoring and visualization of the learning process, it is worth paying attention to TensorBoard and Weights and Biases (wandb). The first is the official visualization tool from TensorFlow, but it also supports logs from PyTorch and other frameworks through a simple API. The second is an independent cloud service with rich capabilities for tracking experiments, comparing metrics, saving artifacts and collaboration.

Of course, this is not a complete list of useful tools, and depending on the specific task and stack, completely different libraries may be useful. The main thing is not to be afraid to experiment, follow the development of the ecosystem and select the optimal tools for your needs. After all, as Donald Knuth said, “premature optimization is the root of all evil,” and this fully applies to the development of neural networks. First you need to make a working prototype, and only then, as necessary, optimize its performance and efficiency. And this is where properly selected frameworks and libraries will come to the rescue.

That's probably all. I hope this article will help you in the difficult but exciting task of optimizing neural networks. Remember that behind every successful project there are months of hard work and many unsuccessful experiments. The main thing is not to give up and continue to look for better solutions.

If you already have experience in accelerating neuron learning or ideas on this topic, be sure to share them in the comments!