Suggested GPU to run the demo code of step2_reward_model_finetuning (DeepSpeed-Chat)

[wenbozhangjs]

I follow the instructions of this page to do step2_reward_model_finetuning with demo code.

On the Google Cloud platform, I create one instance with a single V100(16GB) and another instance with double V100(16GB). I directly use this command bash training_scripts/opt/single_gpu/run_350m.sh but always meet the out-of-memory issues.

Are there any modifications I could do to run this demo code on double V100(16BG)? Or Are there recommendations about which type of GPU I should use to run this demo code successfully?

Appreciate the help!