feat: Log conversations in train and val in a wandb table. #338
Conversation
alexandery-nvidia
changed the title
Signed-off-by: alexandery <[email protected]>
There was a problem hiding this comment.
Thanks for the contribution. I left some high level comments on design. I think to do this properly it should also be done for TB as well so we don't avoid changing the APIs when that comes (should hopefully also not be too difficult to add). TB may not be able to log this structured data format, but we should log what we can to TB.
Perhaps you can propose an API and we can discuss?
nemo_rl/algorithms/grpo.py
Outdated
| @@ -571,6 +572,38 @@ def grpo_train( | |||
| metrics[k] = np.mean(v).item() | |||
| metrics.update(rollout_metrics) | |||
|
|
|||
| # Log conversations to W&B Table | |||
| if master_config["logger"]["wandb_enabled"]: | |||
There was a problem hiding this comment.
It's better to separate the responsibilities of the algorithm and the logger. There was already this PR to add to log the batched data dict to a file, which if anything should be worked so we just do:
logger.log_batched_data_dict(to_log=batched_data_dict)
with a optional tee_file= or something similar to optionally write to disk if desired
This way it's up to the logger to decide how they want to log (so we can support TB too -- in some capacity)
There was a problem hiding this comment.
I've implemented logging the conversations (in a table to wandb if there's a wandb_logger, and also optionally to jsonl) within the Logger class.
There was a problem hiding this comment.
validate returns a BatchedDataDict consisting of the stacked validation dataloader batches so that we can call the same logger.log_conversations both during training as well as for validation. Let me know what you think.
nemo_rl/algorithms/grpo.py
Outdated
| formatted_conversation = "" | ||
| for msg in conversation: | ||
| # The observations from the environment should have role "environment" | ||
| formatted_conversation += f"**{msg['role']}**: {msg['content']}\n\n" # prompts + responses + environment observations |
There was a problem hiding this comment.
This seems like we're re-implementing chat-template. I think we should try to avoid this if possible
Signed-off-by: alexandery <[email protected]>
Signed-off-by: alexandery <[email protected]>
| tokenizer=tokenizer, | ||
| step=step + 1, | ||
| prefix="train", | ||
| jsonl_filename=f"train_data_step{step}.jsonl" |
There was a problem hiding this comment.
i think this should be renamed tee_jsonl to convey that it "also tees to a file"
also below this it is conditioned on if TB is enabled, which I'm not sure I follow. I think we should always tee the json if the path is provided
| @@ -571,6 +572,63 @@ def log_hyperparams(self, params: Dict[str, Any]) -> None: | |||
| for logger in self.loggers: | |||
| logger.log_hyperparams(params) | |||
|
|
|||
| def log_conversations( | |||
There was a problem hiding this comment.
you should push any logic of wandb into the WandbLogger, see how log_metrics does it:
def log_metrics(
self,
metrics: Dict[str, Any],
step: int,
prefix: Optional[str] = "",
step_metric: Optional[str] = None,
) -> None:
"""Log metrics to all enabled backends.
Args:
metrics: Dict of metrics to log
step: Global step value
prefix: Optional prefix for metric names
step_metric: Optional name of a field in metrics to use as step instead
of the provided step value (currently only needed for wandb)
"""
for logger in self.loggers:
logger.log_metrics(metrics, step, prefix, step_metric)
the tee_json logic can live here, but the stuff that's specific to a type of logger should be in that logger.
| ) -> None: | ||
| """Logs conversation data to W&B Table and/or a JSONL file.""" | ||
| # Log to W&B Table | ||
| if self.wandb_logger: |
There was a problem hiding this comment.
can we do the same for the TB logger?
There was a problem hiding this comment.
Out of curiosity, how much does this bloat the wandb and tb binary assets with respect to disk space? If it makes the assets much larger, we should have a flag to enable this. If someone has a large dataset, we'll probably be saving a few times more data IIUC.
What does this PR do ?
Add train and val conversations to a table and log to wandb if enabled.
Issues
List issues that this PR closes (syntax):
Resolves: #339
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information