project-llm-trainer 0.4.6__py3-none-any.whl → 0.4.9__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:
@@ -99,7 +99,6 @@ def _get_ds_full_state_dict_on_rank0(model: DeepSpeedEngine) -> Optional[dict]:
     #             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()
     #

llm_trainer/generate_utils.py

@@ -1,7 +1,6 @@
 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
@@ -109,7 +108,8 @@ def _generate(
         pixel_values: Optional[torch.Tensor] = None,
         tokens_per_image: int = -1,
         suppress_tokens: Optional[List[int]] = None,
-        device: Union[str, torch.device, int]
+        device: Union[str, torch.device, int],
+        reasoning_budget: Optional[int] = None
 ):
     """
     :param model:
@@ -144,6 +144,27 @@ def _generate(
     kv_cache: Optional[KVCache] = None
     generate_tokens = tokens.clone()
 
+    reasoning_start = TrainerTools().tokenizer.reasoning_start
+    reasoning_end = TrainerTools().tokenizer.reasoning_end
+
+    # --- 状态初始化 ---
+    in_reasoning_block = False
+    reasoning_step_count = 0
+    # “冷静期”标志位。当强制结束思考后，在下一步抑制<reasoning>的生成。
+    suppress_reasoning_start_next = False
+
+    if reasoning_budget is not None:
+        prompt_tokens = tokens[0]
+        start_indices = (prompt_tokens == reasoning_start).nonzero(as_tuple=True)[0]
+        end_indices = (prompt_tokens == reasoning_end).nonzero(as_tuple=True)[0]
+
+        last_start_idx = start_indices[-1].item() if len(start_indices) > 0 else -1
+        last_end_idx = end_indices[-1].item() if len(end_indices) > 0 else -1
+
+        if last_start_idx > last_end_idx:
+            in_reasoning_block = True
+            reasoning_step_count = len(prompt_tokens) - 1 - last_start_idx
+
     model.eval()
     with torch.inference_mode():
         for _ in range(max_new_tokens):
@@ -163,6 +184,24 @@ def _generate(
 
             # (batch, vocab_size)
             logits = logits[:, -1, :]
+
+            # --- 推理预算逻辑 ---
+            force_end_reasoning_token = False
+            if reasoning_budget is not None:
+                # 检查是否需要在此步抑制 <reasoning>
+                should_suppress_this_step = suppress_reasoning_start_next
+                suppress_reasoning_start_next = False  # 立即重置标志位
+
+                # 修改: 检查是否超出预算
+                if in_reasoning_block and reasoning_step_count >= reasoning_budget:
+                    force_end_reasoning_token = True
+                    # 设置标志位，在下一步抑制 <reasoning>
+                    suppress_reasoning_start_next = True
+
+                # 如果上一轮设置了抑制标志，则在此轮执行抑制
+                if should_suppress_this_step:
+                    logits[:, reasoning_start] = -float("inf")
+
             # 抑制特殊token输出
             if suppress_tokens and len(suppress_tokens) != 0:
                 logits = _suppress_warper(logits, suppress_tokens)
@@ -175,9 +214,13 @@ 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 force_end_reasoning_token:
+                logits[:] = -float("inf")
+                logits[:, reasoning_end] = 0.0
+
             if multinomial:
                 prob = logits.softmax(dim=-1)
                 # 返回下标
@@ -186,6 +229,18 @@ def _generate(
                 # 返回下标
                 next_token = logits.argmax(dim=-1, keepdim=True)
 
+            if reasoning_budget is not None:
+                current_token_id = next_token.item()
+                if not in_reasoning_block and current_token_id == reasoning_start:
+                    in_reasoning_block = True
+                    reasoning_step_count = 0
+                elif in_reasoning_block:
+                    if current_token_id == reasoning_end:
+                        in_reasoning_block = False
+                        reasoning_step_count = 0
+                    else:
+                        reasoning_step_count += 1
+
             # token, is_full_result
             yield next_token, False
 
@@ -210,11 +265,12 @@ 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,
         device: Union[str, torch.device, int] = None,
+        reasoning_budget: Optional[int] = None
 ):
     device = TrainerTools().parallel.device if not device else device
     encoded_tokens = TrainerTools().tokenizer.encode(prompt, unsqueeze=True, covert_tensor=True).to(device)
@@ -230,7 +286,8 @@ def _streaming_generate(
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
         suppress_tokens=suppress_tokens,
-        device=device
+        device=device,
+        reasoning_budget=reasoning_budget
     )
 
     for (token, is_full_result) in generate_text_iterator:
@@ -245,11 +302,12 @@ 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,
         device: Union[str, torch.device, int] = None,
+        reasoning_budget: Optional[int] = None
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -262,7 +320,8 @@ def streaming_generate(
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
         suppress_tokens=suppress_tokens,
-        device=device
+        device=device,
+        reasoning_budget=reasoning_budget
     )
 
     for (token, is_full_result) in text_iterator:
@@ -278,11 +337,12 @@ 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,
         device: Union[str, torch.device, int] = None,
+        reasoning_budget: Optional[int] = None
 ):
     text_iterator = _streaming_generate(
         model=model,
@@ -295,7 +355,8 @@ def generate(
         suppress_tokens=suppress_tokens,
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
-        device=device
+        device=device,
+        reasoning_budget=reasoning_budget
     )
 
     for (token, is_full_result) in text_iterator:
@@ -382,7 +443,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/train_configs.py

@@ -22,11 +22,11 @@ class DsActivationCheckpointingConfig:
             self,
             *,
             partition_activations: bool = True,
-            cpu_checkpointing: bool = True,
+            cpu_checkpointing: bool = False,
             contiguous_memory_optimization: bool = True,
             number_checkpoints: Optional[int] = None,
-            synchronize_checkpoint_boundary: bool = True,
-            profile: bool = True
+            synchronize_checkpoint_boundary: bool = False,
+            profile: bool = False
     ):
         self.partition_activations =partition_activations
         self.cpu_checkpointing = cpu_checkpointing

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

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

{project_llm_trainer-0.4.6.dist-info → project_llm_trainer-0.4.9.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=7o9oxHUqPJNQESuZz83vHUmV83AkUh19mV9nY6qg4PE,4957
+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=_3TWAt-W8ZIzDZrLHEBR2iiZ3bn4V34WuVvKgCuDtyI,15193
+llm_trainer/generate_utils.py,sha256=tSbA_tLqSq5qJGHSOlPv5T3iRDZkbFg5ZvDAgJ_i_SE,17946
 llm_trainer/grpo_trainer.py,sha256=bZPrxhyPQLAnFzWhI7hhA6fpuKVNwj7nOm9k0ku9aK4,15977
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
@@ -19,17 +19,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=O45-20wRmh-nyTfU-U-XtjbKAoe7boEIsUvWT_NaKx4,3041
-llm_trainer/train_configs.py,sha256=arnet3tIzgVnwshod08F1jE7r4I7e-SIgMy55IagPnE,15971
+llm_trainer/train_configs.py,sha256=gzTXMLUuQexRvqyKIZQ1U6ESa0DELD7hPpYZdrDcyxg,15974
 llm_trainer/trainer.py,sha256=Zy1oesBfsFlDedZ4hn3gcAkTrpi5fr76bFFQikfAkak,25351
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-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,,
+project_llm_trainer-0.4.9.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.9.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.9.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.9.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.9.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.9.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.9.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.9.dist-info/METADATA,sha256=Xn5uE_6i2vFPTFf2O2iJT59czHzHxTAUeb6ZDCaH0-A,195
+project_llm_trainer-0.4.9.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.9.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.9.dist-info/RECORD,,