class LlamaForCausalLM(LlamaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        # 根据配置初始化LLaMA模型，此模型的结构与Meta发布的LLaMA一致，除了不包含最后一层
        self.model = LlamaModel(config)
        # 定义最后一层全连接层
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        ...

    def forward(...):
        # 从LLaMA的模型获取预测的结果，并取最后一个Transformer块的计算结果
        outputs = self.model(...)
        hidden_states = outputs[0]

        # 计算模型输出
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # 取预测结果的前N-1个，使预测的下一个词与标签词对应。
            shift_logits = logits[..., :-1, :].contiguous()
            # 取标签数据的后N-1个，使预测的下一个词与标签词对应。
            shift_labels = labels[..., 1:].contiguous()
            # 创建交叉熵损失，用于计算模型预测结果与标签之间的差异。
            loss_fct = CrossEntropyLoss()
            # 将预测结果和标签进行形状变换，展平为二维张量。第一个维度为样本的数量，第二个维度为词汇表的大小。
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # 计算损失
            loss = loss_fct(shift_logits, shift_labels)

        return CausalLMOutputWithPast(loss=loss, logits=logits, ...)

def train(args):
    strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
    with strategy.model_init_context():
        # 将transformers库中的LlamaForCausalLM模型转化为Lora模型
        model = convert_to_lora_module(LlamaForCausalLM.from_pretrained(args.pretrain), args.lora_rank)

    # 准备分词器
    tokenizer = AutoTokenizer.from_pretrained(args.pretrain, padding_side="right", ...)
    tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, model)

    # 准备优化器
    optim = HybridAdam(model.parameters(), lr=args.lr, clipping_norm=1.0)

    # 准备数据集
    train_dataset = SFTDataset(train_data, tokenizer, max_len)
    eval_dataset = SFTDataset(eval_data, tokenizer, max_len)
    train_dataloader = DataLoader(train_dataset, ...)
    eval_dataloader = DataLoader(eval_dataset, ...)
 
    # 构造训练器并开始训练
    (model, optim) = strategy.prepare((model, optim))
    trainer = SFTTrainer(model=model, ...)
    trainer.fit(...)

class SFTTrainer:
    ...
    def fit(self, ...):
        for epoch in range(self.max_epochs):
            # 设置模型为训练模式
            self.model.train()
            # 对数据集里面的每一个批次进行训练
            for batch_id, batch in enumerate(self.train_dataloader):
                # 执行模型前向计算
                outputs = self.model(batch["input_ids"], batch["attention_mask"], batch["labels"])
                # 下面的代码采用了一种累计梯度的机制，可以让模型在批太小的场景下也能较为稳定的更新
                # 每次计算梯度时，将损失平均一下，再计算梯度
                loss = outputs.loss
                loss = loss / self.accumulation_steps
                # 计算梯度
                self.strategy.backward(loss, self.model, self.optimizer)

                # 当达到定义的累计批次数时更新参数并将梯度置零
                if (batch_id + 1) % self.accumulation_steps == 0:
                    self.strategy.optimizer_step(self.optimizer)  # 更新参数
                    self.optimizer.zero_grad()  # 梯度置零
                    self.scheduler.step()  # 对学习率进行调整

def train(args):
    strategy = ColossalAIStrategy(stage=2, placement_policy='cuda')
    with strategy.model_init_context():
        # 定义基于LLaMA的奖励模型
        model = LlamaRM(pretrained=args.pretrain, lora_rank=args.lora_rank)
    # 以下代码与监督微调部分类似
    tokenizer = LlamaTokenizer.from_pretrained(args.pretrain)
    tokenizer = prepare_llama_tokenizer_and_embedding(tokenizer, model)
    optim = HybridAdam(model.parameters(), lr=5e-6)

    # 定义LogSig损失，LogSig损失是OpenAI在论文Training language models to follow instructions with human feedback中定义的损失函数
    loss_fn = LogSigLoss()

    data = load_dataset(args.dataset)
    train_dataset = RmStaticDataset(data['train'], tokenizer, ...)
    valid_dataset = RmStaticDataset(data['test'], tokenizer, ...)
    train_dataloader = DataLoader(train_dataset, ...)
    eval_dataloader = DataLoader(eval_dataset, ...)

    (model, optim) = strategy.prepare((model, optim))
    trainer = RewardModelTrainer(model=model, ...)
    trainer.fit()

class LogSigLoss(nn.Module):
    def forward(self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor) -> torch.Tensor:
        # 根据OpenAI在论文中的损失计算公式计算损失，见下文的分析
        probs = torch.sigmoid(chosen_reward - reject_reward)
        log_probs = torch.log(probs)
        loss = -log_probs.mean()
        return loss

class RewardModel(LoRAModule):
    def __init__(self, ...) -> None:
        # 转换为Lora模型，以便支持少量参数微调
        self.convert_to_lora()

    def forward(self, sequences: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # 转换为Lora模型，以便支持少量参数微调
        outputs = self.model(sequences, attention_mask=attention_mask)
        last_hidden_states = outputs['last_hidden_state']
        values = self.value_head(last_hidden_states)[:, :-1]
        value = values.mean(dim=1).squeeze(1)    # ensure shape is (B)
        return value

class LlamaRM(RewardModel):
    def __init__(self, ...) -> None:
        # 从预训练模型文件中加载LLaMA的基础模型，使用transformers库的实现，无最后一个线性层
        model = LlamaModel.from_pretrained(pretrained)
        # 定义线性最后层，输出维度为1，即一个数值型的奖励值
        value_head = nn.Linear(model.config.hidden_size, 1)