個人 Copilot：訓練你自己的編碼助手
在不斷發展的程式設計和軟體開發領域，對效率和生產力的追求催生了非凡的創新。其中一項創新便是程式碼生成模型的出現，例如 Codex、StarCoder 和 Code Llama。這些模型在生成類人程式碼片段方面展現了卓越的能力，從而顯示出作為編碼助手的巨大潛力。
然而，雖然這些預訓練模型可以在一系列任務上表現出色，但一個激動人心的可能性正悄然出現：根據你的特定需求定製程式碼生成模型。想象一下，可以在企業範圍內利用個性化的編碼助手。
在這篇博文中，我們將展示我們如何建立了 HugCoder 🤗，一個在 huggingface GitHub 組織的公開倉庫程式碼內容上微調的程式碼 LLM。我們將討論我們的資料收集工作流、訓練實驗以及一些有趣的結果。這將使你能夠基於自己的專有程式碼庫建立個人 copilot。我們還會為你留下幾個該專案的進一步擴充套件方向以供實驗。
讓我們開始吧 🚀
資料收集工作流
我們期望的資料集在概念上很簡單，我們這樣構建它：
| 倉庫名稱 | 倉庫中的檔案路徑 | 檔案內容 |
| --- | --- | --- |
| --- | --- | --- |
使用 Python GitHub API 從 GitHub 抓取程式碼內容非常直接。但是，根據倉庫數量和倉庫內程式碼檔案的數量，人們可能很容易遇到 API 速率限制問題。
為了防止此類問題，我們決定將所有公共倉庫克隆到本地，並從本地提取內容，而不是透過 API。我們使用了 Python 的 multiprocessing 模組來並行下載所有倉庫，如此下載指令碼所示。
一個倉庫通常可能包含非程式碼檔案，如影像、簡報和其他資產。我們對抓取這些不感興趣。我們建立了一個副檔名列表來將它們過濾掉。對於解析除 Jupyter Notebook 之外的程式碼檔案，我們簡單地使用了“utf-8”編碼。對於 notebook，我們只考慮程式碼單元格。
我們還排除了所有與程式碼不直接相關的檔案路徑。這些包括：.git、__pycache__ 和 xcodeproj。
為了使這些內容的序列化相對記憶體友好，我們使用了分塊和 feather 格式。有關完整實現，請參閱此指令碼。
最終的資料集可在 Hub 上找到，它看起來像這樣：
對於這篇部落格，我們考慮了基於 stargazers 數量排名前 10 的 Hugging Face 公開倉庫。它們是以下這些：
['transformers', 'pytorch-image-models', 'datasets', 'diffusers', 'peft', 'tokenizers', 'accelerate', 'text-generation-inference', 'chat-ui', 'deep-rl-class']
這是我們用來生成該資料集的程式碼，而這是 Hub 上的資料集。以下是它的一個快照：
為了降低專案複雜性，我們沒有考慮對資料集進行去重。如果你有興趣為生產應用應用去重技術，這篇博文是在程式碼 LLM 背景下關於該主題的絕佳資源。
微調你自己的個人 Copilot
在本節中，我們將展示如何微調以下模型：bigcode/starcoder (155 億引數)、bigcode/starcoderbase-1b (10 億引數)、Deci/DeciCoder-1b (10 億引數)。我們將在所有實驗中使用單個 A100 40GB Colab Notebook，並採用 🤗 PEFT（引數高效微調）。此外，我們還將展示如何在一臺擁有 8 個 A100 80GB GPU 的機器上，使用 🤗 Accelerate 的 FSDP 整合來完全微調 bigcode/starcoder (155 億引數) 模型。訓練目標是中間填充 (FIM)，即將訓練序列的某些部分移到末尾，然後自迴歸地預測重新排序後的序列。
為什麼選擇 PEFT？全量微調成本高昂。讓我們用一些數字來直觀感受一下：
全量微調所需的最低 GPU 記憶體
- 權重：2 位元組 (混合精度訓練)
- 權重梯度：2 位元組
- 使用 Adam 最佳化器時的狀態：原始 FP32 權重需要 4 位元組 + 一階和二階矩估計需要 8 位元組
- 將上述所有相加，每個引數的成本：每個引數 16 位元組
- 155 億引數模型 -> 248GB GPU 記憶體，這還不包括儲存中間啟用所需的大量記憶體 -> 至少需要 4 塊 A100 80GB GPU
由於硬體要求巨大，我們將使用QLoRA進行引數高效微調。以下是使用 QLoRA 微調 StarCoder 的最低 GPU 記憶體要求：
可訓練引數：110,428,160 || 總引數：15,627,884,544 || 可訓練百分比：0.7066097761926236
- 基礎模型權重：0.5 位元組 * 155.1 億凍結引數 = 7.755 GB
- 介面卡權重：2 位元組 * 1.1 億可訓練引數 = 0.22GB
- 權重梯度：2 位元組 * 1.1 億可訓練引數 = 0.12GB
- 使用 Adam 最佳化器時的狀態：4 位元組 * 1.1 億可訓練引數 * 3 = 1.32GB
- 將以上所有相加 -> 9.51 GB ~ 10GB -> 需要 1 塊 A100 40GB GPU 🤯。之所以需要 A100 40GB GPU，是因為在訓練中，對於 2048 的長序列長度和 4 的批次大小，中間啟用會佔用更高的記憶體。如下所示，所需的 GPU 記憶體為 26GB，可以在 A100 40GB GPU 上容納。此外，A100 GPU 與 Flash Attention 2 的相容性更好。
在上述計算中，我們沒有考慮中間啟用檢查點所需的記憶體，這部分記憶體相當大。我們利用 Flash Attention V2 和梯度檢查點來克服這個問題。
- 對於 QLoRA 加上 Flash Attention V2 和梯度檢查點，在單個 A100 40GB GPU 上，模型佔用的總記憶體為 26 GB，批次大小為 4。
- 對於使用 FSDP 加上 Flash Attention V2 和梯度檢查點的全量微調，每塊 GPU 佔用的記憶體介於 70 GB 到 77.6 GB 之間，每 GPU 批次大小為 1。
請參考模型記憶體使用工具，輕鬆計算在 🤗 Hugging Face Hub 上託管的模型進行訓練和大型模型推理需要多少 vRAM。
全量微調
我們將探討如何使用 PyTorch 全分片資料並行（FSDP）技術，在 8 塊 A100 80GB GPU 上對 bigcode/starcoder（150 億引數）進行全量微調。有關 FSDP 的更多資訊，請參考使用 PyTorch FSDP 微調 Llama 2 70B 和使用 PyTorch 全分片資料並行加速大模型訓練。
資源
- 程式碼庫：連結。它使用了 Transformers 中最近新增的 Flash Attention V2 支援。
- FSDP 配置：fsdp_config.yaml
- 模型：bigcode/stacoder
- 資料集：smangrul/hf-stack-v1
- 微調後的模型：smangrul/peft-lora-starcoder15B-v2-personal-copilot-A100-40GB-colab
啟動訓練的命令在 run_fsdp.sh 中給出。
accelerate launch --config_file "configs/fsdp_config.yaml" train.py \
--model_path "bigcode/starcoder" \
--dataset_name "smangrul/hf-stack-v1" \
--subset "data" \
--data_column "content" \
--split "train" \
--seq_length 2048 \
--max_steps 2000 \
--batch_size 1 \
--gradient_accumulation_steps 2 \
--learning_rate 5e-5 \
--lr_scheduler_type "cosine" \
--weight_decay 0.01 \
--num_warmup_steps 30 \
--eval_freq 100 \
--save_freq 500 \
--log_freq 25 \
--num_workers 4 \
--bf16 \
--no_fp16 \
--output_dir "starcoder-personal-copilot-A100-40GB-colab" \
--fim_rate 0.5 \
--fim_spm_rate 0.5 \
--use_flash_attn
總訓練時間為 9 小時。根據 lambdalabs 的價格，8 塊 A100 80GB GPU 的成本為每小時 12.00 美元，總成本將為 108 美元。
PEFT
我們將探討如何使用 QLoRA 在單個 A100 40GB GPU 上使用 🤗 PEFT 對 bigcode/starcoder（150 億引數）進行微調。有關 QLoRA 和 PEFT 方法的更多資訊，請參考使用 bitsandbytes、4 位量化和 QLoRA 讓 LLM 更易於使用和🤗 PEFT：在低資源硬體上對十億級模型進行引數高效微調。
資源
- 程式碼庫：連結。它使用了 Transformers 中最近新增的 Flash Attention V2 支援。
- Colab notebook：連結。請確保選擇 A100 GPU 並設定高記憶體。
- 模型：bigcode/stacoder
- 資料集：smangrul/hf-stack-v1
- QLoRA 微調模型：smangrul/peft-lora-starcoder15B-v2-personal-copilot-A100-40GB-colab
啟動訓練的命令在 run_peft.sh 中給出。總訓練時間為 12.5 小時。根據 lambdalabs 的價格，每小時成本為 1.10 美元，總成本將為 13.75 美元。這相當不錯 🚀！在成本方面，比全量微調低 7.8 倍。
比較
下圖顯示了 QLoRA 與全量微調的評估損失、訓練損失和學習率排程器。我們觀察到，全量微調的損失略低，收斂速度比 QLoRA 稍快。peft 微調的學習率是全量微調的 10 倍。
為了確保我們的 QLoRA 模型不會導致災難性遺忘，我們在其上運行了 Python Human Eval。以下是我們得到的結果。Pass@1 衡量的是每個問題只考慮一個生成的程式碼候選時的透過率。我們可以觀察到，在 humaneval-python 上的效能，基礎模型 bigcode/starcoder (150 億引數) 和微調後的 PEFT 模型 smangrul/peft-lora-starcoder15B-v2-personal-copilot-A100-40GB-colab 之間是相當的。
| 模型 | Pass@1 |
| bigcode/starcoder | 33.57 |
| smangrul/peft-lora-starcoder15B-v2-personal-copilot-A100-40GB-colab | 33.37 |
現在讓我們看一些定性樣本。在我們的手動分析中，我們注意到 QLoRA 導致了輕微的過擬合，因此我們透過 PEFT 的 add_weighted_adapter 工具建立了一個權重為 0.8 的新加權介面卡來降低其影響。
我們將看兩個程式碼填充的例子，其中模型的任務是填充由 <FILL_ME> 佔位符表示的部分。我們將考慮 GitHub Copilot、QLoRA 微調模型和全量微調模型的填充結果。
定性示例 1
在上面的例子中，GitHub Copilot 的補全方向是正確的，但幫助不大。另一方面，QLoRA 和全量微調模型的補全正確地填充了整個函式呼叫以及必要的引數。然而，它們之後也增加了很多噪音。這可以通過後處理步驟來控制，將補全限制在閉合括號或換行符處。請注意，QLoRA 和全量微調模型都產生了質量相似的結果。
定性示例 2
在上面的第二個例子中，GitHub Copilot 沒有給出任何補全。這可能是因為 🤗 PEFT 是一個較新的庫，尚未成為 Copilot 訓練資料的一部分，這正是我們試圖解決的問題。另一方面，QLoRA 和全量微調模型的補全正確地填充了整個函式呼叫以及必要的引數。再次注意，QLoRA 和全量微調模型都給出了質量相似的生成結果。全量微調模型和 peft 模型的帶有各種示例的推理程式碼分別在 Full_Finetuned_StarCoder_Inference.ipynb 和 PEFT_StarCoder_Inference.ipynb 中提供。
因此，我們可以觀察到兩種變體的生成結果都符合預期。太棒了！🚀
如何在 VS Code 中使用它？
你可以使用 🤗 llm-vscode VS Code 擴充套件，並結合 🤗 Inference EndPoints 託管模型，輕鬆在 VS Code 中配置自定義程式碼補全 LLM。我們將在下面介紹所需的步驟。你可以在推理端點文件中瞭解有關部署端點的更多詳細資訊。
設定推理端點
以下是我們建立自定義推理端點所遵循的步驟截圖。我們使用了我們的 QLoRA 模型，將其匯出為完整的*合併*模型，以便在 transformers 中輕鬆載入。
設定 VS Code 擴充套件
只需按照安裝步驟操作。在設定中，替換下面欄位中的端點，使其指向你部署的 HF 推理端點。
使用起來會像下面這樣：
微調你自己的程式碼聊天助手
到目前為止，我們訓練的模型都是專門為程式碼補全任務訓練的個人 copilot。它們沒有經過訓練來進行對話或問答。Octocoder 和 StarChat 是這類模型的優秀例子。本節簡要介紹如何實現這一點。
資源
- 程式碼庫：連結。它使用了 Transformers 中最近新增的 Flash Attention V2 支援。
- Colab notebook：連結。請確保選擇 A100 GPU 並設定高記憶體。
- 模型：bigcode/stacoderplus
- 資料集：smangrul/code-chat-assistant-v1。混合了
LIMA+GUANACO，並以可直接訓練的格式進行了適當的格式化。 - 訓練好的模型：smangrul/peft-lora-starcoderplus-chat-asst-A100-40GB-colab
LoRA 之舞
如果你曾涉足 Stable Diffusion 模型和 LoRA 來製作你自己的 Dreambooth 模型，你可能熟悉將不同 LoRA 以不同權重組合，或者將一個 LoRA 模型用於與其訓練時不同的基礎模型的概念。在文字/程式碼領域，這仍然是未被探索的領域。我們在這方面進行了實驗，並觀察到了非常有前景的發現。你準備好了嗎？我們開始吧！🚀
混合搭配 LoRA
PEFT 目前支援 3 種組合 LoRA 模型的方式：linear、svd 和 cat。更多詳情請參考 tuners#peft.LoraModel.add_weighted_adapter。
我們的 notebook Dance_of_LoRAs.ipynb 包含了所有的推理程式碼和各種 LoRA 載入組合，比如在 starcoder 上載入聊天助手而不是我們微調時使用的基礎模型 starcodeplus。
在這裡，我們將考慮 2 種能力（聊天/問答和程式碼補全）在 2 種資料分佈（前 10 名 hf 公開程式碼庫和通用程式碼庫）上的表現。這給了我們 4 個維度來進行一些定性評估分析。
首先，讓我們考慮聊天/問答任務。
如果我們停用介面卡，我們觀察到該任務在兩個資料集上都失敗了，因為基礎模型（starcoder）只用於程式碼補全，不適合聊天/問答。啟用 copilot 介面卡的表現與停用情況類似，因為這個 LoRA 也是專門為程式碼補全而微調的。
現在，讓我們啟用 assistant 介面卡。
基於通用程式碼的問答
基於 HF 程式碼的問答
我們可以觀察到，關於 scrapy 的通用問題得到了正確的回答。然而，對於 HF 程式碼相關的問題，它卻失敗了，因為這部分內容不在其預訓練資料中。
現在我們來考慮一下程式碼補全任務。
停用介面卡後，我們觀察到通用的 two-sum 程式碼補全按預期工作。然而，HF 程式碼補全失敗，給出了錯誤的 LoraConfig 引數，因為基礎模型在其預訓練資料中沒有見過它。啟用 assistant 的表現與停用情況類似，因為它是在自然語言對話上訓練的，其中不包含任何 Hugging Face 程式碼庫。
現在，讓我們啟用 copilot 介面卡。
我們可以觀察到，copilot 介面卡在兩種情況下都正確了。因此，它在處理 HF 特定程式碼庫以及通用程式碼庫時，都能按預期進行程式碼補全。
現在，作為使用者，我希望結合 assistant 和 copilot 的能力。這將使我能夠在 IDE 中編碼時使用它進行程式碼補全，同時也可以將它作為一個聊天機器人來回答我關於 API、類、方法、文件的問題。它應該能夠回答諸如“我該如何使用 x”、“請為 Y 寫一個程式碼片段”之類關於我的程式碼庫的問題。
PEFT 允許你透過 add_weighted_adapter 來實現這一點。讓我們建立一個新的介面卡 code_buddy，給予 assistant 和 copilot 介面卡相等的權重。
組合多個介面卡
現在，讓我們看看 code_buddy 在聊天/問答任務上的表現如何。
我們可以觀察到，code_buddy 的表現比單獨使用 assistant 或 copilot 介面卡要好得多！它能夠回答“寫一個程式碼片段”的請求，以展示如何使用特定的 HF 倉庫 API。然而，它也會幻覺出錯誤的連結/解釋，這仍然是 LLM 面臨的一個開放挑戰。
以下是 code_buddy 在程式碼補全任務上的表現。
我們可以觀察到，code_buddy 的表現與專門為此任務微調的 copilot 相當。
將 LoRA 遷移到不同的基礎模型
我們也可以將 LoRA 模型遷移到不同的基礎模型上。我們將採用新鮮出爐的 Octocoder 模型，並將我們上面用 starcoder 基礎模型訓練的 LoRA 應用於其上。請參閱以下 notebook PEFT_Personal_Code_CoPilot_Adapter_Transfer_Octocoder.ipynb 以獲取完整程式碼。
在程式碼補全任務上的表現
我們可以觀察到 `octocoder` 的表現非常出色。它能夠補全 HF 特定的程式碼片段。正如 notebook 中所示，它也能夠補全通用的程式碼片段。
在聊天/問答任務上的表現
由於 Octocoder 經過訓練可以回答問題並進行編碼相關的對話，讓我們看看它是否能使用我們的 LoRA 介面卡來回答 HF 特定的問題。
太棒了！它詳細地正確回答瞭如何建立 LoraConfig 和相關的 peft 模型，並正確使用了模型名稱、資料集名稱以及 LoraConfig 的引數值。停用介面卡後，它無法正確使用 LoraConfig 的 API 或建立 PEFT 模型，這表明這部分內容不在 Octocoder 的訓練資料中。
如何在本地執行它？
我知道，經過這一切，你想在你的程式碼庫上微調 starcoder，並在你的消費級硬體上本地使用它，比如帶 M1 GPU 的 Mac 筆記本、帶 RTX 4090/3090 GPU 的 Windows 電腦……別擔心，我們已經為你準備好了。
我們將使用這個超酷的開源庫 mlc-llm 🔥。具體來說，我們將使用這個 fork pacman100/mlc-llm，它做了一些修改以使其能與 VS Code 的 Hugging Face 程式碼補全擴充套件一起工作。在我的帶 M1 Metal GPU 的 Mac 筆記本上，15B 模型速度非常慢。因此，我們將選擇小一點的模型，訓練一個 PEFT LoRA 版本以及一個 bigcode/starcoderbase-1b 的全量微調版本。訓練的 colab notebook 連結如下：
starcoderbase-1b全量微調和 PEFT LoRA 微調的 Colab notebook：連結
訓練損失、評估損失以及學習率排程如下所示：
現在，我們將詳細介紹如何在本地託管合併後的模型 smangrul/starcoder1B-v2-personal-copilot-merged，並將其與 🤗 llm-vscode VS Code 擴充套件一起使用。
- 克隆倉庫
git clone --recursive https://github.com/pacman100/mlc-llm.git && cd mlc-llm/
- 安裝 mlc-ai 和 mlc-chat（以可編輯模式）
pip install --pre --force-reinstall mlc-ai-nightly mlc-chat-nightly -f https://mlc.ai/wheels
cd python
pip uninstall mlc-chat-nightly
pip install -e "."
- 透過以下命令編譯模型
time python3 -m mlc_llm.build --hf-path smangrul/starcoder1B-v2-personal-copilot-merged --target metal --use-cache=0
- 在
dist/starcoder1B-v2-personal-copilot-merged-q4f16_1/params/mlc-chat-config.json中用以下值更新配置
{
"model_lib": "starcoder7B-personal-copilot-merged-q4f16_1",
"local_id": "starcoder7B-personal-copilot-merged-q4f16_1",
"conv_template": "code_gpt",
- "temperature": 0.7,
+ "temperature": 0.2,
- "repetition_penalty": 1.0,
"top_p": 0.95,
- "mean_gen_len": 128,
+ "mean_gen_len": 64,
- "max_gen_len": 512,
+ "max_gen_len": 64,
"shift_fill_factor": 0.3,
"tokenizer_files": [
"tokenizer.json",
"merges.txt",
"vocab.json"
],
"model_category": "gpt_bigcode",
"model_name": "starcoder1B-v2-personal-copilot-merged"
}
- 執行本地伺服器
python -m mlc_chat.rest --model dist/starcoder1B-v2-personal-copilot-merged-q4f16_1/params --lib-path dist/starcoder1B-v2-personal-copilot-merged-q4f16_1/starcoder1B-v2-personal-copilot-merged-q4f16_1-metal.so
- 在 VS Code 中更改 HF 程式碼補全擴充套件的端點，使其指向本地伺服器
- 在 VS Code 中開啟一個新檔案，貼上下面的程式碼，並將游標放在文件字串引號之間，以便模型嘗試填充文件字串
搞定！⭐️
本文開頭的演示就是這個 1B 模型在我的 Mac 筆記本上本地執行。
結論
在這篇博文中，我們看到了如何微調 starcoder 來建立一個瞭解我們程式碼的個人 copilot。我們稱之為 🤗 HugCoder，因為我們是在 Hugging Face 的程式碼上訓練它的 :) 在介紹了資料收集工作流之後，我們比較了使用 QLoRA 和全量微調的訓練。我們還透過組合不同的 LoRA 進行了實驗，這在文字/程式碼領域仍然是一種未被充分探索的技術。在部署方面，我們研究了使用 🤗 Inference Endpoints 進行遠端推理，並展示了在 VS Code 和 MLC 上執行一個較小模型的裝置端執行。
如果你在自己的程式碼庫上使用這些方法，請告訴我們！
致謝
我們要感謝 Pedro Cuenca、Leandro von Werra、Benjamin Bossan、Sylvain Gugger 和 Loubna Ben Allal 在撰寫這篇博文時提供的幫助。
