# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.

from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

import torch
from torch.nn.utils.rnn import PackedSequence

"""Useful functions to deal with tensor types with other python container types."""


def apply_to_tensors(fn: Callable, container: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:
    """Recursively apply to all tensor in different kinds of container types."""

    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()
            for key, value in x.items():
                od[key] = _apply(value)
            return od
        elif isinstance(x, PackedSequence):
            _apply(x.data)
            return x
        elif isinstance(x, dict):
            return {key: _apply(value) for key, value in x.items()}
        elif isinstance(x, list):
            return [_apply(x) for x in x]
        elif isinstance(x, tuple):
            return tuple(_apply(x) for x in x)
        elif isinstance(x, set):
            return {_apply(x) for x in x}
        else:
            return x

    return _apply(container)


def pack_kwargs(*args: Any, **kwargs: Any) -> Tuple[Tuple[str, ...], Tuple[Any, ...]]:
    """
    Turn argument list into separate key list and value list (unpack_kwargs does the opposite)

    Usage::

        kwarg_keys, flat_args = pack_kwargs(1, 2, a=3, b=4)
        assert kwarg_keys == ("a", "b")
        assert flat_args == (1, 2, 3, 4)
        args, kwargs = unpack_kwargs(kwarg_keys, flat_args)
        assert args == (1, 2)
        assert kwargs == {"a": 3, "b": 4}
    """
    kwarg_keys: List[str] = []
    flat_args: List[Any] = list(args)
    for k, v in kwargs.items():
        kwarg_keys.append(k)
        flat_args.append(v)
    return tuple(kwarg_keys), tuple(flat_args)


def unpack_kwargs(kwarg_keys: Tuple[str, ...], flat_args: Tuple[Any, ...]) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
    """See pack_kwargs."""
    assert len(kwarg_keys) <= len(flat_args), f"too many keys {len(kwarg_keys)} vs. {len(flat_args)}"
    if len(kwarg_keys) == 0:
        return flat_args, {}
    args = flat_args[: -len(kwarg_keys)]
    kwargs = {k: v for k, v in zip(kwarg_keys, flat_args[-len(kwarg_keys) :])}
    return args, kwargs


def split_non_tensors(
    mixed: Union[torch.Tensor, Tuple[Any, ...]]
) -> Tuple[Tuple[torch.Tensor, ...], Optional[Dict[str, List[Any]]]]:
    """
    Split a tuple into a list of tensors and the rest with information
    for later reconstruction.

    Usage::

        x = torch.Tensor([1])
        y = torch.Tensor([2])
        tensors, packed_non_tensors = split_non_tensors((x, y, None, 3))
        assert tensors == (x, y)
        assert packed_non_tensors == {
            "is_tensor": [True, True, False, False],
            "objects": [None, 3],
        }
        recon = unpack_non_tensors(tensors, packed_non_tensors)
        assert recon == (x, y, None, 3)
    """
    if isinstance(mixed, torch.Tensor):
        return (mixed,), None
    tensors: List[torch.Tensor] = []
    packed_non_tensors: Dict[str, List[Any]] = {"is_tensor": [], "objects": []}
    for o in mixed:
        if isinstance(o, torch.Tensor):
            packed_non_tensors["is_tensor"].append(True)
            tensors.append(o)
        else:
            packed_non_tensors["is_tensor"].append(False)
            packed_non_tensors["objects"].append(o)
    return tuple(tensors), packed_non_tensors


def unpack_non_tensors(
    tensors: Tuple[torch.Tensor, ...], packed_non_tensors: Optional[Dict[str, List[Any]]]
) -> Tuple[Any, ...]:
    """See split_non_tensors."""
    if packed_non_tensors is None:
        return tensors
    assert isinstance(packed_non_tensors, dict), type(packed_non_tensors)
    mixed: List[Any] = []
    is_tensor_list = packed_non_tensors["is_tensor"]
    objects = packed_non_tensors["objects"]
    assert len(tensors) + len(objects) == len(is_tensor_list), (
        f"len(tensors) {len(tensors)} len(objects) {len(objects)} " f"len(is_tensor_list) {len(is_tensor_list)}"
    )
    obj_i = tnsr_i = 0
    for is_tensor in is_tensor_list:
        if is_tensor:
            mixed.append(tensors[tnsr_i])
            tnsr_i += 1
        else:
            mixed.append(objects[obj_i])
            obj_i += 1
    return tuple(mixed)

fairscale/utils/containers.py
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.

from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

import torch
from torch.nn.utils.rnn import PackedSequence

"""Useful functions to deal with tensor types with other python container types."""


def apply_to_tensors(fn: Callable, container: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:
    """Recursively apply to all tensor in different kinds of container types."""

    def _apply(x: Union[torch.Tensor, Dict, List, Tuple, Set]) -> Any:
        if torch.is_tensor(x):
            return fn(x)
        elif isinstance(x, OrderedDict):
            od = x.__class__()
            for key, value in x.items():
                od[key] = _apply(value)
            return od
        elif isinstance(x, PackedSequence):
            _apply(x.data)
            return x
        elif isinstance(x, dict):
            return {key: _apply(value) for key, value in x.items()}
        elif isinstance(x, list):
            return [_apply(x) for x in x]
        elif isinstance(x, tuple):
            return tuple(_apply(x) for x in x)
        elif isinstance(x, set):
            return {_apply(x) for x in x}
        else:
            return x

    return _apply(container)


def pack_kwargs(*args: Any, **kwargs: Any) -> Tuple[Tuple[str, ...], Tuple[Any, ...]]:
    """
    Turn argument list into separate key list and value list (unpack_kwargs does the opposite)

    Usage::

        kwarg_keys, flat_args = pack_kwargs(1, 2, a=3, b=4)
        assert kwarg_keys == ("a", "b")
        assert flat_args == (1, 2, 3, 4)
        args, kwargs = unpack_kwargs(kwarg_keys, flat_args)
        assert args == (1, 2)
        assert kwargs == {"a": 3, "b": 4}
    """
    kwarg_keys: List[str] = []
    flat_args: List[Any] = list(args)
    for k, v in kwargs.items():
        kwarg_keys.append(k)
        flat_args.append(v)
    return tuple(kwarg_keys), tuple(flat_args)


def unpack_kwargs(kwarg_keys: Tuple[str, ...], flat_args: Tuple[Any, ...]) -> Tuple[Tuple[Any, ...], Dict[str, Any]]:
    """See pack_kwargs."""
    assert len(kwarg_keys) <= len(flat_args), f"too many keys {len(kwarg_keys)} vs. {len(flat_args)}"
    if len(kwarg_keys) == 0:
        return flat_args, {}
    args = flat_args[: -len(kwarg_keys)]
    kwargs = {k: v for k, v in zip(kwarg_keys, flat_args[-len(kwarg_keys) :])}
    return args, kwargs


def split_non_tensors(
    mixed: Union[torch.Tensor, Tuple[Any, ...]]
) -> Tuple[Tuple[torch.Tensor, ...], Optional[Dict[str, List[Any]]]]:
    """
    Split a tuple into a list of tensors and the rest with information
    for later reconstruction.

    Usage::

        x = torch.Tensor([1])
        y = torch.Tensor([2])
        tensors, packed_non_tensors = split_non_tensors((x, y, None, 3))
        assert tensors == (x, y)
        assert packed_non_tensors == {
            "is_tensor": [True, True, False, False],
            "objects": [None, 3],
        }
        recon = unpack_non_tensors(tensors, packed_non_tensors)
        assert recon == (x, y, None, 3)
    """
    if isinstance(mixed, torch.Tensor):
        return (mixed,), None
    tensors: List[torch.Tensor] = []
    packed_non_tensors: Dict[str, List[Any]] = {"is_tensor": [], "objects": []}
    for o in mixed:
        if isinstance(o, torch.Tensor):
            packed_non_tensors["is_tensor"].append(True)
            tensors.append(o)
        else:
            packed_non_tensors["is_tensor"].append(False)
            packed_non_tensors["objects"].append(o)
    return tuple(tensors), packed_non_tensors


def unpack_non_tensors(
    tensors: Tuple[torch.Tensor, ...], packed_non_tensors: Optional[Dict[str, List[Any]]]
) -> Tuple[Any, ...]:
    """See split_non_tensors."""
    if packed_non_tensors is None:
        return tensors
    assert isinstance(packed_non_tensors, dict), type(packed_non_tensors)
    mixed: List[Any] = []
    is_tensor_list = packed_non_tensors["is_tensor"]
    objects = packed_non_tensors["objects"]
    assert len(tensors) + len(objects) == len(is_tensor_list), (
        f"len(tensors) {len(tensors)} len(objects) {len(objects)} " f"len(is_tensor_list) {len(is_tensor_list)}"
    )
    obj_i = tnsr_i = 0
    for is_tensor in is_tensor_list:
        if is_tensor:
            mixed.append(tensors[tnsr_i])
            tnsr_i += 1
        else:
            mixed.append(objects[obj_i])
            obj_i += 1
    return tuple(mixed)

fairscale/utils/containers.py
diff --git a/fairscale/utils/containers.py b/fairscale/utils/containers.py
index a56985d25..7ea0bb454 100644
--- 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
diff --git a/tests/ci_test_list_3.txt b/tests/ci_test_list_3.txt
index 1986a77c2..d350cf19b 100644
--- 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
index 000000000..d1ec3c1ba
--- /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()

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
```

