Skip to content

Commit

Permalink
add bfloat16 to inference support dtypes
Browse files Browse the repository at this point in the history
to allow running inference tasks using bfloat16
  • Loading branch information
nelyahu committed Sep 12, 2024
1 parent 170b46e commit 4e74b8d
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion deepspeed/inference/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -324,7 +324,7 @@ def _validate_args(self, mpu, replace_with_kernel_inject):
if self._config.checkpoint is not None and not isinstance(self._config.checkpoint, (str, dict)):
raise ValueError(f"checkpoint must be None, str or dict, got {type(self._config.checkpoint)}")

supported_dtypes = [None, torch.half, torch.int8, torch.float]
supported_dtypes = [None, torch.half, torch.int8, torch.float, torch.bfloat16]
if self._config.dtype not in supported_dtypes:
raise ValueError(f"{self._config.dtype} not supported, valid dtype: {supported_dtypes}")

Expand Down

0 comments on commit 4e74b8d

Please sign in to comment.