From 607b491935c73d50bfeaeec2fa250fe0c9c4b419 Mon Sep 17 00:00:00 2001 From: sanmuf Date: Tue, 4 Aug 2026 11:38:36 +0800 Subject: [PATCH] fix: is_offload_optimizer_states_in_train_step not passed correctly in RLVR and agentic pipelines --- roll/pipeline/agentic/agentic_pipeline.py | 1 + roll/pipeline/rlvr/rlvr_pipeline.py | 1 + 2 files changed, 2 insertions(+) diff --git a/roll/pipeline/agentic/agentic_pipeline.py b/roll/pipeline/agentic/agentic_pipeline.py index e4652a7c4..899da1979 100644 --- a/roll/pipeline/agentic/agentic_pipeline.py +++ b/roll/pipeline/agentic/agentic_pipeline.py @@ -314,6 +314,7 @@ def run(self): batch.meta_info["global_step"] = global_step batch.meta_info["_broadcast_non_tensor_batch"] = True batch.meta_info["loss_mask_keys"] = ["response_mask"] + batch.meta_info["is_offload_optimizer_states_in_train_step"] = self.pipeline_config.is_offload_optimizer_states_in_train_step if val_future is not None: val_metrics = val_future.result() diff --git a/roll/pipeline/rlvr/rlvr_pipeline.py b/roll/pipeline/rlvr/rlvr_pipeline.py index 1f3dca2a5..5e4425b23 100644 --- a/roll/pipeline/rlvr/rlvr_pipeline.py +++ b/roll/pipeline/rlvr/rlvr_pipeline.py @@ -570,6 +570,7 @@ def run(self): batch.meta_info["global_step"] = global_step batch.meta_info["_broadcast_non_tensor_batch"] = True batch.meta_info["loss_mask_keys"] = ['response_mask', 'final_response_mask'] + batch.meta_info["is_offload_optimizer_states_in_train_step"] = self.pipeline_config.is_offload_optimizer_states_in_train_step batch.non_tensor_batch['sample_uuid'] = np.array([str(uuid.uuid4()) for _ in range(batch.batch.shape[0])], dtype=object) batch.batch["prompt_id"] = torch.arange(batch.batch.batch_size[0], device=batch.batch.device)