Repository: hpcaitech/ColossalAI
File: applications/ColossalChat/coati/distributed/grpo_consumer.py (Python)

Code before:
                        iter([data_policy_forward]),
                        self.policy_model,
                        criterion=_criterion,
                        optimizer=self.optimizer,
                        return_loss=True,
                        return_outputs=True,
                    )
                    loss = policy_model_outputs["loss"]

                    if self.booster.plugin.stage_manager.is_last_stage():
                        # calculate kl, as we cannot do this inside callback, kl needs be calculate again
                        action_logits = policy_model_outputs["outputs"]["logits"]
                        action_log_probs = calc_action_log_probs(
                            action_logits / self.generate_config["temperature"],
                            input_ids_forward_micro_batch,
                            num_action,
                            self.plugin.shard_config,
                        )
                        per_token_kl = (
                            torch.exp(reference_action_log_probs - action_log_probs)
                            - (reference_action_log_probs - action_log_probs)
                            - 1
                        )
                        kl = torch.sum(per_token_kl * action_mask_forward_micro_batch, dim=-1) / torch.sum(
                            action_mask_forward_micro_batch, dim=-1
                        )
                        kl = all_reduce_mean(kl.mean(), self.plugin)
                        loss = all_reduce_mean(loss, self.plugin)
                        mean_loss.append(loss.data)
                        mean_kl.append(kl)
                else:

                    policy_model_logits = self.policy_model(
                        input_ids=input_ids_forward_micro_batch,
                        attention_mask=attention_mask_forward_micro_batch,
                    ).logits
                    action_log_probs = calc_action_log_probs(
                        policy_model_logits / self.generate_config["temperature"],
                        input_ids_forward_micro_batch,
                        num_action,
                        self.plugin.shard_config,
                    )

                    with torch.no_grad():
                        reference_model_logits = self.reference_model(
                            input_ids=input_ids_forward_micro_batch,
                            attention_mask=attention_mask_forward_micro_batch,
                        ).logits
                    reference_action_log_probs = calc_action_log_probs(
                        reference_model_logits / self.generate_config["temperature"],
                        input_ids_forward_micro_batch,

Review comment (suggestion):
You do not need to recalculate these parameters again. 

Please refer here to use Closure: https://github.com/hpcaitech/ColossalAI/blob/44d4053fec005fe0b06b6bc755fdc962463145df/applications/ColossalChat/coati/trainer/dpo.py#L191-L214

Code after:
                            inputs["advantages"].repeat_interleave(action_log_probs.size(-1), dim=-1),
                            per_token_kl,
                            inputs["action_mask"],
                            loss_mask=inputs["loss_mask"],
                        )
                        return loss

                    policy_model_outputs = self.booster.execute_pipeline(
                        iter([data_policy_forward]),
                        self.policy_model,
                        criterion=_criterion,
                        optimizer=self.optimizer,
                        return_loss=True,
                        return_outputs=True,
                    )
                    loss = policy_model_outputs["loss"]

                    if self.booster.plugin.stage_manager.is_last_stage():
                        if len(kl) > 0:
                            kl = all_reduce_mean(torch.mean(torch.stack(kl)).to(loss.device), self.plugin)
                            mean_kl.append(kl)
                        loss = all_reduce_mean(loss, self.plugin)
                        mean_loss.append(loss.data)
                else:

                    policy_model_logits = self.policy_model(
                        input_ids=input_ids_forward_micro_batch,
                        attention_mask=attention_mask_forward_micro_batch,
                    ).logits
                    action_log_probs = calc_action_log_probs(
                        policy_model_logits / self.generate_config["temperature"],
                        input_ids_forward_micro_batch,
                        num_action,
                        self.plugin.shard_config,
                    )

                    with torch.no_grad():
                        reference_model_logits = self.reference_model(
                            input_ids=input_ids_forward_micro_batch,
                            attention_mask=attention_mask_forward_micro_batch,
                        ).logits
                    reference_action_log_probs = calc_action_log_probs(
                        reference_model_logits / self.generate_config["temperature"],
                        input_ids_forward_micro_batch,
                        num_action,
                        self.plugin.shard_config,
                    )
                    per_token_kl = (
                        torch.exp(reference_action_log_probs - action_log_probs)
                        - (reference_action_log_probs - action_log_probs)
                        - 1