project-llm-trainer 0.4.14__py3-none-any.whl → 0.5.0__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -6,35 +6,11 @@ from torch.optim import Optimizer
 from torch.nn.parallel import DistributedDataParallel as DDP
 
 from .parallel_ds import DsParallel
-from .parallel_fsdp import FsdpParallel
-from .parallel_ddp import DdpParallel
 from .scheduler import LRScheduler
 from .tools import TrainerTools
 
-try:
-    from .dcp import save_dcp, load_dcp, convert_dcp_to_pth
-except:
-    os.environ['ENABLE_DCP'] = "0"
-
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-
-# https://pytorch.org/tutorials/recipes/distributed_checkpoint_recipe.html
-
 DEFAULT_CHECKPOINT_NAME = "checkpoint.pth"
 
-
-def _can_use_dcp(model: nn.Module) -> bool:
-    if os.environ.get('ENABLE_DCP', '1') != '1':
-        return False
-
-    # 如果是fsdp或者ddp，才能使用dcp保存
-    if (isinstance(TrainerTools().parallel, FsdpParallel)
-            or isinstance(TrainerTools().parallel, DdpParallel)):
-        return True
-
-    return False
-
-
 def save_checkpoint(
         model: nn.Module,
         optimizer: Optional[Optimizer] = None,
@@ -43,11 +19,6 @@ def save_checkpoint(
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import save_ds_checkpoint
         save_ds_checkpoint(model, suffix)
-    elif _can_use_dcp(model):
-        save_dcp(model, optimizer, suffix)
-    elif isinstance(model, FSDP):
-        from .fsdp_checkpoint import save_fsdp_checkpoint
-        save_fsdp_checkpoint(model, optimizer, suffix)
     else:
         if TrainerTools().parallel.is_main_process:
             checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
@@ -73,11 +44,6 @@ def load_checkpoint(
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import load_ds_checkpoint
         load_ds_checkpoint(model, load_module_only=load_module_only, suffix=suffix)
-    elif _can_use_dcp(model):
-        load_dcp(model, optimizer, suffix)
-    elif isinstance(model, FSDP):
-        from .fsdp_checkpoint import load_fsdp_checkpoint
-        load_fsdp_checkpoint(model, optimizer, device, suffix)
     else:
         checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
         if suffix:
@@ -99,49 +65,10 @@ def load_checkpoint_for_eval(
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import load_ds_checkpoint_for_eval
         load_ds_checkpoint_for_eval(model)
-    elif _can_use_dcp(model):
-        checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
-
-        # load_dcp方式在cpu上会报错，所以改为先将ckpt转换为pth，然后再加载pth
-        # load_dcp(model, optimizer)
-        pth_name = os.environ.get('EVAL_CHECKPOINT_NAME', checkpoint_name)
-        if suffix:
-            pth_name = f'{pth_name}_{suffix}'
-
-        convert_dcp_to_pth(pth_name)
-
-        if os.path.exists(pth_name):
-            ckpt = torch.load(pth_name, map_location=device, weights_only=True)
-            model.load_state_dict(ckpt['app']['model_state_dict'])
-            # 使用完删除
-            os.remove(pth_name)
     else:
         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, only_rank0=_to is None)
-    elif isinstance(TrainerTools().parallel, FsdpParallel):
-        from .fsdp_checkpoint 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)
-
-
 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 .partition_utils import sync_model_params
 
 from .checkpoint import (
     save_checkpoint,
-    copy_model_params,
     save_steps,
 )
 
+
 class DPOTrainer(Trainer):
     def __init__(
             self,
@@ -37,7 +38,6 @@ class DPOTrainer(Trainer):
 
     def _init_reference_model(self):
         reference_model = self._new_model(self.train_config)
-        copy_model_params(_from=self.train_model, _to=reference_model)
 
         reference_model, _ = TrainerTools().parallel.process(
             model=reference_model,
@@ -50,6 +50,11 @@ class DPOTrainer(Trainer):
         for param in reference_model.parameters():
             param.requires_grad = False
 
+        sync_model_params(
+            _from=self.train_model,
+            _to=reference_model
+        )
+
         return reference_model
 
     def _init_loss(self):
@@ -209,7 +214,6 @@ class DPOTrainer(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/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, 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/eval.py

@@ -5,16 +5,14 @@ from .log import get_log_dir
 from .tools import TrainerTools
 from .train_configs import EvalConfig
 
-
-def _eval_task(
+def submit_gen_task(
         eval_model: torch.nn.Module,
         eval_config: EvalConfig,
         tag,
         prompt,
         pixel_values,
         max_position_embeddings,
-        tokens_per_image,
-        device
+        tokens_per_image
 ):
     log_dir = get_log_dir()
 
@@ -28,33 +26,8 @@ def _eval_task(
         p=eval_config.top_p,
         pixel_values=pixel_values,
         tokens_per_image=tokens_per_image,
-        device=device
+        device=TrainerTools().parallel.device
     )
 
     with open(f'{log_dir}gen.txt', 'a') as f:
         f.write(f"{tag}, gen->{gen_result}\n")
-
-
-def submit_gen_task(
-        eval_model: torch.nn.Module,
-        eval_config: EvalConfig,
-        tag,
-        prompt,
-        pixel_values,
-        max_position_embeddings,
-        tokens_per_image
-):
-    eval_model.to(TrainerTools().parallel.device)
-    _eval_task(
-        eval_model=eval_model,
-        eval_config=eval_config,
-        tag=tag,
-        prompt=prompt,
-        pixel_values=pixel_values,
-        max_position_embeddings=max_position_embeddings,
-        tokens_per_image=tokens_per_image,
-        device=TrainerTools().parallel.device
-    )
-    eval_model.to('cpu')
-
-    # threading.Thread(target=_eval_task, args=args).start()

llm_trainer/generate_utils.py

@@ -1,7 +1,6 @@
 from typing import Union, Optional, List
 from contextlib import nullcontext
 import torch
-from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from llm_model import VlmModel, KVCache
 from .tools import TrainerTools
 from .utils import batch_repeat_image_tok
@@ -131,8 +130,7 @@ def _generate(
         device_type=device,
         dtype=TrainerTools().dtype,
         enabled=True,
-        # fsdp模式，需要将cache_enabled设置为false
-        cache_enabled=False if isinstance(model, FSDP) else None
+        cache_enabled=None
     ) if TrainerTools().use_amp else nullcontext()
 
     if isinstance(model, VlmModel):
@@ -165,7 +163,6 @@ def _generate(
             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):
             # 是否需要截取？？
@@ -386,7 +383,7 @@ def batch_generate(
         device_type=device,
         dtype=TrainerTools().dtype,
         enabled=True,
-        cache_enabled=False if isinstance(model, FSDP) else None
+        cache_enabled=None
     ) if TrainerTools().use_amp else nullcontext()
 
     if isinstance(model, VlmModel):
@@ -403,7 +400,6 @@ def batch_generate(
     end_token = TrainerTools().tokenizer.end
     done = torch.zeros(batch_size, dtype=torch.bool, device=device)
 
-    model.eval()
     with torch.inference_mode():
         for _ in range(max_new_tokens):
             # 只处理未完成的样本

llm_trainer/grpo_trainer.py

@@ -1,5 +1,4 @@
 import time
-import copy
 from typing import Tuple, List, Union, Callable, Optional
 import torch
 from torch.utils.data import Dataset
@@ -16,9 +15,13 @@ from .tools import TrainerTools
 from .generate_utils import batch_generate
 from .log import log
 
+from .partition_utils import (
+    sync_model_params,
+    unwrap_model_for_generation
+)
+
 from .checkpoint import (
     save_checkpoint,
-    copy_model_params,
     save_steps,
 )
 
@@ -39,7 +42,6 @@ class GRPOTrainer(Trainer):
 
         self.reward_func = reward_func
         self.reference_model = self._init_reference_model()
-        self.generate_model = self._init_generate_model()
 
         # 默认使用torch提供的pad_sequence
         # 如果pad_sequence不支持padding_side参数，则将改参数置为False，使用反转的方式
@@ -47,17 +49,20 @@ class GRPOTrainer(Trainer):
 
     def _init_reference_model(self):
         reference_model = self._new_model(self.train_config)
-        reference_model.to('cpu')
-        reference_model.eval()
 
+        reference_model, _ = TrainerTools().parallel.process(
+            model=reference_model,
+            optimizer=None,
+            kwargs=self._init_reference_args(),
+            save_instance=False
+        )
+
+        reference_model.eval()
         for param in reference_model.parameters():
             param.requires_grad = False
 
         return reference_model
 
-    def _init_generate_model(self):
-        return copy.deepcopy(self.reference_model)
-
     def _init_loss(self):
         criterion = GRPOLoss(
             clip_eps=self.train_config.grpo_config.clip_eps,
@@ -163,7 +168,7 @@ class GRPOTrainer(Trainer):
         # [batch*group_size, 1]
         return advantages.unsqueeze(1)  # Add dimension for token-wise operations
 
-    def _generate_completions(self, prompts, group_size: int):
+    def _generate_completions(self, model, prompts, group_size: int):
         pad_token_id = TrainerTools().tokenizer.pad
         device = TrainerTools().parallel.device
 
@@ -181,7 +186,7 @@ class GRPOTrainer(Trainer):
 
         # [batch*group_size, max_prompt_len+max_gen_len]
         outputs: torch.Tensor = batch_generate(
-            model=self.generate_model,
+            model=model,
             tokens=prompt_ids,
             pad_token_id=pad_token_id,
             attention_mask=prompt_masks,
@@ -201,7 +206,7 @@ class GRPOTrainer(Trainer):
 
         return prompt_ids, prompt_masks, completion_ids, completion_masks
 
-    def _generate_rollout_data(self, batch_data: List[dict]):
+    def _generate_rollout_data(self, generate_model, batch_data: List[dict]):
         prompts = [item["prompt"] for item in batch_data]
         answers = [item["answer"] for item in batch_data]
         group_size = self.train_config.grpo_config.group_size
@@ -210,13 +215,13 @@ class GRPOTrainer(Trainer):
         # 修复问题：Inference tensors cannot be saved for backward. To work around you can make a clone to get a normal
         with torch.no_grad():
         # with torch.inference_mode():
-            prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_completions(prompts, group_size)
+            prompt_ids, prompt_mask, completion_ids, completion_mask = self._generate_completions(generate_model, prompts, group_size)
             input_ids = torch.cat([prompt_ids, completion_ids], dim=1)
             attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
             logits_to_keep = completion_ids.shape[1]
 
             # Compute old_log_probs from the current model, with gradients disabled.
-            old_log_probs, _ = self._compute_log_probabilities(self.generate_model, input_ids, attention_mask, logits_to_keep)
+            old_log_probs, _ = self._compute_log_probabilities(generate_model, input_ids, attention_mask, logits_to_keep)
 
             # Compute ref_log_probs from the reference model, which remains static.
             ref_log_probs, _ = self._compute_log_probabilities(self.reference_model, input_ids, attention_mask, logits_to_keep)
@@ -275,12 +280,15 @@ class GRPOTrainer(Trainer):
     def train(self):
         global_steps = 0
         skipping_train = False
-        device = TrainerTools().parallel.device
         aux_loss_coef = self.train_config.loss_config.aux_loss_coef
 
         for epoch in range(self.train_config.n_epochs):
-            copy_model_params(_from=self.train_model, _to=self.reference_model)
-            self.train_model.train()
+            sync_model_params(
+                _from=self.train_model,
+                _to=self.reference_model,
+                mixup_alpha=self.train_config.grpo_config.mixup_alpha
+            )
+
             file_count = len(self.train_config.file_dataset)
 
             for file_idx in range(file_count):
@@ -307,22 +315,13 @@ class GRPOTrainer(Trainer):
                     skipping_train = False
 
                     # start generate
-                    # 使用单独的模型生成数据， 原因是在deepspeed并行训练时，使用train_model生成数据会卡死
-                    self.generate_model.to(device)
-                    self.reference_model.to(device)
-
                     if TrainerTools().parallel.is_main_process:
                         log(f'start generate for batch {batch}/{batch_count_per_file}')
 
                     # 生成数据
-                    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')
-                    self.reference_model.to('cpu')
+                    with unwrap_model_for_generation(self.train_model) as generate_model:
+                        rollout_data = self._generate_rollout_data(generate_model, batch_data)
+
                     torch.cuda.empty_cache()
                     # end generate
 

llm_trainer/loss.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import List, Optional
 import torch
 from torch import nn
 import torch.nn.functional as F

llm_trainer/partition_utils.py

@@ -0,0 +1,146 @@
+from typing import Optional
+from contextlib import contextmanager
+import itertools
+from packaging import version
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from .tools import TrainerTools
+from .parallel_ds import DsParallel
+from .parallel_ddp import DdpParallel
+
+
+@contextmanager
+def unwrap_model_for_generation(model: nn.Module):
+    """
+    Context manager to unwrap distributed or accelerated models for generation tasks.
+
+    Args:
+        model:
+            Model to be unwrapped.
+    Yields:
+        Unwrapped model.
+
+    Example:
+    ```python
+    with unwrap_model_for_generation(model, accelerator) as unwrapped_model:
+        generated_outputs = unwrapped_model.generate(input_ids)
+    ```
+    """
+    if isinstance(TrainerTools().parallel, DsParallel):
+        import deepspeed
+        assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+        if model.zero_optimization_stage() == 3:
+            with deepspeed.zero.GatheredParameters(model.parameters()):
+                _remove_hooks(model)
+                yield unwrap_model(model)
+                _add_hooks(model)
+        else:
+            yield unwrap_model(model)
+    elif isinstance(TrainerTools().parallel, DdpParallel):
+        yield unwrap_model(model)
+    else:
+        yield model
+
+
+def sync_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alpha: float = 1.0):
+    if isinstance(TrainerTools().parallel, DsParallel):
+        _sync_ds_model_params(_from, _to, mixup_alpha)
+    elif isinstance(TrainerTools().parallel, DdpParallel):
+        _sync_ddp_model_params(_from, _to, mixup_alpha)
+    else:
+        _copy_params(_from, _to, mixup_alpha)
+
+
+def unwrap_model(model) -> nn.Module:
+    try:
+        import deepspeed
+        if isinstance(model, deepspeed.DeepSpeedEngine):
+            return model.module
+    except: ...
+
+    if isinstance(model, DDP):
+        return model.module
+
+    return model
+
+
+def _copy_params(model, target_model, mixup_alpha):
+    for target_param, copy_param in zip(target_model.parameters(), model.parameters()):
+        target_param.data.mul_(1.0 - mixup_alpha).add_(copy_param.data, alpha=mixup_alpha)
+
+
+def _sync_ds_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alpha: float = 1.0):
+    import deepspeed
+    assert isinstance(_from, deepspeed.DeepSpeedEngine)
+
+    origin_from = unwrap_model(_from)
+
+    if _from.zero_optimization_stage() == 3:
+        with deepspeed.zero.GatheredParameters(list(origin_from.parameters()) + list(_to.parameters()), modifier_rank=0):
+            if TrainerTools().parallel.is_main_process:
+                _copy_params(origin_from, _to, mixup_alpha)
+    else:
+        _copy_params(origin_from, _to, mixup_alpha)
+
+
+def _sync_ddp_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alpha: float = 1.0):
+    assert isinstance(_from, DDP)
+
+    origin_from = unwrap_model(_from)
+    _copy_params(origin_from, _to, mixup_alpha)
+
+
+def _add_hooks(model: nn.Module) -> None:
+    """Adds the optimizer hooks from a DeepSpeed ZeRO-3 model."""
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+    if not hasattr(model, "optimizer"):  # before the first training step, the model has no optimizer
+        return
+    if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"):
+        optimizer_offload = model.optimizer.parameter_offload
+    elif model.optimizer is not None:
+        optimizer_offload = model.optimizer
+    else:
+        raise RuntimeError("The model optimizer is None, which is not yet supported.")
+    if version.parse(deepspeed.__version__) >= version.parse("0.16.4"):
+        # Account for renaming in https://github.com/deepspeedai/DeepSpeed/pull/6847
+        optimizer_offload._register_deepspeed_module(optimizer_offload.module)
+    else:
+        optimizer_offload._register_hooks_recursively(optimizer_offload.module)
+
+
+def _remove_hooks(model: nn.Module) -> None:
+    """Removes the optimizer hooks from a DeepSpeed ZeRO-3 model."""
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+    if not hasattr(model, "optimizer"):  # before the first training step, the model has no optimizer
+        return
+    if model.optimizer is not None and hasattr(model.optimizer, "parameter_offload"):
+        optimizer_offload = model.optimizer.parameter_offload
+    elif model.optimizer is not None:
+        optimizer_offload = model.optimizer
+    else:
+        raise RuntimeError("The model optimizer is None, which is not yet supported.")
+
+    for param in _iter_params(optimizer_offload.module, recurse=True):
+        param.ds_active_sub_modules.clear()
+
+    for hook in optimizer_offload.forward_hooks:
+        hook.remove()
+    for hook in optimizer_offload.backward_hooks:
+        hook.remove()
+
+    optimizer_offload.forward_hooks = []
+    optimizer_offload.backward_hooks = []
+
+
+def _iter_params(module, recurse=False):
+    return [param for _, param in _get_all_parameters(module, recurse)]
+
+
+def _get_all_parameters(sub_module, recurse=False):
+    return itertools.chain(sub_module.named_parameters(recurse=recurse), sub_module.ds_external_parameters())

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,19 +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 .eval import submit_gen_task
+from .partition_utils import unwrap_model_for_generation
 
 from .train_configs import (
     TrainConfig,
@@ -31,10 +30,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 +44,6 @@ from .log import(
     get_log_dir
 )
 
-from .eval import submit_gen_task
-
 class Trainer:
     def __init__(
             self,
@@ -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.14.dist-info → project_llm_trainer-0.5.0.dist-info}/METADATA

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

project_llm_trainer-0.5.0.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=2MoEGEpoTzx7khO3dPcC2akFLyjtbFFpdJtuB_QQ3OY,17708
+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=A7TYYUbtPf75kjCvWP7yBui4xZBObMk2aPem62YpwpY,6776
+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.0.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.0.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.0.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.0.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.0.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.0.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.0.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.0.dist-info/METADATA,sha256=YDj-N4VL8O_AqNanwfU6Yt38J97p3RgtUSzmwl0Y-GM,195
+project_llm_trainer-0.5.0.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.0.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.0.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/fsdp_checkpoint.py

@@ -1,87 +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
-import torch.distributed as dist
-
-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'])
-
-
-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/parallel_fsdp.py

@@ -1,121 +0,0 @@
-from typing import Optional, Tuple
-import functools
-import torch
-from torch import nn
-from torch.distributed.fsdp import (
-    FullyShardedDataParallel as FSDP,
-    MixedPrecision,
-    ShardingStrategy,
-    BackwardPrefetch,
-    CPUOffload,
-)
-
-from torch.distributed.fsdp.wrap import (
-    size_based_auto_wrap_policy,
-    transformer_auto_wrap_policy,
-    always_wrap_policy,
-    enable_wrap,
-    wrap,
-)
-
-from .parallel import Parallel
-
-class FsdpParallel(Parallel):
-    def __init__(self):
-        super().__init__()
-
-    def process(
-            self,
-            model: nn.Module,
-            optimizer: torch.optim.Optimizer,
-            kwargs: Optional[dict] = None,
-            save_instance: bool = True
-    ) -> Tuple[nn.Module, torch.optim.Optimizer]:
-        """
-            :param model:
-            :param optimizer:
-            :param kwargs:
-                "wrap_policy_num_params" int size_based_auto_wrap_policy的最小参数量
-                "cpu_offload" bool 是否使用cpu卸载
-                "offload_params" bool 是否卸载参数，在cpu_offload为True时生效
-            :param save_instance
-            :return:
-        """
-
-        model.to(self.device)
-
-        if self._use_compile:
-            model = torch.compile(model)
-
-        if self._use_parallel:
-            if 'transformer_layer_cls' in kwargs:
-                auto_wrap_policy = functools.partial(
-                    transformer_auto_wrap_policy,
-                    transformer_layer_cls=kwargs['transformer_layer_cls']
-                )
-            elif 'wrap_policy_num_params' in kwargs:
-                auto_wrap_policy = functools.partial(
-                    size_based_auto_wrap_policy,
-                    min_num_params=kwargs['wrap_policy_num_params']
-                )
-            else:
-                auto_wrap_policy = None
-
-            if 'cpu_offload' in kwargs:
-                offload_params = False
-                if 'offload_params' in kwargs:
-                    offload_params = kwargs['offload_params']
-
-                # 选择配置 cpu_offload，以便在计算中不使用包装参数时将这些参数卸载到 CPU。
-                # 这可以进一步提高内存效率，但代价是主机和设备之间的数据传输开销。
-                cpu_offload = CPUOffload(offload_params=offload_params)
-            else:
-                cpu_offload = None
-
-            if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-                mixed_precision = MixedPrecision(
-                    param_dtype=torch.bfloat16,
-                    # Gradient communication precision.
-                    reduce_dtype=torch.bfloat16,
-                    # Buffer precision.
-                    buffer_dtype=torch.bfloat16,
-                )
-            else:
-                mixed_precision = None
-
-            raw_model = model
-
-            # device_mesh = init_device_mesh("cuda", (self.world_size,))
-            # model = FSDP(
-            #     model,
-            #     auto_wrap_policy=auto_wrap_policy,
-            #     mixed_precision=mixed_precision,
-            #     cpu_offload=cpu_offload,
-            #     device_id=torch.cuda.current_device(),
-            #     device_mesh=device_mesh
-            # )
-
-            model = FSDP(
-                model,
-                sharding_strategy=ShardingStrategy.FULL_SHARD,
-                auto_wrap_policy=auto_wrap_policy,
-                mixed_precision=mixed_precision,
-                cpu_offload=cpu_offload,
-                device_id=torch.cuda.current_device(),
-                process_group=None,
-                # use_orig_params=True,
-                # backward_prefetch=BackwardPrefetch.BACKWARD_PRE,  # bit faster async comms, bit higher memory
-                # limit_all_gathers=False,
-                # forward_prefetch=True,
-            )
-        else:
-            model = model
-            raw_model = model
-
-        if save_instance:
-            self.raw_model = raw_model
-            self.model = model
-
-        return model, optimizer
-
-

project_llm_trainer-0.4.14.dist-info/RECORD

@@ -1,35 +0,0 @@
-llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=ItDzuXVikk-0gWSw-IS7SrODEdlJEb5nZs15dBFkPdk,5793
-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=fprJlbSgtyKmmpytyMOZBs3pcjZA13SeWao0llnLpNQ,4962
-llm_trainer/eval.py,sha256=NDm8PbXLch7xT81xPYPRCNrcrB_Xj5GDJSCxyVwUOp4,1524
-llm_trainer/fsdp_checkpoint.py,sha256=dAHIGHfuvTA6OC0jV9Ls-oD4ZR9CPGa31mjtoh-2dZE,3229
-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
-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=HKzH3nfMT1-SW4Htwa0KqYtMd6FAJcthR5IEo6di8us,8168
-llm_trainer/trainer.py,sha256=j5fDqMzvU6SYwxHsv9wX0UVX4JXS-8eP1AkHgVxKf9U,27119
-llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-project_llm_trainer-0.4.14.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.4.14.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.4.14.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.4.14.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.4.14.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.4.14.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.4.14.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.4.14.dist-info/METADATA,sha256=VMEWVv8pBqFUAhIAiH4_S4ECUHln31gchHLhTtUAM1o,196
-project_llm_trainer-0.4.14.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.4.14.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.4.14.dist-info/RECORD,,