class_docstr: ## Zero3 Layer
  
  Each layer of the model (or a combination of a few consecutive layers) should be wrapped in
  this module.
class_docstr_tok: ["#","#","Zero3","Layer","Each","layer","of","the","model","(","or","a","combination","of","a","few","consecutive","layers",")","should","be","wrapped","in","this","module","."]
docstr: ### Backup the gradients of the current layer
docstr_tok: ["#","#","#","Backup","the","gradients","of","the","current","layer"]
code: def _backup_grads(self):
      """
      ### Backup the gradients of the current layer
      """
      # Skip if there are no trainable parameters
      if self.chunk_size[self.TRAINING_PARAMS_IDX] == 0:
          return
  
      # Use the backup stream to backup the gradients
      with torch.cuda.stream(self.backup_stream):
          # Buffer to store the gradients
          buffer = self._empty(
              (
                  self.world_size
                  * self.chunk_size[self.TRAINING_PARAMS_IDX],
              )
          )
          # Split the continuous buffer into number of nodes. These splits are views of `buffer'.
          buffers = list(
              buffer.split(self.chunk_size[self.TRAINING_PARAMS_IDX])
          )
  
          # Offset of the continuous buffer
          offset = 0
          # Iterate through trainable parameters
          for p in self.param_refs[self.TRAINING_PARAMS_IDX]:
              # Collect gradients
              shape = p._orig_shape  # type: ignore[attr-defined]
              buffer[offset : offset + shape.numel()] = p.grad.view(-1)
              # Update the offset
              offset += shape.numel()
              # Clean the gradients
              p.grad = None
  
          # Empty tensor to accumulate the gradients of the current shard
          grad = self._empty(
              (self.chunk_size[self.TRAINING_PARAMS_IDX],)
          )
          # Accumulate the gradients of each shard. It scatters the buffers across the nodes,
          # and each node accumulates (reduces) the tensors it receives.
          dist.reduce_scatter(grad, buffers)
  
          # Wait for the operation to complete and then clear the references to the buffers
          for b in buffers:
              b.record_stream(self.fetch_stream)
          buffer.record_stream(self.fetch_stream)
          del buffer
          del buffers
  
          # Set the chunk gradients. This is what the optimizer sees.
          self.chunk[self.TRAINING_PARAMS_IDX].grad = grad
          del grad
  
imports: import functools
  from typing import List, Optional, Tuple
  import torch
  import torch.distributed
  from torch import nn
label_desc: Use image node_id 12 for calling the Zero3Layer obj's underlying member method code with example usage: obj._backup_grads() without return types