Repository: Lightning-AI/pytorch-lightning
Pull request: [FIX] ShardedGradScaler should only be set for FP16
## What does this PR do?

Closes #12914

When running with sharded training + bf16, I noticed that the ShardedGradScaler was being passed. This is incorrect as BF16 does not require any scaling logic. Fixed this and added a unit test!

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or minor internal changes/refactors)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?

Make sure you had fun coding 🙃


cc @borda @SeanNaren @awaelchli @rohitgr7 @akihironitta @carmocca @justusschock

Problem statement:
ShardedGradScaler should not be used in BF16 
## 🐛 Bug

The ShardedGradScaler is being passed in when the precision is set to BF16.

```bash
from pytorch_lightning.plugins import ShardedNativeMixedPrecisionPlugin

plugin = ShardedNativeMixedPrecisionPlugin(precision="bf16", device="cuda")

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/sean/pytorch-lightning/pytorch_lightning/plugins/precision/sharded_native_amp.py", line 38, in __init__
    super().__init__(precision, device, scaler=scaler or ShardedGradScaler())
  File "/home/sean/pytorch-lightning/pytorch_lightning/plugins/precision/native_amp.py", line 55, in __init__
    raise MisconfigurationException(f"`precision='bf16'` does not use a scaler, found {scaler}.")
pytorch_lightning.utilities.exceptions.MisconfigurationException: `precision='bf16'` does not use a scaler, found <fairscale.optim.grad_scaler.ShardedGradScaler object at 0x7f29cd024fd0>.
```


cc @SeanNaren @awaelchli @rohitgr7 @akihironitta

Hints from issue discussion:

Commit message: Add CHANGELOG.md

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9c4d8421ee95f..d91e142a7eed5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -147,7 +147,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed an issue to ensure all the checkpoint states are saved in a common filepath with `DeepspeedStrategy` ([#12887](https://github.com/PyTorchLightning/pytorch-lightning/pull/12887))
 
 
--
+- Fixed an issue where sharded grad scaler is passed in when using BF16 with the `ShardedStrategy` ([#12915](https://github.com/PyTorchLightning/pytorch-lightning/pull/12915))
 
 
 -
diff --git a/pytorch_lightning/plugins/precision/sharded_native_amp.py b/pytorch_lightning/plugins/precision/sharded_native_amp.py
index 57840a918a2e1..e40aea8ecf4eb 100644
--- a/pytorch_lightning/plugins/precision/sharded_native_amp.py
+++ b/pytorch_lightning/plugins/precision/sharded_native_amp.py
@@ -35,7 +35,7 @@ def __init__(
                 "You have asked for sharded AMP but you have not installed it."
                 " Install `fairscale` using this guide: https://https://github.com/facebookresearch/fairscale"
             )
-        super().__init__(precision, device, scaler=scaler or ShardedGradScaler())
+        super().__init__(precision, device, scaler=ShardedGradScaler() if scaler is None and precision == 16 else None)
 
     def clip_grad_by_norm(self, optimizer: "OSS", clip_val: Union[int, float]) -> None:
         optimizer.clip_grad_norm(clip_val)
diff --git a/tests/plugins/precision/test_sharded_precision.py b/tests/plugins/precision/test_sharded_precision.py
new file mode 100644
index 0000000000000..fb35cd5189f06
--- /dev/null
+++ b/tests/plugins/precision/test_sharded_precision.py
@@ -0,0 +1,42 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from pytorch_lightning.plugins import ShardedNativeMixedPrecisionPlugin
+from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE
+from tests.helpers.runif import RunIf
+
+ShardedGradScaler = None
+if _FAIRSCALE_AVAILABLE:
+    from fairscale.optim.grad_scaler import ShardedGradScaler
+
+
+@RunIf(fairscale=True)
+@pytest.mark.parametrize(
+    "precision,scaler,expected",
+    [
+        (16, torch.cuda.amp.GradScaler(), torch.cuda.amp.GradScaler),
+        (16, None, ShardedGradScaler),
+        ("bf16", None, None),
+        (32, None, None),
+    ],
+)
+def test_sharded_precision_scaler(precision, scaler, expected):
+    plugin = ShardedNativeMixedPrecisionPlugin(precision=precision, scaler=scaler, device="cuda")
+    if expected:
+        assert isinstance(plugin.scaler, expected)
+    else:
+        assert not plugin.scaler


Code review comments:
diff_hunk: @@ -0,0 +1,42 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from pytorch_lightning.plugins import ShardedNativeMixedPrecisionPlugin
+from pytorch_lightning.utilities import _FAIRSCALE_AVAILABLE
+from tests.helpers.runif import RunIf
+
+ShardedGradScaler = None
+if _FAIRSCALE_AVAILABLE:
+    from fairscale.optim.grad_scaler import ShardedGradScaler
+
+
+@RunIf(fairscale=True)
+@pytest.mark.parametrize(
+    "precision,scaler,expected",
+    [
+        (16, torch.cuda.amp.GradScaler(), torch.cuda.amp.GradScaler),
+        (16, None, ShardedGradScaler),
+        ("bf16", None, None),
path: tests/plugins/precision/test_sharded_precision.py
text: @user1:
As per the failing CI

```suggestion
        pytest.param("bf16", None, None, marks=RunIf(min_torch="1.10")),
```