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

llm_trainer/parallel.py

@@ -0,0 +1,146 @@
+import os
+from typing import Optional, Tuple
+from abc import ABC, abstractmethod
+
+import torch
+from torch import nn
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from .log import log
+
+
+class Parallel(ABC):
+    def __init__(
+            self,
+            init_process_group: bool = True,
+            use_parallel: bool = True,
+            use_compile: bool = False
+    ):
+        self._initialize(init_process_group, use_parallel, use_compile)
+
+    def _initialize(
+            self,
+            init_process_group: bool,
+            use_parallel: bool,
+            use_compile: bool
+    ):
+        self._global_rank: int = int(os.environ.get('RANK', -1))
+        self._local_rank: int = int(os.environ.get('LOCAL_RANK', -1))
+        self._use_parallel: bool = use_parallel and self._global_rank != -1
+        self._use_compile = use_compile
+
+        self._sampler: Optional[DistributedSampler] = None
+
+        self.model: Optional[nn.Module] = None
+        self.raw_model: Optional[nn.Module] = None
+
+        if use_compile:
+            torch.set_float32_matmul_precision('high')
+
+        if self._use_parallel:
+            if init_process_group:
+                dist.init_process_group(backend='nccl')
+
+            self.device: str = f'cuda:{self._local_rank}'
+            self.device_type: str = 'cuda'
+
+            torch.cuda.set_device(self.device)
+
+            log(f'global_rank:{self._global_rank},local_rank:{self._local_rank}, world_size:{self.world_size}')
+        else:
+            device = "cpu"
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                device = "mps"
+
+            self.device: str = device
+            self.device_type: str = device
+
+
+    @abstractmethod
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]: ...
+
+    def process_dataloader(
+            self,
+            dataset: Dataset,
+            data_loader_kwargs: dict,
+            sampler_kwargs: Optional[dict]=None
+    ) -> DataLoader:
+        """
+        :param dataset:
+        :param data_loader_kwargs
+                "batch_size" int,
+                "pin_memory" bool,
+                "collate_fn" collate_fn,
+                "num_workers" int
+                "shuffle" bool
+                "drop_last" bool
+        :param sampler_kwargs:
+                "shuffle" bool
+                "drop_last" bool
+        :return:
+        """
+
+        if self._use_parallel:
+            self._sampler = DistributedSampler(dataset=dataset, **sampler_kwargs)
+            return DataLoader(dataset=dataset, sampler=self._sampler, **data_loader_kwargs)
+
+        return DataLoader(dataset=dataset, **data_loader_kwargs)
+
+    def on_epoch_start(self, epoch):
+        if self._sampler:
+            self._sampler.set_epoch(epoch)
+
+    def on_epoch_end(self, epoch): ...
+
+    def synchronize(self):
+        if self._use_parallel:
+            torch.cuda.synchronize(device=self.device)
+
+    def destroy(self):
+        if self._use_parallel:
+            dist.destroy_process_group()
+
+    # def reduce_loss(self, avg_loss: torch.Tensor, loss: torch.Tensor, batch) -> torch.Tensor:
+    #     if self._use_parallel:
+    #         world_size = dist.get_world_size()
+    #         if world_size < 2:
+    #             return loss.detach()
+    #
+    #         torch.distributed.all_reduce(loss)
+    #         # 整个训练过程的滑动损失均值=在历史平均损失的基础上，加上最新损失再求平均
+    #         avg_loss = (avg_loss * batch + loss.detach()) / (batch + 1)
+    #         return avg_loss
+    #
+    #     return loss.detach()
+
+    @property
+    def parallel_train(self) -> bool:
+        return self._use_parallel
+
+    @property
+    def is_main_process(self) -> bool:
+        if self._use_parallel:
+            return self._global_rank == 0
+
+        return True
+
+    @property
+    def world_size(self) -> int:
+        if self._use_parallel:
+            return dist.get_world_size()
+        return 1
+
+    def wait(self):
+        try:
+            log(f'wait at {self.device}')
+            dist.barrier()
+        except: ...
+        log(f'continue at {self.device}')

llm_trainer/parallel_ddp.py

@@ -0,0 +1,39 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from .parallel import Parallel
+
+
+# python3 -m torch.distributed.launch --nproc_per_node=2 --nnodes=1 gpt.py
+# torchrun --standalone --nproc_per_node=gpu pretrain.py
+#    --standalone 代表单机运行
+#    --nproc_per_node=gpu 代表使用所有可用GPU, 等于号后也可写gpu数量n, 这样会使用前n个GPU
+
+
+class DdpParallel(Parallel):
+    def __init__(self):
+        super().__init__()
+
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
+        model.to(self.device)
+
+        if self._use_compile:
+            model = torch.compile(model)
+
+        if self._use_parallel:
+            # self.model = DDP(module=model, broadcast_buffers=False, find_unused_parameters=True)
+            self.model = DDP(module=model, device_ids=[self._local_rank], output_device=self._local_rank)
+            self.raw_model = self.model.module
+        else:
+            self.model = model
+            self.raw_model = model
+
+        return self.model, optimizer

llm_trainer/parallel_ds.py

@@ -0,0 +1,45 @@
+from typing import Optional, Tuple
+import torch
+from torch import nn
+from .parallel import Parallel
+
+try:
+    import deepspeed
+except: ...
+
+class DsParallel(Parallel):
+    def __init__(self):
+        deepspeed.init_distributed(dist_backend='nccl')
+        super().__init__(init_process_group=False)
+
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
+        """
+        :param model:
+        :param optimizer:
+        :param kwargs:
+            参考deepspeed配置
+        :return:
+        """
+        self.raw_model = model
+
+        model, optim, _, _ = deepspeed.initialize(
+            model=model,
+            optimizer=optimizer,
+            dist_init_required=False,
+            config_params=kwargs
+        )
+
+        self.model = model
+        return model, optim
+
+    def synchronize(self): ...
+
+    def destroy(self): ...
+
+
+

llm_trainer/parallel_fsdp.py

@@ -0,0 +1,115 @@
+from typing import Optional, Tuple
+import functools
+import torch
+from torch import nn
+from torch.distributed.fsdp import (
+    FullyShardedDataParallel as FSDP,
+    MixedPrecision,
+    ShardingStrategy,
+    BackwardPrefetch,
+    CPUOffload,
+)
+
+from torch.distributed.fsdp.wrap import (
+    size_based_auto_wrap_policy,
+    transformer_auto_wrap_policy,
+    always_wrap_policy,
+    enable_wrap,
+    wrap,
+)
+
+from .parallel import Parallel
+
+class FsdpParallel(Parallel):
+    def __init__(self):
+        super().__init__()
+
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
+        """
+        :param model:
+        :param optimizer:
+        :param kwargs:
+            "wrap_policy_num_params" int size_based_auto_wrap_policy的最小参数量
+            "cpu_offload" bool 是否使用cpu卸载
+            "offload_params" bool 是否卸载参数，在cpu_offload为True时生效
+        :return:
+        """
+
+        model.to(self.device)
+
+        if self._use_compile:
+            model = torch.compile(model)
+
+        if self._use_parallel:
+            if 'transformer_layer_cls' in kwargs:
+                auto_wrap_policy = functools.partial(
+                    transformer_auto_wrap_policy,
+                    transformer_layer_cls=kwargs['transformer_layer_cls']
+                )
+            elif 'wrap_policy_num_params' in kwargs:
+                auto_wrap_policy = functools.partial(
+                    size_based_auto_wrap_policy,
+                    min_num_params=kwargs['wrap_policy_num_params']
+                )
+            else:
+                auto_wrap_policy = None
+
+            if 'cpu_offload' in kwargs:
+                offload_params = False
+                if 'offload_params' in kwargs:
+                    offload_params = kwargs['offload_params']
+
+                # 选择配置 cpu_offload，以便在计算中不使用包装参数时将这些参数卸载到 CPU。
+                # 这可以进一步提高内存效率，但代价是主机和设备之间的数据传输开销。
+                cpu_offload = CPUOffload(offload_params=offload_params)
+            else:
+                cpu_offload = None
+
+            if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+                mixed_precision = MixedPrecision(
+                    param_dtype=torch.bfloat16,
+                    # Gradient communication precision.
+                    reduce_dtype=torch.bfloat16,
+                    # Buffer precision.
+                    buffer_dtype=torch.bfloat16,
+                )
+            else:
+                mixed_precision = None
+
+            self.raw_model = model
+
+            # device_mesh = init_device_mesh("cuda", (self.world_size,))
+            # self.model = FSDP(
+            #     model,
+            #     auto_wrap_policy=auto_wrap_policy,
+            #     mixed_precision=mixed_precision,
+            #     cpu_offload=cpu_offload,
+            #     device_id=torch.cuda.current_device(),
+            #     device_mesh=device_mesh
+            # )
+
+            self.model = FSDP(
+                model,
+                sharding_strategy=ShardingStrategy.FULL_SHARD,
+                auto_wrap_policy=auto_wrap_policy,
+                mixed_precision=mixed_precision,
+                cpu_offload=cpu_offload,
+                device_id=torch.cuda.current_device(),
+                process_group=None,
+                # use_orig_params=True,
+                # backward_prefetch=BackwardPrefetch.BACKWARD_PRE,  # bit faster async comms, bit higher memory
+                # limit_all_gathers=False,
+                # forward_prefetch=True,
+            )
+        else:
+            self.model = model
+            self.raw_model = model
+
+        return self.model, optimizer
+
+

llm_trainer/parallel_none.py

@@ -0,0 +1,28 @@
+from typing import Optional, Tuple
+import torch
+from torch import nn
+
+from .parallel import Parallel
+
+class NoneParallel(Parallel):
+    def __init__(self):
+        super().__init__(use_parallel=False)
+
+    def process(
+            self,
+            model: nn.Module,
+            optimizer: torch.optim.Optimizer,
+            kwargs: Optional[dict] = None
+    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
+        model.to(self.device)
+
+        if self._use_compile:
+            model = torch.compile(model)
+
+        self.raw_model = model
+        self.model = model
+
+        return self.model, optimizer
+
+
+

llm_trainer/scheduler.py

@@ -0,0 +1,138 @@
+from abc import ABC, abstractmethod
+import math
+import torch
+from .log import (
+    log,
+    get_log_dir
+)
+
+class LRScheduler(ABC):
+    @property
+    @abstractmethod
+    def cur_steps(self): ...
+
+    @property
+    @abstractmethod
+    def cur_lr(self): ...
+
+    @abstractmethod
+    def update_steps(self, steps): ...
+
+    @abstractmethod
+    def step(self): ...
+
+    @abstractmethod
+    def can_clip_grad(self): ...
+
+
+class WarmupCosineAnnealingLRScheduler(LRScheduler):
+    def __init__(
+            self,
+            *,
+            optimizer: torch.optim.Optimizer,
+            initial_lr: float,
+            min_lr: float,
+            max_lr: float,
+            warmup_iters: int,
+            period: int, # 每个周期的步数
+            period_mul: int = 1, # 周期长度的倍数
+            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._period = period
+        self._period_mul = period_mul
+
+        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
+
+        self.need_log = need_log
+
+
+    @property
+    def cur_steps(self):
+        return self._steps
+
+    @property
+    def cur_lr(self):
+        return self._current_lr
+
+    def update_steps(self, steps):
+        log(f'update step to {steps}')
+        self._steps = steps
+        self._update_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._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 self._optimizer.param_groups:
+                param_group['lr'] = lr
+        else:
+            if not self._cosine_annealing_base_lr:
+                self._cosine_annealing_base_lr = self.cur_lr
+
+            """每步更新学习率"""
+            # 计算当前周期的最大步数
+            T_max = self._period * (self._period_mul ** self.cycle)
+
+            # 更新周期状态
+            self.T_cur += 1
+            if self.T_cur >= T_max:
+                self.cycle += 1
+                self.T_cur = 0  # 重置周期步数
+
+            # 计算并设置新学习率
+            cos_factor = (1 + math.cos(math.pi * self.T_cur / T_max)) / 2
+            lr = self._min_lr + (self._cosine_annealing_base_lr - self._min_lr) * cos_factor
+
+            for param_group in self._optimizer.param_groups:
+                param_group['lr'] = lr
+
+        self._current_lr = lr
+
+        if self.need_log:
+            log(f"step={self.cur_steps},lr={lr}\n", f'{get_log_dir()}lr.txt')
+
+
+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 update_steps(self, steps): ...
+
+    def step(self): ...
+
+    def can_clip_grad(self):
+        return True

llm_trainer/sft_trainer.py

@@ -0,0 +1,39 @@
+from typing import Optional, Tuple, List
+
+from torch.utils.data import Dataset
+
+from .trainer import Trainer
+from .train_configs import TrainConfig, VLMConfig
+from .dataset import LineByLineTextDataset
+from .utils import get_sft_collate_fn
+
+
+class SFTTrainer(Trainer):
+    def __init__(
+            self,
+            *,
+            train_config: TrainConfig,
+            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
+        )
+
+    def _convert_train_args(self) -> Tuple[dict, dict, dict, bool]:
+        sft_collate_fn = get_sft_collate_fn(self.train_config.mask_prompt)
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim = super()._convert_train_args()
+        data_loader_kwargs.update({"collate_fn": sft_collate_fn})
+
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
+
+    def _create_dataset(self, file_path) -> Dataset:
+        max_position_embeddings = self.train_config.model_config.max_position_embeddings
+        if isinstance(self.train_config.model_config, VLMConfig):
+            tokens_per_image = self.train_config.model_config.tokens_per_image
+        else:
+            tokens_per_image = -1
+
+        return LineByLineTextDataset(file_path, max_position_embeddings, tokens_per_image)