torch.compile with FSDP

facebookresearch / fairscale

PyTorch extensions for high performance and large scale training.

Other

3.15k stars 275 forks source link

torch.compile with FSDP #1145

Closed santha96 closed 10 months ago

santha96 commented 11 months ago

Hi, is it possible to use torch.compile with fsdp . if so , what is the right way and when to use torch.compile ?

awgu commented 10 months ago

torch.compile is not supported with Fairscale FSDP. There is preliminary support with PyTorch FSDP when passing use_orig_params=True, where the compute in each FSDP-wrapped module is compiled. More thorough support is under development.

awgu commented 10 months ago

I am marking this as closed because torch.compile support for Fairscale FSDP is not planned.