Training Efficiency

[eslambakr]

Dears,

Thanks for sharing your great work!

I have two questions regarding the training details:

1. It is mentioned in the running script "run_grpo.sh" that we have to determine a certain GPU for the reward. However, I don't see any part of the code that does so. Can you please point out the part that splits the Janus model on X GPUs, then the reward models on a separate GPU?
2. The training time: May I know how long it takes to train the model? From the configuration, it seems you have only trained for 1600 steps, isn't it? I am wondering if it is enough to capture new skills using only 1600 steps, while even the batch size is set to 1? I feel like this number is too small; thus, it would be appreciated if you could elaborate more about the training epochs or steps.

Thanks again for sharing your great work!

[CaraJ7]

Hi @eslambakr , thanks for your interest. As for the questions:

1. Sorry for the confusion. I forgot to correct the script when copying from the previous repo. We do not need to determine a certain GPU for the reward when training T2I-R1. The reward models are on the same device as the training model. The script should be correct now.
2. It takes roughly 18 hours on 8 H800. We observe that 1600 steps are enough to make the difference, as the rewards also converge. Note that the batch size is actually 2 given that the gradient accumulation is set to 2.

[eslambakr]

Hi @CaraJ7
Millions of thanks for your prompt response!
Then, may I ask how you managed to make all these models fit in the GPU memory? It seems you are not using gradient checkpointing, which should make the fitting task even harder on 80 GB.
Actually, I was working on a training code for Janus too, and I can not generate 8 generations simultaneously; the maximum I can get is only 4.
Additionally, indeed, I can not fit any other reward model into the same GPU memory.

So, I am really surprised that you managed to do it.
Thus, it would be highly appreciated if you could share your smart implementation tricks that make it happen.
Also, are u using Lora fine-tuning, or are you fine-tuning the entire model? It seems from the code u are fine-tuning the entire model unless I miss the LORA part.

[CaraJ7]

Actually, we have used gradient checkpointing inside the code. Simply setting the gradient checkpoint to True in the bash file will result in errors, due to its incompatibility with Janus-pro. You could check all the gradient_checkpointing_enable and gradient_checkpointing_disable functions in src/t2i-r1/src/open_r1/trainer/grpo_trainer.py. We disable the gradient checkpointing when we are generating texts or images, since the gradient checkpointing could interfere with the generation. And we enable it after generation. Applying gradient checkpointing could largely reduce the memory cost.

We conduct full fine-tuning of the model instead of LoRA.