trinity.trainer package

Subpackages

• trinity.trainer.tinker package
  • Submodules
    • trinity.trainer.tinker.tinker_trainer module
    • trinity.trainer.tinker.utils module
  • Module contents
• trinity.trainer.verl package
  • Submodules
    • trinity.trainer.verl.dp_actor module
      • DataParallelPPOActor
    • trinity.trainer.verl.fsdp_checkpoint_manager module
      • FSDPCheckpointManager
    • trinity.trainer.verl.fsdp_workers module
    • trinity.trainer.verl.megatron_actor module
    • trinity.trainer.verl.megatron_checkpoint_manager module
    • trinity.trainer.verl.megatron_workers module
    • trinity.trainer.verl.monkey_patch module
      • apply_monkey_patch()
    • trinity.trainer.verl.utils module
      • to_data_proto()
      • compute_data_metrics()
      • get_latest_hf_checkpoint_path()
      • hf_processor()
      • apply_fsdp2()
    • trinity.trainer.verl.verl_config module
      • Data
      • FusedKernelOptions
      • ActorModel
      • Optim
      • WrapPolicy
      • FSDPConfig
      • Checkpoint
      • OverrideTransformerConfig
      • MegatronConfig
      • ProfileConfig
      • Actor
      • Ref
      • Rollout
      • ActorRolloutRef
      • CriticModel
      • Critic
      • RewardModel
      • CustomRewardFunction
      • KL_Ctrl
      • RolloutCorrection
      • Algorithm
      • Trainer
      • veRLConfig
      • load_config()
    • trinity.trainer.verl.verl_trainer module
      • CheckpointMonitor
      • VerlPPOTrainerWrapper
  • Module contents
    • patch_for_transformers_v5()

Submodules

• trinity.trainer.trainer module
  • Trainer
    • Trainer.__init__()
    • Trainer.prepare()
    • Trainer.train()
    • Trainer.train_step()
    • Trainer.need_sync()
    • Trainer.need_save()
    • Trainer.sync_weight()
    • Trainer.save_checkpoint()
    • Trainer.shutdown()
    • Trainer.train_step_num
    • Trainer.is_alive()
    • Trainer.get_actor()
  • TrainEngineWrapper
    • TrainEngineWrapper.prepare()
    • TrainEngineWrapper.train_step_num
    • TrainEngineWrapper.train_step()
    • TrainEngineWrapper.save_checkpoint()
    • TrainEngineWrapper.sync_weight()
    • TrainEngineWrapper.upload_state_dict()
    • TrainEngineWrapper.save_state_dict()
  • get_trainer_wrapper()

Module contents

class trinity.trainer.Trainer(config: Config)

Bases: 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 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 shutdown() → None

async sync_weight() → Dict

Sync the model weight.

async train() → str

Train the model.

async train_step(exps: List[Experience]) → Dict

Train one step.

Returns:

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

Return type:

bool

property train_step_num: int

Get the current training step number.

class trinity.trainer.TrainEngineWrapper

Bases: ABC

A wrapper class to wrap various training engines.

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 checkpoint.

abstractmethod async save_state_dict() → None

Only save the model state dict for Synchronizer.

abstractmethod sync_weight() → None

Sync the model weight.

abstractmethod async train_step(batch_exps: List[Experience]) → Dict

Training one step.

Parameters:

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

Returns:

Metrics of the training step.

Return type:

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.

trinity.trainer.get_trainer_wrapper(config: Config) → TrainEngineWrapper

Get a trainer wrapper.