[RLlib] Removes device infos from state when saving RModules to checkpoints/states.

[simonsays1980]

Why are these changes needed?

When loading an RLModule on CPU from a checkpoint/state that was created from a replica on GPU, an error occurs. This PR fixes this error by forcing the module to save its state in form of numpy.NdArrays.

Related issue number

Closes #43905

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[sven1977]

I'm not sure I understand the exact need for this additional option to the API. My main argument against this would be: When stuff gets stored to a checkpoint, it should be stored in a device-independent fashion. So the issue at hand here is NOT the loading from the checkpoint, but the saving to the checkpoint beforehand, which - I'm guessing - probably happened in torch.cuda tensors, NOT in numpy format.

Can we rather take the opposite approach to keep the mental model of what a checkpoint should be clean? Always save weights (and other tensor/matrix states) as numpy arrays, never as torch or tf tensors. When loading from a checkpoint, the sequence should be something like:
0) RLModule.from_checkpoint(dir=...)

1. load numpy arrays from dir file (pickle?)
2. pack these numpy arrays into a state dict, mapping
3. Call RLModule.set_state(state_dict=...) -> This automatically converts the numpy contents of state_dict into torch tensors (with self.device as the device of the already existing RLModule object), then performs a torch.nn.Module.load_state_dict() operation using these (cuda?) tensors.

[simonsays1980]

I'm not sure I understand the exact need for this additional option to the API. My main argument against this would be: When stuff gets stored to a checkpoint, it should be stored in a device-independent fashion. So the issue at hand here is NOT the loading from the checkpoint, but the saving to the checkpoint beforehand, which - I'm guessing - probably happened in torch.cuda tensors, NOT in numpy format.

Can we rather take the opposite approach to keep the mental model of what a checkpoint should be clean? Always save weights (and other tensor/matrix states) as numpy arrays, never as torch or tf tensors. When loading from a checkpoint, the sequence should be something like: 0) RLModule.from_checkpoint(dir=...)

1. load numpy arrays from dir file (pickle?)
2. pack these numpy arrays into a state dict, mapping
3. Call RLModule.set_state(state_dict=...) -> This automatically converts the numpy contents of state_dict into torch tensors (with self.device as the device of the already existing RLModule object), then performs a torch.nn.Module.load_state_dict() operation using these (cuda?) tensors.

I agree with your argument that we should ensure that checkpointing is device-independent. This should be the cleanest way of doing this. We should investigate, where exactly this device-dependent checkpointing takes place and fix the problem there.

I am, however, not so sure, if 3. describes how the workflow runs right now. Here is what makes me wonder: If the torch.nn.Module.load_state_dict() would use the self.device then it should not matter, if the state contains numpy.arrays or torch.tensor (with their own device attribute) as it uses always the module's device. But no matter what - checkpointing always numpy.array will avoid this error anyways.

[sven1977]

Remove this line?

[sven1977]

Remove this line?

[sven1977]

Perfect! This should work.

[sven1977]

Looks good! Thanks for this important fix @simonsays1980 !

Just two nits on the docstrings.

[sven1977]

Suggested change

	def load_state(
	self,
	dir: Union[str, pathlib.Path],
	) -> None:
	def load_state(self, dir: Union[str, pathlib.Path]) -> None: