將 TRL 用於 LLaMA 模型

我們已經開始推出在 trl 中使用 Meta LLaMA 模型的示例（原始 LLaMA 模型請參閱 Meta 的 LLaMA 釋出）。

高效的訓練策略

即使是訓練最小的 LLaMA 模型也需要巨大的記憶體。簡單計算一下：在 bf16 精度下，每個引數使用 2 位元組（在 fp32 中為 4 位元組），此外，像 Adam 最佳化器還會額外使用 8 位元組（更多資訊請參閱 Transformers 中的效能文件）。因此，一個 7B 引數的模型僅僅為了載入到記憶體就需要 (2+8)*7B=70GB，並且在計算注意力分數等中間值時可能需要更多記憶體。所以，你甚至無法在一塊 80GB 的 A100 上進行這樣的訓練。你可以使用一些技巧，比如更高效的最佳化器或半精度訓練，來擠佔更多記憶體，但遲早會耗盡。

另一種選擇是使用引數高效微調（PEFT）技術，例如 peft 庫，它可以對以 8 位載入的模型執行低秩適應（LoRA）。更多關於 peft + trl 的資訊，請參閱 Peft 整合文件。

以 8 位載入模型可以大幅減少記憶體佔用，因為每個引數的權重只需要一個位元組（例如，7B 的 LlaMa 模型在記憶體中佔用 7GB）。LoRA 不是直接訓練原始權重，而是在一些特定層（通常是注意力層）之上新增小型的介面卡層；因此，可訓練引數的數量大大減少。

在這種情況下，一個經驗法則是為每十億引數分配約 1.2-1.4GB 記憶體（取決於批次大小和序列長度），以容納整個微調設定。這使得在低成本的情況下微調更大的模型（在 NVIDIA A100 80GB 上可達 50-60B 規模的模型）成為可能。

現在我們可以將非常大的模型裝入單個 GPU，但訓練可能仍然非常慢。在這種情況下，最簡單的策略是資料並行：我們將相同的訓練設定複製到不同的 GPU 上，併為每個 GPU 傳遞不同的批次。這樣，你可以並行化模型的前向/後向傳播，並隨 GPU 數量進行擴充套件。

我們使用 transformers.Trainer 或 accelerate，它們都支援資料並行，無需任何程式碼更改，只需在使用 torchrun 或 accelerate launch 呼叫指令碼時傳遞引數即可。下面分別是使用 accelerate 和 torchrun 在單臺機器上用 8 個 GPU 執行訓練指令碼的示例。

accelerate launch --multi_gpu --num_machines 1  --num_processes 8 my_accelerate_script.py
torchrun --nnodes 1  --nproc_per_node 8 my_torch_script.py

監督式微調

在我們開始訓練獎勵模型並用強化學習（RL）來調優模型之前，如果模型在我們感興趣的領域已經表現良好，會很有幫助。在我們的案例中，我們希望它能回答問題，而對於其他用例，我們可能希望它遵循指令，這時指令調優是一個好主意。實現這一目標的最簡單方法是繼續使用領域或任務中的文字，以語言建模的目標來訓練語言模型。StackExchange 資料集非常龐大（超過 1000 萬條指令），所以我們可以輕鬆地在其子集上訓練語言模型。

在進行 RLHF 之前微調模型並沒有什麼特別之處——我們只是應用了預訓練中的因果語言建模目標。為了高效地使用資料，我們使用了一種稱為打包（packing）的技術：我們不在批次中的每個樣本中只放一個文字，然後填充到最長文字或模型的最大上下文長度，而是用 EOS 標記將大量文字連線起來，然後切出與上下文大小相同的塊來填充批次，無需任何填充。

透過這種方法，訓練效率更高，因為每個透過模型的標記都參與了訓練，而填充標記通常會從損失計算中被掩碼掉。如果你的資料不多，更擔心偶爾會截斷一些超出上下文的標記，你也可以使用經典的資料載入器。

# load model in 8bit
model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        load_in_8bit=True,
        device_map={"": Accelerator().local_process_index}
    )
model = prepare_model_for_kbit_training(model)

# add LoRA to model
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, config)

我們使用因果語言建模目標訓練模型幾千步，然後儲存模型。由於我們將用不同的目標再次調優模型，我們將介面卡權重與原始模型權重合並。

免責宣告： 由於 LLaMA 的許可證限制，我們僅釋出此部分及後續章節中模型檢查點的介面卡權重。您可以透過填寫 Meta AI 的表格申請基礎模型權重的訪問許可權，然後執行此指令碼將其轉換為 🤗 Transformers 格式。請注意，在 v4.28 釋出之前，您還需要從原始碼安裝 🤗 Transformers。

現在我們已經為任務微調了模型，我們準備好訓練一個獎勵模型了。

獎勵建模與人類偏好

原則上，我們可以直接使用人工標註透過 RLHF 微調模型。然而，這將要求我們在每次最佳化迭代後將一些樣本傳送給人類進行評分。這既昂貴又緩慢，因為收斂需要大量的訓練樣本，並且人類閱讀和標註速度存在固有的延遲。

一個效果很好的技巧是，在 RL 迴圈開始前，在收集到的人工標註資料上訓練一個獎勵模型，而不是直接使用反饋。獎勵模型的目標是模仿人類如何評價一段文字。建立獎勵模型有幾種可能的策略：最直接的方法是預測標註（例如，評分或“好”/“壞”的二元值）。實際上，更好的做法是預測兩個示例的排名，其中獎勵模型會看到給定提示 x 的兩個候選答案 (y_k, y_j)，並必須預測哪一個會被人類標註者評價更高。

使用 StackExchange 資料集，我們可以根據分數推斷出兩個答案中哪一個更受使用者青睞。有了這些資訊和上面定義的損失函式，我們就可以透過新增一個自定義的損失函式來修改 transformers.Trainer。

class RewardTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        rewards_j = model(input_ids=inputs["input_ids_j"],  attention_mask=inputs["attention_mask_j"])[0]
        rewards_k = model(input_ids=inputs["input_ids_k"], attention_mask=inputs["attention_mask_k"])[0]
        loss = -nn.functional.logsigmoid(rewards_j - rewards_k).mean()
        if return_outputs:
            return loss, {"rewards_j": rewards_j, "rewards_k": rewards_k}
        return loss

我們利用一個包含 10 萬對候選答案的子集，並在一個包含 5 萬對的留出集上進行評估。我們使用適度的批次大小 4，使用 LoRA peft 介面卡，以 BF16 精度和 Adam 最佳化器訓練 Llama 模型一個週期。我們的 LoRA 配置如下：

peft_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
)

如下一節所述，生成的介面卡可以合併到凍結的模型中，並儲存以供後續下游使用。

從人類反饋中進行強化學習

有了微調好的語言模型和獎勵模型，我們現在可以執行 RL 迴圈了。它大致遵循三個步驟：

1. 根據提示生成響應，
2. 使用獎勵模型對響應進行評分，
3. 使用評分執行強化學習策略最佳化步驟。

查詢和響應提示在被分詞並傳遞給模型之前，會按照以下模板進行格式化：

Question: <Query>

Answer: <Response>

在 SFT、RM 和 RLHF 階段都使用了相同的模板。我們再次利用 peft 進行記憶體高效的訓練，這在 RLHF 上下文中提供了額外的優勢。在這裡，參考模型和策略共享相同的基礎，即 SFT 模型，我們以 8 位載入並在訓練期間凍結。我們僅使用 PPO 最佳化策略的 LoRA 權重，同時共享基礎模型的權重。

for epoch, batch in tqdm(enumerate(ppo_trainer.dataloader)):
    question_tensors = batch["input_ids"]

	# sample from the policy and to generate responses
    response_tensors = ppo_trainer.generate(
        question_tensors,
        return_prompt=False,
        length_sampler=output_length_sampler,
        **generation_kwargs,
    )
    batch["response"] = tokenizer.batch_decode(response_tensors, skip_special_tokens=True)

    # Compute sentiment score
    texts = [q + r for q, r in zip(batch["query"], batch["response"])]
    pipe_outputs = sentiment_pipe(texts, **sent_kwargs)
    rewards = [torch.tensor(output[0]["score"] - script_args.reward_baseline) for output in pipe_outputs]

    # Run PPO step
    stats = ppo_trainer.step(question_tensors, response_tensors, rewards)
	# Log stats to Wandb
    ppo_trainer.log_stats(stats, batch, rewards)

有關其餘細節和評估，請參閱我們關於 StackLLaMA 的部落格文章。