Using tensorboard instead of wandb

[pratikkumar018]

Hi Team,
I wanted to use tensorboard instead of wandb. Is there any flag or way i can quickly make this switch. I am running on offline servers.

[wuxibin89]

@pratikkumar018 Hi, we don't support tensorboard for now. Can you help to contribute? it's super easy to replace wandb with tensorboard:
https://github.com/OpenLLMAI/OpenRLHF/blob/main/openrlhf/trainer/ppo_trainer.py#L362-L370

[pratikkumar018]

Sure will check. Thanks

[pratikkumar018]

Hi Here is the code I have written for tensorboard

https://github.com/OpenLLMAI/OpenRLHF/blob/main/openrlhf/trainer/ppo_trainer.py#L147

if self.strategy.args.use_tensorboard and self.strategy.is_rank_0():
              from torch.utils.tensorboard import SummaryWriter
              current_time = datetime.now().strftime("%Y%m%d-%H%M%S")
              self._writer=SummaryWriter('{}/{}'.format(log_dir, current_time))

OpenRLHF/openrlhf/trainer/ppo_trainer.py

Line 371 in 4e15591

                logs = {
                    "train/%s" % k: v
                    for k, v in {
                        **logs_dict,
                    }.items()
                }
                                
                for k, v in logs.items():
                    print(k,v)
                    print(global_step)
                    self._writer.add_scalar(k, v, global_step=global_step)
                self._writer.add_scalar('global_step', global_step, global_step=global_step)

OpenRLHF/openrlhf/trainer/ppo_trainer.py

Line 195 in 4e15591

if self.strategy.args.use_tensorboard and self.strategy.is_rank_0():
            self._writer.close()```

[pratikkumar018]

@wuxibin89 can you please help me with this for ppo_ray as well. I tried finding couldnt locate the file where it needs to be done. I need to use this for bigger model . Or since it is already inheriting PPO trainer changing it there would be enough?

[wuxibin89]

ppo_ray inherits PPO trainer, so it should be enough.

[pratikkumar018]

Thanks.

[wuxibin89]

@pratikkumar018 Hi, can you contribute your tensorboard support with a PR? We're very appreciate it.