Fix first_micro_batch_rollout_time

jiangyangmu · The tunix Authors · commit e4bcbfe5588a · 2025-11-20T17:53:32.000-08:00
PiperOrigin-RevId: 834982334
diff --git a/tests/perf/export_test.py b/tests/perf/export_test.py
@@ -161,7 +161,7 @@ def test_export_grpo_metrics_rollout_1_actor_2_reference_2(
         "perf/global_step_time": 1.3,
         "perf/weight_sync_time": 0.1,
         "perf/rollout_idle_time": 0.69,
-        "perf/first_micro_batch_rollout_time": 0.3,
+        "perf/first_micro_batch_rollout_time": 0.41,
         "perf/sum/rollout_time": 0.11,
         "perf/sum/refer_inference_time": 0.11,
         "perf/sum/actor_train_grad_time": 0.11,
diff --git a/tunix/perf/export.py b/tunix/perf/export.py
@@ -202,7 +202,7 @@ def _grpo_metrics_rollout_1_actor_2_reference_2(
     ]
 
     first_micro_batch_rollout_time: float = (
-        rollout_spans[0].begin - global_step_group.begin
+        rollout_spans[0].end - global_step_group.begin
     )
 
     # append [0.0] to make size equal to micro batch

Original file line number	Diff line number	Diff line change
`@@ -202,7 +202,7 @@ def _grpo_metrics_rollout_1_actor_2_reference_2(`
`202`	`202`	`]`
`203`	`203`
`204`	`204`	`first_micro_batch_rollout_time: float = (`
`205`		`- rollout_spans[0].begin - global_step_group.begin`
	`205`	`+ rollout_spans[0].end - global_step_group.begin`
`206`	`206`	`)`
`207`	`207`
`208`	`208`	`# append [0.0] to make size equal to micro batch`