trinity.manager

Subpackages

• trinity.manager.config_registry
  • Submodules
  • trinity.manager.config_registry.algorithm_config_manager module
    • set_algorithm_type()
    • set_repeat_times()
    • set_sample_strategy()
    • set_expert_data_ratio_in_sample_strategy()
    • set_advantage_fn()
    • set_gamma_in_advantage_fn()
    • set_lam_in_advantage_fn()
    • set_epsilon_in_advantage_fn()
    • set_opmd_baseline_in_advantage_fn()
    • set_tau_in_advantage_fn()
    • set_kl_loss_fn()
    • set_kl_coef_in_kl_loss_fn()
    • set_kl_penalty_fn()
    • set_adaptive_in_kl_penalty_fn()
    • set_kl_coef_in_kl_penalty_fn()
    • set_policy_loss_fn()
    • set_clip_range_in_policy_loss_fn()
    • set_use_token_level_loss_in_policy_loss_fn()
    • set_beta_in_policy_loss_fn()
    • set_label_smoothing_in_policy_loss_fn()
    • set_tau_in_policy_loss_fn()
    • set_mu_in_policy_loss_fn()
    • set_entropy_loss_fn()
    • set_entropy_coef_in_entropy_loss_fn()
  • trinity.manager.config_registry.buffer_config_manager module
    • set_total_epochs()
    • set_explore_batch_size()
    • set_train_batch_size()
    • check_train_batch_size()
    • set_buffer_max_retry_times()
    • set_max_retry_interval()
    • set_taskset_path()
    • check_taskset_path()
    • set_taskset_args()
    • set_eval_tasksets()
    • set_default_workflow_type()
    • set_default_eval_workflow_type()
    • set_default_reward_fn_type()
    • set_system_prompt()
    • set_reply_prefix()
    • set_storage_type()
    • set_use_priority_queue()
    • set_reuse_cooldown_time()
    • set_priority_fn()
    • set_priority_decay()
    • set_experience_buffer_path()
    • check_experience_buffer_path()
    • set_dpo_dataset_kwargs()
    • set_sft_warmup_dataset_path()
    • check_sft_warmup_dataset_path()
    • set_sft_warmup_dataset_args()
    • set_sft_warmup_steps()
  • trinity.manager.config_registry.config_registry module
    • ConfigRegistry
      • ConfigRegistry.__init__()
      • ConfigRegistry.set_unfinished_fields()
      • ConfigRegistry.default_config
      • ConfigRegistry.get()
      • ConfigRegistry.get_check_func()
      • ConfigRegistry.get_configs()
      • ConfigRegistry.register_config()
      • ConfigRegistry.register_check()
  • trinity.manager.config_registry.explorer_config_manager module
    • explorer_visible()
    • set_runner_per_model()
    • set_max_timeout()
    • set_explorer_max_retry_times()
    • set_eval_interval()
    • set_bench_on_latest_checkpoint()
    • set_engine_type()
    • set_engine_num()
    • set_tensor_parallel_size()
    • check_tensor_parallel_size()
    • set_use_v1()
    • set_enforce_eager()
    • set_enable_prefix_caching()
    • set_enable_chunked_prefill()
    • set_gpu_memory_utilization()
    • set_dtype()
    • set_seed()
    • set_enable_thinking()
    • set_enable_openai_api()
    • set_enable_auto_tool_choice()
    • set_tool_call_parser()
    • set_reasoning_parser()
    • set_auxiliary_models()
    • check_auxiliary_models()
    • set_sync_method()
    • set_sync_interval()
    • set_sync_timeout()
  • trinity.manager.config_registry.model_config_manager module
    • set_total_gpu_num()
    • set_trainer_gpu_num()
    • set_project()
    • set_exp_name()
    • set_checkpoint_root_dir()
    • check_checkpoint_root_dir()
    • set_monitor_type()
    • set_model_path()
    • check_model_path()
    • set_critic_model_path()
    • set_max_response_tokens()
    • set_max_model_len()
    • set_node_num()
    • set_gpu_per_node()
  • trinity.manager.config_registry.trainer_config_manager module
    • use_critic()
    • set_trainer_type()
    • set_save_interval()
    • set_enable_preview()
    • set_actor_grad_clip()
    • set_training_args()
    • set_ppo_epochs()
    • set_training_strategy()
    • use_fsdp()
    • set_param_offload()
    • set_optimizer_offload()
    • set_forward_prefetch()
    • set_resume_mode()
    • set_resume_from_path()
    • check_resume_from_path()
    • set_impl_backend()
    • set_critic_warmup()
    • set_total_training_steps()
    • set_default_hdfs_dir()
    • set_remove_previous_ckpt_in_save()
    • set_del_local_ckpt_after_load()
    • set_max_actor_ckpt_to_keep()
    • set_max_critic_ckpt_to_keep()
    • set_norm_adv_by_std_in_grpo()
    • set_use_kl_in_reward()
    • set_kl_penalty()
    • set_kl_ctrl_type()
    • set_kl_ctrl_coef()
    • set_horizon()
    • set_target_kl()
    • set_actor_ppo_micro_batch_size_per_gpu()
    • set_ref_log_prob_micro_batch_size_per_gpu()
    • set_actor_ulysses_sequence_parallel_size()
    • set_actor_entropy_from_logits_with_chunking()
    • set_actor_entropy_checkpointing()
    • set_actor_lr()
    • set_actor_warmup_style()
    • set_actor_lr_warmup_steps_ratio()
    • set_actor_checkpoint()
    • set_critic_lr()
    • set_critic_warmup_style()
    • set_critic_lr_warmup_steps_ratio()
    • set_critic_grad_clip()
    • set_critic_cliprange_value()
    • set_critic_ppo_micro_batch_size_per_gpu()
    • set_critic_ulysses_sequence_parallel_size()
    • set_critic_checkpoint()
  • Module contents

Submodules

trinity.manager.config_manager module

class trinity.manager.config_manager.ConfigManager

Bases: object

__init__()

reset_session_state()

maintain_session_state()

get_configs(*config_names: str, columns_spec: List[int] | None = None)

beginner_mode()

expert_mode()

generate_config()

run_config(parent, yaml_config: str) → None

trinity.manager.manager module

Data manager.

class trinity.manager.manager.CacheManager(config: Config, check_config: bool = False)

Bases: object

A Manager class for managing the cache dir.

__init__(config: Config, check_config: bool = False)

save_explorer(current_task_index: int, current_step: int) → None

load_explorer() → dict

save_trainer(current_step: int) → None

load_trainer() → dict

trinity.manager.synchronizer module

A centralized synchronizer for coordinating explorer and trainer.

class trinity.manager.synchronizer.Synchronizer(config: Config, module_ref: ActorHandle)

Bases: object

A central component to manage synchronization of models and states between the trainer and one or more explorers in a distributed training setup.

trainer_status

Current status of the trainer (e.g., running, waiting).

explorer_status_counts

Dictionary tracking the number of explorers in each status.

_ready_condition

Async condition variable for signaling state changes.

model_state_dict

The latest model weights.

model_version

Version number of the current model.

checkpoint_shard_counter

Tracks how many shards are received from trainer for a specific train step.

__init__(config: Config, module_ref: ActorHandle)

async add_module(module_ref: ActorHandle) → None

Adds a module to be tracked by the synchronizer.

Parameters:

module_ref – The Ray actor handle of the module to track.

async set_trainer_status(status: RunningStatus)

Update the status of the trainer.

get_trainer_status() → RunningStatus

Get the current status of the trainer.

async set_explorer_status(status: RunningStatus, old_status: RunningStatus | None = None)

Update the status count for an explorer.

Parameters:

• status – New status of the explorer.

• old_status – Previous status if changing from one to another.

get_explorer_status_counts() → Dict[RunningStatus, int]

Return the current status counts for all explorers.

async set_model_state_dict_with_step_num(step_num: int | None = None, world_size: int | None = None) → int

Load and set the model state dictionary from a checkpoint at a specific step.

Parameters:

• step_num – Training step number corresponding to the checkpoint.

• world_size – Number of shards expected for this checkpoint.

Returns:

The updated model version (step number).

async set_model_state_dict(model_state_dict: dict | None, trainer_step: int)

Set the new model state and update the version.

Parameters:

• model_state_dict – The PyTorch model state dictionary.

• trainer_step – Step number associated with this model version.

get_model_state_dict()

Return the current model state and its version.

get_state_dict_meta()

Return metadata about the model state (names, data types, shapes).

Returns:

(name, dtype, shape).

Return type:

List of tuples

async setup_weight_sync_group(master_address: str, master_port: int, state_dict_meta: List | None = None)

Notify the explorer actor to setup weight sync group.

This is used to initialize NCCL-based synchronization for distributed training.

Parameters:

• master_address – IP address of the master node.

• master_port – Port used for synchronization.

• state_dict_meta – Metadata of the model parameters.

async wait_new_model_state_dict(current_version: int, no_wait: bool = False) → int

Wait until a new model state is available.

Parameters:

current_version – Current model version known to one explorer.

Returns:

The new model version after it has been updated.

async ready_to_nccl_sync(module: str, trainer_step: int | None = None) → int | None

Prepare for NCCL-based synchronization between modules.

Only supports one explorer currently.

Parameters:

• module – Either ‘trainer’ or ‘explorer’.

• trainer_step – Optional step number from the trainer.

Returns:

The model version if both sides are ready; otherwise None.

classmethod get_actor(config: Config | None = None, namespace: str | None = None)

Get or create a remote Ray actor for the Synchronizer.

Parameters:

• config – Optional configuration to use for creating the actor.

• namespace – Optional Ray namespace for the actor.

Returns:

A reference to the Synchronizer actor.

Module contents

class trinity.manager.CacheManager(config: Config, check_config: bool = False)

Bases: object

A Manager class for managing the cache dir.

__init__(config: Config, check_config: bool = False)

save_explorer(current_task_index: int, current_step: int) → None

load_explorer() → dict

save_trainer(current_step: int) → None

load_trainer() → dict

class trinity.manager.Synchronizer(config: Config, module_ref: ActorHandle)

Bases: object

A central component to manage synchronization of models and states between the trainer and one or more explorers in a distributed training setup.

trainer_status

Current status of the trainer (e.g., running, waiting).

explorer_status_counts

Dictionary tracking the number of explorers in each status.

_ready_condition

Async condition variable for signaling state changes.

model_state_dict

The latest model weights.

model_version

Version number of the current model.

checkpoint_shard_counter

Tracks how many shards are received from trainer for a specific train step.

__init__(config: Config, module_ref: ActorHandle)

async add_module(module_ref: ActorHandle) → None

Adds a module to be tracked by the synchronizer.

Parameters:

module_ref – The Ray actor handle of the module to track.

async set_trainer_status(status: RunningStatus)

Update the status of the trainer.

get_trainer_status() → RunningStatus

Get the current status of the trainer.

async set_explorer_status(status: RunningStatus, old_status: RunningStatus | None = None)

Update the status count for an explorer.

Parameters:

• status – New status of the explorer.

• old_status – Previous status if changing from one to another.

get_explorer_status_counts() → Dict[RunningStatus, int]

Return the current status counts for all explorers.

async set_model_state_dict_with_step_num(step_num: int | None = None, world_size: int | None = None) → int

Load and set the model state dictionary from a checkpoint at a specific step.

Parameters:

• step_num – Training step number corresponding to the checkpoint.

• world_size – Number of shards expected for this checkpoint.

Returns:

The updated model version (step number).

async set_model_state_dict(model_state_dict: dict | None, trainer_step: int)

Set the new model state and update the version.

Parameters:

• model_state_dict – The PyTorch model state dictionary.

• trainer_step – Step number associated with this model version.

get_model_state_dict()

Return the current model state and its version.

get_state_dict_meta()

Return metadata about the model state (names, data types, shapes).

Returns:

(name, dtype, shape).

Return type:

List of tuples

async setup_weight_sync_group(master_address: str, master_port: int, state_dict_meta: List | None = None)

Notify the explorer actor to setup weight sync group.

This is used to initialize NCCL-based synchronization for distributed training.

Parameters:

• master_address – IP address of the master node.

• master_port – Port used for synchronization.

• state_dict_meta – Metadata of the model parameters.

async wait_new_model_state_dict(current_version: int, no_wait: bool = False) → int

Wait until a new model state is available.

Parameters:

current_version – Current model version known to one explorer.

Returns:

The new model version after it has been updated.

async ready_to_nccl_sync(module: str, trainer_step: int | None = None) → int | None

Prepare for NCCL-based synchronization between modules.

Only supports one explorer currently.

Parameters:

• module – Either ‘trainer’ or ‘explorer’.

• trainer_step – Optional step number from the trainer.

Returns:

The model version if both sides are ready; otherwise None.

classmethod get_actor(config: Config | None = None, namespace: str | None = None)

Get or create a remote Ray actor for the Synchronizer.

Parameters:

• config – Optional configuration to use for creating the actor.

• namespace – Optional Ray namespace for the actor.

Returns:

A reference to the Synchronizer actor.