[RLlib] Use PyTorch vectorized max() and sum() in SampleBatch.__init__ when possible

[cassidylaidlaw]

Why are these changes needed?

Currently, there are two lines in SampleBatch.__init__ where the Python builtins sum and max are used on seq_lens:

1. self.max_seq_len = max(seq_lens_)
2. self.count = sum(self[SampleBatch.SEQ_LENS])

However, if the seq_lens are a PyTorch tensor on the GPU, this is incredibly slow. I believe this is because each element of seq_lens has to be fetched independently from the GPU when iterating over the tensor for sum and max. Thus, I have changed the two lines to use seq_lens_.max().item() and self[SampleBatch.SEQ_LENS].sum().item() for PyTorch tensors.

This significantly speeds up training with models that require state. For instance, consider the following training run:

rllib train --run PPO --env CartPole-v0 \
--config '{"train_batch_size": 2000, "sgd_minibatch_size": 1000, "num_sgd_iter": 100, "framework": "torch", "num_workers": 10, "num_gpus": 1, "model": {"use_lstm": true}, "rollout_fragment_length": 1}' \
--stop '{"training_iteration": 1}'

Before this PR, it takes 22.7 s to run, but after only 13.5 s. If look at result["timers"]["learn_time_ms"], we can also see that it is 8042 before the PR but 3162 after. Thus in some cases we're getting more than a 2x speedup for SGD!

Related issue number

I haven't opened an issue.

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

[stale]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

• If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[cassidylaidlaw]

It looks like the only tests that are failing in CI are ones that are flaky. Can somebody review this?

[cassidylaidlaw]

Updated the PR and tests seem to be passing again. Can anyone look over and/or merge? @sven1977 @gjoliver @avnishn @ArturNiederfahrenhorst @smorad @maxpumperla @kouroshHakha @krfricke

[kouroshHakha]

sound reasonable to me.

[cassidylaidlaw]

Another bump—will this be merged soon?

[stale]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

• If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[cassidylaidlaw]

Is there anything I can do to get this merged? @kouroshHakha

[kouroshHakha]

cc @gjoliver @sven1977

[gjoliver]

cool. looks like an awesome change for torch.
does this make things slower for TF? just want to have an idea.

[cassidylaidlaw]

It shouldn't make things any slower in TF. The PyTorch vectorized max/sum are only used if the tensors are from torch, otherwise the code runs the same as before the PR.