project-llm-trainer 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

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/train_configs.py

@@ -67,7 +67,7 @@ class DsFp16Config:
     initial_scale_power: int = 16
     hysteresis: int = 2
     min_loss_scale: int = 1
-    fp16_opt_level: Optional[str] = '02'
+    fp16_opt_level: Optional[str] = 'O2'
 
 
 @dataclass(kw_only=True)
@@ -77,9 +77,9 @@ class DsBf16Config:
 
 @dataclass(kw_only=True)
 class DsConfig:
-    zero_config: Optional[DsZeROConfig] = DsZero3Config()
-    fp16_config: Optional[DsFp16Config] = DsFp16Config()
-    bf16_config: Optional[DsBf16Config] = DsBf16Config()
+    zero_config: Optional[DsZeROConfig] = field(default_factory=DsZero3Config)
+    fp16_config: Optional[DsFp16Config] = field(default_factory=DsFp16Config)
+    bf16_config: Optional[DsBf16Config] = field(default_factory=DsBf16Config)
     gradient_clipping: Optional[float] = 1.0
     activation_checkpointing: Optional[DsActivationCheckpointingConfig] = None
 
@@ -224,14 +224,14 @@ class TrainConfig:
     model_config: Union[ModelConfig, VLMConfig]
 
     file_dataset: FileDataset
-    data_loader_config: DataLoaderConfig = DataLoaderConfig()
+    data_loader_config: DataLoaderConfig = field(default_factory=DataLoaderConfig)
     image_tags_file_dataset: Optional[FileDataset] = None
 
-    loss_config: LossConfig = LossConfig()
-    lr_config: LrConfig = LrConfig()
+    loss_config: LossConfig = field(default_factory=LossConfig)
+    lr_config: LrConfig = field(default_factory=LrConfig)
 
-    ds_config: DsConfig = DsConfig()
-    fsdp_config: FsdpConfig = FsdpConfig()
+    ds_config: DsConfig = field(default_factory=DsConfig)
+    fsdp_config: FsdpConfig = field(default_factory=FsdpConfig)
 
     kd_config: Optional[KDConfig] = None
     dpo_config: Optional[DPOConfig] = None
@@ -241,7 +241,7 @@ class TrainConfig:
     gradient_accumulation_steps: int = 0
     eval_batch_interval: int = 100
 
-    eval_config: EvalConfig = EvalConfig()
+    eval_config: EvalConfig = field(default_factory=EvalConfig)
     pixel_values_provider: Optional[Callable[[list[str]], torch.Tensor]] = None
 
     init_state_dict: Optional[Mapping[str, Any]] = 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.11.dist-info → project_llm_trainer-0.4.13.dist-info}/METADATA

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

{project_llm_trainer-0.4.11.dist-info → project_llm_trainer-0.4.13.dist-info}/RECORD

@@ -2,7 +2,7 @@ llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
 llm_trainer/checkpoint.py,sha256=yZcExxneN2yzvWxRiK-pstMWs35LV7GiOfqcLq-S6vc,5745
 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/dpo_trainer.py,sha256=djBhvI_ixTV1nLNg84tgCpfV--pu6IRiOhO28V-aANQ,11425
 llm_trainer/ds_checkpoint.py,sha256=x_tjgJR47P8gVwV4qAnTUCGwx7eVq2Epw0vOVV7fkYo,4925
 llm_trainer/eval.py,sha256=NDm8PbXLch7xT81xPYPRCNrcrB_Xj5GDJSCxyVwUOp4,1524
 llm_trainer/fsdp_checkpoint.py,sha256=lqZFzHyWyfzuCq_81kQNtJd2qaiMeY1N5BCEMnrJTBw,3192
@@ -19,17 +19,17 @@ 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=4sM96SOgwcn6jBGtbG5-qDZbJjiHVB6l7FWqdq7hbj0,7979
-llm_trainer/trainer.py,sha256=pUtJVRosn54j1hn76CFAptJcAsrDo59H6p8NMkg2zt4,25521
+llm_trainer/train_configs.py,sha256=HKzH3nfMT1-SW4Htwa0KqYtMd6FAJcthR5IEo6di8us,8168
+llm_trainer/trainer.py,sha256=j5fDqMzvU6SYwxHsv9wX0UVX4JXS-8eP1AkHgVxKf9U,27119
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.4.11.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.4.11.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.4.11.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.4.11.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.4.11.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.4.11.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.4.11.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.4.11.dist-info/METADATA,sha256=JEZo2-np0t_K-J6yapyAXsArpvYTmrSNGDsdy32kWas,196
-project_llm_trainer-0.4.11.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.4.11.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.4.11.dist-info/RECORD,,
+project_llm_trainer-0.4.13.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.13.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.13.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.13.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.13.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.13.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.13.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.13.dist-info/METADATA,sha256=hiW-7qgWuPizKVz4cU8mEHoqiuT6ZqNlCBb7nwVfFQ4,196
+project_llm_trainer-0.4.13.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.13.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.13.dist-info/RECORD,,