trinity.trainer.verl.fsdp_workers module

The main entry point to run the PPO algorithm. Modified from https://github.com/volcengine/verl/blob/v0.4.1/verl/workers/fsdp_workers.py

class trinity.trainer.verl.fsdp_workers.ActorRolloutRefWorker(*args, **kwargs)[source]

Bases: Worker

This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy or a hybrid engine based on the config.rollout

__init__(config: DictConfig, role: str)[source]

Initialize the worker with environment settings and device configuration.

Parameters:: cuda_visible_devices (str, optional) – CUDA visible devices configuration. Defaults to None.

init_model()[source]

setup_weight_sync_group()[source]

sync_weight()[source]

upload_state_dict(trainer_step: int)[source]

set_algorithm(algo_config: AlgorithmConfig)[source]

update_actor(data: DataProto)[source]

compute_log_prob(data: DataProto)[source]

compute_ref_log_prob(data: DataProto)[source]

save_checkpoint(local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None, model_state_dict_only=False)[source]

load_checkpoint(local_path, hdfs_path=None, del_local_after_load=False)[source]

clear_optimizer_state()[source]

wait_on_save_thread() → None[source]

class trinity.trainer.verl.fsdp_workers.CriticWorker(*args, **kwargs)[source]

Bases: Worker

__init__(config)[source]

Initialize the worker with environment settings and device configuration.

Parameters:: cuda_visible_devices (str, optional) – CUDA visible devices configuration. Defaults to None.

init_model()[source]

compute_values(data: DataProto)[source]

update_critic(data: DataProto)[source]

save_checkpoint(local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None)[source]

load_checkpoint(local_path, hdfs_path=None, del_local_after_load=True)[source]

clear_optimizer_state()[source]

wait_on_save_thread() → None[source]