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

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)
@@ -178,6 +217,10 @@ def _generate(
             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
 
@@ -215,6 +270,7 @@ def _streaming_generate(
         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:
@@ -250,6 +307,7 @@ def streaming_generate(
         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:
@@ -283,6 +342,7 @@ def generate(
         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:

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

llm_trainer/trainer.py

@@ -540,13 +540,17 @@ class Trainer:
                         if gradient_accumulation_steps > 1:
                             loss = loss / gradient_accumulation_steps
 
-                        loss_accumulation += loss.detach()
+                        loss_accumulation += loss.detach().item()
                         self._backward_loss(loss)
 
                         if need_update_grad:
+                            loss_tensor = torch.tensor(loss_accumulation, device=TrainerTools().parallel.device)
+
                             # todo check all_reduce??
                             if TrainerTools().parallel.parallel_train:
-                                dist.all_reduce(loss_accumulation, dist.ReduceOp.AVG)
+                                dist.all_reduce(loss_tensor, dist.ReduceOp.AVG)
+
+                            final_log_loss = loss_tensor.item()
 
                             # ds模式已经集成gradient_clipping
                             if not isinstance(TrainerTools().parallel, DsParallel) and self.lr_scheduler.can_clip_grad():
@@ -560,7 +564,7 @@ class Trainer:
                                 epoch_tag=f'epoch: {epoch}',
                                 file_tag=f'file: {file_idx + 1}/{file_count}',
                                 batch_tag=f'batch: {batch}/{batch_count_per_file}',
-                                loss=loss_accumulation.item()
+                                loss=final_log_loss
                             )
                             # reset to default
                             loss_accumulation = 0.0

{project_llm_trainer-0.4.7.dist-info → project_llm_trainer-0.4.10.dist-info}/METADATA

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

{project_llm_trainer-0.4.7.dist-info → project_llm_trainer-0.4.10.dist-info}/RECORD

@@ -6,7 +6,7 @@ llm_trainer/dpo_trainer.py,sha256=rC_I5ipesSlP3gFK_SG2GB8NbgJAMu4K7KLxkAS-aRY,13
 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=BmjpCrus_jvJ3SM2KS1bQNzJWAFnpJ9mI28iBWXZpvo,15206
+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/trainer.py,sha256=Zy1oesBfsFlDedZ4hn3gcAkTrpi5fr76bFFQikfAkak,25351
+llm_trainer/train_configs.py,sha256=gzTXMLUuQexRvqyKIZQ1U6ESa0DELD7hPpYZdrDcyxg,15974
+llm_trainer/trainer.py,sha256=pUtJVRosn54j1hn76CFAptJcAsrDo59H6p8NMkg2zt4,25521
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-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,,
+project_llm_trainer-0.4.10.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.10.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.10.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.10.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.10.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.10.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.10.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.10.dist-info/METADATA,sha256=zrHUkQPm7Zox2CSeYN5HBqedZebXuZAQgZVj0O24U6I,196
+project_llm_trainer-0.4.10.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.10.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.10.dist-info/RECORD,,