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

Code before:
        self.policy_model, self.optimizer, _, _, self.lr_scheduler = self.booster.boost(
            self.policy_model, self.optimizer, lr_scheduler=self.lr_scheduler
        )
        self.reference_model, *_ = self.booster.boost(self.reference_model)

    def step(self, step_idx: int, **kwargs) -> Optional[float]:
        """
        Step data from policy model:
            [{
                "input_ids": torch.Tensor,
                "attention_mask": torch.Tensor,
                "action_mask": torch.Tensor,
                "action_log_probs": torch.Tensor,
            },
            ...]
        Format:
            [batch_size, num_of_generation, prompt_length + response_length] --- <PAD>...<PAD><PROMPT>...<PROMPT><RESPONSE>...<RESPONSE><PAD>...<PAD>.
        """

        # Reshape to [batch_size x num_of_generation, prompt_length + response_length]
        data = {k: v.view(-1, v.size(-1)) for k, v in kwargs.items()}
        action_mask = data["action_mask"]
        num_action = action_mask.shape[1]
        old_action_log_probs = data["action_log_probs"]
        response_length = torch.sum(action_mask, dim=1).to(torch.float32)
        forward_batch_size = self.training_config.get("forward_micro_batch_size", data["input_ids"].size(0))

        need_update = (step_idx + 1) % self.num_microbatches == 0
        # Gradient must be synchronized if zero2 is enabled. https://github.com/hpcaitech/ColossalAI/blob/44d4053fec005fe0b06b6bc755fdc962463145df/colossalai/booster/plugin/hybrid_parallel_plugin.py#L1500
        ctx = (
            nullcontext()
            if need_update or self.booster.plugin.zero_stage == 2
            else self.booster.no_sync(self.policy_model, self.optimizer)
        )
        with ctx:
            reward_group = self.reward_model(
                data["input_ids"], gt_answer=data["gt_answer"], response_idx=data["response_idx"]
            )

            reward = torch.tensor([value[0] for value in reward_group]).to(data["input_ids"].device)
            format_reward = torch.tensor([value[1] for value in reward_group]).to(data["input_ids"].device)
            acc_reward = torch.tensor([value[2] for value in reward_group]).to(data["input_ids"].device)

            # [batch_size, num_generations]

            group_reward = reward.view(-1, self.num_generations)
            reward_mean = group_reward.mean(dim=1)
            # [batch_size x num_generations]
            reward_mean = reward_mean.repeat_interleave(self.num_generations, dim=0)
            reward_std = group_reward.std(dim=1).repeat_interleave(self.num_generations, dim=0)
            # [batch_size x num_generations]

Review comment (performance):
It is better to add assert to check that the `forward_batch_size` should be divided by microbatchsize * num_generation 

Code after:
        self.policy_model, self.optimizer, _, _, self.lr_scheduler = self.booster.boost(
            self.policy_model, self.optimizer, lr_scheduler=self.lr_scheduler
        )
        self.reference_model, *_ = self.booster.boost(self.reference_model)

    def step(self, step_idx: int, **kwargs) -> Optional[float]:
        """
        Step data from policy model:
            [{
                "input_ids": torch.Tensor,
                "attention_mask": torch.Tensor,
                "action_mask": torch.Tensor,
                "action_log_probs": torch.Tensor,
            },
            ...]
        Format:
            [batch_size, num_of_generation, prompt_length + response_length] --- <PAD>...<PAD><PROMPT>...<PROMPT><RESPONSE>...<RESPONSE><PAD>...<PAD>.
        """

        # Reshape to [batch_size x num_of_generation, prompt_length + response_length]
        data = {k: v.view(-1, v.size(-1)) for k, v in kwargs.items()}
        action_mask = data["action_mask"]
        num_action = action_mask.shape[1]
        old_action_log_probs = data["action_log_probs"]
        response_length = torch.sum(action_mask, dim=1).to(torch.float32)
        forward_batch_size = self.training_config.get("train_microbatch_size", data["input_ids"].size(0))

        need_update = (step_idx + 1) % self.num_microbatches == 0
        # Gradient must be synchronized if zero2 is enabled. https://github.com/hpcaitech/ColossalAI/blob/44d4053fec005fe0b06b6bc755fdc962463145df/colossalai/booster/plugin/hybrid_parallel_plugin.py#L1500
        ctx = (
            nullcontext()
            if need_update or self.booster.plugin.zero_stage == 2
            else self.booster.no_sync(self.policy_model, self.optimizer)
        )
        with ctx:
            reward_group = self.reward_model(
                data["input_ids"], gt_answer=data["gt_answer"], response_idx=data["response_idx"]
            )

            reward = torch.tensor([value[0] for value in reward_group]).to(data["input_ids"].device)
            format_reward = torch.tensor([value[1] for value in reward_group]).to(data["input_ids"].device)
            acc_reward = torch.tensor([value[2] for value in reward_group]).to(data["input_ids"].device)

            # [batch_size, num_generations]

            group_reward = reward.view(-1, self.num_generations)
            reward_mean = group_reward.mean(dim=1)
            # [batch_size x num_generations]
            reward_mean = reward_mean.repeat_interleave(self.num_generations, dim=0)
            reward_std = group_reward.std(dim=1).repeat_interleave(self.num_generations, dim=0)
            # [batch_size x num_generations]