Repository: facebookresearch/fairscale

Issue:
FSDP: excessive GPU memory due to extra process groups
We have found this issue related to process group and its associated NCCL buffers. Fix is coming.


Discussion:

Resolving patch:
diff --git a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
--- a/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
+++ b/fairscale/nn/data_parallel/fully_sharded_data_parallel.py
@@ -24,7 +24,12 @@
 from fairscale.nn.misc import FlattenParamsWrapper
 from fairscale.nn.wrap import auto_wrap, default_auto_wrap_policy, enable_wrap
 from fairscale.utils.containers import apply_to_tensors
-from fairscale.utils.parallel import chunk_and_pad, enable_pytorch_sync_bn, validate_process_group
+from fairscale.utils.parallel import (
+    chunk_and_pad,
+    enable_pytorch_sync_bn,
+    get_process_group_cached,
+    validate_process_group,
+)
 from fairscale.utils.params import broadcast_object, calc_grad_norm, recursive_copy_to_device
 from fairscale.utils.reduce_scatter_bucketer import ReduceScatterBucketer
 from fairscale.utils.state_dict import replace_by_prefix_
@@ -233,7 +238,7 @@ def __init__(
     ):
         init_start = time.time()
         super().__init__()
-        self.process_group = process_group or dist.new_group()
+        self.process_group = process_group or get_process_group_cached()
         self.rank = self.process_group.rank()
         self.world_size = self.process_group.size()
         self.reshard_after_forward = reshard_after_forward
@@ -1816,13 +1821,11 @@ def wrap_bn_only_policy(module: nn.Module, recurse: bool, unwrapped_params: int)
                 module, tuple(default_auto_wrap_policy.EXCLUDE_WRAP_MODULES)  # type: ignore
             )
 
-    pg = None
+    pg = process_group
     if single_rank_pg:
         # No sharding with this single member group.
         my_rank = dist.get_rank()
-        pg = dist.new_group(ranks=[my_rank])
-    else:
-        pg = process_group
+        pg = get_process_group_cached(ranks=[my_rank])
 
     if fsdp_config is None:
         fsdp_config = {
diff --git a/fairscale/utils/parallel.py b/fairscale/utils/parallel.py
--- a/fairscale/utils/parallel.py
+++ b/fairscale/utils/parallel.py
@@ -5,7 +5,7 @@
 
 """Useful functions for parallel training."""
 
-from typing import List
+from typing import List, Optional
 
 import torch
 import torch.distributed as dist
@@ -58,18 +58,44 @@ def enable_pytorch_sync_bn(module: torch.nn.Module) -> None:
             layer._specify_ddp_gpu_num(1)  # type: ignore
 
 
-def get_global_group() -> None:
+def get_process_group_cached(ranks: Optional[List[int]] = None) -> ProcessGroup:
     """
-    Singleton PyTorch distributed group.
-    Inspired by https://github.com/pytorch/fairseq
+    Singleton PyTorch distributed group cache. Inspired by the code from fairseq.
 
-    For FSDP, it is important to use a global group, otherwise, inner FSDP instances
-    will not share the gradient reduction bucket buffer with the root instance, end up using
-    more GPU memory.
+    For FSDP, it is important to use the same group between outer and inner FSDP instances,
+    otherwise, inner FSDP instances will not share the gradient reduction bucket buffer with
+    the root instance. This will result in increased GPU memory utilization.
+
+    Each separate process group also uses separate NCCL library instances, which will have
+    a significant effect on GPU memory use if too many process groups are created and used.
+    Setting NCCL_BUFFSIZE=102400 env variable is a useful technique to check if the NCCL
+    memory is causing GPU OOM. Note, the NCCL buffers are not allocated
+    through the PyTorch caching allocator, therefore, you may see GPU OOM even when
+    torch.cuda.reserved_memory() is still way below the total amount of GPU memory.
+
+    Extra process groups can also reduce training speed (observed on VISSL models).
+
+    Args:
+        ranks (Optional[List[int]]):
+            Ranks requested in the target group. None for all ranks.
+            Default: None
+
+    Returns:
+        (ProcessGroup):
+            Return the requested process group. Throws RuntimeError if torch.distributed module is not yet initialized.
     """
-    if dist.is_initialized():
-        if not hasattr(get_global_group, "_global_group"):
-            get_global_group._global_group = dist.new_group()  # type: ignore
-        return get_global_group._global_group  # type: ignore
-    else:
-        return None
+    if not dist.is_initialized():
+        raise RuntimeError("torch.distributed is not yet initialized but process group is requested.")
+
+    if not hasattr(get_process_group_cached, "_global_group_cache"):
+        get_process_group_cached._global_group_cache = {}  # type: ignore
+    cache = get_process_group_cached._global_group_cache  # type: ignore
+
+    if ranks is None:
+        ranks = list(range(dist.get_world_size()))
+
+    ranks_set = frozenset(ranks)  # take care of ordering and duplicates in the ranks list.
+    if ranks_set not in cache:
+        cache[ranks_set] = dist.new_group(list(ranks_set))
+
+    return cache[ranks_set]


Test patch:
diff --git a/tests/nn/data_parallel/test_fsdp_memory.py b/tests/nn/data_parallel/test_fsdp_memory.py
--- a/tests/nn/data_parallel/test_fsdp_memory.py
+++ b/tests/nn/data_parallel/test_fsdp_memory.py
@@ -21,7 +21,7 @@
 from fairscale.nn import checkpoint_wrapper
 from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
 from fairscale.nn.data_parallel import auto_wrap_bn
-from fairscale.utils.parallel import get_global_group
+from fairscale.utils.parallel import get_process_group_cached
 from fairscale.utils.testing import (
     dist_init,
     dump_all_tensors,
@@ -33,7 +33,7 @@
 
 
 def to_fsdp(module, fsdp_config):
-    return FSDP(module, process_group=get_global_group(), **fsdp_config)
+    return FSDP(module, process_group=get_process_group_cached(), **fsdp_config)
 
 
 def get_cur_mem(rank, result, prefix):
@@ -199,78 +199,78 @@ def test_fsdp_memory(fsdp, ckpt):
             "iter 0: start": 3,
             "iter 0: after fwd": 340,
             "iter 0: after loss": 340,
-            "iter 0: after bwd": 66,
-            "iter 0: after step": 68,
+            "iter 0: after bwd": 16,
+            "iter 0: after step": 18,
             "iter 0: done": 5,
             "iter 1: start": 5,
             "iter 1: after fwd": 342,
             "iter 1: after loss": 342,
-            "iter 1: after bwd": 68,
-            "iter 1: after step": 68,
+            "iter 1: after bwd": 18,
+            "iter 1: after step": 18,
             "iter 1: done": 5,
             "iter 2: start": 5,
             "iter 2: after fwd": 342,
             "iter 2: after loss": 342,
-            "iter 2: after bwd": 68,
-            "iter 2: after step": 68,
+            "iter 2: after bwd": 18,
+            "iter 2: after step": 18,
             "iter 2: done": 5,
             "iter 3: start": 5,
             "iter 3: after fwd": 342,
             "iter 3: after loss": 342,
-            "iter 3: after bwd": 68,
-            "iter 3: after step": 68,
+            "iter 3: after bwd": 18,
+            "iter 3: after step": 18,
             "iter 3: done": 5,
         },
         ("fsdp_amp_default", "no_ckpt"): {
             "iter 0: start": 28,
             "iter 0: after fwd": 630,
             "iter 0: after loss": 630,
-            "iter 0: after bwd": 104,
-            "iter 0: after step": 131,
+            "iter 0: after bwd": 67,
+            "iter 0: after step": 93,
             "iter 0: done": 54,
             "iter 1: start": 54,
             "iter 1: after fwd": 657,
             "iter 1: after loss": 657,
-            "iter 1: after bwd": 131,
-            "iter 1: after step": 131,
+            "iter 1: after bwd": 93,
+            "iter 1: after step": 93,
             "iter 1: done": 54,
             "iter 2: start": 54,
             "iter 2: after fwd": 657,
             "iter 2: after loss": 657,
-            "iter 2: after bwd": 131,
-            "iter 2: after step": 131,
+            "iter 2: after bwd": 93,
+            "iter 2: after step": 93,
             "iter 2: done": 54,
             "iter 3: start": 54,
             "iter 3: after fwd": 657,
             "iter 3: after loss": 657,
-            "iter 3: after bwd": 131,
-            "iter 3: after step": 131,
+            "iter 3: after bwd": 93,
+            "iter 3: after step": 93,
             "iter 3: done": 54,
         },
         ("fsdp_amp_compute_dtype32", "no_ckpt"): {
             "iter 0: start": 28,
             "iter 0: after fwd": 657,
             "iter 0: after loss": 657,
-            "iter 0: after bwd": 117,
-            "iter 0: after step": 143,
+            "iter 0: after bwd": 67,
+            "iter 0: after step": 93,
             "iter 0: done": 54,
             "iter 1: start": 54,
             "iter 1: after fwd": 684,
             "iter 1: after loss": 684,
-            "iter 1: after bwd": 143,
-            "iter 1: after step": 143,
+            "iter 1: after bwd": 93,
+            "iter 1: after step": 93,
             "iter 1: done": 54,
             "iter 2: start": 54,
             "iter 2: after fwd": 684,
             "iter 2: after loss": 684,
-            "iter 2: after bwd": 143,
-            "iter 2: after step": 143,
+            "iter 2: after bwd": 93,
+            "iter 2: after step": 93,
             "iter 2: done": 54,
             "iter 3: start": 54,
             "iter 3: after fwd": 684,
             "iter 3: after loss": 684,
-            "iter 3: after bwd": 143,
-            "iter 3: after step": 143,
+            "iter 3: after bwd": 93,
+            "iter 3: after step": 93,
             "iter 3: done": 54,
         },
         ("ddp", "ckpt"): {
@@ -303,78 +303,78 @@ def test_fsdp_memory(fsdp, ckpt):
             "iter 0: start": 3,
             "iter 0: after fwd": 51,
             "iter 0: after loss": 51,
-            "iter 0: after bwd": 66,
-            "iter 0: after step": 68,
+            "iter 0: after bwd": 16,
+            "iter 0: after step": 18,
             "iter 0: done": 5,
             "iter 1: start": 5,
             "iter 1: after fwd": 53,
             "iter 1: after loss": 53,
-            "iter 1: after bwd": 68,
-            "iter 1: after step": 68,
+            "iter 1: after bwd": 18,
+            "iter 1: after step": 18,
             "iter 1: done": 5,
             "iter 2: start": 5,
             "iter 2: after fwd": 53,
             "iter 2: after loss": 53,
-            "iter 2: after bwd": 68,
-            "iter 2: after step": 68,
+            "iter 2: after bwd": 18,
+            "iter 2: after step": 18,
             "iter 2: done": 5,
             "iter 3: start": 5,
             "iter 3: after fwd": 53,
             "iter 3: after loss": 53,
-            "iter 3: after bwd": 68,
-            "iter 3: after step": 68,
+            "iter 3: after bwd": 18,
+            "iter 3: after step": 18,
             "iter 3: done": 5,
         },
         ("fsdp_amp_default", "ckpt"): {
             "iter 0: start": 28,
             "iter 0: after fwd": 52,
             "iter 0: after loss": 52,
-            "iter 0: after bwd": 104,
-            "iter 0: after step": 131,
+            "iter 0: after bwd": 67,
+            "iter 0: after step": 93,
             "iter 0: done": 54,
             "iter 1: start": 54,
             "iter 1: after fwd": 79,
             "iter 1: after loss": 79,
-            "iter 1: after bwd": 131,
-            "iter 1: after step": 131,
+            "iter 1: after bwd": 93,
+            "iter 1: after step": 93,
             "iter 1: done": 54,
             "iter 2: start": 54,
             "iter 2: after fwd": 79,
             "iter 2: after loss": 79,
-            "iter 2: after bwd": 131,
-            "iter 2: after step": 131,
+            "iter 2: after bwd": 93,
+            "iter 2: after step": 93,
             "iter 2: done": 54,
             "iter 3: start": 54,
             "iter 3: after fwd": 79,
             "iter 3: after loss": 79,
-            "iter 3: after bwd": 131,
-            "iter 3: after step": 131,
+            "iter 3: after bwd": 93,
+            "iter 3: after step": 93,
             "iter 3: done": 54,
         },
         ("fsdp_amp_compute_dtype32", "ckpt"): {
             "iter 0: start": 28,
             "iter 0: after fwd": 52,
             "iter 0: after loss": 52,
-            "iter 0: after bwd": 117,
-            "iter 0: after step": 143,
+            "iter 0: after bwd": 67,
+            "iter 0: after step": 93,
             "iter 0: done": 54,
             "iter 1: start": 54,
             "iter 1: after fwd": 79,
             "iter 1: after loss": 79,
-            "iter 1: after bwd": 143,
-            "iter 1: after step": 143,
+            "iter 1: after bwd": 93,
+            "iter 1: after step": 93,
             "iter 1: done": 54,
             "iter 2: start": 54,
             "iter 2: after fwd": 79,
             "iter 2: after loss": 79,
-            "iter 2: after bwd": 143,
-            "iter 2: after step": 143,
+            "iter 2: after bwd": 93,
+            "iter 2: after step": 93,
             "iter 2: done": 54,
             "iter 3: start": 54,
             "iter 3: after fwd": 79,
             "iter 3: after loss": 79,
-            "iter 3: after bwd": 143,
-            "iter 3: after step": 143,
+            "iter 3: after bwd": 93,
+            "iter 3: after step": 93,
             "iter 3: done": 54,
         },
     }[(fsdp, ckpt)]
diff --git a/tests/nn/moe/test_moe_layer.py b/tests/nn/moe/test_moe_layer.py
--- a/tests/nn/moe/test_moe_layer.py
+++ b/tests/nn/moe/test_moe_layer.py
@@ -62,7 +62,7 @@ def expert_params(device):
     expert = torch.nn.Linear(model_dim, model_dim)
     moe = MOELayer(gate, expert).to(device)
     for p in expert.parameters():
-        assert p.expert is True
+        assert p.expert is True, str(p.expert)
 
 
 @pg_test()
@@ -77,9 +77,9 @@ def forward(device):
     expert.weight = torch.nn.Parameter(torch.eye(model_dim))
     moe = MOELayer(gate, expert).to(device)
     output = moe(input)
-    assert output.shape == input.shape
+    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
     # Re-assembled output should match input due to identity expert.
-    assert torch.allclose(input, output)
+    torch.testing.assert_allclose(input, output)
 
 
 @pg_test()
@@ -99,11 +99,13 @@ def forward_multi(device):
         experts += [expert]
     moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
     output = moe(input)
-    assert output.shape == input.shape
+    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
     # 90% of the input should have gone to an expert
-    assert len(output.nonzero(as_tuple=False)) / output.numel() > 0.90
+    assert (
+        len(output.nonzero(as_tuple=False)) / output.numel() > 0.90
+    ), f"{len(output.nonzero(as_tuple=False))} / {output.numel()}"
     # Except for zeros, re-assembled output should match input due to identity expert.
-    assert torch.allclose(input, torch.where(output > 0, output, input))
+    torch.testing.assert_allclose(input, torch.where(output > 0, output, input))
 
 
 # Test Gate which round-robin routes tokens to experts
@@ -115,7 +117,7 @@ def __init__(self, model_dim, num_experts):
 
     def forward(self, input):
         s = input.shape[0]
-        assert s % self.num_experts == 0
+        assert s % self.num_experts == 0, f"{s} % {self.num_experts} != 0"
         capacity = 2 * s // self.num_experts
         output = torch.zeros(s, self.num_experts, capacity, dtype=input.dtype, device=input.device)
         for i in range(s):
@@ -136,12 +138,12 @@ def forward_routing(device):
     expert.weight = torch.nn.Parameter(torch.eye(model_dim) * scale)
     moe = MOELayer(gate, expert).to(device)
     output = moe(input)
-    assert output.shape == input.shape
+    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
     # Verify that each token was sent to the correct expert by checking its scale.
     t = input.shape[1]
     for i in range(t):
         expert = i % num_experts
-        assert torch.allclose(input[:, i] * (expert + 1), output[:, i])
+        torch.testing.assert_allclose(input[:, i] * (expert + 1), output[:, i])
 
 
 @pg_test()
@@ -161,12 +163,12 @@ def forward_routing_multi(device):
         experts += [expert]
     moe = MOELayer(gate, torch.nn.ModuleList(experts)).to(device)
     output = moe(input)
-    assert output.shape == input.shape
+    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
     # Verify that each token was sent to the correct expert by checking its scale.
     t = input.shape[1]
     for i in range(t):
         expert = i % num_experts
-        assert torch.allclose(input[:, i] * (expert + 1), output[:, i])
+        torch.testing.assert_allclose(input[:, i] * (expert + 1), output[:, i])
 
 
 @pg_test()
@@ -182,7 +184,7 @@ def backward(device):
     expert.weight = torch.nn.Parameter(torch.eye(model_dim))
     moe = MOELayer(gate, expert).to(device)
     output = moe(input)
-    assert output.shape == input.shape
+    assert output.shape == input.shape, f"{output.shape} != {input.shape}"
     output = loss(output, input)
     output.backward()
-    assert torch.allclose(expert.weight.grad, torch.zeros_like(expert.weight))
+    torch.testing.assert_allclose(expert.weight.grad, torch.zeros_like(expert.weight))