trending/miles
1
0
Fork 0
mirror of https://github.com/radixark/miles.git synced 2026-10-02 07:14:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

[profile] add perf stats on rollout

Browse Source
This commit is contained in:
Zilin Zhu
2025-08-28 04:48:22 -07:00
parent 0070e590fe
commit fd31dc7563
1 changed files with 41 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
slime/ray/buffer.py
+41 -19
View File
@@ -1,6 +1,7 @@
import logging
from pathlib import Path
from typing import Union
from time import time
import wandb
import ray
@@ -16,25 +17,6 @@ logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)
def log_eval_data(rollout_id, args, data):
log_dict = {}
for key in data.keys():
rewards = data[key]["rewards"]
log_dict[f"eval/{key}"] = sum(rewards) / len(rewards)
if "truncated" in data[key]:
truncated = data[key]["truncated"]
log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated)
print(f"eval {rollout_id}: {log_dict}")
if args.use_wandb:
log_dict["eval/step"] = (
rollout_id
if not args.wandb_always_use_train_step
else rollout_id * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
)
wandb.log(log_dict)
@ray.remote
class RolloutController:
"""The class to run rollout and convert rollout data to training data."""
@@ -60,6 +42,7 @@ class RolloutController:
def generate(self, rollout_id):
self.rollout_id = rollout_id
start_time = time()
if self.args.load_debug_rollout_data:
data = torch.load(
open(self.args.load_debug_rollout_data.format(rollout_id=rollout_id), "rb"),
@@ -85,6 +68,7 @@ class RolloutController:
path,
)
data = self._convert_samples_to_train_data(data)
log_rollout_data(rollout_id, self.args, data, time() - start_time)
return Box(ray.put(data))
def eval(self, rollout_id):
@@ -174,3 +158,41 @@ class RolloutController:
def load(self, rollout_id=None):
self.data_source.load(rollout_id)
def log_eval_data(rollout_id, args, data):
log_dict = {}
for key in data.keys():
rewards = data[key]["rewards"]
log_dict[f"eval/{key}"] = sum(rewards) / len(rewards)
if "truncated" in data[key]:
truncated = data[key]["truncated"]
log_dict[f"eval/{key}-truncated_ratio"] = sum(truncated) / len(truncated)
print(f"eval {rollout_id}: {log_dict}")
if args.use_wandb:
log_dict["eval/step"] = (
rollout_id
if not args.wandb_always_use_train_step
else rollout_id * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
)
wandb.log(log_dict)
def log_rollout_data(rollout_id, args, data, rollout_time):
if args.load_debug_rollout_data:
return
log_dict = {}
response_lengths = [sum(loss_mask) for loss_mask in data["loss_masks"]]
log_dict["perf/rollout_time"] = rollout_time
log_dict["perf/tokens_per_gpu_per_sec"] = sum(response_lengths) / rollout_time / args.rollout_num_gpus
log_dict["perf/longest_sample_tokens_per_sec"] = max(response_lengths) / rollout_time
print(f"perf {rollout_id}: {log_dict}")
if args.use_wandb:
log_dict["perf/step"] = (
rollout_id
if not args.wandb_always_use_train_step
else rollout_id * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
)
wandb.log(log_dict)