Repository: mosaicml/composer

Issue:
Trainer crashes when using FP16 + DeepSpeed
```
Traceback (most recent call last):
  File "/home/mosaicml/composer/examples/run_composer_trainer.py", line 66, in <module>
    main()
  File "/home/mosaicml/composer/examples/run_composer_trainer.py", line 62, in main
    trainer.fit()
  File "/home/mosaicml/composer/composer/trainer/trainer.py", line 979, in fit
    self._train_loop()
  File "/home/mosaicml/composer/composer/trainer/trainer.py", line 1118, in _train_loop
    total_loss = self._train_batch(use_grad_scaling)
  File "/home/mosaicml/composer/composer/trainer/trainer.py", line 1217, in _train_batch
    total_loss = self._train_microbatches(microbatches)
  File "/home/mosaicml/composer/composer/trainer/trainer.py", line 1295, in _train_microbatches
    self._train_microbatch(use_grad_scaling, current_batch_size, total_loss, is_final_microbatch)
  File "/home/mosaicml/composer/composer/trainer/trainer.py", line 1333, in _train_microbatch
    with get_precision_context(self.state.precision):
  File "/usr/lib/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/home/mosaicml/composer/composer/core/precision.py", line 60, in get_precision_context
    raise ValueError(f"Unsupported precision: {precision}")
ValueError: Unsupported precision: Precision.FP16
```


Discussion:
This looks to have been introduced in https://github.com/mosaicml/composer/pull/832 which refactored the `precision_context` manager code, but did not catch the `FP16` case.

Resolving patch:
diff --git a/composer/core/precision.py b/composer/core/precision.py
--- a/composer/core/precision.py
+++ b/composer/core/precision.py
@@ -35,6 +35,12 @@ class Precision(StringEnum):
 def get_precision_context(precision: Union[str, Precision]) -> Generator[None, None, None]:
     """Returns a context manager to automatically cast to a specific precision.
 
+    .. warning::
+
+        :attr:`.Precision.FP16` is only supported when using DeepSpeed, as PyTorch does not
+        natively support this precision. When this function is invoked with :attr:`.Precision.FP16`,
+        the precision context will be a no-op.
+
     Args:
         precision (str or Precision): Precision for the context
     """
@@ -47,6 +53,11 @@ def get_precision_context(precision: Union[str, Precision]) -> Generator[None, N
         else:
             # Yield here to avoid warnings about cuda not being available
             yield
+    elif precision == Precision.FP16:
+        # No-op if FP16. FP16 is only supported by DeepSpeed, which is configured via the `deepspeed_config`
+        # DeepSpeed ignores `get_precision_context`. The Trainer init validates that Precision.FP16 is used
+        # only when using DeepSpeed.
+        yield
     elif precision == Precision.AMP:
         # Retain compatibility with PyTorch < 1.10
         with torch.cuda.amp.autocast(True):
diff --git a/composer/trainer/trainer.py b/composer/trainer/trainer.py
--- a/composer/trainer/trainer.py
+++ b/composer/trainer/trainer.py
@@ -725,6 +725,9 @@ def __init__(
         if isinstance(precision, str):
             precision = Precision(precision)
 
+        if not self.deepspeed_enabled and precision == Precision.FP16:
+            raise ValueError("FP16 precision is only supported when training with DeepSpeed.")
+
         # optimizers and schedulers
         if not optimizers:
             optimizers = DecoupledSGDW(list(model.parameters()), lr=0.1)
diff --git a/composer/trainer/trainer_hparams.py b/composer/trainer/trainer_hparams.py
--- a/composer/trainer/trainer_hparams.py
+++ b/composer/trainer/trainer_hparams.py
@@ -450,9 +450,6 @@ def validate(self):
             if isinstance(self.device, CPUDeviceHparams):
                 raise ValueError("Training on CPUs is not supported with DeepSpeed.")
 
-        elif self.precision == Precision.FP16:
-            raise ValueError("FP16 precision is only supported when training with DeepSpeed.")
-
         world_size = dist.get_world_size()
 
         if self.train_batch_size % world_size != 0:


Test patch:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -1,5 +1,6 @@
 # Copyright 2021 MosaicML. All Rights Reserved.
 
+import contextlib
 import os
 import pathlib
 from copy import deepcopy
@@ -40,6 +41,34 @@ def config(self, rank_zero_seed: int):
             'seed': rank_zero_seed,
         }
 
+    @pytest.mark.gpu
+    @pytest.mark.parametrize("precision", list(Precision))
+    def test_precision(self, config, precision: Precision):
+        config['precision'] = precision
+        config['device'] = 'gpu'
+
+        if precision == Precision.BF16:
+            pytest.importorskip("torch", minversion="1.10", reason="BF16 precision requires PyTorch 1.10+")
+
+        with pytest.raises(ValueError) if precision == Precision.FP16 else contextlib.nullcontext():
+            Trainer(**config)
+
+    @pytest.mark.gpu
+    @pytest.mark.parametrize("precision", list(Precision))
+    def test_trainer_with_deepspeed(self, config, precision: Precision):
+        config['deepspeed_config'] = {}
+        config['precision'] = precision
+        config['device'] = 'gpu'
+
+        if precision == Precision.BF16:
+            pytest.importorskip("torch", minversion="1.10", reason="BF16 precision requires PyTorch 1.10+")
+
+        trainer = Trainer(**config)
+
+        assert trainer.deepspeed_enabled
+
+        trainer.fit()
+
     def test_init(self, config):
         trainer = Trainer(**config)
         assert isinstance(trainer, Trainer)