How to save val_data predictions during each evaluation of `trainer.fit`?

[dalek-who]

During the trainer.fit process, how can I access and save the prediction results (for example, the final State returned) of the val_data for each evaluation step?

[ultmaster]

I see there are a few approaches.

1. Save to a dedicated file folder manually in agent rollout.
2. Use emit_object, which will save the state as a dedicated span in the store. And you can use store.query_spans APIs later to fetch that saved object.
3. You can also automate (2) by adding logics after each validation finishes to auto query all the spans from all rollouts; assemble them and save to a dedicated place.

[dalek-who]

You can also automate (2) by adding logics after each validation finishes to auto query all the spans from all rollouts; assemble them and save to a dedicated place.

Where should I add this logic? Do I need to modify the agent-lightning source code?

[ultmaster]

If you are doing RL training, current verl implementation are designed for using as templates. You can copy them out and modify them anyway you want.

[dalek-who]

Or do we have any hooks like on_validation_end?

[ultmaster]

No. The VERL algorithms are designed for copying out for arbitrary changes.

on_validation_end

But I can note this for future refactor of VERL.

[beanie00]

In my case, I modified the _validate function in agentlightning/verl/trainer to save the test rollout data.

def _dump_generations(self, inputs, outputs, scores, data_sources, dump_path):
        """Dump rollout/validation samples as JSONL."""
        os.makedirs(dump_path, exist_ok=True)
        filename = os.path.join(dump_path, f"{self.global_steps}.jsonl")

        n = len(inputs)
        base_data = {
            "input": inputs,
            "output": outputs,
            "score": scores,
            "data_sources": data_sources,
            "step": [self.global_steps] * n,
        }

        with open(filename, "w") as f:
            for i in range(n):
                entry = {k: v[i] for k, v in base_data.items()}
                f.write(json.dumps(entry, ensure_ascii=False) + "\n")

        print(f"Dumped generations to {filename}")

def _validate(self):
        # Lists to collect samples
        sample_scores = []

        test_data = next(iter(self.val_dataloader))
        test_batch = DataProto.from_single_dict(test_data)

        self.async_rollout_manager.wake_up()
        self.agent_mode_daemon.set_up_data_and_server(
            test_batch.non_tensor_batch,
            self.async_rollout_manager.server_addresses,
            is_train=False,
        )
        self.agent_mode_daemon.run_until_all_finished()
        test_metrics, test_input_ids, test_responses, sample_scores, data_sources = self.agent_mode_daemon.get_test_data_batch_and_metrics()
        self.agent_mode_daemon.clear_data_and_server()
        self.async_rollout_manager.sleep()

        # dump generations
        val_data_dir = self.config.trainer.get("validation_data_dir", None)

        if val_data_dir:
            # Store original inputs
            sample_inputs = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in test_input_ids]

            # Store generated outputs
            sample_outputs = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in test_responses]

            self._dump_generations(
                inputs=sample_inputs,
                outputs=sample_outputs,
                scores=sample_scores,
                data_sources=data_sources,
                dump_path=val_data_dir,
            )

        return test_metrics

and also modified the get_test_metrics function in agentlightning/verl/daemon.py to get_test_data_batch_and_metrics, and additionally added tracking for input_ids_list, response_ids_list, score_list and data_source_list.

def get_test_data_batch_and_metrics(self):
        """Calculates and returns metrics for a validation run."""
        assert not self.is_train, "This method should only be called during validation."
        assert len(self._completed_rollouts_v0) == self._total_tasks_queued

        sample_stat_list: List[Dict[str, Any]] = []
        sample_stat_list_by_source: Dict[str, List[Dict[str, Any]]] = defaultdict(
            list
        )  # FIXME: Evaluate whether grouping stats by source is actually needed.
        data_source_list: List[str] = []
        score_list: List[float] = []

        input_ids_list: List[List[int]] = []
        response_ids_list: List[List[int]] = []

        for rollout_id, rollout in self._completed_rollouts_v0.items():
            final_reward_raw: Optional[float] = rollout.final_reward
            final_reward = self._fillna_reward(rollout)
            if not rollout.triplets:
                print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
                sample_stat_list.append({"reward": final_reward, "has_reward": final_reward_raw is not None})
                continue

            input_ids_list.extend([triplet.prompt.get("token_ids", []) for triplet in rollout.triplets])
            response_ids_list.extend([triplet.response.get("token_ids", []) for triplet in rollout.triplets])
            response_length_list = [len(triplet.response.get("token_ids", [])) for triplet in rollout.triplets]
            score_list.append(final_reward)

            if "data_source" in self._task_id_to_original_sample[rollout_id]:
                # When a test sample includes a 'data_source' field, record per-source statistics for test results.
                # TODO: This is a flawed design. We should have a better way to handle this.
                data_source = self._task_id_to_original_sample[rollout_id]["data_source"]
                sample_stat_list_by_source[data_source].append(
                    {
                        "sum_response_length": np.sum(response_length_list),
                        "mean_response_length": np.mean(response_length_list) if response_length_list else 0,
                        "turn_count": len(rollout.triplets),
                        "reward": final_reward,
                        "has_reward": final_reward_raw is not None,
                    }
                )
                data_source_list.append(data_source)
            else:
                data_source_list.append("")
...

[dalek-who]

Save to a dedicated file folder manually in agent rollout.

Within agl.LitAgent.rollout, is it possible to access epoch and global_step during trainer.fit?

[ultmaster]

Not natively. Those variables are controlled by the algorithm and only algorithm knows what they mean.

You can however add them to the metadata of rollout and access them via the metadata.

[dalek-who]

You can however add them to the metadata of rollout and access them via the metadata.

How to add them to the metadata?

[ultmaster]

        for i in range(num_samples):
            data_id = str(uuid.uuid4())
            original_sample = {key: data[key][i] for key in keys}
            original_sample["data_id"] = data_id
            data_id_to_original_sample[data_id] = original_sample

            # For training, each sample is rolled out multiple times
            # Data ID is different from Rollout ID, as one data can have multiple rollouts.
            for _ in range(rollouts_per_sample):
                # !!! here !!!
                task_metadata = {"data_id": data_id, "is_train": is_train, "epoch": some_epoch, "batch_index": some_batch_index}
                if self.mode == "v0":
                    # Queue immediately
                    rollout_id = await self.server.queue_task(
                        sample=_to_native(original_sample),
                        mode="train" if is_train else "val",
                        resources_id=resources_id,
                        metadata=task_metadata,
                    )

                    # Store original sample data to reconstruct batch information later
                    self._task_id_to_original_sample[rollout_id] = original_sample
                    self._total_tasks_queued += 1
                else:
                    # Collect tasks to enqueue in batch and queue them later
                    enqueue_rollout_requests.append(
                        EnqueueRolloutRequest(
                            input=_to_native(original_sample),
                            mode="train" if is_train else "val",
                            resources_id=resources_id,
                            config=RolloutConfig(
                                unresponsive_seconds=self.llm_timeout_seconds,
                                timeout_seconds=self.llm_timeout_seconds,
                            ),
                            metadata=task_metadata,
                        )
                    )

        if self.mode == "v1":
            # Enqueue all the tasks in a single batch
            rollouts = await self.store.enqueue_many_rollouts(enqueue_rollout_requests)
            self._task_id_to_original_sample.update(
                {
                    # Recover the original data and store it for later use.
                    rollout.rollout_id: data_id_to_original_sample[cast(Dict[str, Any], rollout.metadata)["data_id"]]
                    for rollout in rollouts
                }
            )
            self._total_tasks_queued += len(rollouts)

[dalek-who]

@ultmaster Thanks for the explanation. I'm considering where it would be most appropriate to modify the code (the source code of agentlightning or the source code of verl).

By the way, if I want to use a modified version of verl instead of the pip-installed verl, what is the most recommended approach:

1. Modify the verl source code in the Python library directory.
2. Instead of using pip, directly clone the source code of both agentlightning and verl into my project, and make the necessary modifications (in this case, aside from modifying verl, should I also change all instances of "import verl" in the agentlightning source code to import my locally modified verl?).
3. Any other best practices?

[ultmaster]

By the way, if I want to use a modified version of verl instead of the pip-installed verl, what is the most recommended approach:

If you are talking about verl package (https://github.com/volcengine/verl), please ask the verl community. We have no recommendations.

If you are talking about agentlightning/verl, I recommend copying out the folder.

If you are asking about my personal opinion, you can install both agent-lightning and verl from source code.