refactor: make compile configurable for logprobs computation

Author: gitlost-murali

src/forge/actors/reference_model.py

@@ -99,6 +99,10 @@ def __post_init__(self):
         self.rank = current_rank().rank
         self.size = math.prod(current_size().values())
 
+        self.compute_log_probs = compute_logprobs
+        if self.compile.enable:
+            self.compute_log_probs = torch.compile(self.compute_log_probs)
+
         env = {
             "RANK": str(self.rank),
             "LOCAL_RANK": str(self.rank),
@@ -193,11 +197,11 @@ async def forward(
             response_tokens = input_ids[:, max_req_tokens:]
             if parallel_dims.tp_enabled and isinstance(logits, DTensor):
                 with loss_parallel():
-                    logprobs = compute_logprobs(logits, response_tokens, align=True)
+                    logprobs = self.compute_log_probs(logits, response_tokens)
 
                 logprobs = logprobs.to_local()
             else:
-                logprobs = compute_logprobs(logits, response_tokens)
+                logprobs = self.compute_log_probs(logits, response_tokens)
             t.step("compute_logprobs")
             t.stop()
             return logprobs

src/forge/util/ops.py

@@ -9,7 +9,6 @@
 from torch.distributed.tensor import DTensor
 
 
-@torch.compile
 def compute_logprobs(
     logits: torch.Tensor | DTensor,
     input_ids: torch.Tensor,