project-llm-trainer 0.4.12__py3-none-any.whl → 0.4.14__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -129,10 +129,10 @@ def copy_model_params(
 
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import get_ds_model_params
-        state_dict = get_ds_model_params(_from)
+        state_dict = get_ds_model_params(_from, only_rank0=_to is None)
     elif isinstance(TrainerTools().parallel, FsdpParallel):
         from .fsdp_checkpoint import get_fsdp_model_params
-        state_dict = get_fsdp_model_params(_from)
+        state_dict = get_fsdp_model_params(_from, only_rank0=_to is None)
     elif isinstance(_from, DDP):
         state_dict = _from.module.state_dict()
     else:

llm_trainer/dpo_trainer.py

@@ -6,7 +6,6 @@ import torch.distributed as dist
 import torch.nn.functional as F
 
 from .parallel_ds import DsParallel
-from .parallel_fsdp import FsdpParallel
 from .trainer import Trainer
 from .train_configs import TrainConfig
 from .dataset import DPODataset
@@ -53,52 +52,6 @@ class DPOTrainer(Trainer):
 
         return reference_model
 
-    def _init_reference_args(self):
-        if isinstance(TrainerTools().parallel, DsParallel) and self.train_config.ds_config:
-            parallel_kwargs = {
-                'gradient_accumulation_steps': 1,
-                'train_micro_batch_size_per_gpu': 1
-            }
-
-            if self.train_config.ds_config.zero_config:
-                zero_optimization = {'stage': 0}
-                parallel_kwargs['zero_optimization'] = zero_optimization
-
-
-            if (self.train_config.ds_config.bf16_config is not None
-                    and self.train_config.ds_config.bf16_config.enabled):
-                bf16_config = self.train_config.ds_config.bf16_config
-                bf16 = {
-                    'enabled': bf16_config.enabled
-                }
-                parallel_kwargs['bf16'] = bf16
-            elif self.train_config.ds_config.fp16_config:
-                fb16_config = self.train_config.ds_config.fp16_config
-                fp16 = {
-                    'enabled': fb16_config.enabled,
-                    'loss_scale': fb16_config.loss_scale,
-                    'loss_scale_window': fb16_config.loss_scale_window,
-                    'initial_scale_power': fb16_config.initial_scale_power,
-                    'hysteresis': fb16_config.hysteresis,
-                    'min_loss_scale': fb16_config.min_loss_scale
-                }
-
-                if fb16_config.fp16_opt_level is not None:
-                    fp16['fp16_opt_level'] = fb16_config.fp16_opt_level
-
-                parallel_kwargs['fp16'] = fp16
-        elif isinstance(TrainerTools().parallel, FsdpParallel) and self.train_config.fsdp_config:
-            parallel_kwargs = {
-                'transformer_layer_cls': self.train_config.fsdp_config.transformer_layer_cls,
-                'wrap_policy_num_params': self.train_config.fsdp_config.wrap_policy_num_params,
-                'cpu_offload': self.train_config.fsdp_config.cpu_offload,
-                'offload_params': self.train_config.fsdp_config.offload_params
-            }
-        else:
-            parallel_kwargs = None
-
-        return parallel_kwargs
-
     def _init_loss(self):
         criterion = DPOLoss(
             beta=self.train_config.dpo_config.loss_beta,

llm_trainer/ds_checkpoint.py

@@ -105,7 +105,7 @@ def _get_ds_full_state_dict_on_rank0(model: DeepSpeedEngine) -> Optional[dict]:
     # return state_dict_on_rank_0 if TrainerTools().parallel.is_main_process else None
 
 
-def get_ds_model_params(model: nn.Module):
+def get_ds_model_params(model: nn.Module, only_rank0=False):
     """
         从一个正在运行的 DeepSpeedEngine 中高效地提取完整的 FP32 state_dict，
         兼容 ZeRO Stages 0, 1, 2, 3。
@@ -117,7 +117,7 @@ def get_ds_model_params(model: nn.Module):
 
     # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
     # 我们需要将其广播给所有进程。
-    if TrainerTools().parallel.world_size > 1:
+    if not only_rank0 and TrainerTools().parallel.world_size > 1:
         # 准备一个列表，rank 0 有数据，其他 rank 是占位符
         object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
         # 执行广播，这个操作是阻塞的，会同步所有进程

llm_trainer/fsdp_checkpoint.py

@@ -66,7 +66,7 @@ def _get_fsdp_full_state_dict_on_rank0(model: nn.Module) -> Optional[dict]:
     return None
 
 
-def get_fsdp_model_params(model: nn.Module):
+def get_fsdp_model_params(model: nn.Module, only_rank0=False):
     """
         从一个 FSDP 包装的模型中高效地提取完整的 FP32 state_dict。
         这个函数会聚合所有分片的参数，并确保所有 rank 都收到一个完整的副本。
@@ -76,7 +76,7 @@ def get_fsdp_model_params(model: nn.Module):
 
     # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
     # 我们需要将其广播给所有进程。
-    if TrainerTools().parallel.world_size > 1:
+    if not only_rank0 and TrainerTools().parallel.world_size > 1:
         # 准备一个列表，rank 0 有数据，其他 rank 是占位符
         object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
         # 执行广播，这个操作是阻塞的，会同步所有进程

llm_trainer/trainer.py

@@ -178,7 +178,7 @@ class Trainer:
 
     def _init_eval_model(self) -> Optional[nn.Module]:
         if TrainerTools().parallel.is_main_process:
-            return self._new_model(self.train_config).to('cpu')
+            return self._new_model(self.train_config).to(device='cpu', dtype=TrainerTools().dtype)
 
         return None
 
@@ -337,6 +337,34 @@ class Trainer:
 
         return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
 
+    def _init_reference_args(self) -> dict:
+        parallel_kwargs, _, _, _ = self._convert_train_args()
+
+        if parallel_kwargs and isinstance(TrainerTools().parallel, DsParallel):
+            # reference to https://github.com/huggingface/trl/blob/main/trl/models/utils.py:prepare_deepspeed
+            # if model is not None:
+            #     hidden_size = (
+            #         max(model.config.hidden_sizes)
+            #         if getattr(model.config, "hidden_sizes", None)
+            #         else getattr(model.config, "hidden_size", None)
+            #     )
+            #     if hidden_size is not None and stage == 3:
+            #         # Note that `stage3_prefetch_bucket_size` can produce DeepSpeed messages like: `Invalidate trace cache
+            #         # @ step 0: expected module 1, but got module 0`
+            #         # This is expected and is not an error, see: https://github.com/microsoft/DeepSpeed/discussions/4081
+            #         config_kwargs.update(
+            #             {
+            #                 "zero_optimization.reduce_bucket_size": hidden_size * hidden_size,
+            #                 "zero_optimization.stage3_param_persistence_threshold": 10 * hidden_size,
+            #                 "zero_optimization.stage3_prefetch_bucket_size": 0.9 * hidden_size * hidden_size,
+            #             }
+            #         )
+
+            if parallel_kwargs['zero_optimization']['stage'] != 3:
+                parallel_kwargs['zero_optimization']['stage'] = 0
+
+        return parallel_kwargs
+
     def _create_dataset(self, file_idx) -> Tuple[Dataset, str]:
         file_path = self.train_config.file_dataset[file_idx]
         max_position_embeddings = self.train_config.model_config.max_position_embeddings

{project_llm_trainer-0.4.12.dist-info → project_llm_trainer-0.4.14.dist-info}/METADATA

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

{project_llm_trainer-0.4.12.dist-info → project_llm_trainer-0.4.14.dist-info}/RECORD

@@ -1,11 +1,11 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=yZcExxneN2yzvWxRiK-pstMWs35LV7GiOfqcLq-S6vc,5745
+llm_trainer/checkpoint.py,sha256=ItDzuXVikk-0gWSw-IS7SrODEdlJEb5nZs15dBFkPdk,5793
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
 llm_trainer/dcp.py,sha256=PkD97DyrOtoTKn4FJsfL3VqAy4dxufgjdzJEz8-Cnoc,3635
-llm_trainer/dpo_trainer.py,sha256=34E2b-t0GZYutaw6bESgARe9C12PUMWcY4aGZ34eAZU,13576
-llm_trainer/ds_checkpoint.py,sha256=x_tjgJR47P8gVwV4qAnTUCGwx7eVq2Epw0vOVV7fkYo,4925
+llm_trainer/dpo_trainer.py,sha256=djBhvI_ixTV1nLNg84tgCpfV--pu6IRiOhO28V-aANQ,11425
+llm_trainer/ds_checkpoint.py,sha256=fprJlbSgtyKmmpytyMOZBs3pcjZA13SeWao0llnLpNQ,4962
 llm_trainer/eval.py,sha256=NDm8PbXLch7xT81xPYPRCNrcrB_Xj5GDJSCxyVwUOp4,1524
-llm_trainer/fsdp_checkpoint.py,sha256=lqZFzHyWyfzuCq_81kQNtJd2qaiMeY1N5BCEMnrJTBw,3192
+llm_trainer/fsdp_checkpoint.py,sha256=dAHIGHfuvTA6OC0jV9Ls-oD4ZR9CPGa31mjtoh-2dZE,3229
 llm_trainer/generate_utils.py,sha256=tSbA_tLqSq5qJGHSOlPv5T3iRDZkbFg5ZvDAgJ_i_SE,17946
 llm_trainer/grpo_trainer.py,sha256=bZPrxhyPQLAnFzWhI7hhA6fpuKVNwj7nOm9k0ku9aK4,15977
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
@@ -20,16 +20,16 @@ llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,17
 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=HKzH3nfMT1-SW4Htwa0KqYtMd6FAJcthR5IEo6di8us,8168
-llm_trainer/trainer.py,sha256=pUtJVRosn54j1hn76CFAptJcAsrDo59H6p8NMkg2zt4,25521
+llm_trainer/trainer.py,sha256=j5fDqMzvU6SYwxHsv9wX0UVX4JXS-8eP1AkHgVxKf9U,27119
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.4.12.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.4.12.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.4.12.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.4.12.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.4.12.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.4.12.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.4.12.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.4.12.dist-info/METADATA,sha256=W-HeRGlXi3bFsKIVE1FyQAh4Lcvo0SOXMNu-9YnACKQ,196
-project_llm_trainer-0.4.12.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.4.12.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.4.12.dist-info/RECORD,,
+project_llm_trainer-0.4.14.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.14.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.14.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.14.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.14.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.14.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.14.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.14.dist-info/METADATA,sha256=VMEWVv8pBqFUAhIAiH4_S4ECUHln31gchHLhTtUAM1o,196
+project_llm_trainer-0.4.14.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.14.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.14.dist-info/RECORD,,