Add norm_grad_sample for RMSNorm (#755)

Summary:
Pull Request resolved: #755

Llama3 model has RMSNorm and currently, we use functorch to support FGC for RMSNorm. This causes FSDP to rely on the root node for all_gather call of the RMSNorm layers. Adding norm_grad_sample method for RMSNorm to support layer-wise FSDP for this layer.

Reviewed By: HuanyuZhang

Differential Revision: D74334633

fbshipit-source-id: 399e98422066a8b9a5c5794538810c8f0d1c2de2

Author: aparna-aketi

opacus/grad_sample/__init__.py

@@ -27,6 +27,7 @@
 from .gsm_exp_weights import GradSampleModuleExpandedWeights
 from .gsm_no_op import GradSampleModuleNoOp
 from .instance_norm import compute_instance_norm_grad_sample  # noqa
+from .rms_norm import compute_rms_norm_grad_sample  # noqa
 from .layer_norm import compute_layer_norm_grad_sample  # noqa
 from .linear import compute_linear_grad_sample  # noqa
 from .utils import (

opacus/grad_sample/rms_norm.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# 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.
+
+
+from typing import Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from opacus.utils.tensor_utils import sum_over_all_but_batch_and_last_n
+
+from .utils import register_grad_sampler
+
+
+@register_grad_sampler(nn.RMSNorm)
+def compute_rms_norm_grad_sample(
+    layer: nn.RMSNorm,
+    activations: List[torch.Tensor],
+    backprops: torch.Tensor,
+) -> Dict[nn.Parameter, torch.Tensor]:
+    """
+    Computes per sample gradients for RMSNorm
+
+    Args:
+        layer: Layer
+        activations: Activations
+        backprops: Backpropagations
+    """
+    activations = activations[0]
+    ret = {}
+    if layer.weight.requires_grad:
+        ret[layer.weight] = sum_over_all_but_batch_and_last_n(
+            F.rms_norm(activations, layer.normalized_shape, eps=layer.eps) * backprops,
+            layer.weight.dim(),
+        )
+    return ret

opacus/tests/grad_samples/rms_norm_test.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# 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 hypothesis.strategies as st
+import torch
+import torch.nn as nn
+from hypothesis import given, settings
+
+from .common import GradSampleHooks_test
+
+
+class RMSNorm_test(GradSampleHooks_test):
+    @given(
+        N=st.integers(1, 4),
+        Z=st.integers(1, 4),
+        H=st.integers(1, 3),
+        W=st.integers(5, 10),
+        input_dim=st.integers(2, 4),
+        norm_dim=st.integers(1, 3),
+    )
+    @settings(deadline=60000)
+    def test_input_norm(
+        self, N: int, Z: int, W: int, H: int, input_dim: int, norm_dim: int
+    ):
+        if norm_dim >= input_dim:
+            return
+        normalized_shape, x_shape = self.get_x_shape_and_norm_shape(
+            H, N, W, Z, input_dim, norm_dim
+        )
+
+        norm = nn.RMSNorm(normalized_shape, elementwise_affine=True)
+        x = torch.randn(x_shape)
+        self.run_test(x, norm, batch_first=True, ew_compatible=False)
+
+    @staticmethod
+    def get_x_shape_and_norm_shape(H, N, W, Z, input_dim, norm_dim):
+        if norm_dim == 1:
+            normalized_shape = W
+            if input_dim == 2:
+                x_shape = [N, W]
+            if input_dim == 3:
+                x_shape = [N, Z, W]
+            if input_dim == 4:
+                x_shape = [N, Z, H, W]
+        elif norm_dim == 2:
+            if input_dim == 3:
+                normalized_shape = [Z, W]
+                x_shape = [N, Z, W]
+            if input_dim == 4:
+                normalized_shape = [H, W]
+                x_shape = [N, Z, H, W]
+        elif norm_dim == 3:
+            normalized_shape = [Z, H, W]
+            x_shape = [N, Z, H, W]
+        return normalized_shape, x_shape