Draft: Feat: Enable torch compile

[jiemingz]

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[terrykong]

Could you add this key to all the configs/recipes?

[terrykong]

Is this possible to unit test?

[terrykong]

@jiemingz is the only thing blocking this PR the seq-packing change since we need static shapes for torch.compile?

[terrykong]

Dependent on #300

[SahilJain314]

Dtensor sequence packing has been merged. @ahmadki to support max-padding packed sequences in DTensor to enable torch.compile (fixed seqlen).

[ahmadki]

Dtensor sequence packing has been merged. @ahmadki to support max-padding packed sequences in DTensor to enable torch.compile (fixed seqlen).

tracking here

[StrongerXi]

Could you try model.compile() instead? That should fix the _orig_mod issue. This is also the recommended way of compiling a model now. We'll work on throwing warnings and publicizing to raise awareness on this.

[euronymous-aithal]

@yuki-97 can you please review this and take this forward.
(QQ: @ffrujeri and @joyang-nv not sure if Automodel path already enables torch.compile if so we should close this )

[joyang-nv]

@euronymous-aithal Hi Ashwath, automodel indeed has torch.compile support already. But we have never enabled within nemo RL. I think it is proper to make it fixed with FP8 dtensor is ready @RayenTian is working on it, since FP8 dtensor train requires to enable torch.compile. We still have encounter torch.compile bug when TP>1. Still investigation.

[terrykong]

Closing since this is w.r.t. v1 and we are currently centralizing efforts around v2 . closing this and reassigned #4 to @joyang-nv for now