trinity.trainer package

trinity.trainer package#

Subpackages#

Submodules#

Module contents#

class trinity.trainer.Trainer(config: Config)[源代码]#

基类：object

Consume the experience and train the model.

__init__(config: Config) → None[源代码]#
classmethod get_actor(config: Config)[源代码]#

Get a Ray actor for the trainer.

async get_weight_sync_info() → Tuple[str, int, List] | None[源代码]#

Get rendezvous info for NCCL weight sync group setup.

Returns (master_address, master_port, state_dict_meta) from the trainer's GPU worker rank 0. Called by Synchronizer before coordinating NCCL group creation.

async is_alive() → bool[源代码]#

Check if the trainer is alive.

need_save() → bool[源代码]#

Whether to save the checkpoint.

async need_sync() → bool[源代码]#

Whether to sync the model weight.

async prepare() → None[源代码]#

Prepare the trainer.

async save_checkpoint(block_until_saved: bool = False, save_as_hf: bool = False) → Dict[源代码]#
async setup_weight_sync_group(master_address: str, master_port: int, world_size: int, group_name: str, timeout: int) → None[源代码]#

Join the NCCL weight sync group. Called by Synchronizer.

async shutdown() → None[源代码]#
async sync_weight() → Dict[源代码]#

Sync the model weight.

async teardown_weight_sync_group() → None[源代码]#

Destroy the NCCL weight sync group. Called by Synchronizer.

async train() → str[源代码]#

Train the model.

async train_step(exps: List[Experience]) → Dict[源代码]#

Train one step.

返回:

Whether to continue training. Dict: Metrics of the training step.

返回类型:

bool

property train_step_num: int#

Get the current training step number.

class trinity.trainer.TrainEngineWrapper[源代码]#

基类：ABC

A wrapper class to wrap various training engines.

abstractmethod async get_weight_sync_info() → Tuple[str, int, List] | None[源代码]#

Get (master_address, master_port, state_dict_meta) for NCCL group setup.

abstractmethod async prepare() → None[源代码]#

Do some preparation before training started.

abstractmethod async save_checkpoint(block_until_saved: bool = False, save_as_hf: bool = False) → None[源代码]#

Save the whole checkpoint (Including model, optimizer, and other states).

abstractmethod async save_state_dict() → None[源代码]#

Only save the model state dict for Synchronizer. (For CHECKPOINT sync method)

abstractmethod async setup_weight_sync_group(master_address: str, master_port: int, world_size: int, group_name: str, timeout: int) → None[源代码]#

Join the NCCL weight sync group.

abstractmethod sync_weight_nccl() → None[源代码]#

Sync the model weight by NCCL. (For NCCL sync method)

abstractmethod async teardown_weight_sync_group() → None[源代码]#

Tear down the NCCL weight sync group.

abstractmethod async train_step(batch_exps: List[Experience]) → Dict[源代码]#

Training one step.

参数:

batch_exps (List[Experience]) -- A batch of experiences to train.

返回:

Metrics of the training step.

返回类型:

Dict

abstract property train_step_num: int#

Get the current training step number.

abstractmethod async upload_state_dict() → None[源代码]#

Upload the state dict to Synchronizer. (For MEMORY sync method)

async wait_for_save() → None[源代码]#

Wait for any pending background save operations to complete.

Default implementation is a no-op. Override in subclasses that use background save threads to ensure the checkpoint iteration file is written before the trainer exits.

trinity.trainer.get_latest_hf_checkpoint_path(config: Config) → str | None[源代码]#

Return the latest HF checkpoint path for a verl trainer config.

trinity.trainer.get_trainer_wrapper(config: Config) → TrainEngineWrapper[源代码]#

Get a trainer wrapper.

trinity.trainer.is_verl_legacy() → bool[源代码]#

Return True when the installed verl package is < 0.8 (legacy backend).