PyTorch torch.compile, scripting, tracing for faster computation (specifically training)

[albertz]

For the first few steps, it could run without tracing/scripting, but then it could enable it and from then on use the Torch graph directly (very similar to TF computation graph). So maybe some option like torch_jit_enable_after_step: int.

Tracing should always work for static sizes and no (or always same) control flow.

Questions:

• How to handle dynamic dims? With tracing for ONNX, this is somehow possible, but how to use this directly?
• How to handle dynamic control flow? Maybe we can use scripting just for the loop, and tracing otherwise? With RF, we already try to make control flow explicit, so there we should be able to do all automatically?

Instead of using the old-style scripting/tracing, probably better would be to use torch.compile now. I'm not sure if the questions above are relevant there as well.

This might be even more helpful for RF, due to the overhead of RF (although the overhead has greatly been reduced now, see #1402), but it would just be as helpful for pure PT as well. This feature would not be specific to RF in any way, except that RF maybe can provide special handling of control flow for this. For RF, some other aspects might become relevant, so #1491 is specifically about this.

[albertz]

Some links:

https://pytorch.org/docs/stable/generated/torch.jit.trace.html
https://pytorch.org/docs/stable/generated/torch.jit.ignore.html#torch.jit.ignore
https://pytorch.org/docs/stable/jit.html#mixing-tracing-and-scripting
https://github.com/rwth-i6/returnn/wiki/PyTorch-optimizations

[albertz]

Just now, I found about torch.export, which references functorch.experimental.control_flow for control flow:

https://pytorch.org/tutorials/intermediate/torch_export_tutorial.html
https://pytorch.org/executorch/stable/tutorials/export-to-executorch-tutorial.html