Implementation:Hpcaitech ColossalAI Trainer Base

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

Overview

Abstract base trainer classes for ColossalChat, providing SLTrainer for supervised learning tasks (pretraining, SFT, DPO, reward model) and OLTrainer for online learning tasks (RLHF-PPO).

Description

This module defines two abstract base classes. SLTrainer provides a simple epoch-based training loop that calls _train and _eval for each epoch, managed through a Booster and Plugin from ColossalAI. OLTrainer implements an on-policy reinforcement learning loop with separate collect and update phases, an experience buffer, callback lifecycle management (fit, episode, make_experience, learn_epoch, learn_batch), and checkpoint saving at configurable intervals.

OLTrainer uses context managers (_fit_ctx and _episode_ctx) for clean callback dispatching around the fit and episode boundaries. The update phase clears the experience buffer after each episode, consistent with on-policy RL requirements.

Usage

Subclass SLTrainer for supervised fine-tuning, reward model training, or DPO training. Subclass OLTrainer for online RL algorithms like PPO. Both require implementing abstract methods for the actual training logic, evaluation, and data preparation.

Code Reference

Source Location

Repository: Hpcaitech_ColossalAI
File: applications/ColossalChat/coati/trainer/base.py
Lines: 1-218

Signature

class SLTrainer(ABC):
    def __init__(
        self,
        booster: Booster,
        max_epochs: int,
        model: nn.Module,
        optimizer: Optimizer,
        plugin: Plugin,
        start_epoch: int = 0,
    ) -> None: ...

    @abstractmethod
    def _train(self, epoch): ...
    @abstractmethod
    def _eval(self, epoch): ...
    @abstractmethod
    def _before_fit(self): ...
    def fit(self, *args, **kwargs): ...

class OLTrainer(ABC):
    def __init__(
        self,
        actor_booster: Booster,
        critic_booster: Booster,
        data_buffer: NaiveExperienceBuffer,
        sample_buffer: bool,
        dataloader_pin_memory: bool,
        callbacks: List[Callable] = [],
    ) -> None: ...

    def fit(
        self,
        num_episodes: int,
        num_collect_steps: int,
        num_update_steps: int,
        *args, **kwargs,
    ): ...

    @abstractmethod
    def _make_experience(self, collect_step: int): ...
    @abstractmethod
    def _learn(self, update_step: int): ...
    @abstractmethod
    def _setup_update_phrase_dataload(self): ...
    @abstractmethod
    def _save_checkpoint(self, episode: int = 0): ...

Import

from coati.trainer.base import SLTrainer, OLTrainer

I/O Contract

Inputs (SLTrainer)

Name	Type	Required	Description
booster	Booster	Yes	ColossalAI Booster instance for distributed training
max_epochs	int	Yes	Maximum number of training epochs
model	nn.Module	Yes	The model to train
optimizer	Optimizer	Yes	The optimizer for training
plugin	Plugin	Yes	ColossalAI Plugin for parallelism strategy
start_epoch	int	No	Starting epoch, non-zero when resuming from checkpoint (default 0)

Inputs (OLTrainer)

Name	Type	Required	Description
actor_booster	Booster	Yes	Booster for the actor model
critic_booster	Booster	Yes	Booster for the critic model
data_buffer	NaiveExperienceBuffer	Yes	Buffer for collecting experience data
sample_buffer	bool	Yes	Whether to sample from buffer during update phase
dataloader_pin_memory	bool	Yes	Whether to pin memory for data loaders
callbacks	List[Callable]	No	List of callback instances (default [])

Outputs

Name	Type	Description
fit return	None	Both trainers run the training loop to completion

Usage Examples

from coati.trainer.base import SLTrainer
from colossalai.booster import Booster

class MyRewardTrainer(SLTrainer):
    def _before_fit(self, train_loader, eval_loader):
        self.train_loader = train_loader
        self.eval_loader = eval_loader

    def _train(self, epoch):
        for batch in self.train_loader:
            loss = self.model(batch)
            self.booster.backward(loss, self.optimizer)
            self.optimizer.step()
            self.optimizer.zero_grad()

    def _eval(self, epoch):
        self.model.eval()
        # evaluation logic

trainer = MyRewardTrainer(
    booster=booster, max_epochs=3,
    model=model, optimizer=optimizer, plugin=plugin,
)
trainer.fit(train_loader=train_dl, eval_loader=eval_dl)

Related Pages

Page Connections

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

Principle

Implementation

Heuristic

Environment