project-llm-trainer 0.4.13__py3-none-any.whl → 0.4.15__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -119,29 +119,6 @@ def load_checkpoint_for_eval(
         load_checkpoint(model, None, device, suffix=suffix)
 
 
-def copy_model_params(
-        _from: 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)
-    elif isinstance(TrainerTools().parallel, FsdpParallel):
-        from .fsdp_checkpoint import get_fsdp_model_params
-        state_dict = get_fsdp_model_params(_from)
-    elif isinstance(_from, DDP):
-        state_dict = _from.module.state_dict()
-    else:
-        state_dict = _from.state_dict()
-
-    if _to and state_dict:
-        _to.load_state_dict(state_dict)
-
-
 def save_steps(global_steps: int, lr_scheduler: Optional[LRScheduler] = None):
     # 暂时只保存主进程的
     if TrainerTools().parallel.is_main_process:

llm_trainer/dpo_trainer.py

@@ -12,13 +12,14 @@ from .dataset import DPODataset
 from .loss import DPOLoss
 from .tools import TrainerTools
 from .utils import get_dpo_collate_fn
+from .model_params import copy_model_params
 
 from .checkpoint import (
     save_checkpoint,
-    copy_model_params,
     save_steps,
 )
 
+
 class DPOTrainer(Trainer):
     def __init__(
             self,

llm_trainer/ds_checkpoint.py

@@ -2,11 +2,7 @@ import os
 from typing import Optional
 from glob import glob
 import shutil
-import torch
 from torch import nn
-import torch.distributed as dist
-
-from .tools import TrainerTools
 
 try:
     import deepspeed
@@ -65,65 +61,3 @@ def load_ds_checkpoint_for_eval(model: nn.Module):
     ckpt_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
     state_dict = get_fp32_state_dict_from_zero_checkpoint(ckpt_dir)
     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 ---
-    # 只调用一次 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):
-    """
-        从一个正在运行的 DeepSpeedEngine 中高效地提取完整的 FP32 state_dict，
-        兼容 ZeRO Stages 0, 1, 2, 3。
-        包含了对 ZeRO-3 中分片参数的正确处理。
-    """
-
-    assert isinstance(model, DeepSpeedEngine)
-    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]
-
-    return state_dict
-

llm_trainer/ds_model_params.py

@@ -0,0 +1,72 @@
+from typing import Optional
+from torch import nn
+import torch.distributed as dist
+
+from .tools import TrainerTools
+
+try:
+    import deepspeed
+    from deepspeed import DeepSpeedEngine
+    from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+except: ...
+
+
+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 ---
+    # 只调用一次 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, only_rank0=False):
+    """
+        从一个正在运行的 DeepSpeedEngine 中高效地提取完整的 FP32 state_dict，
+        兼容 ZeRO Stages 0, 1, 2, 3。
+        包含了对 ZeRO-3 中分片参数的正确处理。
+    """
+
+    assert isinstance(model, DeepSpeedEngine)
+    state_dict = _get_ds_full_state_dict_on_rank0(model)
+
+    # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
+    # 我们需要将其广播给所有进程。
+    if not only_rank0 and 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]
+
+    return state_dict

llm_trainer/fsdp_checkpoint.py

@@ -4,7 +4,6 @@ import torch
 from torch import nn
 from torch.optim import Optimizer
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-import torch.distributed as dist
 
 from .tools import TrainerTools
 
@@ -51,37 +50,3 @@ def load_fsdp_checkpoint(
 
         if optimizer:
             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):
-    """
-        从一个 FSDP 包装的模型中高效地提取完整的 FP32 state_dict。
-        这个函数会聚合所有分片的参数，并确保所有 rank 都收到一个完整的副本。
-    """
-
-    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]
-
-    return state_dict

llm_trainer/fsdp_model_params.py

@@ -0,0 +1,39 @@
+from typing import Optional
+from torch import nn
+import torch.distributed as dist
+
+from .tools import TrainerTools
+
+
+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, only_rank0=False):
+    """
+        从一个 FSDP 包装的模型中高效地提取完整的 FP32 state_dict。
+        这个函数会聚合所有分片的参数，并确保所有 rank 都收到一个完整的副本。
+    """
+
+    state_dict = _get_fsdp_full_state_dict_on_rank0(model)
+
+    # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
+    # 我们需要将其广播给所有进程。
+    if not only_rank0 and 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]
+
+    return state_dict

llm_trainer/grpo_trainer.py

@@ -15,10 +15,10 @@ from .loss import GRPOLoss
 from .tools import TrainerTools
 from .generate_utils import batch_generate
 from .log import log
+from .model_params import copy_model_params
 
 from .checkpoint import (
     save_checkpoint,
-    copy_model_params,
     save_steps,
 )
 

llm_trainer/model_params.py

@@ -0,0 +1,28 @@
+from typing import Optional
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from .tools import TrainerTools
+from .parallel_ds import DsParallel
+from .parallel_fsdp import FsdpParallel
+
+def copy_model_params(
+        _from: nn.Module,
+        _to: Optional[nn.Module]
+):
+    """
+        必须在所有rank上调用，非rank0, _to可以设置为None
+    """
+    if isinstance(TrainerTools().parallel, DsParallel):
+        from .ds_model_params import get_ds_model_params
+        state_dict = get_ds_model_params(_from, only_rank0=_to is None)
+    elif isinstance(TrainerTools().parallel, FsdpParallel):
+        from .fsdp_model_params import get_fsdp_model_params
+        state_dict = get_fsdp_model_params(_from, only_rank0=_to is None)
+    elif isinstance(_from, DDP):
+        state_dict = _from.module.state_dict()
+    else:
+        state_dict = _from.state_dict()
+
+    if _to and state_dict:
+        _to.load_state_dict(state_dict)

llm_trainer/trainer.py

@@ -14,6 +14,8 @@ from .parallel_fsdp import FsdpParallel
 from .tools import TrainerTools
 from .loss import LMLoss, KDLoss
 from .dataset import TextDataset
+from .model_params import copy_model_params
+from .eval import submit_gen_task
 
 from .train_configs import (
     TrainConfig,
@@ -31,10 +33,10 @@ from .scheduler import (
 from .checkpoint import (
     load_checkpoint,
     save_checkpoint,
-    copy_model_params,
     load_steps,
     save_steps,
 )
+
 from .utils import (
     set_seed,
     pretrain_collate_fn,
@@ -45,8 +47,6 @@ from .log import(
     get_log_dir
 )
 
-from .eval import submit_gen_task
-
 class Trainer:
     def __init__(
             self,

{project_llm_trainer-0.4.13.dist-info → project_llm_trainer-0.4.15.dist-info}/METADATA

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

{project_llm_trainer-0.4.13.dist-info → project_llm_trainer-0.4.15.dist-info}/RECORD

@@ -1,15 +1,18 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=yZcExxneN2yzvWxRiK-pstMWs35LV7GiOfqcLq-S6vc,5745
+llm_trainer/checkpoint.py,sha256=AvUC1JLxuahKtg3VNW20VHIE3iIjpaMHIi_pyyDYVJ0,5043
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
 llm_trainer/dcp.py,sha256=PkD97DyrOtoTKn4FJsfL3VqAy4dxufgjdzJEz8-Cnoc,3635
-llm_trainer/dpo_trainer.py,sha256=djBhvI_ixTV1nLNg84tgCpfV--pu6IRiOhO28V-aANQ,11425
-llm_trainer/ds_checkpoint.py,sha256=x_tjgJR47P8gVwV4qAnTUCGwx7eVq2Epw0vOVV7fkYo,4925
+llm_trainer/dpo_trainer.py,sha256=o5lYxt6yVMCvoBqW_yTu9l6Ff-xjEu-CwdPVttu3H8E,11447
+llm_trainer/ds_checkpoint.py,sha256=wz48HoLBBt8QGO1tXfvJwrXoiGtPG_gjwHfEqARllso,2175
+llm_trainer/ds_model_params.py,sha256=Nwmv0YcBtO6ynC0dXallAD1rWkN22-elGfVjLaWp2Yg,2988
 llm_trainer/eval.py,sha256=NDm8PbXLch7xT81xPYPRCNrcrB_Xj5GDJSCxyVwUOp4,1524
-llm_trainer/fsdp_checkpoint.py,sha256=lqZFzHyWyfzuCq_81kQNtJd2qaiMeY1N5BCEMnrJTBw,3192
+llm_trainer/fsdp_checkpoint.py,sha256=xsm71s9WeTaBvBvv6CbuGpwkmX3V6i3xmBcMTDfGxKc,1770
+llm_trainer/fsdp_model_params.py,sha256=MRjrs9zmMl-61a1l6188Ij5PSalzztOSp8E4evDvJXo,1541
 llm_trainer/generate_utils.py,sha256=tSbA_tLqSq5qJGHSOlPv5T3iRDZkbFg5ZvDAgJ_i_SE,17946
-llm_trainer/grpo_trainer.py,sha256=bZPrxhyPQLAnFzWhI7hhA6fpuKVNwj7nOm9k0ku9aK4,15977
+llm_trainer/grpo_trainer.py,sha256=1gZXiL1pogLFecFQUGj9zCU_k66ryVjZciYyd8J5ph4,15998
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=Yv3fsaVuZ5AhnGPJOr5vEMb_tM2urR6mCb4DBbrHHI8,6030
+llm_trainer/model_params.py,sha256=2f2W9KRCjyqSfEwxI3w5f6TPZaqq25WzY-nEc7aJxcs,970
 llm_trainer/parallel.py,sha256=DQu8GqEFxD99HQ6hKuIxxyKi-05dMO33eMhImYlPuOI,4468
 llm_trainer/parallel_ddp.py,sha256=Pob9vUlBZnkL4oP1Re11kFob7nufMSE96pn7m7fuOEM,1345
 llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,1145
@@ -20,16 +23,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=HKzH3nfMT1-SW4Htwa0KqYtMd6FAJcthR5IEo6di8us,8168
-llm_trainer/trainer.py,sha256=j5fDqMzvU6SYwxHsv9wX0UVX4JXS-8eP1AkHgVxKf9U,27119
+llm_trainer/trainer.py,sha256=95ARdNDfalhZ7Ug-fDj3qIhWEiZQeX9n5WANhijIRLE,27140
 llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.4.13.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.4.13.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.4.13.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.4.13.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.4.13.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.4.13.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.4.13.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.4.13.dist-info/METADATA,sha256=hiW-7qgWuPizKVz4cU8mEHoqiuT6ZqNlCBb7nwVfFQ4,196
-project_llm_trainer-0.4.13.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.4.13.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.4.13.dist-info/RECORD,,
+project_llm_trainer-0.4.15.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.4.15.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.4.15.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.4.15.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.4.15.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.4.15.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.4.15.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.4.15.dist-info/METADATA,sha256=5sveZ3kkRMVCz9dI5_NI64o9tFBVsJhHhun9vwzzL9Q,196
+project_llm_trainer-0.4.15.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.4.15.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.4.15.dist-info/RECORD,,