AWS Trainium & Inferentia 文件
使用 LoRA 微調 Llama 3 8B
並獲得增強的文件體驗
開始使用
使用 LoRA 微調 Llama 3 8B
注意：本教程的完整指令碼可在此處下載。
本教程將教你如何在 AWS Trainium 上微調像 Llama 3 這樣的開源 LLM。在我們的示例中，我們將利用 Optimum Neuron、Transformers 和 Datasets 庫。
您將學習如何
雖然在本教程中我們將使用 Llama-3 8B，但完全可以使用其他模型，只需切換 model_id 即可。
1. 設定 AWS 環境
在開始本教程之前，你需要設定你的環境
- 建立一個 AWS Trainium 例項。你需要一個
trn1.32xlarge例項，它包含 16 個 Neuron 裝置。 你可以按照這篇指南來建立一個。 - 確保你已經登入到 Hugging Face Hub
huggingface-cli login --token YOUR_TOKEN
- 檢查你是否可以訪問該模型。一些開源模型是受限的，這意味著使用者需要向模型所有者申請才能使用模型權重。這裡我們將訓練 Llama-3 8B，有兩種可能性
- 官方受限倉庫：
meta-llama/Meta-Llama-3-8B - 非官方不受限倉庫：
NousResearch/Meta-Llama-3-8B
- 克隆 Optimum Neuron 倉庫，其中包含本教程中描述的完整指令碼：
git clone https://github.com/huggingface/optimum-neuron.git- 確保你安裝了
training附加包，以獲取所有必要的依賴項
python -m pip install .[training]
2. 載入並準備資料集
在本教程中，我們將使用 Dolly，這是一個開源的指令遵循記錄資料集，其類別在 InstructGPT 論文中有概述，包括頭腦風暴、分類、封閉式問答、生成、資訊提取、開放式問答和摘要。
示例
{
"instruction": "What is world of warcraft",
"context": "",
"response": (
"World of warcraft is a massive online multi player role playing game. "
"It was released in 2004 by bizarre entertainment"
)
}我們可以使用 🤗 Datasets 庫中的 load_dataset() 方法非常輕鬆地載入 dolly 資料集。
from datasets import load_dataset
from random import randrange
# Load dataset from the hub
dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
print(f"dataset size: {len(dataset)}")
print(dataset[randrange(len(dataset))])
# dataset size: 15011為了對我們的模型進行指令微調，我們需要
將我們的結構化示例轉換為透過指令描述的任務集合
（可選）將多個示例打包到一個序列中以進行更高效的訓練。換句話說，我們將多個示例堆疊成一個示例，並用 EOS 標記分隔它們。
我們可以手動完成這些操作，但我們將改用 NeuronSFTTrainer。
3. 使用 NeuronSFTTrainer 在 AWS Trainium 上對 Llama 進行監督式微調
通常，你會使用 SFTConfig 和 SFTTrainer 類來對基於 PyTorch 的 Transformer 模型進行監督式微調。
在這裡，我們將改用 NeuronSFTConfig 和 NeuronSFTTrainer。這些類複製了 trl 庫中的類，同時確保它們在 Neuron 核心上正常工作。
格式化我們的資料集
有多種方法可以將資料集提供給 NeuronSFTTrainer，其中一種是提供一個格式化函式。對於不打包示例的 dolly，它看起來如下
def format_dolly(examples):
output_text = []
for i in range(len(examples["instruction"])):
instruction = f"### Instruction\n{examples['instruction'][i]}"
context = f"### Context\n{examples['context'][i]}" if len(examples["context"][i]) > 0 else None
response = f"### Answer\n{examples['response'][i]}"
prompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])
output_text.append(prompt)
return output_text準備模型
由於 Llama-3 8B 是一個大模型，即使使用分散式訓練，它也無法在一個 trn1.32xlarge 例項上容納。要僅使用一個 Trainium 例項實際微調一個 8B 模型，我們需要同時使用 LoRA 和分散式訓練。
如果你想了解更多關於分散式訓練的資訊，可以檢視文件。
在這裡，我們將結合 LoRA 使用張量並行。我們的訓練程式碼將如下所示
from peft import LoraConfig
from optimum.neuron import NeuronSFTConfig, NeuronSFTTrainer
# Define the tensor_parallel_size
tensor_parallel_size = 2
dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
model_id = "meta-llama/Meta-Llama-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_id)
config = LoraConfig(
r=16,
lora_alpha=16,
lora_dropout=0.05,
target_modules=[
"q_proj",
"gate_proj",
"v_proj",
"o_proj",
"k_proj",
"up_proj",
"down_proj"
],
bias="none",
task_type="CAUSAL_LM",
)
# training_args is an instance of NeuronTrainingArguments
args = training_args.to_dict()
sft_config = NeuronSFTConfig(
max_seq_length=1024,
packing=False,
**args,
)
trainer = NeuronSFTTrainer(
args=sft_config,
model=model,
peft_config=config,
tokenizer=tokenizer,
train_dataset=dataset,
formatting_func=format_dolly,
)
# Start training
trainer.train()
trainer.save_model() # Saves the tokenizer too for easy upload這裡的關鍵點是
- 我們定義了一個
LoraConfig，它指定了哪些層應該有介面卡，以及這些介面卡的超引數。 - 我們從常規的
NeuronTrainingArguments建立一個NeuronSFTConfig。這裡我們指定不打包我們的示例，並且最大序列長度應為1024，這意味著每個示例都將被填充或截斷到1024的長度。 - 我們使用
NeuronSFTTrainer來進行訓練。它將接收模型以及lora_config、sft_config和format_dolly，並準備用於監督式微調的資料集和模型。
4. 啟動訓練
我們準備了一個名為 sft_lora_finetune_llm.py 的指令碼，總結了本教程中提到的所有內容。
PyTorch Neuron 使用 torch_xla。它在訓練迴圈執行期間延遲評估操作，這意味著它在後臺構建一個符號圖，並且只有在列印張量、將其傳輸到 CPU 或呼叫 xm.mark_step() 時，圖才會在硬體上執行。在執行過程中，根據控制流可能會構建多個圖，按順序編譯每個圖可能需要時間。為了緩解這個問題，Neuron SDK 提供了 neuron_parallel_compile，這是一個執行快速試執行的工具，它會構建所有圖並並行編譯它們。這個步驟通常被稱為預編譯。
預編譯
在 AWS Trainium 上訓練模型時，我們首先需要用我們的訓練引數來編譯模型。
為了簡化這一步，我們添加了一個模型快取倉庫，它允許我們使用來自 Hugging Face Hub 的預編譯模型來跳過編譯步驟。這很有用，因為它能比實際訓練時更快地編譯模型，因為編譯可以並行化。但要注意：模型配置的任何更改都可能導致新的編譯，這可能會導致一些快取未命中。
要了解更多關於快取系統以及如何建立自己的私有快取倉庫的資訊，請檢視這篇指南。
編譯命令只需將你的指令碼作為輸入傳遞給 neuron_parallel_compile 工具即可
#!/bin/bash
set -ex
export NEURON_FUSE_SOFTMAX=1
export NEURON_RT_ASYNC_EXEC_MAX_INFLIGHT_REQUESTS=3
export MALLOC_ARENA_MAX=64 # limit the CPU allocation to avoid potential crashes
export NEURON_CC_FLAGS="--model-type=transformer --distribution-strategy=llm-training --enable-saturate-infinity --cache_dir=/home/ubuntu/cache_dir_neuron/"
PROCESSES_PER_NODE=8
TP_DEGREE=2
PP_DEGREE=1
BS=1
GRADIENT_ACCUMULATION_STEPS=8
LOGGING_STEPS=1
MODEL_NAME="meta-llama/Meta-Llama-3-8B"
OUTPUT_DIR=dolly_llama_output
if [ "$NEURON_EXTRACT_GRAPHS_ONLY" = "1" ]; then
MAX_STEPS=10
NUM_EPOCHS=1
else
MAX_STEPS=-1
NUM_EPOCHS=3
fi
torchrun --nproc_per_node $PROCESSES_PER_NODE docs/source/training_tutorials/sft_lora_finetune_llm.py \
--model_id $MODEL_NAME \
--num_train_epochs $NUM_EPOCHS \
--do_train \
--learning_rate 5e-5 \
--warmup_ratio 0.03 \
--max_steps $MAX_STEPS \
--per_device_train_batch_size $BS \
--per_device_eval_batch_size $BS \
--gradient_accumulation_steps $GRADIENT_ACCUMULATION_STEPS \
--gradient_checkpointing true \
--bf16 \
--tensor_parallel_size $TP_DEGREE \
--pipeline_parallel_size $PP_DEGREE \
--logging_steps $LOGGING_STEPS \
--save_total_limit 1 \
--output_dir $OUTPUT_DIR \
--lr_scheduler_type "constant" \
--overwrite_output_dir為方便起見，我們將這個 shell 指令碼儲存到一個檔案中，sft_lora_finetune_llm.sh。你現在可以將其傳遞給 neuron_parallel_compile 工具來觸發編譯
neuron_parallel_compile bash docs/source/training_tutorials/sft_lora_finetune_llm.sh
注意：在編譯結束時，可能會出現一個 `FileNotFoundError` 訊息。你可以安全地忽略它，因為一些編譯快取已經建立好了。
這個預編譯階段會執行 10 個訓練步驟，以確保編譯器已經編譯了所有必要的圖。這通常足以累積和編譯實際訓練期間需要的所有圖。
注意：沒有快取的編譯可能需要一段時間。它還會在編譯期間在 `dolly_llama_output` 目錄中建立一些虛擬檔案，你之後需要將它們刪除。
# remove dummy artifacts which are created by the precompilation command
rm -rf dolly_llama_output實際訓練
編譯完成後，我們可以用類似的命令開始我們的實際訓練，我們只需要移除對 `neuron_parallel_compile` 的使用。
我們將使用 `torchrun` 來啟動我們的訓練指令碼。`torchrun` 是一個可以自動將 PyTorch 模型分佈到多個加速器上的工具。我們可以將加速器的數量作為 `nproc_per_node` 引數與我們的超引數一起傳遞。
與編譯命令的區別在於，我們更改了變數 `max_steps=10` 和 `num_train_epochs=3`。
使用與預編譯步驟中相同的命令啟動訓練，但不帶 `neuron_parallel_compile`
bash docs/source/training_tutorials/sft_lora_finetune_llm.sh
就是這樣，我們成功地在 AWS Trainium 上訓練了 Llama-3 8B！
但在我們分享和測試我們的模型之前，我們需要整合我們的模型。由於我們在訓練期間使用了張量並行，我們儲存了檢查點的分片版本。我們現在需要將它們整合起來。
整合檢查點併合並模型
Optimum CLI 提供了一種非常簡單的方法，透過 `optimum neuron consolidate [sharded_checkpoint] [output_dir]` 命令來實現這一點。
optimum-cli neuron consolidate dolly_llama_output dolly_llama_output
這將建立一個 `adapter_model.safetensors` 檔案，即我們在上一步中訓練的 LoRA 介面卡權重。我們現在可以重新載入模型並將其合併，以便可以載入它進行評估。
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig
MODEL_NAME = 'meta-llama/Meta-Llama-3-8B'
ADAPTER_PATH = 'dolly_llama_output'
MERGED_MODEL_PATH = 'dolly_llama'
# Load base odel
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
# Load adapter configuration and model
adapter_config = PeftConfig.from_pretrained(ADAPTER_PATH)
finetuned_model = PeftModel.from_pretrained(model, ADAPTER_PATH, config=adapter_config)
print("Saving tokenizer")
tokenizer.save_pretrained(MERGED_MODEL_PATH)
print("Saving model")
finetuned_model = finetuned_model.merge_and_unload()
finetuned_model.save_pretrained(MERGED_MODEL_PATH)
這個步驟可能需要幾分鐘。我們現在有了一個目錄，其中包含評估微調模型所需的所有檔案。
5. 評估和測試微調後的 Llama 模型
與訓練一樣，為了能夠在 AWS Trainium 或 AWS Inferentia2 上執行推理，我們需要編譯我們的模型。在這種情況下，我們將使用我們的 Trainium 例項進行推理測試，但你可以切換到 Inferentia2（`inf2.24xlarge`）進行推理。
Optimum Neuron 實現了類似於 Transformers AutoModel 的類，以便於推理使用。我們將使用 `NeuronModelForCausalLM` 類來載入我們原始的 transformers 檢查點，並將其轉換為 neuron。
from optimum.neuron import NeuronModelForCausalLM
from transformers import AutoTokenizer
compiler_args = {"num_cores": 2, "auto_cast_type": 'fp16'}
input_shapes = {"batch_size": 1, "sequence_length": 2048}
tokenizer = AutoTokenizer.from_pretrained("dolly_llama")
model = NeuronModelForCausalLM.from_pretrained(
"dolly_llama",
export=True,
**compiler_args,
**input_shapes)注意：推理編譯可能需要長達 25 分鐘。幸運的是，你只需要執行一次。與訓練前完成的預編譯步驟一樣，如果你更改執行推理的硬體，例如從 Trainium 移動到 Inferentia2，你也需要執行此編譯步驟。編譯是引數和硬體特定的。
# COMMENT IN if you want to save the compiled model
# model.save_pretrained("compiled_dolly_llama_output")我們現在可以測試推理了，但必須確保我們將輸入格式化為我們用於微調的提示格式。因此，我們建立了一個輔助方法，它接受一個包含 `instruction` 和可選的 `context` 的 `dict`。
def format_dolly_inference(sample):
instruction = f"### Instruction\n{sample['instruction']}"
context = f"### Context\n{sample['context']}" if "context" in sample else None
response = f"### Answer\n"
prompt = "\n\n".join([i for i in [instruction, context, response] if i is not None])
return prompt
def generate(sample):
prompt = format_dolly_inference(sample)
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(
**inputs,
max_new_tokens=512,
do_sample=True,
temperature=0.9,
top_k=50,
top_p=0.9
)
return tokenizer.decode(outputs[0], skip_special_tokens=False)[len(prompt):]讓我們測試一下推理。首先我們測試沒有上下文的情況。
注意：在 AWS Trainium 上使用 2 個核心進行推理，速度預計不會非常快。對於推理，我們建議使用 Inferentia2。
prompt = {
"instruction": "Can you tell me something about AWS?"
}
res = generate(prompt)
print(res)AWS 代表 Amazon Web Services。AWS 是亞馬遜提供的一套遠端計算服務。其中最廣泛使用的包括 Amazon Elastic Compute Cloud（Amazon EC2），它在雲中提供可調整大小的計算容量；Amazon Simple Storage Service（Amazon S3），它是一種物件儲存服務；以及 Amazon Elastic Block Store（Amazon EBS），它旨在為 AWS 例項提供高效能、持久的塊儲存卷。AWS 還提供其他服務，例如 AWS Identity and Access Management（IAM），一項使組織能夠控制對其 AWS 資源訪問的服務，以及 AWS Key Management Service（AWS KMS），它幫助客戶建立和控制加密金鑰的使用。
這看起來是正確的。現在，讓我們新增一些上下文，例如在 RAG 應用中你會做的那樣
prompt = {
"instruction": "How can I train models on AWS Trainium?",
"context": "🤗 Optimum Neuron is the interface between the 🤗 Transformers library and AWS Accelerators including [AWS Trainium](https://aws.amazon.com/machine-learning/trainium/?nc1=h_ls) and [AWS Inferentia](https://aws.amazon.com/machine-learning/inferentia/?nc1=h_ls). It provides a set of tools enabling easy model loading, training and inference on single- and multi-Accelerator settings for different downstream tasks."
}
res = generate(prompt)
print(res)你可以使用 Optimum Neuron 介面在 AWS Trainium 上訓練模型。
太棒了，我們的模型也正確地使用了提供的上下文。我們完成了。恭喜你在 AWS Trainium 上微調了 Llama。