Implementation:Hpcaitech ColossalAI Trainer PerformanceEvaluator

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Reinforcement Learning, Performance Profiling, Training
Last Updated	2026-02-09 00:00 GMT

Overview

A callback for the ColossalChat on-policy trainer that evaluates model performance by measuring throughput, TFLOPS, and timing breakdowns for both experience generation and learning phases.

Description

PerformanceEvaluator extends the Callback base class and tracks timing and FLOP metrics for the full on-policy RLHF pipeline. It uses Timer instances to separately measure experience making duration, learning duration, and overall episode duration. FLOP calculations account for actor generation (with variable sequence lengths), actor/critic/initial/reward model forward passes during experience making, and actor/critic forward-backward passes (with optional gradient checkpoint overhead) during learning.

The evaluator supports skipping initial episodes for warmup via ignore_episodes, aggregates metrics across distributed workers using all_reduce_mean, and writes a detailed performance summary to a file via save_eval_result_rank_0. It also includes utility functions get_world_size, divide (safe division), all_reduce_mean, and save_eval_result_rank_0.

Usage

Pass a PerformanceEvaluator instance as a callback to the OLTrainer. It automatically collects timing data during the training loop and writes a summary file at fit end. Provide model parameter counts, optional training config metadata, and a save path for the results file.

Code Reference

Source Location

Repository: Hpcaitech_ColossalAI
File: applications/ColossalChat/coati/trainer/callbacks/performance_evaluator.py
Lines: 1-191

Signature

class PerformanceEvaluator(Callback):
    def __init__(
        self,
        actor_num_params: int,
        critic_num_params: int,
        initial_model_num_params: int,
        reward_model_num_params: int,
        enable_grad_checkpoint: bool = False,
        ignore_episodes: int = 0,
        train_config: Optional[dict] = None,
        save_path: Optional[str] = None,
    ) -> None: ...

def get_world_size() -> int: ...
def save_eval_result_rank_0(s: str, save_path: str, **kwargs) -> None: ...
def divide(x: float, y: float) -> float: ...
def all_reduce_mean(x: float, world_size: int) -> float: ...

class Timer:
    def start(self) -> None: ...
    def end(self) -> None: ...
    def reset(self) -> None: ...

Import

from coati.trainer.callbacks.performance_evaluator import PerformanceEvaluator, Timer

I/O Contract

Inputs

Name	Type	Required	Description
actor_num_params	int	Yes	Number of parameters in the actor model
critic_num_params	int	Yes	Number of parameters in the critic model
initial_model_num_params	int	Yes	Number of parameters in the initial (reference) model
reward_model_num_params	int	Yes	Number of parameters in the reward model
enable_grad_checkpoint	bool	No	Whether gradient checkpointing is enabled (default False)
ignore_episodes	int	No	Number of initial episodes to skip for warmup (default 0)
train_config	dict	No	Training configuration metadata for logging (default None)
save_path	str	No	File path to write the performance summary (default None)

Outputs

Name	Type	Description
return	None	Performance summary is written to save_path file on rank 0

Usage Examples

from coati.trainer.callbacks.performance_evaluator import PerformanceEvaluator

evaluator = PerformanceEvaluator(
    actor_num_params=7_000_000_000,
    critic_num_params=7_000_000_000,
    initial_model_num_params=7_000_000_000,
    reward_model_num_params=7_000_000_000,
    enable_grad_checkpoint=True,
    ignore_episodes=1,
    train_config={"batch_size": 8, "lr": 1e-5},
    save_path="./performance_results.txt",
)

# Pass as callback to OLTrainer
trainer = PPOTrainer(
    ...,
    callbacks=[evaluator],
)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment