github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix

pull/6024/head
wangbluo 2024-08-22 03:00:40 +00:00
parent d77e66a577
commit 971b16a74f
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/booster/plugin/low_level_zero_plugin.py
View File

@ -63,7 +63,7 @@ class OptimizerParamCheckState(enum.Enum):
class LowLevelZeroModel(ModelWrapper, AMPModelMixin): class LowLevelZeroModel(ModelWrapper, AMPModelMixin):
def __init__( def __init__(
self, module: nn.Module, precision: str, overlap_allgather: bool = False, use_fp8: bool = False self, module: nn.Module, precision: str, overlap_allgather: bool = False, cast_inputs: bool = True, use_fp8: bool = False
) -> None: ) -> None:
super().__init__(module) super().__init__(module)
self.dtype = None self.dtype = None
@ -77,7 +77,7 @@ class LowLevelZeroModel(ModelWrapper, AMPModelMixin):
self.module = module self.module = module
self.convert_fn = None self.convert_fn = None
self.use_fp8 = use_fp8 self.use_fp8 = use_fp8
if self.dtype is not None: if self.dtype is not None and cast_inputs:
self.convert_fn = partial(_convert_floating_point, dtype=self.dtype) self.convert_fn = partial(_convert_floating_point, dtype=self.dtype)
self.overlap_allgather = overlap_allgather self.overlap_allgather = overlap_allgather
self.op_hooks = [] self.op_hooks = []
@ -342,6 +342,7 @@ class LowLevelZeroPlugin(DPPluginBase):
cpu_offload: bool = False, cpu_offload: bool = False,
master_weights: bool = True, master_weights: bool = True,
verbose: bool = False, verbose: bool = False,
cast_inputs: bool = True,
fp8_communication: bool = False, fp8_communication: bool = False,
use_fp8: bool = False, use_fp8: bool = False,
) -> None: ) -> None:
@ -372,6 +373,8 @@ class LowLevelZeroPlugin(DPPluginBase):
self.lora_enabled = False self.lora_enabled = False
self.verbose = verbose self.verbose = verbose
self.logger = get_dist_logger() self.logger = get_dist_logger()
self.cast_inputs = cast_inputs
self.use_fp8 = use_fp8 self.use_fp8 = use_fp8
# set class name with stage, for better error message # set class name with stage, for better error message
setattr(self.__class__, "__name__", f"LowLevelZeroPlugin_ZeRO-{stage}") setattr(self.__class__, "__name__", f"LowLevelZeroPlugin_ZeRO-{stage}")
@ -490,6 +493,7 @@ class LowLevelZeroPlugin(DPPluginBase):
model, model,
self.precision, self.precision,
overlap_allgather=self.zero_optim_kwargs["overlap_allgather"], overlap_allgather=self.zero_optim_kwargs["overlap_allgather"],
cast_inputs=self.cast_inputs,
use_fp8=self.use_fp8, use_fp8=self.use_fp8,
) )