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

llm_trainer/scheduler.py

@@ -0,0 +1,220 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+import math
+import torch
+from .log import Logger
+
+class LRScheduler(ABC):
+    @property
+    @abstractmethod
+    def cur_steps(self): ...
+
+    @property
+    @abstractmethod
+    def cur_lr(self): ...
+
+    @abstractmethod
+    def step(self): ...
+
+    @abstractmethod
+    def can_clip_grad(self): ...
+
+    @abstractmethod
+    def get_ckpt_dict(self) -> dict: ...
+
+    @abstractmethod
+    def restore_ckpt_dict(self, ckpt: dict): ...
+
+
+class WarmupCosineAnnealingLRScheduler(LRScheduler):
+    def __init__(
+            self,
+            *,
+            optimizer: torch.optim.Optimizer,
+            warmup_iters: int,
+            initial_lr: float,
+            min_lr: float,
+            max_lr: float,
+            cosine_annealing_period: int, # 每个周期的步数
+            cosine_annealing_period_mul: int = 0, # 周期长度的倍数
+            param_group_indices: Optional[List[int]] = None,
+            need_log: bool = False
+    ):
+        super().__init__()
+
+        self._optimizer = optimizer
+        self._initial_lr = initial_lr
+        self._min_lr = min_lr
+        self._max_lr = max_lr
+        self._warmup_iters = warmup_iters
+
+        self._cosine_annealing_period = cosine_annealing_period
+        self._cosine_annealing_period_mul = cosine_annealing_period_mul
+
+        self.param_group_indices = param_group_indices
+
+        self.T_cur = 0  # 当前周期内已走过的步数
+        self.cycle = 0  # 当前周期编号
+
+        if warmup_iters != 0:
+            self._lr_increment = (max_lr - initial_lr) / warmup_iters
+        else:
+            self._lr_increment = 0
+
+        self._steps = -1
+        self._current_lr = initial_lr
+        self._cosine_annealing_base_lr = None
+
+        if need_log:
+            self.logger = Logger('lr.txt')
+        else:
+            self.logger = None
+
+    @property
+    def cur_steps(self):
+        return self._steps
+
+    @property
+    def cur_lr(self):
+        return self._current_lr
+
+    def step(self):
+        self._steps += 1
+        self._update_lr()
+
+    def can_clip_grad(self):
+        return self._steps > self._warmup_iters
+
+    def _update_lr(self):
+        if self.param_group_indices is None:
+            target_groups = self._optimizer.param_groups
+        else:
+            target_groups = [self._optimizer.param_groups[i] for i in self.param_group_indices]
+
+        # 如果period_mul是0，则认为没有周期，超过余弦退火总步数，则一直保持最小lr
+        if self._cosine_annealing_period_mul == 0 and self._steps >= self._cosine_annealing_period + self._warmup_iters:
+            lr = self._min_lr
+            for param_group in target_groups:
+                param_group['lr'] = lr
+        elif self._steps <= self._warmup_iters:
+            # Warmup: adjust learning rate linearly
+            # (max_lr - initial_lr) / warmup_iters
+            lr = self._initial_lr + self._steps * self._lr_increment
+            for param_group in target_groups:
+                param_group['lr'] = lr
+        else:
+            if not self._cosine_annealing_base_lr:
+                self._cosine_annealing_base_lr = self.cur_lr
+
+            """每步更新学习率"""
+            # 计算当前周期的最大步数
+            T_max = self._cosine_annealing_period * (max(self._cosine_annealing_period_mul, 1) ** self.cycle)
+
+            # 更新周期状态
+            self.T_cur += 1
+            calc_t = self.T_cur
+
+            if self.T_cur >= T_max:
+                if self._cosine_annealing_period_mul == 0:
+                    self.T_cur = T_max
+                    calc_t = T_max
+                else:
+                    self.cycle += 1
+                    self.T_cur = 0
+                    calc_t = T_max
+
+            # 计算并设置新学习率
+            cos_factor = (1 + math.cos(math.pi * calc_t / T_max)) / 2
+            lr = self._min_lr + (self._cosine_annealing_base_lr - self._min_lr) * cos_factor
+
+            for param_group in target_groups:
+                param_group['lr'] = lr
+
+        self._current_lr = lr
+
+        if self.logger:
+            self.logger.log(f"step: {self.cur_steps}, lr: {lr}", log_to_console=False)
+
+    def get_ckpt_dict(self) -> dict:
+        return {
+            'cur_lr': self._current_lr,
+            'lr_steps': self.cur_steps,
+            'cosine_annealing_base_lr': self._cosine_annealing_base_lr,
+            't_cur': self.T_cur,
+            'cycle': self.cycle,
+        }
+
+    def restore_ckpt_dict(self, ckpt: dict):
+        if 'cur_lr' in ckpt:
+            self._current_lr = ckpt['cur_lr']
+
+        if 'lr_steps' in ckpt:
+            self._steps = ckpt['lr_steps']
+
+        if 'cosine_annealing_base_lr' in ckpt:
+            self._cosine_annealing_base_lr = ckpt['cosine_annealing_base_lr']
+
+        if 't_cur' in ckpt:
+            self.T_cur = ckpt['t_cur']
+
+        if 'cycle' in ckpt:
+            self.cycle = ckpt['cycle']
+
+        self._update_lr()
+
+
+class NoneLRScheduler(LRScheduler):
+    def __init__(self, initial_lr):
+        self._current_lr = initial_lr
+
+    @property
+    def cur_steps(self):
+        return -1
+
+    @property
+    def cur_lr(self):
+        return self._current_lr
+
+    def step(self): ...
+
+    def can_clip_grad(self):
+        return True
+
+    def get_ckpt_dict(self) -> dict:
+        return {'cur_lr': self._current_lr}
+
+    def restore_ckpt_dict(self, ckpt: dict):
+        if 'cur_lr' in ckpt:
+            self._current_lr = ckpt['cur_lr']
+
+
+class CompositeLRScheduler(LRScheduler):
+    def __init__(self, schedulers: List[LRScheduler]):
+        self.schedulers = schedulers
+
+    @property
+    def cur_steps(self):
+        return self.schedulers[0].cur_steps if self.schedulers else 0
+
+    @property
+    def cur_lr(self):
+        return self.schedulers[0].cur_lr if self.schedulers else 0.0
+
+    def step(self):
+        for scheduler in self.schedulers:
+            scheduler.step()
+
+    def can_clip_grad(self):
+        return all(s.can_clip_grad() for s in self.schedulers)
+
+    def get_ckpt_dict(self) -> dict:
+        ckpt = {}
+        for i, scheduler in enumerate(self.schedulers):
+            ckpt[f'scheduler_{i}'] = scheduler.get_ckpt_dict()
+        return ckpt
+
+    def restore_ckpt_dict(self, ckpt: dict):
+        for i, scheduler in enumerate(self.schedulers):
+            key = f'scheduler_{i}'
+            if key in ckpt:
+                scheduler.restore_ckpt_dict(ckpt[key])

llm_trainer/sft_trainer.py

@@ -0,0 +1,97 @@
+from typing import Optional, Tuple, List
+from torch.utils.data import Dataset
+
+from llm_model import (
+    VLMConfig,
+    LlmModel,
+    VlmModel
+)
+
+from .base_trainer import BaseTrainer
+from .train_configs import TrainConfig
+from .dataset import SFTDataset
+from .utils import get_sft_collate_fn
+from .tools import TrainerTools
+
+
+class SFTTrainer(BaseTrainer):
+    def __init__(
+            self,
+            *,
+            train_config: TrainConfig,
+            eval_prompts: List[str],
+            eval_image_tags: Optional[List[str]] = None
+    ):
+        self.sft_config = train_config.sft_config
+        self.pixel_values_provider = self.sft_config.pixel_values_provider
+        self.eval_image_tags = eval_image_tags
+
+        super().__init__(
+            train_config=train_config,
+            eval_prompts=eval_prompts,
+            kd_config=self.sft_config.kd_config,
+            gradient_accumulation_steps=self.sft_config.gradient_accumulation_steps
+        )
+
+        if isinstance(train_config.model_config, VLMConfig):
+            self.pixel_values_provider = self.sft_config.pixel_values_provider
+        else:
+            self.pixel_values_provider = None
+
+    def _new_model(self, train_config: TrainConfig):
+        if isinstance(train_config.model_config, VLMConfig):
+            return VlmModel(train_config.model_config)
+        else:
+            return LlmModel(train_config.model_config)
+
+    def _check_freeze_llm_model(self, model):
+        # freeze llm model for vlm training
+        if self.sft_config.freeze_llm_model:
+            for name, param in model.named_parameters():
+                if not any(sub_module in name for sub_module in ['multi_modal_projector']):
+                    param.requires_grad = False
+
+            # model.embed_tokens.eval()
+            # model.layers.eval()
+            # model.head_norm.eval()
+            # model.lm_head.eval()
+
+    def _convert_train_args(self) -> Tuple[dict, dict, dict]:
+        sft_collate_fn = get_sft_collate_fn(self.sft_config.mask_prompt)
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs = super()._convert_train_args()
+        data_loader_kwargs.update({"collate_fn": sft_collate_fn})
+
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs
+
+    def _get_pixel_values(self, batch_data):
+        if self.pixel_values_provider and 'image_tags' in batch_data:
+            image_tags = batch_data['image_tags']
+            return self.pixel_values_provider(image_tags).to(TrainerTools().parallel.device)
+
+        return None
+
+    def _get_eval_pixel_values_and_tokens_count(self, eval_idx):
+        if not self.eval_image_tags:
+            return None, None
+
+        eval_image_tag = self.eval_image_tags[eval_idx]
+        if isinstance(self.train_config.model_config, VLMConfig) and self.pixel_values_provider and eval_image_tag:
+            return self.pixel_values_provider([eval_image_tag]), self.train_config.model_config.tokens_per_image
+
+        return None, None
+
+    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
+
+        image_tag_file_path = None
+        tokens_per_image = -1
+
+        if isinstance(self.train_config.model_config, VLMConfig):
+            if self.sft_config.image_tags_file_dataset:
+                image_tag_file_path = self.sft_config.image_tags_file_dataset[file_idx]
+
+            if self.train_config.model_config.tokens_per_image:
+                tokens_per_image = self.train_config.model_config.tokens_per_image
+
+        return SFTDataset(file_path, block_size, image_tag_file_path, tokens_per_image), file_path

llm_trainer/tokenizer.py

@@ -0,0 +1,162 @@
+import os
+import warnings
+from typing import List, Dict, Union
+from transformers import AutoTokenizer
+import torch
+
+
+class Tokenizer:
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(os.environ['TOKEN_DIR'])
+
+        self.text_end = '</s>'
+
+        self.text_pad = '<pad>'
+        self.text_unk = '<unk>'
+
+        self.text_user = '<user>'
+        self.text_assistant = '<assistant>'
+
+        self.text_think_start = '<think>'
+        self.text_think_end = '</think>'
+
+        self.text_answer_start = '<answer>'
+        self.text_answer_end = '</answer>'
+
+        self.text_system = '<system>'
+
+        self.text_image = '<image>'
+
+        self.end = self.tokenizer.convert_tokens_to_ids(self.text_end)
+
+        self.pad = self.tokenizer.convert_tokens_to_ids(self.text_pad)
+        self.unk = self.tokenizer.convert_tokens_to_ids(self.text_unk)
+
+        self.user = self.tokenizer.convert_tokens_to_ids(self.text_user)
+        self.assistant = self.tokenizer.convert_tokens_to_ids(self.text_assistant)
+
+        self.think_start = self.tokenizer.convert_tokens_to_ids(self.text_think_start)
+        self.think_end = self.tokenizer.convert_tokens_to_ids(self.text_think_end)
+
+        self.answer_start = self.tokenizer.convert_tokens_to_ids(self.text_answer_start)
+        self.answer_end = self.tokenizer.convert_tokens_to_ids(self.text_answer_end)
+
+        self.system = self.tokenizer.convert_tokens_to_ids(self.text_system)
+        self.image = self.tokenizer.convert_tokens_to_ids(self.text_image)
+
+        self.vocab_size = len(self.tokenizer)
+
+    def encode(
+            self,
+            text: str,
+            unsqueeze: bool = False,
+            covert_tensor: bool = False
+    ) -> Union[torch.Tensor, List[int]]:
+        # [x,x,x]
+        encoded = self.tokenizer.encode(text, add_special_tokens=False)
+
+        if unsqueeze:
+            # tensor: [[x,x,x]]
+            return torch.tensor(encoded, dtype=torch.long).unsqueeze(0)
+        else:
+            # tensor: # [x,x,x]
+            if covert_tensor:
+                return torch.tensor(encoded, dtype=torch.long)
+
+            return encoded
+
+    def batch_encode(
+            self,
+            text: List[str],
+            padding = False,
+            truncation = False,
+            covert_tensor: bool = False,
+            return_attention_mask: bool = False
+    ) -> Union[torch.Tensor, List[List[int]]]:
+        encoded = self.tokenizer(
+            text,
+            padding=padding,
+            truncation=truncation,
+            return_attention_mask=return_attention_mask
+        )['input_ids']
+
+        if covert_tensor:
+            encoded = torch.tensor(encoded, dtype=torch.long)
+
+        return encoded
+
+    def decode(
+            self,
+            token: Union[torch.Tensor, List[int]],
+            skip_special_tokens: bool = False
+    ) -> str:
+        return self.tokenizer.decode(token, skip_special_tokens=skip_special_tokens)
+
+    def batch_decode(
+            self,
+            tokens: Union[torch.Tensor, List[int], List[List[int]]],
+            skip_special_tokens: bool = False
+    ) -> List[str]:
+        return self.tokenizer.batch_decode(tokens, skip_special_tokens=skip_special_tokens)
+
+    def encode_to_token(self, text: str, unsqueeze=True, covert_tensor=True):
+        warnings.warn('encode_to_token is deprecated. Please use `encode` instead.')
+        return self.encode(text, unsqueeze, covert_tensor)
+
+    def decode_to_text(self, token: torch.Tensor, skip_special_tokens: bool = False) -> str:
+        warnings.warn('decode_to_text is deprecated. Please use `decode` instead.')
+        return self.decode(token.squeeze(0), skip_special_tokens)
+
+    def apply_chat_template(
+            self,
+            conversations: List[Dict[str, str]],
+            tokenizer: bool = True,
+            add_answer_tag_for_assistant: bool = True,
+            unsqueeze=False,
+            covert_tensor=False
+    ):
+        """
+            [
+                {"role":"system", "content":"system prompt"},
+                {"role":"user", "content":"hello?"},
+                {"role":"assistant", "content":"hello"},
+                {"role":"user", "content":"hello hello?"},
+                {"role":"assistant", "think":"thinking", "content":"hello hello"},
+            ]
+            <system>{system_prompt}</s><user>hello?</s><assistant>hello</s><user>hello hello?</s><assistant><think>thinking</think><answer>hello hello</answer></s>
+        """
+
+        chat_template = ''
+        support_roles = {'system': self.text_system, 'user': self.text_user, 'assistant': self.text_assistant}
+        for conversation in conversations:
+            role = conversation['role']
+            if role in support_roles:
+                content = conversation['content']
+                if add_answer_tag_for_assistant and role == 'assistant':
+                    content = f"{self.text_answer_start}{content}{self.text_answer_end}"
+
+                if 'think' in conversation:
+                    content = f"{self.text_think_start}{conversation['think']}{self.text_think_end}{content}"
+
+                chat_template = f"{chat_template}{support_roles[role]}{content}{self.text_end}"
+
+        if tokenizer:
+            return self.encode(chat_template, unsqueeze, covert_tensor)
+
+        return chat_template
+
+    def get_special_tokens_dict(self):
+        return {
+            self.text_end: self.end,
+            self.text_pad: self.pad,
+            self.text_unk: self.unk,
+            self.text_user: self.user,
+            self.text_assistant: self.assistant,
+            self.text_think_start: self.think_start,
+            self.text_think_end: self.think_end,
+            self.text_answer_start: self.answer_start,
+            self.text_answer_end: self.answer_end,
+            self.text_system: self.system,
+            self.text_image: self.image,
+        }
+

llm_trainer/tools.py

@@ -0,0 +1,116 @@
+import os
+from abc import ABC, abstractmethod
+import torch
+from .tokenizer import Tokenizer
+from .parallel import DsParallel, DdpParallel, NoneParallel
+from .log import Logger
+
+
+parallel_types = {
+    'ds': DsParallel,
+    'ddp': DdpParallel,
+    'none': NoneParallel
+}
+
+dtypes = {
+    'float': torch.float,
+    'float16': torch.float16,
+    'float32': torch.float32,
+    'float64': torch.float64
+}
+
+class TrainerTools:
+    def __init__(self):
+        if not hasattr(TrainerTools, "_first_init"):
+            TrainerTools._first_init = True
+
+            self.parallel = self._new_parallel()
+
+            self.tokenizer = Tokenizer()
+            self.use_amp = 'cuda' in self.parallel.device and not isinstance(self.parallel, DsParallel)
+
+            Logger.std_log(f'word_size={self.parallel.world_size}, use_amp={self.use_amp}')
+
+    def _new_parallel(self):
+        parallel_type = os.environ.get('PARALLEL_TYPE', 'none')
+        Logger.std_log(f'parallel_type={parallel_type}')
+        return parallel_types[parallel_type]()
+
+    def __new__(cls, *args, **kwargs):
+        if not hasattr(TrainerTools, "_instance"):
+            TrainerTools._instance = object.__new__(cls)
+
+        return TrainerTools._instance
+
+
+class FileDataset(ABC):
+    @abstractmethod
+    def __len__(self) -> int: ...
+
+    @abstractmethod
+    def __getitem__(self, idx) -> str: ...
+
+
+def estimate_data_size(
+        file_dataset: FileDataset,
+        block_size: int,
+        type: str
+) -> int:
+    """
+    估计数据集大小
+    """
+    data_size = 0
+    files_count = len(file_dataset)
+
+    if type == 'sft':
+        from .dataset import SFTDataset
+        for idx in range(files_count):
+            dataset = SFTDataset(file_dataset[idx], block_size)
+            data_size += len(dataset)
+    elif type == 'dpo':
+        from .dataset import DPODataset
+        for idx in range(files_count):
+            dataset = DPODataset(file_dataset[idx], block_size)
+            data_size += len(dataset)
+    elif type == 'grpo' or type == 'ppo':
+        from .dataset import RLDataset
+        for idx in range(files_count):
+            dataset = RLDataset(file_dataset[idx])
+            data_size += len(dataset)
+    else:
+        from .dataset import PretrainDataset
+        for idx in range(files_count):
+            dataset = PretrainDataset(
+                file_dataset[idx],
+                block_size,
+                block_size
+            )
+            data_size += len(dataset)
+
+    return data_size
+
+
+def extract_policy_weights_from_ppo(model_config, ppo_weights):
+    from llm_model import LlmModel
+    from .ppo_trainer import PolicyAndValueModelWrapper, ValueModel
+
+    policy_model = LlmModel(model_config)
+    value_model = ValueModel(LlmModel(model_config))
+
+    wrapper = PolicyAndValueModelWrapper(policy_model, value_model)
+    wrapper.load_state_dict(ppo_weights)
+
+    return wrapper.policy_model.state_dict()
+
+
+def extract_value_weights_from_ppo(model_config, ppo_weights):
+    from llm_model import LlmModel
+    from .ppo_trainer import PolicyAndValueModelWrapper, ValueModel
+
+    policy_model = LlmModel(model_config)
+    value_model = ValueModel(LlmModel(model_config))
+
+    wrapper = PolicyAndValueModelWrapper(policy_model, value_model)
+    wrapper.load_state_dict(ppo_weights)
+
+    return wrapper.value_model.state_dict()