SiriusNEO opened a new pull request, #14542: URL: https://github.com/apache/tvm/pull/14542
### Introduction This PR introduces the high-level reverse-mode automatic differentiation pass `Gradient` for Relax. It's the core component when we are trying training or fine-tuning in Relax IR. Before upstreaming, this work is actively iterated and maintained in many forks like [mlc](https://github.com/mlc-ai/relax) and [relax-training](https://github.com/ACMClass-TVM-20/relax-training). Now it reaches a relatively stable version and it's time for us to upstream this important work to the unity branch. The Python side API: - `Gradient(func_name: str, require_grads: Optional[Union[Var, List[Var]]] = None, target_index: int = 0) -> tvm.ir.transform.Pass` It will transform the given funcion in the IRModule, and adds a new function that calculates the gradient with regard to the function's output. ### Examples ``` @I.ir_module class Module: @R.function def main( x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32") ) -> R.Tensor((), dtype="float32"): with R.dataflow(): lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y) # use R.sum to reduce the tensor to a scalar lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False) R.output(lv2) return lv2 After = relax.transform.Gradient("main")(Module) ``` Then the transformed module `After` will be ``` @I.ir_module class After: @R.function def main( x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32") ) -> R.Tensor((), dtype="float32"): with R.dataflow(): lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y) lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False) R.output(lv2) return lv2 @R.function def main_adjoint( x: R.Tensor((3, 3), dtype="float32"), y: R.Tensor((3, 3), dtype="float32") ) -> R.Tuple( R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((3, 3), dtype="float32"), R.Tensor((3, 3), dtype="float32")), ): with R.dataflow(): # original bindings lv1: R.Tensor((3, 3), dtype="float32") = R.add(x, y) lv2: R.Tensor((), dtype="float32") = R.sum(lv1, axis=None, keepdims=False) # bindings w.r.t. intermediate variables lv2_adjoint: R.Tensor((), dtype="float32") = R.ones((), dtype="float32") lv1_adjoint: R.Tensor((3, 3), dtype="float32") = R.broadcast_to( lv2_adjoint, (3, 3) ) # bindings w.r.t. parameters x_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint y_adjoint: R.Tensor((3, 3), dtype="float32") = lv1_adjoint R.output(lv2, x_adjoint, y_adjoint) # return value: (orig_return_values, tuple(adjoints)) return (lv2, (x_adjoint, y_adjoint)) ``` Here we specify the target function `main` by its name. We let the `require_grads` be default value (`None`) so it will calculate all inputs' adjoints (`x_adjoint`, `y_adjoint`) and return them. We let the `target_index` be default value `0` so it will take the unique return value `lv2` as the target (the scalar we start to differentiate and propagating adjoints) of AD. ### Links More details can be found in the following related links: - [Relax Training APIs Tutorial and Examples](https://github.com/mlc-ai/mlc-training) - Specially focus on [the document](https://github.com/mlc-ai/mlc-training/blob/main/tutorial/Reverse_Mode_Automatic_Differentiation_in_Relax.md) of this pass. - https://github.com/tlc-pack/relax/issues/413 - Previous PR in mlc: https://github.com/mlc-ai/relax/pull/103 - Our talk @ TVMCon23: [Cross Platform Training Using Automatic Differentiation on Relax IR](https://www.youtube.com/watch?v=bsfpzMufyNw) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
