在這個AI大爆炸的時代，大型語言模型（LLMs）已經成爲推動機器學習應用的超級引擎。但是，訓練這些龐然大物需要巨大的計算資源。想象一下，如果我們能在世界各地分散的設備上，高效地訓練這些模型，那將會怎樣?這就是OpenDiLoCo帶給我們的驚喜。

QQ截圖20240715114010.jpg

傳統的分佈式訓練方法，需要頻繁的通信和大量的帶寬，這限制了訓練的規模和效率。而DiLoCo（分佈式低通信）訓練方法，通過減少通信需求，讓LLMs的全球訓練成爲可能。

OpenDiLoCo是一個開源框架，它實現了DiLoCo訓練方法，並用Hivemind庫提供了一個可擴展的、去中心化的訓練框架。這個框架的厲害之處在於，它在全球範圍內，跨越兩大洲、三個國家進行了模型訓練，同時保持了90-95%的計算利用率。

主要特性:

  • 資源的動態增減:訓練過程中，可動態調整計算資源，新設備和集羣可以在訓練中途加入或退出。

  • 容錯性:去中心化訓練中，某些設備可能不可靠。通過Hivemind的容錯訓練，即使設備隨時不可用，訓練過程也不會停止。

  • 點對點通信:沒有主節點，所有通信均以點對點方式進行。

研究人員們不僅復現了DiLoCo的實驗，還將其擴展到了十億參數模型的規模。他們通過消融研究，展示了DiLoCo算法在計算效率和擴展性方面的優勢。更令人印象深刻的是，他們證明了DiLoCo的梯度可以在不降低性能的情況下，使用FP16進行全歸約。

QQ截圖20240715114048.jpg

主要貢獻:

  • 複製與擴展:成功復現了DiLoCo的原始實驗，並將其擴展到十億參數模型的規模。

  • 開源實現:基於Hivemind庫構建的可擴展實現，使去中心化訓練對廣泛的開發者和研究者開放。

  • 全球去中心化訓練:通過在兩大洲和三個國家之間的模型訓練，展示了OpenDiLoCo的實際潛力，同時保持了90-95%的計算利用率。

  • 效率洞察:通過消融研究提供了算法可擴展性和計算效率的寶貴見解。

實驗結果:

通過複製DiLoCo的主要實驗結果，Prime Intellect公司證明了其方法的有效性。在C4數據集上使用150M參數的模型進行語言建模任務的訓練，展示了DiLoCo在通信需求降低500倍的情況下，與基線性能相當。

DeepMind的原始DiLoCo論文僅針對最多4億參數的模型進行了實驗。在這項工作中，Prime Intellect公司將該方法擴展到了1.1億參數的模型，並採用了與TinyLlama相同的超參數。

爲了展示OpenDiLoCo在不同大陸的去中心化訓練功能，實驗使用了位於加拿大、芬蘭和美國兩個不同州的四個DiLoCo工作節點，每個節點都配備了八個H100GPU。

Prime Intellect公司成功復現了DiLoCo的主要實驗結果，並將該方法擴展到原始工作參數大小的三倍，並在現實世界的去中心化訓練環境中展示了其應用。

未來，該公司計劃將DiLoCo擴展到更多分佈式工作節點上的更大模型，並探索可能提高穩定性和收斂速度的模型合併技術，以及通過實施異步權重平均通信方法來減少計算空閒時間。

論文地址：https://arxiv.org/pdf/2407.07852