trl-强化学习训练-grpo训练为例

一、定义

1.trl + lora+ transformers 训练模型
2.部署与预测
3.模型合并
4.vllm 部署
4. 代码讲解
5. trl 训练器参数讲解

二、训练

1.trl + lora+ transformers 训练模型
训练数据格式要求：包含字段：prompt ，不需要ans
如：
注意： 数据的处理，被封装在了grpo trainer 优化器的compute_loss 中。
自定义输入：

# Load and prep dataset
SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""

XML_COT_FORMAT = """\
<reasoning>
{reasoning}
</reasoning>
<answer>
{answer}
</answer>
"""

def extract_xml_answer(text: str) -> str:
    answer = text.split("<answer>")[-1]
    answer = answer.split("</answer>")[0]
    return answer.strip()

def extract_hash_answer(text: str):
    if "####" not in text:
        return None
    return text.split("####")[1].strip()


# uncomment middle messages for 1-shot prompting
def get_gsm8k_questions(split = "train"):
    data = load_dataset('E://openaigsm8k/main')[split].select(range(200)) # type: ignore
    data = data.map(lambda x: { # type: ignore   自定义输入格式
        'prompt': [
            {'role': 'system', 'content': SYSTEM_PROMPT},
            {'role': 'user', 'content': x['question']}
        ],
        'answer': extract_hash_answer(x['answer'])
    }) # type: ignore
    return data # type: ignore

dataset = get_gsm8k_questions()

训练：

import argparse
from dataclasses import dataclass, field
from typing import Optional

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer

from trl import GRPOConfig, GRPOTrainer, ModelConfig, ScriptArguments, TrlParser, get_peft_config


@dataclass
class GRPOScriptArguments(ScriptArguments):
    """
    Script arguments for the GRPO training script.

    Args:
        reward_model_name_or_path (`str` or `None`):
            Reward model id of a pretrained model hosted inside a model repo on huggingface.co or local path to a
            directory containing model weights saved using [`~transformers.PreTrainedModel.save_pretrained`].
    """

    reward_model_name_or_path: Optional[str] = field(
        default=None,
        metadata={
            "help": "Reward model id of a pretrained model hosted inside a model repo on huggingface.co or "
            "local path to a directory containing model weights saved using `PreTrainedModel.save_pretrained`."
        },
    )


def main(script_args, training_args, model_args):
    # Load a pretrained model
    model = AutoModelForCausalLM.from_pretrained(
        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
    )

    def reward_len(completions, **kwargs):
        return [-abs(20 - len(completion)) for completion in completions]
    
    # Load the dataset
    dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config)

    # Initialize the GRPO trainer
    trainer = GRPOTrainer(
        model=model,
        reward_funcs=reward_len,
        args=training_args,
        train_dataset=dataset[script_args.dataset_train_split].select(range(100)),
        eval_dataset=dataset[script_args.dataset_test_split].select(range(100)) if training_args.eval_strategy != "no" else None,
        processing_class=tokenizer,
        peft_config=get_peft_config(model_args),
    )

    # Train and push the model to the Hub
    trainer.train()

    # Save and push to hub
    trainer.save_model(training_args.output_dir)
    # if training_args.push_to_hub:
    #     trainer.push_to_hub(dataset_name=script_args.dataset_name)


def make_parser(subparsers: argparse._SubParsersAction = None):
    dataclass_types = (GRPOScriptArguments, GRPOConfig, ModelConfig)
    if subparsers is not None:
        parser = subparsers.add_parser("grpo", help="Run the GRPO training script", dataclass_types=dataclass_types)
    else:
        parser = TrlParser(dataclass_types)
    return parser


if __name__ == "__main__":
    parser = make_parser()
    script_args, training_args, model_args = parser.parse_args_and_config()
    main(script_args, training_args, model_args)

    '''
    CUDA_VISIBLE_DEVICES=0 python test.py --model_name_or_path /home/grpo_test//Qwen2.5-0.5B-Instruct --dataset_name /home/grpo_test/trl-libtldr \
        --learning_rate 2.0e-4 --num_train_epochs 1 --per_device_train_batch_size 2 \
        --gradient_accumulation_steps 2 --eval_strategy no --logging_steps 2 --use_peft 1 --lora_r 32 --lora_alpha 16 --output_dir Qwen2-0.5B-GRPO
    '''
#加速使用vllm 训练

2.部署与预测

from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer
model = AutoPeftModelForCausalLM.from_pretrained(
    "/home/grpo_test/qwen2-grpo-lora", # YOUR MODEL YOU USED FOR TRAINING
    load_in_4bit = True,
)
tokenizer = AutoTokenizer.from_pretrained("/home/grpo_test/qwen2-grpo-lora")

SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""

text = tokenizer.apply_chat_template([
    {"role" : "user", "content" : "Calculate pi."},
], tokenize = False, add_generation_prompt = True)

inputs = tokenizer(
[
   text,
], return_tensors = "pt").to("cuda")

from transformers import TextStreamer
text_streamer = TextStreamer(tokenizer)
outputs = model.generate(**inputs, max_new_tokens = 128)
print(outputs)
outputs = outputs[:, inputs['input_ids'].shape[1]:]
print(tokenizer.decode(outputs[0]))


'''
As an artificial intelligence language model, I do not calculate or execute numerical operations. However, I can provide you with the mathematical formula to calculate pi. The approximate value of pi is 3.14159, and the full mathematical formula is:

π = ∑ (2 × sin(1/n) × n^2) / (n × (n + 1) × (2n + 1))

Please note that this is an approximate value, and the accuracy can vary depending on the precision required.
'''

3.模型合并

from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer
lora_model = AutoPeftModelForCausalLM.from_pretrained(
    "/home/grpo_test/qwen2-grpo-lora", # YOUR MODEL YOU USED FOR TRAINING
    torch_dtype="auto",
    device_map="cuda:0"
)

model2 = lora_model.merge_and_unload()
print(model2)
model2.save_pretrained("merged-model")


from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model_name = "merged-model"

model1 = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="cuda:0"
)
print(model1)


model1_dict = dict()
model2_dict = dict()

for k, v in model1.named_parameters():
    model1_dict[k] = v
        
for k, v in model2.named_parameters():
    model2_dict[k] = v

keys = model1_dict.keys()
keys2 = model2_dict.keys()
print(set(keys) == set(keys2))
for k in keys:
    print(f'{k}: {torch.allclose(model1_dict[k], model2_dict[k])},')

4.vllm 加载

from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

max_model_len, tp_size = 8000, 1
model_name = "merged-model"
#prompt = [{"role": "user", "content": "你好"}]

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
llm = LLM(
    model=model_name,
    tensor_parallel_size=tp_size,
    max_model_len=max_model_len,
    trust_remote_code=True,
    enforce_eager=True,
    # GLM-4-9B-Chat-1M 如果遇见 OOM 现象，建议开启下述参数
    dtype=torch.float16,
    gpu_memory_utilization=0.35      #降低gpu 显存
    # enable_chunked_prefill=True,
    # max_num_batched_tokens=8192
)

SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""
#自定义输入
prompt = tokenizer.apply_chat_template([
    {"role" : "system", "content" : SYSTEM_PROMPT},
    {"role" : "user", "content" : "How many r's are in strawberry?"},
], tokenize = False, add_generation_prompt = True)


sampling_params = SamplingParams(temperature=0.95, max_tokens = 1000)

inputs = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
outputs = llm.generate(prompts=inputs, sampling_params=sampling_params)
print(outputs[0].outputs[0].text)

4. 代码讲解
   1. 数据处理部分

#数据处理   已包含在compute_loss中。
device = self.accelerator.device
prompts = [x["prompt"] for x in inputs]   #获取prompt 内容
#maybe_apply_chat_template 将输入转换为指定格式的输入。后台会自动判定，是否需要应用template 模板
prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]

#转向量化
prompt_inputs = self.processing_class(
    prompts_text, return_tensors="pt", padding=True, padding_side="left", add_special_tokens=False
)
#gpu化
prompt_inputs = super()._prepare_inputs(prompt_inputs)

if self.max_prompt_length is not None:
    prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -self.max_prompt_length :]
    prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -self.max_prompt_length :]

def maybe_apply_chat_template(
    example: dict[str, list[dict[str, str]]],
    tokenizer: PreTrainedTokenizer,
    tools: Optional[list[Union[dict, Callable]]] = None,
) -> dict[str, str]:
    r"""
    Example:
    ```python
    >>> from transformers import AutoTokenizer
    >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-128k-instruct")
    >>> example = {
    ...     "prompt": [{"role": "user", "content": "What color is the sky?"}],
    ...     "completion": [{"role": "assistant", "content": "It is blue."}]
    ... }
    >>> apply_chat_template(example, tokenizer)
    {'prompt': '<|user|>\nWhat color is the sky?<|end|>\n<|assistant|>\n', 'completion': 'It is blue.<|end|>\n<|endoftext|>'}
    ```
    """
    if is_conversational(example):    #判定是不是问答模式，即 [{"role": "user", "content": "What color is the sky?"}] 
        return apply_chat_template(example, tokenizer, tools)
    else:
        return example

2. 损失部分

#参考模型生成
with torch.inference_mode():
    if self.ref_model is not None:
        ref_per_token_logps = get_per_token_logps(self.ref_model, prompt_completion_ids, num_logits_to_keep)
    else:
        with self.accelerator.unwrap_model(model).disable_adapter():    # 关闭适配器，原生模型作为参考模型
            ref_per_token_logps = get_per_token_logps(model, prompt_completion_ids, num_logits_to_keep) 
 
 #计算模型与参考模型的KL 散度
# Compute the KL divergence between the model and the reference model
per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1

 # Compute the rewards    计算奖励
prompts = [prompt for prompt in prompts for _ in range(self.num_generations)]

rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device)
for i, (reward_func, reward_processing_class) in enumerate(
    zip(self.reward_funcs, self.reward_processing_classes)
):
    if isinstance(reward_func, PreTrainedModel):
        if is_conversational(inputs[0]):
            messages = [{"messages": p + c} for p, c in zip(prompts, completions)]
            texts = [apply_chat_template(x, reward_processing_class)["text"] for x in messages]
        else:
            texts = [p + c for p, c in zip(prompts, completions)]
        reward_inputs = reward_processing_class(
            texts, return_tensors="pt", padding=True, padding_side="right", add_special_tokens=False
        )
        reward_inputs = super()._prepare_inputs(reward_inputs)
        with torch.inference_mode():
            rewards_per_func[:, i] = reward_func(**reward_inputs).logits[:, 0]  # Shape (B*G,)
    else:
        # Repeat all input columns (but "prompt" and "completion") to match the number of generations
        reward_kwargs = {key: [] for key in inputs[0].keys() if key not in ["prompt", "completion"]}
        for key in reward_kwargs:
            for example in inputs:
                # Repeat each value in the column for `num_generations` times
                reward_kwargs[key].extend([example[key]] * self.num_generations)
        output_reward_func = reward_func(prompts=prompts, completions=completions, **reward_kwargs)
        rewards_per_func[:, i] = torch.tensor(output_reward_func, dtype=torch.float32, device=device)

# Sum the rewards from all reward functions      计算出奖励
rewards = rewards_per_func.sum(dim=1)

# Compute grouped-wise rewards          计算组内奖励均值和方差
mean_grouped_rewards = rewards.view(-1, self.num_generations).mean(dim=1)
std_grouped_rewards = rewards.view(-1, self.num_generations).std(dim=1)

# Normalize the rewards to compute the advantages       标准化奖励来计算优势
mean_grouped_rewards = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
std_grouped_rewards = std_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
advantages = (rewards - mean_grouped_rewards) / (std_grouped_rewards + 1e-4)       #优势标准化

#计算最后的损失函数
# x - x.detach() allows for preserving gradients from x      #x - x.detach()允许从x中保留梯度
per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
per_token_loss = -(per_token_loss - self.beta * per_token_kl)
loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()

#per_token_logps - per_token_logps.detach() 的理解：策略比例 采用梯度代替

1. per_token_logps - per_token_logps.detach()的目的是在计算损失函数时，排除那些不需要梯度更新的部分，通常用于避免反向传播到某些特定的操作或变量。