Merge branch 'main' into language-reward-feature

Author: casteryh

apps/grpo/main.py

@@ -439,23 +439,21 @@ async def continuous_rollouts():
                 input_ids[i, :max_req_tokens] = episode.request_tensor
                 input_ids[i, max_req_tokens:] = episode.response_tensor
 
-                # drop episodes if
-                # 1> reward std-dev is very small (including all 0s and all 1s)
-                # 2> response is potentially truncated (response_len >= max_res_tokens)
-                rewards = [e.reward for e in episodes]
-                rewards_std = torch.std(torch.tensor(rewards))
-                max_response_len = max(
-                    e.completion.token_ids.shape[0] for e in episodes
-                )
-                drop = rewards_std < 1e-3 or max_response_len >= max_res_tokens
-                record_metric(
-                    "main/continuous_rollouts/dropped_episodes",
-                    1 if drop else 0,
-                    Reduce.SUM,
-                )
-                if drop:
-                    del input_ids, episodes
-                    continue
+            # drop episodes if
+            # 1> reward std-dev is very small (including all 0s and all 1s)
+            # 2> response is potentially truncated (response_len >= max_res_tokens)
+            rewards = [e.reward for e in episodes]
+            rewards_std = torch.std(torch.tensor(rewards))
+            max_response_len = max(e.completion.token_ids.shape[0] for e in episodes)
+            drop = rewards_std < 1e-3 or max_response_len >= max_res_tokens
+            record_metric(
+                "main/continuous_rollouts/dropped_episodes",
+                1 if drop else 0,
+                Reduce.SUM,
+            )
+            if drop:
+                del input_ids, episodes
+                continue
 
             t.step("reward_evaluation")
 

src/forge/controller/launcher.py

@@ -248,8 +248,8 @@ async def launch_mast_job(self):
             scheduler_args={
                 "hpcIdentity": "hyper_monarch",
                 "hpcJobOncall": "monarch",
-                "hpcClusterUuid": "MastProdCluster",
-                "rmAttribution": "pytorch4all_clients_approved",
+                "hpcClusterUuid": "MastGenAICluster",
+                "rmAttribution": "msl_infra_hw_enab_agentrl",
             },
             appdef=self.build_appdef(),
             workspace=Workspace(

src/forge/controller/service/service.py

@@ -285,8 +285,8 @@ async def _migrate_remaining_requests(self, failed_replica: Replica):
             return
 
         # Distribute requests among healthy replicas
-        for i, request in enumerate(migrated_requests):
-            target_replica = healthy_replicas[i % len(healthy_replicas)]
+        for request in migrated_requests:
+            target_replica = self._default_router.get_replica(healthy_replicas)
             await target_replica.enqueue_request(request)
 
             # Update session mapping if needed