ColossalAI/applications/ColossalChat/coati/ray/callbacks/performance_evaluator.py

from time import time
from typing import Optional

import torch
import torch.distributed as dist
from coati.experience_maker import Experience

from .base import MakerCallback, TrainerCallback


def get_world_size() -> int:
    if dist.is_initialized():
        return dist.get_world_size()
    return 1


def print_rank_0(*args, **kwargs) -> None:
    if not dist.is_initialized() or dist.get_rank() == 0:
        print(*args, **kwargs)


@torch.no_grad()
def all_reduce_mean(x: float, world_size: int) -> float:
    if world_size == 1:
        return x
    tensor = torch.tensor([x], device=torch.cuda.current_device())
    dist.all_reduce(tensor)
    tensor = tensor / world_size
    return tensor.item()


class Timer:
    def __init__(self) -> None:
        self.start_time: Optional[float] = None
        self.duration: float = 0.0

    def start(self) -> None:
        self.start_time = time()

    def end(self) -> None:
        self.duration += time() - self.start_time

    def reset(self) -> None:
        self.duration = 0.0


class ExperienceMakerPerformanceEvaluator(MakerCallback):
    def __init__(
        self, actor_num_params: int, critic_num_params: int, initial_model_num_params: int, reward_model_num_params: int
    ) -> None:
        super().__init__()
        self.world_size = get_world_size()
        self.actor_num_params = actor_num_params
        self.critic_num_params = critic_num_params
        self.initial_model_num_params = initial_model_num_params
        self.reward_model_num_params = reward_model_num_params

        self.batch_timer = Timer()
        self.send_timer = Timer()
        self.make_experience_timer = Timer()
        self.total_samples: int = 0
        self.make_experience_flop: int = 0

        print_rank_0(
            f"ExperienceMaker actor: {actor_num_params/1024**3:.2f}B, critic: {critic_num_params/1024**3:.2f}B, initial model: {initial_model_num_params/1024**3:.2f}B, reward model: {reward_model_num_params/1024**3:.2f}B, world size: {self.world_size}"
        )

    def on_make_experience_start(self) -> None:
        self.make_experience_timer.start()

    def on_make_experience_end(self, experience: Experience) -> None:
        self.make_experience_timer.end()

        batch_size, seq_len = experience.sequences.shape

        self.total_samples += batch_size

        # actor generate
        num_actions = experience.action_mask.size(1)
        input_len = seq_len - num_actions
        total_seq_len = (input_len + seq_len - 1) * num_actions / 2
        self.make_experience_flop += self.actor_num_params * batch_size * total_seq_len * 2
        # actor forward
        self.make_experience_flop += self.actor_num_params * batch_size * seq_len * 2
        # critic forward
        self.make_experience_flop += self.critic_num_params * batch_size * seq_len * 2
        # initial model forward
        self.make_experience_flop += self.initial_model_num_params * batch_size * seq_len * 2
        # reward model forward
        self.make_experience_flop += self.reward_model_num_params * batch_size * seq_len * 2

    def on_send_start(self) -> None:
        self.send_timer.start()

    def on_send_end(self) -> None:
        self.send_timer.end()

    def on_batch_start(self) -> None:
        self.batch_timer.start()

    def on_batch_end(self) -> None:
        self.batch_timer.end()

    def on_loop_end(self) -> None:
        avg_make_experience_duration = all_reduce_mean(self.make_experience_timer.duration, self.world_size)
        avg_overall_duration = all_reduce_mean(self.batch_timer.duration, self.world_size)
        avg_send_duration = all_reduce_mean(self.send_timer.duration, self.world_size)

        avg_throughput = self.total_samples * self.world_size / (avg_overall_duration + 1e-12)
        avg_make_experience_tflops = self.make_experience_flop / 1e12 / (avg_make_experience_duration + 1e-12)
        avg_time_per_sample = (avg_overall_duration + 1e-12) / (self.total_samples * self.world_size)
        avg_make_experience_time_per_sample = (avg_make_experience_duration + 1e-12) / (
            self.total_samples * self.world_size
        )
        avg_send_time_per_sample = (avg_send_duration + 1e-12) / (self.total_samples * self.world_size)

        print_rank_0(
            "Making Experience Performance Summary:\n"
            + f"Throughput: {avg_throughput:.3f} samples/sec\n"
            + f"TFLOPS per GPU: {avg_make_experience_tflops:.3f}\n"
            + f"Sample time (overall): {avg_time_per_sample:.3f} s\n"
            + f"Sample time (make experience): {avg_make_experience_time_per_sample:.3f} s, {avg_make_experience_time_per_sample/avg_time_per_sample*100:.2f}%\n"
            + f"Sample time (send): {avg_send_time_per_sample:.3f} s, {avg_send_time_per_sample/avg_time_per_sample*100:.2f}%\n"
        )


class TrainerPerformanceEvaluator(TrainerCallback):
    def __init__(
        self,
        actor_num_params: int,
        critic_num_params: int,
        enable_grad_checkpoint: bool = False,
        ignore_first_episodes: int = 1,
    ) -> None:
        super().__init__()
        self.world_size = get_world_size()
        self.actor_num_params = actor_num_params
        self.critic_num_params = critic_num_params
        self.enable_grad_checkpoint = enable_grad_checkpoint
        self.ignore_first_episodes = ignore_first_episodes
        self.ignore_this_episode = False

        self.episode_timer = Timer()
        self.batch_timer = Timer()
        self.update_timer = Timer()
        self.total_samples: int = 0
        self.learn_flop: int = 0

        print_rank_0(
            f"Trainer actor: {self.actor_num_params/1024**3:.2f}B, critic: {self.critic_num_params/1024**3:.2f}B, world size: {self.world_size}"
        )

    def on_episode_start(self, episodes: int) -> None:
        self.ignore_this_episode = episodes < self.ignore_first_episodes
        if self.ignore_this_episode:
            return
        self.episode_timer.start()

    def on_episode_end(self, episodes: int) -> None:
        if self.ignore_this_episode:
            return
        self.episode_timer.end()

    def on_batch_start(self) -> None:
        if self.ignore_this_episode:
            return
        self.batch_timer.start()

    def on_batch_end(self, metrics: dict, experience: Experience) -> None:
        if self.ignore_this_episode:
            return
        self.batch_timer.end()

        batch_size, seq_len = experience.sequences.shape

        self.total_samples += batch_size

        # actor forward-backward, 3 means forward(1) + backward(2)
        self.learn_flop += self.actor_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint))
        # critic forward-backward
        self.learn_flop += self.critic_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint))

    def on_update_start(self) -> None:
        if self.ignore_this_episode:
            return
        self.update_timer.start()

    def on_update_end(self) -> None:
        if self.ignore_this_episode:
            return
        self.update_timer.end()

    def on_fit_end(self) -> None:
        if self.total_samples == 0:
            print_rank_0("No samples are collected, skip trainer performance evaluation")
            return
        avg_train_duration = all_reduce_mean(self.batch_timer.duration, self.world_size)
        avg_update_duration = all_reduce_mean(self.update_timer.duration, self.world_size)
        avg_episode_duration = all_reduce_mean(self.episode_timer.duration, self.world_size)

        avg_throughput = self.total_samples * self.world_size / (avg_episode_duration + 1e-12)
        avg_learn_tflops = self.learn_flop / 1e12 / (avg_train_duration + 1e-12)
        avg_time_per_sample = (avg_episode_duration + 1e-12) / (self.total_samples * self.world_size)
        avg_train_time_per_sample = (avg_train_duration + 1e-12) / (self.total_samples * self.world_size)
        avg_update_time_per_sample = (avg_update_duration + 1e-12) / (self.total_samples * self.world_size)

        print_rank_0(
            "Learning Performance Summary:\n"
            + f"Throughput: {avg_throughput:.3f} samples/sec\n"
            + f"TFLOPS per GPU: {avg_learn_tflops:.3f}\n"
            + f"Sample time (overall): {avg_time_per_sample:.3f} s\n"
            + f"Sample time (train): {avg_train_time_per_sample:.3f} s, {avg_train_time_per_sample/avg_time_per_sample*100:.2f}%\n"
            + f"Sample time (update): {avg_update_time_per_sample:.3f} s, {avg_update_time_per_sample/avg_time_per_sample*100:.2f}%\n"
        )