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

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.15.dist-info/RECORD

@@ -1,38 +0,0 @@
-llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-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=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=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=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
-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=95ARdNDfalhZ7Ug-fDj3qIhWEiZQeX9n5WANhijIRLE,27140
-llm_trainer/utils.py,sha256=-ivhMF0d999va13S1wt2uBvtVw8Nvr3uBzhaUFKL04Q,6826
-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,,