Zero fp32 and fp16 parameter grads.
(self, set_grads_to_None=False)
| 262 | raise RuntimeError("FP16_Optimizer should be deserialized using load_state_dict().") |
| 263 | |
| 264 | def zero_grad(self, set_grads_to_None=False): |
| 265 | """ |
| 266 | Zero fp32 and fp16 parameter grads. |
| 267 | """ |
| 268 | # In principle, only the .grad attributes of the model params need to be zeroed, |
| 269 | # because gradients are copied into the FP32 master params. However, we zero |
| 270 | # all gradients owned by the optimizer, just to be safe: |
| 271 | for group in self.optimizer.param_groups: |
| 272 | for p in group['params']: |
| 273 | if set_grads_to_None: |
| 274 | p.grad = None |
| 275 | else: |
| 276 | if p.grad is not None: |
| 277 | p.grad.detach_() |
| 278 | p.grad.zero_() |
| 279 | |
| 280 | # Zero fp16 gradients owned by the model: |
| 281 | for fp16_group in self.fp16_groups: |
| 282 | for param in fp16_group: |
| 283 | if set_grads_to_None: |
| 284 | param.grad = None |
| 285 | else: |
| 286 | if param.grad is not None: |
| 287 | param.grad.detach_() # as in torch.optim.optimizer.zero_grad() |
| 288 | param.grad.zero_() |
| 289 | |
| 290 | def _check_overflow(self): |
| 291 | params = [] |