Fix optimizer load state dict copy tensor by reference

[shaltielshmid]

Fixes: #1172

[shaltielshmid]

Also fixes #1176 with the conditional state being copied to the right device.

[NiklasGustafsson]

No need to dispose the old momentum buffer if it's actually moved?

[shaltielshmid]

Good catch. Will fix.

[shaltielshmid]

Question: Do you want me to write something like:

if (momentum_buffer.device_type != _parameter.device_type || momentum_buffer.device_index != _parameter.device_index) {
    using var copy = result;
    result = copy.to(_parameter.device);
}

[shaltielshmid]

(That's not handled in the regular Optimizer.to(Device device) function. See for example here)

[NiklasGustafsson]

I think it's sufficient to check the Handle before and after the call to `to'. I think.

[NiklasGustafsson]

Maybe not.

[shaltielshmid]

Haha
In that case, should I go with my proposed solution?

Also, does that mean we should update all the Optimizer.to(...) functions?

Maybe we should have a parameter dispose_if_moved to the to function to make it generic?

[NiklasGustafsson]

It's handled in the Module variants of to(). Follow that pattern.

[NiklasGustafsson]

Maybe it's worth writing an internal utility function that has the 'move or not move' predicate, so it can be reused.

[shaltielshmid]

I checked the Module variants of to(), and there is no dispose going on there, but there they check the device to see if it changed. I'll follow that pattern.

What do we want to do about all the cases where the .to() function is called but no dispose being handled, like here or here etc.

[shaltielshmid]

@NiklasGustafsson What do you think of my solution?
I added a disposeAfter flag to the to(..) function which dispose the input tensor after moving to the new tensor.
Turns out we didn't need to check if the tensor actually moved, because we get a new handle regardless.

I wrote a unit test attempting to move and cast tensors to the same and different types/devices, and made sure the behavior was as expected.

What do you think?

If this solution is good, I'll wait for this to be merged and then use it in the other PR.

[NiklasGustafsson]

@shaltielshmid -- here's a tip: Once you have created a PR, try to consolidate a number of commits before pushing again. Every push starts a new build on the CI/CD pipeline, which wastes resources.

[shaltielshmid]

Of course, I apologize.
I generally try to make a few commits, but only push when I complete changes.

[NiklasGustafsson]

Why do we have to set copy to true here?

[shaltielshmid]

If the State Dictionary is actually from another existing optimizer, then the tensors are copied by reference and we don't want two different optimizers sharing a state tensor.

[shaltielshmid]

For example:

var lin1 = torch.nn.Linear(10, 10);
var optim1 = torch.optim.Adam(lin1.parameters(), 0.05f);
var lin2 = torch.nn.Linear(10, 10);
var optim2 = torch.optim.Adam(lin2.parameters());
optim2.load_state_dict(optim1.state_dict())

[NiklasGustafsson]

In that case, shouldn't the 'copy' argument be passed down from where you know whether it's necessary or not?

[shaltielshmid]

I can add it in, if you think it's the right way to it.

The way I see it, the StateDictionary is an object which is created only by calling Optimizer.state_dict(), and the source Optimizer is the one which manages the tensors that are in the dictionary. If they are to be disposed, then the Optimizer would handle them. And therefore, the new Optimizer should receive a fresh copy of the tensors to manage itself.

[shaltielshmid]

Meaning, the StateDictionary is just an wrapper interface for accessing the values of an Optimizer.
I think we don't ever want two Optimizers to share the same tensor handle.

That being said, if you disagree and would like me to add a copy parameter to the load_state_dict function, no problem.

[NiklasGustafsson]

Got it. I think that's probably right. If we later conceive of a case where you do want to share state, we can add an argument later and set the default to 'true'

[NiklasGustafsson]

You have run the unit tests on CUDA, right?

[shaltielshmid]

Yup

[NiklasGustafsson]

Merged