Why are these changes needed?

We found the performance of LSTMs in Rllib to be extremely slow compared to other methods, with a single training iteration of PPO taking 179 seconds (compared to ~9 seconds with a similarly-sized MLP network). This made RNNs/LSTMs, as well as some transformer implementations, completely unusable for our purposes.

However, when profiling, we found this was primarily due to a very slow copy operation:

Further investigation revealed that most of this runtime was spent copying theinfos dict. We determined that the root cause was inconsistent handling of the dictionary inrnn_sequencing. While the non-recurrent implementation stores the list of dictionaries as a NumPy array of objects,rnn_sequencing instead stores it as a Python list:

We applied a one-line fix to make this behavior consistent and store the list as NumPy array:

# old and slowf_pad= [None]*length# new and fastf_pad=np.full([length],None,dtype=f.dtype)

This causes thecopy function to perform a shallow copy, drastically improving performance by ~6x to around 29 seconds:

However, we found that the training loop was still spending a lot of time in rnn_sequencing. We traced this down to a slow element-wise copy into an array. We instead replaced this with a vectorized copy:

# old and slowforseq_offsetinrange(len_):f_pad[seq_base+seq_offset]=f[i]i+=1# new and fastf_pad[seq_base :seq_base+len_]=f[i :i+len_]i+=len_

This further improved performance (in this sample, we also removed the 1-time summary logging, but did not remove this in this PR):

Combined, these changes reduced the run time of the training step from 179 seconds to 15 seconds, approximately a 12x speedup and competitive with training an MLP.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e.,git commit -s) in this PR.
• I've runscripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed forhttps://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 indoc/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 athttps://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(