We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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 :)

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `benchmarks/oss.py`
Content:
```
1 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
2 
3 
4 import argparse
5 from enum import Enum
6 import importlib
7 import logging
8 import math
9 import shutil
10 import tempfile
11 import time
12 from typing import Any, List, Optional, cast
13 
14 import numpy as np
15 import torch
16 import torch.autograd.profiler as profiler
17 import torch.distributed as dist
18 import torch.multiprocessing as mp
19 import torch.nn as nn
20 from torch.nn.parallel import DistributedDataParallel as DDP
21 from torch.utils.data import BatchSampler, DataLoader, Sampler
22 from torch.utils.data.distributed import DistributedSampler
23 from torchvision.datasets import MNIST
24 from torchvision.transforms import ToTensor
25 
26 from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP
27 from fairscale.optim import OSS
28 
29 OPTIM = torch.optim.RMSprop
30 TEMPDIR = tempfile.gettempdir()
31 
32 
33 def dist_init(rank, world_size, backend):
34     logging.info(f"Using backend: {backend}")
35     dist.init_process_group(backend=backend, init_method="tcp://localhost:29501", rank=rank, world_size=world_size)
36 
37 
38 def get_problem(rank, world_size, batch_size, device, model_name: str):
39     # Select the desired model on the fly
40     logging.info(f"Using {model_name} for benchmarking")
41     model = getattr(importlib.import_module("torchvision.models"), model_name)(pretrained=False).to(device)
42 
43     # Data setup, duplicate the grey channels to get pseudo color
44     def collate(inputs: List[Any]):
45         return {
46             "inputs": torch.stack([i[0] for i in inputs]).repeat(1, 3, 1, 1).to(device),
47             "label": torch.tensor([i[1] for i in inputs]).to(device),
48         }
49 
50     dataset = MNIST(transform=ToTensor(), download=False, root=TEMPDIR)
51     sampler: Sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
52     batch_sampler = BatchSampler(sampler, batch_size, drop_last=True)
53     dataloader = DataLoader(dataset=dataset, batch_sampler=batch_sampler, collate_fn=collate)
54 
55     loss_fn = nn.CrossEntropyLoss()
56     return model, dataloader, loss_fn
57 
58 
59 class OptimType(str, Enum):
60     vanilla = "pytorch"
61     oss_ddp = "oss_ddp"
62     oss_sharded_ddp = "oss_sharded_ddp"
63     everyone = "everyone"
64 
65 
66 def train(
67     rank: int,
68     args: argparse.Namespace,
69     backend: str = "gloo",
70     optim_type: OptimType = OptimType.vanilla,
71     check_regression: bool = True,
72 ):
73     logging.basicConfig(level=logging.INFO if not args.debug else logging.DEBUG)
74 
75     # DDP
76     dist_init(rank=rank, world_size=args.world_size, backend=backend)
77 
78     # Setup
79     if not args.cpu:
80         torch.cuda.set_device(rank)
81         torch.cuda.manual_seed(0)
82     torch.manual_seed(0)  # also sets the cuda seed
83     np.random.seed(0)
84 
85     if backend == "nccl":
86         torch.backends.cudnn.deterministic = True
87         torch.backends.cudnn.benchmark = False
88 
89     device = torch.device("cpu") if args.cpu else torch.device(rank)
90     model, dataloader, loss_fn = get_problem(rank, args.world_size, args.batch_size, device, args.torchvision_model)
91 
92     # Shard the optimizer
93     optimizer: Optional[torch.optim.Optimizer] = None
94     model = cast(nn.Module, model)
95 
96     if optim_type == OptimType.oss_sharded_ddp:
97         model = ShardedDDP(
98             model,
99             optimizer=OPTIM,
100             optimizer_params={"lr": 1e-4, "momentum": 0.9},
101             world_size=args.world_size,
102             broadcast_buffers=True,
103         )
104         optimizer = model.sharded_optimizer
105 
106     else:
107         model = DDP(model, device_ids=[rank], find_unused_parameters=False)  # type: ignore
108         optimizer = (
109             OSS(params=model.parameters(), optim=OPTIM, lr=1e-4, momentum=0.9)
110             if optim_type == OptimType.oss_ddp
111             else OPTIM(model.parameters(), lr=1e-4, momentum=0.9)
112         )
113     optimizer = cast(torch.optim.Optimizer, optimizer)
114 
115     # Reset the memory use counter
116     if not args.cpu:
117         torch.cuda.reset_peak_memory_stats(rank)
118         torch.cuda.synchronize(rank)
119 
120     # Standard training loop
121     training_start = time.monotonic()
122     model.train()
123 
124     measurements = []
125     final_loss: Optional[float] = -1.0
126     need_profiling = args.profile
127 
128     for epoch in range(args.epochs):
129         n_items = 0
130         epoch_runtime = 0.0
131 
132         for batch in dataloader:
133             batch__start = time.monotonic()
134 
135             def closure():
136                 model.zero_grad()
137                 if args.debug and rank == 0 and next(model.parameters()).grad is not None:
138                     logging.debug(
139                         "\nbefore:  param {} -- grad {}".format(
140                             next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()
141                         )
142                     )
143 
144                 outputs = model(batch["inputs"])
145                 loss = loss_fn(outputs, batch["label"])
146                 loss.backward()
147 
148                 if optim_type == OptimType.oss_sharded_ddp:
149                     model.reduce()
150 
151                 if args.debug and rank == 0 and next(model.parameters()).grad is not None:
152                     logging.debug(
153                         "after BW: param {} -- grad {}".format(
154                             next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()
155                         )
156                     )
157                 return loss
158 
159             if need_profiling and not args.cpu:
160                 logging.info("Profiling the run")
161                 with profiler.profile(use_cuda=True, record_shapes=True, profile_memory=True) as prof:  # type: ignore
162                     with profiler.record_function("batch"):
163                         final_loss = optimizer.step(closure)
164                         logging.info("profiling done")
165 
166                 if rank == 0:
167                     prof.export_chrome_trace(f"{optim_type}_trace.json")
168 
169                 need_profiling = False  # only profile once
170 
171             else:
172                 final_loss = optimizer.step(closure)
173 
174             if args.debug and rank == 0:
175                 logging.debug("buffer: {}".format(next(model.buffers()).norm().item()))
176                 logging.debug(
177                     "after update: param {} -- grad {}".format(
178                         next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()
179                     )
180                 )
181 
182             n_items += args.batch_size
183 
184             batch_end = time.monotonic()
185             epoch_runtime += batch_end - batch__start
186 
187         if optim_type == OptimType.oss_ddp or optim_type == OptimType.oss_sharded_ddp:
188             # Check the checkpointing in the case of the OSS optimizer
189             # Memory usage could spill over from there
190             optimizer = cast(OSS, optimizer)
191             optimizer.consolidate_state_dict()
192             if dist.get_rank() == 0:
193                 _ = optimizer.state_dict()
194                 logging.info("... State dict collected")
195 
196         measurements.append(n_items / epoch_runtime)
197         if dist.get_rank() == 0:
198             logging.info(f"Epoch {epoch} - processed {measurements[-1]:.2f} img per sec. Loss {final_loss:.3f}")
199 
200     max_memory = -1.0
201     if not args.cpu:
202         torch.cuda.synchronize(rank)
203         max_memory = torch.cuda.max_memory_allocated(rank) / 2 ** 20
204         logging.info(f"[{dist.get_rank()}] : Peak memory {max_memory:.1f}MiB")
205 
206     training_stop = time.monotonic()
207     img_per_sec = n_items / (training_stop - training_start) * args.epochs
208     max_memory = torch.cuda.max_memory_allocated(rank) / 2 ** 20
209 
210     logging.info(f"[{dist.get_rank()}] : Training done. {img_per_sec:.2f} img per sec inc. checkpoint")
211     logging.info(f"[{dist.get_rank()}] : Peak memory {max_memory:.1f}MiB")
212 
213     # Compute the mean and average img per second
214     mean = sum(measurements) / len(measurements)
215     diff = map(lambda x: pow(x - mean, 2.0), measurements)
216     std = math.sqrt(sum(diff) / (len(measurements) - 1)) if args.epochs > 2 else -1
217     logging.info(f"[{dist.get_rank()}] : Mean speed: {mean:.2f} +/- {std:.2f}")
218 
219     if check_regression and dist.get_rank() == 0:
220         assert (mean + 3.0 * std) > args.reference_speed, "Speed regression detected"
221         assert max_memory < 1.05 * args.reference_memory, "Memory use regression detected"
222         assert abs(cast(float, final_loss) - args.reference_loss) < 1e-3, "Loss regression detected"
223 
224         logging.info("[Regression Test] VALID")
225 
226     dist.destroy_process_group()  # type: ignore
227 
228 
229 if __name__ == "__main__":
230     parser = argparse.ArgumentParser(
231         description="Benchmark the optimizer state sharding, on a typical computer vision workload"
232     )
233     parser.add_argument("--world_size", action="store", default=2, type=int)
234     parser.add_argument("--epochs", action="store", default=10, type=int)
235     parser.add_argument("--batch_size", action="store", default=256, type=int)
236     parser.add_argument("--check_regression", action="store_true", default=False)
237     parser.add_argument("--reference_speed", action="store", default=1430, type=float)
238     parser.add_argument("--reference_memory", action="store", default=1220, type=float)
239     parser.add_argument("--reference_loss", action="store", default=0.006, type=float)
240     parser.add_argument(
241         "--optim_type", type=OptimType, choices=[o.value for o in OptimType], default=OptimType.everyone
242     )
243     parser.add_argument("--gloo", action="store_true", default=False)
244     parser.add_argument("--profile", action="store_true", default=False)
245     parser.add_argument("--cpu", action="store_true", default=False)
246     parser.add_argument("--torchvision_model", type=str, help="Any torchvision model name (str)", default="resnet101")
247     parser.add_argument("--debug", action="store_true", default=False)
248 
249     args = parser.parse_args()
250 
251     logging.basicConfig(level=logging.INFO if not args.debug else logging.DEBUG)
252     logging.info(f"Benchmark arguments: {args}")
253 
254     backend = "nccl" if (not args.gloo or not torch.cuda.is_available()) and not args.cpu else "gloo"
255 
256     # Download dataset once for all processes
257     dataset, tentatives = None, 0
258     while dataset is None and tentatives < 5:
259         try:
260             dataset = MNIST(transform=None, download=True, root=TEMPDIR)
261         except (RuntimeError, EOFError) as e:
262             if isinstance(e, RuntimeError):
263                 # Corrupted data, erase and restart
264                 shutil.rmtree(TEMPDIR + "/MNIST")
265 
266             logging.warning("Failed loading dataset: ", e)
267             tentatives += 1
268 
269     if dataset is None:
270         logging.error("Could not download MNIST dataset")
271         exit(-1)
272     else:
273         logging.info("Dataset downloaded")
274 
275     # Benchmark the different configurations, via multiple processes
276     if args.optim_type == OptimType.vanilla or args.optim_type == OptimType.everyone:
277         logging.info("\n*** Benchmark vanilla optimizer")
278         mp.spawn(
279             train,
280             args=(args, backend, OptimType.vanilla, False,),  # no regression check
281             nprocs=args.world_size,
282             join=True,
283         )
284 
285     if args.optim_type == OptimType.oss_ddp or args.optim_type == OptimType.everyone:
286         logging.info("\n*** Benchmark OSS with DDP")
287         mp.spawn(
288             train, args=(args, backend, OptimType.oss_ddp, args.check_regression), nprocs=args.world_size, join=True,
289         )
290 
291     if args.optim_type == OptimType.oss_sharded_ddp or args.optim_type == OptimType.everyone:
292         logging.info("\n*** Benchmark OSS with ShardedDDP")
293         mp.spawn(
294             train,
295             args=(
296                 args,
297                 backend,
298                 OptimType.oss_sharded_ddp,
299                 False,
300             ),  # FIXME: @lefaudeux - SDP should give the same results
301             nprocs=args.world_size,
302             join=True,
303         )
304
```
--- 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()

```

diff --git a/benchmarks/oss.py b/benchmarks/oss.py
--- a/benchmarks/oss.py
+++ b/benchmarks/oss.py
@@ -140,9 +140,15 @@
                             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 @@
     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()

{"golden_diff": "diff --git a/benchmarks/oss.py b/benchmarks/oss.py\n--- a/benchmarks/oss.py\n+++ b/benchmarks/oss.py\n@@ -140,9 +140,15 @@\n                             next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()\n                         )\n                     )\n+                if not args.cpu and args.amp:\n+                    # Automatically computes the FW pass in half precision\n+                    with torch.cuda.amp.autocast():\n+                        outputs = model(batch[\"inputs\"])\n+                        loss = loss_fn(outputs, batch[\"label\"])\n+                else:\n+                    outputs = model(batch[\"inputs\"])\n+                    loss = loss_fn(outputs, batch[\"label\"])\n \n-                outputs = model(batch[\"inputs\"])\n-                loss = loss_fn(outputs, batch[\"label\"])\n                 loss.backward()\n \n                 if optim_type == OptimType.oss_sharded_ddp:\n@@ -244,7 +250,8 @@\n     parser.add_argument(\"--profile\", action=\"store_true\", default=False)\n     parser.add_argument(\"--cpu\", action=\"store_true\", default=False)\n     parser.add_argument(\"--torchvision_model\", type=str, help=\"Any torchvision model name (str)\", default=\"resnet101\")\n-    parser.add_argument(\"--debug\", action=\"store_true\", default=False)\n+    parser.add_argument(\"--debug\", action=\"store_true\", default=False, help=\"Display additional debug information\")\n+    parser.add_argument(\"--amp\", action=\"store_true\", default=False, help=\"Activate torch AMP\")\n \n     args = parser.parse_args()\n", "issue": "Add mixed precision autocast to benchmarking script \n## \ud83d\ude80 Feature\r\n\r\nIt 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!\r\n\r\n## Motivation\r\n\r\nIntegrating OSS/ShardedDDP into PL :)\n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n\n\nimport argparse\nfrom enum import Enum\nimport importlib\nimport logging\nimport math\nimport shutil\nimport tempfile\nimport time\nfrom typing import Any, List, Optional, cast\n\nimport numpy as np\nimport torch\nimport torch.autograd.profiler as profiler\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nimport torch.nn as nn\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data import BatchSampler, DataLoader, Sampler\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torchvision.datasets import MNIST\nfrom torchvision.transforms import ToTensor\n\nfrom fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP\nfrom fairscale.optim import OSS\n\nOPTIM = torch.optim.RMSprop\nTEMPDIR = tempfile.gettempdir()\n\n\ndef dist_init(rank, world_size, backend):\n    logging.info(f\"Using backend: {backend}\")\n    dist.init_process_group(backend=backend, init_method=\"tcp://localhost:29501\", rank=rank, world_size=world_size)\n\n\ndef get_problem(rank, world_size, batch_size, device, model_name: str):\n    # Select the desired model on the fly\n    logging.info(f\"Using {model_name} for benchmarking\")\n    model = getattr(importlib.import_module(\"torchvision.models\"), model_name)(pretrained=False).to(device)\n\n    # Data setup, duplicate the grey channels to get pseudo color\n    def collate(inputs: List[Any]):\n        return {\n            \"inputs\": torch.stack([i[0] for i in inputs]).repeat(1, 3, 1, 1).to(device),\n            \"label\": torch.tensor([i[1] for i in inputs]).to(device),\n        }\n\n    dataset = MNIST(transform=ToTensor(), download=False, root=TEMPDIR)\n    sampler: Sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)\n    batch_sampler = BatchSampler(sampler, batch_size, drop_last=True)\n    dataloader = DataLoader(dataset=dataset, batch_sampler=batch_sampler, collate_fn=collate)\n\n    loss_fn = nn.CrossEntropyLoss()\n    return model, dataloader, loss_fn\n\n\nclass OptimType(str, Enum):\n    vanilla = \"pytorch\"\n    oss_ddp = \"oss_ddp\"\n    oss_sharded_ddp = \"oss_sharded_ddp\"\n    everyone = \"everyone\"\n\n\ndef train(\n    rank: int,\n    args: argparse.Namespace,\n    backend: str = \"gloo\",\n    optim_type: OptimType = OptimType.vanilla,\n    check_regression: bool = True,\n):\n    logging.basicConfig(level=logging.INFO if not args.debug else logging.DEBUG)\n\n    # DDP\n    dist_init(rank=rank, world_size=args.world_size, backend=backend)\n\n    # Setup\n    if not args.cpu:\n        torch.cuda.set_device(rank)\n        torch.cuda.manual_seed(0)\n    torch.manual_seed(0)  # also sets the cuda seed\n    np.random.seed(0)\n\n    if backend == \"nccl\":\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n\n    device = torch.device(\"cpu\") if args.cpu else torch.device(rank)\n    model, dataloader, loss_fn = get_problem(rank, args.world_size, args.batch_size, device, args.torchvision_model)\n\n    # Shard the optimizer\n    optimizer: Optional[torch.optim.Optimizer] = None\n    model = cast(nn.Module, model)\n\n    if optim_type == OptimType.oss_sharded_ddp:\n        model = ShardedDDP(\n            model,\n            optimizer=OPTIM,\n            optimizer_params={\"lr\": 1e-4, \"momentum\": 0.9},\n            world_size=args.world_size,\n            broadcast_buffers=True,\n        )\n        optimizer = model.sharded_optimizer\n\n    else:\n        model = DDP(model, device_ids=[rank], find_unused_parameters=False)  # type: ignore\n        optimizer = (\n            OSS(params=model.parameters(), optim=OPTIM, lr=1e-4, momentum=0.9)\n            if optim_type == OptimType.oss_ddp\n            else OPTIM(model.parameters(), lr=1e-4, momentum=0.9)\n        )\n    optimizer = cast(torch.optim.Optimizer, optimizer)\n\n    # Reset the memory use counter\n    if not args.cpu:\n        torch.cuda.reset_peak_memory_stats(rank)\n        torch.cuda.synchronize(rank)\n\n    # Standard training loop\n    training_start = time.monotonic()\n    model.train()\n\n    measurements = []\n    final_loss: Optional[float] = -1.0\n    need_profiling = args.profile\n\n    for epoch in range(args.epochs):\n        n_items = 0\n        epoch_runtime = 0.0\n\n        for batch in dataloader:\n            batch__start = time.monotonic()\n\n            def closure():\n                model.zero_grad()\n                if args.debug and rank == 0 and next(model.parameters()).grad is not None:\n                    logging.debug(\n                        \"\\nbefore:  param {} -- grad {}\".format(\n                            next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()\n                        )\n                    )\n\n                outputs = model(batch[\"inputs\"])\n                loss = loss_fn(outputs, batch[\"label\"])\n                loss.backward()\n\n                if optim_type == OptimType.oss_sharded_ddp:\n                    model.reduce()\n\n                if args.debug and rank == 0 and next(model.parameters()).grad is not None:\n                    logging.debug(\n                        \"after BW: param {} -- grad {}\".format(\n                            next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()\n                        )\n                    )\n                return loss\n\n            if need_profiling and not args.cpu:\n                logging.info(\"Profiling the run\")\n                with profiler.profile(use_cuda=True, record_shapes=True, profile_memory=True) as prof:  # type: ignore\n                    with profiler.record_function(\"batch\"):\n                        final_loss = optimizer.step(closure)\n                        logging.info(\"profiling done\")\n\n                if rank == 0:\n                    prof.export_chrome_trace(f\"{optim_type}_trace.json\")\n\n                need_profiling = False  # only profile once\n\n            else:\n                final_loss = optimizer.step(closure)\n\n            if args.debug and rank == 0:\n                logging.debug(\"buffer: {}\".format(next(model.buffers()).norm().item()))\n                logging.debug(\n                    \"after update: param {} -- grad {}\".format(\n                        next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()\n                    )\n                )\n\n            n_items += args.batch_size\n\n            batch_end = time.monotonic()\n            epoch_runtime += batch_end - batch__start\n\n        if optim_type == OptimType.oss_ddp or optim_type == OptimType.oss_sharded_ddp:\n            # Check the checkpointing in the case of the OSS optimizer\n            # Memory usage could spill over from there\n            optimizer = cast(OSS, optimizer)\n            optimizer.consolidate_state_dict()\n            if dist.get_rank() == 0:\n                _ = optimizer.state_dict()\n                logging.info(\"... State dict collected\")\n\n        measurements.append(n_items / epoch_runtime)\n        if dist.get_rank() == 0:\n            logging.info(f\"Epoch {epoch} - processed {measurements[-1]:.2f} img per sec. Loss {final_loss:.3f}\")\n\n    max_memory = -1.0\n    if not args.cpu:\n        torch.cuda.synchronize(rank)\n        max_memory = torch.cuda.max_memory_allocated(rank) / 2 ** 20\n        logging.info(f\"[{dist.get_rank()}] : Peak memory {max_memory:.1f}MiB\")\n\n    training_stop = time.monotonic()\n    img_per_sec = n_items / (training_stop - training_start) * args.epochs\n    max_memory = torch.cuda.max_memory_allocated(rank) / 2 ** 20\n\n    logging.info(f\"[{dist.get_rank()}] : Training done. {img_per_sec:.2f} img per sec inc. checkpoint\")\n    logging.info(f\"[{dist.get_rank()}] : Peak memory {max_memory:.1f}MiB\")\n\n    # Compute the mean and average img per second\n    mean = sum(measurements) / len(measurements)\n    diff = map(lambda x: pow(x - mean, 2.0), measurements)\n    std = math.sqrt(sum(diff) / (len(measurements) - 1)) if args.epochs > 2 else -1\n    logging.info(f\"[{dist.get_rank()}] : Mean speed: {mean:.2f} +/- {std:.2f}\")\n\n    if check_regression and dist.get_rank() == 0:\n        assert (mean + 3.0 * std) > args.reference_speed, \"Speed regression detected\"\n        assert max_memory < 1.05 * args.reference_memory, \"Memory use regression detected\"\n        assert abs(cast(float, final_loss) - args.reference_loss) < 1e-3, \"Loss regression detected\"\n\n        logging.info(\"[Regression Test] VALID\")\n\n    dist.destroy_process_group()  # type: ignore\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(\n        description=\"Benchmark the optimizer state sharding, on a typical computer vision workload\"\n    )\n    parser.add_argument(\"--world_size\", action=\"store\", default=2, type=int)\n    parser.add_argument(\"--epochs\", action=\"store\", default=10, type=int)\n    parser.add_argument(\"--batch_size\", action=\"store\", default=256, type=int)\n    parser.add_argument(\"--check_regression\", action=\"store_true\", default=False)\n    parser.add_argument(\"--reference_speed\", action=\"store\", default=1430, type=float)\n    parser.add_argument(\"--reference_memory\", action=\"store\", default=1220, type=float)\n    parser.add_argument(\"--reference_loss\", action=\"store\", default=0.006, type=float)\n    parser.add_argument(\n        \"--optim_type\", type=OptimType, choices=[o.value for o in OptimType], default=OptimType.everyone\n    )\n    parser.add_argument(\"--gloo\", action=\"store_true\", default=False)\n    parser.add_argument(\"--profile\", action=\"store_true\", default=False)\n    parser.add_argument(\"--cpu\", action=\"store_true\", default=False)\n    parser.add_argument(\"--torchvision_model\", type=str, help=\"Any torchvision model name (str)\", default=\"resnet101\")\n    parser.add_argument(\"--debug\", action=\"store_true\", default=False)\n\n    args = parser.parse_args()\n\n    logging.basicConfig(level=logging.INFO if not args.debug else logging.DEBUG)\n    logging.info(f\"Benchmark arguments: {args}\")\n\n    backend = \"nccl\" if (not args.gloo or not torch.cuda.is_available()) and not args.cpu else \"gloo\"\n\n    # Download dataset once for all processes\n    dataset, tentatives = None, 0\n    while dataset is None and tentatives < 5:\n        try:\n            dataset = MNIST(transform=None, download=True, root=TEMPDIR)\n        except (RuntimeError, EOFError) as e:\n            if isinstance(e, RuntimeError):\n                # Corrupted data, erase and restart\n                shutil.rmtree(TEMPDIR + \"/MNIST\")\n\n            logging.warning(\"Failed loading dataset: \", e)\n            tentatives += 1\n\n    if dataset is None:\n        logging.error(\"Could not download MNIST dataset\")\n        exit(-1)\n    else:\n        logging.info(\"Dataset downloaded\")\n\n    # Benchmark the different configurations, via multiple processes\n    if args.optim_type == OptimType.vanilla or args.optim_type == OptimType.everyone:\n        logging.info(\"\\n*** Benchmark vanilla optimizer\")\n        mp.spawn(\n            train,\n            args=(args, backend, OptimType.vanilla, False,),  # no regression check\n            nprocs=args.world_size,\n            join=True,\n        )\n\n    if args.optim_type == OptimType.oss_ddp or args.optim_type == OptimType.everyone:\n        logging.info(\"\\n*** Benchmark OSS with DDP\")\n        mp.spawn(\n            train, args=(args, backend, OptimType.oss_ddp, args.check_regression), nprocs=args.world_size, join=True,\n        )\n\n    if args.optim_type == OptimType.oss_sharded_ddp or args.optim_type == OptimType.everyone:\n        logging.info(\"\\n*** Benchmark OSS with ShardedDDP\")\n        mp.spawn(\n            train,\n            args=(\n                args,\n                backend,\n                OptimType.oss_sharded_ddp,\n                False,\n            ),  # FIXME: @lefaudeux - SDP should give the same results\n            nprocs=args.world_size,\n            join=True,\n        )\n", "path": "benchmarks/oss.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.\n\n\nimport argparse\nfrom enum import Enum\nimport importlib\nimport logging\nimport math\nimport shutil\nimport tempfile\nimport time\nfrom typing import Any, List, Optional, cast\n\nimport numpy as np\nimport torch\nimport torch.autograd.profiler as profiler\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nimport torch.nn as nn\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils.data import BatchSampler, DataLoader, Sampler\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torchvision.datasets import MNIST\nfrom torchvision.transforms import ToTensor\n\nfrom fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP\nfrom fairscale.optim import OSS\n\nOPTIM = torch.optim.RMSprop\nTEMPDIR = tempfile.gettempdir()\n\n\ndef dist_init(rank, world_size, backend):\n    logging.info(f\"Using backend: {backend}\")\n    dist.init_process_group(backend=backend, init_method=\"tcp://localhost:29501\", rank=rank, world_size=world_size)\n\n\ndef get_problem(rank, world_size, batch_size, device, model_name: str):\n    # Select the desired model on the fly\n    logging.info(f\"Using {model_name} for benchmarking\")\n    model = getattr(importlib.import_module(\"torchvision.models\"), model_name)(pretrained=False).to(device)\n\n    # Data setup, duplicate the grey channels to get pseudo color\n    def collate(inputs: List[Any]):\n        return {\n            \"inputs\": torch.stack([i[0] for i in inputs]).repeat(1, 3, 1, 1).to(device),\n            \"label\": torch.tensor([i[1] for i in inputs]).to(device),\n        }\n\n    dataset = MNIST(transform=ToTensor(), download=False, root=TEMPDIR)\n    sampler: Sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)\n    batch_sampler = BatchSampler(sampler, batch_size, drop_last=True)\n    dataloader = DataLoader(dataset=dataset, batch_sampler=batch_sampler, collate_fn=collate)\n\n    loss_fn = nn.CrossEntropyLoss()\n    return model, dataloader, loss_fn\n\n\nclass OptimType(str, Enum):\n    vanilla = \"pytorch\"\n    oss_ddp = \"oss_ddp\"\n    oss_sharded_ddp = \"oss_sharded_ddp\"\n    everyone = \"everyone\"\n\n\ndef train(\n    rank: int,\n    args: argparse.Namespace,\n    backend: str = \"gloo\",\n    optim_type: OptimType = OptimType.vanilla,\n    check_regression: bool = True,\n):\n    logging.basicConfig(level=logging.INFO if not args.debug else logging.DEBUG)\n\n    # DDP\n    dist_init(rank=rank, world_size=args.world_size, backend=backend)\n\n    # Setup\n    if not args.cpu:\n        torch.cuda.set_device(rank)\n        torch.cuda.manual_seed(0)\n    torch.manual_seed(0)  # also sets the cuda seed\n    np.random.seed(0)\n\n    if backend == \"nccl\":\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\n\n    device = torch.device(\"cpu\") if args.cpu else torch.device(rank)\n    model, dataloader, loss_fn = get_problem(rank, args.world_size, args.batch_size, device, args.torchvision_model)\n\n    # Shard the optimizer\n    optimizer: Optional[torch.optim.Optimizer] = None\n    model = cast(nn.Module, model)\n\n    if optim_type == OptimType.oss_sharded_ddp:\n        model = ShardedDDP(\n            model,\n            optimizer=OPTIM,\n            optimizer_params={\"lr\": 1e-4, \"momentum\": 0.9},\n            world_size=args.world_size,\n            broadcast_buffers=True,\n        )\n        optimizer = model.sharded_optimizer\n\n    else:\n        model = DDP(model, device_ids=[rank], find_unused_parameters=False)  # type: ignore\n        optimizer = (\n            OSS(params=model.parameters(), optim=OPTIM, lr=1e-4, momentum=0.9)\n            if optim_type == OptimType.oss_ddp\n            else OPTIM(model.parameters(), lr=1e-4, momentum=0.9)\n        )\n    optimizer = cast(torch.optim.Optimizer, optimizer)\n\n    # Reset the memory use counter\n    if not args.cpu:\n        torch.cuda.reset_peak_memory_stats(rank)\n        torch.cuda.synchronize(rank)\n\n    # Standard training loop\n    training_start = time.monotonic()\n    model.train()\n\n    measurements = []\n    final_loss: Optional[float] = -1.0\n    need_profiling = args.profile\n\n    for epoch in range(args.epochs):\n        n_items = 0\n        epoch_runtime = 0.0\n\n        for batch in dataloader:\n            batch__start = time.monotonic()\n\n            def closure():\n                model.zero_grad()\n                if args.debug and rank == 0 and next(model.parameters()).grad is not None:\n                    logging.debug(\n                        \"\\nbefore:  param {} -- grad {}\".format(\n                            next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()\n                        )\n                    )\n                if not args.cpu and args.amp:\n                    # Automatically computes the FW pass in half precision\n                    with torch.cuda.amp.autocast():\n                        outputs = model(batch[\"inputs\"])\n                        loss = loss_fn(outputs, batch[\"label\"])\n                else:\n                    outputs = model(batch[\"inputs\"])\n                    loss = loss_fn(outputs, batch[\"label\"])\n\n                loss.backward()\n\n                if optim_type == OptimType.oss_sharded_ddp:\n                    model.reduce()\n\n                if args.debug and rank == 0 and next(model.parameters()).grad is not None:\n                    logging.debug(\n                        \"after BW: param {} -- grad {}\".format(\n                            next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()\n                        )\n                    )\n                return loss\n\n            if need_profiling and not args.cpu:\n                logging.info(\"Profiling the run\")\n                with profiler.profile(use_cuda=True, record_shapes=True, profile_memory=True) as prof:  # type: ignore\n                    with profiler.record_function(\"batch\"):\n                        final_loss = optimizer.step(closure)\n                        logging.info(\"profiling done\")\n\n                if rank == 0:\n                    prof.export_chrome_trace(f\"{optim_type}_trace.json\")\n\n                need_profiling = False  # only profile once\n\n            else:\n                final_loss = optimizer.step(closure)\n\n            if args.debug and rank == 0:\n                logging.debug(\"buffer: {}\".format(next(model.buffers()).norm().item()))\n                logging.debug(\n                    \"after update: param {} -- grad {}\".format(\n                        next(model.parameters()).norm().item(), next(model.parameters()).grad.norm().item()\n                    )\n                )\n\n            n_items += args.batch_size\n\n            batch_end = time.monotonic()\n            epoch_runtime += batch_end - batch__start\n\n        if optim_type == OptimType.oss_ddp or optim_type == OptimType.oss_sharded_ddp:\n            # Check the checkpointing in the case of the OSS optimizer\n            # Memory usage could spill over from there\n            optimizer = cast(OSS, optimizer)\n            optimizer.consolidate_state_dict()\n            if dist.get_rank() == 0:\n                _ = optimizer.state_dict()\n                logging.info(\"... State dict collected\")\n\n        measurements.append(n_items / epoch_runtime)\n        if dist.get_rank() == 0:\n            logging.info(f\"Epoch {epoch} - processed {measurements[-1]:.2f} img per sec. Loss {final_loss:.3f}\")\n\n    max_memory = -1.0\n    if not args.cpu:\n        torch.cuda.synchronize(rank)\n        max_memory = torch.cuda.max_memory_allocated(rank) / 2 ** 20\n        logging.info(f\"[{dist.get_rank()}] : Peak memory {max_memory:.1f}MiB\")\n\n    training_stop = time.monotonic()\n    img_per_sec = n_items / (training_stop - training_start) * args.epochs\n    max_memory = torch.cuda.max_memory_allocated(rank) / 2 ** 20\n\n    logging.info(f\"[{dist.get_rank()}] : Training done. {img_per_sec:.2f} img per sec inc. checkpoint\")\n    logging.info(f\"[{dist.get_rank()}] : Peak memory {max_memory:.1f}MiB\")\n\n    # Compute the mean and average img per second\n    mean = sum(measurements) / len(measurements)\n    diff = map(lambda x: pow(x - mean, 2.0), measurements)\n    std = math.sqrt(sum(diff) / (len(measurements) - 1)) if args.epochs > 2 else -1\n    logging.info(f\"[{dist.get_rank()}] : Mean speed: {mean:.2f} +/- {std:.2f}\")\n\n    if check_regression and dist.get_rank() == 0:\n        assert (mean + 3.0 * std) > args.reference_speed, \"Speed regression detected\"\n        assert max_memory < 1.05 * args.reference_memory, \"Memory use regression detected\"\n        assert abs(cast(float, final_loss) - args.reference_loss) < 1e-3, \"Loss regression detected\"\n\n        logging.info(\"[Regression Test] VALID\")\n\n    dist.destroy_process_group()  # type: ignore\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(\n        description=\"Benchmark the optimizer state sharding, on a typical computer vision workload\"\n    )\n    parser.add_argument(\"--world_size\", action=\"store\", default=2, type=int)\n    parser.add_argument(\"--epochs\", action=\"store\", default=10, type=int)\n    parser.add_argument(\"--batch_size\", action=\"store\", default=256, type=int)\n    parser.add_argument(\"--check_regression\", action=\"store_true\", default=False)\n    parser.add_argument(\"--reference_speed\", action=\"store\", default=1430, type=float)\n    parser.add_argument(\"--reference_memory\", action=\"store\", default=1220, type=float)\n    parser.add_argument(\"--reference_loss\", action=\"store\", default=0.006, type=float)\n    parser.add_argument(\n        \"--optim_type\", type=OptimType, choices=[o.value for o in OptimType], default=OptimType.everyone\n    )\n    parser.add_argument(\"--gloo\", action=\"store_true\", default=False)\n    parser.add_argument(\"--profile\", action=\"store_true\", default=False)\n    parser.add_argument(\"--cpu\", action=\"store_true\", default=False)\n    parser.add_argument(\"--torchvision_model\", type=str, help=\"Any torchvision model name (str)\", default=\"resnet101\")\n    parser.add_argument(\"--debug\", action=\"store_true\", default=False, help=\"Display additional debug information\")\n    parser.add_argument(\"--amp\", action=\"store_true\", default=False, help=\"Activate torch AMP\")\n\n    args = parser.parse_args()\n\n    logging.basicConfig(level=logging.INFO if not args.debug else logging.DEBUG)\n    logging.info(f\"Benchmark arguments: {args}\")\n\n    backend = \"nccl\" if (not args.gloo or not torch.cuda.is_available()) and not args.cpu else \"gloo\"\n\n    # Download dataset once for all processes\n    dataset, tentatives = None, 0\n    while dataset is None and tentatives < 5:\n        try:\n            dataset = MNIST(transform=None, download=True, root=TEMPDIR)\n        except (RuntimeError, EOFError) as e:\n            if isinstance(e, RuntimeError):\n                # Corrupted data, erase and restart\n                shutil.rmtree(TEMPDIR + \"/MNIST\")\n\n            logging.warning(\"Failed loading dataset: \", e)\n            tentatives += 1\n\n    if dataset is None:\n        logging.error(\"Could not download MNIST dataset\")\n        exit(-1)\n    else:\n        logging.info(\"Dataset downloaded\")\n\n    # Benchmark the different configurations, via multiple processes\n    if args.optim_type == OptimType.vanilla or args.optim_type == OptimType.everyone:\n        logging.info(\"\\n*** Benchmark vanilla optimizer\")\n        mp.spawn(\n            train,\n            args=(args, backend, OptimType.vanilla, False,),  # no regression check\n            nprocs=args.world_size,\n            join=True,\n        )\n\n    if args.optim_type == OptimType.oss_ddp or args.optim_type == OptimType.everyone:\n        logging.info(\"\\n*** Benchmark OSS with DDP\")\n        mp.spawn(\n            train, args=(args, backend, OptimType.oss_ddp, args.check_regression), nprocs=args.world_size, join=True,\n        )\n\n    if args.optim_type == OptimType.oss_sharded_ddp or args.optim_type == OptimType.everyone:\n        logging.info(\"\\n*** Benchmark OSS with ShardedDDP\")\n        mp.spawn(\n            train,\n            args=(\n                args,\n                backend,\n                OptimType.oss_sharded_ddp,\n                False,\n            ),  # FIXME: @lefaudeux - SDP should give the same results\n            nprocs=args.world_size,\n            join=True,\n        )\n", "path": "benchmarks/oss.py"}]}