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

llm_trainer/tokenizer.py

@@ -26,8 +26,8 @@ class Tokenizer:
         self.text_user = '<user>'
         self.text_assistant = '<assistant>'
 
-        self.text_reasoning_start = '<reasoning>'
-        self.text_reasoning_end = '</reasoning>'
+        self.text_think_start = '<think>'
+        self.text_think_end = '</think>'
 
         self.text_answer_start = '<answer>'
         self.text_answer_end = '</answer>'
@@ -47,8 +47,8 @@ class Tokenizer:
             additional_special_tokens = [
                 AddedToken(self.text_user, lstrip=False, rstrip=False),
                 AddedToken(self.text_assistant, lstrip=False, rstrip=False),
-                AddedToken(self.text_reasoning_start, lstrip=False, rstrip=False),
-                AddedToken(self.text_reasoning_end, lstrip=False, rstrip=False),
+                AddedToken(self.text_think_start, lstrip=False, rstrip=False),
+                AddedToken(self.text_think_end, lstrip=False, rstrip=False),
                 AddedToken(self.text_answer_start, lstrip=False, rstrip=False),
                 AddedToken(self.text_answer_end, lstrip=False, rstrip=False),
                 AddedToken(self.text_system, lstrip=False, rstrip=False),
@@ -69,8 +69,8 @@ class Tokenizer:
         self.user = self.tokenizer.convert_tokens_to_ids(self.text_user)
         self.assistant = self.tokenizer.convert_tokens_to_ids(self.text_assistant)
 
-        self.reasoning_start = self.tokenizer.convert_tokens_to_ids(self.text_reasoning_start)
-        self.reasoning_end = self.tokenizer.convert_tokens_to_ids(self.text_reasoning_end)
+        self.think_start = self.tokenizer.convert_tokens_to_ids(self.text_think_start)
+        self.think_end = self.tokenizer.convert_tokens_to_ids(self.text_think_end)
 
         self.answer_start = self.tokenizer.convert_tokens_to_ids(self.text_answer_start)
         self.answer_end = self.tokenizer.convert_tokens_to_ids(self.text_answer_end)
@@ -140,9 +140,9 @@ class Tokenizer:
                 {"role":"user", "content":"hello?"},
                 {"role":"assistant", "content":"hello"},
                 {"role":"user", "content":"hello hello?"},
-                {"role":"assistant", "reasoning":"thinking", "content":"hello hello"},
+                {"role":"assistant", "think":"thinking", "content":"hello hello"},
             ]
-            <system>{system_prompt}</s><user>hello?</s><assistant>hello</s><user>hello hello?</s><assistant><reasoning>thinking</reasoning><answer>hello hello</answer></s>
+            <system>{system_prompt}</s><user>hello?</s><assistant>hello</s><user>hello hello?</s><assistant><think>thinking</think><answer>hello hello</answer></s>
         """
 
         chat_template = ''
@@ -154,8 +154,8 @@ class Tokenizer:
                 if add_answer_tag_for_assistant and role == 'assistant':
                     content = f"{self.text_answer_start}{content}{self.text_answer_end}"
 
-                if 'reasoning' in conversation:
-                    content = f"{self.text_reasoning_start}{conversation['reasoning']}{self.text_reasoning_end}{content}"
+                if 'think' in conversation:
+                    content = f"{self.text_think_start}{conversation['think']}{self.text_think_end}{content}"
 
                 chat_template = f"{chat_template}{support_roles[role]}{content}{self.text_end}"
 

llm_trainer/tools.py

@@ -3,7 +3,6 @@ from abc import ABC, abstractmethod
 import torch
 from .tokenizer import Tokenizer
 from .parallel_ds import DsParallel
-from .parallel_fsdp import FsdpParallel
 from .parallel_ddp import DdpParallel
 from .parallel_none import NoneParallel
 from .log import log
@@ -11,7 +10,6 @@ from .log import log
 
 parallel_types = {
     'ds': DsParallel,
-    'fsdp': FsdpParallel,
     'ddp': DdpParallel,
     'none': NoneParallel
 }

llm_trainer/train_configs.py

@@ -1,8 +1,7 @@
-from typing import Optional, Union, Set, Type, Callable, List, Mapping, Any
+from typing import Optional, Union, Callable, List, Mapping, Any
 from dataclasses import dataclass, field
 
 import torch
-from torch import nn
 from llm_model import ModelConfig, VLMConfig
 from .tools import FileDataset
 
@@ -33,6 +32,9 @@ class DsZeROConfig:
     reduce_bucket_size: Optional[Union[str, int]] = 5e8
     contiguous_gradients: Optional[bool] = True
 
+@dataclass(kw_only=True)
+class DsZero0Config(DsZeROConfig):
+    stage: int = field(default=0, init=False)
 
 @dataclass(kw_only=True)
 class DsZero1Config(DsZeROConfig):
@@ -84,26 +86,6 @@ class DsConfig:
     activation_checkpointing: Optional[DsActivationCheckpointingConfig] = None
 
 
-@dataclass(kw_only=True)
-class FsdpConfig:
-    """
-        fsdp训练模式配置项
-        Args:
-            transformer_layer_cls (`Set[Type[nn.Module]]`, *optional*, default is None):
-                提供transformer层的类
-            wrap_policy_num_params (`int`, *optional*, default is -1):
-                size_based_auto_wrap_policy的min_num_params参数，-1不生效该策略
-            cpu_offload (`bool`, *optional*, default is False):
-                是否使用cpu卸载
-            offload_params (`bool`, default is False):
-                是否卸载参数，在cpu_offload为True时生效
-    """
-    transformer_layer_cls: Optional[Set[Type[nn.Module]]] = None
-    wrap_policy_num_params: int = -1
-    cpu_offload: bool = False
-    offload_params: bool = False
-
-
 @dataclass(kw_only=True)
 class DataLoaderConfig:
     """
@@ -157,6 +139,7 @@ class GRPOConfig:
     clip_eps: float = 0.2
     kl_weight: float = 0.01
     group_size: int = 12
+    mixup_alpha: float = 1.0
     gen_max_new_tokens: Optional[int] = None
     gen_temperature: Optional[float] = None
     gen_k: Optional[int] = None
@@ -210,8 +193,6 @@ class TrainConfig:
                 每隔多少个batch进行模型eval
             lr_config (`LrConfig`):
                 lr配置项
-            fsdp_config: (`FsdpConfig`):
-                fsdp训练模式配置项
             data_loader_config: (`DataLoaderConfig`):
                 data loader配置项
             kd_config: (`KDConfig`, *Optional*, default is None):
@@ -231,7 +212,6 @@ class TrainConfig:
     lr_config: LrConfig = field(default_factory=LrConfig)
 
     ds_config: DsConfig = field(default_factory=DsConfig)
-    fsdp_config: FsdpConfig = field(default_factory=FsdpConfig)
 
     kd_config: Optional[KDConfig] = None
     dpo_config: Optional[DPOConfig] = None

llm_trainer/trainer.py

@@ -1,21 +1,18 @@
 import time
 from contextlib import nullcontext
-from typing import Optional, Tuple, List, Dict, Any, Union
+from typing import Optional, Tuple, List, Dict, Any
 
 import torch
-from torch import nn
 import torch.distributed as dist
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.utils.data import Dataset
 from llm_model import LlmModel, VlmModel
 
 from .parallel_ds import DsParallel
-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 .partition_utils import unwrap_model_for_generation
 
 from .train_configs import (
     TrainConfig,
@@ -78,7 +75,6 @@ class Trainer:
 
         self.train_model, self.optimizer = self._init_train_model_and_optim(initial_lr, parallel_kwargs, use_ds_optim)
         self.lr_scheduler = self._init_lr_scheduler(initial_lr)
-        self.eval_model: Optional[nn.Module] = self._init_eval_model()
 
         self.criterion, self.kd_loss = self._init_loss()
 
@@ -86,9 +82,7 @@ class Trainer:
             device_type=TrainerTools().parallel.device_type,
             dtype=TrainerTools().dtype,
             enabled=True,
-            # fsdp模式，需要将cache_enabled设置为false
-            # https://www.zhihu.com/question/642793891
-            cache_enabled=False if isinstance(self.train_model, FSDP) else None
+            cache_enabled=None
         ) if TrainerTools().use_amp else nullcontext()
 
         load_checkpoint(
@@ -176,12 +170,6 @@ class Trainer:
 
         return model, optim
 
-    def _init_eval_model(self) -> Optional[nn.Module]:
-        if TrainerTools().parallel.is_main_process:
-            return self._new_model(self.train_config).to(device='cpu', dtype=TrainerTools().dtype)
-
-        return None
-
     def _init_lr_scheduler(self, initial_lr: float) -> LRScheduler:
         if self.train_config.lr_config.enable_lr_scheduler:
             min_lr = self.train_config.lr_config.min_lr
@@ -313,13 +301,6 @@ class Trainer:
                     activation_checkpointing['number_checkpoints'] = activation_checkpointing_config.number_checkpoints
 
                 parallel_kwargs['activation_checkpointing'] = activation_checkpointing
-        elif isinstance(TrainerTools().parallel, FsdpParallel) and self.train_config.fsdp_config:
-            parallel_kwargs = {
-                'transformer_layer_cls': self.train_config.fsdp_config.transformer_layer_cls,
-                'wrap_policy_num_params': self.train_config.fsdp_config.wrap_policy_num_params,
-                'cpu_offload': self.train_config.fsdp_config.cpu_offload,
-                'offload_params': self.train_config.fsdp_config.offload_params
-            }
 
         dataloader_args = self.train_config.data_loader_config
         data_loader_kwargs = {
@@ -441,54 +422,35 @@ class Trainer:
 
         raise e
 
-    def _on_batch_end(
-            self,
-            tag: str
-    ):
-        copy_model_params(_from=self.train_model, _to=self.eval_model)
+    def _eval(self, tag: str):
+        with unwrap_model_for_generation(self.train_model) as generate_model:
+            if TrainerTools().parallel.is_main_process:
+                generate_model.eval()
+                eval_prompt, eval_image_tag = self._get_eval_data()
+
+                if isinstance(self.train_config, VLMConfig) and self.pixel_values_provider and eval_image_tag:
+                    eval_pixel_values = self.pixel_values_provider([eval_image_tag])
+                else:
+                    eval_pixel_values = None
+
+                submit_gen_task(
+                    generate_model,
+                    self.train_config.eval_config,
+                    tag=tag,
+                    prompt=eval_prompt,
+                    pixel_values=eval_pixel_values,
+                    max_position_embeddings=self.train_config.model_config.max_position_embeddings,
+                    tokens_per_image=self.tokens_per_image
+                )
+                generate_model.train()
 
-        if TrainerTools().parallel.is_main_process:
-            eval_prompt, eval_image_tag = self._get_eval_data()
-            if isinstance(self.train_config, VLMConfig) and self.pixel_values_provider and eval_image_tag:
-                eval_pixel_values = self.pixel_values_provider([eval_image_tag])
-            else:
-                eval_pixel_values = None
-
-            submit_gen_task(
-                self.eval_model,
-                self.train_config.eval_config,
-                tag=f'sign:batch/{tag}',
-                prompt=eval_prompt,
-                pixel_values=eval_pixel_values,
-                max_position_embeddings=self.train_config.model_config.max_position_embeddings,
-                tokens_per_image=self.tokens_per_image
-            )
         TrainerTools().parallel.wait()
 
-    def _on_epoch_end(
-            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_config, VLMConfig) and self.pixel_values_provider and eval_image_tag:
-                eval_pixel_values = self.pixel_values_provider([eval_image_tag])
-            else:
-                eval_pixel_values = None
-
-            submit_gen_task(
-                self.eval_model,
-                self.train_config.eval_config,
-                tag=f'sign:epoch/{tag}',
-                prompt=eval_prompt,
-                pixel_values=eval_pixel_values,
-                max_position_embeddings=self.train_config.model_config.max_position_embeddings,
-                tokens_per_image=self.tokens_per_image
-            )
+    def _on_batch_end(self, tag: str):
+        self._eval(f'sign:batch/{tag}')
 
-        TrainerTools().parallel.wait()
+    def _on_epoch_end(self, tag: str):
+        self._eval(f'sign:epoch/{tag}')
 
     def _on_file_start(
             self,
@@ -574,7 +536,6 @@ class Trainer:
                         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_tensor, dist.ReduceOp.AVG)
 

llm_trainer/utils.py

@@ -1,5 +1,4 @@
 import random
-from typing import Tuple, Optional
 import torch
 from torch.nn.utils.rnn import pad_sequence
 import torch.nn.functional as F

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

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

project_llm_trainer-0.5.1.dist-info/RECORD

@@ -0,0 +1,33 @@
+llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
+llm_trainer/checkpoint.py,sha256=xTmmQSJ_jQDVSTT3km1p_8eRrc7yE_dEsi92z9OX5ec,3251
+llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
+llm_trainer/dpo_trainer.py,sha256=wMREatLt0I8Ajdm_sI2U8Zj-IN1L6txP9s_tH1oI3-s,11431
+llm_trainer/ds_checkpoint.py,sha256=wz48HoLBBt8QGO1tXfvJwrXoiGtPG_gjwHfEqARllso,2175
+llm_trainer/eval.py,sha256=fjASCILU3fSPJxo9cP3rIXEEnkc5ZlUyHqXlZtUiHrw,888
+llm_trainer/generate_utils.py,sha256=CbJ3mfAD6DkQ0GUHcJQ1AK02m-ocwmd-BPXEpiwvNNQ,14933
+llm_trainer/grpo_trainer.py,sha256=qiC3KwxYPSB9UKqyk4eSRvORP3b6GM-2ozqI8u3QvI0,15568
+llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
+llm_trainer/loss.py,sha256=NZCQeUXnLSj__mmDflE8g89KgE0emAJXIab0IERCLno,6023
+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_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,676
+llm_trainer/partition_utils.py,sha256=xzv8kwlbKp3dai2pBwX89gN5ymeHk1bGbTkGru5H-UM,5167
+llm_trainer/scheduler.py,sha256=Xz8HhwoRMjRe41sf_NHhpZfkTlEs0I2MYusvMY6hCVw,3531
+llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,1783
+llm_trainer/tokenizer.py,sha256=SSpgXtb0e1NtQqRW0gCq09TTZi47umggy-Fh5EMHKJg,6708
+llm_trainer/tools.py,sha256=yF17lp6oOfLe2XJeKDQ1juZcbv-6vFamJSLwEeArduA,2975
+llm_trainer/train_configs.py,sha256=m57W71SI5VCCU9aJ_nJkB-3AJrSGiNXmV28rdpuYmLg,7332
+llm_trainer/trainer.py,sha256=zTJVyY1cAjJdTkyXCOy2ZPVP18SOMLdWhD54Mz2JRe4,25314
+llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
+project_llm_trainer-0.5.1.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.1.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.1.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.1.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.1.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.1.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.1.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.1.dist-info/METADATA,sha256=x-Bobn0EH7wyKznJydUeVLK9sdIrkBmDYDbEpyG4pKc,195
+project_llm_trainer-0.5.1.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.1.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.1.dist-info/RECORD,,

llm_trainer/dcp.py

@@ -1,93 +0,0 @@
-import os
-from typing import Optional, Dict, Any
-from torch import nn
-from torch.optim import Optimizer
-import torch.distributed.checkpoint as dcp
-from torch.distributed.checkpoint.stateful import Stateful
-from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
-from torch.distributed.checkpoint.format_utils import dcp_to_torch_save, torch_save_to_dcp
-
-DEFAULT_CHECKPOINT_DIR = "checkpoint"
-
-class AppState(Stateful):
-    def __init__(self, model: nn.Module, optimizer: Optimizer):
-        self.model = model
-        self.optimizer = optimizer
-
-    def state_dict(self) -> Dict[str, Any]:
-        model_state_dict, optimizer_state_dict = get_state_dict(self.model, self.optimizer)
-        return {
-            'model_state_dict': model_state_dict,
-            'optim_state_dict': optimizer_state_dict
-        }
-
-    def load_state_dict(self, state_dict: Dict[str, Any]):
-        set_state_dict(
-            model=self.model,
-            optimizers=self.optimizer,
-            model_state_dict=state_dict['model_state_dict'],
-            optim_state_dict=state_dict['optim_state_dict']
-        )
-
-
-def save_dcp(
-        model: nn.Module,
-        optimizer: Optimizer,
-        suffix: Optional[str] = None
-):
-    checkpoint_id = os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR)
-    if suffix:
-        checkpoint_id = f"{checkpoint_id}_{suffix}"
-
-    state_dict = {'app': AppState(model, optimizer)}
-
-    # fs_storage_writer = dcp.FileSystemWriter(checkpoint_id, overwrite=True)
-    # dcp.save(state_dict=state_dict, storage_writer=fs_storage_writer)
-    dcp.save(state_dict=state_dict, checkpoint_id=checkpoint_id)
-
-
-def load_dcp(
-        model: nn.Module,
-        optimizer: Optional[Optimizer] = None,
-        suffix: Optional[str] = None
-):
-    checkpoint_id = os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR)
-    if suffix:
-        checkpoint_id = f"{checkpoint_id}_{suffix}"
-
-    if os.path.exists(checkpoint_id):
-        state_dict = {'app': AppState(model, optimizer)}
-        # AppState帮助加载到state_dict中, 然后加载到model中
-        dcp.load(state_dict=state_dict, checkpoint_id=checkpoint_id)
-
-        # if isinstance(model, FSDP):
-        #     state_dict = {'app': AppState(model, optimizer)}
-        #     # AppState帮助加载到state_dict中, 然后加载到model中
-        #     dcp.load(state_dict=state_dict, checkpoint_id=checkpoint_id)
-        # else:
-        #     state_dict = {"model_state_dict": model.state_dict()}
-        #
-        #     if optimizer:
-        #         state_dict.update({'optim_state_dict': optimizer.state_dict()})
-        #
-        #     # since no progress group is initialized, DCP will disable any collectives.
-        #     # 加载到state_dict中，然后通过model.load_state_dict加载到model中
-        #     dcp.load(
-        #         state_dict=state_dict,
-        #         checkpoint_id=checkpoint_id,
-        #     )
-        #
-        #     model.load_state_dict(state_dict["model_state_dict"])
-        #     if optimizer:
-        #         optimizer.load_state_dict(state_dict["optim_state_dict"])
-
-def convert_dcp_to_pth(pth_path: str):
-    dcp_path = os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR)
-    if os.path.exists(dcp_path):
-        # convert dcp model to torch.save (assumes checkpoint was generated as above)
-        dcp_to_torch_save(dcp_path, pth_path)
-
-def convert_pth_to_dcp(pth_path: str):
-    if os.path.exists(pth_path):
-        # converts the torch.save model back to DCP
-        torch_save_to_dcp(pth_path, os.environ.get('DIST_CHECKPOINT_DIR', DEFAULT_CHECKPOINT_DIR))

llm_trainer/ds_model_params.py

@@ -1,72 +0,0 @@
-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

@@ -1,52 +0,0 @@
-import os
-from typing import Optional, Union, Tuple
-import torch
-from torch import nn
-from torch.optim import Optimizer
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-
-from .tools import TrainerTools
-
-DEFAULT_CHECKPOINT_NAME = "checkpoint.pth"
-
-def save_fsdp_checkpoint(
-        model: nn.Module,
-        optimizer: Optional[Optimizer] = None,
-        suffix: Optional[str] = None
-):
-    # 未经过测试 参考：https://doc.hfai.high-flyer.cn/haiscale/haiscale_fsdp.html
-    # 是否使用rank0_only=True？
-    with FSDP.summon_full_params(
-            module=model,
-            rank0_only=True,
-            writeback=False,
-            offload_to_cpu=True
-    ):
-        if TrainerTools().parallel.is_main_process:
-            checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
-            if suffix:
-                checkpoint_name = f"{checkpoint_name}_{suffix}"
-
-            ckpt = {'model_state_dict': model.state_dict()}
-            if optimizer:
-                ckpt.update({'optim_state_dict': optimizer.state_dict()})
-
-            torch.save(ckpt, checkpoint_name)
-
-
-def load_fsdp_checkpoint(
-        model: nn.Module,
-        optimizer: Optional[Optimizer] = None,
-        device: Optional[Union[torch.device, str]] = None,
-        suffix: Optional[str] = None
-):
-    checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
-    if suffix:
-        checkpoint_name = f"{checkpoint_name}_{suffix}"
-
-    with FSDP.summon_full_params(module=model):
-        state_dict = torch.load(checkpoint_name, weights_only=True, map_location=device)
-        model.load_state_dict(state_dict['model_state_dict'])
-
-        if optimizer:
-            optimizer.load_state_dict(state_dict['optim_state_dict'])

llm_trainer/fsdp_model_params.py

@@ -1,39 +0,0 @@
-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/model_params.py

@@ -1,28 +0,0 @@
-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)