Repository: facebookresearch/fairscale

Issue:
[FSDP] Frozen weights support
## 🐛 Bug: Frozen weight support in FSDP

This bug is a follow up to issue https://github.com/facebookresearch/fairscale/issues/610. I encountered a new test case in which a scenario of linear evaluation fails.

### Context (same as previous issue)

In the context of VISSL, we are training a SWAV model with FSDP and then want to evaluate this model using linear evaluation, again using FSDP to wrap our model.

In the context of linear evaluation, only the linear layer on top of the "trunk" will be learned: the representations of the trunk, fed to the linear layer, will stay constant, and so we set `requires_grad=False` on the parameters of the trunk.

### New symptom

The initial test case raised in https://github.com/facebookresearch/fairscale/issues/610 is green, but I discovered that nesting FSDP blocks in the trunk (the part that is frozen) is leading to a new assertion being raised:

`ERROR: expected to be in states [<TrainingState.BACKWARD_PRE: 3>] but current state is TrainingState.IDLE`

## Command

## To Reproduce

This script should help you reproduce the issue: simply run the following script with the options `-f -l` where:
- `-f` means "enable FSDP"
- `-l` means "freeze the trunk"

**Important notes**:
- The issue only appears if both `-f` and `-l` are used together
- The issue disappear if I comment the line `self.trunk = FSDP(self.trunk)` and so it seems that the issue is linked to the nesting of FSDP blocks when layers are frozen

```
import argparse

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
from fairscale.nn.data_parallel import FullyShardedDataParallel as FSDP
from torch.nn.parallel import DistributedDataParallel


class Model(nn.Module):
    def __init__(self, with_fsdp: bool):
        super().__init__()
        self.trunk = nn.Sequential(
            self._create_block(3, 64, with_fsdp),
            self._create_block(64, 64, with_fsdp),
        )
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
            nn.Flatten(),
            nn.Linear(64, 10),
        )
        if with_fsdp:
            self.trunk = FSDP(self.trunk)
            self.head = FSDP(self.head)

    def forward(self, x):
        return self.head(self.trunk(x))

    def _create_block(self, in_channels: int, out_channels: int, with_fsdp: bool):
        block = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3),
            nn.ReLU(inplace=True),
        )
        if with_fsdp:
            block = FSDP(block)
        return block


def create_model(with_fsdp: bool):
    return Model(with_fsdp)


def _distributed_worker(gpu_id: int, with_fsdp: bool, with_linear: bool):
    torch.cuda.set_device(gpu_id)
    dist.init_process_group(
        backend="nccl", init_method="tcp://127.0.0.1:9099", world_size=2, rank=gpu_id
    )

    torch.manual_seed(0)
    torch.backends.cudnn.deterministic = True
    batch = torch.randn(size=(2, 3, 224, 224)).cuda()

    model = create_model(with_fsdp)
    model = model.cuda()

    if with_linear:
        for name, param in model.named_parameters():
            if "trunk" in name:
                param.requires_grad = False

    if with_fsdp:
        model = FSDP(model)
    else:
        model = DistributedDataParallel(model, device_ids=[gpu_id])

    if gpu_id == 0:
        print(model)

    target = torch.LongTensor([0, 1]).cuda()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

    for iteration in range(3):
        out = model(batch)
        fake_loss = criterion(out, target)
        if gpu_id == 0:
            print("Loss", iteration, ":", fake_loss.item())
        optimizer.zero_grad()
        fake_loss.backward()
        optimizer.step()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--fsdp", action="store_const", const=True, default=False)
    parser.add_argument(
        "-l", "--linear", action="store_const", const=True, default=False
    )
    args = parser.parse_args()
    mp.spawn(_distributed_worker, (args.fsdp, args.linear), nprocs=2)
```

## Environment

I use pytorch 1.6.0+cu102, with the following modules enabled: `module load cuda/11.0 cudnn/v8.0.3.33-cuda.11.0 NCCL/2.8.3-1-cuda.11.0`.

My version of fairscale is 0.3.6.

## Additional context

CC: @prigoyal @min-xu-ai 



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
@@ -1301,15 +1301,19 @@ def _remove_shard_bwd_hook(fsdp_module: FullyShardedDataParallel) -> None:
             if isinstance(m, FullyShardedDataParallel):
                 _remove_shard_bwd_hook(m)
                 m._pre_backward_hook_has_run = False
-                if m._has_params:
-                    if any(p.requires_grad for p in m.params):
+                # Note: m.parameters() should not be an empty list. FSDP
+                # wrapping modules without weights is not tested at the moment.
+                if any(p.requires_grad for p in m.parameters()):
+                    if m._has_params:
                         m.assert_state(TrainingState.BACKWARD_POST)
                     else:
-                        # Unlikely case, should only happens if `m` has params but none of the
-                        # params has `requires_grad==True`.
-                        m.assert_state(TrainingState.IDLE)
+                        m.assert_state(TrainingState.BACKWARD_PRE)
                 else:
-                    m.assert_state(TrainingState.BACKWARD_PRE)
+                    # Unlikely case. When m and its children has no params
+                    # with `requires_grad==True`, then m's pre-backward and
+                    # post-backward hooks aren't called by autograd. Therefore,
+                    # it is in IDLE state.
+                    m.assert_state(TrainingState.IDLE)
                 m.training_state = TrainingState.IDLE
 
     @torch.no_grad()


Test patch:
diff --git a/tests/nn/data_parallel/test_fsdp_freezing_weights.py b/tests/nn/data_parallel/test_fsdp_freezing_weights.py
--- a/tests/nn/data_parallel/test_fsdp_freezing_weights.py
+++ b/tests/nn/data_parallel/test_fsdp_freezing_weights.py
@@ -10,6 +10,7 @@
 """ Test FSDP with some params frozen. """
 
 
+from enum import Enum
 import tempfile
 
 import pytest
@@ -38,16 +39,51 @@ def forward(self, x):
         return self.head(self.trunk(x))
 
 
-def _create_model(with_fsdp):
-    model = Model()
-    if with_fsdp:
-        model.trunk = FSDP(model.trunk)
-        model.head = FSDP(model.head)
+class NestedTrunkModel(nn.Module):
+    def __init__(self, with_fsdp):
+        super().__init__()
+        self.trunk = nn.Sequential(self._create_block(3, 64, with_fsdp), self._create_block(64, 64, with_fsdp),)
+        self.head = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)), nn.Flatten(), nn.Linear(64, 10),)
+        if with_fsdp:
+            self.trunk = FSDP(self.trunk)
+            self.head = FSDP(self.head)
+
+    def forward(self, x):
+        return self.head(self.trunk(x))
+
+    def _create_block(self, in_channels, out_channels, with_fsdp):
+        block = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3), nn.ReLU(inplace=True),)
+        if with_fsdp:
+            block = FSDP(block)
+        return block
+
+
+def _create_model(with_fsdp, with_nested_trunk):
+    if with_nested_trunk:
+        model = NestedTrunkModel(with_fsdp)
+    else:
+        model = Model()
+        if with_fsdp:
+            model.trunk = FSDP(model.trunk)
+            model.head = FSDP(model.head)
     return model
 
 
+class FreezingMethod(str, Enum):
+    GradToNone = "grad_to_none"
+    RequiresGrad = "requires_grad"
+
+
 def _distributed_worker(
-    gpu_id, world_size, with_fsdp, freezing_method, tempfile_name, unused, rank_0_output, expected_state
+    gpu_id,
+    world_size,
+    with_fsdp,
+    with_nested_trunk,
+    freezing_method,
+    tempfile_name,
+    unused,
+    rank_0_output,
+    expected_state,
 ):
     torch.cuda.set_device(gpu_id)
 
@@ -59,12 +95,11 @@ def _distributed_worker(
     torch.backends.cudnn.deterministic = True
     batch = torch.randn(size=(2, 3, 224, 224)).cuda()
 
-    model = _create_model(with_fsdp)
+    model = _create_model(with_fsdp, with_nested_trunk)
     model = model.cuda()
 
     # freezing the trunk using requires_grad.
-    assert freezing_method in ["requires_grad", "grad_to_none"]
-    if freezing_method == "requires_grad":
+    if freezing_method == FreezingMethod.RequiresGrad:
         for param in model.trunk.parameters():
             param.requires_grad = False
 
@@ -86,7 +121,7 @@ def _distributed_worker(
         print("Loss", iteration, ":", fake_loss.item())
         optimizer.zero_grad()
         fake_loss.backward()
-        if freezing_method == "grad_to_none":
+        if freezing_method == FreezingMethod.GradToNone:
             for param in model.trunk.parameters():
                 param.grad = None
         optimizer.step()
@@ -118,21 +153,30 @@ def temp_files():
 
 
 @skip_if_single_gpu
-def test_freezing_weights(temp_files):
+@pytest.mark.parametrize("nested_trunk", ["nested_trunk", "simple_trunk"])
+def test_freezing_weights(temp_files, nested_trunk):
+    with_nested_trunk = nested_trunk == "nested_trunk"
+
     world_size = 2
     # DDP
-    fsdp = False
-    freezing_method = "requires_grad"
-    mp.spawn(_distributed_worker, (world_size, fsdp, freezing_method) + temp_files[0:3] + (None,), nprocs=world_size)
+    with_fsdp = False
+    freezing_method = FreezingMethod.RequiresGrad
+    mp.spawn(
+        _distributed_worker,
+        (world_size, with_fsdp, with_nested_trunk, freezing_method) + temp_files[0:3] + (None,),
+        nprocs=world_size,
+    )
     # FSDP, case 1 and 2.
-    fsdp = True
+    with_fsdp = True
     expected_state = torch.load(temp_files[2])
     temp_file_idx = 3
-    for freezing_method in ["requires_grad", "grad_to_none"]:
+    for freezing_method in [FreezingMethod.RequiresGrad, FreezingMethod.GradToNone]:
         print(f"Testing FSDP with freezing method {freezing_method}")
         mp.spawn(
             _distributed_worker,
-            (world_size, fsdp, freezing_method) + temp_files[temp_file_idx : temp_file_idx + 3] + (expected_state,),
+            (world_size, with_fsdp, with_nested_trunk, freezing_method)
+            + temp_files[temp_file_idx : temp_file_idx + 3]
+            + (expected_state,),
             nprocs=world_size,
         )
         temp_file_idx += 3