Performance Optimization and Comparative Analysis of Generative AI Models on Advanced Accelerators

Technical Summary: Performance Optimization and Comparative Analysis of Generative AI Models on Advanced Accelerators

Problem Statement
Despite the impressive capabilities of Generative AI models, including Large Language Models (LLMs) and diffusion models, their deployment faces significant hurdles. These challenges include substantial memory requirements, extended inference latency, high computational demands, and elevated hardware costs. These issues are further exacerbated when evaluating models across heterogeneous platforms, where variations in numerical formats, memory bandwidths, and software stacks interact complexly with model architectures and workload characteristics. The paper addresses the need for a systematic study to optimize performance and compare these models across diverse downstream tasks on modern High-Performance Computing (HPC) systems and advanced accelerators.

Methodology
The authors conducted a comprehensive evaluation across three distinct HPC systems: Perlmutter (NERSC, featuring NVIDIA A100 GPUs), Expanse (SDSC, featuring NVIDIA V100, A100, and H100 GPUs), and Voyager (SDSC, featuring Intel Gaudi accelerators). The study focused on three primary areas:

1. Mixed-Precision Post-Training Quantization (PTQ): The authors introduced a novel sensitivity-aware mixed-precision PTQ framework designed to optimize LLMs (specifically TinyLlama-1.1B) across NVIDIA A100 and Intel Gaudi HPU accelerators.

   • Sensitivity Analysis: The framework employs a two-phase approach. It computes sensitivity scores at the transformer layer level by analyzing weight magnitudes (excluding biases and normalization parameters) and applying a unified pruning threshold based on a sparsity target (30%).
   • Quantization Strategy: Sensitive layers are assigned higher precision, while robust layers are assigned lower precision. The study evaluated "fake" quantization (simulated in FP32) and "real" quantization using FP16/INT8 for A100 and BF16/FP8 (via Intel Neural Compressor) for Gaudi.
   • Datasets: Evaluation was performed on WikiText-2 (language modeling), HellaSwag (commonsense reasoning), and BoolQ (reading comprehension).

2. Fine-Tuning Performance on Gaudi: The study assessed the fine-tuning of three open-source LLMs (Llama3.1-8B, Gemma2-27B, and Llama3.3-70B) on Voyager using the Optimum-Habana library.

   • Parallelization: The authors compared Distributed Data Parallel (DDP), DeepSpeed (ZeRO3), and Fully Sharded Data Parallel (FSDP) strategies.
   • Optimization: The impact of Flash Attention (FA) and the generational differences between Gaudi1, Gaudi2, and Gaudi3 were analyzed.

3. Diffusion Model Training: A Super-Resolution model, utilized for enhancing astrophysical images of galactic dust, was trained on Gaudi nodes and compared against NVIDIA GPUs (V100, A100, H100) on Expanse. The model was parallelized using PyTorch DDP and executed in SynapseAI Lazy mode.

Key Results

• Quantization Performance: Real quantization on A100 and Gaudi2 achieved compression ratios between 2x and 6x with minimal accuracy loss. For instance, on the WikiText-2 task, Gaudi2 achieved a 5.93x compression ratio with a perplexity of 12.39 (compared to 12.02 for FP32), while A100 achieved 3.77x compression with a perplexity of 12.02. Throughput improvements were significant, with Gaudi2 showing up to 34,081 tokens/sec in simulated quantization scenarios.
• Fine-Tuning Efficiency:
  • Generational Gains: Gaudi2 cards demonstrated a 2.5x to 3x runtime reduction compared to Gaudi1, attributed to the increase from 8 to 24 Tensor Processing Cores (TPCs). Gaudi3 offered an additional 1.8x to 2x performance gain over Gaudi2.
  • Flash Attention: Enabling Flash Attention on Gaudi1 and Gaudi2 resulted in a 10% to 20% performance improvement by minimizing data movement between memory types.
  • Sharding Overhead: For the 27B parameter Gemma2 model, DDP decisively outperformed model sharding methods (DeepSpeed and FSDP). The authors found that while sharding reduces memory footprint, it introduces significant overhead that degrades performance. Consequently, sharding should only be employed when a model cannot fit within the HBM of a single card.
• Diffusion Model Scaling:
  • NVIDIA GPUs: Exhibited nearly perfect scaling with a ~2x speed-up from V100 to A100 and another ~2x from A100 to H100.
  • Intel Gaudi: Showed a 4x speed-up from Gaudi1 to Gaudi2. However, this trend did not continue linearly from Gaudi2 to Gaudi3, indicating a plateau in scaling efficiency for this specific workload.

Significance and Claims
The paper claims to provide a systematic, comparative analysis of Generative AI performance across heterogeneous HPC environments, offering practical insights for deployment. Its primary contributions include:

1. A novel sensitivity-aware mixed-precision quantization framework that balances model accuracy and compression across different hardware architectures.
2. Empirical evidence demonstrating that for models fitting within single-card memory, standard DDP is superior to complex sharding strategies like DeepSpeed or FSDP, which introduce unnecessary overhead.
3. Detailed benchmarking of the evolution of Intel Gaudi accelerators, highlighting specific performance gains (and limitations) in generational upgrades and the efficacy of Flash Attention in this ecosystem.
4. Validation of diffusion model training capabilities on Gaudi accelerators, providing a baseline for scientific computing applications such as astrophysical image enhancement.

The work emphasizes that while advanced accelerators offer substantial performance gains, optimal deployment requires careful consideration of model size, memory constraints, and the specific trade-offs between quantization, parallelization strategies, and hardware generations.