How to compile a RNN loop once?

username: proger
cooked: <p>Hi,</p>
<p>I’d like to use the feature of torch.compile that unrolls for loops to implement RNNs. When using multiple identical layers of the same RNN I’ve noticed compilation time grows proportional to the number of layers: there is no reuse of the code which uses a lot of time and memory.</p>
<p>In practice this means I can’t compile a reasonably large RNN successfully.</p>
<p>Example model with complation metrics: <a href="https://gist.github.com/proger/746ceb4aec2e61cfac4a34f1787be976/eb686ff61c863f9fc230095cc205f02693e0fde8" class="inline-onebox" rel="noopener nofollow ugc">compile_rnns.py · GitHub</a></p>
<p>I tried looping through layers in eager mode and only using the compiled forward function, however the function was specialized for each layer anyway: <a href="https://gist.github.com/proger/746ceb4aec2e61cfac4a34f1787be976/revisions#diff-0776dda4c2dddb4121e0093cc461f4c961927f1ee452ea47caa6410630486f96" class="inline-onebox" rel="noopener nofollow ugc">Revisions · compile_rnns.py · GitHub</a></p>
<p>TORCH_LOGS=guards mention the following guards. <code>___check_obj_id</code> looks suspicious, does it mean the specialization is to the particular tensor?</p>
<pre><code class="lang-auto">GUARDS:
___check_type_id(L['T'], 94350706369312)                                      
L['T'] == 1024                                                                
hasattr(L['x'], '_dynamo_dynamic_indices') == False           # _dynamo/variables/builder.py:1410 in wrap_fx_proxy_cls 
___check_obj_id(L['input'], 140194674075456)                
L['input'].training == True                                                   
___check_obj_id(L['forget'], 140194674076224)               
L['forget'].training == True                                                  
hasattr(L['hidden'], '_dynamo_dynamic_indices') == False      # _dynamo/variables/builder.py:1410 in wrap_fx_proxy_cls 
___check_obj_id(L['output'], 140194674077184)               
L['output'].training == True                                                  
utils_device.CURRENT_DEVICE == None                           # _dynamo/output_graph.py:374 in init_ambient_guards
___skip_backend_check() or ___current_backend() == ___lookup_backend(140196638382528)  # _dynamo/output_graph.py:380 in init_ambient_guards
check_tensor(L['x'], Tensor, DispatchKeySet(CUDA, BackendSelect, ADInplaceOrView, AutogradCUDA), torch.float32, device=0, requires_grad=True, size=[8, 1024, 64], stride=[65536, 64, 1])  # _dynamo/variables/builder.py:1410 in wrap_fx_proxy_cls 
check_tensor(L['hidden'], Tensor, DispatchKeySet(CUDA, BackendSelect, ADInplaceOrView, AutogradCUDA), torch.float32, device=0, requires_grad=False, size=[8, 64], stride=[64, 1])  # _dynamo/variables/builder.py:1410 in wrap_fx_proxy_cls 
</code></pre>
<p>Could you advise on the best tactic to successfully compile long loops?</p>
<p>Thanks</p>

username: proger
cooked: <p>Turns out the hunch to think about <code>___check_obj_id</code> was in the right direction. The code I wrote was passing <code>nn.Module</code> objects as arguments. I changed the arguments to have type <code>nn.Parameter</code>, which made dynamo produce more <code>check_tensor</code> guards.</p>
<p>The result is available here: <a href="https://gist.github.com/proger/3543f2e1a369ba0a4aaf16a1625272fa" class="inline-onebox" rel="noopener nofollow ugc">compile_rnns3.py · GitHub</a></p>
<p>The compiled RNN is about 10x faster than eager for sequence length 128 and about 8x faster for sequence length 256. Compilation crashed due to stack explosion somewhere in torch.fx on longer sequences.</p>
<p>Btw love torch.export.</p>