add-support-for-non-trainable-params

[nikmo33]

Adds support for non-trainable parameters in the strategy classes and operations.
Currently there is an assumption that all parameters are trainable in two places -

Assume there is a one to one mapping between parameter keys and optimiser keys
When updating parameters during densification all parameters are replaced with trainable counterparts.
This PR supports non trainable parameters by -

Asserting all trainable parameters have a 1-1 mapping with optimisers (this is inferred via the requires grad flag)
Persisting the retain grad flags when parameters are replaced during densification
Adds some tests to ensure that all parameters are densified and requires grad flag is persisted

[liruilong940607]

Add a msg to explain. sth like assert not param.requires_grad, f"param {name} does not present in the optimizer, it's requires_grad should be False, but found True"

[liruilong940607]

This breaks the API of this _update_param_with_optimizer function. Though it is kinda ok because it's an internal function, there seem to be a simple way that wouldn't break it:

Instead of

def param_fn(name: str, p: Tensor, requires_grad: bool) -> Tensor:
        return torch.nn.Parameter(p[sel], requires_grad=requires_grad)

We could do

def param_fn(name: str, p: Tensor) -> Tensor:
        return torch.nn.Parameter(p[sel], requires_grad=p.requires_grad)

[liruilong940607]

Happy to merge it after resolving these minor comments. And thanks for writting the test!

[liruilong940607]

LGTM!