Add a context manager for activation sharding.

fairscale/nn/model_parallel/utils.py

@@ -23,6 +23,9 @@
 from typing import Tuple
 
 import torch
+from torch.autograd.graph import saved_tensors_hooks
+from .initialize import get_model_parallel_rank, get_model_parallel_world_size
+from .mappings import get_model_parallel_group
 
 
 def ensure_divisibility(numerator: int, denominator: int) -> None:
@@ -76,3 +79,44 @@ def vocab_range_from_per_partition_vocab_size(
     def vocab_range_from_global_vocab_size(global_vocab_size: int, rank: int, world_size: int) -> Tuple[int, int]:
         per_partition_vocab_size = divide_and_check_no_remainder(global_vocab_size, world_size)
         return VocabUtility.vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank, world_size)
+
+
+def _pack_over_mp(tensor):
+    mp_world_size = get_model_parallel_world_size()
+    if mp_world_size == 1:
+        return tensor # no-op for mp=1
+    full_tensor_shape = list(tensor.shape)
+    shard = tensor.view(-1).chunk(mp_world_size, dim=0)[get_model_parallel_rank()]
+    shard = shard.detach().clone().contiguous()  # clone to explicitly release memory of the full tensor
+    del tensor
+    return shard, full_tensor_shape
+
+
+def _unpack_over_mp(sharded_tensor):
+    sharded_tensor, full_tensor_shape = sharded_tensor
+    mp_world_size = get_model_parallel_world_size()
+    if mp_world_size == 1:
+        return sharded_tensor # no-op for mp=1
+    full_tensor = torch.empty(
+        *full_tensor_shape,
+        dtype=sharded_tensor.dtype,
+        device=sharded_tensor.device)
+
+    torch.distributed.all_gather_into_tensor(
+        full_tensor.view(-1), sharded_tensor, group=get_model_parallel_group()
+    )
+
+    return full_tensor
+
+
+class shard_over_mp_group(saved_tensors_hooks):
+    """Context manager for activatoin sharding.
+
+    This context manager shard tensors saved by autograd over the
+    model parallel group in the forward pass and unshards them
+    in the backward pass. Useful to remove redundancy in the
+    long-living activation tensors.
+    """
+
+    def __init__(self):
+        super().__init__(_pack_over_mp, _unpack_over_mp)