project-llm-trainer 0.4__py3-none-any.whl → 0.4.2__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -121,8 +121,12 @@ def load_checkpoint_for_eval(
 
 def copy_model_params(
         _from: nn.Module,
-        _to: nn.Module
+        _to: Optional[nn.Module]
 ):
+    """
+        必须在所有rank上调用，非rank0, _to可以设置为None
+    """
+
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import get_ds_model_params
         state_dict = get_ds_model_params(_from)
@@ -134,7 +138,7 @@ def copy_model_params(
     else:
         state_dict = _from.state_dict()
 
-    if state_dict:
+    if _to and state_dict:
         _to.load_state_dict(state_dict)
 
 

llm_trainer/ds_checkpoint.py

@@ -67,45 +67,47 @@ def load_ds_checkpoint_for_eval(model: nn.Module):
     model.load_state_dict(state_dict)
 
 
+def _get_ds_full_state_dict_on_rank0(model: DeepSpeedEngine) -> Optional[dict]:
+    """
+        可以在任意rank上调用，然后只有rank0有值
+    """
+
+    if model.zero_optimization_stage() != 3:
+        if TrainerTools().parallel.is_main_process:
+            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
+
+
 def get_ds_model_params(model: nn.Module):
     """
         从一个正在运行的 DeepSpeedEngine 中高效地提取完整的 FP32 state_dict，
         兼容 ZeRO Stages 0, 1, 2, 3。
-        这个版本包含了对 ZeRO-3 中非分片参数的正确处理。
+        包含了对 ZeRO-3 中分片参数的正确处理。
     """
 
     assert isinstance(model, DeepSpeedEngine)
-    zero_stage = model.zero_optimization_stage()
-    state_dict = None
-
-    if TrainerTools().parallel.is_main_process:
-        if zero_stage == 3:
-            # ZeRO-3: Rank 0 聚合参数来构建完整的 state_dict
-            state_dict = {}
-            for param in model.module.parameters():
-                # 关键检查：判断参数是否被 ZeRO-3 分片管理
-                if hasattr(param, 'ds_id'):
-                    # 这是被分片的参数，使用 GatheredParameters 聚合
-                    with deepspeed.zero.GatheredParameters(param, modifier_rank=0):
-                        # .clone() 创建一个独立副本, .to('cpu') 移动到CPU, .to(torch.float32) 确保类型
-                        state_dict[param.ds_name] = param.data.to(torch.float32).cpu().clone()
-                else:
-                    # 这是未被分片的参数 (e.g., tied weights, buffers), 直接从 Rank 0 复制
-                    state_dict[param.ds_name] = param.data.to(torch.float32).cpu().clone()
-        else:  # zero_stage in [0, 1, 2]
-            # 在这些 stage，rank 0 已经有完整的模型。
-            # 我们从 model_engine.module 获取原始模型状态。
-            state_dict = {k: v.cpu().clone() for k, v in model.module.state_dict().items()}
+    state_dict = _get_ds_full_state_dict_on_rank0(model)
 
     # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
     # 我们需要将其广播给所有进程。
     if TrainerTools().parallel.world_size > 1:
         # 准备一个列表，rank 0 有数据，其他 rank 是占位符
         object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
-
         # 执行广播，这个操作是阻塞的，会同步所有进程
         dist.broadcast_object_list(object_list, src=0)
-
         # 所有进程从列表中获取广播后的 state_dict 副本
         state_dict = object_list[0]
 

llm_trainer/eval.py

@@ -1,7 +1,6 @@
 import torch
 
 from .generate_utils import generate
-from .checkpoint import copy_model_params
 from .log import get_log_dir
 from .tools import TrainerTools
 from .train_configs import EvalConfig
@@ -37,7 +36,6 @@ def _eval_task(
 
 
 def submit_gen_task(
-        train_model: torch.nn.Module,
         eval_model: torch.nn.Module,
         eval_config: EvalConfig,
         tag,
@@ -46,13 +44,6 @@ def submit_gen_task(
         max_position_embeddings,
         tokens_per_image
 ):
-    try:
-        copy_model_params(_from=train_model, _to=eval_model)
-    except Exception as e:
-        if isinstance(e, KeyboardInterrupt):
-            raise e
-        return
-
     eval_model.to(TrainerTools().parallel.device)
     _eval_task(
         eval_model=eval_model,

llm_trainer/fsdp_checkpoint.py

@@ -53,6 +53,18 @@ def load_fsdp_checkpoint(
             optimizer.load_state_dict(state_dict['optim_state_dict'])
 
 
+def _get_fsdp_full_state_dict_on_rank0(model: nn.Module) -> Optional[dict]:
+    """
+        可以在任意rank上调用，然后只有rank0有值
+    """
+
+    from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+    with FSDP.summon_full_params(model, writeback=False, offload_to_cpu=True):
+        if TrainerTools().parallel.is_main_process:
+            return {k: v.clone() for k, v in model.state_dict().items()}
+
+    return None
+
 
 def get_fsdp_model_params(model: nn.Module):
     """
@@ -60,31 +72,15 @@ def get_fsdp_model_params(model: nn.Module):
         这个函数会聚合所有分片的参数，并确保所有 rank 都收到一个完整的副本。
     """
 
-    # FSDP 要求在所有 rank 上都调用 summon_full_params，即使我们只在 rank 0 上操作。
-    # writeback=False: 我们只读取参数，不写回，可以节省开销。
-    # offload_to_cpu=True: 直接将聚合后的参数卸载到 CPU，避免在 GPU 上产生大的峰值内存，
-    #                      并为我们省去了 .cpu() 的步骤。这是一个非常有用的优化。
-    # rank0_only=False: 为了让 offload_to_cpu 在所有 rank 上都生效，这里通常设为 False。
-    #                   我们稍后通过 get_rank() 来确保只有 rank 0 实际构建字典。
-    with FSDP.summon_full_params(model, writeback=False, offload_to_cpu=True):
-
-        state_dict = None
-        if TrainerTools().parallel.is_main_process:
-            # 在这个 with 块内部, model.state_dict() 会返回一个在 CPU 上的、完整的状态字典。
-            # 因为我们设置了 offload_to_cpu=True。
-            # 我们使用 .clone() 来确保我们得到的是一个独立的副本，
-            # 尽管 offload_to_cpu 已经帮我们处理了大部分情况。
-            state_dict = {k: v.clone() for k, v in model.state_dict().items()}
+    state_dict = _get_fsdp_full_state_dict_on_rank0(model)
 
     # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
     # 我们需要将其广播给所有进程。
     if TrainerTools().parallel.world_size > 1:
         # 准备一个列表，rank 0 有数据，其他 rank 是占位符
         object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
-
         # 执行广播，这个操作是阻塞的，会同步所有进程
         dist.broadcast_object_list(object_list, src=0)
-
         # 所有进程从列表中获取广播后的 state_dict 副本
         state_dict = object_list[0]
 

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

@@ -31,6 +31,7 @@ from .scheduler import (
 from .checkpoint import (
     load_checkpoint,
     save_checkpoint,
+    copy_model_params,
     load_steps,
     save_steps,
 )
@@ -416,6 +417,8 @@ class Trainer:
             self,
             tag: str
     ):
+        copy_model_params(_from=self.train_model, _to=self.eval_model)
+
         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:
@@ -424,7 +427,6 @@ class Trainer:
                 eval_pixel_values = None
 
             submit_gen_task(
-                self.train_model,
                 self.eval_model,
                 self.train_config.eval_config,
                 tag=f'sign:batch/{tag}',
@@ -439,6 +441,8 @@ class Trainer:
             self,
             tag: str
     ):
+        copy_model_params(_from=self.train_model, _to=self.eval_model)
+
         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:
@@ -447,7 +451,6 @@ class Trainer:
                 eval_pixel_values = None
 
             submit_gen_task(
-                self.train_model,
                 self.eval_model,
                 self.train_config.eval_config,
                 tag=f'sign:epoch/{tag}',

{project_llm_trainer-0.4.dist-info → project_llm_trainer-0.4.2.dist-info}/METADATA

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

project_llm_trainer-0.4.2.dist-info/RECORD

@@ -0,0 +1,35 @@
+llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
+llm_trainer/checkpoint.py,sha256=yZcExxneN2yzvWxRiK-pstMWs35LV7GiOfqcLq-S6vc,5745
+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/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/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
+llm_trainer/parallel_ddp.py,sha256=Pob9vUlBZnkL4oP1Re11kFob7nufMSE96pn7m7fuOEM,1345
+llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,1145
+llm_trainer/parallel_fsdp.py,sha256=cQOdY8ou6m8OsR06PpFVn6GiyZlK9nefkcGyszUOIJk,4055
+llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,676
+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=hOn-z8kOd67RTuaaNMmdQjlw7N5LIZRHjSt5frpA1xI,25355
+llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
+project_llm_trainer-0.4.2.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.2.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.2.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.2.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.2.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.2.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.2.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.2.dist-info/METADATA,sha256=CBDzoyiYlDPzgmffGgQIMy134eDLLOyRwBoLeUzgQ2g,195
+project_llm_trainer-0.4.2.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.2.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.2.dist-info/RECORD,,

project_llm_trainer-0.4.dist-info/RECORD

@@ -1,35 +0,0 @@
-llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=GPaSGvnLCGMgsIA_vfjuw34tTQY26EuNwu7c08fhJHQ,5638
-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=H0BxYQixOWKRC20t55cFqNDTPzalD3AGTVt-owIB0_4,4488
-llm_trainer/eval.py,sha256=CsB3TpSVwhYVS9SP4Kuj_JhFUUvLcZUkvd8hvEIkPDU,1782
-llm_trainer/fsdp_checkpoint.py,sha256=xPQnAfXbx1SRKcVDLLgOtVrqjk0CjIRleVY0ZrwOAJU,3876
-llm_trainer/generate_utils.py,sha256=4iM0vyc_1C_iTL31GlS9PR4eZtYaELPRZ02KDSPZA9U,15158
-llm_trainer/grpo_trainer.py,sha256=fqLT48ORSCece_e8dpyt8J7EarDuTnGoJ_eHk7Oy-1k,16177
-llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
-llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
-llm_trainer/parallel.py,sha256=DQu8GqEFxD99HQ6hKuIxxyKi-05dMO33eMhImYlPuOI,4468
-llm_trainer/parallel_ddp.py,sha256=Pob9vUlBZnkL4oP1Re11kFob7nufMSE96pn7m7fuOEM,1345
-llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,1145
-llm_trainer/parallel_fsdp.py,sha256=cQOdY8ou6m8OsR06PpFVn6GiyZlK9nefkcGyszUOIJk,4055
-llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,676
-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=DujZR1KOHyP3EHR8uIQPEsnX_5b7YC9Cto_eH7zxWqc,25256
-llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.4.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.4.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.4.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.4.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.4.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.4.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.4.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.4.dist-info/METADATA,sha256=-xxg-UyXn5MhW5OdYGFUcL5DtbIkgnQoUZS5b5bcEio,193
-project_llm_trainer-0.4.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.4.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.4.dist-info/RECORD,,