[air] Horovod: Use Torch.encode_data if torch is imported

[krfricke]

Signed-off-by: Kai Fricke [email protected]

Why are these changes needed?

Horovod with Tune does not work out of the box for GPU checkpoints as they get deserialized on the non-GPU trainer worker, leading to errors. With this PR, we detect if torch is imported and a tensor is supplied in the Horovod backend. If so, we use the torch backend to serialize the data.

Related issue number

Closes #28439

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 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 :(

[amogkam]

Thanks @krfricke, lgtm as a stopgap fix! But ultimately we should refactor the checkpoint encoding/decoding logic out of the Backends and into the framework-specific checkpoints.

Then, when saving a Torch model via TorchCheckpoint.from_model(), the same encode/decode logic will apply regardless of if I'm using TorchTrainer or HorovodTrainer.

Made an issue to track this here: #28462

[amogkam]

Do we need this? I think if torch is installed, it should be safe to always use the TorchBackend for encoding/decoding (even if the data dict does not contain a tensor). I'm worried in the worst case contains_tensor can lead to a lot of recursion.

[krfricke]

I originally didn't have it in, but thought the overhead wouldn't be as bad. But I agree, since it only concerns an internal communication channel and the intermediate objects are not exposed to the user, we can just do this always when torch is loaded. Updated the PR

[krfricke]

Turns out we do need it as torch.save seems to silently fail if a full model is passed (and not a state dict).

I think it should be fine - a similar lookup has to be in pickling after all, and in most cases it should finish early.