r"""A wrapper of GradManager's :meth:`~.GradManager.backward`, used to scale ``y``'s grad and unscale parameters' grads. Args: gm: The to be wrapped GradManager. y: Same as GradManager backward's ``y``. dy: Same as GradManager backward's ``dy``. W
(
self,
gm: GradManager,
y: Union[Tensor, List[Tensor]] = None,
dy: Union[Tensor, List[Tensor]] = None,
*,
unscale_grad: bool = True,
update_scale: bool = "if_unscale_grad"
)
| 75 | self._found_non_finite = False |
| 76 | |
| 77 | def backward( |
| 78 | self, |
| 79 | gm: GradManager, |
| 80 | y: Union[Tensor, List[Tensor]] = None, |
| 81 | dy: Union[Tensor, List[Tensor]] = None, |
| 82 | *, |
| 83 | unscale_grad: bool = True, |
| 84 | update_scale: bool = "if_unscale_grad" |
| 85 | ): |
| 86 | r"""A wrapper of GradManager's :meth:`~.GradManager.backward`, used to scale |
| 87 | ``y``'s grad and unscale parameters' grads. |
| 88 | |
| 89 | Args: |
| 90 | gm: The to be wrapped GradManager. |
| 91 | y: Same as GradManager backward's ``y``. |
| 92 | dy: Same as GradManager backward's ``dy``. Will be multiplied |
| 93 | by ``scale_factor``. |
| 94 | unscale_grad: Whether do :meth:`unscale` at the same time. Could be |
| 95 | ``False`` if needs to accumulate grads. |
| 96 | update_scale: Same as :meth:`unscale`'s ``update``. Will be ignored |
| 97 | if ``unscale_grad`` is ``False``. |
| 98 | """ |
| 99 | # These checks should be consistent with GradManager's |
| 100 | if y is None: |
| 101 | ys = [] |
| 102 | elif isinstance(y, (tuple, list)): |
| 103 | ys = y |
| 104 | else: |
| 105 | ys = [y] |
| 106 | if dy is None: |
| 107 | dys = [full_like(y, self.scale_factor) for y in ys] |
| 108 | elif isinstance(dy, (tuple, list)): |
| 109 | dys = [dy_ * self.scale_factor for dy_ in dy] |
| 110 | else: |
| 111 | dys = [dy * self.scale_factor] |
| 112 | |
| 113 | gm.backward(y=ys, dy=dys) |
| 114 | |
| 115 | if unscale_grad: |
| 116 | self.unscale(gm.attached_tensors()) |
| 117 | if update_scale: |
| 118 | self.update() |
| 119 | |
| 120 | def unscale(self, grad_tensors: Iterable[Tensor]): |
| 121 | r"""Unscale all ``grad_tensors``'s grad. |
no test coverage detected