project-llm-trainer 0.3.5__py3-none-any.whl → 0.4__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -3,6 +3,7 @@ from typing import Optional, Union, Tuple
 import torch
 from torch import nn
 from torch.optim import Optimizer
+from torch.nn.parallel import DistributedDataParallel as DDP
 
 from .parallel_ds import DsParallel
 from .parallel_fsdp import FsdpParallel
@@ -44,39 +45,22 @@ def save_checkpoint(
         save_ds_checkpoint(model, suffix)
     elif _can_use_dcp(model):
         save_dcp(model, optimizer, suffix)
+    elif isinstance(model, FSDP):
+        from .fsdp_checkpoint import save_fsdp_checkpoint
+        save_fsdp_checkpoint(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}"
+        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()}
+            raw_model = model if not isinstance(model, DDP) else model.module
+            ckpt = {'model_state_dict': raw_model.state_dict()}
 
-                    if optimizer:
-                        ckpt.update({'optim_state_dict': optimizer.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)
+            torch.save(ckpt, checkpoint_name)
 
 
 def load_checkpoint(
@@ -91,26 +75,20 @@ def load_checkpoint(
         load_ds_checkpoint(model, load_module_only=load_module_only, suffix=suffix)
     elif _can_use_dcp(model):
         load_dcp(model, optimizer, suffix)
+    elif isinstance(model, FSDP):
+        from .fsdp_checkpoint import load_fsdp_checkpoint
+        load_fsdp_checkpoint(model, optimizer, device, 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'])
+        state_dict = torch.load(checkpoint_name, weights_only=True, map_location=device)
+        raw_model = model.module if isinstance(model, DDP) else model
+        raw_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'])
+        if optimizer:
+            optimizer.load_state_dict(state_dict['optim_state_dict'])
 
 
 def load_checkpoint_for_eval(
@@ -141,6 +119,25 @@ def load_checkpoint_for_eval(
         load_checkpoint(model, None, device, suffix=suffix)
 
 
+def copy_model_params(
+        _from: nn.Module,
+        _to: nn.Module
+):
+    if isinstance(TrainerTools().parallel, DsParallel):
+        from .ds_checkpoint import get_ds_model_params
+        state_dict = get_ds_model_params(_from)
+    elif isinstance(TrainerTools().parallel, FsdpParallel):
+        from .fsdp_checkpoint import get_fsdp_model_params
+        state_dict = get_fsdp_model_params(_from)
+    elif isinstance(_from, DDP):
+        state_dict = _from.module.state_dict()
+    else:
+        state_dict = _from.state_dict()
+
+    if state_dict:
+        _to.load_state_dict(state_dict)
+
+
 def save_steps(global_steps: int, lr_scheduler: Optional[LRScheduler] = None):
     # 暂时只保存主进程的
     if TrainerTools().parallel.is_main_process:

llm_trainer/dpo_trainer.py

@@ -16,7 +16,7 @@ from .utils import get_dpo_collate_fn
 
 from .checkpoint import (
     save_checkpoint,
-    load_checkpoint_for_eval,
+    copy_model_params,
     save_steps,
 )
 
@@ -37,23 +37,18 @@ class DPOTrainer(Trainer):
         self.reference_model = self._init_reference_model()
 
     def _init_reference_model(self):
-        parallel = TrainerTools().new_parallel()
-
         reference_model = self._new_model(self.train_config)
-        if self.train_config.init_state_dict:
-            reference_model.load_state_dict(self.train_config.init_state_dict, strict=False)
-            self.train_config.init_state_dict = None
-        else:
-            load_checkpoint_for_eval(model=reference_model, device=parallel.device)
+        copy_model_params(_from=self.train_model, _to=reference_model)
 
-        reference_model, _ = parallel.process(
+        reference_model, _ = TrainerTools().parallel.process(
             model=reference_model,
             optimizer=None,
-            kwargs=self._init_reference_args()
+            kwargs=self._init_reference_args(),
+            save_instance=False
         )
 
-        parallel.raw_model.eval()
-        for param in parallel.raw_model.parameters():
+        reference_model.eval()
+        for param in reference_model.parameters():
             param.requires_grad = False
 
         return reference_model
@@ -267,7 +262,7 @@ class DPOTrainer(Trainer):
                             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)
+                                torch.nn.utils.clip_grad_norm_(self._get_trainable_params(self.train_model), 1.0)
 
                             self._step()
 

llm_trainer/ds_checkpoint.py

@@ -2,8 +2,14 @@ import os
 from typing import Optional
 from glob import glob
 import shutil
+import torch
 from torch import nn
+import torch.distributed as dist
+
+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: ...
@@ -59,3 +65,49 @@ def load_ds_checkpoint_for_eval(model: nn.Module):
     ckpt_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
     state_dict = get_fp32_state_dict_from_zero_checkpoint(ckpt_dir)
     model.load_state_dict(state_dict)
+
+
+def get_ds_model_params(model: nn.Module):
+    """
+        从一个正在运行的 DeepSpeedEngine 中高效地提取完整的 FP32 state_dict，
+        兼容 ZeRO Stages 0, 1, 2, 3。
+        这个版本包含了对 ZeRO-3 中非分片参数的正确处理。
+    """
+
+    assert isinstance(model, DeepSpeedEngine)
+    zero_stage = model.zero_optimization_stage()
+    state_dict = None
+
+    if TrainerTools().parallel.is_main_process:
+        if zero_stage == 3:
+            # ZeRO-3: Rank 0 聚合参数来构建完整的 state_dict
+            state_dict = {}
+            for param in model.module.parameters():
+                # 关键检查：判断参数是否被 ZeRO-3 分片管理
+                if hasattr(param, 'ds_id'):
+                    # 这是被分片的参数，使用 GatheredParameters 聚合
+                    with deepspeed.zero.GatheredParameters(param, modifier_rank=0):
+                        # .clone() 创建一个独立副本, .to('cpu') 移动到CPU, .to(torch.float32) 确保类型
+                        state_dict[param.ds_name] = param.data.to(torch.float32).cpu().clone()
+                else:
+                    # 这是未被分片的参数 (e.g., tied weights, buffers), 直接从 Rank 0 复制
+                    state_dict[param.ds_name] = param.data.to(torch.float32).cpu().clone()
+        else:  # zero_stage in [0, 1, 2]
+            # 在这些 stage，rank 0 已经有完整的模型。
+            # 我们从 model_engine.module 获取原始模型状态。
+            state_dict = {k: v.cpu().clone() for k, v in model.module.state_dict().items()}
+
+    # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
+    # 我们需要将其广播给所有进程。
+    if TrainerTools().parallel.world_size > 1:
+        # 准备一个列表，rank 0 有数据，其他 rank 是占位符
+        object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
+
+        # 执行广播，这个操作是阻塞的，会同步所有进程
+        dist.broadcast_object_list(object_list, src=0)
+
+        # 所有进程从列表中获取广播后的 state_dict 副本
+        state_dict = object_list[0]
+
+    return state_dict
+

llm_trainer/eval.py

@@ -1,9 +1,7 @@
-import time
-
 import torch
 
 from .generate_utils import generate
-from .checkpoint import load_checkpoint_for_eval
+from .checkpoint import copy_model_params
 from .log import get_log_dir
 from .tools import TrainerTools
 from .train_configs import EvalConfig
@@ -21,27 +19,6 @@ def _eval_task(
 ):
     log_dir = get_log_dir()
 
-    # 当eval_model不是独立model时可以尝试这个
-    # if isinstance(eval_model, FSDP):
-    #     with FSDP.summon_full_params(module=eval_model, writeback=False, recurse=False):
-    #         gen = generate(
-    #             eval_model,
-    #             prompt=prompt,
-    #             max_position_embeddings=max_position_embeddings,
-    #             max_new_tokens=max_new_tokens,
-    #             # temperature=None,
-    #             # k=None,
-    #             # p=None,
-    #             device='cpu',
-    #             item_callback=lambda item: write_temp(item)
-    #         )
-
-    # ---------
-    try:
-        load_checkpoint_for_eval(eval_model, device=device)
-    except:
-        return
-
     gen_result = generate(
         eval_model,
         prompt=prompt,
@@ -60,6 +37,7 @@ def _eval_task(
 
 
 def submit_gen_task(
+        train_model: torch.nn.Module,
         eval_model: torch.nn.Module,
         eval_config: EvalConfig,
         tag,
@@ -68,8 +46,13 @@ def submit_gen_task(
         max_position_embeddings,
         tokens_per_image
 ):
-    # 等待1s，防止deepspeed模式下，找不到checkpoint问题
-    time.sleep(1)
+    try:
+        copy_model_params(_from=train_model, _to=eval_model)
+    except Exception as e:
+        if isinstance(e, KeyboardInterrupt):
+            raise e
+        return
+
     eval_model.to(TrainerTools().parallel.device)
     _eval_task(
         eval_model=eval_model,

llm_trainer/fsdp_checkpoint.py

@@ -0,0 +1,91 @@
+import os
+from typing import Optional, Union, Tuple
+import torch
+from torch import nn
+from torch.optim import Optimizer
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+import torch.distributed as dist
+
+from .tools import TrainerTools
+
+DEFAULT_CHECKPOINT_NAME = "checkpoint.pth"
+
+def save_fsdp_checkpoint(
+        model: nn.Module,
+        optimizer: Optional[Optimizer] = None,
+        suffix: Optional[str] = None
+):
+    # 未经过测试 参考：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)
+
+
+def load_fsdp_checkpoint(
+        model: nn.Module,
+        optimizer: Optional[Optimizer] = None,
+        device: Optional[Union[torch.device, str]] = None,
+        suffix: Optional[str] = None
+):
+    checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
+    if suffix:
+        checkpoint_name = f"{checkpoint_name}_{suffix}"
+
+    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'])
+
+
+
+def get_fsdp_model_params(model: nn.Module):
+    """
+        从一个 FSDP 包装的模型中高效地提取完整的 FP32 state_dict。
+        这个函数会聚合所有分片的参数，并确保所有 rank 都收到一个完整的副本。
+    """
+
+    # FSDP 要求在所有 rank 上都调用 summon_full_params，即使我们只在 rank 0 上操作。
+    # writeback=False: 我们只读取参数，不写回，可以节省开销。
+    # offload_to_cpu=True: 直接将聚合后的参数卸载到 CPU，避免在 GPU 上产生大的峰值内存，
+    #                      并为我们省去了 .cpu() 的步骤。这是一个非常有用的优化。
+    # rank0_only=False: 为了让 offload_to_cpu 在所有 rank 上都生效，这里通常设为 False。
+    #                   我们稍后通过 get_rank() 来确保只有 rank 0 实际构建字典。
+    with FSDP.summon_full_params(model, writeback=False, offload_to_cpu=True):
+
+        state_dict = None
+        if TrainerTools().parallel.is_main_process:
+            # 在这个 with 块内部, model.state_dict() 会返回一个在 CPU 上的、完整的状态字典。
+            # 因为我们设置了 offload_to_cpu=True。
+            # 我们使用 .clone() 来确保我们得到的是一个独立的副本，
+            # 尽管 offload_to_cpu 已经帮我们处理了大部分情况。
+            state_dict = {k: v.clone() for k, v in model.state_dict().items()}
+
+    # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
+    # 我们需要将其广播给所有进程。
+    if TrainerTools().parallel.world_size > 1:
+        # 准备一个列表，rank 0 有数据，其他 rank 是占位符
+        object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
+
+        # 执行广播，这个操作是阻塞的，会同步所有进程
+        dist.broadcast_object_list(object_list, src=0)
+
+        # 所有进程从列表中获取广播后的 state_dict 副本
+        state_dict = object_list[0]
+
+    return state_dict

llm_trainer/grpo_trainer.py

@@ -17,7 +17,7 @@ from .generate_utils import batch_generate
 
 from .checkpoint import (
     save_checkpoint,
-    load_checkpoint_for_eval,
+    copy_model_params,
     save_steps,
 )
 
@@ -44,9 +44,6 @@ class GRPOTrainer(Trainer):
         # 如果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 = self._new_model(self.train_config)
 
@@ -296,7 +293,7 @@ class GRPOTrainer(Trainer):
         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)
+            copy_model_params(_from=self.train_model, _to=self.reference_model)
             self.train_model.train()
             file_count = len(self.train_config.file_dataset)
 
@@ -325,11 +322,11 @@ class GRPOTrainer(Trainer):
 
                     # start generate
                     # 使用单独的模型生成数据， 原因是在deepspeed并行训练时，使用train_model生成数据会卡死
-                    self.generate_model.to(TrainerTools().parallel.device)
-                    self.reference_model.to(TrainerTools().parallel.device)
+                    self.generate_model.to(device)
+                    self.reference_model.to(device)
 
                     # 保存了train_model checkpoint后，这里保证生成模型使用的参数是最新
-                    load_checkpoint_for_eval(self.generate_model, TrainerTools().parallel.device)
+                    copy_model_params(_from=self.train_model, _to=self.generate_model)
                     # 生成数据
                     rollout_data = self._generate_rollout_data(batch_data)
 
@@ -355,7 +352,7 @@ class GRPOTrainer(Trainer):
                             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)
+                                torch.nn.utils.clip_grad_norm_(self._get_trainable_params(self.train_model), 1.0)
 
                             self._step()
 

llm_trainer/parallel.py

@@ -64,7 +64,8 @@ class Parallel(ABC):
             self,
             model: nn.Module,
             optimizer: torch.optim.Optimizer,
-            kwargs: Optional[dict] = None
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
     ) -> Tuple[nn.Module, torch.optim.Optimizer]: ...
 
     def process_dataloader(

llm_trainer/parallel_ddp.py

@@ -21,7 +21,8 @@ class DdpParallel(Parallel):
             self,
             model: nn.Module,
             optimizer: torch.optim.Optimizer,
-            kwargs: Optional[dict] = None
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
     ) -> Tuple[nn.Module, torch.optim.Optimizer]:
         model.to(self.device)
 
@@ -30,10 +31,14 @@ class DdpParallel(Parallel):
 
         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
+            model = DDP(module=model, device_ids=[self._local_rank], output_device=self._local_rank)
+            raw_model = model.module
         else:
+            model = model
+            raw_model = model
+
+        if save_instance:
             self.model = model
-            self.raw_model = model
+            self.raw_model = raw_model
 
-        return self.model, optimizer
+        return model, optimizer

llm_trainer/parallel_ds.py

@@ -16,16 +16,20 @@ class DsParallel(Parallel):
             self,
             model: nn.Module,
             optimizer: torch.optim.Optimizer,
-            kwargs: Optional[dict] = None
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
     ) -> Tuple[nn.Module, torch.optim.Optimizer]:
         """
-        :param model:
-        :param optimizer:
-        :param kwargs:
-            参考deepspeed配置
-        :return:
+            :param model:
+            :param optimizer:
+            :param kwargs:
+                参考deepspeed配置
+            :param save_instance
+            :return:
         """
-        self.raw_model = model
+
+        if save_instance:
+            self.raw_model = model
 
         model, optim, _, _ = deepspeed.initialize(
             model=model,
@@ -34,7 +38,9 @@ class DsParallel(Parallel):
             config_params=kwargs
         )
 
-        self.model = model
+        if save_instance:
+            self.model = model
+
         return model, optim
 
     def synchronize(self): ...

llm_trainer/parallel_fsdp.py

@@ -28,16 +28,18 @@ class FsdpParallel(Parallel):
             self,
             model: nn.Module,
             optimizer: torch.optim.Optimizer,
-            kwargs: Optional[dict] = None
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
     ) -> 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:
+            :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时生效
+            :param save_instance
+            :return:
         """
 
         model.to(self.device)
@@ -81,10 +83,10 @@ class FsdpParallel(Parallel):
             else:
                 mixed_precision = None
 
-            self.raw_model = model
+            raw_model = model
 
             # device_mesh = init_device_mesh("cuda", (self.world_size,))
-            # self.model = FSDP(
+            # model = FSDP(
             #     model,
             #     auto_wrap_policy=auto_wrap_policy,
             #     mixed_precision=mixed_precision,
@@ -93,7 +95,7 @@ class FsdpParallel(Parallel):
             #     device_mesh=device_mesh
             # )
 
-            self.model = FSDP(
+            model = FSDP(
                 model,
                 sharding_strategy=ShardingStrategy.FULL_SHARD,
                 auto_wrap_policy=auto_wrap_policy,
@@ -107,9 +109,13 @@ class FsdpParallel(Parallel):
                 # forward_prefetch=True,
             )
         else:
+            model = model
+            raw_model = model
+
+        if save_instance:
+            self.raw_model = raw_model
             self.model = model
-            self.raw_model = model
 
-        return self.model, optimizer
+        return model, optimizer
 
 

llm_trainer/parallel_none.py

@@ -12,17 +12,19 @@ class NoneParallel(Parallel):
             self,
             model: nn.Module,
             optimizer: torch.optim.Optimizer,
-            kwargs: Optional[dict] = None
+            kwargs: Optional[dict] = None,
+            save_instance: bool = True
     ) -> 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
+        if save_instance:
+            self.raw_model = model
+            self.model = model
 
-        return self.model, optimizer
+        return model, optimizer
 
 
 

llm_trainer/tools.py

@@ -28,7 +28,7 @@ class TrainerTools:
         if not hasattr(TrainerTools, "_first_init"):
             TrainerTools._first_init = True
 
-            self.parallel = self.new_parallel()
+            self.parallel = self._new_parallel()
 
             self.tokenizer = Tokenizer(os.environ.get('TOKENIZERS_TYPE', 'zh_llama'))
             self.use_amp = 'cuda' in self.parallel.device and not isinstance(self.parallel, DsParallel)
@@ -43,7 +43,7 @@ class TrainerTools:
                 f' use_amp={self.use_amp},'
                 f' dtype={self.dtype}')
 
-    def new_parallel(self):
+    def _new_parallel(self):
         parallel_type = os.environ.get('PARALLEL_TYPE', 'none')
         log(f'parallel_type={parallel_type}')
         return parallel_types[parallel_type]()

llm_trainer/train_configs.py

@@ -422,7 +422,8 @@ class TrainConfig:
             kd_config: Optional[KDConfig] = None,
             pixel_values_provider: Optional[Callable[[list[str]], torch.Tensor]] = None,
             init_state_dict: Optional[Mapping[str, Any]] = None,
-            eval_config: EvalConfig = EvalConfig()
+            eval_config: EvalConfig = EvalConfig(),
+            freeze_llm_model: bool = False
     ):
         self.n_epochs = n_epochs
         self.batch_size = batch_size
@@ -443,5 +444,6 @@ class TrainConfig:
         self.pixel_values_provider = pixel_values_provider
         self.init_state_dict = init_state_dict
         self.eval_config = eval_config
+        self.freeze_llm_model = freeze_llm_model
 
 

llm_trainer/trainer.py

@@ -116,6 +116,10 @@ class Trainer:
         else:
             return LlmModel(train_config.model_config)
 
+    def _get_trainable_params(self, model):
+        freeze_llm_model = self.train_config.freeze_llm_model
+        return model.parameters() if not freeze_llm_model else filter(lambda p: p.requires_grad, model.parameters())
+
     def _init_train_model_and_optim(
             self,
             initial_lr: float,
@@ -128,10 +132,24 @@ class Trainer:
             model.load_state_dict(self.train_config.init_state_dict, strict=False)
             self.train_config.init_state_dict = None
 
+        # freeze llm model for vlm training
+        if self.train_config.freeze_llm_model:
+            for name, param in model.named_parameters():
+                if not any(sub_module in name for sub_module in ['vision_tower', 'multi_modal_projector']):
+                    param.requires_grad = False
+
+            model.embed_tokens.eval()
+            model.layers.eval()
+            model.head_norm.eval()
+            model.lm_head.eval()
+
         if TrainerTools().parallel.is_main_process:
             total_params = sum(p.numel() for p in model.parameters())
             log(f"Total number of parameters: {total_params:,}")
 
+            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+            log(f"Trainable number of parameters: {trainable_params:,}")
+
             total_size_bytes = total_params * 4
             total_size_mb = total_size_bytes / (1024 * 1024)
             log(f"Total size of the model: {total_size_mb:.2f} MB")
@@ -139,13 +157,13 @@ class Trainer:
         if use_ds_optim:
             import deepspeed
             origin_optim = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                model.parameters(),
+                self._get_trainable_params(model),
                 lr=initial_lr,
                 weight_decay=self.train_config.lr_config.weight_decay
             )
         else:
             origin_optim = torch.optim.AdamW(
-                model.parameters(),
+                self._get_trainable_params(model),
                 lr=initial_lr,
                 weight_decay=self.train_config.lr_config.weight_decay
             )
@@ -406,6 +424,7 @@ class Trainer:
                 eval_pixel_values = None
 
             submit_gen_task(
+                self.train_model,
                 self.eval_model,
                 self.train_config.eval_config,
                 tag=f'sign:batch/{tag}',
@@ -428,6 +447,7 @@ class Trainer:
                 eval_pixel_values = None
 
             submit_gen_task(
+                self.train_model,
                 self.eval_model,
                 self.train_config.eval_config,
                 tag=f'sign:epoch/{tag}',
@@ -529,7 +549,7 @@ class Trainer:
                             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)
+                                torch.nn.utils.clip_grad_norm_(self._get_trainable_params(self.train_model), 1.0)
 
                             self._step()
 

{project_llm_trainer-0.3.5.dist-info → project_llm_trainer-0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: project_llm_trainer
-Version: 0.3.5
+Version: 0.4
 Summary: LLM and VLM trainer
 Author: qibin
 Author-email: qibin0506@gmail.com

project_llm_trainer-0.4.dist-info/RECORD

@@ -0,0 +1,35 @@
+llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
+llm_trainer/checkpoint.py,sha256=GPaSGvnLCGMgsIA_vfjuw34tTQY26EuNwu7c08fhJHQ,5638
+llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
+llm_trainer/dcp.py,sha256=PkD97DyrOtoTKn4FJsfL3VqAy4dxufgjdzJEz8-Cnoc,3635
+llm_trainer/dpo_trainer.py,sha256=rC_I5ipesSlP3gFK_SG2GB8NbgJAMu4K7KLxkAS-aRY,13406
+llm_trainer/ds_checkpoint.py,sha256=H0BxYQixOWKRC20t55cFqNDTPzalD3AGTVt-owIB0_4,4488
+llm_trainer/eval.py,sha256=CsB3TpSVwhYVS9SP4Kuj_JhFUUvLcZUkvd8hvEIkPDU,1782
+llm_trainer/fsdp_checkpoint.py,sha256=xPQnAfXbx1SRKcVDLLgOtVrqjk0CjIRleVY0ZrwOAJU,3876
+llm_trainer/generate_utils.py,sha256=4iM0vyc_1C_iTL31GlS9PR4eZtYaELPRZ02KDSPZA9U,15158
+llm_trainer/grpo_trainer.py,sha256=fqLT48ORSCece_e8dpyt8J7EarDuTnGoJ_eHk7Oy-1k,16177
+llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
+llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
+llm_trainer/parallel.py,sha256=DQu8GqEFxD99HQ6hKuIxxyKi-05dMO33eMhImYlPuOI,4468
+llm_trainer/parallel_ddp.py,sha256=Pob9vUlBZnkL4oP1Re11kFob7nufMSE96pn7m7fuOEM,1345
+llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,1145
+llm_trainer/parallel_fsdp.py,sha256=cQOdY8ou6m8OsR06PpFVn6GiyZlK9nefkcGyszUOIJk,4055
+llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,676
+llm_trainer/scheduler.py,sha256=Xz8HhwoRMjRe41sf_NHhpZfkTlEs0I2MYusvMY6hCVw,3531
+llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,1783
+llm_trainer/tokenizer.py,sha256=A7TYYUbtPf75kjCvWP7yBui4xZBObMk2aPem62YpwpY,6776
+llm_trainer/tools.py,sha256=O45-20wRmh-nyTfU-U-XtjbKAoe7boEIsUvWT_NaKx4,3041
+llm_trainer/train_configs.py,sha256=arnet3tIzgVnwshod08F1jE7r4I7e-SIgMy55IagPnE,15971
+llm_trainer/trainer.py,sha256=DujZR1KOHyP3EHR8uIQPEsnX_5b7YC9Cto_eH7zxWqc,25256
+llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
+project_llm_trainer-0.4.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.dist-info/METADATA,sha256=-xxg-UyXn5MhW5OdYGFUcL5DtbIkgnQoUZS5b5bcEio,193
+project_llm_trainer-0.4.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.dist-info/RECORD,,

project_llm_trainer-0.3.5.dist-info/RECORD

@@ -1,34 +0,0 @@
-llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=Dlkcit0o7Gx6S9QUrIrVp2pTurP9X0zVA7w7ImSuVQU,6049
-llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
-llm_trainer/dcp.py,sha256=PkD97DyrOtoTKn4FJsfL3VqAy4dxufgjdzJEz8-Cnoc,3635
-llm_trainer/dpo_trainer.py,sha256=kDpzHrxP3qWbBxDGi9Rkus1kw8P3-bAtw9IyuYINgk0,13625
-llm_trainer/ds_checkpoint.py,sha256=_svpzqRaa43--DKPputoXAelc6X9vPM0gNQu-hlh6NI,2153
-llm_trainer/eval.py,sha256=sCvdYnqWWf5_nuDQN5BHb_YivXLOQW-V0ET9mPu0tPU,2389
-llm_trainer/generate_utils.py,sha256=4iM0vyc_1C_iTL31GlS9PR4eZtYaELPRZ02KDSPZA9U,15158
-llm_trainer/grpo_trainer.py,sha256=M6vp6QjxhBQVaw3e_3BJ4earuezQNKQ3JeZfQLBaSLQ,16370
-llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
-llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
-llm_trainer/parallel.py,sha256=2VJtW3Gq2c1yS_LdcrNhk7B12prFwBmFnKhvV8FS2d8,4428
-llm_trainer/parallel_ddp.py,sha256=Gz-3LZ6LKmqlNwxrnGRC4uKoqoSxCvp9JHejIBSQp3c,1238
-llm_trainer/parallel_ds.py,sha256=W_PkczyAlgffCRcQadN-Pf7H7HM7TU26v5W63jKELFM,990
-llm_trainer/parallel_fsdp.py,sha256=u9XbbVTzcsMcaf-aQFrC_QwWsDRGoEpRmgvu1cKNtgk,3887
-llm_trainer/parallel_none.py,sha256=a6tt3aBmCq5rSP7n2I-sF-hsZ992BbLbpbxutDCFJfs,607
-llm_trainer/scheduler.py,sha256=Xz8HhwoRMjRe41sf_NHhpZfkTlEs0I2MYusvMY6hCVw,3531
-llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,1783
-llm_trainer/tokenizer.py,sha256=A7TYYUbtPf75kjCvWP7yBui4xZBObMk2aPem62YpwpY,6776
-llm_trainer/tools.py,sha256=AhfjN9oln5Pyif1SgCWwgQg-Q5acTCd9xpz4L26QUjA,3039
-llm_trainer/train_configs.py,sha256=cadfo8RgxNUR-L3ZLyjiRXTQvhjUl4A1qENaq-ol8h4,15878
-llm_trainer/trainer.py,sha256=5DgDzg0TReZrXsIaM6A4DzeJnzePNybGdfoVSDybQ2U,24308
-llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.3.5.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.3.5.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.3.5.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.3.5.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.3.5.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.3.5.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.3.5.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.3.5.dist-info/METADATA,sha256=jfnJI_XqE7U89-8tLEGPmLpuzwp-3qw-aERIgV8GJpk,195
-project_llm_trainer-0.3.5.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.3.5.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.3.5.dist-info/RECORD,,