project-llm-trainer 0.3__py3-none-any.whl

llm_trainer/grpo_trainer.py

@@ -0,0 +1,393 @@
+import time
+import copy
+from typing import Tuple, List, Union, Callable, Optional
+import torch
+from torch.utils.data import Dataset
+from torch.nn.utils.rnn import pad_sequence
+import torch.distributed as dist
+import torch.nn.functional as F
+
+from llm_model import LlmModel
+
+from .parallel_ds import DsParallel
+from .trainer import Trainer
+from .train_configs import TrainConfig
+from .dataset import GRPORolloutDataset
+from .loss import GRPOLoss
+from .tools import TrainerTools
+from .generate_utils import batch_generate
+
+from .checkpoint import (
+    save_checkpoint,
+    load_checkpoint_for_eval,
+    save_steps,
+)
+
+class GRPOTrainer(Trainer):
+    def __init__(
+            self,
+            *,
+            train_config: TrainConfig,
+            reward_func: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], List[float]],
+            eval_prompts: List[str],
+            eval_image_tags: Optional[List[int]] = None
+    ):
+        super().__init__(
+            train_config=train_config,
+            eval_prompts=eval_prompts,
+            eval_image_tags=eval_image_tags
+        )
+
+        self.reward_func = reward_func
+        self.reference_model = self._init_reference_model()
+        self.generate_model = self._init_generate_model()
+
+        # 默认使用torch提供的pad_sequence
+        # 如果pad_sequence不支持padding_side参数，则将改参数置为False，使用反转的方式
+        self._use_origin_pad_sequence = True
+
+        # 保存一下train model的checkpoint，方便下面reference_model使用
+        save_checkpoint(self.train_model, self.optimizer)
+
+    def _init_reference_model(self):
+        reference_model = LlmModel(self.train_config.model_config)
+
+        device = 'cpu' # TrainerTools().parallel.device
+        reference_model.to(device)
+        # load_checkpoint_for_eval(model=reference_model, device=device)
+
+        reference_model.eval()
+        for param in reference_model.parameters():
+            param.requires_grad = False
+
+        return reference_model
+
+    def _init_generate_model(self):
+        return copy.deepcopy(self.reference_model)
+        # generate_model = LlmModel(self.train_config.model_config)
+        #
+        # device = 'cpu' #TrainerTools().parallel.device
+        # generate_model.to(device)
+        # # load_checkpoint_for_eval(model=generate_model, device=device)
+        #
+        # generate_model.eval()
+        # for param in generate_model.parameters():
+        #     param.requires_grad = False
+        #
+        # return generate_model
+
+    def _init_loss(self):
+        criterion = GRPOLoss(
+            clip_eps=self.train_config.grpo_config.clip_eps,
+            kl_weight=self.train_config.grpo_config.kl_weight
+        )
+
+        return criterion, None
+
+    def _convert_train_args(self) -> Tuple[dict, dict, dict, bool]:
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim = super()._convert_train_args()
+        data_loader_kwargs.update({"collate_fn": lambda x: x})
+
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
+
+    def _create_dataset(self, file_path) -> Dataset:
+        return GRPORolloutDataset(file_path)
+
+    def _calc_loss(self, inputs, attention_mask, logits, labels): ...
+
+    def _left_pad_sequence(
+            self,
+            sequences: Union[torch.Tensor, List[torch.Tensor]],
+            padding_value: float,
+    ) -> torch.Tensor:
+        if self._use_origin_pad_sequence:
+            try:
+                return pad_sequence(sequences, batch_first=True, padding_value=padding_value, padding_side='left')
+            except:
+                self._use_origin_pad_sequence = False
+                return self._left_pad_sequence(sequences, padding_value)
+        else:
+            # 反转每个序列的顺序（如 [1,2,3] → [3,2,1]）
+            reversed_sequences = [seq.flip(dims=(0,)) for seq in sequences]
+            # 使用默认的右侧填充
+            padded_reversed = pad_sequence(reversed_sequences, batch_first=True, padding_value=padding_value)
+            # 再次反转序列顺序，恢复原始方向（填充在左侧）
+            return padded_reversed.flip(dims=(1,))
+
+    def _selective_log_softmax(self, logits, input_ids):
+        # Convert raw logits into log probabilities along the vocabulary axis.
+        # [batch_size, seq_len, vocab_size]
+        log_probs = F.log_softmax(logits, dim=-1)
+
+        # Reshape input_ids from (batch_size, seq_len) to (batch_size, seq_len, 1) for gathering.
+        # Then, gather the log probability for each token in input_ids.
+        selected_log_probs = log_probs.gather(dim=-1, index=input_ids.unsqueeze(-1))
+
+        # Remove the extra last dimension to get back to shape (batch_size, seq_len).
+        return selected_log_probs.squeeze(-1)
+
+    def _compute_log_probabilities(
+            self,
+            model,
+            input_ids,
+            attention_mask,
+            logits_to_keep
+    ):
+        # prompt部分[1, 2, 3]
+        # 生成模型生成的内容是[4, 5]，logits_to_keep=2
+        # 则下面的输入 [1, 2, 3, 4, 5], 正常情况下输出是[2, 3, 4, 5, 6]
+        # logits_to_keep=2，时输出[5, 6]
+        # 但是我们想要的[4, 5]部分
+        # 所以需要logits_to_keep=2+1，输出[4, 5, 6]
+
+        # [batch_size, total_seq_len, vocab_size]
+        outputs = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            logits_to_keep=logits_to_keep + 1
+        )
+
+        # [batch_size, total_seq_len - 1, vocab_size]
+        logits = outputs['logits'][:, :-1, :]
+
+        input_ids = input_ids[:, -logits_to_keep:]
+        logits = logits[:, -logits_to_keep:, :]
+
+        # Compute and return the log probabilities for the selected tokens.
+        return self._selective_log_softmax(logits, input_ids), outputs['aux_loss']
+
+    def _compute_group_relative_advantages(self, rewards):
+        group_size = self.train_config.grpo_config.group_size
+
+        # Reshape rewards to group by prompt
+        # [batch, group_size]
+        rewards_by_group = rewards.view(-1, group_size)
+
+        # Compute mean and standard deviation for each prompt group
+        # [batch]
+        group_means = rewards_by_group.mean(dim=1)
+        group_stds = rewards_by_group.std(dim=1)
+
+        # Expand the means and stds to match the original flat rewards tensor shape
+        # [batch*group_size]
+        expanded_means = group_means.repeat_interleave(group_size)
+        expanded_stds = group_stds.repeat_interleave(group_size)
+
+        # Normalize rewards to get advantages
+        # [batch*group_size]
+        advantages = (rewards - expanded_means) / (expanded_stds + 1e-4)
+
+        # [batch*group_size, 1]
+        return advantages.unsqueeze(1)  # Add dimension for token-wise operations
+
+    def _generate_completions(self, prompts, group_size: int):
+        pad_token_id = TrainerTools().tokenizer.pad
+        device = TrainerTools().parallel.device
+
+        # 左边添加pad，对齐prompt长度
+        # [batch, max_prompt_len]
+        prompt_ids = self._left_pad_sequence(prompts, padding_value=pad_token_id)
+        prompt_ids = prompt_ids.to(device)
+
+        prompt_len = prompt_ids.shape[1]
+
+        # [batch*group_size, max_prompt_len]
+        prompt_ids = prompt_ids.repeat_interleave(group_size, 0)
+        # [batch*group_size, max_prompt_len]
+        prompt_masks = prompt_ids != pad_token_id
+
+        # [batch*group_size, max_prompt_len+max_gen_len]
+        outputs: torch.Tensor = batch_generate(
+            # model=self.train_model,
+            model=self.generate_model,
+            tokens=prompt_ids,
+            pad_token_id=pad_token_id,
+            attention_mask=prompt_masks,
+            max_position_embeddings=self.train_config.model_config.max_position_embeddings,
+            max_new_tokens=self.train_config.grpo_config.gen_max_new_tokens,
+            temperature=self.train_config.grpo_config.gen_temperature,
+            k=self.train_config.grpo_config.gen_k,
+            p=self.train_config.grpo_config.gen_p,
+            device=device,
+            suppress_tokens=self.train_config.grpo_config.gen_suppress_tokens
+        )
+
+        # [batch*group_size, max_gen_len]
+        completion_ids = outputs[:, prompt_len:]
+        # [batch*group_size, max_gen_len]
+        completion_masks = (completion_ids != pad_token_id).int()
+
+        return prompt_ids, prompt_masks, completion_ids, completion_masks
+
+    def _generate_rollout_data(self, batch_data: List[dict]):
+        prompts = [item["prompt"] for item in batch_data]
+        answers = [item["answer"] for item in batch_data]
+        group_size = self.train_config.grpo_config.group_size
+
+        # 使用no_grad替换inference_mode
+        # 修复问题：Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal
+        with torch.no_grad():
+        # with torch.inference_mode():
+            prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_completions(prompts, group_size)
+            input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
+            attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
+            logits_to_keep = completion_ids.shape[1]
+
+            # Compute old_log_probs from the current model, with gradients disabled.
+            old_log_probs, _ = self._compute_log_probabilities(self.generate_model, input_ids, attention_mask, logits_to_keep)
+
+            # Compute ref_log_probs from the reference model, which remains static.
+            ref_log_probs, _ = self._compute_log_probabilities(self.reference_model, input_ids, attention_mask, logits_to_keep)
+
+        repeated_prompts = [p for p in prompts for _ in range(group_size)]
+        repeated_answers = [a for a in answers for _ in range(group_size)]
+
+        return {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask,
+            'completion_mask': completion_mask,
+            'old_log_probs': old_log_probs,
+            'ref_log_probs': ref_log_probs,
+            'completion_ids': completion_ids,
+            'repeated_prompts': repeated_prompts,
+            'repeated_answers': repeated_answers,
+            'logits_to_keep': logits_to_keep
+        }
+
+    def _maximize_grpo_objective(self, rollout_data):
+        device = TrainerTools().parallel.device
+
+        input_ids = rollout_data['input_ids']
+        attention_mask = rollout_data['attention_mask']
+        completion_mask = rollout_data['completion_mask']
+        old_log_probs = rollout_data['old_log_probs']
+        ref_log_probs = rollout_data['ref_log_probs']
+        logits_to_keep = rollout_data['logits_to_keep']
+        completion_ids = rollout_data['completion_ids']
+        repeated_prompts = rollout_data['repeated_prompts']
+        repeated_answers = rollout_data['repeated_answers']
+
+        # [batch*group_size]
+        rewards = torch.tensor(
+            self.reward_func(repeated_prompts, completion_ids, repeated_answers),
+            dtype=torch.float32,
+            device=device
+        )
+
+        # [batch*group_size, 1]
+        advantages = self._compute_group_relative_advantages(rewards)
+
+        # Compute current log probabilities
+        log_probs, aux_loss = self._compute_log_probabilities(self.train_model, input_ids, attention_mask, logits_to_keep)
+
+        loss = self.criterion(
+            log_probs=log_probs,
+            old_log_probs=old_log_probs,
+            ref_log_probs=ref_log_probs,
+            completion_mask=completion_mask,
+            advantages=advantages
+        )
+
+        return loss, aux_loss
+
+    def train(self):
+        global_steps = 0
+        skipping_train = False
+        device = TrainerTools().parallel.device
+        aux_loss_coef = self.train_config.loss_config.aux_loss_coef
+
+        for epoch in range(self.train_config.n_epochs):
+            load_checkpoint_for_eval(model=self.reference_model, device=device)
+            self.train_model.train()
+            file_count = len(self.train_config.file_dataset)
+
+            for file_idx in range(file_count):
+                file_path = self.train_config.file_dataset[file_idx]
+                dataset = self._create_dataset(file_path)
+
+                train_data_loader = TrainerTools().parallel.process_dataloader(
+                    dataset=dataset,
+                    data_loader_kwargs=self.data_loader_kwargs,
+                    sampler_kwargs=self.sampler_kwargs
+                )
+
+                last_ckpt_batch = 0
+                batch_count_per_file = len(train_data_loader)
+
+                TrainerTools().parallel.on_epoch_start(epoch)
+                self._on_file_start(epoch, file_path)
+
+                for batch, batch_data in enumerate(train_data_loader):
+                    global_steps += 1
+                    if global_steps < self.last_global_steps:
+                        skipping_train = True
+                        continue
+
+                    skipping_train = False
+
+                    # start generate
+                    # 使用单独的模型生成数据， 原因是在deepspeed并行训练时，使用train_model生成数据会卡死
+                    self.generate_model.to(TrainerTools().parallel.device)
+                    self.reference_model.to(TrainerTools().parallel.device)
+
+                    # 保存了train_model checkpoint后，这里保证生成模型使用的参数是最新
+                    load_checkpoint_for_eval(self.generate_model, TrainerTools().parallel.device)
+                    # 生成数据
+                    rollout_data = self._generate_rollout_data(batch_data)
+
+                    # 卸载到cpu上，等待下次使用时再to gpu
+                    self.generate_model.to('cpu')
+                    self.reference_model.to('cpu')
+                    torch.cuda.empty_cache()
+                    # end generate
+
+                    try:
+                        for grpo_step in range(self.train_config.grpo_config.grpo_steps):
+                            with self.ctx:
+                                loss, aux_loss = self._maximize_grpo_objective(rollout_data)
+                                if aux_loss_coef and aux_loss:
+                                    loss += aux_loss_coef * aux_loss
+
+                            self._backward_loss(loss)
+
+                            if TrainerTools().parallel.parallel_train:
+                                dist.all_reduce(loss, dist.ReduceOp.AVG)
+
+                            # ds模式已经集成gradient_clipping
+                            if not isinstance(TrainerTools().parallel, DsParallel) and self.lr_scheduler.can_clip_grad():
+                                # clip grad
+                                self.scalar.unscale_(self.optimizer)
+                                torch.nn.utils.clip_grad_norm_(self.train_model.parameters(), 1.0)
+
+                            self._step()
+
+                            self._log_loss(
+                                epoch_tag=f'epoch: {epoch}',
+                                file_tag=f'file: {file_idx + 1}/{file_count}',
+                                batch_tag=f'batch: {batch}/{batch_count_per_file}',
+                                loss=loss.detach().item()
+                            )
+                    except Exception as e:
+                        self._on_exception(e, epoch, batch)
+                    finally:
+                        save_steps(global_steps=global_steps, lr_scheduler=self.lr_scheduler)
+
+                        if (batch - last_ckpt_batch) >= self.train_config.eval_batch_interval:
+                            save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                            last_ckpt_batch = batch
+                            self._on_batch_end(tag=f'epoch:{epoch}/batch:{batch}')
+
+                        try:
+                            del loss
+                        except UnboundLocalError: ...
+
+            # end epoch
+            if not skipping_train:
+                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                save_steps(global_steps=global_steps, lr_scheduler=self.lr_scheduler)
+                TrainerTools().parallel.on_epoch_end(epoch)
+                self._on_epoch_end(tag=f'epoch:{epoch}')
+
+        # 等待checkpoint保存完成
+        time.sleep(10)
+        TrainerTools().parallel.destroy()

llm_trainer/log.py

@@ -0,0 +1,16 @@
+import time, os
+
+def get_log_dir() -> str:
+    log_dir = os.environ['LOG_DIR']
+    if not os.path.exists(log_dir):
+        os.mkdir(log_dir)
+
+    return f'{log_dir}/' if not log_dir.endswith('/') else log_dir
+
+def log(msg: str, log_file=None):
+    cur_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
+    if not log_file:
+        print(f'[{cur_time}] {msg}')
+    else:
+        with open(log_file, 'a') as f:
+            f.write(f"[{cur_time}] {msg}")

llm_trainer/loss.py

@@ -0,0 +1,171 @@
+from typing import List, Optional, Tuple
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class LMLoss(nn.Module):
+    """
+    llm loss
+    """
+    def __init__(
+            self,
+            ignore_index: int = -100,
+            *,
+            critical_tokens: Optional[List[int]] = None,
+            critical_alpha: float = 1.0,
+            vocab_size: int = 0
+    ):
+        super().__init__()
+        self.ignore_index = ignore_index
+        self.critical_tokens = critical_tokens
+        self.critical_alpha = critical_alpha
+
+        if critical_tokens and vocab_size > 0:
+            self.register_buffer('weights', torch.ones(vocab_size))
+            # 为关键token设置权重
+            self.weights[self.critical_tokens] = critical_alpha
+
+
+    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        # logits shape (batch, seq_len, vocab_size)
+        # labels shape (batch, seq_len)
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+
+        logits = shift_logits.reshape(-1, logits.shape[-1])
+        targets = shift_labels.reshape(-1)
+
+        ce_loss = F.cross_entropy(
+            logits,
+            targets,
+            ignore_index=self.ignore_index,
+            weight=self.weights.to(logits.device, dtype=logits.dtype) if self.critical_tokens else None
+        )
+
+        # 添加额外惩罚项（可选）
+        # if self.critical_tokens:
+        #     crit_mask = torch.isin(targets, torch.tensor(self.critical_tokens).to(targets.device))
+        #     crit_logits = logits[crit_mask]
+        #     crit_targets = targets[crit_mask]
+        #     extra_loss = F.cross_entropy(crit_logits, crit_targets, ignore_index=self.ignore_index)
+        #     return ce_loss + extra_loss * (self.critical_alpha - 1)  # 增强惩罚
+
+        return ce_loss
+
+
+class KDLoss(nn.Module):
+    """
+    Language Model Knowledge Distillation Loss
+    https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/models/loss.py#L266
+    """
+    def __init__(self, ignore_index: int = -100):
+        super().__init__()
+        self.ignore_index = ignore_index
+
+    def forward(self, logits: torch.Tensor, teacher_logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        teacher_probs = F.softmax(teacher_logits, dim=-1, dtype=torch.float32)
+        inf_mask = torch.isinf(logits)
+
+        logprobs = F.log_softmax(logits, dim=-1, dtype=torch.float32)
+        prod_probs = torch.masked_fill(teacher_probs * logprobs, inf_mask, 0)
+
+        x = torch.sum(prod_probs, dim=-1).view(-1)
+        mask = (labels != self.ignore_index).int()
+
+        distil_loss = -torch.sum(x * mask.view(-1), dim=0) / torch.sum(mask.view(-1), dim=0)
+
+        return distil_loss
+
+
+class DPOLoss(nn.Module):
+    def __init__(
+            self,
+            beta: float,
+            label_smoothing: float = 0.0,
+            ipo: bool = False
+    ):
+        super().__init__()
+        self.beta = beta
+        self.label_smoothing = label_smoothing
+        self.ipo = ipo
+
+    def forward(
+            self,
+            policy_logps: torch.Tensor,
+            reference_logps: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size = reference_logps.shape[0]
+        ref_chosen_probs = reference_logps[:batch_size//2]
+        ref_reject_probs = reference_logps[batch_size//2:]
+        policy_chosen_probs = policy_logps[:batch_size//2]
+        policy_reject_probs = policy_logps[batch_size//2:]
+
+        pi_logratios = policy_chosen_probs - policy_reject_probs
+        ref_logratios = ref_chosen_probs - ref_reject_probs
+        logits = pi_logratios - ref_logratios
+
+        if self.ipo:
+            losses = (logits - 1 / (2 * self.beta)) ** 2  # Eq. 17 of https://arxiv.org/pdf/2310.12036v2.pdf
+        else:
+            # Eq. 3 https://ericmitchell.ai/cdpo.pdf; label_smoothing=0 gives original DPO (Eq. 7 of https://arxiv.org/pdf/2305.18290.pdf)
+            losses = (
+                    -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
+                    - F.logsigmoid(-self.beta * logits) * self.label_smoothing
+            )
+
+        loss = losses.mean()
+        # chosen_rewards = self.beta * (policy_chosen_probs - ref_chosen_probs).detach()
+        # rejected_rewards = self.beta * (policy_reject_probs - ref_reject_probs).detach()
+
+        return loss
+
+
+class GRPOLoss(nn.Module):
+    def __init__(
+            self,
+            clip_eps: float,
+            kl_weight: float
+    ):
+        super().__init__()
+        self.clip_eps = clip_eps
+        self.kl_weight = kl_weight
+
+    def forward(
+            self,
+            log_probs: torch.Tensor,
+            old_log_probs: torch.Tensor,
+            ref_log_probs: torch.Tensor,
+            completion_mask: torch.Tensor,
+            advantages: torch.Tensor
+    ) -> torch.Tensor:
+        # Compute policy ratio
+        ratio = torch.exp(log_probs - old_log_probs)
+
+        # Compute surrogate loss with clipping
+        surrogate1 = ratio * advantages
+        surrogate2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
+        surrogate_loss = torch.min(surrogate1, surrogate2)
+
+        # Compute KL divergence penalty
+        kl_div = torch.exp(ref_log_probs - log_probs) - (ref_log_probs - log_probs) - 1
+
+        # Combine losses
+        per_token_loss = surrogate_loss - self.kl_weight * kl_div
+        loss = -((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
+
+        return loss
+
+        # kl = self._approx_kl_divergence(
+        #     log_probs=log_probs,
+        #     ref_log_probs=ref_log_probs,
+        #     mask=mask,
+        # )
+        #
+        # ratio = (log_probs - old_log_probs).exp()
+        # surr1 = ratio * advantages
+        # surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages
+        # loss = -torch.min(surr1, surr2) + self.kl_weight * kl
+        #
+        # loss = self._masked_mean(loss, mask, dim=-1).mean()
+        # return loss, kl.mean()