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

llm_trainer/__init__.py

@@ -0,0 +1,6 @@
+from .trainer import Trainer
+from .sft_trainer import SFTTrainer
+from .dpo_trainer import DPOTrainer
+from .grpo_trainer import GRPOTrainer
+from .tools import TrainerTools, FileDataset, estimate_data_size
+from .generate_utils import generate, streaming_generate

llm_trainer/checkpoint.py

@@ -0,0 +1,161 @@
+import os
+from typing import Optional, Union, Tuple
+import torch
+from torch import nn
+from torch.optim import Optimizer
+
+from .parallel_ds import DsParallel
+from .parallel_fsdp import FsdpParallel
+from .parallel_ddp import DdpParallel
+from .scheduler import LRScheduler
+from .tools import TrainerTools
+
+try:
+    from .dcp import save_dcp, load_dcp, convert_dcp_to_pth
+except:
+    os.environ['ENABLE_DCP'] = "0"
+
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+
+# https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html
+
+DEFAULT_CHECKPOINT_NAME = "checkpoint.pth"
+
+
+def _can_use_dcp(model: nn.Module) -> bool:
+    if os.environ.get('ENABLE_DCP', '1') != '1':
+        return False
+
+    # 如果是fsdp或者ddp，才能使用dcp保存
+    if (isinstance(TrainerTools().parallel, FsdpParallel)
+            or isinstance(TrainerTools().parallel, DdpParallel)):
+        return True
+
+    return False
+
+
+def save_checkpoint(
+        model: nn.Module,
+        optimizer: Optional[Optimizer] = None,
+        suffix: Optional[str] = None
+):
+    if isinstance(TrainerTools().parallel, DsParallel):
+        from .ds_checkpoint import save_ds_checkpoint
+        save_ds_checkpoint(model, suffix)
+    elif _can_use_dcp(model):
+        save_dcp(model, optimizer, suffix)
+    else:
+        if isinstance(model, FSDP):
+            # 未经过测试 参考：https://doc.hfai.high-flyer.cn/haiscale/haiscale_fsdp.html
+            # 是否使用rank0_only=True？
+            with FSDP.summon_full_params(
+                    module=model,
+                    rank0_only=True,
+                    writeback=False,
+                    offload_to_cpu=True
+            ):
+                if TrainerTools().parallel.is_main_process:
+                    checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
+                    if suffix:
+                        checkpoint_name = f"{checkpoint_name}_{suffix}"
+
+                    ckpt = {'model_state_dict': model.state_dict()}
+
+                    if optimizer:
+                        ckpt.update({'optim_state_dict': optimizer.state_dict()})
+
+                    torch.save(ckpt, checkpoint_name)
+        else:
+            if TrainerTools().parallel.is_main_process:
+                checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
+                if suffix:
+                    checkpoint_name = f"{checkpoint_name}_{suffix}"
+
+                ckpt = {'model_state_dict': TrainerTools().parallel.raw_model.state_dict()}
+
+                if optimizer:
+                    ckpt.update({'optim_state_dict': optimizer.state_dict()})
+
+                torch.save(ckpt, checkpoint_name)
+
+
+def load_checkpoint(
+        model: nn.Module,
+        optimizer: Optional[Optimizer] = None,
+        device: Optional[Union[torch.device, str]] = None,
+        load_module_only: bool = False,
+        suffix: Optional[str] = None
+):
+    if isinstance(TrainerTools().parallel, DsParallel):
+        from .ds_checkpoint import load_ds_checkpoint
+        load_ds_checkpoint(model, load_module_only=load_module_only, suffix=suffix)
+    elif _can_use_dcp(model):
+        load_dcp(model, optimizer, suffix)
+    else:
+        checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
+        if suffix:
+            checkpoint_name = f"{checkpoint_name}_{suffix}"
+
+        if os.path.exists(checkpoint_name):
+            # 未经过测试，else的逻辑经过测试在fsdp下也没问题
+            if isinstance(model, FSDP):
+                with FSDP.summon_full_params(module=model):
+                    state_dict = torch.load(checkpoint_name, weights_only=True, map_location=device)
+                    model.load_state_dict(state_dict['model_state_dict'])
+
+                    if optimizer:
+                        optimizer.load_state_dict(state_dict['optim_state_dict'])
+            else:
+                state_dict = torch.load(checkpoint_name, weights_only=True, map_location=device)
+                model.load_state_dict(state_dict['model_state_dict'])
+
+                if optimizer:
+                    optimizer.load_state_dict(state_dict['optim_state_dict'])
+
+
+def load_checkpoint_for_eval(
+        model: nn.Module,
+        device: Optional[Union[torch.device, str]] = None,
+        suffix: Optional[str] = None
+):
+    if isinstance(TrainerTools().parallel, DsParallel):
+        from .ds_checkpoint import load_ds_checkpoint_for_eval
+        load_ds_checkpoint_for_eval(model)
+    elif _can_use_dcp(model):
+        checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
+
+        # load_dcp方式在cpu上会报错，所以改为先将ckpt转换为pth，然后再加载pth
+        # load_dcp(model, optimizer)
+        pth_name = os.environ.get('EVAL_CHECKPOINT_NAME', checkpoint_name)
+        if suffix:
+            pth_name = f'{pth_name}_{suffix}'
+
+        convert_dcp_to_pth(pth_name)
+
+        if os.path.exists(pth_name):
+            ckpt = torch.load(pth_name, map_location=device, weights_only=True)
+            model.load_state_dict(ckpt['app']['model_state_dict'])
+            # 使用完删除
+            os.remove(pth_name)
+    else:
+        load_checkpoint(model, None, device, suffix=suffix)
+
+
+def save_steps(global_steps: int, lr_scheduler: Optional[LRScheduler] = None):
+    # 暂时只保存主进程的
+    if TrainerTools().parallel.is_main_process:
+        steps_checkpoint_name = f"{os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)}.steps"
+        ckpt = {'global_steps': global_steps, 'lr_steps': lr_scheduler.cur_steps}
+        torch.save(ckpt, steps_checkpoint_name)
+
+
+def load_steps(
+        default_global_steps: int = 0,
+        default_lr_steps: int = 0
+) -> Tuple[Optional[int], Optional[int]]:
+    steps_checkpoint_name = f"{os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)}.steps"
+    if os.path.exists(steps_checkpoint_name):
+        ckpt = torch.load(steps_checkpoint_name, weights_only=True)
+        return ckpt['global_steps'], ckpt['lr_steps']
+
+    return default_global_steps, default_lr_steps

llm_trainer/dataset.py

@@ -0,0 +1,140 @@
+import os.path
+
+import torch
+from torch.utils.data import Dataset
+import pickle
+
+from .tools import TrainerTools
+from .utils import extra_image_tag_and_repeat_image_tok
+
+
+def _try_load_pkl(file_path: str):
+    tokens = None
+    try:
+        with open(file_path, 'rb') as f:
+            tokens = pickle.load(f)
+    except Exception as e:
+        raise e
+    finally:
+        return tokens
+
+
+class TextDataset(Dataset):
+    """
+    适用于pretrain阶段
+    """
+    def __init__(self, file_path, block_size, stride):
+        super().__init__()
+
+        self.input_ids = []
+
+        tokens = _try_load_pkl(file_path)
+        if not tokens:
+            cache_file = f'{file_path}.cache'
+            if os.path.exists(cache_file):
+                tokens = _try_load_pkl(cache_file)
+            else:
+                tokens = []
+                with open(file_path, 'r') as f:
+                    for line in f:
+                        tokens.extend(TrainerTools().tokenizer.encode(line))
+
+                with open(cache_file, 'wb') as f:
+                    pickle.dump(tokens, f)
+
+        for i in range(0, len(tokens) - block_size + 1, stride):
+            self.input_ids.append(tokens[i:i+block_size])
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def __getitem__(self, item):
+        return torch.tensor(self.input_ids[item]).long()
+
+
+class LineByLineTextDataset(Dataset):
+    """
+    适用于sft阶段
+    """
+    def __init__(self, file_path, max_len, tokens_per_image=-1):
+        super().__init__()
+
+        self.max_len = max_len
+        self.tokens_per_image = tokens_per_image
+        self.input_ids = []
+
+        tokens = _try_load_pkl(file_path)
+        if not tokens:
+            cache_file = f'{file_path}.cache'
+            if os.path.exists(cache_file):
+                tokens = _try_load_pkl(cache_file)
+            else:
+                tokens = []
+                with open(file_path, 'r') as f:
+                    for line in f:
+                        tokens.append(TrainerTools().tokenizer.encode(line))
+
+                with open(cache_file, 'wb') as f:
+                    pickle.dump(tokens, f)
+
+        self.input_ids = tokens
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def __getitem__(self, item):
+        inputs = self.input_ids[item]
+        if self.tokens_per_image != -1:
+            inputs, image_tag = extra_image_tag_and_repeat_image_tok(inputs, self.tokens_per_image)
+        else:
+            image_tag = None
+
+        inputs = inputs[:self.max_len]
+
+        return {'inputs': torch.tensor(inputs).long(), 'image_tag': image_tag}
+
+
+class DPODataset(Dataset):
+    def __init__(self, file_path, max_len):
+        self.max_len = max_len
+        self.chosen_ids = []
+        self.rejected_ids = []
+
+        # [{'chosen': xxx, 'rejected': xxx} ...]
+        tokens = _try_load_pkl(file_path)
+        for token in tokens:
+            self.chosen_ids.append(token['chosen'])
+            self.rejected_ids.append(token['rejected'])
+
+    def __len__(self):
+        return len(self.chosen_ids)
+
+    def __getitem__(self, item):
+        chosen_id = self.chosen_ids[item]
+        rejected_id = self.rejected_ids[item]
+
+        return {'chosen': chosen_id[:self.max_len], 'rejected': rejected_id[:self.max_len]}
+
+
+class GRPORolloutDataset(Dataset):
+    def __init__(self, file_path):
+        self.questions = []
+        self.answers = []
+
+        # [{'question': xxx, 'answer': ''}]
+        tokens = _try_load_pkl(file_path)
+        for token in tokens:
+            self.questions.append(token['prompt'])
+            self.answers.append(token['answer'])
+
+    def __len__(self):
+        return len(self.questions)
+
+    def __getitem__(self, item):
+        question = self.questions[item]
+        answer = self.answers[item]
+
+        return {
+            'prompt': torch.tensor(question).long(),
+            'answer': torch.tensor(answer).long()
+        }

llm_trainer/dcp.py

@@ -0,0 +1,93 @@
+import os
+from typing import Optional, Dict, Any
+from torch import nn
+from torch.optim import Optimizer
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint.stateful import Stateful
+from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
+from torch.distributed.checkpoint.format_utils import dcp_to_torch_save, torch_save_to_dcp
+
+DEFAULT_CHECKPOINT_DIR = "checkpoint"
+
+class AppState(Stateful):
+    def __init__(self, model: nn.Module, optimizer: Optimizer):
+        self.model = model
+        self.optimizer = optimizer
+
+    def state_dict(self) -> Dict[str, Any]:
+        model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
+        return {
+            'model_state_dict': model_state_dict,
+            'optim_state_dict': optimizer_state_dict
+        }
+
+    def load_state_dict(self, state_dict: Dict[str, Any]):
+        set_state_dict(
+            model=self.model,
+            optimizers=self.optimizer,
+            model_state_dict=state_dict['model_state_dict'],
+            optim_state_dict=state_dict['optim_state_dict']
+        )
+
+
+def save_dcp(
+        model: nn.Module,
+        optimizer: Optimizer,
+        suffix: Optional[str] = None
+):
+    checkpoint_id = os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR)
+    if suffix:
+        checkpoint_id = f"{checkpoint_id}_{suffix}"
+
+    state_dict = {'app': AppState(model, optimizer)}
+
+    # fs_storage_writer = dcp.FileSystemWriter(checkpoint_id, overwrite=True)
+    # dcp.save(state_dict=state_dict, storage_writer=fs_storage_writer)
+    dcp.save(state_dict=state_dict, checkpoint_id=checkpoint_id)
+
+
+def load_dcp(
+        model: nn.Module,
+        optimizer: Optional[Optimizer] = None,
+        suffix: Optional[str] = None
+):
+    checkpoint_id = os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR)
+    if suffix:
+        checkpoint_id = f"{checkpoint_id}_{suffix}"
+
+    if os.path.exists(checkpoint_id):
+        state_dict = {'app': AppState(model, optimizer)}
+        # AppState帮助加载到state_dict中, 然后加载到model中
+        dcp.load(state_dict=state_dict, checkpoint_id=checkpoint_id)
+
+        # if isinstance(model, FSDP):
+        #     state_dict = {'app': AppState(model, optimizer)}
+        #     # AppState帮助加载到state_dict中, 然后加载到model中
+        #     dcp.load(state_dict=state_dict, checkpoint_id=checkpoint_id)
+        # else:
+        #     state_dict = {"model_state_dict": model.state_dict()}
+        #
+        #     if optimizer:
+        #         state_dict.update({'optim_state_dict': optimizer.state_dict()})
+        #
+        #     # since no progress group is initialized, DCP will disable any collectives.
+        #     # 加载到state_dict中，然后通过model.load_state_dict加载到model中
+        #     dcp.load(
+        #         state_dict=state_dict,
+        #         checkpoint_id=checkpoint_id,
+        #     )
+        #
+        #     model.load_state_dict(state_dict["model_state_dict"])
+        #     if optimizer:
+        #         optimizer.load_state_dict(state_dict["optim_state_dict"])
+
+def convert_dcp_to_pth(pth_path: str):
+    dcp_path = os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR)
+    if os.path.exists(dcp_path):
+        # convert dcp model to torch.save (assumes checkpoint was generated as above)
+        dcp_to_torch_save(dcp_path, pth_path)
+
+def convert_pth_to_dcp(pth_path: str):
+    if os.path.exists(pth_path):
+        # converts the torch.save model back to DCP
+        torch_save_to_dcp(pth_path, os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR))