How to use torch.compile in DeepSpeed?

[noob-ctrl]

If I want to use the new feature of Pytorch2.0——torch.compile, what should I do? Where should I put the following code or just pass a command line parameter?

model = torch.compile(model)

[c-maxey]

I was able to get it to work by placing the compile decorator @torch.compile above each forward function of my models. Just a heads up, you'll see a large increase in training time for a first pass or two, then you should see some speedup.

[noob-ctrl]

@c-maxey Sorry, I still don't quite get it, can you show me a simple example？

[c-maxey]

@c-maxey Sorry, I still don't quite get it, can you show me a simple example？

Hey no problem. Where your forward function is for your model, make it looks like the following:

@torch.compile def forward(self, )

[loadams]

Hi @noob-ctrl - we don't fully support all of torch 2's features just yet. @c-maxey has good insight on what you can do here, but there should be better instructions on how to leverage torch 2 and deepspeed in the future.

[Cherryjingyao]

@c-maxey Sorry, I still don't quite get it, can you show me a simple example？

Hey no problem. Where your forward function is for your model, make it looks like the following:

@torch.compile def forward(self, )

did this work? I use this way, but got a bug. I guess maybe the my cuda version is 11.4 which can't surpport torch2.0 compile?