Support logging for non-scalar metrics
- Dominant language
- Python
- Stars
- 4.8k
- Forks
- 487
- PR merge metrics
- No merged PRs in 30d
Description
### 🚀 The feature, motivation, and pitch
AccelerateRLTrainer.evaluate() logs a table of generated eval outputs and metrics to the metrics tracker.
If I understand correctly, only scalar metrics are currently supported.
This feature would allow non-scalar metrics to be logged.
#### Use cases:
Allow passthrough logging of non-scalar prompt metadata:
```
def eval_metrics(samples, prompts, outputs, **kwargs):
return kwargs
trlx.train(
model_name,
config=config,
samples=train_samples, # type: ignore
rewards=train_rewards if training_method == "ilql" else None, # type: ignore
eval_prompts=eval_eval_prompts, # type: ignore
metric_fn=eval_metrics,
).model
```
#### Implementation suggestion:
Modify the mean_metrics calculation (below) to only calculate means for values that can be successfully cast to float tensors.
https://github.com/CarperAI/trlx/blob/0dce99d96b7d70b6a9114129d8e38bf6c80eb653/trlx/trainer/accelerate_base_trainer.py#L428-L430
### Alternatives
_No response_
### Additional context
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.