project-llm-trainer 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

llm_trainer/dpo_trainer.py

@@ -69,19 +69,29 @@ class DPOTrainer(Trainer):
                 zero_optimization = {'stage': 0}
                 parallel_kwargs['zero_optimization'] = zero_optimization
 
-            if self.train_config.ds_config.fp16_config:
+
+            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 }
+                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
-
-            if self.train_config.ds_config.bf16_config:
-                bf16_config = self.train_config.ds_config.bf16_config
-                bf16 = { 'enabled': bf16_config.enabled }
-                parallel_kwargs['bf16'] = bf16
         elif isinstance(TrainerTools().parallel, FsdpParallel) and self.train_config.fsdp_config:
             parallel_kwargs = {
                 'transformer_layer_cls': self.train_config.fsdp_config.transformer_layer_cls,
@@ -257,7 +267,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/grpo_trainer.py

@@ -355,7 +355,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/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
             )
@@ -529,7 +547,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.4.dist-info → project_llm_trainer-0.3.6.dist-info}/METADATA

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

{project_llm_trainer-0.3.4.dist-info → project_llm_trainer-0.3.6.dist-info}/RECORD

@@ -2,11 +2,11 @@ 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=7Bf6snWcu2fT8QRDI1CSzmrc7Cog6JauIeK2KoW_f8I,13135
+llm_trainer/dpo_trainer.py,sha256=rEhoVN4gPweX5NYKZaEH7jgWav4w6OQ2x-QRocahYjg,13640
 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/grpo_trainer.py,sha256=1oH0argbpITlzAEkGKW8F9kZPr67bcb95FGOVpP8XTM,16385
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
 llm_trainer/parallel.py,sha256=2VJtW3Gq2c1yS_LdcrNhk7B12prFwBmFnKhvV8FS2d8,4428
@@ -18,17 +18,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=AhfjN9oln5Pyif1SgCWwgQg-Q5acTCd9xpz4L26QUjA,3039
-llm_trainer/train_configs.py,sha256=cadfo8RgxNUR-L3ZLyjiRXTQvhjUl4A1qENaq-ol8h4,15878
-llm_trainer/trainer.py,sha256=5DgDzg0TReZrXsIaM6A4DzeJnzePNybGdfoVSDybQ2U,24308
+llm_trainer/train_configs.py,sha256=arnet3tIzgVnwshod08F1jE7r4I7e-SIgMy55IagPnE,15971
+llm_trainer/trainer.py,sha256=2cO-MwWJsgPbTisOp_HVIdA0SVodFZx3M8lafarnLdw,25188
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.3.4.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.3.4.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.3.4.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.3.4.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.3.4.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.3.4.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.3.4.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.3.4.dist-info/METADATA,sha256=Y8XjOGdQb7VxN5QKHyKICkkOzjGcXJuI6hPziULJNfc,195
-project_llm_trainer-0.3.4.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.3.4.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.3.4.dist-info/RECORD,,
+project_llm_trainer-0.3.6.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.3.6.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.3.6.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.3.6.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.3.6.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.3.6.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.3.6.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.3.6.dist-info/METADATA,sha256=1ClKvVThd4g8uToJQevDXmjAI8gbVYzDfYImWXHFRqI,195
+project_llm_trainer-0.3.6.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.3.6.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.3.6.dist-info/RECORD,,