[Relax][Training] Optimizer API

[Ubospica]

This PR supports the new relax.training.optimizer.Optimizer API, which is the abstraction of optimizers in the relax training workflow. Specific optimizers should inherit it and override its methods.

See https://github.com/ACMClass-TVM-20/AD-Example/blob/main/optimizer_api_example.py for examples of its usage.

relax.optimizer is added as an alias for relax.training.optimizer.

The optimizer would be like:

class Optimizer:
    def __init__(self, params: Union[Var, List[Var]]) -> None:
        pass

    @property
    def state(self):
        """Return the state of the optimizer that would be updated in the optimization process."""
		pass

    def get_function(self) -> Function:
        """In any implementation of Optimizer, we will use blockbuilder to build a optimizer
        function that executes the update of parameters and the optimizer state.

        The optimizer function will take in a tuple of parameters, a tuple of gradients of
        parameters, and a tuple of optimizer states. It will return a tuple of updated parameters,
        and a tuple of optimizer states.
        """
        pass

Now we provide three common optimizers in module tvm.relax.training.optimizer:

• SGD
• SGD with momentum or Nesterov momentum
• Adam

[Ubospica]

mypy reports error on this line:

error: Cannot find implementation or library stub for module named "numpy"

Now we uses # type: ignore to suppress the error.

[SiriusNEO]

Would it be great if we change the name of test files from test_optimizer to test_training_optimizer?

[MasterJH5574]

See https://github.com/ACMClass-TVM-20/AD-Example/blob/main/optimizer_api_example.py for examples of its usage.

Seems I have no access to this repo.

[Ubospica]

See ACMClass-TVM-20/AD-Example@main/optimizer_api_example.py for examples of its usage.

Seems I have no access to this repo.

@MasterJH5574 That repo has been set to public.

[Hzfengsy]

One general comment: should optimizer be a Python-only data structure or write it with C++?

[Ubospica]

One general comment: should optimizer be a Python-only data structure or write it with C++?

The main difficulty of putting optimizer in C++ is that optimizers need numpy to construct optimizer states. Now we don't find good utils to construct tensors, like a zeros array with given shape, in C++.

[Hzfengsy]

We have it in Relay and we can make a sperate one for Relax

[Hzfengsy]

Thanks @Ubospica for the continuous improving

[Ubospica]

Thank @SiriusNEO @MasterJH5574 @Hzfengsy @spectrometerHBH for the review!