訓練大型AI模型（如Transformer和語言模型）已成爲AI領域不可或缺的關鍵環節，但也面臨着高昂的計算成本、內存消耗和能源需求。例如，OpenAI的GPT-3擁有1750億個參數，需要數週的GPU訓練。這種巨大的資源需求限制了這項技術在大規模計算資源充足的組織中的應用，同時也加劇了人們對能源效率和環境影響的擔憂。解決這些挑戰對於確保AI發展的更廣泛可及性和可持續性至關重要。

傳統訓練方法效率低下，亟需創新解決方案

大型模型訓練效率低下的主要原因在於其對稠密矩陣的依賴，這需要大量的內存和計算能力。現代GPU對優化的低精度或低秩操作支持有限，進一步加劇了這些需求。儘管已經提出了一些方法，如矩陣分解和啓發式秩降低，來緩解這些問題，但它們在實際應用中仍然受到限制。例如，GaLore能夠在單批次設置下進行訓練，但存在不切實際的運行時開銷。同樣，採用低秩適配器的LTE在大型任務上的收斂性方面也存在問題。目前缺乏一種能夠同時降低內存使用、計算成本和訓練時間，而不損害性能的方法，這使得創新解決方案的需求變得迫切。

image.png

CoMERA框架:通過自適應張量優化實現高效訓練

來自奧爾巴尼大學（紐約州立大學）、加州大學聖巴巴拉分校、亞馬遜Alexa AI和Meta的研究人員共同推出了一種名爲CoMERA(Computing-and Memory-Efficient training method via Rank-Adaptive tensor optimization)的新型框架。該框架結合了內存效率和計算速度，通過自適應秩張量壓縮技術來實現。與傳統方法僅關注壓縮不同，CoMERA採用多目標優化方法來平衡壓縮比和模型精度。它利用張量化嵌入和先進的張量網絡收縮來優化GPU利用率，從而減少運行時開銷，同時保持強大的性能。該框架還引入了CUDA圖，以最大程度地減少GPU操作期間的內核啓動延遲，這是傳統張量壓縮方法中的一個主要瓶頸。

image.png

CoMERA的基礎是自適應張量表示，它允許模型層根據資源約束動態調整其秩。通過修改張量秩，該框架可以在不損害神經網絡操作完整性的情況下實現壓縮。這種動態優化是通過一個兩階段的訓練過程實現的:

早期階段:專注於穩定收斂。

後期階段:微調秩以滿足特定的壓縮目標。

在一個六編碼器Transformer模型中，CoMERA在其早期階段實現了高達43倍的壓縮比，而在其後期優化階段，壓縮比更是高達361倍。此外，與GaLore相比，它將內存消耗降低了9倍，每輪訓練速度提高了2-3倍。

多項測試結果表明CoMERA性能卓越

在應用於MNLI數據集上訓練的Transformer模型時，CoMERA將模型大小從256MB縮小到低至3.2MB，同時保持了精度。在諸如DLRM的大規模推薦系統中，CoMERA將模型壓縮了99倍，並使峯值內存使用量減少了7倍。該框架還在預訓練CodeBERT（一個特定領域的大型語言模型）方面表現出色，獲得了4.23倍的整體壓縮比，並在某些訓練階段實現了2倍的加速。這些結果突顯了其處理各種任務和架構的能力，擴展了其在各個領域的適用性。

CoMERA框架的關鍵優勢總結

這項研究的主要結論如下:

CoMERA爲特定層實現了高達361倍的壓縮比，爲整個模型實現了99倍的壓縮比，大大降低了存儲和內存需求。

該框架將Transformer和推薦系統的每輪訓練時間縮短了2-3倍，節省了計算資源和時間。

通過使用張量化表示和CUDA圖，CoMERA將峯值內存消耗減少了7倍，使得在較小的GPU上進行訓練成爲可能。

CoMERA的方法支持包括Transformer和大型語言模型在內的多種架構，同時保持或提高了精度。

通過降低訓練所需的能源和資源，CoMERA有助於實現更可持續的AI實踐，並使更廣泛的受衆能夠使用尖端模型。