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

llm_trainer/ds_checkpoint.py

@@ -77,18 +77,33 @@ def _get_ds_full_state_dict_on_rank0(model: DeepSpeedEngine) -> Optional[dict]:
             return {k: v.cpu().clone() for k, v in model.module.state_dict().items()}
         return None
 
-    # ZeRO-3
-    state_dict_on_rank_0 = {}
-    for param_name, param in model.module.named_parameters():
-        if hasattr(param, 'ds_id'):
-            with deepspeed.zero.GatheredParameters(param, modifier_rank=0):
-                if TrainerTools().parallel.is_main_process:
-                    state_dict_on_rank_0[param_name] = param.data.to(torch.float32).cpu().clone()
-        else:
-            if TrainerTools().parallel.is_main_process:
-                state_dict_on_rank_0[param_name] = param.data.to(torch.float32).cpu().clone()
-
-    return state_dict_on_rank_0 if TrainerTools().parallel.is_main_process else None
+    # --- ZeRO-3 ---
+    # 只调用一次 GatheredParameters，传入所有参数
+    with deepspeed.zero.GatheredParameters(model.parameters(), modifier_rank=0):
+        if TrainerTools().parallel.is_main_process:
+            # 在这个 'with' 代码块内，rank 0 上的 model.module 拥有完整的参数
+            # 所以我们可以像操作普通模型一样直接调用 state_dict()
+            full_state_dict = model.module.state_dict()
+
+            # 将其克隆到 CPU 并返回
+            return {k: v.cpu().clone() for k, v in full_state_dict.items()}
+
+    # 其他 rank 执行到这里时，上下文结束，直接返回 None
+    return None
+
+    # # ZeRO-3
+    # state_dict_on_rank_0 = {}
+    # for param_name, param in model.module.named_parameters():
+    #     if hasattr(param, 'ds_id'):
+    #         with deepspeed.zero.GatheredParameters(param, modifier_rank=0):
+    #             if TrainerTools().parallel.is_main_process:
+    #                 state_dict_on_rank_0[param_name] = param.data.to(torch.float32).cpu().clone()
+    #     else:
+    #         print("22222222")
+    #         if TrainerTools().parallel.is_main_process:
+    #             state_dict_on_rank_0[param_name] = param.data.to(torch.float32).cpu().clone()
+    #
+    # return state_dict_on_rank_0 if TrainerTools().parallel.is_main_process else None
 
 
 def get_ds_model_params(model: nn.Module):

llm_trainer/generate_utils.py

@@ -244,7 +244,7 @@ def streaming_generate(
         max_position_embeddings: int,
         max_new_tokens: int,
         temperature: Optional[float] = 1.0,
-        k: Optional[int] = 50,
+        k: Optional[int] = None,
         p: Optional[float] = 1.0,
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,

llm_trainer/trainer.py

@@ -139,10 +139,10 @@ class Trainer:
                 if not any(sub_module in name for sub_module in ['multi_modal_projector']):
                     param.requires_grad = False
 
-            model.embed_tokens.eval()
-            model.layers.eval()
-            model.head_norm.eval()
-            model.lm_head.eval()
+            # 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())
@@ -421,7 +421,7 @@ class Trainer:
 
         if TrainerTools().parallel.is_main_process:
             eval_prompt, eval_image_tag = self._get_eval_data()
-            if isinstance(self.train_model, VlmModel) and self.pixel_values_provider and eval_image_tag:
+            if isinstance(self.train_config, VLMConfig) and self.pixel_values_provider and eval_image_tag:
                 eval_pixel_values = self.pixel_values_provider([eval_image_tag])
             else:
                 eval_pixel_values = None
@@ -445,7 +445,7 @@ class Trainer:
 
         if TrainerTools().parallel.is_main_process:
             eval_prompt, eval_image_tag = self._get_eval_data()
-            if isinstance(self.train_model, VlmModel) and self.pixel_values_provider and eval_image_tag:
+            if isinstance(self.train_config, VLMConfig) and self.pixel_values_provider and eval_image_tag:
                 eval_pixel_values = self.pixel_values_provider([eval_image_tag])
             else:
                 eval_pixel_values = None

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

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

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

@@ -3,10 +3,10 @@ llm_trainer/checkpoint.py,sha256=yZcExxneN2yzvWxRiK-pstMWs35LV7GiOfqcLq-S6vc,574
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
 llm_trainer/dcp.py,sha256=PkD97DyrOtoTKn4FJsfL3VqAy4dxufgjdzJEz8-Cnoc,3635
 llm_trainer/dpo_trainer.py,sha256=rC_I5ipesSlP3gFK_SG2GB8NbgJAMu4K7KLxkAS-aRY,13406
-llm_trainer/ds_checkpoint.py,sha256=nchGocJE2oJnQ_KNN1kw-BkOAEIyTtO8SJt41cuN_xM,4232
+llm_trainer/ds_checkpoint.py,sha256=7o9oxHUqPJNQESuZz83vHUmV83AkUh19mV9nY6qg4PE,4957
 llm_trainer/eval.py,sha256=NDm8PbXLch7xT81xPYPRCNrcrB_Xj5GDJSCxyVwUOp4,1524
 llm_trainer/fsdp_checkpoint.py,sha256=lqZFzHyWyfzuCq_81kQNtJd2qaiMeY1N5BCEMnrJTBw,3192
-llm_trainer/generate_utils.py,sha256=RpAIjN0fvyTkMk9b9x7YE6c5GiiE3x5YGyPaa4R_BjA,15191
+llm_trainer/generate_utils.py,sha256=_3TWAt-W8ZIzDZrLHEBR2iiZ3bn4V34WuVvKgCuDtyI,15193
 llm_trainer/grpo_trainer.py,sha256=bZPrxhyPQLAnFzWhI7hhA6fpuKVNwj7nOm9k0ku9aK4,15977
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
@@ -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=aoZYL5U4Z5axXBMM_DHgzIzJ89YbU9xUQ56jppcT65c,25339
+llm_trainer/trainer.py,sha256=Zy1oesBfsFlDedZ4hn3gcAkTrpi5fr76bFFQikfAkak,25351
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-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,,
+project_llm_trainer-0.4.6.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.6.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.6.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.6.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.6.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.6.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.6.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.6.dist-info/METADATA,sha256=PcLGKG5luK4XZFXbVRjiBrXNq9EzRvWCgZ6K1cxMzlo,195
+project_llm_trainer-0.4.6.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.6.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.6.dist-info/RECORD,,