Accelerating PyTorch Distributed Fine-Tuning with Intel Technologies

Hugging Face has demonstrated that distributing PyTorch fine-tuning jobs across a cluster of Intel Xeon Scalable CPU servers can significantly reduce training time, providing a viable alternative to GPU-based training for transfer learning tasks. By leveraging Intel's Ice Lake architecture and optimized software libraries, training speedups of up to 3x were achieved when scaling from one to four nodes.

Hardware Acceleration via Intel Ice Lake

For optimal performance, the setup utilizes Intel servers based on the Ice Lake architecture. This hardware supports specific features that accelerate deep learning operations:

  • Intel AVX-512: Advanced Vector Extensions for high-performance computing.
  • Intel Vector Neural Network Instructions (VNNI): Specialized instructions designed to speed up neural network inference and training.

These capabilities are available across major cloud providers, including Amazon EC2 (M6i and C6i instances), Azure (Dv5, Dsv5, Ddv5, Ddsv5, Edv5, and Edsv5 series), and Google Cloud Platform (N2 Compute Engine VMs).

Software Optimization Stack

To fully utilize the hardware capabilities of Intel CPUs, two primary software components are integrated into the PyTorch workflow:

Intel Extension for PyTorch

The Intel extension for PyTorch provides out-of-the-box speedups for both training and inference by enabling PyTorch to leverage AVX-512 and VNNI.

Intel oneAPI Collective Communications Library (oneCCL)

Distributed training often faces networking bottlenecks when large models exchange state information. The Intel oneAPI Collective Communications Library (oneCCL) is used as the communication backend for torch.distributed to efficiently handle communication patterns like all-reduce, which are critical for keeping nodes in sync during distributed training.

Performance Benchmarks

Using a BERT model fine-tuned on the MRPC dataset (part of the GLUE benchmark), Hugging Face measured the training time across different cluster sizes. The baseline was established on a single node (Amazon EC2 c6i.16xlarge instance).

MRPC Dataset Results

Cluster Size Training Time Speedup
1 Node 7m 46s Baseline
2 Nodes 4m 39s 1.7x
4 Nodes 2m 36s 3x

QQP Dataset Results

For the Quora Question Pairs (QQP) task, which involves a much larger dataset of over 400,000 training samples, the speedup remained consistent:

Cluster Size Training Time Speedup
1 Node 11h 22m Baseline
2 Nodes 6h 38m 1.71x
4 Nodes 3h 51m 2.95x

Implementation Requirements

Setting up the distributed environment requires specific configuration steps to ensure compatibility and performance:

  • Infrastructure: Identical instances with password-less SSH configured between the master and worker nodes, and all TCP ports open for internal oneCCL communication.
  • Dependency Matching: PyTorch and the Intel extension for PyTorch must have matching versions (e.g., PyTorch 1.9.0 and torch_ipex 1.9.0).
  • Script Modifications: Training scripts must be updated to import torch_ccl, handle the local rank via environment variables (such as PMI_RANK), and configure the master node address and port for the CCL backend.
  • Execution: Jobs are launched using mpirun to specify the number of processes and processes per node.

Sources