project-llm-trainer 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

llm_trainer/dpo_trainer.py

@@ -5,8 +5,6 @@ from torch.utils.data import Dataset
 import torch.distributed as dist
 import torch.nn.functional as F
 
-from llm_model import LlmModel
-
 from .parallel_ds import DsParallel
 from .parallel_fsdp import FsdpParallel
 from .trainer import Trainer
@@ -41,7 +39,7 @@ class DPOTrainer(Trainer):
     def _init_reference_model(self):
         parallel = TrainerTools().new_parallel()
 
-        reference_model = LlmModel(self.train_config.model_config)
+        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

llm_trainer/grpo_trainer.py

@@ -7,8 +7,6 @@ from torch.nn.utils.rnn import pad_sequence
 import torch.distributed as dist
 import torch.nn.functional as F
 
-from llm_model import LlmModel
-
 from .parallel_ds import DsParallel
 from .trainer import Trainer
 from .train_configs import TrainConfig
@@ -50,7 +48,7 @@ class GRPOTrainer(Trainer):
         save_checkpoint(self.train_model, self.optimizer)
 
     def _init_reference_model(self):
-        reference_model = LlmModel(self.train_config.model_config)
+        reference_model = self._new_model(self.train_config)
 
         device = 'cpu' # TrainerTools().parallel.device
         reference_model.to(device)
@@ -64,7 +62,7 @@ class GRPOTrainer(Trainer):
 
     def _init_generate_model(self):
         return copy.deepcopy(self.reference_model)
-        # generate_model = LlmModel(self.train_config.model_config)
+        # generate_model = self._new_model(self.train_config)
         #
         # device = 'cpu' #TrainerTools().parallel.device
         # generate_model.to(device)

llm_trainer/sft_trainer.py

@@ -32,11 +32,15 @@ class SFTTrainer(Trainer):
     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
+
+        image_tag_file_path = None
+        tokens_per_image = -1
+
         if isinstance(self.train_config.model_config, VLMConfig):
-            image_tag_file_path = self.train_config.image_tags_file_dataset[file_idx]
-            tokens_per_image = self.train_config.model_config.tokens_per_image
-        else:
-            image_tag_file_path = None
-            tokens_per_image = -1
+            if self.train_config.image_tags_file_dataset:
+                image_tag_file_path = self.train_config.image_tags_file_dataset[file_idx]
+
+            if self.train_config.model_config.tokens_per_image:
+                tokens_per_image = self.train_config.model_config.tokens_per_image
 
         return LineByLineTextDataset(file_path, max_position_embeddings, image_tag_file_path, tokens_per_image), file_path

llm_trainer/trainer.py

@@ -1,6 +1,6 @@
 import time
 from contextlib import nullcontext
-from typing import Optional, Tuple, List, Dict, Any
+from typing import Optional, Tuple, List, Dict, Any, Union
 
 import torch
 from torch import nn
@@ -110,16 +110,19 @@ class Trainer:
             self.pixel_values_provider = None
             self.tokens_per_image = -1
 
+    def _new_model(self, train_config: TrainConfig):
+        if isinstance(train_config.model_config, VLMConfig):
+            return VlmModel(train_config.model_config)
+        else:
+            return LlmModel(train_config.model_config)
+
     def _init_train_model_and_optim(
             self,
             initial_lr: float,
             parallel_kwargs: dict,
             use_ds_optim: bool
     ):
-        if isinstance(self.train_config.model_config, VLMConfig):
-            model = VlmModel(self.train_config.model_config)
-        else:
-            model = LlmModel(self.train_config.model_config)
+        model = self._new_model(self.train_config)
 
         if self.train_config.init_state_dict:
             model.load_state_dict(self.train_config.init_state_dict, strict=False)
@@ -156,10 +159,7 @@ class Trainer:
 
     def _init_eval_model(self) -> Optional[nn.Module]:
         if TrainerTools().parallel.is_main_process:
-            if isinstance(self.train_config.model_config, VLMConfig):
-                return VlmModel(self.train_config.model_config).to('cpu')
-            else:
-                return LlmModel(self.train_config.model_config).to('cpu')
+            return self._new_model(self.train_config).to('cpu')
 
         return None
 
@@ -400,7 +400,7 @@ class Trainer:
     ):
         if TrainerTools().parallel.is_main_process:
             eval_prompt, eval_image_tag = self._get_eval_data()
-            if isinstance(self.train_config.model_config, VLMConfig) and eval_image_tag:
+            if isinstance(self.train_model, VlmModel) and self.pixel_values_provider and eval_image_tag:
                 eval_pixel_values = self.pixel_values_provider([eval_image_tag])
             else:
                 eval_pixel_values = None
@@ -422,7 +422,7 @@ class Trainer:
     ):
         if TrainerTools().parallel.is_main_process:
             eval_prompt, eval_image_tag = self._get_eval_data()
-            if isinstance(self.train_config.model_config, VLMConfig) and eval_image_tag:
+            if isinstance(self.train_model, VlmModel) and eval_image_tag:
                 eval_pixel_values = self.pixel_values_provider([eval_image_tag])
             else:
                 eval_pixel_values = None

{project_llm_trainer-0.3.2.dist-info → project_llm_trainer-0.3.3.dist-info}/METADATA

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

{project_llm_trainer-0.3.2.dist-info → project_llm_trainer-0.3.3.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=q3JZ1iKzmiuwUV-DTrSXUea2d39g6f5x1oUuF1QzBGA,13173
+llm_trainer/dpo_trainer.py,sha256=7Bf6snWcu2fT8QRDI1CSzmrc7Cog6JauIeK2KoW_f8I,13135
 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=_k9pik-kpbE8g9taQyG9w3dTLAHilgVBTUa4Y90Wae4,16414
+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
@@ -15,20 +15,20 @@ llm_trainer/parallel_ds.py,sha256=W_PkczyAlgffCRcQadN-Pf7H7HM7TU26v5W63jKELFM,99
 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=WWmg8YOwr-w90otmeMjXvK9sa_DSPKlfgAPg3kHyRF4,1672
+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=153F8FzsKh6k9XLm9i6JzmwN4Vwva5mWr9rVoge_3bY,24353
+llm_trainer/trainer.py,sha256=tUi0Xcwgci_Y4T_I3CHR6pcsqURsXIMKbpAXB4DgRWo,24277
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.3.2.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.3.2.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.3.2.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.3.2.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.3.2.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.3.2.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.3.2.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.3.2.dist-info/METADATA,sha256=NQpGh0Xy09euhzVTSBcC6m5P23ATvRKQ-zmkE0o__6g,195
-project_llm_trainer-0.3.2.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.3.2.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.3.2.dist-info/RECORD,,
+project_llm_trainer-0.3.3.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.3.3.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.3.3.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.3.3.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.3.3.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.3.3.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.3.3.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.3.3.dist-info/METADATA,sha256=plDfgI4_qj6tHaJvEhmhmRFmKt0JJpU37JNDjkDjnRY,195
+project_llm_trainer-0.3.3.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.3.3.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.3.3.dist-info/RECORD,,