prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <issue>
  The GPU load is unbalanced
  https://github.com/pytorch/examples/blob/2ee8d43dbe420be152fd5ce0d80b43b419a0e352/distributed/ddp-tutorial-series/multigpu_torchrun.py#L39
  When I run the code and resume from a existed .pt file. The memory usage of GPU0 is significantly higher than other GPUs. 
  It can be solved by adding a parameter "map_location".
  `snapshot = torch.load(snapshot_path, map_location=torch.device('cuda', int(os.environ["LOCAL_RANK"])))`
  ## My Environment
  cudatoolkit 10.2
  pytorch 12.1
  
  
  </issue>
  <code>
  [start of distributed/ddp-tutorial-series/multinode.py]
  1 import torch
  2 import torch.nn.functional as F
  3 from torch.utils.data import Dataset, DataLoader
  4 from datautils import MyTrainDataset
  5 
  6 import torch.multiprocessing as mp
  7 from torch.utils.data.distributed import DistributedSampler
  8 from torch.nn.parallel import DistributedDataParallel as DDP
  9 from torch.distributed import init_process_group, destroy_process_group
  10 import os
  11 
  12 
  13 def ddp_setup():
  14     init_process_group(backend="nccl")
  15 
  16 class Trainer:
  17     def __init__(
  18         self,
  19         model: torch.nn.Module,
  20         train_data: DataLoader,
  21         optimizer: torch.optim.Optimizer,
  22         save_every: int,
  23         snapshot_path: str,
  24     ) -> None:
  25         self.local_rank = int(os.environ["LOCAL_RANK"])
  26         self.global_rank = int(os.environ["RANK"])
  27         self.model = model.to(self.local_rank)
  28         self.train_data = train_data
  29         self.optimizer = optimizer
  30         self.save_every = save_every
  31         self.epochs_run = 0
  32         self.snapshot_path = snapshot_path
  33         if os.path.exists(snapshot_path):
  34             print("Loading snapshot")
  35             self._load_snapshot(snapshot_path)
  36 
  37         self.model = DDP(self.model, device_ids=[self.local_rank])
  38 
  39     def _load_snapshot(self, snapshot_path):
  40         snapshot = torch.load(snapshot_path)
  41         self.model.load_state_dict(snapshot["MODEL_STATE"])
  42         self.epochs_run = snapshot["EPOCHS_RUN"]
  43         print(f"Resuming training from snapshot at Epoch {self.epochs_run}")
  44 
  45     def _run_batch(self, source, targets):
  46         self.optimizer.zero_grad()
  47         output = self.model(source)
  48         loss = F.cross_entropy(output, targets)
  49         loss.backward()
  50         self.optimizer.step()
  51 
  52     def _run_epoch(self, epoch):
  53         b_sz = len(next(iter(self.train_data))[0])
  54         print(f"[GPU{self.global_rank}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}")
  55         self.train_data.sampler.set_epoch(epoch)
  56         for source, targets in self.train_data:
  57             source = source.to(self.local_rank)
  58             targets = targets.to(self.local_rank)
  59             self._run_batch(source, targets)
  60 
  61     def _save_snapshot(self, epoch):
  62         snapshot = {
  63             "MODEL_STATE": self.model.module.state_dict(),
  64             "EPOCHS_RUN": epoch,
  65         }
  66         torch.save(snapshot, self.snapshot_path)
  67         print(f"Epoch {epoch} | Training snapshot saved at {self.snapshot_path}")
  68 
  69     def train(self, max_epochs: int):
  70         for epoch in range(self.epochs_run, max_epochs):
  71             self._run_epoch(epoch)
  72             if self.local_rank == 0 and epoch % self.save_every == 0:
  73                 self._save_snapshot(epoch)
  74 
  75 
  76 def load_train_objs():
  77     train_set = MyTrainDataset(2048)  # load your dataset
  78     model = torch.nn.Linear(20, 1)  # load your model
  79     optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
  80     return train_set, model, optimizer
  81 
  82 
  83 def prepare_dataloader(dataset: Dataset, batch_size: int):
  84     return DataLoader(
  85         dataset,
  86         batch_size=batch_size,
  87         pin_memory=True,
  88         shuffle=False,
  89         sampler=DistributedSampler(dataset)
  90     )
  91 
  92 
  93 def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = "snapshot.pt"):
  94     ddp_setup()
  95     dataset, model, optimizer = load_train_objs()
  96     train_data = prepare_dataloader(dataset, batch_size)
  97     trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path)
  98     trainer.train(total_epochs)
  99     destroy_process_group()
  100 
  101 
  102 if __name__ == "__main__":
  103     import argparse
  104     parser = argparse.ArgumentParser(description='simple distributed training job')
  105     parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')
  106     parser.add_argument('save_every', type=int, help='How often to save a snapshot')
  107     parser.add_argument('--batch_size', default=32, help='Input batch size on each device (default: 32)')
  108     args = parser.parse_args()
  109     
  110     main(args.save_every, args.total_epochs, args.batch_size)
  111
  [end of distributed/ddp-tutorial-series/multinode.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: diff --git a/distributed/ddp-tutorial-series/multinode.py b/distributed/ddp-tutorial-series/multinode.py
  --- a/distributed/ddp-tutorial-series/multinode.py
  +++ b/distributed/ddp-tutorial-series/multinode.py
  @@ -37,7 +37,8 @@
           self.model = DDP(self.model, device_ids=[self.local_rank])
   
       def _load_snapshot(self, snapshot_path):
  -        snapshot = torch.load(snapshot_path)
  +        loc = f"cuda:{self.gpu_id}"
  +        snapshot = torch.load(snapshot_path, map_location=loc)
           self.model.load_state_dict(snapshot["MODEL_STATE"])
           self.epochs_run = snapshot["EPOCHS_RUN"]
           print(f"Resuming training from snapshot at Epoch {self.epochs_run}")
  
verification_info: {"golden_diff": "diff --git a/distributed/ddp-tutorial-series/multinode.py b/distributed/ddp-tutorial-series/multinode.py\n--- a/distributed/ddp-tutorial-series/multinode.py\n+++ b/distributed/ddp-tutorial-series/multinode.py\n@@ -37,7 +37,8 @@\n         self.model = DDP(self.model, device_ids=[self.local_rank])\n \n     def _load_snapshot(self, snapshot_path):\n-        snapshot = torch.load(snapshot_path)\n+        loc = f\"cuda:{self.gpu_id}\"\n+        snapshot = torch.load(snapshot_path, map_location=loc)\n         self.model.load_state_dict(snapshot[\"MODEL_STATE\"])\n         self.epochs_run = snapshot[\"EPOCHS_RUN\"]\n         print(f\"Resuming training from snapshot at Epoch {self.epochs_run}\")\n", "issue": "The GPU load is unbalanced\nhttps://github.com/pytorch/examples/blob/2ee8d43dbe420be152fd5ce0d80b43b419a0e352/distributed/ddp-tutorial-series/multigpu_torchrun.py#L39\r\nWhen I run the code and resume from a existed .pt file. The memory usage of GPU0 is significantly higher than other GPUs. \r\nIt can be solved by adding a parameter \"map_location\".\r\n`snapshot = torch.load(snapshot_path, map_location=torch.device('cuda', int(os.environ[\"LOCAL_RANK\"])))`\r\n## My Environment\r\ncudatoolkit 10.2\r\npytorch 12.1\r\n\n", "before_files": [{"content": "import torch\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom datautils import MyTrainDataset\n\nimport torch.multiprocessing as mp\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.distributed import init_process_group, destroy_process_group\nimport os\n\n\ndef ddp_setup():\n    init_process_group(backend=\"nccl\")\n\nclass Trainer:\n    def __init__(\n        self,\n        model: torch.nn.Module,\n        train_data: DataLoader,\n        optimizer: torch.optim.Optimizer,\n        save_every: int,\n        snapshot_path: str,\n    ) -> None:\n        self.local_rank = int(os.environ[\"LOCAL_RANK\"])\n        self.global_rank = int(os.environ[\"RANK\"])\n        self.model = model.to(self.local_rank)\n        self.train_data = train_data\n        self.optimizer = optimizer\n        self.save_every = save_every\n        self.epochs_run = 0\n        self.snapshot_path = snapshot_path\n        if os.path.exists(snapshot_path):\n            print(\"Loading snapshot\")\n            self._load_snapshot(snapshot_path)\n\n        self.model = DDP(self.model, device_ids=[self.local_rank])\n\n    def _load_snapshot(self, snapshot_path):\n        snapshot = torch.load(snapshot_path)\n        self.model.load_state_dict(snapshot[\"MODEL_STATE\"])\n        self.epochs_run = snapshot[\"EPOCHS_RUN\"]\n        print(f\"Resuming training from snapshot at Epoch {self.epochs_run}\")\n\n    def _run_batch(self, source, targets):\n        self.optimizer.zero_grad()\n        output = self.model(source)\n        loss = F.cross_entropy(output, targets)\n        loss.backward()\n        self.optimizer.step()\n\n    def _run_epoch(self, epoch):\n        b_sz = len(next(iter(self.train_data))[0])\n        print(f\"[GPU{self.global_rank}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}\")\n        self.train_data.sampler.set_epoch(epoch)\n        for source, targets in self.train_data:\n            source = source.to(self.local_rank)\n            targets = targets.to(self.local_rank)\n            self._run_batch(source, targets)\n\n    def _save_snapshot(self, epoch):\n        snapshot = {\n            \"MODEL_STATE\": self.model.module.state_dict(),\n            \"EPOCHS_RUN\": epoch,\n        }\n        torch.save(snapshot, self.snapshot_path)\n        print(f\"Epoch {epoch} | Training snapshot saved at {self.snapshot_path}\")\n\n    def train(self, max_epochs: int):\n        for epoch in range(self.epochs_run, max_epochs):\n            self._run_epoch(epoch)\n            if self.local_rank == 0 and epoch % self.save_every == 0:\n                self._save_snapshot(epoch)\n\n\ndef load_train_objs():\n    train_set = MyTrainDataset(2048)  # load your dataset\n    model = torch.nn.Linear(20, 1)  # load your model\n    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)\n    return train_set, model, optimizer\n\n\ndef prepare_dataloader(dataset: Dataset, batch_size: int):\n    return DataLoader(\n        dataset,\n        batch_size=batch_size,\n        pin_memory=True,\n        shuffle=False,\n        sampler=DistributedSampler(dataset)\n    )\n\n\ndef main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str = \"snapshot.pt\"):\n    ddp_setup()\n    dataset, model, optimizer = load_train_objs()\n    train_data = prepare_dataloader(dataset, batch_size)\n    trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path)\n    trainer.train(total_epochs)\n    destroy_process_group()\n\n\nif __name__ == \"__main__\":\n    import argparse\n    parser = argparse.ArgumentParser(description='simple distributed training job')\n    parser.add_argument('total_epochs', type=int, help='Total epochs to train the model')\n    parser.add_argument('save_every', type=int, help='How often to save a snapshot')\n    parser.add_argument('--batch_size', default=32, help='Input batch size on each device (default: 32)')\n    args = parser.parse_args()\n    \n    main(args.save_every, args.total_epochs, args.batch_size)\n", "path": "distributed/ddp-tutorial-series/multinode.py"}]}