Last updates

Author: joecummings

apps/grpo/main.py

@@ -196,9 +196,9 @@ async def train_step(self, batch: list[list[Episode]]):
 
         mask = response != pad_id
         loss = self.loss(logprobs, ref_logprobs, advantages, mask)
-        self.optimizer.zero_grad()
         loss.backward()
         self.optimizer.step()
+        self.optimizer.zero_grad(set_to_none=True)
 
         return loss.item()
 
@@ -447,7 +447,7 @@ async def continuous_training():
             if batch is None:
                 await asyncio.sleep(0.1)
             else:
-                loss = sum(await trainer.train_step.call(batch))
+                loss = await trainer.train_step.choose(batch)
                 training_step += 1
                 mlogger.log("loss/training_step", loss, training_step)
                 await trainer.push_weights.call(policy_version)

src/forge/actors/policy.py

@@ -243,9 +243,6 @@ async def generate(self, prompt: str, priority: int = 0) -> RequestOutput:
         Returns:
             RequestOutput: vLLM class with the generated response.
         """
-        return await self._generate(prompt, priority)
-
-    async def _generate(self, prompt: str, priority: int = 0) -> RequestOutput:
         self.request_id += 1 % sys.maxsize
         request_id = str(self.request_id)  # implement from a counter