how to enable torch.compile when use Megatron-LM as training backend?

[robotsp]

As described, torch.compile seems not compatible with megatron-lm, especially when I train the model with PP/TP/EP etc. How Can I Enable torch.compile for kernel optimization?

[robotsp]

@HuangJoJo