Hello everyone,
I’m training a deep reinforcement learning agent that leverages GRU. In my training setup the sequence length is set to the length of the longest episode.
Now that I’m trying out torch.compile(), every epoch compiles a new backwards graph. After 40 epochs, 266 backwards graphs were compiled.
Is my assumption correct, that torch.compile (in the context of recurrent layers) should only be used if the sequence length is fixed?
There’s an argument called torch.compile(..., dynamic=True) but you should try it out on master and expect some bugs
Thanks for the heads up! Setting the dynamic argument is not working on the current release. I’ll wait for newer versions.
I revisited this issue today and tried the recent nightly version of PyTorch. The rercurrent model successfully compiles under dynamic=True and mode=max-autotune.
Before compilation an A100 needed 6.84s on average per PPO Update. After compilation, there is no notable speed up 6.87s.
I conducted a training run without GRU Cell. Without compiling the model, I measured 5.56s and with compilation 5.25s.
My code also features TransformerXL training. If I compile this model, the forward pass outputs infinity. To debug this issue, I compiled every single submodule on its own and did not run into this exception. So far I did not measure any performance gains.
Please, does anyone have any updates on this? I am using PyTorch 2.1.0 and I am facing the same issue.