Support torch.compile

[tohtana]

This issue aims to track the progress, bugs, and requests regarding the support of torch.compile.
@tjruwase

Tasks

• ZeRO3 support (Enable torch.compile with ZeRO (Experimental) #4878, we currently break the graph to make communication collectives work)
• Pipeline parallel support ([REQUEST] torch.compile + DeepSpeed #4677)
• Activation checkpointing with Megatron-DeepSpeed (Enable torch.compile Megatron-DeepSpeed#322)
• Pass HuggingFace CI tests

PRs

• Enable torch.compile with ZeRO (Experimental) #4878: The initial proposal of the official torch.compile support

Issues

• [REQUEST] torch.compile + DeepSpeed #4677
• How to use torch.compile in DeepSpeed? #3375
• [BUG] torch.compile doesn't work with stage 2 on 32 GPUs #3032

[tjruwase]

Thanks @tohtana!
@stas00, FYI

[zaptrem]

Hello, quick related question to #4878 : We can only compile certain parts of our model so just compiling the whole thing as the new option does wouldn't work. Is there a path for this sort of model?

[tohtana]

Hi @zaptrem, thank you for your question!

You can pass a function to compile the model with set_compiler_fn. I think you can define a function that compiles only the necessary part of the model and pass it to this API.

You'll find an example of the API usage here:
https://github.com/microsoft/DeepSpeed/blob/master/tests/unit/runtime/compile/test_load_config.py

[BramVanroy]

Is still actively being worked on?