project-llm-trainer 0.4__py3-none-any.whl → 0.4.1__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/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.1.dist-info}/METADATA

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

{project_llm_trainer-0.4.dist-info → project_llm_trainer-0.4.1.dist-info}/RECORD

@@ -1,11 +1,11 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=GPaSGvnLCGMgsIA_vfjuw34tTQY26EuNwu7c08fhJHQ,5638
+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=H0BxYQixOWKRC20t55cFqNDTPzalD3AGTVt-owIB0_4,4488
-llm_trainer/eval.py,sha256=CsB3TpSVwhYVS9SP4Kuj_JhFUUvLcZUkvd8hvEIkPDU,1782
-llm_trainer/fsdp_checkpoint.py,sha256=xPQnAfXbx1SRKcVDLLgOtVrqjk0CjIRleVY0ZrwOAJU,3876
+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=4iM0vyc_1C_iTL31GlS9PR4eZtYaELPRZ02KDSPZA9U,15158
 llm_trainer/grpo_trainer.py,sha256=fqLT48ORSCece_e8dpyt8J7EarDuTnGoJ_eHk7Oy-1k,16177
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
@@ -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=DujZR1KOHyP3EHR8uIQPEsnX_5b7YC9Cto_eH7zxWqc,25256
+llm_trainer/trainer.py,sha256=hOn-z8kOd67RTuaaNMmdQjlw7N5LIZRHjSt5frpA1xI,25355
 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,,
+project_llm_trainer-0.4.1.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.1.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.1.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.1.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.1.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.1.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.1.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.1.dist-info/METADATA,sha256=9z1AB745r7BzQHNc3j-3N2nOdB9ZRUYsxcM42QoSb1o,195
+project_llm_trainer-0.4.1.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.1.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.1.dist-info/RECORD,,