[rl] Module forward call in RL crashes when using AC and torch.compile

[daniellepintz]

Bug description

Error:

    File "/torchtitan/titan-rl/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1058, in compile_wrapper
      raise e.with_traceback(None) from e.__cause__  # User compiler error
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  torch._dynamo.exc.Unsupported: HOP: Unsafe side effect
    Higher Order Operator: torch.utils.checkpoint.checkpoint
    Explanation: Mutating a variable from outside the scope of this HOP is not supported.
    Hint: If the HOP is activation checkpointing (torch.utils.checkpoint.checkpoint), this points to a side effect in forward method. Eager activation checkpointing replays that side-effect while recomputing the forward in the backward. If you are ok with side-effect not replayed in the backward, try setting `torch._dynamo.config.skip_fwd_side_effects_in_bwd_under_checkpoint = True`

Command to repro:

python torchtitan/experiments/rl/simple_grpo_sum_digits.py --module rl --config rl_grpo_qwen3_0_6b --trainer.ac_config.mode selective

Versions

latest PT nightly