[BUG]: Fetching wrong data from pretrained dataloader when ptx_coef is not zero in Staging 3 training

[yynil]

🐛 Describe the bug

When the pretrained data loss is used to train the prompts model in stage-3, the pretrained data loader fetch 3 times data from the dataloader to get input_ids, labels and attention_mask.
The following is the current implementation in ppo.py.

        # ptx loss
        if self.ptx_coef != 0:
            ptx = next(iter(self.pretrain_dataloader))['input_ids'].to(torch.cuda.current_device())
            label = next(iter(self.pretrain_dataloader))['labels'].to(torch.cuda.current_device())[:, 1:]
            attention_mask = next(iter(self.pretrain_dataloader))['attention_mask'].to(torch.cuda.current_device())
            ptx_log_probs = self.actor.get_base_model()(ptx, attention_mask=attention_mask)['logits'][..., :-1, :]
            ptx_loss = self.ptx_loss_fn(ptx_log_probs.view(-1, ptx_log_probs.size(-1)), label.view(-1))
            actor_loss = ptx_loss * self.ptx_coef + actor_loss * (1 - self.ptx_coef)

It's an obvious bug here but if you guys are using exact same length of data, that bug will not terminate your training procedure. But the loss is totally wrong since ptx,label and attention_mask are from different data batch.

Environment

Pytorch 1.12+cuda 11.6+CentOS 7

[E4546]

GUIDELINES AND RULES FOR WRITING ASSIGNMETS
• Complete assignments carefully and thoughtfully with attention to details and quality
of work.
Complete assignments on your own without the help of third parties and to the best of
your capability.
Complete and turn in assignments by the specific due date.
Talk to your lecturer about problems with understanding of assignments.
Use other people's copyrighted works without approval or without adequate reference
to the title of the work and the name of the auth is forbidden
Each assignment should be of a minimum of 1000 words and of a maximum of 3000
words.

ASSESSMENT OF THE ASSIGNMENTS
Each assignment holds 100% and it will be assessed based on the following criteria:

1. Qulity of the work: The overall effort you put in writing the assignment includes
   the use of images, draw illustrative schematics, how coherent your writing is and how
   you discuss a point and deliver it in an understandable way will represent the quality
   of your work (25%)

According to previous
please carefully write One thousand five hundred (1500) words with references books by ApA style citation

About the

( Communication between cells is a process coordinated by specific signal transduction
pathways. Based on your understanding write an overview about the process of signal transduction clarifying the definition of signal transduction, describing the steps of the whole
process, mentioning the differences between signal molecules, receptor molecules and second messengers. (Schematic drawing must be included).