Add per-sample gradient norm computation as a functionality

Summary:
Per-sample gradient norm is computed for Ghost Clipping, but it can be useful generally. Exposed it as a functionality.


```
...

loss.backward()
per_sample_norms  = model.per_sample_gradient_norms

```

Differential Revision: D68634969

Author: EnayatUllah

opacus/grad_sample/grad_sample_module_fast_gradient_clipping.py

@@ -18,8 +18,6 @@
 import logging
 from typing import List
 
-import torch
-import torch.nn as nn
 from opacus.grad_sample.functorch import ft_compute_per_sample_gradient
 from opacus.grad_sample.grad_sample_module import (
     GradSampleModule,
@@ -28,6 +26,9 @@
 )
 from opacus.utils.module_utils import requires_grad, trainable_parameters
 
+import torch
+import torch.nn as nn
+
 
 logger = logging.getLogger(__name__)
 logger.disabled = True
@@ -120,6 +121,7 @@ def __init__(
         self.trainable_parameters = [p for _, p in trainable_parameters(self._module)]
         self.max_grad_norm = max_grad_norm
         self.use_ghost_clipping = use_ghost_clipping
+        self._per_sample_gradient_norms = None
 
     def get_clipping_coef(self) -> torch.Tensor:
         """Get per-example gradient scaling factor for clipping."""
@@ -131,6 +133,7 @@ def get_norm_sample(self) -> torch.Tensor:
         norm_sample = torch.stack(
             [param._norm_sample for param in self.trainable_parameters], dim=0
         ).norm(2, dim=0)
+        self.per_sample_gradient_norms = norm_sample
         return norm_sample
 
     def capture_activations_hook(
@@ -231,3 +234,16 @@ def capture_backprops_hook(
         if len(module.activations) == 0:
             if hasattr(module, "max_batch_len"):
                 del module.max_batch_len
+
+    @property
+    def per_sample_gradient_norms(self) -> torch.Tensor:
+        if self._per_sample_gradient_norms is not None:
+            return self._per_sample_gradient_norms
+        else:
+            raise AttributeError(
+                "per_sample_gradient_norms is not set. Please call forward and backward on the model before accessing this property."
+            )
+
+    @per_sample_gradient_norms.setter
+    def per_sample_gradient_norms(self, value):
+        self._per_sample_gradient_norms = value

opacus/optimizers/optimizer_fast_gradient_clipping.py

@@ -146,6 +146,7 @@ def zero_grad(self, set_to_none: bool = False):
 
         for p in self.params:
             p.grad_sample = None
+            p._per_sample_gradient_norms = None
 
             if not self._is_last_step_skipped:
                 p.summed_grad = None