We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG]: GradAccumDataloader will ignore the last iter.
### 🐛 Describe the bug

I used the list as dataloader for testing.

```
class GradAccumDataloader:
    """A wrapper for dataloader to enable gradient accumulation by dropping the last incomplete steps.

    Note:
        The dataloader would drop the last incomplete steps for gradient accumulation.
        For example, if a dataloader has 10 batches of data and accumulate size is 4. The model parameters will
        be updated only twice at step 4 and step 8. The last two batches of data do not form a complete 4-step cycle.
        Thus, they will be automatically skipped by this class. If the dataloader is not standard PyTorch dataloader,
        (e.g. Dali dataloader), this class will automatically consume (load data for nothing) the remaining 2 batches.

    Args:
        dataloader (``Iterable``): Your dataloader object for gradient accumulation.
        accumulate_size (int): The number of steps to accumulate gradients.
    """

    def __init__(self, dataloader: Iterable, accumulate_size: int) -> None:
        self.dataloader = dataloader
        self.consume_remain_data = True
        self.steps_per_epoch = len(dataloader) - len(dataloader) % accumulate_size

    def __getattr__(self, __name: str) -> Any:
        return getattr(self.dataloader, __name)

    def __len__(self) -> int:
        return self.steps_per_epoch

    def __iter__(self) -> Iterable:
        self._cur_step = 0
        self._dataiter = iter(self.dataloader)
        return self

    def __next__(self) -> Union[Tensor, Tuple[Tensor]]:
        if self._cur_step < self.steps_per_epoch:
            self._cur_step += 1

            if self._cur_step == self.steps_per_epoch and self.consume_remain_data:
                # this is to handle non standard pytorch dataloader
                # such as dali dataloader
                while True:
                    try:
                        _ = next(self._dataiter)
                    except StopIteration:
                        break
            return next(self._dataiter)
        else:
            raise StopIteration

a = list(range(1,11))
print(len(a), a)
a_ = GradAccumDataloaderTorch(a, accumulate_size=4)
for i, val in enumerate(a_):
    print(i, val)
```

the output is:
```
10 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
0 1
1 2
2 3
3 4
4 5
5 6
6 7
```



### Environment

collosal : pip install colossalai==0.1.4+torch1.10cu11.1 -f https://release.colossalai.org
ubuntu 20.04

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/engine/gradient_accumulation/_gradient_accumulation.py`
Content:
```
1 #!/usr/bin/env python
2 # -*- encoding: utf-8 -*-
3 
4 from typing import Union
5 import torch.nn as nn
6 from torch import Tensor
7 from typing import Iterable, Any, Tuple
8 from colossalai.nn.optimizer import ColossalaiOptimizer
9 from torch.nn.parallel.distributed import DistributedDataParallel
10 from torch.optim import Optimizer
11 from torch.optim.lr_scheduler import _LRScheduler
12 from torch.utils.data import DataLoader
13 from colossalai.utils import conditional_context
14 from colossalai.engine import BaseGradientHandler
15 
16 
17 class GradAccumOptimizer(ColossalaiOptimizer):
18     """A wrapper for the optimizer to enable gradient accumulation by skipping the steps 
19     before accumulation size is reached.
20 
21     Args:
22         optim (:class:`torch.optim.Optimizer`): Your optimizer object for gradient accumulation.
23         accumulate_size (int): The number of steps to accumulate gradients.
24         model (:class:`torch.nn.Module`):
25             Your model object to check if it is DistributedDataParallel for special handling of no_sync() context.
26     """
27 
28     def __init__(self, optim: Optimizer, accumulate_size: int, model: nn.Module = None):
29         super().__init__(optim)
30         self.accumulate_size = accumulate_size
31         self.accumulate_step = 0
32 
33         # handle pytorch ddp auto all reduce
34         self.model = model
35         self.is_torch_ddp = isinstance(self.model, DistributedDataParallel)
36 
37     def zero_grad(self, *args, **kwargs) -> None:
38         """
39         Set all gradients to zero.
40 
41         Args:
42             *args: positional arguments for the optimizer wrapped
43             **kwargs: keyword arguments for the optimizer wrapped
44         """
45 
46         if self.accumulate_step == 0:
47             self.optim.zero_grad(*args, **kwargs)
48 
49     def step(self, *args, **kwargs) -> None:
50         """
51         Update the model parameters.
52 
53         Args:
54             *args: positional arguments for the optimizer wrapped
55             **kwargs: keyword arguments for the optimizer wrapped
56         """
57 
58         if self.accumulate_step < self.accumulate_size:
59             return None
60         else:
61             self.accumulate_step = 0
62             return self.optim.step(*args, **kwargs)
63 
64     def clip_grad_norm(self, model: nn.Module, max_norm: float) -> None:
65         """
66         Clip gradients by norm.
67 
68         Args:
69             model (:class:`torch.nn.Module`): a torch module instance
70             max_norm (float): the max norm for gradient clipping
71         """
72 
73         if self.accumulate_step < self.accumulate_size:
74             pass
75         else:
76             self.optim.clip_grad_norm(model, max_norm)
77 
78     def backward(self, loss: Tensor) -> None:
79         """Execute backward pass.
80 
81         Args:
82             loss (:class:`torch.Tensor`): the loss value.
83         """
84 
85         self.accumulate_step += 1
86 
87         if self.is_torch_ddp:
88             no_sync = self.accumulate_step < self.accumulate_size
89             with conditional_context(self.model.no_sync(), enable=no_sync):
90                 scaled_loss = loss / self.accumulate_size
91                 self.optim.backward(scaled_loss)
92         else:
93             scaled_loss = loss / self.accumulate_size
94             self.optim.backward(scaled_loss)
95 
96     def backward_by_grad(self, tensor: Tensor, grad: Tensor) -> None:
97         """Execute backward pass given the gradients of the output.
98 
99         Args:
100             loss (:class:`torch.Tensor`): the loss value.
101             grad (:class:`torch.Tensor`): the output gradient.
102         """
103 
104         self.accumulate_step += 1
105         no_sync = self.is_torch_ddp and self.accumulate_step < self.accumulate_size
106 
107         if no_sync:
108             with self.model.no_sync():
109                 self.optim.backward_by_grad(tensor, grad)
110         else:
111             self.optim.backward_by_grad(tensor, grad)
112 
113 
114 class GradAccumDataloader:
115     """A wrapper for dataloader to enable gradient accumulation by dropping the last incomplete steps.
116 
117     Note:
118         The dataloader would drop the last incomplete steps for gradient accumulation.
119         For example, if a dataloader has 10 batches of data and accumulate size is 4. The model parameters will
120         be updated only twice at step 4 and step 8. The last two batches of data do not form a complete 4-step cycle.
121         Thus, they will be automatically skipped by this class. If the dataloader is not standard PyTorch dataloader,
122         (e.g. Dali dataloader), this class will automatically consume (load data for nothing) the remaining 2 batches.
123 
124     Args:
125         dataloader (``Iterable``): Your dataloader object for gradient accumulation.
126         accumulate_size (int): The number of steps to accumulate gradients.
127     """
128 
129     def __init__(self, dataloader: Iterable, accumulate_size: int) -> None:
130         self.dataloader = dataloader
131         self.consume_remain_data = not isinstance(dataloader, DataLoader)
132         self.steps_per_epoch = len(dataloader) - len(dataloader) % accumulate_size
133 
134     def __getattr__(self, __name: str) -> Any:
135         return getattr(self.dataloader, __name)
136 
137     def __len__(self) -> int:
138         return self.steps_per_epoch
139 
140     def __iter__(self) -> Iterable:
141         self._cur_step = 0
142         self._dataiter = iter(self.dataloader)
143         return self
144 
145     def __next__(self) -> Union[Tensor, Tuple[Tensor]]:
146         if self._cur_step < self.steps_per_epoch:
147             self._cur_step += 1
148 
149             if self._cur_step == self.steps_per_epoch and self.consume_remain_data:
150                 # this is to handle non standard pytorch dataloader
151                 # such as dali dataloader
152                 while True:
153                     try:
154                         _ = next(self._dataiter)
155                     except StopIteration:
156                         break
157             return next(self._dataiter)
158         else:
159             raise StopIteration
160 
161 
162 class GradAccumLrSchedulerByStep(_LRScheduler):
163     """A wrapper for the LR scheduler to enable gradient accumulation by skipping the steps 
164     before accumulation size is reached.
165 
166     Args:
167         lr_scheduler (:class:`torch.optim.lr_scheduler._LRScheduler`):
168             Your ``lr_scheduler`` object for gradient accumulation.
169         accumulate_size (int): The number of steps to accumulate gradients.
170     """
171 
172     def __init__(self, lr_scheduler: _LRScheduler, accumulate_size: int) -> None:
173         self.lr_scheduler = lr_scheduler
174         self.accumulate_size = accumulate_size
175         self.accumulate_step = 0
176 
177     @staticmethod
178     def compute_effective_steps_per_epoch(dataloader: Iterable, accumulate_size: int) -> int:
179         """
180         Computes the number of effective training iterations. An effective iteration is defined
181         as the the aggregation of <accumulate_size> iterations. For examples, if accumulate_size = 4,
182         then 4 iterations are considered as one effective iteration.
183 
184         Args:
185             dataloader (``Iterable``): Your dataloader object for gradient accumulation.
186             accumulate_size (int): The number of steps to accumulate gradients.
187 
188         """
189         return len(dataloader) // accumulate_size
190 
191     def __getattr__(self, __name: str) -> Any:
192         return getattr(self.lr_scheduler, __name)
193 
194     def step(self, *args, **kwargs) -> None:
195         """
196         Update the learning rate.
197 
198         Args:
199             *args: positional arguments for the lr scheduler wrapped.
200             **kwargs: keyword arguments for the lr scheduler wrapped.
201         """
202         self.accumulate_step += 1
203         if self.accumulate_step < self.accumulate_size:
204             pass
205         else:
206             self.accumulate_step = 0
207             self.lr_scheduler.step(*args, **kwargs)
208 
209     def get_lr(self) -> Tensor:
210         """
211         Compute the next learning rate.
212 
213         Returns:
214             Tensor: the upcoming learning rate.
215         """
216 
217         return self.lr_scheduler.get_lr()
218 
219     def get_last_lr(self) -> Tensor:
220         """
221         Returns the current learning rate.
222 
223         Returns:
224             Tensor: the current learning rate.
225         """
226 
227         return self.lr_scheduler.get_last_lr()
228 
229     def print_lr(self, *args, **kwargs) -> None:
230         """
231         Print he learning rate.
232 
233         Args:
234             *args: positional arguments for the lr scheduler wrapped.
235             **kwargs: keyword arguments for the lr scheduler wrapped.
236         """
237         self.lr_scheduler.print_lr(*args, **kwargs)
238 
239     def state_dict(self) -> dict:
240         """
241         Returns the states of the lr scheduler as dictionary.
242 
243         Returns:
244             dict: the states of the lr scheduler.
245         """
246         return self.lr_scheduler.state_dict()
247 
248     def load_state_dict(self, state_dict: dict) -> None:
249         """
250         Load the states of the lr scheduler from a dictionary object.
251 
252         Returns:
253             dict: the states of the lr scheduler.
254         """
255         self.lr_scheduler.load_state_dict(state_dict)
256 
257 
258 class GradAccumGradientHandler:
259     r"""A wrapper for the gradient handler to enable gradient accumulation by skipping the steps
260     before accumulation size is reached.
261 
262     Args:
263         grad_handler (:class:`colossalai.engine.BaseGradientHandler`):
264             Your ``gradient_handler`` object for gradient accumulation, would be called when achieving `accumulate_size`.
265         accumulate_size (int): The number of steps to accumulate gradients.
266 
267     More details about ``gradient_handlers`` could be found in
268     `Gradient_handler <https://github.com/hpcaitech/ColossalAI/tree/main/colossalai/engine/gradient_handler>`_.
269 
270     """
271 
272     def __init__(self, grad_handler: BaseGradientHandler, accumulate_size: int) -> None:
273         assert isinstance(grad_handler, BaseGradientHandler), \
274             f'expected grad_handler to be type BaseGradientHandler, but got {type(grad_handler)}'
275         self.grad_handler = grad_handler
276         self.accumulate_size = accumulate_size
277         self.accumulate_step = 0
278 
279     def handle_gradient(self) -> None:
280         """
281         Handle gradients reduction only in the last gradient accumulation step.
282         """
283 
284         self.accumulate_step += 1
285         if self.accumulate_step < self.accumulate_size:
286             pass
287         else:
288             self.accumulate_step = 0
289             self.grad_handler.handle_gradient()
290
```
--- 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/colossalai/engine/gradient_accumulation/_gradient_accumulation.py b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py
--- a/colossalai/engine/gradient_accumulation/_gradient_accumulation.py
+++ b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py
@@ -145,6 +145,7 @@
     def __next__(self) -> Union[Tensor, Tuple[Tensor]]:
         if self._cur_step < self.steps_per_epoch:
             self._cur_step += 1
+            data = next(self._dataiter)
 
             if self._cur_step == self.steps_per_epoch and self.consume_remain_data:
                 # this is to handle non standard pytorch dataloader
@@ -154,7 +155,7 @@
                         _ = next(self._dataiter)
                     except StopIteration:
                         break
-            return next(self._dataiter)
+            return data
         else:
             raise StopIteration

{"golden_diff": "diff --git a/colossalai/engine/gradient_accumulation/_gradient_accumulation.py b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py\n--- a/colossalai/engine/gradient_accumulation/_gradient_accumulation.py\n+++ b/colossalai/engine/gradient_accumulation/_gradient_accumulation.py\n@@ -145,6 +145,7 @@\n     def __next__(self) -> Union[Tensor, Tuple[Tensor]]:\n         if self._cur_step < self.steps_per_epoch:\n             self._cur_step += 1\n+            data = next(self._dataiter)\n \n             if self._cur_step == self.steps_per_epoch and self.consume_remain_data:\n                 # this is to handle non standard pytorch dataloader\n@@ -154,7 +155,7 @@\n                         _ = next(self._dataiter)\n                     except StopIteration:\n                         break\n-            return next(self._dataiter)\n+            return data\n         else:\n             raise StopIteration\n", "issue": "[BUG]: GradAccumDataloader will ignore the last iter.\n### \ud83d\udc1b Describe the bug\r\n\r\nI used the list as dataloader for testing.\r\n\r\n```\r\nclass GradAccumDataloader:\r\n    \"\"\"A wrapper for dataloader to enable gradient accumulation by dropping the last incomplete steps.\r\n\r\n    Note:\r\n        The dataloader would drop the last incomplete steps for gradient accumulation.\r\n        For example, if a dataloader has 10 batches of data and accumulate size is 4. The model parameters will\r\n        be updated only twice at step 4 and step 8. The last two batches of data do not form a complete 4-step cycle.\r\n        Thus, they will be automatically skipped by this class. If the dataloader is not standard PyTorch dataloader,\r\n        (e.g. Dali dataloader), this class will automatically consume (load data for nothing) the remaining 2 batches.\r\n\r\n    Args:\r\n        dataloader (``Iterable``): Your dataloader object for gradient accumulation.\r\n        accumulate_size (int): The number of steps to accumulate gradients.\r\n    \"\"\"\r\n\r\n    def __init__(self, dataloader: Iterable, accumulate_size: int) -> None:\r\n        self.dataloader = dataloader\r\n        self.consume_remain_data = True\r\n        self.steps_per_epoch = len(dataloader) - len(dataloader) % accumulate_size\r\n\r\n    def __getattr__(self, __name: str) -> Any:\r\n        return getattr(self.dataloader, __name)\r\n\r\n    def __len__(self) -> int:\r\n        return self.steps_per_epoch\r\n\r\n    def __iter__(self) -> Iterable:\r\n        self._cur_step = 0\r\n        self._dataiter = iter(self.dataloader)\r\n        return self\r\n\r\n    def __next__(self) -> Union[Tensor, Tuple[Tensor]]:\r\n        if self._cur_step < self.steps_per_epoch:\r\n            self._cur_step += 1\r\n\r\n            if self._cur_step == self.steps_per_epoch and self.consume_remain_data:\r\n                # this is to handle non standard pytorch dataloader\r\n                # such as dali dataloader\r\n                while True:\r\n                    try:\r\n                        _ = next(self._dataiter)\r\n                    except StopIteration:\r\n                        break\r\n            return next(self._dataiter)\r\n        else:\r\n            raise StopIteration\r\n\r\na = list(range(1,11))\r\nprint(len(a), a)\r\na_ = GradAccumDataloaderTorch(a, accumulate_size=4)\r\nfor i, val in enumerate(a_):\r\n    print(i, val)\r\n```\r\n\r\nthe output is:\r\n```\r\n10 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]\r\n0 1\r\n1 2\r\n2 3\r\n3 4\r\n4 5\r\n5 6\r\n6 7\r\n```\r\n\r\n\r\n\r\n### Environment\r\n\r\ncollosal : pip install colossalai==0.1.4+torch1.10cu11.1 -f https://release.colossalai.org\r\nubuntu 20.04\n", "before_files": [{"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n\nfrom typing import Union\nimport torch.nn as nn\nfrom torch import Tensor\nfrom typing import Iterable, Any, Tuple\nfrom colossalai.nn.optimizer import ColossalaiOptimizer\nfrom torch.nn.parallel.distributed import DistributedDataParallel\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler\nfrom torch.utils.data import DataLoader\nfrom colossalai.utils import conditional_context\nfrom colossalai.engine import BaseGradientHandler\n\n\nclass GradAccumOptimizer(ColossalaiOptimizer):\n    \"\"\"A wrapper for the optimizer to enable gradient accumulation by skipping the steps \n    before accumulation size is reached.\n\n    Args:\n        optim (:class:`torch.optim.Optimizer`): Your optimizer object for gradient accumulation.\n        accumulate_size (int): The number of steps to accumulate gradients.\n        model (:class:`torch.nn.Module`):\n            Your model object to check if it is DistributedDataParallel for special handling of no_sync() context.\n    \"\"\"\n\n    def __init__(self, optim: Optimizer, accumulate_size: int, model: nn.Module = None):\n        super().__init__(optim)\n        self.accumulate_size = accumulate_size\n        self.accumulate_step = 0\n\n        # handle pytorch ddp auto all reduce\n        self.model = model\n        self.is_torch_ddp = isinstance(self.model, DistributedDataParallel)\n\n    def zero_grad(self, *args, **kwargs) -> None:\n        \"\"\"\n        Set all gradients to zero.\n\n        Args:\n            *args: positional arguments for the optimizer wrapped\n            **kwargs: keyword arguments for the optimizer wrapped\n        \"\"\"\n\n        if self.accumulate_step == 0:\n            self.optim.zero_grad(*args, **kwargs)\n\n    def step(self, *args, **kwargs) -> None:\n        \"\"\"\n        Update the model parameters.\n\n        Args:\n            *args: positional arguments for the optimizer wrapped\n            **kwargs: keyword arguments for the optimizer wrapped\n        \"\"\"\n\n        if self.accumulate_step < self.accumulate_size:\n            return None\n        else:\n            self.accumulate_step = 0\n            return self.optim.step(*args, **kwargs)\n\n    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> None:\n        \"\"\"\n        Clip gradients by norm.\n\n        Args:\n            model (:class:`torch.nn.Module`): a torch module instance\n            max_norm (float): the max norm for gradient clipping\n        \"\"\"\n\n        if self.accumulate_step < self.accumulate_size:\n            pass\n        else:\n            self.optim.clip_grad_norm(model, max_norm)\n\n    def backward(self, loss: Tensor) -> None:\n        \"\"\"Execute backward pass.\n\n        Args:\n            loss (:class:`torch.Tensor`): the loss value.\n        \"\"\"\n\n        self.accumulate_step += 1\n\n        if self.is_torch_ddp:\n            no_sync = self.accumulate_step < self.accumulate_size\n            with conditional_context(self.model.no_sync(), enable=no_sync):\n                scaled_loss = loss / self.accumulate_size\n                self.optim.backward(scaled_loss)\n        else:\n            scaled_loss = loss / self.accumulate_size\n            self.optim.backward(scaled_loss)\n\n    def backward_by_grad(self, tensor: Tensor, grad: Tensor) -> None:\n        \"\"\"Execute backward pass given the gradients of the output.\n\n        Args:\n            loss (:class:`torch.Tensor`): the loss value.\n            grad (:class:`torch.Tensor`): the output gradient.\n        \"\"\"\n\n        self.accumulate_step += 1\n        no_sync = self.is_torch_ddp and self.accumulate_step < self.accumulate_size\n\n        if no_sync:\n            with self.model.no_sync():\n                self.optim.backward_by_grad(tensor, grad)\n        else:\n            self.optim.backward_by_grad(tensor, grad)\n\n\nclass GradAccumDataloader:\n    \"\"\"A wrapper for dataloader to enable gradient accumulation by dropping the last incomplete steps.\n\n    Note:\n        The dataloader would drop the last incomplete steps for gradient accumulation.\n        For example, if a dataloader has 10 batches of data and accumulate size is 4. The model parameters will\n        be updated only twice at step 4 and step 8. The last two batches of data do not form a complete 4-step cycle.\n        Thus, they will be automatically skipped by this class. If the dataloader is not standard PyTorch dataloader,\n        (e.g. Dali dataloader), this class will automatically consume (load data for nothing) the remaining 2 batches.\n\n    Args:\n        dataloader (``Iterable``): Your dataloader object for gradient accumulation.\n        accumulate_size (int): The number of steps to accumulate gradients.\n    \"\"\"\n\n    def __init__(self, dataloader: Iterable, accumulate_size: int) -> None:\n        self.dataloader = dataloader\n        self.consume_remain_data = not isinstance(dataloader, DataLoader)\n        self.steps_per_epoch = len(dataloader) - len(dataloader) % accumulate_size\n\n    def __getattr__(self, __name: str) -> Any:\n        return getattr(self.dataloader, __name)\n\n    def __len__(self) -> int:\n        return self.steps_per_epoch\n\n    def __iter__(self) -> Iterable:\n        self._cur_step = 0\n        self._dataiter = iter(self.dataloader)\n        return self\n\n    def __next__(self) -> Union[Tensor, Tuple[Tensor]]:\n        if self._cur_step < self.steps_per_epoch:\n            self._cur_step += 1\n\n            if self._cur_step == self.steps_per_epoch and self.consume_remain_data:\n                # this is to handle non standard pytorch dataloader\n                # such as dali dataloader\n                while True:\n                    try:\n                        _ = next(self._dataiter)\n                    except StopIteration:\n                        break\n            return next(self._dataiter)\n        else:\n            raise StopIteration\n\n\nclass GradAccumLrSchedulerByStep(_LRScheduler):\n    \"\"\"A wrapper for the LR scheduler to enable gradient accumulation by skipping the steps \n    before accumulation size is reached.\n\n    Args:\n        lr_scheduler (:class:`torch.optim.lr_scheduler._LRScheduler`):\n            Your ``lr_scheduler`` object for gradient accumulation.\n        accumulate_size (int): The number of steps to accumulate gradients.\n    \"\"\"\n\n    def __init__(self, lr_scheduler: _LRScheduler, accumulate_size: int) -> None:\n        self.lr_scheduler = lr_scheduler\n        self.accumulate_size = accumulate_size\n        self.accumulate_step = 0\n\n    @staticmethod\n    def compute_effective_steps_per_epoch(dataloader: Iterable, accumulate_size: int) -> int:\n        \"\"\"\n        Computes the number of effective training iterations. An effective iteration is defined\n        as the the aggregation of <accumulate_size> iterations. For examples, if accumulate_size = 4,\n        then 4 iterations are considered as one effective iteration.\n\n        Args:\n            dataloader (``Iterable``): Your dataloader object for gradient accumulation.\n            accumulate_size (int): The number of steps to accumulate gradients.\n\n        \"\"\"\n        return len(dataloader) // accumulate_size\n\n    def __getattr__(self, __name: str) -> Any:\n        return getattr(self.lr_scheduler, __name)\n\n    def step(self, *args, **kwargs) -> None:\n        \"\"\"\n        Update the learning rate.\n\n        Args:\n            *args: positional arguments for the lr scheduler wrapped.\n            **kwargs: keyword arguments for the lr scheduler wrapped.\n        \"\"\"\n        self.accumulate_step += 1\n        if self.accumulate_step < self.accumulate_size:\n            pass\n        else:\n            self.accumulate_step = 0\n            self.lr_scheduler.step(*args, **kwargs)\n\n    def get_lr(self) -> Tensor:\n        \"\"\"\n        Compute the next learning rate.\n\n        Returns:\n            Tensor: the upcoming learning rate.\n        \"\"\"\n\n        return self.lr_scheduler.get_lr()\n\n    def get_last_lr(self) -> Tensor:\n        \"\"\"\n        Returns the current learning rate.\n\n        Returns:\n            Tensor: the current learning rate.\n        \"\"\"\n\n        return self.lr_scheduler.get_last_lr()\n\n    def print_lr(self, *args, **kwargs) -> None:\n        \"\"\"\n        Print he learning rate.\n\n        Args:\n            *args: positional arguments for the lr scheduler wrapped.\n            **kwargs: keyword arguments for the lr scheduler wrapped.\n        \"\"\"\n        self.lr_scheduler.print_lr(*args, **kwargs)\n\n    def state_dict(self) -> dict:\n        \"\"\"\n        Returns the states of the lr scheduler as dictionary.\n\n        Returns:\n            dict: the states of the lr scheduler.\n        \"\"\"\n        return self.lr_scheduler.state_dict()\n\n    def load_state_dict(self, state_dict: dict) -> None:\n        \"\"\"\n        Load the states of the lr scheduler from a dictionary object.\n\n        Returns:\n            dict: the states of the lr scheduler.\n        \"\"\"\n        self.lr_scheduler.load_state_dict(state_dict)\n\n\nclass GradAccumGradientHandler:\n    r\"\"\"A wrapper for the gradient handler to enable gradient accumulation by skipping the steps\n    before accumulation size is reached.\n\n    Args:\n        grad_handler (:class:`colossalai.engine.BaseGradientHandler`):\n            Your ``gradient_handler`` object for gradient accumulation, would be called when achieving `accumulate_size`.\n        accumulate_size (int): The number of steps to accumulate gradients.\n\n    More details about ``gradient_handlers`` could be found in\n    `Gradient_handler <https://github.com/hpcaitech/ColossalAI/tree/main/colossalai/engine/gradient_handler>`_.\n\n    \"\"\"\n\n    def __init__(self, grad_handler: BaseGradientHandler, accumulate_size: int) -> None:\n        assert isinstance(grad_handler, BaseGradientHandler), \\\n            f'expected grad_handler to be type BaseGradientHandler, but got {type(grad_handler)}'\n        self.grad_handler = grad_handler\n        self.accumulate_size = accumulate_size\n        self.accumulate_step = 0\n\n    def handle_gradient(self) -> None:\n        \"\"\"\n        Handle gradients reduction only in the last gradient accumulation step.\n        \"\"\"\n\n        self.accumulate_step += 1\n        if self.accumulate_step < self.accumulate_size:\n            pass\n        else:\n            self.accumulate_step = 0\n            self.grad_handler.handle_gradient()\n", "path": "colossalai/engine/gradient_accumulation/_gradient_accumulation.py"}], "after_files": [{"content": "#!/usr/bin/env python\n# -*- encoding: utf-8 -*-\n\nfrom typing import Union\nimport torch.nn as nn\nfrom torch import Tensor\nfrom typing import Iterable, Any, Tuple\nfrom colossalai.nn.optimizer import ColossalaiOptimizer\nfrom torch.nn.parallel.distributed import DistributedDataParallel\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import _LRScheduler\nfrom torch.utils.data import DataLoader\nfrom colossalai.utils import conditional_context\nfrom colossalai.engine import BaseGradientHandler\n\n\nclass GradAccumOptimizer(ColossalaiOptimizer):\n    \"\"\"A wrapper for the optimizer to enable gradient accumulation by skipping the steps \n    before accumulation size is reached.\n\n    Args:\n        optim (:class:`torch.optim.Optimizer`): Your optimizer object for gradient accumulation.\n        accumulate_size (int): The number of steps to accumulate gradients.\n        model (:class:`torch.nn.Module`):\n            Your model object to check if it is DistributedDataParallel for special handling of no_sync() context.\n    \"\"\"\n\n    def __init__(self, optim: Optimizer, accumulate_size: int, model: nn.Module = None):\n        super().__init__(optim)\n        self.accumulate_size = accumulate_size\n        self.accumulate_step = 0\n\n        # handle pytorch ddp auto all reduce\n        self.model = model\n        self.is_torch_ddp = isinstance(self.model, DistributedDataParallel)\n\n    def zero_grad(self, *args, **kwargs) -> None:\n        \"\"\"\n        Set all gradients to zero.\n\n        Args:\n            *args: positional arguments for the optimizer wrapped\n            **kwargs: keyword arguments for the optimizer wrapped\n        \"\"\"\n\n        if self.accumulate_step == 0:\n            self.optim.zero_grad(*args, **kwargs)\n\n    def step(self, *args, **kwargs) -> None:\n        \"\"\"\n        Update the model parameters.\n\n        Args:\n            *args: positional arguments for the optimizer wrapped\n            **kwargs: keyword arguments for the optimizer wrapped\n        \"\"\"\n\n        if self.accumulate_step < self.accumulate_size:\n            return None\n        else:\n            self.accumulate_step = 0\n            return self.optim.step(*args, **kwargs)\n\n    def clip_grad_norm(self, model: nn.Module, max_norm: float) -> None:\n        \"\"\"\n        Clip gradients by norm.\n\n        Args:\n            model (:class:`torch.nn.Module`): a torch module instance\n            max_norm (float): the max norm for gradient clipping\n        \"\"\"\n\n        if self.accumulate_step < self.accumulate_size:\n            pass\n        else:\n            self.optim.clip_grad_norm(model, max_norm)\n\n    def backward(self, loss: Tensor) -> None:\n        \"\"\"Execute backward pass.\n\n        Args:\n            loss (:class:`torch.Tensor`): the loss value.\n        \"\"\"\n\n        self.accumulate_step += 1\n\n        if self.is_torch_ddp:\n            no_sync = self.accumulate_step < self.accumulate_size\n            with conditional_context(self.model.no_sync(), enable=no_sync):\n                scaled_loss = loss / self.accumulate_size\n                self.optim.backward(scaled_loss)\n        else:\n            scaled_loss = loss / self.accumulate_size\n            self.optim.backward(scaled_loss)\n\n    def backward_by_grad(self, tensor: Tensor, grad: Tensor) -> None:\n        \"\"\"Execute backward pass given the gradients of the output.\n\n        Args:\n            loss (:class:`torch.Tensor`): the loss value.\n            grad (:class:`torch.Tensor`): the output gradient.\n        \"\"\"\n\n        self.accumulate_step += 1\n        no_sync = self.is_torch_ddp and self.accumulate_step < self.accumulate_size\n\n        if no_sync:\n            with self.model.no_sync():\n                self.optim.backward_by_grad(tensor, grad)\n        else:\n            self.optim.backward_by_grad(tensor, grad)\n\n\nclass GradAccumDataloader:\n    \"\"\"A wrapper for dataloader to enable gradient accumulation by dropping the last incomplete steps.\n\n    Note:\n        The dataloader would drop the last incomplete steps for gradient accumulation.\n        For example, if a dataloader has 10 batches of data and accumulate size is 4. The model parameters will\n        be updated only twice at step 4 and step 8. The last two batches of data do not form a complete 4-step cycle.\n        Thus, they will be automatically skipped by this class. If the dataloader is not standard PyTorch dataloader,\n        (e.g. Dali dataloader), this class will automatically consume (load data for nothing) the remaining 2 batches.\n\n    Args:\n        dataloader (``Iterable``): Your dataloader object for gradient accumulation.\n        accumulate_size (int): The number of steps to accumulate gradients.\n    \"\"\"\n\n    def __init__(self, dataloader: Iterable, accumulate_size: int) -> None:\n        self.dataloader = dataloader\n        self.consume_remain_data = not isinstance(dataloader, DataLoader)\n        self.steps_per_epoch = len(dataloader) - len(dataloader) % accumulate_size\n\n    def __getattr__(self, __name: str) -> Any:\n        return getattr(self.dataloader, __name)\n\n    def __len__(self) -> int:\n        return self.steps_per_epoch\n\n    def __iter__(self) -> Iterable:\n        self._cur_step = 0\n        self._dataiter = iter(self.dataloader)\n        return self\n\n    def __next__(self) -> Union[Tensor, Tuple[Tensor]]:\n        if self._cur_step < self.steps_per_epoch:\n            self._cur_step += 1\n            data = next(self._dataiter)\n\n            if self._cur_step == self.steps_per_epoch and self.consume_remain_data:\n                # this is to handle non standard pytorch dataloader\n                # such as dali dataloader\n                while True:\n                    try:\n                        _ = next(self._dataiter)\n                    except StopIteration:\n                        break\n            return data\n        else:\n            raise StopIteration\n\n\nclass GradAccumLrSchedulerByStep(_LRScheduler):\n    \"\"\"A wrapper for the LR scheduler to enable gradient accumulation by skipping the steps \n    before accumulation size is reached.\n\n    Args:\n        lr_scheduler (:class:`torch.optim.lr_scheduler._LRScheduler`):\n            Your ``lr_scheduler`` object for gradient accumulation.\n        accumulate_size (int): The number of steps to accumulate gradients.\n    \"\"\"\n\n    def __init__(self, lr_scheduler: _LRScheduler, accumulate_size: int) -> None:\n        self.lr_scheduler = lr_scheduler\n        self.accumulate_size = accumulate_size\n        self.accumulate_step = 0\n\n    @staticmethod\n    def compute_effective_steps_per_epoch(dataloader: Iterable, accumulate_size: int) -> int:\n        \"\"\"\n        Computes the number of effective training iterations. An effective iteration is defined\n        as the the aggregation of <accumulate_size> iterations. For examples, if accumulate_size = 4,\n        then 4 iterations are considered as one effective iteration.\n\n        Args:\n            dataloader (``Iterable``): Your dataloader object for gradient accumulation.\n            accumulate_size (int): The number of steps to accumulate gradients.\n\n        \"\"\"\n        return len(dataloader) // accumulate_size\n\n    def __getattr__(self, __name: str) -> Any:\n        return getattr(self.lr_scheduler, __name)\n\n    def step(self, *args, **kwargs) -> None:\n        \"\"\"\n        Update the learning rate.\n\n        Args:\n            *args: positional arguments for the lr scheduler wrapped.\n            **kwargs: keyword arguments for the lr scheduler wrapped.\n        \"\"\"\n        self.accumulate_step += 1\n        if self.accumulate_step < self.accumulate_size:\n            pass\n        else:\n            self.accumulate_step = 0\n            self.lr_scheduler.step(*args, **kwargs)\n\n    def get_lr(self) -> Tensor:\n        \"\"\"\n        Compute the next learning rate.\n\n        Returns:\n            Tensor: the upcoming learning rate.\n        \"\"\"\n\n        return self.lr_scheduler.get_lr()\n\n    def get_last_lr(self) -> Tensor:\n        \"\"\"\n        Returns the current learning rate.\n\n        Returns:\n            Tensor: the current learning rate.\n        \"\"\"\n\n        return self.lr_scheduler.get_last_lr()\n\n    def print_lr(self, *args, **kwargs) -> None:\n        \"\"\"\n        Print he learning rate.\n\n        Args:\n            *args: positional arguments for the lr scheduler wrapped.\n            **kwargs: keyword arguments for the lr scheduler wrapped.\n        \"\"\"\n        self.lr_scheduler.print_lr(*args, **kwargs)\n\n    def state_dict(self) -> dict:\n        \"\"\"\n        Returns the states of the lr scheduler as dictionary.\n\n        Returns:\n            dict: the states of the lr scheduler.\n        \"\"\"\n        return self.lr_scheduler.state_dict()\n\n    def load_state_dict(self, state_dict: dict) -> None:\n        \"\"\"\n        Load the states of the lr scheduler from a dictionary object.\n\n        Returns:\n            dict: the states of the lr scheduler.\n        \"\"\"\n        self.lr_scheduler.load_state_dict(state_dict)\n\n\nclass GradAccumGradientHandler:\n    r\"\"\"A wrapper for the gradient handler to enable gradient accumulation by skipping the steps\n    before accumulation size is reached.\n\n    Args:\n        grad_handler (:class:`colossalai.engine.BaseGradientHandler`):\n            Your ``gradient_handler`` object for gradient accumulation, would be called when achieving `accumulate_size`.\n        accumulate_size (int): The number of steps to accumulate gradients.\n\n    More details about ``gradient_handlers`` could be found in\n    `Gradient_handler <https://github.com/hpcaitech/ColossalAI/tree/main/colossalai/engine/gradient_handler>`_.\n\n    \"\"\"\n\n    def __init__(self, grad_handler: BaseGradientHandler, accumulate_size: int) -> None:\n        assert isinstance(grad_handler, BaseGradientHandler), \\\n            f'expected grad_handler to be type BaseGradientHandler, but got {type(grad_handler)}'\n        self.grad_handler = grad_handler\n        self.accumulate_size = accumulate_size\n        self.accumulate_step = 0\n\n    def handle_gradient(self) -> None:\n        \"\"\"\n        Handle gradients reduction only in the last gradient accumulation step.\n        \"\"\"\n\n        self.accumulate_step += 1\n        if self.accumulate_step < self.accumulate_size:\n            pass\n        else:\n            self.accumulate_step = 0\n            self.grad_handler.handle_gradient()\n", "path": "colossalai/engine/gradient_accumulation/_gradient_accumulation.py"}]}