Apply suggestions from code review

Co-authored-by: Aryan <aryan@huggingface.co>
This commit is contained in:
Sayak Paul
2024-12-26 08:00:04 +05:30
committed by GitHub
parent 0c0a0a5e4a
commit 0309f3bf67
+3 -3
View File
@@ -340,7 +340,7 @@ class Trainer:
pin_memory=self.args.pin_memory, pin_memory=self.args.pin_memory,
) )
def sort_out_weight_dtype(self, accelerator): def _get_training_dtype(self, accelerator):
weight_dtype = torch.float32 weight_dtype = torch.float32
if accelerator.state.deepspeed_plugin: if accelerator.state.deepspeed_plugin:
# DeepSpeed is handling precision, use what's in the DeepSpeed config # DeepSpeed is handling precision, use what's in the DeepSpeed config
@@ -392,8 +392,8 @@ class Trainer:
# TODO(aryan): handle torch dtype from accelerator vs model dtype; refactor # TODO(aryan): handle torch dtype from accelerator vs model dtype; refactor
self.state.weight_dtype = weight_dtype self.state.weight_dtype = weight_dtype
if self.args.mixed_precision != _INVERSE_DTYPE_MAP[weight_dtype]: if self.args.mixed_precision != _INVERSE_DTYPE_MAP[weight_dtype]:
logger.info( logger.warning(
f"`mixed_precision` was set to {_INVERSE_DTYPE_MAP[weight_dtype]} which different from what was initially passed ({self.args.mixed_precision})." f"`mixed_precision` was set to {_INVERSE_DTYPE_MAP[weight_dtype]} which is different from configured argument ({self.args.mixed_precision})."
) )
self.args.mixed_precision = _INVERSE_DTYPE_MAP[weight_dtype] self.args.mixed_precision = _INVERSE_DTYPE_MAP[weight_dtype]
self.transformer.to(dtype=weight_dtype) self.transformer.to(dtype=weight_dtype)