project-llm-trainer 0.4.1__py3-none-any.whl → 0.4.3__py3-none-any.whl

llm_trainer/generate_utils.py

@@ -1,6 +1,7 @@
 from typing import Union, Optional, List
 from contextlib import nullcontext
 import torch
+import torch.distributed as dist
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from llm_model import VlmModel, KVCache
 from .tools import TrainerTools

llm_trainer/grpo_trainer.py

@@ -14,6 +14,7 @@ from .dataset import GRPORolloutDataset
 from .loss import GRPOLoss
 from .tools import TrainerTools
 from .generate_utils import batch_generate
+from .log import log
 
 from .checkpoint import (
     save_checkpoint,
@@ -46,12 +47,9 @@ class GRPOTrainer(Trainer):
 
     def _init_reference_model(self):
         reference_model = self._new_model(self.train_config)
-
-        device = 'cpu' # TrainerTools().parallel.device
-        reference_model.to(device)
-        # load_checkpoint_for_eval(model=reference_model, device=device)
-
+        reference_model.to('cpu')
         reference_model.eval()
+
         for param in reference_model.parameters():
             param.requires_grad = False
 
@@ -59,17 +57,6 @@ class GRPOTrainer(Trainer):
 
     def _init_generate_model(self):
         return copy.deepcopy(self.reference_model)
-        # generate_model = self._new_model(self.train_config)
-        #
-        # device = 'cpu' #TrainerTools().parallel.device
-        # generate_model.to(device)
-        # # load_checkpoint_for_eval(model=generate_model, device=device)
-        #
-        # generate_model.eval()
-        # for param in generate_model.parameters():
-        #     param.requires_grad = False
-        #
-        # return generate_model
 
     def _init_loss(self):
         criterion = GRPOLoss(
@@ -194,7 +181,6 @@ class GRPOTrainer(Trainer):
 
         # [batch*group_size, max_prompt_len+max_gen_len]
         outputs: torch.Tensor = batch_generate(
-            # model=self.train_model,
             model=self.generate_model,
             tokens=prompt_ids,
             pad_token_id=pad_token_id,
@@ -325,10 +311,14 @@ class GRPOTrainer(Trainer):
                     self.generate_model.to(device)
                     self.reference_model.to(device)
 
-                    # 保存了train_model checkpoint后，这里保证生成模型使用的参数是最新
-                    copy_model_params(_from=self.train_model, _to=self.generate_model)
+                    if TrainerTools().parallel.is_main_process:
+                        log(f'start generate for batch {batch}/{batch_count_per_file}')
+
                     # 生成数据
-                    rollout_data = self._generate_rollout_data(batch_data)
+                    with torch.no_grad():
+                        # 保存了train_model checkpoint后，这里保证生成模型使用的参数是最新
+                        copy_model_params(_from=self.train_model, _to=self.generate_model)
+                        rollout_data = self._generate_rollout_data(batch_data)
 
                     # 卸载到cpu上，等待下次使用时再to gpu
                     self.generate_model.to('cpu')
@@ -337,6 +327,9 @@ class GRPOTrainer(Trainer):
                     # end generate
 
                     try:
+                        if TrainerTools().parallel.is_main_process:
+                            log(f'start train for batch {batch}/{batch_count_per_file}')
+
                         for grpo_step in range(self.train_config.grpo_config.grpo_steps):
                             with self.ctx:
                                 loss, aux_loss = self._maximize_grpo_objective(rollout_data)

llm_trainer/trainer.py

@@ -136,7 +136,7 @@ class Trainer:
         # 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']):
+                if not any(sub_module in name for sub_module in ['multi_modal_projector']):
                     param.requires_grad = False
 
             model.embed_tokens.eval()

{project_llm_trainer-0.4.1.dist-info → project_llm_trainer-0.4.3.dist-info}/METADATA

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

{project_llm_trainer-0.4.1.dist-info → project_llm_trainer-0.4.3.dist-info}/RECORD

@@ -6,8 +6,8 @@ llm_trainer/dpo_trainer.py,sha256=rC_I5ipesSlP3gFK_SG2GB8NbgJAMu4K7KLxkAS-aRY,13
 llm_trainer/ds_checkpoint.py,sha256=nchGocJE2oJnQ_KNN1kw-BkOAEIyTtO8SJt41cuN_xM,4232
 llm_trainer/eval.py,sha256=NDm8PbXLch7xT81xPYPRCNrcrB_Xj5GDJSCxyVwUOp4,1524
 llm_trainer/fsdp_checkpoint.py,sha256=lqZFzHyWyfzuCq_81kQNtJd2qaiMeY1N5BCEMnrJTBw,3192
-llm_trainer/generate_utils.py,sha256=4iM0vyc_1C_iTL31GlS9PR4eZtYaELPRZ02KDSPZA9U,15158
-llm_trainer/grpo_trainer.py,sha256=fqLT48ORSCece_e8dpyt8J7EarDuTnGoJ_eHk7Oy-1k,16177
+llm_trainer/generate_utils.py,sha256=RpAIjN0fvyTkMk9b9x7YE6c5GiiE3x5YGyPaa4R_BjA,15191
+llm_trainer/grpo_trainer.py,sha256=bZPrxhyPQLAnFzWhI7hhA6fpuKVNwj7nOm9k0ku9aK4,15977
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
 llm_trainer/parallel.py,sha256=DQu8GqEFxD99HQ6hKuIxxyKi-05dMO33eMhImYlPuOI,4468
@@ -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=arnet3tIzgVnwshod08F1jE7r4I7e-SIgMy55IagPnE,15971
-llm_trainer/trainer.py,sha256=hOn-z8kOd67RTuaaNMmdQjlw7N5LIZRHjSt5frpA1xI,25355
+llm_trainer/trainer.py,sha256=aoZYL5U4Z5axXBMM_DHgzIzJ89YbU9xUQ56jppcT65c,25339
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.4.1.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.4.1.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.4.1.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.4.1.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.4.1.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.4.1.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.4.1.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.4.1.dist-info/METADATA,sha256=9z1AB745r7BzQHNc3j-3N2nOdB9ZRUYsxcM42QoSb1o,195
-project_llm_trainer-0.4.1.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.4.1.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.4.1.dist-info/RECORD,,
+project_llm_trainer-0.4.3.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.3.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.3.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.3.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.3.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.3.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.3.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.3.dist-info/METADATA,sha256=kmmc6L6SE9iBvNHutWpeb0TocGX5vixhvHHLS4ltqec,195
+project_llm_trainer-0.4.3.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.3.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.3.dist-info/RECORD,,