project-llm-trainer 0.4.4__py3-none-any.whl → 0.4.7__py3-none-any.whl

llm_trainer/ds_checkpoint.py

@@ -69,7 +69,7 @@ def load_ds_checkpoint_for_eval(model: nn.Module):
 
 def _get_ds_full_state_dict_on_rank0(model: DeepSpeedEngine) -> Optional[dict]:
     """
-        可以在任意rank上调用，然后只有rank0有值
+        需要在所有rank上调用，然后只有rank0有值
     """
 
     if model.zero_optimization_stage() != 3:
@@ -77,18 +77,32 @@ 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:
+    #         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

@@ -175,7 +175,7 @@ def _generate(
             if k and k != 0:
                 logits = _top_k_warper(logits, k, device)
 
-            if p and p < 1:
+            if p and 0 < p <= 1:
                 logits = _top_p_warper(logits, p)
 
             if multinomial:
@@ -210,7 +210,7 @@ def _streaming_generate(
         max_new_tokens: int,
         temperature: Optional[float] = 1.0,
         k: Optional[int] = None,
-        p: Optional[float] = 1.0,
+        p: Optional[float] = None,
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
@@ -244,8 +244,8 @@ def streaming_generate(
         max_position_embeddings: int,
         max_new_tokens: int,
         temperature: Optional[float] = 1.0,
-        k: Optional[int] = 50,
-        p: Optional[float] = 1.0,
+        k: Optional[int] = None,
+        p: Optional[float] = None,
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
@@ -278,7 +278,7 @@ def generate(
         max_new_tokens: int,
         temperature: Optional[float] = 1.0,
         k: Optional[int] = None,
-        p: Optional[float] = 1.0,
+        p: Optional[float] = None,
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
@@ -382,7 +382,7 @@ def batch_generate(
             if k and k != 0:
                 logits = _top_k_warper(logits, k, device)
 
-            if p and p < 1:
+            if p and 0 < p <= 1:
                 logits = _top_p_warper(logits, p)
 
             prob = logits.softmax(dim=-1)

llm_trainer/trainer.py

@@ -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.4.dist-info → project_llm_trainer-0.4.7.dist-info}/METADATA

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

{project_llm_trainer-0.4.4.dist-info → project_llm_trainer-0.4.7.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=x_tjgJR47P8gVwV4qAnTUCGwx7eVq2Epw0vOVV7fkYo,4925
 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=BmjpCrus_jvJ3SM2KS1bQNzJWAFnpJ9mI28iBWXZpvo,15206
 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=Aynp-lT1-lHHJ2S7FboHb57M4bPe51lLSgM6SOICQOo,25347
+llm_trainer/trainer.py,sha256=Zy1oesBfsFlDedZ4hn3gcAkTrpi5fr76bFFQikfAkak,25351
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.4.4.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.4.4.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.4.4.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.4.4.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.4.4.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.4.4.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.4.4.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.4.4.dist-info/METADATA,sha256=7cpmCfNcYyufr2GgoiK3UoelmgWQZh5BOi6rCXrxxkA,195
-project_llm_trainer-0.4.4.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.4.4.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.4.4.dist-info/RECORD,,
+project_llm_trainer-0.4.7.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.7.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.7.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.7.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.7.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.7.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.7.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.7.dist-info/METADATA,sha256=u4_cQkQaH9QKqG_XcWiXzGHD5rnrzqHjvJWQvgVnkZQ,195
+project_llm_trainer-0.4.7.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.7.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.7.dist-info/RECORD,,