Loss refactor (#699)

Co-authored-by: Felipe Mello <felipemello@fb.com>

Author: felipemello1

apps/grpo/llama3_8b.yaml

@@ -39,6 +39,7 @@ generator:
     max_tokens: ${max_res_tokens}
     temperature: 1.0
     top_p: 1.0
+    logprobs: 1  # returns log probabilities for sampled tokens
 
 # Trainer configuration
 trainer:

apps/grpo/main.py

@@ -24,65 +24,16 @@
 from forge.observability.metrics import record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
 from forge.rl import collate, ComputeAdvantages, Episode, RewardActor
+from forge.rl.loss import GRPOLoss
 from forge.types import LauncherConfig, ProvisionerConfig
 from forge.util.checkpoint import drop_weights
 from forge.util.config import parse
 from forge.util.logging import get_logger
-from forge.util.ops import compute_logprobs
 from omegaconf import DictConfig, OmegaConf
 
 logger = get_logger("INFO")
 
 
-# TODO (T245547773): Consolidate with SimpleGRPOLoss in losses/grpo_loss.py
-# Currently duplicated because of function signature differences:
-# - This function takes logits + response, computes logprobs internally
-# - SimpleGRPOLoss takes pre-computed logprobs
-# - TitanTrainer passes logits, so would need wrapper or signature change
-# Consider refactoring TitanTrainer's loss interface to standardize this.
-def simple_grpo_loss(
-    logits: torch.Tensor,
-    response: torch.Tensor,
-    ref_logprobs: torch.Tensor,
-    advantages: torch.Tensor,
-    padding_mask: torch.Tensor,
-    beta: float = 1e-6,
-) -> torch.Tensor:
-    logprobs: torch.Tensor = compute_logprobs(logits, response)
-    kl = torch.exp(ref_logprobs - logprobs) - (ref_logprobs - logprobs) - 1
-    per_token_policy_loss = torch.exp(logprobs - logprobs.detach()) * advantages
-
-    # Compute mean KL per valid token
-    mean_kl = (
-        ((kl * padding_mask).sum(dim=1)) / (padding_mask.sum(dim=1).clamp(min=1.0))
-    ).mean()
-
-    # Compute mean policy loss per valid token
-    mean_policy_loss = (
-        ((per_token_policy_loss * padding_mask).sum(dim=1))
-        / (padding_mask.sum(dim=1).clamp(min=1.0))
-    ).mean()
-
-    # Compute loss using the means (mathematically equivalent)
-    loss = -(mean_policy_loss - beta * mean_kl)
-
-    # Log metrics
-    # TODO: Better design - have loss function return all metrics as a dict,
-    # then record them in rl_trainer so all training metrics are in one namespace
-    # and we avoid doing .item here, which is not compile friendly
-    record_metric("grpo_loss/kl_divergence_mean", mean_kl.item(), Reduce.MEAN)
-    record_metric(
-        "grpo_loss/kl_divergence_max", (kl * padding_mask).max().item(), Reduce.MAX
-    )
-    record_metric(
-        "grpo_loss/policy_gradient_loss", mean_policy_loss.item(), Reduce.MEAN
-    )
-    record_metric("grpo_loss/total_loss", loss.item(), Reduce.MEAN)
-    record_metric("grpo_loss/advantage_mean", advantages.mean().item(), Reduce.MEAN)
-    record_metric("grpo_loss/advantage_std", advantages.std().item(), Reduce.MEAN)
-    return loss
-
-
 async def main(cfg: DictConfig):
     """Main GRPO training loop with rollout and training processes."""
     # Convert OmegaConf config to plain dict
@@ -116,8 +67,32 @@ async def main(cfg: DictConfig):
         backend_config=metric_logging_cfg, run_config=run_config_for_logging
     )
 
+    # ---- Setup loss function ---- #
+    loss_fn = GRPOLoss(
+        clip_low=0.2,
+        clip_high=0.28,
+        beta=0.1,
+        agg_type="fixed_horizon",
+    )
+
+    # Fail-fast: Check loss/ref_model compatibility before spawning actors
+    uses_ref_model = cfg.get("services", {}).get("ref_model") is not None
+    if uses_ref_model and not isinstance(loss_fn, GRPOLoss):
+        raise ValueError(
+            f"ref_model is configured but {type(loss_fn).__name__} does not use ref_logprobs. "
+            "Either remove the ref_model service config or use GRPOLoss with beta > 0."
+        )
+    if isinstance(loss_fn, GRPOLoss) and loss_fn.beta > 0 and not uses_ref_model:
+        raise ValueError(
+            f"GRPOLoss with beta={loss_fn.beta} requires ref_logprobs, but ref_model is not configured. "
+            "Either add ref_model to services config or set beta=0."
+        )
+
     # ---- Setup services ---- #
 
+    async def noop():
+        return None
+
     (
         dataloader,
         generator,
@@ -130,13 +105,17 @@ async def main(cfg: DictConfig):
         DatasetActor.options(**cfg.actors.dataset).as_actor(**cfg.dataset),
         Generator.options(**cfg.services.generator).as_service(**cfg.generator),
         TitanTrainer.options(**cfg.actors.trainer).as_actor(
-            **cfg.trainer, loss=simple_grpo_loss
+            **cfg.trainer, loss=loss_fn
         ),
         ReplayBuffer.options(**cfg.actors.replay_buffer).as_actor(
             **cfg.replay_buffer, collate=collate
         ),
         ComputeAdvantages.options(**cfg.actors.compute_advantages).as_actor(),
-        ReferenceModel.options(**cfg.services.ref_model).as_service(**cfg.ref_model),
+        (
+            ReferenceModel.options(**cfg.services.ref_model).as_service(**cfg.ref_model)
+            if uses_ref_model
+            else noop()
+        ),
         RewardActor.options(**cfg.services.reward_actor).as_service(
             reward_functions=[MathReward(), ThinkingReward()]
         ),
@@ -187,7 +166,34 @@ async def continuous_rollouts():
                 (group_size, max_req_tokens + max_res_tokens),
                 dtype=torch.long,
             )
+            seq_len = max_req_tokens + max_res_tokens
+
             for i, response in enumerate(responses):
+                # Validate logprobs exist
+                if response.logprobs is None:
+                    raise ValueError(
+                        "Completion.logprobs is None. "
+                        "Ensure Generator returns logprobs by setting 'logprobs: 1' in sampling_params config."
+                    )
+
+                # Prepare generator_logprobs
+                # Shift by -1 to align with next-token prediction
+                actual_response_len = response.token_ids.shape[0]
+                generator_logprobs = torch.zeros(seq_len, dtype=response.logprobs.dtype)
+                generator_logprobs[
+                    max_req_tokens : max_req_tokens + actual_response_len
+                ] = response.logprobs
+                generator_logprobs = torch.roll(generator_logprobs, shifts=-1, dims=0)
+                generator_logprobs[-1] = 0.0
+
+                # Prepare loss_mask
+                response_mask = torch.zeros(seq_len, dtype=torch.float32)
+                response_mask[max_req_tokens : max_req_tokens + actual_response_len] = (
+                    1.0
+                )
+                loss_mask = torch.roll(response_mask, shifts=-1, dims=0)
+                loss_mask[-1] = 0.0
+
                 episode = Episode(
                     episode_id=str(uuid.uuid4()),
                     pad_id=pad_id,
@@ -197,7 +203,10 @@ async def continuous_rollouts():
                     request=prompt,
                     response=response.text,
                     completion=response,
+                    generator_logprobs=generator_logprobs,
+                    loss_mask=loss_mask,
                 )
+
                 (
                     episode.reward_breakdown,
                     episode.reward,
@@ -263,21 +272,33 @@ async def continuous_rollouts():
 
             t.step("reward_evaluation")
 
-            ref_logprobs = await ref_model.forward.route(
-                input_ids, max_req_tokens, return_logprobs=True
-            )
-            t.step("reference_model_calculate_logprobs")
+            # Compute ref_logprobs only if ref_model is configured
+            if ref_model is not None:
+                ref_logprobs = await ref_model.forward.route(
+                    input_ids, return_logprobs=True
+                )
+                t.step("reference_model_calculate_logprobs")
+
+                for i, episode in enumerate(episodes):
+                    episode.ref_logprobs = ref_logprobs[i]  # [seq_len]
 
-            for i, episode in enumerate(episodes):
-                episode.ref_logprobs = ref_logprobs[i]
-            del ref_logprobs, input_ids
+                del ref_logprobs
+
+            del input_ids
 
             advantages = await compute_advantages.compute.call_one(episodes)
             for episode, advantage in zip(episodes, advantages):
                 episode.advantage = advantage
                 await replay_buffer.add.call_one(episode)
 
-                sample = episode.to_dict(exclude=["ref_logprobs", "completion"])
+                sample = episode.to_dict(
+                    exclude=[
+                        "completion",
+                        "loss_mask",
+                        "generator_logprobs",
+                        "ref_logprobs",
+                    ]
+                )
                 sample["score"] = sample["reward"]
                 record_metric(
                     "main_samples/continuous_rollouts/sample_table",

apps/grpo/qwen3_1_7b.yaml

@@ -43,6 +43,7 @@ generator:
     max_tokens: ${max_res_tokens}
     temperature: 1.0
     top_p: 1.0
+    logprobs: 1  # returns log probabilities for sampled tokens
 
 # Trainer configuration
 trainer:

apps/grpo/qwen3_8b.yaml

@@ -39,6 +39,7 @@ generator:
     max_tokens: ${max_res_tokens}
     temperature: 1.0
     top_p: 1.0
+    logprobs: 1  # returns log probabilities for sampled tokens
 
 # Trainer configuration
 trainer:

src/forge/actors/reference_model.py

@@ -15,10 +15,9 @@
 from forge.controller import ForgeActor
 from forge.observability.metrics import record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
-from forge.util.ops import compute_logprobs
+from forge.rl.loss import compute_logprobs, create_shifted_targets
 from monarch.actor import current_rank, current_size, endpoint
 from torch.distributed.tensor import DTensor
-from torch.distributed.tensor.parallel import loss_parallel
 from torchtitan.config.job_config import (
     Checkpoint,
     Comm,
@@ -96,9 +95,9 @@ def __post_init__(self):
         self.rank = current_rank().rank
         self.size = math.prod(current_size().values())
 
-        self.compute_log_probs = compute_logprobs
+        self.compute_logprobs = compute_logprobs
         if self.compile.enable:
-            self.compute_log_probs = torch.compile(self.compute_log_probs)
+            self.compute_logprobs = torch.compile(self.compute_logprobs)
 
         env = {
             "RANK": str(self.rank),
@@ -128,23 +127,20 @@ async def setup(self):
 
     @endpoint
     async def forward(
-        self, input_ids: torch.Tensor, max_req_tokens: int, return_logprobs: bool
+        self, input_ids: torch.Tensor, return_logprobs: bool = True
     ) -> torch.Tensor:
         """
         Args:
-            input_ids (torch.Tensor): input token ids with shape [group_size, req + res length].
-            max_req_tokens (int): maximum request length.
+            input_ids (torch.Tensor): input token ids with shape [group_size, seq_len].
             return_logprobs (bool): whether to return log probabilities instead of raw logits.
 
             return_logprobs flag significantly impacts the amount of data transferred to the caller:
-            - When False: Returns logits with shape [group_size, req + res_length, vocab_size].
+            - When False: Returns logits with shape [group_size, seq_len, vocab_size].
               This includes the full vocabulary distribution for each token position.
 
-            - When True: Returns log probabilities with shape [group_size, req_length].
-              This only includes probabilities for the request tokens, significantly reducing memory
-              usage and transfer overhead.
+            - When True: Returns log probabilities with shape [group_size, seq_len].
+              Prompt positions will have logprobs = 0.
         """
-        # Record reference model metrics
         record_metric("reference_perf/forward/count_forward_passes", 1, Reduce.SUM)
 
         t = Tracer("reference_perf/forward", timer="gpu", track_memory=True)
@@ -175,24 +171,16 @@ async def forward(
                 with self.engine.maybe_enable_amp:
                     with torch.inference_mode():
                         logits = self.model(input_ids)
-        self.step += 1
 
-        if not return_logprobs:
-            if isinstance(logits, DTensor):
-                logits = logits.full_tensor()
-            t.stop()
-            return logits
-        else:
-            response_tokens = input_ids[:, max_req_tokens:]
-            if parallel_dims.tp_enabled and isinstance(logits, DTensor):
-                with loss_parallel():
-                    logprobs = self.compute_log_probs(logits, response_tokens)
-
-                # loss_parallel produces Replicated output - to_local() returns the full tensor
-                logprobs = logprobs.to_local()
-            else:
-                if isinstance(logits, DTensor):
-                    logits = logits.full_tensor()
-                logprobs = self.compute_log_probs(logits, response_tokens)
-            t.stop()
-            return logprobs
+                        if return_logprobs:
+                            target_ids = create_shifted_targets(input_ids)
+                            logprobs, _ = self.compute_logprobs(logits, target_ids)
+
+        out = logprobs if return_logprobs else logits
+
+        if isinstance(out, DTensor):
+            out = out.full_tensor()
+
+        self.step += 1
+        t.stop()
+        return out

src/forge/actors/trainer/titan.py

@@ -19,6 +19,7 @@
 from forge.data.utils import batch_to_device
 from forge.observability.metrics import record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
+from forge.rl.loss import create_shifted_targets
 from monarch.actor import endpoint
 from torch import Tensor
 from torch.distributed.checkpoint._nested_dict import flatten_state_dict
@@ -122,15 +123,34 @@ def forward_backward(
         model_parts = self.engine.model_parts
         parallel_dims = self.engine.parallel_dims
         optional_context_parallel_ctx = None
+
+        # Create shifted target_ids for next-token prediction
+        # target_ids[i] = input_ids[i+1], with loss_mask applied
+        targets["target_ids"] = create_shifted_targets(
+            inputs["tokens"], targets.get("loss_mask")
+        )
+
         if parallel_dims.pp_enabled:
             raise NotImplementedError("PP not implemented yet")
         else:
             with self.engine.train_context(optional_context_parallel_ctx):
                 assert len(model_parts) == 1
                 with self.engine.maybe_enable_amp:
                     logits = model_parts[0](**inputs)
-                    loss = self.loss(logits, **targets)
-                del logits  # Free to before bwd to avoid peaking memory
+                    loss_output = self.loss(logits, **targets)
+                    loss = loss_output.loss
+
+                # Record metrics from loss output
+                for metric in loss_output.metrics:
+                    value = (
+                        metric.value.item()
+                        if isinstance(metric.value, torch.Tensor)
+                        else metric.value
+                    )
+                    record_metric(metric.key, value, metric.reduction, metric.timestamp)
+
+                # Free to before bwd to avoid peaking memory
+                del logits, loss_output.metrics
                 loss.backward()
         self._accumulated_microbatches += 1
         return loss