prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  DD-PPO does not all reduce gradients
  ## 🐛 Bug
  
  DD-PPO does not all reduce gradients during the backward call, because `reducer.prepare_for_backward` is not being called during training process.
  
  The problem is in this line: https://github.com/facebookresearch/habitat-api/blob/v0.1.4/habitat_baselines/rl/ddppo/algo/ddppo.py#L96
  
  ```
  class DecentralizedDistributedMixin:
  
      ...
      def before_backward(self, loss):
           # ...
           self.reducer.prepare_for_backward(..)
  
  
  # Mixin goes second that way the PPO __init__ will still be called
  class DDPPO(PPO, DecentralizedDistributedMixin):  
      # Here PPO and Mixin both have "before_backward" method, 
      # DDPPO will call PPO's not the Mixin's.
      pass
  ```
  
  And here is a quick fix:
  ```
  class DecentralizedDistributedMixin:
      ...
  
  
  # Mixin goes second that way the PPO __init__ will still be called
  class DDPPO(PPO, DecentralizedDistributedMixin):  
  
        # Move before_backward to DDPPO
        def before_backward(self, loss):
           # ...
           self.reducer.prepare_for_backward(..)
  ```
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `habitat_baselines/rl/ddppo/algo/ddppo.py`
  Content:
  ```
  1 #!/usr/bin/env python3
  2 
  3 # Copyright (c) Facebook, Inc. and its affiliates.
  4 # This source code is licensed under the MIT license found in the
  5 # LICENSE file in the root directory of this source tree.
  6 
  7 from typing import Tuple
  8 
  9 import torch
  10 import torch.distributed as distrib
  11 
  12 from habitat_baselines.common.rollout_storage import RolloutStorage
  13 from habitat_baselines.rl.ppo import PPO
  14 
  15 EPS_PPO = 1e-5
  16 
  17 
  18 def distributed_mean_and_var(
  19     values: torch.Tensor,
  20 ) -> Tuple[torch.Tensor, torch.Tensor]:
  21     r"""Computes the mean and variances of a tensor over multiple workers.
  22 
  23     This method is equivalent to first collecting all versions of values and
  24     then computing the mean and variance locally over that
  25 
  26     :param values: (*,) shaped tensors to compute mean and variance over.  Assumed
  27                         to be solely the workers local copy of this tensor,
  28                         the resultant mean and variance will be computed
  29                         over _all_ workers version of this tensor.
  30     """
  31     assert distrib.is_initialized(), "Distributed must be initialized"
  32 
  33     world_size = distrib.get_world_size()
  34     mean = values.mean()
  35     distrib.all_reduce(mean)
  36     mean /= world_size
  37 
  38     sq_diff = (values - mean).pow(2).mean()
  39     distrib.all_reduce(sq_diff)
  40     var = sq_diff / world_size
  41 
  42     return mean, var
  43 
  44 
  45 class DecentralizedDistributedMixin:
  46     def _get_advantages_distributed(
  47         self, rollouts: RolloutStorage
  48     ) -> torch.Tensor:
  49         advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]
  50         if not self.use_normalized_advantage:
  51             return advantages
  52 
  53         mean, var = distributed_mean_and_var(advantages)
  54 
  55         return (advantages - mean) / (var.sqrt() + EPS_PPO)
  56 
  57     def init_distributed(self, find_unused_params: bool = True) -> None:
  58         r"""Initializes distributed training for the model
  59 
  60         1. Broadcasts the model weights from world_rank 0 to all other workers
  61         2. Adds gradient hooks to the model
  62 
  63         :param find_unused_params: Whether or not to filter out unused parameters
  64                                    before gradient reduction.  This *must* be True if
  65                                    there are any parameters in the model that where unused in the
  66                                    forward pass, otherwise the gradient reduction
  67                                    will not work correctly.
  68         """
  69         # NB: Used to hide the hooks from the nn.Module,
  70         # so they don't show up in the state_dict
  71         class Guard:
  72             def __init__(self, model, device):
  73                 if torch.cuda.is_available():
  74                     self.ddp = torch.nn.parallel.DistributedDataParallel(
  75                         model, device_ids=[device], output_device=device
  76                     )
  77                 else:
  78                     self.ddp = torch.nn.parallel.DistributedDataParallel(model)
  79 
  80         self._ddp_hooks = Guard(self.actor_critic, self.device)
  81         self.get_advantages = self._get_advantages_distributed
  82 
  83         self.reducer = self._ddp_hooks.ddp.reducer
  84         self.find_unused_params = find_unused_params
  85 
  86     def before_backward(self, loss):
  87         super().before_backward(loss)
  88 
  89         if self.find_unused_params:
  90             self.reducer.prepare_for_backward([loss])
  91         else:
  92             self.reducer.prepare_for_backward([])
  93 
  94 
  95 # Mixin goes second that way the PPO __init__ will still be called
  96 class DDPPO(PPO, DecentralizedDistributedMixin):
  97     pass
  98
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/habitat_baselines/rl/ddppo/algo/ddppo.py b/habitat_baselines/rl/ddppo/algo/ddppo.py
  --- a/habitat_baselines/rl/ddppo/algo/ddppo.py
  +++ b/habitat_baselines/rl/ddppo/algo/ddppo.py
  @@ -92,6 +92,5 @@
               self.reducer.prepare_for_backward([])
   
   
  -# Mixin goes second that way the PPO __init__ will still be called
  -class DDPPO(PPO, DecentralizedDistributedMixin):
  +class DDPPO(DecentralizedDistributedMixin, PPO):
       pass
  
verification_info: {"golden_diff": "diff --git a/habitat_baselines/rl/ddppo/algo/ddppo.py b/habitat_baselines/rl/ddppo/algo/ddppo.py\n--- a/habitat_baselines/rl/ddppo/algo/ddppo.py\n+++ b/habitat_baselines/rl/ddppo/algo/ddppo.py\n@@ -92,6 +92,5 @@\n             self.reducer.prepare_for_backward([])\n \n \n-# Mixin goes second that way the PPO __init__ will still be called\n-class DDPPO(PPO, DecentralizedDistributedMixin):\n+class DDPPO(DecentralizedDistributedMixin, PPO):\n     pass\n", "issue": "DD-PPO does not all reduce gradients\n## \ud83d\udc1b Bug\r\n\r\nDD-PPO does not all reduce gradients during the backward call, because `reducer.prepare_for_backward` is not being called during training process.\r\n\r\nThe problem is in this line: https://github.com/facebookresearch/habitat-api/blob/v0.1.4/habitat_baselines/rl/ddppo/algo/ddppo.py#L96\r\n\r\n```\r\nclass DecentralizedDistributedMixin:\r\n\r\n    ...\r\n    def before_backward(self, loss):\r\n         # ...\r\n         self.reducer.prepare_for_backward(..)\r\n\r\n\r\n# Mixin goes second that way the PPO __init__ will still be called\r\nclass DDPPO(PPO, DecentralizedDistributedMixin):  \r\n    # Here PPO and Mixin both have \"before_backward\" method, \r\n    # DDPPO will call PPO's not the Mixin's.\r\n    pass\r\n```\r\n\r\nAnd here is a quick fix:\r\n```\r\nclass DecentralizedDistributedMixin:\r\n    ...\r\n\r\n\r\n# Mixin goes second that way the PPO __init__ will still be called\r\nclass DDPPO(PPO, DecentralizedDistributedMixin):  \r\n\r\n      # Move before_backward to DDPPO\r\n      def before_backward(self, loss):\r\n         # ...\r\n         self.reducer.prepare_for_backward(..)\r\n```\r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) Facebook, Inc. and its affiliates.\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Tuple\n\nimport torch\nimport torch.distributed as distrib\n\nfrom habitat_baselines.common.rollout_storage import RolloutStorage\nfrom habitat_baselines.rl.ppo import PPO\n\nEPS_PPO = 1e-5\n\n\ndef distributed_mean_and_var(\n    values: torch.Tensor,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    r\"\"\"Computes the mean and variances of a tensor over multiple workers.\n\n    This method is equivalent to first collecting all versions of values and\n    then computing the mean and variance locally over that\n\n    :param values: (*,) shaped tensors to compute mean and variance over.  Assumed\n                        to be solely the workers local copy of this tensor,\n                        the resultant mean and variance will be computed\n                        over _all_ workers version of this tensor.\n    \"\"\"\n    assert distrib.is_initialized(), \"Distributed must be initialized\"\n\n    world_size = distrib.get_world_size()\n    mean = values.mean()\n    distrib.all_reduce(mean)\n    mean /= world_size\n\n    sq_diff = (values - mean).pow(2).mean()\n    distrib.all_reduce(sq_diff)\n    var = sq_diff / world_size\n\n    return mean, var\n\n\nclass DecentralizedDistributedMixin:\n    def _get_advantages_distributed(\n        self, rollouts: RolloutStorage\n    ) -> torch.Tensor:\n        advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]\n        if not self.use_normalized_advantage:\n            return advantages\n\n        mean, var = distributed_mean_and_var(advantages)\n\n        return (advantages - mean) / (var.sqrt() + EPS_PPO)\n\n    def init_distributed(self, find_unused_params: bool = True) -> None:\n        r\"\"\"Initializes distributed training for the model\n\n        1. Broadcasts the model weights from world_rank 0 to all other workers\n        2. Adds gradient hooks to the model\n\n        :param find_unused_params: Whether or not to filter out unused parameters\n                                   before gradient reduction.  This *must* be True if\n                                   there are any parameters in the model that where unused in the\n                                   forward pass, otherwise the gradient reduction\n                                   will not work correctly.\n        \"\"\"\n        # NB: Used to hide the hooks from the nn.Module,\n        # so they don't show up in the state_dict\n        class Guard:\n            def __init__(self, model, device):\n                if torch.cuda.is_available():\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(\n                        model, device_ids=[device], output_device=device\n                    )\n                else:\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(model)\n\n        self._ddp_hooks = Guard(self.actor_critic, self.device)\n        self.get_advantages = self._get_advantages_distributed\n\n        self.reducer = self._ddp_hooks.ddp.reducer\n        self.find_unused_params = find_unused_params\n\n    def before_backward(self, loss):\n        super().before_backward(loss)\n\n        if self.find_unused_params:\n            self.reducer.prepare_for_backward([loss])\n        else:\n            self.reducer.prepare_for_backward([])\n\n\n# Mixin goes second that way the PPO __init__ will still be called\nclass DDPPO(PPO, DecentralizedDistributedMixin):\n    pass\n", "path": "habitat_baselines/rl/ddppo/algo/ddppo.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\n# Copyright (c) Facebook, Inc. and its affiliates.\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Tuple\n\nimport torch\nimport torch.distributed as distrib\n\nfrom habitat_baselines.common.rollout_storage import RolloutStorage\nfrom habitat_baselines.rl.ppo import PPO\n\nEPS_PPO = 1e-5\n\n\ndef distributed_mean_and_var(\n    values: torch.Tensor,\n) -> Tuple[torch.Tensor, torch.Tensor]:\n    r\"\"\"Computes the mean and variances of a tensor over multiple workers.\n\n    This method is equivalent to first collecting all versions of values and\n    then computing the mean and variance locally over that\n\n    :param values: (*,) shaped tensors to compute mean and variance over.  Assumed\n                        to be solely the workers local copy of this tensor,\n                        the resultant mean and variance will be computed\n                        over _all_ workers version of this tensor.\n    \"\"\"\n    assert distrib.is_initialized(), \"Distributed must be initialized\"\n\n    world_size = distrib.get_world_size()\n    mean = values.mean()\n    distrib.all_reduce(mean)\n    mean /= world_size\n\n    sq_diff = (values - mean).pow(2).mean()\n    distrib.all_reduce(sq_diff)\n    var = sq_diff / world_size\n\n    return mean, var\n\n\nclass DecentralizedDistributedMixin:\n    def _get_advantages_distributed(\n        self, rollouts: RolloutStorage\n    ) -> torch.Tensor:\n        advantages = rollouts.returns[:-1] - rollouts.value_preds[:-1]\n        if not self.use_normalized_advantage:\n            return advantages\n\n        mean, var = distributed_mean_and_var(advantages)\n\n        return (advantages - mean) / (var.sqrt() + EPS_PPO)\n\n    def init_distributed(self, find_unused_params: bool = True) -> None:\n        r\"\"\"Initializes distributed training for the model\n\n        1. Broadcasts the model weights from world_rank 0 to all other workers\n        2. Adds gradient hooks to the model\n\n        :param find_unused_params: Whether or not to filter out unused parameters\n                                   before gradient reduction.  This *must* be True if\n                                   there are any parameters in the model that where unused in the\n                                   forward pass, otherwise the gradient reduction\n                                   will not work correctly.\n        \"\"\"\n        # NB: Used to hide the hooks from the nn.Module,\n        # so they don't show up in the state_dict\n        class Guard:\n            def __init__(self, model, device):\n                if torch.cuda.is_available():\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(\n                        model, device_ids=[device], output_device=device\n                    )\n                else:\n                    self.ddp = torch.nn.parallel.DistributedDataParallel(model)\n\n        self._ddp_hooks = Guard(self.actor_critic, self.device)\n        self.get_advantages = self._get_advantages_distributed\n\n        self.reducer = self._ddp_hooks.ddp.reducer\n        self.find_unused_params = find_unused_params\n\n    def before_backward(self, loss):\n        super().before_backward(loss)\n\n        if self.find_unused_params:\n            self.reducer.prepare_for_backward([loss])\n        else:\n            self.reducer.prepare_for_backward([])\n\n\nclass DDPPO(DecentralizedDistributedMixin, PPO):\n    pass\n", "path": "habitat_baselines/rl/ddppo/algo/ddppo.py"}]}