internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: facebookresearch/fairscale

Issue:
Add mixed precision autocast to benchmarking script 
## 🚀 Feature

It would be nice to make sure OSS/ShardedDDP plays nice with torch autocast! Adding it to the oss.py benchmark will allow us to ensure there is no glaring issues. I don't mind working on this!

## Motivation

Integrating OSS/ShardedDDP into PL :)


Discussion:
This seems like a unit test, not a benchmark?

For unit test, it would be nice to make sure we run a test only when its dependent files change in a PR or commit. But that is certainly outside and independent of this issue.
I think we should have the functionality in both! I've been using the oss.py benchmark script primarily but should make its way into a unit test. What would you like the unit test to replicate/integrate into? was thinking this: https://github.com/facebookresearch/fairscale/blob/master/tests/nn/data_parallel/test_sharded_ddp.py
Have a branch which is super simple, just adds autocast wrappers to the forward passes: https://github.com/SeanNaren/fairscale/tree/feature/autocast_benchmark

Once other functionality has merged, I'll open the PR :)

Resolving patch:
diff --git a/benchmarks/oss.py b/benchmarks/oss.py
--- a/benchmarks/oss.py
+++ b/benchmarks/oss.py
@@ -140,9 +140,15 @@ def closure():
                             next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()
                         )
                     )
+                if not args.cpu and args.amp:
+                    # Automatically computes the FW pass in half precision
+                    with torch.cuda.amp.autocast():
+                        outputs = model(batch["inputs"])
+                        loss = loss_fn(outputs, batch["label"])
+                else:
+                    outputs = model(batch["inputs"])
+                    loss = loss_fn(outputs, batch["label"])
 
-                outputs = model(batch["inputs"])
-                loss = loss_fn(outputs, batch["label"])
                 loss.backward()
 
                 if optim_type == OptimType.oss_sharded_ddp:
@@ -244,7 +250,8 @@ def closure():
     parser.add_argument("--profile", action="store_true", default=False)
     parser.add_argument("--cpu", action="store_true", default=False)
     parser.add_argument("--torchvision_model", type=str, help="Any torchvision model name (str)", default="resnet101")
-    parser.add_argument("--debug", action="store_true", default=False)
+    parser.add_argument("--debug", action="store_true", default=False, help="Display additional debug information")
+    parser.add_argument("--amp", action="store_true", default=False, help="Activate torch AMP")
 
     args = parser.parse_args()
 


Test patch:
Repository: facebookresearch/fairscale

Issue:
an error importing AMP in fairscale when using a Pytorch version less than 1.6
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

## Command

## To Reproduce

<img width="737" alt="Screenshot 2020-11-23 at 08 21 28" src="https://user-images.githubusercontent.com/12861981/99949945-c2f11300-2d73-11eb-87f2-17a02e64da75.png">


Steps to reproduce the behavior:

<!-- If you were running a command, post the exact command that you were running -->

1.
2.
3.

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from the
environment collection script from PyTorch
(or fill out the checklist below manually).

You can run the script with:
```
# For security purposes, please check the contents of collect_env.py before running it.
python -m torch.utils.collect_env
```

 - PyTorch Version (e.g., 1.0):
 - OS (e.g., Linux):
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:
 - GPU models and configuration:
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->



Discussion:
I"ll add an import guard, sorry about that

Resolving patch:
diff --git a/fairscale/optim/__init__.py b/fairscale/optim/__init__.py
--- a/fairscale/optim/__init__.py
+++ b/fairscale/optim/__init__.py
@@ -6,11 +6,16 @@
 """
 :mod:`fairscale.optim` is a package implementing various torch optimization algorithms.
 """
+import logging
+
+from .adascale import AdaScale
+from .oss import OSS
 
 try:
     from .adam import Adam, Precision
 except ImportError:  # pragma: no cover
     pass  # pragma: no cover
-from .adascale import AdaScale
-from .grad_scaler import GradScaler
-from .oss import OSS
+try:
+    from .grad_scaler import GradScaler
+except ImportError:
+    logging.warning("Torch AMP is not available on this platform")


Test patch:
Repository: facebookresearch/fairscale

Issue:
[ShardedGradScaler] Handle optimizers not being OSS instances
## 🚀 Feature
Some frameworks (Classy for instance) change the optimizers to enable different non-pytorch-compliant features. ShardedGradScaler asserts on that

## Motivation
Enable ShardedDDP + AMP on Classy-like frameworks

## Pitch
Remove the assert, replace by a one time warning

## Alternatives
Not doing anything

## Additional context
@mannatsingh 


Discussion:

Resolving patch:
diff --git a/fairscale/optim/grad_scaler.py b/fairscale/optim/grad_scaler.py
--- a/fairscale/optim/grad_scaler.py
+++ b/fairscale/optim/grad_scaler.py
@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import logging
 from typing import Dict
 
 import torch
@@ -31,9 +32,16 @@ class ShardedGradScaler(TorchGradScaler):
 
     def __init__(self) -> None:
         super().__init__()
+        self.display_warning = True
 
     def unscale_(self, optimizer: Optimizer) -> None:
-        assert isinstance(optimizer, OSS), "ShardedGradScaler is to be used in combination with a sharded optimizer"
+        # Could be a mistake, this scaler is supposed to work with ZeroRedundancyOptimizer only
+        if self.display_warning and not isinstance(optimizer, OSS):
+            logging.warning(
+                "ShardedGradScaler is to be used in combination with a sharded optimizer, this could not be checked"
+            )
+
+        self.display_warning = False  # Only warn once
 
         # Call the upstream unscale_ method which will only act on this rank's gradients
         super().unscale_(optimizer)


Test patch:
Repository: facebookresearch/fairscale

Issue:
[bug] pip package 0.1.3 fails to install
Both `python3.8 -m pip install fairscale --verbose` and `python3.8 -m pip install fairscale --no-build-isolation --verbose` trigger

```
ninja: error: '/tmp/pip-install-zqe21k7a/fairscale_4066f1fa225242299ead5fd852fd2ce8/fairscale/clib/fused_adam_cuda/fused_adam_cuda.cpp', needed by '/t
mp/pip-install-zqe21k7a/fairscale_4066f1fa225242299ead5fd852fd2ce8/build/temp.linux-x86_64-3.8/fairscale/clib/fused_adam_cuda/fused_adam_cuda.o', missin
g and no known rule to make it
```

Cloning the repository and running `python3.8 -m pip install . --verbose` from within the directory works fine.


Discussion:
oh, I was the one doing this release, I may have done something wrong. Testing locally I cannot reproduce though, @msbaines  would you have an idea ? I remember seeing this issue discussed before
Maybe related to https://github.com/facebookresearch/fairscale/issues/190
ok, I just checked the package and these files are missing indeed, I'm not sure why the pip install works out for me, maybe because of caching. Trying to publish a new package asap

Resolving patch:
diff --git a/fairscale/__init__.py b/fairscale/__init__.py
--- a/fairscale/__init__.py
+++ b/fairscale/__init__.py
@@ -3,7 +3,7 @@
 # This source code is licensed under the BSD license found in the
 # LICENSE file in the root directory of this source tree.
 
-__version__ = "0.1.3"
+__version__ = "0.1.4"
 
 ################################################################################
 # Import most common subpackages
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -65,6 +65,8 @@ def find_version(version_file_path):
         python_requires=">=3.6",
         author="Facebook AI Research",
         author_email="todo@fb.com",
+        long_description="FairScale is a PyTorch extension library for high performance and large scale training on one or multiple machines/nodes. This library extends basic PyTorch capabilities while adding new experimental ones.",
+        long_description_content_type="text/markdown",
         classifiers=[
             "Programming Language :: Python :: 3.6",
             "Programming Language :: Python :: 3.7",


Test patch:
Repository: facebookresearch/fairscale

Issue:
[OSS] When loading a state, the insertion order for the params may not match the params_groups
## 🐛 Bug
(FB Only f250383189) Guess is that the issue comes from the parameters not being inserted in the right order in the sharded optimizer state. [See](https://github.com/facebookresearch/fairscale/blob/master/fairscale/optim/oss.py#L403)

Would need fixing + repro in a unit test


Discussion:

Resolving patch:
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -391,16 +391,16 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
 
         # NOTE: PyTorch 1.5 does not index linearly but with the id(params) at saving time
         # we work around that here by using the fact that the params are ordered as in the param_groups
+        pytorch15_index_redirect = {k: i for i, k in enumerate(state_dict["state"].keys())}
 
-        for i_param, (key, value) in enumerate(state_dict["state"].items()):
-            param = self.index_to_param[i_param]
+        for key, value in state_dict["state"].items():
+            param = self.index_to_param[pytorch15_index_redirect[key]]
 
             # Populate the sharded optimizer state on the fly
             if self.param_to_rank[param] != self.rank:
                 state_dict["state"][key] = None
 
-            if key in self.index_to_param:
-                param = self.index_to_param[i_param]
+            else:
 
                 # Only add this state to the sharded optimizer if it owns this param
                 for pg in self.optim.param_groups:


Test patch:
Repository: facebookresearch/fairscale

Issue:
[OSS] Catch the device being changed in between init and step()
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

## Pitch

<!-- A clear and concise description of what you want to happen. -->

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:

Resolving patch:
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -220,6 +220,12 @@ def step(self, closure: Optional[Callable[[], float]] = None, **kwargs: Any) ->
         # Sync oss param_groups attributes in case they've been updated by a scheduler.
         OSS._sync_param_groups(self.param_groups, self.optim.param_groups)
 
+        # Catch a possible change of devices in between OSS construction and step()
+        if self._default_device.type != self.param_groups[0]["params"][0].device.type:
+            logging.info("OSS detected that the parameter changed devices, re-allocating buffers")
+            self._clear_cache()
+            self.refresh_trainable()
+
         # Run the optimizer step on this shard only:
         if closure is not None:
             loss = self.optim.step(closure=closure, **kwargs)  # type: ignore
@@ -591,3 +597,9 @@ def _setup_flat_buffers(self) -> None:
                 else:
                     # This rank has an empty shard, that's fine
                     self.buckets[device].append(torch.zeros(0, device=device))
+
+        # Clear the buffer keys which are not in use anymore (could be that the devices changed)
+        devices_in_use = list(self.per_device_params.keys())
+        devices_to_pop = list(filter(lambda x: x not in devices_in_use, self.buckets.keys()))
+        for d in devices_to_pop:
+            self.buckets.pop(d)


Test patch:
diff --git a/tests/optim/test_oss.py b/tests/optim/test_oss.py
--- a/tests/optim/test_oss.py
+++ b/tests/optim/test_oss.py
@@ -154,6 +154,18 @@ def step(self, closure=None, kwarg=[]):
         assert kwarg == [5]
         assert x == torch.tensor([0.9], device=DEVICE)
 
+    @skip_if_no_cuda
+    def test_device_change(self):
+        x = torch.nn.Linear(1, 1).to("cpu")
+        o = optim.OSS(x.parameters(), torch.optim.SGD, lr=0.1)
+
+        # Move the model to device after OSS was constructed
+        x.to(DEVICE)
+        x(torch.zeros((1), device=DEVICE)).backward()
+
+        # Check that OSS detects that the device changed
+        o.step()
+
     def test_step_with_extra_inner_key(self):
         class SGDWithNewKey(torch.optim.SGD):
             # Dummy optimizer which adds a new key to the param groups
Repository: facebookresearch/fairscale

Issue:
[FSDP] use case of wrapping module without any weights
This question came up during this PR https://github.com/facebookresearch/fairscale/pull/657

It is unclear there is a valid use case for it. If not, we may want to add a check and give user a clear error msg about it.


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
@@ -1317,18 +1317,15 @@ 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
-                # 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:
                         m.assert_state(TrainingState.BACKWARD_PRE)
                 else:
-                    # 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.
+                    # Unlikely case. When `m` and its children has no params or has params but
+                    # none 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
 


Test patch:
Repository: facebookresearch/fairscale

Issue:
OSS detected that the parameter changed devices, re-allocating buffers 
## ❓ Questions and Help
For each iteration, I get this message "OSS detected that the parameter changed devices, re-allocating buffers ". Is this normal? if not, what should be done? 


Discussion:
cc @blefaudeux .  I think it happen when you model's compute graph change between iterations. This can cause a slowdown in training but should not cause correctness problem.
@gongbudaizhe would you happen to move your model to and from GPU at each iteration, or change the compute graph (make parameters trainable and no-trainable) ? The message should be limited to these cases, if not this is a bug and I would love more information to repro ! 

To give some insight, OSS has to change the way the parameters are stored, for efficiency purposes, and moving the model to cpu and back breaks that.
Thanks for your prompt reply. I am not aware of any operations that move model to and from GPU, or change the compute graph. 

The model used is MoCO (don't worry, the code below is concise).  Do you have any clue?

https://github.com/facebookresearch/moco/blob/master/moco/builder.py
Ah, I know MoCo, that's a good link ! 
So that's probably unrelated though, I'm guessing that the model is moved to GPU *after* the wrap of the optimizer ? If you wrap the optimizer last (after the model is moved to GPU), then the error should go away and the training will speed up a fair bit. I'll update the code to make it easier though, I get that this can be a hindrance, one sec

Resolving patch:
diff --git a/fairscale/optim/oss.py b/fairscale/optim/oss.py
--- a/fairscale/optim/oss.py
+++ b/fairscale/optim/oss.py
@@ -462,10 +462,12 @@ def refresh_trainable(self) -> None:
         of some parameters changed.
         """
 
+        # Make sure that we capture the current default device
+        self._default_device = list(self._per_device_params.keys())[0]
+
         # Create the optim which will work on the param shard
         if not hasattr(self, "optim"):
             self._clear_cache()
-            self._default_device = list(self._per_device_params.keys())[0]
             self.optim = self._optim_constructor(self.partition_parameters()[self.rank], **self._optim_defaults)
             OSS._sync_param_groups(self.optim.param_groups, self.param_groups)
 


Test patch:
diff --git a/tests/optim/test_oss.py b/tests/optim/test_oss.py
--- a/tests/optim/test_oss.py
+++ b/tests/optim/test_oss.py
@@ -151,6 +151,9 @@ def test_device_change(self):
         # Check that OSS detects that the device changed
         o.step()
 
+        # Check that the default device has been updated
+        assert o._default_device.type == DEVICE
+
     def test_step_with_extra_inner_key(self):
         class SGDWithNewKey(torch.optim.SGD):
             # Dummy optimizer which adds a new key to the param groups
Repository: facebookresearch/fairscale

Issue:
shard_model does not respect shard_count.
## 🐛 Bug

`autoshard.shard_model()` does not respect `shard_count`. See [comments](https://github.com/facebookresearch/fairscale/pull/817/files#r733237403) in https://github.com/facebookresearch/fairscale/pull/817 for more context and how to reproduce. 


Discussion:

Resolving patch:
diff --git a/fairscale/experimental/nn/auto_shard.py b/fairscale/experimental/nn/auto_shard.py
--- a/fairscale/experimental/nn/auto_shard.py
+++ b/fairscale/experimental/nn/auto_shard.py
@@ -52,8 +52,7 @@ def _split_nodes(traced_graph_module: torch.fx.GraphModule, shard_count: int = 3
     # Find the total number of params in the model and
     # the number of params per shard we are aiming for.
     for name, module in traced_graph_module.named_modules():
-        if "." in name:
-            continue
+        name = name.replace(".", "_")
         param_count[name] = sum([x.numel() for x in module.parameters()])
     logging.info(f"Total number of params are {param_count['']}")
     per_shard_param = param_count[""] // shard_count


Test patch:
diff --git a/tests/experimental/nn/test_auto_shard.py b/tests/experimental/nn/test_auto_shard.py
--- a/tests/experimental/nn/test_auto_shard.py
+++ b/tests/experimental/nn/test_auto_shard.py
@@ -140,9 +140,7 @@ def test_dynaimc_conditionals_auto_wrapped():
 
     model = BranchedNetwork(features)
     sharded_model = shard_model(model, 3)
-    # TODO(ehotaj): There might be a bug in our split code because we shard the
-    # model into 10 shards even though we specify 3 shards above.
-    assert len(sharded_model) == 10
+    assert len(sharded_model) == 3
 
     input_ = torch.randn(3, features)
     model_output = model(input_)
Repository: facebookresearch/fairscale

Issue:
[FSDP] reshard_after_forward = False for root instance
```
if self._is_root:
    # Buffers stay on GPU, and don't get sharded. Since _cast_buffers
    # applies recursively, we only call this from the root instance.
    self._cast_buffers()
            
    # Don't free the full params for the outer-most (root) instance,
    # since those params will be needed immediately after for the
    # backward pass.
    self.reshard_after_forward = False

    # Due to the use of streams, we need to make sure the previous
    # ``optim.step()`` is done before we all-gather parameters.
    self._wait_for_previous_optim_step()
```
According to the following code (from fully_sharded_data_parallel.py), we can find that reshard_after_forward is set to False for root instance (for better performance). However, when we use autowrap() to wrap model, we only wrap leaf nodes, i.e., every leaf is an FSDP root instance. In this case, every FSDP instance's reshard_after_forward will be set to False even if we set reshard_after_forward to True (in initialization). However, due to a memory issue, I still want to reshard after forward when wrapping with autowrap(). 
I have a few questions,
1. Could I just comment "self.reshard_after_forward = False" as an work around solution?
2. Could you add a parameter in __init__() of FullyShardedDataParallel to decide whether self.reshard_after_forward need to be set to False for root instance?
3. For the FSDP wrapping using autowrap(), is the performance affected by the reshard_after_forward? Since the FSDP instances in this case are not outer-most (root) instances, these params will not be needed immediately for the backward pass


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
@@ -307,7 +307,7 @@ def __init__(
         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
+        self.reshard_after_forward = self._orig_reshard_after_forward = reshard_after_forward
         self.mixed_precision = mixed_precision
         self.fp32_reduce_scatter = fp32_reduce_scatter
         self.flatten_parameters = flatten_parameters
@@ -1091,6 +1091,7 @@ def _reset_lazy_init(self) -> None:
             if hasattr(p, "_fp32_shard"):
                 del p._fp32_shard  # reset _init_param_attributes
         self._output_pre_backward_hook_registered: Optional[List] = None
+        self.reshard_after_forward = self._orig_reshard_after_forward
 
     def _lazy_init(self) -> None:
         """Initialization steps that should happen lazily, typically right


Test patch: