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

Code before:
            f"Consumer{self.rank} num_update: {self.num_update_per_episode}, num_recv: {self.num_recv_per_update}, nmb: {self.num_microbatches}"
        )
        for episode in range(self.num_episodes):
            with tqdm(
                range(self.num_update_per_episode),
                desc=f"Episode {episode} with rollout step(s)",
                disable=self.rank != 0,
            ) as pbar:
                for step in pbar:
                    i = 0
                    allow_sync_model = False
                    for _ in range(self.num_recv_per_update):
                        # receive data from producers
                        for r in range(self.num_producers):
                            print(f"[T{dist.get_rank()}] Recv data episode {episode} step {step} from {r}")
                            self.buffer.extend(
                                unbind_batch(
                                    ray_broadcast_tensor_dict(
                                        None, src=0, device=self.device, group_name=f"sync_data_{r}"
                                    )
                                )
                            )
                        while len(self.buffer) >= self.dp_size * self.minibatch_size:
                            batches = self.buffer[
                                self.dp_rank * self.minibatch_size : (self.dp_rank + 1) * self.minibatch_size
                            ]
                            batch = bind_batch(batches)
                            batch = post_recv(batch)
                            loss, excessive_prompts_idx = self.step(i, pbar, **batch)

                            if excessive_prompts_idx is not None:
                                excessive_prompts = [
                                    self.buffer[idx + self.dp_rank * self.minibatch_size]
                                    for idx in excessive_prompts_idx
                                ]
                                self.buffer = excessive_prompts + self.buffer[self.dp_size * self.minibatch_size :]
                            else:
                                self.buffer = self.buffer[self.dp_size * self.minibatch_size :]
                            if loss is not None:
                                allow_sync_model = True
                                pbar.set_postfix({"loss": loss})
                            i += 1
                    if self.lr_scheduler is not None:
                        self.lr_scheduler.step()
                    if (step + 1) % self.save_interval == 0 or (step + 1) == self.num_update_per_episode:
                        if self.rank == 0:
                            print(f"Start saving policy model at step {step + 1}.")
                        save_path = os.path.join(self.save_dir, f"modeling-episode-{episode}-step-{step + 1}")
                        self.booster.save_model(self.policy_model, save_path, shard=True)
                        if self.rank == 0:
                            print(f"Saved model checkpoint at step {step + 1} in folder {save_path}")

Review comment (performance):
pad batch here is necessary when `imbs` is smaller than `tMbs`, samples may have differ in size, need to pad before stacking

Code after:
            f"Consumer{self.rank} num_update: {self.num_update_per_episode}, num_recv: {self.num_recv_per_update}, nmb: {self.num_microbatches}"
        )
        for episode in range(self.num_episodes):
            with tqdm(
                range(self.num_update_per_episode),
                desc=f"Episode {episode} with rollout step(s)",
                disable=self.rank != 0,
            ) as pbar:
                for step in pbar:
                    i = 0
                    allow_sync_model = False
                    for _ in range(self.num_recv_per_update):
                        # receive data from producers
                        for r in range(self.num_producers):
                            print(f"[T{dist.get_rank()}] Recv data episode {episode} step {step} from {r}")
                            self.buffer.extend(
                                unbind_batch(
                                    ray_broadcast_tensor_dict(
                                        None, src=0, device=self.device, group_name=f"sync_data_{r}"
                                    )
                                )
                            )
                        while len(self.buffer) >= self.dp_size * self.minibatch_size:
                            batches = self.buffer[
                                self.dp_rank * self.minibatch_size : (self.dp_rank + 1) * self.minibatch_size
                            ]
                            batch = bind_batch(batches)
                            batch = post_recv(batch)
                            loss, excessive_prompts_idx = self.step(i, pbar, **batch)

                            if excessive_prompts_idx is not None:
                                excessive_prompts = [self.buffer[idx] for idx in excessive_prompts_idx]
                                self.buffer = excessive_prompts + self.buffer[self.dp_size * self.minibatch_size :]
                            else:
                                self.buffer = self.buffer[self.dp_size * self.minibatch_size :]
                            if loss is not None:
                                allow_sync_model = True
                                pbar.set_postfix({"loss": loss})
                            i += 1
                    if self.lr_scheduler is not None:
                        self.lr_scheduler.step()
                    if (step + 1) % self.save_interval == 0 or (step + 1) == self.num_update_per_episode:
                        if self.rank == 0:
                            print(f"Start saving policy model at step {step + 1}.")
                        save_path = os.path.join(self.save_dir, f"modeling-episode-{episode}-step-{step + 1}")
                        self.booster.save_model(self.policy_model, save_path, shard=True)
                        if self.rank == 0:
                            print(f"Saved model checkpoint at step {step + 1} in folder {save_path}")

                    if episode != self.num_episodes - 1 or step != self.num_update_per_episode - 1:
                        if allow_sync_model: