Repository: facebookresearch/fairscale

Issue:
Auto wrapping Huggingface models with FullyShardedDataParallel can break them
## 🐛 Bug

If you auto wrap a HuggingFace model with `FullyShardedDataParallel`, depending on what sub-modules are wrapped, and how the output of the modules are accessed in the model, calling `forward` can result in exceptions. 

## What happens

Some sub-modules of the HuggingFace model output a child implementation of [`ModelOutput`](https://github.com/huggingface/transformers/blob/8b240a06617455eae59e1116af6a1a016664e963/src/transformers/file_utils.py#L1884), which is itself a descendent of `OrderedDict`.

An example is, for instance, [`CausalLMOutputWithCrossAttentions` ](https://github.com/huggingface/transformers/blob/8b240a06617455eae59e1116af6a1a016664e963/src/transformers/modeling_outputs.py#L375).

You can access attributes of `ModelOutput` instance derivatives in three ways, e.g. like this:

- `hidden_states = transformer_outputs[0]`
- `hidden_states = transformer_outputs.hidden_states`
- `hidden_states = transformer_outputs["hidden_states"]`

In the HuggingFace library these three types of accessing attributes are used interchangeably.

Now, when a sub-module, which outputs a `ModelOutput` derived instance,  is wrapped with `FullyShardedDataParallel`, this output instance is converted in to an `OrderedDict`, in the [`apply_to_tensors`](https://github.com/facebookresearch/fairscale/blob/fecb665b812b6bfc38442e1fb1557e21508917f4/fairscale/utils/containers.py#L21) function, in `utils/containers.py`, since the `ModelOutput` derives from `OrderedDict`! :

```
        elif isinstance(x, OrderedDict):
            od = OrderedDict()
            for key, value in x.items():
                od[key] = _apply(value)
            return od
```

After it has been converted into a plain `OrderedDict` by `FullyShardedDataParallel`, when the output of the submodule is, for instance, accessed using an index (`transformer_outputs[0]`) or attribute accessor (`transformer_outputs. hidden_states`), an exception is thrown, because the `OrderedDict` instance doesn't know how to handle that. 

## A simple solution

A simple solution is to replace the `OrderedDict` instantiation above with `od = x.__class__()`, thus:

```
        elif isinstance(x, OrderedDict):
            od = x.__class__()
            for key, value in x.items():
                od[key] = _apply(value)
            return od
```

In this way we keep on using the original class, with its extended means to access attributes, which is used in the HuggingFace library.


## To Reproduce

Steps to reproduce the behavior:

1. Use a HuggingFace model, e.g. `GPT2LMHeadModel`
```
from transformers import GPT2Config, GPT2LMHeadModel
```

2. Build the model 
```
        # You could also try, e.g. "gpt2-large"
        model_config = GPT2Config.from_pretrained("gpt2")
        self._model = GPT2LMHeadModel(model_config)

        self._model.to(self._device)
```

3. Auto wrap the model
```
                wrap_policy = functools.partial(default_auto_wrap_policy,
                                                min_num_params=1e8)

                fsdp_params = {
                    "mixed_precision": False,
                    "move_params_to_cpu": False
                }

                with enable_wrap(wrapper_cls=FullyShardedDDP, auto_wrap_policy=wrap_policy, **fsdp_params):
                    self._training_model = auto_wrap(self._model)

                # Finally wrap root module
                self._training_model = FullyShardedDDP(self._training_model, **fsdp_params)
```

4. Use this in a training loop

This results in errors similar to this one where the `OrderedDict` accessed in invalid ways:
```
... transformers/models/gpt2/modeling_gpt2.py", line 1057, in forward
    hidden_states = transformer_outputs[0]
KeyError: 0
```

## Expected behavior

No exception should occur, `FullyShardedDataParallel` should not break calling the HuggingFace model.

## Environment

```
$ pip show fairscale
Name: fairscale
Version: 0.4.3

$ pip show transformers
Name: transformers
Version: 4.12.5

$ pip show torch
Name: torch
Version: 1.10.0

$ python --version
Python 3.7.4
```



Discussion:
^ Updated the issue text, by accident published it early.
Is there a way to use a custom wrapping policy instead of the default_auto_wrap_policy function in this case? FSDP can't easily detect what is a safe boundary for wrapping a sub-module. 

Resolving patch:
diff --git a/fairscale/utils/containers.py b/fairscale/utils/containers.py
--- a/fairscale/utils/containers.py
+++ b/fairscale/utils/containers.py
@@ -19,7 +19,7 @@ def _apply(x: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:
         if torch.is_tensor(x):
             return fn(x)
         elif isinstance(x, OrderedDict):
-            od = OrderedDict()
+            od = x.__class__()
             for key, value in x.items():
                 od[key] = _apply(value)
             return od


Test patch:
diff --git a/tests/ci_test_list_3.txt b/tests/ci_test_list_3.txt
--- a/tests/ci_test_list_3.txt
+++ b/tests/ci_test_list_3.txt
@@ -22,3 +22,4 @@ tests/optim/test_oss.py
 tests/optim/test_oss_adascale.py
 tests/optim/test_ddp_adascale.py
 tests/experimental/nn/data_parallel/test_gossip.py
+tests/nn/data_parallel/test_fsdp_hf_transformer_eval.py
diff --git a/tests/nn/data_parallel/test_fsdp_hf_transformer_eval.py b/tests/nn/data_parallel/test_fsdp_hf_transformer_eval.py
new file mode 100644
--- /dev/null
+++ b/tests/nn/data_parallel/test_fsdp_hf_transformer_eval.py
@@ -0,0 +1,122 @@
+from collections import OrderedDict
+from dataclasses import dataclass
+import tempfile
+import unittest
+
+import torch
+from torch import nn
+
+from fairscale.nn import FullyShardedDataParallel as FSDP
+from fairscale.nn import auto_wrap, enable_wrap
+from fairscale.utils.testing import dist_init
+
+
+def wrap_transformer_only(module, recurse, **kwargs):
+    if recurse:
+        return True
+    else:
+        return isinstance(module, nn.Transformer)
+
+
+class ModelOutput(OrderedDict):
+    # Partially adapted from the HF transformers lib in order to simulate the behavior
+
+    def to_tuple(self):
+        return tuple(self[k] for k in self.keys())
+
+    def __post_init__(self):
+        class_fields = getattr(self, "__dataclass_fields__")
+
+        for field in class_fields:
+            v = getattr(self, field)
+            if v is not None:
+                self[field] = v
+
+    def __getitem__(self, k):
+        if isinstance(k, str):
+            inner_dict = {k: v for (k, v) in self.items()}
+            return inner_dict[k]
+        else:
+            return self.to_tuple()[k]
+
+
+@dataclass
+class TransformerOutput(ModelOutput):
+    output: torch.FloatTensor = None
+
+
+class TransformerWithCustomOutput(nn.Transformer):  # type: ignore[name-defined]
+    def forward(self, *args, **kwargs):
+        output = super().forward(*args, **kwargs)
+
+        return TransformerOutput(output=output)
+
+
+class TransformerWithLMHead(nn.Module):
+    def __init__(self, d_vocab=100, d_model=16):
+        super().__init__()
+        self.d_vocab = d_vocab
+        self.d_model = d_model
+
+        self.embed_tokens = nn.Embedding(d_vocab, d_model)
+
+        self.transformer = TransformerWithCustomOutput(
+            d_model, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=64
+        )
+
+        self.output_proj = nn.Linear(d_model, d_vocab)
+
+    def generate_random_sequences(self, seq_len=20, batch_size=2):
+        source_seq = torch.randint(high=self.d_vocab, size=(seq_len, batch_size))
+        target_seq = torch.randint(high=self.d_vocab, size=(seq_len, batch_size))
+
+        return source_seq, target_seq
+
+    def forward(self, source_seq, target_seq):
+        source_embeddings = self.embed_tokens(source_seq)
+        target_embeddings = self.embed_tokens(target_seq)
+
+        output = self.transformer(source_embeddings, target_embeddings)
+
+        # Using integer key here, just like in Huggingface transformer lib
+        return self.output_proj(output[0])
+
+
+class TestHFTransformersAutoWrap(unittest.TestCase):
+    def setUp(self) -> None:
+        if not torch.cuda.is_available():
+            raise unittest.SkipTest("CUDA not available, skipping test")
+
+        torch.cuda.set_device(0)
+
+        _, filename = tempfile.mkstemp()
+        _, filename_rpc = tempfile.mkstemp()
+
+        dist_init(0, 1, filename, filename_rpc)
+
+        self.device = torch.device("cuda")
+
+        print("Build model ...")
+        self.model = TransformerWithLMHead()
+        self.model.to(self.device)
+
+    def test_auto_wrap_hf_model(self):
+
+        with enable_wrap(wrapper_cls=FSDP, auto_wrap_policy=wrap_transformer_only):
+            self.model = auto_wrap(self.model)
+
+        self.model = FSDP(self.model)
+
+        self.assertTrue(isinstance(self.model.transformer, FSDP), "Transformer should have been wrapped with FSDP")
+
+        source_seq, target_seq = self.model.generate_random_sequences()
+        source_seq = source_seq.to(self.device)
+        target_seq = target_seq.to(self.device)
+
+        print("Evaluating model ...")
+        # This should not fail
+        self.model(source_seq, target_seq)
+
+
+if __name__ == "__main__":
+    unittest.main()