[RL] Changes to enable compilation for trainer

[Lucaskabela]

Summary

In this PR, we enable naive, JIT style torch.compile for the RL policy trainer. This is the first step towards speeding up the trainer model. Changes are:

1. Wiring through compilation config:

• Added TrainerCompileConfig dataclass with enable (bool) and backend (str, default "eager") fields
• Added compile field to PolicyTrainer.Config with compile and aot_eager backend
• Added apply_compile() method that calls .compile(backend=..., fullgraph=True) on each transformer
  layer -> This is crticial, as torch.compile() results in weight name change which breaks the weight transfer

2. config_registry.py — Enable compile by default in configs

• Both rl_grpo_qwen3_0_6b and rl_grpo_qwen3_debug configs now set
  compile=TrainerCompileConfig(enable=True, backend="aot_eager")

3. vllm_compat/models/attention.py — Make attention operation compile-compatible

• Moved the FlashAttnWithBackward autograd function out of the forward() method (nested classes
  can't be traced by the compiler) into a module-level FlashAttnVarlenFunction
• Registered the flash-attention forward as a torch.library.custom_op (rl::flash_attn_varlen_fwd)
  with a fake implementation, so AOT Autograd can trace through it with FakeTensors
• Simplified the call site in VLLMCompatibleFlashAttention.forward() to use the new function

Test Plan

python torchtitan/experiments/rl/simple_grpo_sum_digits.py --module rl --config rl_grpo_qwen3_0_6b --hf_assets_path=torchtitan/experiments/rl/example_checkpoint/Qwen3-0.6B

Results in the same losses as on main - the timing is now like:

Main

[actor=<root>] Step  9 | Loss: +0.0031 | Reward: +0.450 | Correct: 29/40 | Avg tokens: 100 | Logprob diff: mean=-1.1882e-01, max=1.0228e+01 | Time: 21.4s
[actor=<root>]   Step Timing | Generator: 2.3s | Trainer: 15.3s | WeightSync: 3.8s
[actor=<root>] Cumulative Timing | Generator: 21.3s | Trainer: 150.2s | WeightSync: 33.1s | Total: 204.6s
[actor=<root>] RL Training complete
[actor=<root>] Evaluating post-training performance...
[actor=<root>] Eval: Accuracy=55% (11/20) Format=95% (19/20)
[actor=<root>] ================================================================================
[actor=<root>] Pre-training:  Accuracy=40% (8/20) Format=50% (10/20)
[actor=<root>] Post-training: Accuracy=55% (11/20) Format=95% (19/20)
[actor=<root>] ================================================================================```

Changes

[actor=<root>] Step  9 | Loss: +0.0031 | Reward: +0.450 | Correct: 29/40 | Avg tokens: 100 | Logprob diff: mean=-1.1882e-01, max=1.0228e+01 | Time: 13.9s
[actor=<root>]   Step Timing | Generator: 2.4s | Trainer: 6.5s | WeightSync: 5.0s
[actor=<root>] Cumulative Timing | Generator: 21.7s | Trainer: 88.5s | WeightSync: 30.6s | Total: 140.9s
[actor=<root>] RL Training complete
[actor=<root>] Evaluating post-training performance...
[actor=<root>] Eval: Accuracy=55% (11/20) Format=95% (19/20)
[actor=<root>] ================================================================================
[actor=<root>] Pre-training:  Accuracy=40% (8/20) Format=50% (10/20)
[actor=<root>] Post-training: Accuracy=55% (11/20) Format=95% (19/20)
[actor=<root>] ================================================================================
(vllm) [lucaskabela@devgpu007.eag6 ~/torchtitan (lucaskabela/enable_trainer_compile_03_10)]$ pytho

So we save ~60s of runtime e2e via compilation in this manner without affecting our logits/accuracy

[tianyu-l]

if we go with pytorch varlen, do we still need to worry about this file? cc @wwwjn

[Lucaskabela]

This is post rebase - yes we do need these changes for the following reasons:

1. Moving torch.autograd.Function out of the inner context - this is so compile can trace it
2. _flash_attn_varlen_fwd custom op - this is so torch can trace it; without custom op, we don't know how shape will propogate through this

[liangel-02]

if ur going with pytorch varlen i think u can directly call varlen_attn (which is already a custom op) pytorch's varlen calls the upstream flash attention impl instead of vllm's flash attention

[Lucaskabela]

Sorry I misunderstood this question (I thought this was related to the deletion of vllm_attention directory).

We will need to move the autograd.Function out but yes we won't need the custom op in that case

[tianyu-l]

let's not merge this, let use pytorch varlen attention in #2364

cc @zhxchen17 to unblock

[tianyu-l]

@wwwjn could you work with @sanketpurandare to land PP + DTensor so that we can remove this ad hoc function?

[wwwjn]

Sure, let me follow up with @sanketpurandare

[tianyu-l]

we switched to a central impl in #2615, but I understand that this is doing model.compile() not torch.compile(model)

@fegin can we universally switch to model.compile style (and it should still work with DCP)?