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

llm_trainer/dpo_trainer.py

@@ -0,0 +1,311 @@
+from typing import Tuple, List, Optional
+import gc
+import torch
+from torch.utils.data import Dataset
+from itertools import islice
+
+from .base_trainer import BaseTrainer
+from .train_configs import TrainConfig
+from .dataset import DPODataset
+from .loss import DPOLoss
+from .tools import TrainerTools
+from .utils import (
+    autocast,
+    get_dpo_collate_fn,
+    log_softmax,
+    disable_dropout_in_model
+)
+
+from .checkpoint import (
+    save_checkpoint,
+    save_steps,
+)
+from .log import Logger
+
+
+class DPOTrainer(BaseTrainer):
+    def __init__(
+            self,
+            *,
+            train_config: TrainConfig,
+            eval_prompts: List[str]
+    ):
+        self.dpo_config = train_config.dpo_config
+        super().__init__(
+            train_config=train_config,
+            eval_prompts=eval_prompts,
+            gradient_accumulation_steps=self.dpo_config.gradient_accumulation_steps
+        )
+        self.ref_model = self._init_ref_model()
+
+    def _init_ref_model(self):
+        ref_model = self._new_model(self.train_config)
+
+        if self.dpo_config.ref_model_checkpoint:
+            ref_model.load_state_dict(self.dpo_config.ref_model_checkpoint)
+            self.dpo_config.ref_model_checkpoint = {}
+
+        ref_model.eval()
+        for param in ref_model.parameters():
+            param.requires_grad = False
+
+        ref_model, _ = TrainerTools().parallel.process(
+            model=ref_model,
+            optimizer=None,
+            kwargs=self._init_ref_model_args(),
+            save_instance=False
+        )
+
+        return ref_model
+
+    def _new_model(self, train_config: TrainConfig):
+        model = super()._new_model(train_config)
+        disable_dropout_in_model(model)
+        return model
+
+    def _init_loss(self):
+        criterion = DPOLoss(
+            beta=self.dpo_config.loss_beta,
+            label_smoothing=self.dpo_config.loss_label_smoothing,
+            ipo=self.dpo_config.loss_ipo
+        )
+
+        return criterion, None
+
+    def _convert_train_args(self) -> Tuple[dict, dict, dict]:
+        dpo_collate_fn = get_dpo_collate_fn(self.dpo_config.mask_prompt)
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs = super()._convert_train_args()
+        data_loader_kwargs.update({"collate_fn": dpo_collate_fn})
+
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs
+
+    def _create_dataset(self, file_idx) -> Tuple[Dataset, str]:
+        file_path = self.train_config.file_dataset[file_idx]
+        block_size = self.train_config.dataset_block_size
+        return DPODataset(file_path, block_size), file_path
+
+    def _calc_loss(self, inputs, attention_mask, logits, labels): ...
+
+    def _logprobs(self, logits, labels):
+        """
+        Calculate the average log probabilities for a batch of sequences.
+
+        Args:
+            logits (torch.Tensor): Logits from the model with shape (B, T, V)
+            labels (torch.Tensor): Ground truth labels with shape (B, T).
+
+        Returns:
+            torch.Tensor: Average log probabilities for each sequence in the batch.
+                          Shape is (B,) representing the mean log probability for each sequence.
+        """
+        loss_masks = (labels != -100)
+
+        logits = logits[:, :-1, :]
+        labels = labels[:, 1:].clone()
+        loss_masks = loss_masks[:, 1:]
+
+        # dummy token; we'll ignore the losses on these tokens later
+        labels[labels == -100] = 0
+
+        # Gather the log probabilities for the actual labels
+        per_token_logps = log_softmax(logits, labels)
+
+        # Apply the mask to set log-probs of padding tokens to 0
+        logprobs_sums = (per_token_logps * loss_masks).sum(-1)
+        logprobs_means = (per_token_logps * loss_masks).sum(-1) / loss_masks.sum(-1).clamp(min=1.0)
+
+        return logprobs_sums, logprobs_means
+
+    def train(self):
+        # 梯度累积步数
+        gradient_accumulation_steps = max(1, self.gradient_accumulation_steps)
+
+        loss_accumulation = 0.0
+        aux_loss_accumulation = 0.0
+        nll_loss_accumulation = 0.0
+        batches_accumulated = 0
+
+        aux_loss_coef = self.train_config.loss_config.aux_loss_coef
+        nll_loss_coef = self.dpo_config.nll_loss_coef
+
+        for epoch in range(self.resume_epoch, self.train_config.n_epochs):
+            self.train_model.train()
+            file_count = len(self.train_config.file_dataset)
+            start_file_idx = self.resume_file_idx if epoch == self.resume_epoch else 0
+
+            for file_idx in range(start_file_idx, file_count):
+                dataset, file_path = self._create_dataset(file_idx)
+                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)
+
+                skip_batches = 0
+                if epoch == self.resume_epoch and file_idx == self.resume_file_idx:
+                    skip_batches = self.resume_batch_idx
+                    if skip_batches > 0 and TrainerTools().parallel.is_main_process:
+                        Logger.std_log(f"Fast forwarding {skip_batches} batches in {file_path}...")
+
+                data_iterator = iter(train_data_loader)
+                if skip_batches > 0:
+                    data_iterator = islice(data_iterator, skip_batches, None)
+                    last_ckpt_batch = skip_batches
+
+                for batch, batch_data in enumerate(data_iterator):
+                    batch = skip_batches + batch
+
+                    # 是否需要更新梯度
+                    if gradient_accumulation_steps > 1:
+                        need_update_grad = (batch + 1) % gradient_accumulation_steps == 0 or batch == batch_count_per_file - 1
+                    else:
+                        need_update_grad = True
+
+                    try:
+                        chosen_inputs: torch.Tensor = batch_data['chosen_inputs'].to(TrainerTools().parallel.device)
+                        chosen_labels: torch.Tensor = batch_data['chosen_labels'].to(TrainerTools().parallel.device)
+
+                        rejected_inputs: torch.Tensor = batch_data['rejected_inputs'].to(TrainerTools().parallel.device)
+                        rejected_labels: torch.Tensor = batch_data['rejected_labels'].to(TrainerTools().parallel.device)
+
+                        chosen_attention_masks: torch.Tensor = chosen_inputs != TrainerTools().tokenizer.pad
+                        rejected_attention_masks: torch.Tensor = rejected_inputs != TrainerTools().tokenizer.pad
+
+                        # 在batch维度concat
+                        # [chosen, chosen, reject, reject]
+                        concat_inputs = torch.concat([chosen_inputs, rejected_inputs], dim=0)
+                        concat_labels = torch.concat([chosen_labels, rejected_labels], dim=0)
+                        concat_attention_masks = torch.concat([chosen_attention_masks, rejected_attention_masks], dim=0)
+
+                        if TrainerTools().parallel.parallel_train:
+                            self.train_model.require_backward_grad_sync = need_update_grad
+
+                        with autocast(TrainerTools().parallel.device_type):
+                            policy_outputs = self.train_model(concat_inputs, attention_mask=concat_attention_masks)
+                            policy_logprobs_sums, policy_logprobs_means = self._logprobs(policy_outputs['logits'], concat_labels)
+
+                            with torch.no_grad():
+                                ref_outputs = self.ref_model(concat_inputs, attention_mask=concat_attention_masks)
+                                ref_logprobs_sums, _ = self._logprobs(ref_outputs['logits'], concat_labels)
+
+                            policy_chosen_logps = policy_logprobs_sums[:chosen_inputs.shape[0]]
+                            policy_rejected_logps = policy_logprobs_sums[chosen_inputs.shape[0]:]
+
+                            ref_chosen_logps = ref_logprobs_sums[:chosen_inputs.shape[0]]
+                            ref_rejected_logps = ref_logprobs_sums[chosen_inputs.shape[0]:]
+
+                            nll_loss = -policy_logprobs_means[:chosen_inputs.shape[0]].mean()
+
+                            # calc loss
+                            loss = self.criterion(
+                                policy_chosen_logps,
+                                policy_rejected_logps,
+                                ref_chosen_logps,
+                                ref_rejected_logps
+                            )
+
+                            if aux_loss_coef and policy_outputs.get('aux_loss'):
+                                aux_loss = aux_loss_coef * policy_outputs.get('aux_loss')
+                            else:
+                                aux_loss = torch.tensor(0.0, device=loss.device, dtype=loss.dtype)
+
+                            if nll_loss_coef and nll_loss:
+                                nll_loss = nll_loss_coef * nll_loss
+                            else:
+                                nll_loss = torch.tensor(0.0, device=loss.device, dtype=loss.dtype)
+
+                        if gradient_accumulation_steps > 1:
+                            loss = loss / gradient_accumulation_steps
+                            aux_loss = aux_loss / gradient_accumulation_steps
+                            nll_loss = nll_loss / gradient_accumulation_steps
+
+                        total_loss = loss + aux_loss + nll_loss
+                        self._backward_loss(total_loss)
+
+                        loss_accumulation += total_loss.detach().item()
+                        aux_loss_accumulation += aux_loss.detach().item()
+                        nll_loss_accumulation += nll_loss.detach().item()
+
+                        batches_accumulated += 1
+
+                        if need_update_grad:
+                            self._apply_grad_clipping()
+                            self._apply_step()
+
+                            avg_loss, avg_aux_loss, avg_nll_loss = self._avg_loss(
+                                losses=[
+                                    loss_accumulation,
+                                    aux_loss_accumulation,
+                                    nll_loss_accumulation,
+                                ],
+                                gradient_accumulation_steps=gradient_accumulation_steps,
+                                batches_accumulated=batches_accumulated
+                            )
+
+                            self._log(
+                                keys={
+                                    'epoch': epoch,
+                                    'file': f'{file_idx + 1}/{file_count}',
+                                    'batch': f'{batch + 1}/{batch_count_per_file}',
+                                },
+                                values={
+                                    'loss': avg_loss,
+                                    'moe_aux_loss': avg_aux_loss,
+                                    'nll_loss': avg_nll_loss
+                                }
+                            )
+
+                            # reset to default
+                            loss_accumulation = 0.0
+                            aux_loss_accumulation = 0.0
+                            nll_loss_accumulation = 0.0
+                            batches_accumulated = 0
+
+                            if (batch - last_ckpt_batch) >= self.train_config.eval_config.eval_batch_interval:
+                                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                                save_steps(
+                                    epoch=epoch,
+                                    file_idx=file_idx,
+                                    batch_idx=batch + 1,
+                                    lr_scheduler=self.lr_scheduler
+                                )
+
+                                last_ckpt_batch = batch
+                                self._on_batch_end(tag=f'epoch:{epoch}/batch:{batch}')
+                    except Exception as e:
+                        self._on_exception(e, epoch, batch)
+
+                # 一个文件训练结束后，清理内存
+                del train_data_loader
+                del dataset
+                if hasattr(TrainerTools().parallel, '_sampler'):
+                    TrainerTools().parallel._sampler = None
+
+                gc.collect()
+                torch.cuda.empty_cache()
+
+            # end epoch
+
+            # reset resume state
+            self.resume_file_idx = 0
+            self.resume_batch_idx = 0
+
+            save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+            save_steps(
+                epoch=epoch + 1,
+                file_idx=0,
+                batch_idx=0,
+                lr_scheduler=self.lr_scheduler
+            )
+
+            TrainerTools().parallel.on_epoch_end(epoch)
+            self._on_epoch_end(tag=f'epoch:{epoch}')
+
+        TrainerTools().parallel.destroy()
+

llm_trainer/ds_checkpoint.py

@@ -0,0 +1,72 @@
+import os
+from glob import glob
+from typing import Optional
+import shutil
+import torch
+from torch import nn
+from .tools import TrainerTools
+
+try:
+    import deepspeed
+    from deepspeed import DeepSpeedEngine
+    from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+except: ...
+
+"""
+函数	功能	是否加载模型到内存	是否保存到文件 主要用途
+get_fp32_state_dict_from_zero_checkpoint	从 ZeRO 检查点提取 FP32 状态字典	否	否	获取模型权重，用于推理、迁移等
+load_state_dict_from_zero_checkpoint	从 ZeRO 检查点加载模型和优化器状态	是	否	恢复训练状态，继续训练
+convert_zero_checkpoint_to_fp32_state_dict	将 ZeRO 检查点转换为独立的 FP32 状态字典文件	否	是	创建可移植的 FP32 权重文件，用于部署、分享等
+"""
+
+def save_ds_checkpoint(
+        model: nn.Module,
+        extra_module: Optional[nn.Module] = None
+):
+    assert isinstance(model, DeepSpeedEngine)
+    ckpt_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
+
+    try:
+        # 包括model、optimizer等状态
+        model.save_checkpoint(save_dir=ckpt_dir)
+    except: ...
+
+    # 只在main rank上执行
+    if TrainerTools().parallel.is_main_process:
+        if extra_module:
+            torch.save(extra_module.state_dict(), os.path.join(ckpt_dir, "extra_module_state_dict.pt"))
+
+        # 最多保存多少checkpoint，默认为2
+        max_to_keep = int(os.environ.get('CKPT_MAX_TO_KEEP', '2'))
+        # 删除历史checkpoint
+        ckpt_paths = glob(os.path.join(ckpt_dir, "global_*"))
+        if len(ckpt_paths) > max_to_keep:
+            # 按修改时间排序，找到最旧的目录
+            oldest_ckpt = sorted(ckpt_paths, key=os.path.getmtime)[0]
+            try:
+                shutil.rmtree(oldest_ckpt)
+            except: ...
+
+    TrainerTools().parallel.wait('remove old ds checkpoint')
+
+
+def load_ds_checkpoint(
+        model: nn.Module,
+        load_module_only: bool = False,
+        extra_module: Optional[nn.Module] = None
+):
+    assert isinstance(model, DeepSpeedEngine)
+    ckpt_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
+
+    # 包括model、optimizer等状态
+    if os.path.exists(ckpt_dir):
+        model.load_checkpoint(
+            load_dir=ckpt_dir,
+            load_module_only=load_module_only
+        )
+
+        path = os.path.join(ckpt_dir, "extra_module_state_dict.pt")
+        if os.path.exists(path):
+            state = torch.load(path, map_location=TrainerTools().parallel.device, weights_only=True)
+            extra_module.load_state_dict(state)
+

llm_trainer/eval.py

@@ -0,0 +1,33 @@
+import os
+import torch
+
+from .generate_utils import generate
+from .tools import TrainerTools
+from .train_configs import TrainConfig
+from .log import _get_log_dir
+
+def submit_gen_task(
+        eval_model: torch.nn.Module,
+        train_config: TrainConfig,
+        tag,
+        prompt,
+        pixel_values,
+        tokens_per_image
+):
+    tokens = TrainerTools().tokenizer.encode(prompt, unsqueeze=True, covert_tensor=True)
+    max_new_tokens = max(train_config.eval_config.max_seq_len - tokens.shape[1], 0)
+
+    gen_result = generate(
+        eval_model,
+        prompt=tokens,
+        max_new_tokens=max_new_tokens,
+        temperature=train_config.eval_config.temperature,
+        k=train_config.eval_config.top_k,
+        p=train_config.eval_config.top_p,
+        pixel_values=pixel_values,
+        tokens_per_image=tokens_per_image,
+        device=TrainerTools().parallel.device
+    )
+
+    with open(os.path.join(_get_log_dir(), 'gen.txt'), 'a') as f:
+        f.write(f"{tag}, gen->{gen_result}\n")