Loss updates

Author: joecummings

apps/grpo/main.py

@@ -49,24 +49,21 @@ def compute_logprobs(
 
 class SimpleGRPOLoss(nn.Module):
     """Simplified GRPO Loss for simplified single step updates
-    Copied from https://github.com/pytorch/torchtune/blob/main/torchtune/dev/grpo/loss.py.
+    Inspired by the Hugging Face TRL implementation:
+        https://github.com/huggingface/trl/blob/417915a3e4d3e3bc8d7b196594308b8eabf928be/trl/trainer/grpo_trainer.py#L1624.
     """
 
     def __init__(self, beta: float = 0.1):
         super().__init__()
         self.beta = beta
 
     def forward(self, logprobs, ref_logprobs, advantages, padding_mask):
-        per_token_kl = (
-            torch.exp(ref_logprobs.detach() - logprobs)
-            - (ref_logprobs.detach() - logprobs)
-            - 1
-        )
+        kl = torch.exp(ref_logprobs - logprobs) - (ref_logprobs - logprobs) - 1
         per_token_policy_loss = torch.exp(logprobs - logprobs.detach()) * advantages
-        per_token_loss = -(per_token_policy_loss - self.beta * per_token_kl)
+        per_token_loss = -(per_token_policy_loss - self.beta * kl)
         loss = (
-            (per_token_loss * padding_mask).sum(dim=1)
-            / (padding_mask.sum(dim=1) + 1e-8)
+            ((per_token_loss * padding_mask).sum(dim=1))
+            / (padding_mask.sum(dim=1).clamp(min=1.0))
         ).mean()
         return loss
 
@@ -211,21 +208,21 @@ def _qwen3_hf_to_vllm(self, saved_sd):
         return load_sd
 
     @endpoint
-    async def train_step(self, batch: list[Episode]):
-        batch = batch[self.dp_rank]
-        pad_id = batch[0].pad_id
+    async def train_step(self, batch: list[list[Episode]]):
+        microbatch = batch[self.dp_rank]
+        pad_id = microbatch[0].pad_id
 
         # prepare batch
-        request = [e.request_tensor for e in batch]
+        request = [e.request_tensor for e in microbatch]
         request = torch.stack(request).to(self.device)  # [b x s]
 
-        response = [e.response_tensor for e in batch]
+        response = [e.response_tensor for e in microbatch]
         response = torch.stack(response).to(self.device)  # [b x s]
 
-        ref_logprobs = [e.ref_logprobs for e in batch]
+        ref_logprobs = [e.ref_logprobs for e in microbatch]
         ref_logprobs = torch.stack(ref_logprobs).to(self.device).squeeze()  # [b x s]
 
-        advantages = [e.advantage for e in batch]
+        advantages = [e.advantage for e in microbatch]
         advantages = torch.tensor(advantages).to(self.device).unsqueeze(-1)  # [b x 1]
         del batch
 
@@ -522,10 +519,10 @@ async def continuous_training():
         )
 
 
-@parse
-def recipe_main(cfg: DictConfig) -> None:
-    asyncio.run(main(cfg))
+if __name__ == "__main__":
 
+    @parse
+    def _main(cfg):
+        asyncio.run(main(cfg))
 
-if __name__ == "__main__":
-    recipe_main()
+    _main()  # @parse grabs the cfg from CLI

src/forge/actors/policy.py

@@ -11,15 +11,8 @@
 from collections.abc import Mapping
 from copy import copy
 from dataclasses import asdict, dataclass, field, fields
-from typing import Dict, List
 
 import torch
-
-from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
-
-from forge.data.sharding import VLLMSharding
-from forge.interfaces import Policy as PolicyInterface
-from forge.types import ProcessConfig
 from monarch.actor import current_rank, endpoint, ProcMesh
 from torchstore import MultiProcessStore
 from torchstore._state_dict_utils import DELIM
@@ -44,6 +37,12 @@
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.worker.worker_base import WorkerWrapperBase
 
+from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
+
+from forge.data.sharding import VLLMSharding
+from forge.interfaces import Policy as PolicyInterface
+from forge.types import ProcessConfig
+
 
 @dataclass
 class SamplingConfig: