大規模な言語視覚モデルのトレーニングを 1) 加速および 2) 安定化するための新しい方法を紹介します。 1) トレーニングの高速化に向けて、1B パラメータ CLIP ViT-Huge の bfloat16 トレーニングのパフォーマンスを 0.1 パーセンテージ ポイント以内で一致させながら、13-25% のスピードアップを提供する int8 量子化トレーニング用の線形レイヤーである SwitchBack を導入します。これまでで最大の int8 トレーニング。 GPU が float8 をサポートすることはめったにないため、主な焦点は int8 ですが、シミュレーションを通じて float8 のトレーニングも分析します。 SwitchBack は float8 で効果的であることが証明されていますが、ネットワークがトレーニングおよび初期化されている場合、標準的な手法も成功することを示しています。 2) 安定したトレーニングに向けて、損失スパイクを分析し、2 乗勾配が AdamW の 2 次モーメント推定器によって過小評価された後、1 ~ 8 回の反復で一貫して発生することを発見しました。その結果、AdamW-Adafactor ハイブリッドをお勧めします。これは、CLIP ViT-Huge モデルをトレーニングする際の損失スパイクを回避し、勾配クリッピングよりも優れているため、StableAdamW と呼ばれます。
We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vision models. 1) Towards accelerating training, we introduce SwitchBack, a linear layer for int8 quantized training which provides a speed-up of 13-25% while matching the performance of bfloat16 training within 0.1 percentage points for the 1B parameter CLIP ViT-Huge -- the largest int8 training to date. Our main focus is int8 as GPU support for float8 is rare, though we also analyze float8 training through simulation. While SwitchBack proves effective for float8, we show that standard techniques are also successful if the network is trained and initialized so that large feature magnitudes are discouraged, which we accomplish via layer-scale initialized with zeros. 2) Towards stable training, we analyze loss spikes and find they consistently occur 1-8 iterations after the squared gradients become under-estimated by their AdamW second moment estimator. As a result, we recommend an AdamW-Adafactor hybrid, which we refer to as StableAdamW because it avoids loss spikes when training a CLIP ViT-Huge model and outperforms gradient clipping.