Zenn
🚄

基盤モデルを加速させる3つのハードウェア技術

2024/03/09に公開
AI
アーキテクチャ
メモリ
高速化
tech

対象読者

  • 機械学習モデルの大規模化に興味のある方
  • 機械学習のインフラに興味のある方
  • 機械学習関連のハードウェア技術動向を追いかけたい方

はじめに

ハードウェア技術の重要性

基盤モデルは、数億から数千億のパラメータを持つ巨大なモデルであるため、学習には膨大な計算リソースが必要です。例えば、2020年発表のGPT-3は1750億個のパラメータを持っており[1]、学習には1000台以上のGPU（NVIDIA V100）を使用した1ヶ月以上におよぶ計算が必要だったでしょう[2]。こうした大規模な学習を効率的に行うには、高性能なハードウェア技術が不可欠です。

具体的には、以下の３つが求められています。

  • メモリ性能の向上
  • 高速インターコネクトによる分散学習の効率化
  • 量子化技術による演算量の削減

本記事ではこれら、基盤モデルの学習を加速する3つの重要な技術について、わかりやすく解説します。

基盤モデルを加速させる3つのハードウェア技術

メモリ性能の向上

メモリ帯域幅の重要性

メモリ帯域幅は、単位時間あたりにメモリとプロセッサ間で転送できるデータ量を表します。基盤モデルの学習では、大量のパラメータを頻繁にメモリとプロセッサ間で読み書きする必要があるため、高いメモリ帯域幅が求められます。そして現状、基盤モデルのアプリケーションでは、この通信がボトルネックになっています[3]。

メモリ帯域幅を向上させるには、メモリインターフェースの幅を広げたり、クロック周波数を上げたりする必要があります。HBMでは、1024ビット幅のインターフェースを採用することで、高いメモリ帯域幅を実現しています。また、次世代（第5世代）のHBM3eでは、1TB/s以上のメモリ帯域幅が実現され、基盤モデルのパラメータ読み書きをさらに高速化できます。

HBM3Eの仕組み
image (3).jpg

HBMの利点と次世代のHBM3eへの期待

大規模な基盤モデルの学習には、高速かつ大容量のメモリが必要です。従来のGDDR（Graphics Double Data Rate）メモリでは、容量と帯域幅が不足するため、HBM（High Bandwidth Memory）が使用されています。

HBMは、Si貫通電極をもちいた3D積層技術により、メモリチップをGPUパッケージ上に直接配置することで、高速なメモリアクセスを実現します。現在のHBM3では、最大819GB/sのメモリ帯域幅と最大24GBの容量を提供しています[4]。

次世代のHBM3eでは、さらなる大容量化と高速化が期待されており、最大1,280GB/sのメモリ帯域幅と最大36GBの容量が実現される見込みです[5]。これにより、基盤モデルの学習を大幅に加速することが可能になるでしょう。

TSV（Through silicon via)技術

image (5).jpg

image (1).jpg

高速インターコネクトによる分散学習の効率化

NVLinkの仕組みとメリット

基盤モデルの学習では、1台のGPUでは処理しきれないほどの計算量が必要となるため、複数のGPUを用いた分散学習が一般的です。分散学習では、モデルのパラメータを複数のGPUに分割し、並列に処理することで学習を高速化します。その際、GPUの間で大量のデータをやり取りする必要があるため、高速なインターコネクトが重要になります。

NVIDIAのNVLinkは、GPUを専用のケーブルで直接接続し、第4世代では最大900GB/sの帯域幅を提供します[7]。これにより、GPUメモリ間のデータ転送を高速化し、分散学習をより効率的におこなえます。また、NVLinkはCPUを介さずにGPU間で直接通信できるため、レイテンシを最小限に抑えることが可能です。

image (7).jpg

分散学習におけるボトルネックの解消

分散学習を効率的に行うには、高速なインターコネクトだけでなく、通信量を削減するアルゴリズムも重要です。例えば、勾配圧縮では、勾配のスパース性を利用して、重要な勾配のみを選択的に通信することで、データ量を削減します[8]。また、量子化を適用することで、勾配や重みのビット数を削減し、通信量を抑えることができます。

これらの技術とNVLinkを組み合わせることで、分散学習におけるボトルネックを解消し、基盤モデルの学習を大幅に加速できます。

Tensor コアによる疎行列の圧縮
image (8).jpg

量子化技術による演算量の削減

量子化技術の概要とメリット

量子化技術は、モデルのパラメータを低ビット化することで、メモリ使用量と演算量を削減する手法です。通常、ニューラルネットワークの重みは、32ビットの浮動小数点数（FP32）で表現されますが、これを16ビットの半精度浮動小数点数（FP16）や8ビットの整数（INT8）で表現することで、メモリ使用量を削減できます。また、低ビット演算ユニットを使用することで、演算量も削減できます。これにより、精度を維持しつつ、学習の高速化と省メモリ化を実現できます。

image (9).jpg

量子化に適したモデルアーキテクチャと学習手法

超低ビット量子化を適用する際は、量子化に適したモデルアーキテクチャの設計が重要です。例えば、バイナリニューラルネットワーク（BNN）は、重みを1ビットで表現するモデルであり、メモリ使用量と演算量を大幅に削減できます[11]。とくにFPGAなどの組み込み系では、超低ビット量子化は重要になってきます。

また、量子化による精度低下を抑える学習手法も重要です。量子化認識トレーニング（QAT：Quantization-aware training）は、学習時の段階から量子化の影響を考慮することで、量子化による精度低下を最小限に抑えます。これらの技術により、超低ビット量子化を適用しつつ、高い精度を維持することが可能になります。

今後の展望

  • ハードウェアとソフトウェアの協調による基盤モデルの発展

基盤モデルの学習を加速するには、ハードウェア技術の進歩だけでなく、ソフトウェア側の工夫も重要です。量子化技術やモデル並列化など、ソフトウェアとハードウェアの協調により、基盤モデルの学習効率を最大限に引き出すことができます。また、ハードウェアの特性を考慮したモデルアーキテクチャの設計や、ハードウェアに最適化された学習アルゴリズムの開発など、ソフトウェアとハードウェアの密な連携が求められます。こうした取り組みにより、より大規模で高性能なモデルが実現されることを期待しています。

P.S.

間違っているところがあれば、教えてくださると大変ありがたいです！

さいごにAI/半導体/XRの情報収集およびアウトプット用にXのアカウントを開設したので、ぜひフォローよろしくお願いします！

https://twitter.com/ayame_willow

脚注

  1. https://arxiv.org/abs/2005.14165 ↩︎

  2. 推測です ↩︎

  3. https://medium.com/riselab/ai-and-memory-wall-2cb4265cb0b8 より ↩︎

  4. https://semiconductor.samsung.com/jp/dram/hbm/hbm2e-flashbolt/ より ↩︎

  5. https://news.samsung.com/global/samsung-develops-industry-first-36gb-hbm3e-12h-dram より ↩︎

  6. https://investors.micron.com/news-releases/news-release-details/micron-commences-volume-production-industry-leading-hbm3e より ↩︎

  7. https://www.nvidia.com/ja-jp/data-center/nvlink/ より ↩︎

  8. https://blogs.nvidia.co.jp/2020/05/27/sparsity-ai-inference/ より ↩︎

  9. https://speakerdeck.com/joisino/shen-ceng-moderunogao-su-hua?slide=29 より ↩︎

  10. https://www.nvidia.com/ja-jp/data-center/tensor-cores/ より ↩︎

  11. https://www.slideshare.net/kentotajiri/ss-77136469 より ↩︎

Discussion