Support mixed and low precision training (#764)

Summary:

We add support for mixed and low precision training in Opacus. 

Mixed precision training is supported iwth "hooks", "ghost" grad_sample_modes. 

Low-precision trainig is additionally supported with "functorch", "ew"

Why there is no functorch support for mixed precision trainig: The backward pass with functorch performs both a forward and backward pass to compute per-sample gradients. The forrwad pass happens outside of the torch.amp context, so it cannot handle mixed precision. 

Support for low and mixed precision training is GPU dependent.

Differential Revision: D72415906

Author: iden-kalemaj

opacus/grad_sample/conv.py

@@ -40,6 +40,10 @@ def compute_conv_grad_sample(
         backprops: Backpropagations
     """
     activations = activations[0]
+
+    if activations.dtype != backprops.dtype:
+        activations = activations.to(backprops.dtype)
+
     n = activations.shape[0]
     if n == 0:
         # Empty batch

opacus/grad_sample/embedding.py

@@ -51,7 +51,10 @@ def compute_embedding_grad_sample(
             .reshape(batch_size, -1, layer.embedding_dim)
         )
         grad_sample = torch.zeros(
-            batch_size, *layer.weight.shape, device=layer.weight.device
+            batch_size,
+            *layer.weight.shape,
+            device=layer.weight.device,
+            dtype=backprops.dtype
         )
         grad_sample.scatter_add_(
             1, index, backprops.reshape(batch_size, -1, layer.embedding_dim)
@@ -65,7 +68,13 @@ def compute_embedding_grad_sample(
 def compute_embeddingbag_gradsampler(layer, inputs, backprops):
     index, offset = inputs
     batch_size = offset.shape[0]
-    gsm = torch.zeros(batch_size, layer.num_embeddings, layer.embedding_dim)
+    gsm = torch.zeros(
+        batch_size,
+        layer.num_embeddings,
+        layer.embedding_dim,
+        device=index.device,
+        dtype=backprops.dtype,
+    )
 
     for i in range(batch_size):
         begin = offset[i]

opacus/grad_sample/embedding_norm_sample.py

@@ -131,15 +131,18 @@ def compute_embedding_norm_sample(
     # Sum gradients over new index positions and compute squared gradient norms
     num_unique_paired_indices = unique_paired_indices.size(0)
     summed_gradients = torch.zeros(
-        num_unique_paired_indices, grad_values.size(-1), device=device
+        num_unique_paired_indices,
+        grad_values.size(-1),
+        device=device,
+        dtype=grad_values.dtype,
     )
     summed_gradients = summed_gradients.index_add(
         0, new_index_positions.to(device), grad_values
     )
     sqr_gradient_sum = torch.sum(summed_gradients**2, dim=1)
 
     # Scatter add the squared sums back to their respective rows
-    result = torch.zeros(nrows, device=device)
+    result = torch.zeros(nrows, device=device, dtype=grad_values.dtype)
     unique_batch_ids = unique_paired_indices[:, 0].to(device)
     result.scatter_add_(0, unique_batch_ids, sqr_gradient_sum)
 

opacus/grad_sample/functorch.py

@@ -105,9 +105,10 @@ def ft_compute_per_sample_gradient(layer, activations, backprops):
     if not hasattr(layer, "ft_compute_sample_grad"):
         prepare_layer(layer)
 
-    per_sample_grads = layer.ft_compute_sample_grad(
-        parameters, activations[0], backprops
-    )
+    activations = activations[0]
+    if activations.dtype != backprops.dtype and activations.is_floating_point():
+        activations = activations.to(backprops.dtype)
+    per_sample_grads = layer.ft_compute_sample_grad(parameters, activations, backprops)
 
     ret = {}
     for i_p, p in enumerate(parameters):

opacus/grad_sample/linear.py

@@ -39,6 +39,9 @@ def compute_linear_grad_sample(
         backprops: Backpropagations
     """
     activations = activations[0]
+    if activations.dtype != backprops.dtype:
+        activations = activations.to(backprops.dtype)
+
     ret = {}
     if layer.weight.requires_grad:
         gs = torch.einsum("n...i,n...j->nij", backprops, activations)
@@ -61,6 +64,9 @@ def compute_linear_norm_sample(
         backprops: Backpropagations
     """
     activations = activations[0]
+    if activations.dtype != backprops.dtype:
+        activations = activations.to(backprops.dtype)
+
     ret = {}
 
     if backprops.dim() == 2:

opacus/optimizers/optimizer.py

@@ -134,7 +134,7 @@ def _generate_noise(
         `2^53` (easy to break) but with `n=2`, we get `2^159`, which is hard
         enough for an attacker to break.
     """
-    zeros = torch.zeros(reference.shape, device=reference.device)
+    zeros = torch.zeros(reference.shape, device=reference.device, dtype=reference.dtype)
     if std == 0:
         return zeros
     # TODO: handle device transfers: generator and reference tensor
@@ -164,6 +164,7 @@ def _generate_noise(
             size=reference.shape,
             device=reference.device,
             generator=generator,
+            dtype=reference.dtype,
         )
 
 
@@ -450,6 +451,14 @@ def clip_and_accumulate(self):
         for p in self.params:
             _check_processed_flag(p.grad_sample)
             grad_sample = self._get_flat_grad_sample(p)
+
+            # gradients should match the dtype of the optimizer parameters
+            # for mixed precision, optimizer parameters are usually in FP32
+            # lower precision grads will be cast up to FP32
+            if grad_sample.dtype != p.dtype:
+                grad_sample = grad_sample.to(p.dtype)
+            if per_sample_clip_factor.dtype != p.dtype:
+                per_sample_clip_factor = per_sample_clip_factor.to(p.dtype)
             grad = torch.einsum("i,i...", per_sample_clip_factor, grad_sample)
 
             if p.summed_grad is not None: