Allow BF16 dtype support on CPU (pytorch#1218)

weifengpy · Jul 26, 2024 · e101420 · e101420
1 parent 651b18c
commit e101420
Show file tree

Hide file tree

Showing 5 changed files with 4 additions and 6 deletions.
diff --git a/docs/source/deep_dives/recipe_deepdive.rst b/docs/source/deep_dives/recipe_deepdive.rst
@@ -137,7 +137,7 @@ Initialize recipe state including seed, device, dtype, metric loggers, relevant
     def __init__(...):
 
         self._device = utils.get_device(device=params.device)
-        self._dtype = utils.get_dtype(dtype=params.dtype)
+        self._dtype = utils.get_dtype(dtype=params.dtype, device=self._device)
         ...
 
 Load checkpoint, update recipe state from checkpoint, initialize components and load state dicts from checkpoint

diff --git a/recipes/eleuther_eval.py b/recipes/eleuther_eval.py
@@ -182,7 +182,7 @@ def __init__(self, cfg: DictConfig) -> None:
 
     def setup(self) -> None:
         self._device = utils.get_device(device=self._cfg.device)
-        self._dtype = utils.get_dtype(dtype=self._cfg.dtype)
+        self._dtype = utils.get_dtype(dtype=self._cfg.dtype, device=self._device)
         self._limit = self._cfg.limit
         self._tasks = list(self._cfg.tasks)
         self._quantizer = config.instantiate(self._cfg.quantizer)

diff --git a/recipes/generate.py b/recipes/generate.py
@@ -36,7 +36,7 @@ class InferenceRecipe:
 
     def __init__(self, cfg: DictConfig) -> None:
         self._device = utils.get_device(device=cfg.device)
-        self._dtype = utils.get_dtype(dtype=cfg.dtype)
+        self._dtype = utils.get_dtype(dtype=cfg.dtype, device=self._device)
         self._quantizer = config.instantiate(cfg.quantizer)
         self._quantization_mode = utils.get_quantizer_mode(self._quantizer)
 

diff --git a/recipes/quantize.py b/recipes/quantize.py
@@ -69,7 +69,7 @@ class QuantizationRecipe:
 
     def __init__(self, cfg: DictConfig) -> None:
         self._device = utils.get_device(device=cfg.device)
-        self._dtype = utils.get_dtype(dtype=cfg.dtype)
+        self._dtype = utils.get_dtype(dtype=cfg.dtype, device=self._device)
         self._quantizer = config.instantiate(cfg.quantizer)
         self._quantization_mode = utils.get_quantizer_mode(self._quantizer)
         utils.set_seed(seed=cfg.seed)

diff --git a/torchtune/utils/precision.py b/torchtune/utils/precision.py
@@ -101,8 +101,6 @@ def get_dtype(
             f"Dtype {torch_dtype} must be one of {', '.join(list(PRECISION_STR_TO_DTYPE.keys()))} for finetuning."
         )
 
-    # TODO (rohan-varma): prefer to use get_default_device() here to figure out whether user is training on
-    # CPU or GPU, but it is not supported in versions of torch we test.
     if (
         torch_dtype == torch.bfloat16
         and device != torch.device("cpu")