要對代碼下手了，加油(? ?_?)?
作者在他們自己的設備上展現了推理的評估結果，受第三方評估認證，EAGLE為目前最快的投機方法（雖然加速度是評估投機解碼方法的主要指標，但其他點也值得關注。比如PLD和Lookahead無需額外參數，更容易和許多模型進行集成），所有用來評估的方法都和Spec-Bench對齊。
設備：一臺NVIDIA GeForce RTX 3090 GPU(24GB) ，帶12個CPU核
測試環境：Pytorch 2.0.1，CUDA 11.8
環境設置：Vicuna-7B-v1.3，貪心解碼，FP16精度，批量大小為1
EAGLE-2會利用草稿模型打出的信心分數去近似接受率，動態調整草稿的樹形架構，進一步提升性能。EAGLE-2在兩塊RTX 3060 GPU上的推理速度，比原始的投機解碼在一塊A100 GPU上的推理速度要快。
設備：一臺NVIDIA A100 GPU(80GB)?，帶64個CPU核（比V100提升20倍的AI計算性能，實驗室沒有－_－，相較于RTX 4090、A40，更適合大模型AI訓練）（看人家怎么說的，在8張RTX 3090 GPU上也能訓練，1-2天完事兒，“so even the GPU poor can afford it”）（下面3張圖看著樂呵一下，俺只有3090(≧?≦)ゞ）
測試環境：Pytorch 2.0.1，CUDA 11.4
實驗設置：貪心解碼，FP16精度，批量大小為1
?
?
EAGLE已被集成到許多主流的LLM服務框架中了，比如Intel Extension for Transformers、vLLM、SGLang等等。
相應GitHub庫的更新如下：2023.12.8發布EAGLE v1.0，2024.1.17起支持Mixtral-8x7B-Instruct、2024.2.25經第三方評估認證為最快的投機方法、2024.6.27發布EAGLE-2、2024.8.8起支持QWen-2（阿里巴巴集團QWen團隊開發的第二代LLM系列，旨在提升自然語言處理和生成任務的性能）。不久的將來將發布EAGLE-3（2025.3.19更新：太牛了，我代碼還沒復現完它就出了）。
我選用的基礎模型主要是Llama-2-chat-7B，該模型是由Meta（原Facebook）基于Transformer架構開發的開源LLM，有70億參數，屬于Llama-2系列中的較小版本，專為對話任務微調，適合交互式應用。
?配置并安裝
git clone https://github.com/SafeAILab/EAGLE.git
cd EAGLE
pip install -r requirements.txt在安裝requirements.txt中的包時，要使用合適的python版本，比如我試過python3.11報錯說沒有“python3.11/site-packages/torch/include/torch/目錄”，而3.9可以絲滑安裝。
而且需要提前設置一下pip鏡像，參考博客，不然慢不慢另說，還可能找不到某些版本的包而發生報錯。
為什么要cd到這個目錄呢？其實后面可以發現，許多文件里到處都是相對路徑（🤯）也不是說一定要在這個目錄吧，只是要改成對應的正確路徑。
下載EAGLE權重
作者提供了各式目標模型對應的EAGLE參數的Hugging Face網址。對于在上面托管的模型，點開hugging face界面的“Use this model”-》Transformers庫，一般可以看到兩種使用transformer庫的加載方法（下圖microsoft/DialoGPT-small對應的界面這樣）
一種使用高級封裝pipeline，自動處理tokenization（文本切分）、模型推理、解碼這些步驟，簡單不靈活；一種手動加載模型和分詞器，可自由調整參數，可擴展性強，代碼復雜。?
1. 使用pipeline作為高級封裝
# Use a pipeline as a high-level helper
from transformers import pipelinepipe = pipeline("text-generation", model="yuhuili/EAGLE-Vicuna-7B-v1.3")
pipeline是transformer庫提供的一個高級封裝，可自動處理tokenization（文本切分）、模型推理、解碼這些步驟；"text-generation"任務代表使用自回歸文本生成模型，比如GPT類模型（Vicuna也在其中）；model="yuhuili/EAGLE-Vicuna-7B-v1.3"告訴pipeline下載該模型。加載后使用方法就是
result=pipe("Hello, how are you?")
print(result)pipe()直接輸入文本，返回生成的文本結果，通常是一個包含文本輸出的列表。簡單易用，適用于對性能要求不高時的快速部署。但缺乏靈活性，無法自定義tokenizer或model的參數，例如溫度、最大長度等，且pipeline默認會嘗試自動優化加載方式，可能消耗額外顯存。
2. 手動加載模型和分詞器
# Load model directly
from transformers import AutoTokenizer, AutoModelForCausalLMtokenizer = AutoTokenizer.from_pretrained("yuhuili/EAGLE-Vicuna-7B-v1.3")
model = AutoModelForCausalLM.from_pretrained("yuhuili/EAGLE-Vicuna-7B-v1.3")
tokenizer分詞器負責將輸入的文本轉換為token，用于模型計算，并將模型的輸出轉換回人類可讀的文本。從Hugging Face服務器下載該模型的權重和架構時，AutoModelForCausalLM適用于因果語言模型，如GPT、Vicuna這類基于自回歸生成的Transformer模型。使用時需要手動處理輸入輸出
input_text = "Hello, how are you?"
input_ids = tokenizer(input_text, return_tensors="pt").input_ids # 將文本轉換為 token ID（張量格式）
output = model.generate(input_ids, max_length=50) # 讓模型生成文本，指定最大生成長度為50
result = tokenizer.decode(output[0], skip_special_tokens=True) # 將生成的 token ID序列轉換成字符串
print(result)優點是高度可控，可以自由調整generate()里的參數，如溫度、top_k、top_p等，提高生成質量，可以優化顯存占用（如啟用torch_dtype=torch.float16或device_map="auto"），適合大規模任務，可擴展性強，可與LoRA、DeepSpeed、FSDP等優化技術結合，可擴展性強。但代碼復雜度高，默認不會自動優化顯存占用，加載大模型時可能會超出GPU負載。?
但是，要真直接運行這代碼，常會報錯(≧?≦)ゞ，比如EAGLE系列模型的就說沒分詞器啦~聊天小模型microsoft/DialoGPT-small就說輸入內容格式不對啦~因為這代碼是網站自動生成的，適用于通常情況。就以yuhuili/EAGLE-Vicuna-7B-v1.3的情況來說吧，點開“Files and versions”，下面確實沒有分詞器（可以拿上圖比對一下）
咱就下載一下EAGLE weight🙄，直接運行下邊這個得了，下載條拉滿即可。
git lfs clone https://huggingface.co/yuhuili/EAGLE-llama2-chat-7B?另外作者說了一下，目標模型是QWen時，應采用bf16而非fp16以避免數字溢出，兩者都是16位浮點數格式，但指數和位數分布分配不同，BF16為8位指數7位尾數，數值范圍接近FP32，FP16為5位指數10位尾數。草稿模型的訓練數據集為ShareGPT，數據全英文，如果想要將其用在非英文，比如中文的數據上，需要用相應的數據進行訓練。在EAGLE的基礎上，EAGLE-2無需額外的訓練，直接使用相同的權重。倉庫提供的推理代碼會自動分配模型權重，在多個GPU上加載模型，使得超出單個GPU內存的模型也能跑起來。（自動分布式嗎？還挺牛的！看代碼咋實現）
用UI體驗
作者提供了網絡端口，運行下列命令即可啟動（damn！還得事先下載llama2-chat-7B）。等模型被完全加載，終端就會輸出網址，點進去就跳到瀏覽器體驗啦。（好好好顯存不夠，換臺新的服務器。每當俺在一臺新服務器上注冊賬號后，需要設置免密登錄，安裝Miniconda、pip，禁止自動激活base環境、安裝擴展，scp文件，有時還得加速這個加速內個🙄）（還得小心代碼中給你指定設備的情況）
python -m eagle.application.webui --ea-model-path [path of EAGLE weight]\ --base-model-path [path of the original model]\--model-type ["llama-2-chat","vicuna","mixtral","llama-3-instruct"]\--total-token [int]total-token這個選項代表草稿token的數量，如果用的是較小的模型或先進的GPUs（沒有:p）的話這個值可以大點。反正就根據具體的硬件和模型做調整，以達更好的性能吧。設成-1的話EAGLE-2會自動配置的（😭模型啥都能自動了而本菜雞怎么活）。
3090GPU的服務器掛掉了，我換了臺A40 GPU的服務器，這時倉庫提供的requirement.txt中，torch和accelerate版本（2.0.1和0.21.0）會出現“設備映射”的問題，代碼里有的地方又用的舊版本，而且輕易更換版本容易沖突，總之各種問題！再換臺服務器吧😭
代碼運行報錯說讓你安裝什么包、什么版本，就照著報錯提示去安裝得了。
TODO: 怪我寫的太慢，eagle3都出來了我還在這磨蹭！注意到eagle3中添加了一個參數”draft_vocab_size“，這難道可以控制草稿長度？后面看看！然后在eagle3版本的webui.py中有下面這么句代碼，光從字面意思上看也很迷惑啊！加個“not”吧。
use_eagle3=args.no_eagle3,特別解釋一下-m選項，用于以模塊方式運行Python腳本，告訴python查找eagle.application.webui這個模塊并運行它（會去哪找eagle模塊呢？當前目錄以及PYTHON中），而無需手動cd進入目錄再執行python webui.py，更靈活（嗎？）
所以直接運行可能會遇到一個問題，如果你不是正處于EAGLE這一級目錄（誰知道經歷了前面一堆亂七八糟的時候跑到了哪個鬼目錄），又在PYTHONPATH找不到eagle.application.webui這個模塊的話，就會報錯“Error while finding module specification for 'eagle.application.webui' (ModuleNotFoundError: No module named 'eagle')”。要么執行“cd xxx/EAGLE”回到正確目錄，要么執行下列指令
export PYTHONPATH=xxx/EAGLE:$PYTHONPATH（哪個辦法好我也是反復橫跳，最后結論是覺得前者麻煩，限死了終端執行目錄，腳本文件存放位置。我傾向于后者，直接在終端執行的話，就臨時添加一下，換個終端便沒了，寫到~/.bashrc文件里，source一下永久生效吧也行，下次換別的哪個模塊的話改一下）
用第一種辦法就每次在終端打那么老長的指令，也累，可以直接去改 eagle/application/webui.py下解析參數的代碼，設成默認值，指令在“webui”后截斷即可
python -m eagle.application.webui --ea-model-path ../EAGLE-llama2-chat-7B/ --base-model-path ../Llama-2-7b-chat-hf/ --model-type llama-2-chat --total-token 2?感謝！下面看看彈出來的界面（兩張圖才截全，說真的，作者的界面都做得好好看???）
我問這個負責任有道德的模型，為啥它有的字體顯示橙色，它說它只會生成文本，沒能力調顏色或用視覺特效，可能是設備、瀏覽器、平臺或什么別的的鍋。作者用gradio包做的網頁，只需看上面的勾選框和解釋即可知，橙色高亮部分是EAGLE-2正確猜測的token，上圖便是將total-token設成2的結果。
這個紛繁復雜的交互式網頁，是作者通過“ea_generate”流式返回模型每次前向的結果，以快速響應用戶請求的。其中每次前向的結果output_ids的形狀為(batch_size, sequence_length)，都存到text里，其中第一個token由原始模型生成，存到naive_text里，剩下的token便是EAGLE-2的功勞。所以高亮哪些token呢？找text中naive_text里沒有的。代碼簡陋一點看就是
for output_ids in model.ea_generate(input_ids, ...):...# decode_ids是截至本輪所有新生的tokendecode_ids = output_ids[0, input_len:].tolist() # 去掉輸入部分decode_ids = truncate_list(decode_ids, model.tokenizer.eos_token_id) # 截斷終止符后面的部分...text = model.tokenizer.decode(decode_ids, skip_special_tokens=True,spaces_between_special_tokens=False,clean_up_tokenization_spaces=True, )naive_text.append(model.tokenizer.decode(output_ids[0, cu_len], # 本輪首個tokenskip_special_tokens=True,spaces_between_special_tokens=False,clean_up_tokenization_spaces=True, ))cu_len = output_ids.shape[1]# 將text中naive_token里沒有的內容高亮colored_text = highlight_text(text, naive_text, "orange")...用代碼體驗
下面用“eagenerate”，一次性返回完整的token序列（體驗感不如上面的好，有時會懷疑它卡了或者又要報錯），就像用Hugging Face的“generate”那樣，如下
from eagle.model.ea_model import EaModel
# EaModel，來自eagle.model.ea_model模塊，一個用于NLP任務的模型類，支持從預訓練模型加載權重
from fastchat.model import get_conversation_template
# get_conversation_template，來自fastchat.model，用于獲取對話模板（如vicuna）
import torchdef warmup(model):# 按照目標模型類型創建對話模板conv = get_conversation_template(args.model_type)if args.model_type == "llama-2-chat":# Llama 2 Chat版本需要一個系統提示詞，確保其回答安全無偏見符合道德，其他模型可能就無需這種額外約束了sys_p = "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information."conv.system_message = sys_pelif args.model_type == "mixtral":conv = get_conversation_template("llama-2-chat")conv.system_message = ''conv.sep2 = "</s>" # 特定結束符your_message="who are you?"# 將用戶輸入“hello”作為第一個角色（通常是用戶）的話加入對話conv.append_message(conv.roles[0], your_message)# 給第二個角色（通常是AI模型）留一個空的響應位置，等待模型生成conv.append_message(conv.roles[1], None)# get_prompt()負責將對話格式化成適合EaModel處理的輸入文本prompt = conv.get_prompt()if args.model_type == "llama-2-chat":prompt += " "# 分詞器將prompt轉換成token idinput_ids=model.tokenizer([prompt]).input_ids# 再轉換成PyTorch張量，并轉移到GPU提高推理效率input_ids = torch.as_tensor(input_ids).cuda()# 進行文本生成output_ids = model.eagenerate(input_ids,temperature=0.5,max_new_tokens=512) # eagenerate一次性返回完整的token序列output=model.tokenizer.decode(output_ids[0])print(output)# 使用命令行參數，添加參數解析（包已內置于python中）
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--ea_model_path",type=str,default="EAGLE-llama2-chat-7B",help="The path of EAGLE weight. This can be a local folder or a Hugging Face repo ID（<組織名或用戶名>/<模型名>）."
)
parser.add_argument("--base_model_path",type=str,default="Llama-2-7b-chat-hf",help="path of the original model. a local folder or a Hugging Face repo ID"
)
parser.add_argument("--load_in_8bit",action="store_true", # 如果提供該參數，則值為True，否則默認為Falsehelp="use 8-bit quantization"
)
parser.add_argument("--load_in_4bit",action="store_true",help="use 4-bit quantization"
)
parser.add_argument("--model_type",type=str,default="llama-2-chat",choices=["llama-2-chat","vicuna","mixtral","llama-3-instruct"]
)
parser.add_argument("--total_token",type=int,default=-1,help=" the number of draft tokens"
)
parser.add_argument("--max_new_token",type=int,default=512,help="the maximum number of new generated tokens",
)
args = parser.parse_args()model = EaModel.from_pretrained(base_model_path=args.base_model_path,ea_model_path=args.ea_model_path,total_token=args.total_token,torch_dtype=torch.float16,low_cpu_mem_usage=True,load_in_4bit=args.load_in_4bit,load_in_8bit=args.load_in_8bit,device_map="auto",
)# 讓模型進入推理模式，防止dropout等影響推理
model.eval()warmup(model)eagle.application.webui里有解析參數的代碼，但eagle.model.ea_model里沒有，得在作者提供的調用代碼里再額外填上，另外Vicuna、LLaMA2-Chat、LLaMA3-Instruct都是聊天模型，咱需要使用對應正確的聊天模板，不然會從模型里產生異常輸出，影響EAGLE性能。我仿照webui.py的代碼改了改。（不要兩種體驗同時用，一旦內存不夠了它就給你報奇奇怪怪的錯誤）（一個來回的對話，而且比較久才回你）
訓練
生成訓練數據
先說一下文件調用結構，在原項目中，allocation.py將待處理的數據平均分配，各GPU執行各自任務，并行調用ge_data_all_llama2chat.py，使用llama2-chat-7B進行數據預處理、加載數據集、推理并保存結果。數據集的存儲格式如下
new_examples = {"conversation": [], # 存儲對話文本"input_ids": [], # 輸入token ID"loss_mask": [] # 標記哪些token需要計算loss
}預測結果被存儲成.ckpt(checkpoint)文件，存儲內容如下：
input_ids: Tensor, shape=torch.Size([151]), dtype=torch.int64
hidden_state: Tensor, shape=torch.Size([151, 4096]), dtype=torch.float16
loss_mask: Tensor, shape=torch.Size([151]), dtype=torch.int64其中最重要的信息就是
hidden_state_big = outs_big.hidden_states[-1] # 提取最后一層隱藏狀態作者沒有提供原數據集，咱去hugging face上隨便找了個——shareGPT/computer_en_26k.jsonl。檢查了一下，總共20692條數據，單拎第一條出來看看
{"conversation_id": "vtu3ZfW", "category": "Program and Code", "conversation": [{"human": "Give Landing page HTML and CSS with bootstrap framework that takes email address and phone number in the lead generation form", "assistant": "Here is an example of a basic landing page HTML structure using the Bootstrap framework:\n...server."}]
}得據此修改ge_data_all_llama2chat.py中用到的鍵名，另外這個數據集沒有清理的好，“category”的值存在字符串和列表混用的現象，得自己檢查一遍，檢查代碼如下：
import pandas as pd
original_file = '/home/xxx/EAGLE_test/original_data/computer_en_26k.jsonl'
cleaned_file = '/home/xxx/EAGLE_test/original_data/computer_en_26k_cleaned.jsonl'try:df = pd.read_json(original_file, lines=True)print(df.head())
except Exception as e:print(f"Error loading JSON: {e}")# 檢查 'category' 列是否存在類型不一致
print(df['category'].apply(type).value_counts())# 遍歷 'category' 列并查找出現 list 的位置
for index, value in enumerate(df['category']):if isinstance(value, list):print(f"First occurrence of list in 'category' at row {index}")print(value)def clean_category(value):if isinstance(value, list):return ' '.join(value)elif isinstance(value, str):return valueelse:return str(value)df['category'] = df['category'].apply(clean_category)
df.to_json(cleaned_file, orient='records', lines=True)用原代碼中讀取json文件的方式讀取jsonl文件也行得通。額外說一下，jsonl文件的每行都是一個JSON（JavaScript Object Notation）格式數據。JSON是一種輕量級的數據交換格式，使用鍵值對存儲數據，類似于Python的dict或JavaScipt的對象（Object）。.json文件通常是配置文件、模型索引、數據存儲或API響應。
我摘取了前68條數據，在一臺服務器的4張GPU上進行處理，運行指令（參數同樣可以在文件中設置默認值）
python -m eagle.ge_data.allocation --outdir [path of data]得到的訓練數據如下
4個進程并行時的輸出、文件命名方式都沒有太大更改，原代碼還是蠻清晰的，但有些邏輯錯誤，到后面就會發現他媽的loss_mask全零🤯！主要是修改ge_data_all_llama2chat.py中的映射代碼
def preprocess_function(examples):new_examples = {"conversation":[], # 存儲對話文本"input_ids": [], # token ID"loss_mask": [] # 標記哪些token需要計算loss}# 獲取LLaMA-2對話模板conv = get_conversation_template("llama-2-chat")# 設定AI助手的形為準則sys_p="You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. " \"Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. " \"Please ensure that your responses are socially unbiased and positive in nature.\n\n" \"If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. " \"If you don't know the answer to a question, please don't share false information."conv.system_message=sys_p# 遍歷數據集中的所有對話num_of_examples = len(examples['conversation_id'])for i in range(num_of_examples):source= examples['conversation'][i]conv.messages = []# 處理human和assistant的每輪對話for j, sentence in enumerate(source):if "human" in sentence and "assistant" in sentence:conv.append_message(conv.roles[0], sentence["human"])conv.append_message(conv.roles[1], " " + sentence["assistant"])else:print(f"Warning: Invalid or incomplete dialogue at index {j}: {sentence}")# 獲取最終格式化后的對話文本conversation=conv.get_prompt()# 將對話轉成token IDinput_ids = tokenizer(conversation,return_tensors="pt", # 結果返回PyTorch張量max_length=2048,truncation=True,).input_ids[0]# 創建一個形同input_ids的張量loss_mask，初始值全為1，默認所有token都會計算lossloss_mask=torch.ones_like(input_ids)# AI回復前的分隔符：“[/INST] ”sep = conv.sep + conv.roles[1] + " "# 將對話拆分為不同輪次，conv.sep2：“ </s><s>”turns = conversation.split(conv.sep2)# 忽略開始符<s>的loss計算cur_len = 1loss_mask[: cur_len] = 0# 處理每輪對話for j, turn in enumerate(turns):# print(f"本輪內容：{turn}") # 兩輪對話三個turnif turn == "":breakturn_len = len(tokenizer(turn).input_ids) # 當前輪對話的token長度parts = turn.split(sep) # 拆分成用戶輸入和AI回復if len(parts) != 2:breakparts[0] += sepinst_len = len(tokenizer(parts[0]).input_ids) - 2 # 去掉開始符和結尾空格# 忽略（第一輪對話還額外有提示詞部分）用戶輸入部分的loss計算loss_mask[cur_len : cur_len + inst_len] = 0cur_len += (turn_len + 2)if j != 0 and not tokenizer.legacy:# print("是新版，</s><s>只占2個token")cur_len -= 1# print(tokenizer.decode(input_ids[cur_len-2:cur_len]))# 忽略padding位置的loss計算（實際上似乎沒有填充）loss_mask[cur_len:] = 0# 把格式化后的對話、token ID、loss_mask存入new_examplesnew_examples["conversation"].append(conversation)new_examples["input_ids"].append(input_ids[None,:])new_examples["loss_mask"].append(loss_mask[None,:])# print(f"loss_mask是什么類型啊現在？{type(new_examples['loss_mask'])}") # listreturn new_examples而且用這么點數據訓練肯定是不夠的，后面會出現在訓練集上的準確率嘎嘎提高，而在測試集上表現平平的情況，也就是過擬合了。干脆將jsonl文件中所有數據一并處理了，總20692條樣本
wc -l path/to/file好的，顯存不夠了😭原本我想換一臺更加空閑的服務器，但已經換了幾次服務器了，深知要想轉到一臺全新服務器有多麻煩，然后就想用docker把實驗依賴的所有環境之類的一起打包，獲取鏡像、實例化容器，在新服務器上拉取巴拉巴拉……但是家人們，我這條菜狗又困在網絡這一關了😭docker學習進程再次擱置。
更換策略，啟用4-bit NF4量化，節省顯存，如下更改代碼
'''啟用4-bit NF4量化，最節省顯存'''
quantization_config = BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_compute_dtype=torch.bfloat16, # 若GPU不支持bfloat16，改用torch.float16bnb_4bit_use_double_quant=True,bnb_4bit_quant_type="nf4", # Normalized Float 4
)
bigmodel = AutoModelForCausalLM.from_pretrained(bigname,quantization_config=quantization_config,device_map="auto"
)# "auto"會自動將模型的不同層分配到可用的GPU上，以實現模型并行
# bigmodel = AutoModelForCausalLM.from_pretrained(bigname, device_map="auto", torch_dtype=torch.float16)?額外說一下，在AutoModelForCausalLM.from_pretrained中，device_map參數用于控制模型在多個GPU上的分配方式。“auto”和“balanced”是兩種不同的分配策略，核心區別在于顯存分配的智能程度和均勻性。
device_map="auto"時按層順序分配，模型從第一層開始依次分配到當前顯存最充足的GPU，直到占滿該GPU的顯存，再切換到下一個GPU，實現簡單分配快，但可能導致現存利用率不均衡，某些GPU顯存剩余較多，對顯存碎片化較敏感。適用于模型層之間顯存占用差異較小（如大部分Transformer模型）和GPU顯存容量相同（下圖我用的服務器的顯存情況可能就不太符合？）的場景。（“gpustat -i 1”每隔1s刷新一次GPU使用信息）
device_map="balanced"時則是顯存均勻分配，盡可能讓每個GPU的顯存占用接近均等，計算所有層的顯存需求后，動態規劃最優分配方案。其現存利用率高，可減少OOM風險，但分配計算開銷稍大（首次加載略慢），適合模型層顯存需求差異大（如MoE）和GPU顯存容量不同的場景，對模型結構的適應性要求較高了。
建議就是優先嘗試“balanced”，尤其是面對OOM時，若仍失敗，就結合量化（4-bit）或手動分配，超大模型還可考慮offload_to_cpu（如accelerate庫的device_map="auto"就支持CPU卸載）
（每個進程處理5173條數據）（如果還是OOM運行失敗，但輸出文件下的東西看著有模有樣的，不要被騙啦！像我的，OOM時輸出文件大小71G）（文件夾大小查看指令如下）
du -sh /path/to/directory?allocation分配不同GPU并行，每個GPU上運行ge_data_all_llama2chat，讓多個CPU進行數據處理，ctrl+z暫停指令（c才是終止）時，顯存不一定釋放，可用如下指令篩選一下自己的進程
ps aux | grep "$USER" # 篩選屬于當前用戶的所有進程或直接批量終止
ps aux | grep "my_ge_data_all_llama2chat.py" | grep -v "grep" | awk '{print $2}' | xargs kill -9?記得在代碼中常加入下列指令：
torch.cuda.empty_cache()為什么需要這個PyTorch提供的函數？
《PyTorch的顯存管理機制》：PyTorch會使用CUDA內存緩存機制來加速張量分配和計算，當釋放張量時（如del tensor或變量超出作用域）時，PyTorch不會立即將顯存歸還給系統，而是保留在內部的緩存池中，當后續需要重新分配新張量時，優先從緩存中復用顯存，而非重新向CUDA申請，以便快速重用，減少顯存分配的開銷。具體地，PyTorch使用分塊內存池（Block-based Memory Pool）管理顯存，顯存被劃分為不同大小的塊，分配張量時PyTorch會尋找大小最匹配的塊，減少碎片化，釋放張量時，塊會被標記為空閑，供后續使用。但這種緩存機制就可能會導致顯存占用看起來比實際需求更高。?
《顯存碎片化問題》若頻繁分配和釋放不同大小的張量，顯存可能變得碎片化，即剩余顯存被分割為小塊，無法滿足大塊請求，此時即使顯存總量足夠，PyTorch也可能因找不到足夠大的連續顯存塊而報CUDA out of memory錯誤。PyTorch的顯存管理機制正是通過塊復用、合并相鄰空閑塊、按大小分類管理等方式減少碎片化。
empty_cache()的作用就在于強制釋放PyTorch緩存的未使用顯存，使其歸還給系統，減少顯存碎片化，提高顯存利用率。
另外進一步優化顯存管理的方法還有復用張量（避免在循環中反復創建臨時張量），還有按照報錯“torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 23.68 GiB total capacity; 1.06 GiB already allocated; 12.06 MiB free; 1.13 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. ?See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF”提示的那樣，如下調整分配策略，減少每次分配的內存塊最大大小，減少碎片化：
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:32"可以通過如下代碼監控顯存使用
'''檢查顯存占用情況'''
def print_memory():allocated = torch.cuda.memory_allocated() / 1024**3reserved = torch.cuda.memory_reserved() / 1024**3print(f"Allocated: {allocated:.2f}GB | Reserved: {reserved:.2f}GB")?終于不報錯了，可能在1h的推理等待后，成功輸出文件的大小在138G，友友們可以幫我檢查一下！
訓練自回歸頭
接下來用前面獲取的checkpoint訓練自回歸頭。訓練代碼和設置可參考eagle/train/
期間有幾點說明，.safetensors是一種專門用于存儲神經網絡權重的二進制文件格式（不能直接用文本編輯器打開，需用safetensors庫解析），由Hugging Face團隊開發，作為PyTorch的.pt或.bin文件更安全高效的替代方案。具體地，
1.更安全。其僅支持張量數據的存儲，即使部分文件損壞，仍可加載其他部分，存儲結構固定，而.bin（通常是torch.save生成的PyTorch權重文件）可以包含任意Python對象，若加載不受信任的.bin文件，可能導致代碼執行漏洞（Pickle反序列化攻擊），若文件損壞，可能整個模型都無法加載，若PyTorch版本不同，可能導致序列化格式變化，影響兼容性。
2. 更快。其支持按需加載（lazy loading），避免不必要數據傳輸，而.bin必須一次性全部加載
3. 更節約內存。其通過零拷貝（zero-copy）方式映射到內存，而.bin需要將數據從磁盤拷貝到內存，增加了額外開銷
存儲權重數據的二進制.safetensors文件通常有索引文件.index.json，內容如下：（本節在預訓練模型上較關心的正是它的語言建模頭lm_head）
{"metadata": {"total_size": 13476835328},"weight_map": {"lm_head.weight": "model-00002-of-00002.safetensors","model.embed_tokens.weight": "model-00001-of-00002.safetensors","model.layers.0.input_layernorm.weight": "model-00001-of-00002.safetensors","model.layers.0.mlp.down_proj.weight": "model-00001-of-00002.safetensors",...}
}這些文件在原先下載的預訓練模型下都配好了👌?
將上圖高亮的.safetensors文件中部分張量的鍵名和形狀輸出，有如下圖
再簡要介紹一下W&B（這種工具本菜狗才發現???）。它是一個機器學習實驗管理平臺，提供了強大的工具來追蹤、可視化和優化模型訓練過程，允許用戶記錄和管理實驗的超參數、損失函數、評估指標等，支持實現可視化訓練曲線，并為團隊提供協作和版本控制功能。通過與常見機器學習框架（如TensorFlow、PyTorch）無縫集成，W&B可幫助開發者更高效地管理實驗、提高模型性能，并加速團隊的研究與開發進程。！！！想用它得用點魔法訪問網站！！！
1. 在Python環境中安裝W&B
pip install wandb2. 在W&B官網注冊一個賬號，并登錄
wandb login <API key>
# 然后終端會輸出：wandb: Appending key for api.wandb.ai to your netrc file: /xx/xx/.netrc3. 在代碼中導入wandb庫并初始化一個新的實驗
import wandb
wandb.init(project="your_project_name")4. 記錄實驗參數（config）
wandb.config.batch_size = 32
wandb.config.epochs = 10
wandb.config.learning_rate = 0.001?5. 記錄訓練過程中的數據（log追蹤）
# 假設在訓練過程中記錄每個epoch的損失和準確率
for epoch in range(10):loss = train_one_epoch()accuracy = evaluate_model()wandb.log({"epoch": epoch, "loss": loss, "accuracy": accuracy})6. 若想在訓練結束后保存模型，也可將其與W&B關聯，進行版本控制
wandb.save('path_to_your_model.h5')?7. 結束實驗
wandb.finish()每次運行完wandb.init()后，W&B會自動創建一個新頁面，記錄該實驗的所有信息 。
關于BF16：Bfloat16是一種16-bit浮點數格式，主要由Google提出，并在其TPU和一些NVIDIA A100等加速器中得到了廣泛使用。與FP16相比，其在數值范圍和精度表示上做了不同的權衡，FP16使用5位指數位，10位尾數位；BF16使用8位指數位，7位尾數位，其指數范圍和FP32相當，使其能表示更廣泛的數值范圍，于神經網絡訓練中的梯度計算和大數值表示非常有用。
面對OOM問題時還可以用一個PyTorch的梯度檢查點（gradient checkpointing）技術，以顯著減少顯存占用（↓50%~70%），但會略增加計算時間（↑20%~30%）。
問題背景：訓練深度神經網絡時，前向傳播計算的中間結果需要保存，以便在反向傳播時計算梯度，這些中間結果就會占用大量GPU顯存，尤其是大模型（如Llama-2-7B）。
解決方法：只保留部分關鍵中間結果，其余的在反向傳播時重新計算，用計算時間換顯存。
from torch.utils.checkpoint import checkpoint # 導入檢查點功能def _forward(self, x):return self.layer2(self.layer1(x)) # 模型的計算邏輯（如Transformer層的堆疊）def forward(self, x):return checkpoint(self._forward, x) # 使用檢查點包裝前向計算'''關鍵點解釋：
1. checkpoint(func, *args)func：要優化的前向計算函數（如self._forward）*args：傳給func的輸入（如x）作用：PyTorch不會保存func的中間激活值，而是在反向傳播時重新計算它們
2. self._forward(x)是你模型的實際前向計算邏輯（如nn.Module的forward方法）
'''''' 進階操作如下 '''# 只對顯存占用高的部分使用檢查點
def forward(self, x):x = checkpoint(self.layer1, x)x = self.layer2(x) # layer2正常計算return x# 調整檢查點頻率
def forward(self, x):if self.training: # 僅在訓練時使用檢查點return checkpoint(self._forward, x)else: # 推理時不使用，避免額外計算return self._forward(x)?事先說一下我們將要處理的數據，原始數據集shareGPT/computer_en_26k.jsonl有20692條對話樣本，把里面的“category”類處理一下，清洗后的數據集文件cleaned.jsonl文件中樣本數是一樣的（.jsonl文件每行是一個JSON樣本，可用wc -l file_path進行查看），處理成20692個ckpt文件，從中抽取0.95用作訓練的話，訓練樣本數就是19657。然后accelerator檢測到服務器上有3塊GPU，咱為了避免OOM設置了較小的批處理大小，為2，故每個訓練epoch的進度條顯示有“/3277”，測試時顯示“/173”。
為什么說OOM時老說要減少批處理大小？（原先通過LLM推理獲取數據集時批處理大小就是1了）加載模型本身需要固定顯存，例如LLaMA-2-7B的4-bit量化后，7B參數，每個0.5個字節，再加些七七八八的，實際推理時的經驗值為6-8GB；然后每增加一個樣本，前向傳播的中間結果就會線性增加顯存占用，故總顯存≈模型參數+批處理大小×單樣本激活值顯存。
好不容易準備好數據集，運行下列代碼進行訓練（我和網絡和OOM問題不共戴天！）
accelerate launch -m --mixed_precision=bf16 eagle.train.main --tmpdir [path of data]\
--cpdir [path of checkpoints] --configpath [path of config file]?跑了超4h，約10min/訓練epoch（明白為啥原代碼只20個epoch了），可以將代碼放到后臺開始運行，即使電腦關閉也不影響服務器繼續當牛做馬。首先創建一個新會話，運行下列指令后便會進入一個新終端，此時再打開另一個終端窗口，輸入tmux ls便能看到包含它在內的會話列表。
tmux new -s <會話名>在tmux會話中輸入要在后臺運行的指令。如果想滾動查看窗口內容，可以配置一下鼠標支持，即在~/.tmux.conf中添加如下配置后重載（tmux source-file ~/.tmux.conf）
set -g mouse on # 允許鼠標滾動和選擇面板?創建會話后鍵入Ctrl+B, 松手，再按D就能從前臺分離會話，通過如下指令又能重新連接
tmux attach -t <會話名>創建會話后鍵入Ctrl+B，松手，再按C就能新建又一個窗口，執行exit就能退出當前窗口，當最后一個窗口也被關閉時，整個會話會被自動刪除。
利用tmux就能方向地關閉電腦了嗎？No，還記得前面提到的，為了使用W&B，做了到本地代理的端口映射嗎？這個Clash還不能關！或者說我設置了透明代理，代理IP就是本機的IPv4地址，那本機就不能關！我們的確將程序放在服務器上由tmux托管，一般本機關閉并不影響程序本體的繼續運行，但這里一旦關機，目標代理IP無法連接，網絡請求就會失敗。W&B還挺好用的，又自動給你畫圖，又記錄日志的，我們的服務器又不能科學上網，自己的這破電腦就一直開著唄。
?區分一下幾個指標?
交叉熵損失：（plogp是 目標概率分布 與 模型預測的對數概率 的乘積，loss_mask用于忽略無效位置的損失）用于衡量模型預測的token分布和真實token分布的差異，適用于分類任務
out_head = head(predict)
out_logp = nn.LogSoftmax(dim=2)(out_head)
plogp = target_p * out_logp
cross_entropy_loss = -torch.sum(torch.sum(loss_mask * plogp, 2)) / (loss_mask.sum() + 1e-5)平滑L1損失：計算模型預測的隱藏狀態predict與目標隱藏狀態target之間的差異
# 平滑L1損失（又稱Huber損失）（介于絕對誤差損失和均方誤差之間的損失函數）
criterion = nn.SmoothL1Loss(reduction="none")
...
smooth_L1_loss = criterion(predict, target)
smooth_L1_loss = torch.sum(torch.mean(loss_mask * smooth_L1_loss, 2)) / (loss_mask.sum() + 1e-5)Top-k準確率：top-1即標準準確率，指示預測的最高概率token是否正確，top-2/3表示前2/3個token是否包含正確答案，用于評估語言模型的預測能力
def top_kaccuracy(output, target, topk=(1, 2, 3)):
'''
預測頭輸出output的形狀為(bs, num_classes)
目標輸出target的形狀為(bs,)
'''with torch.no_grad():maxk = max(topk) # 3_, pred = output.topk(maxk, 1, True, True) # 取預測分數前3名，分數忽略，索引保留pred = pred.t() # pred形狀轉置為(maxk, bs)correct = pred.eq(target.view(1, -1).expend_as(pred)?關于deepspeed：它是由Microsoft開發的深度學習優化庫，旨在通過高效的內存管理、分布式訓練和混合精度訓練等技術，顯著提高大規模模型的訓練效率。利用零冗余優化器（ZeRO）減少顯存占用，支持訓練數十億參數的超大模型（如GPT-3），還支持數據并行、模型并行和多機多卡訓練，能在有限的硬件資源上加速訓練過程，并廣泛用于NLP、CV等領域的大規模深度學習任務。
很抱歉寫到這里我打算放棄了，有點戛然而止的感覺，而且前面寫的也亂七八糟。主要是在完成主線任務的時候本人太菜了會遇到很多支線問題，而且原論文提供的代碼里也存在一些很明顯的錯誤，會讓人不禁懷疑這是否是一篇high level的論文，然后實驗室掛了兩臺有3090的服務器成了“最后一根稻草”，我想過換另一臺顯存很局限的3090服務器，或一臺有A40的服務器、或拿師兄買的4090服務器，但停下來審視一番后還是算了，其實也算走完了訓練這一步，最后也是在整理和修改原作者的代碼。我想再去探索一下別的科研方向，很感謝eagle的作者🙇?，研究了代碼之后我對論文的理解會更加深刻，也會感慨現在的論文能被接收很不容易，同時在和代碼搏斗的這個月里俺也學到了很多東西，這就夠了。最后最后，謝謝友友們能看到這里🌼
