Calculates tokens per second for actors. (#1034)

* Added actor tokens_per_second

* Fixed bug where one timing was using time.time() and the other was using time.perf_counter().

* Updated code

* Cleaned up PR.

Author: finbarrtimbers

open_instruct/grpo_fast.py

@@ -1828,6 +1828,9 @@ def data_preparation_thread(
                 **reward_metrics,
             }
 
+            total_tokens = result.token_statistics.num_prompt_tokens + result.token_statistics.num_response_tokens
+            metrics["val/actor_tokens_per_second"] = total_tokens / result.token_statistics.generation_time
+
         if args.save_traces:
             traces = {
                 "scores": scores.tolist(),
@@ -2287,8 +2290,8 @@ def one_training_step(
         "val/num_total_tokens": num_total_tokens,
         "val/num_step_tokens": num_step_tokens,
         "epoch": episode / args.num_samples_per_prompt_rollout / len(train_dataset),
-        "tokens_per_second_overall": num_total_tokens / total_training_time if total_training_time > 0 else 0,
-        "tokens_per_second_step": num_step_tokens / step_time if step_time > 0 else 0,
+        "learner_tokens_per_second_overall": num_total_tokens / total_training_time,
+        "learner_tokens_per_second_step": num_step_tokens / step_time,
         "time/total": step_time,
         "time/training": train_timer.duration,
         "time/saving": save_time,
@@ -2374,6 +2377,12 @@ def maybe_evaluate(
         }
         if "time/generation" in eval_generate_metrics:
             eval_metrics["eval/generation_time"] = eval_generate_metrics["time/generation"]
+
+        total_tokens = (
+            eval_result.token_statistics.num_prompt_tokens + eval_result.token_statistics.num_response_tokens
+        )
+        eval_metrics["eval/actor_tokens_per_second"] = total_tokens / eval_result.token_statistics.generation_time
+
         print_rich_single_line_metrics(eval_metrics)
 
         table = {}

open_instruct/vllm_utils3.py

@@ -497,7 +497,7 @@ def process_from_queue(self, timeout: float = 60.0):
                 self._prefetch_future.result()
 
             self._poll_tool_futures(self.tracking, self.llm_engine.tokenizer)
-            current_time = time.time()
+            current_time = time.perf_counter()
             if self.llm_engine.has_unfinished_requests():
                 for output in [o for o in self.llm_engine.step() if o.finished]:
                     # Fix the index field for all sub-requests
@@ -824,7 +824,7 @@ def _poll_tool_futures(self, tracking, tokenizer):
                 tracking["pending_tool_futures"].pop(req_id, None)
 
                 complete_output = tracking["concat_outputs"][req_id].outputs[0]
-                current_time = time.time()
+                current_time = time.perf_counter()
                 self._finalize_sub_request(req_id, last_output, complete_output, current_time)
                 # Don't add to dict_keys_to_delete since we already removed it
                 continue