trinity.manager package

Subpackages

• trinity.manager.config_registry package
  • 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_sft_loss_agg_mode()
      • 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_total_steps()
      • set_explore_batch_size()
      • get_train_batch_size()
      • get_train_batch_size_per_gpu()
      • 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_storage_type()
      • set_enable_replay_buffer()
      • set_reuse_cooldown_time()
      • set_priority_fn()
      • parse_priority_fn_args()
      • set_priority_fn_args()
      • set_experience_buffer_path()
      • check_experience_buffer_path()
      • set_dpo_dataset_kwargs()
      • set_sft_dataset_kwargs()
    • trinity.manager.config_registry.config_registry module
      • ConfigRegistry
    • trinity.manager.config_registry.explorer_config_manager module
      • explorer_visible()
      • set_runner_per_model()
      • set_eval_interval()
      • set_engine_num()
      • set_tensor_parallel_size()
      • check_tensor_parallel_size()
      • set_enforce_eager()
      • set_enable_prefix_caching()
      • set_enable_chunked_prefill()
      • set_gpu_memory_utilization()
      • set_dtype()
      • set_seed()
      • set_enable_thinking()
      • set_enable_history()
      • 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_style()
      • 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_log_level()
      • set_model_path()
      • check_model_path()
      • set_critic_model_path()
      • set_max_prompt_tokens()
      • set_min_response_tokens()
      • set_max_response_tokens()
      • set_max_model_len()
      • set_node_num()
      • set_gpu_per_node()
    • trinity.manager.config_registry.trainer_config_manager module
      • set_trainer_gpu_num_display()
      • use_critic()
      • set_trainer_type()
      • set_save_interval()
      • set_enable_preview()
      • set_actor_grad_clip()
      • set_training_args()
      • set_ppo_epochs()
      • set_training_strategy()
      • set_param_offload()
      • set_grad_offload()
      • set_optimizer_offload()
      • set_forward_prefetch()
      • set_offload_policy()
      • set_reshard_after_forward()
      • set_tensor_model_parallel_size()
      • set_pipeline_model_parallel_size()
      • set_virtual_pipeline_model_parallel_size()
      • set_expert_model_parallel_size()
      • set_expert_tensor_parallel_size()
      • set_context_parallel_size()
      • set_sequence_parallel()
      • set_use_distributed_optimizer()
      • set_use_dist_checkpointing()
      • set_dist_checkpointing_path()
      • set_use_mbridge()
      • set_recompute_granularity()
      • set_recompute_modules()
      • set_recompute_method()
      • set_recompute_num_layers()
      • set_resume_mode()
      • set_resume_from_path()
      • check_resume_from_path()
      • set_impl_backend()
      • set_critic_warmup()
      • set_default_hdfs_dir()
      • 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_ppo_max_token_len_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_save_checkpoint()
      • set_actor_load_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_save_checkpoint()
      • set_critic_load_checkpoint()
  • Module contents

Submodules

• trinity.manager.config_manager module
  • ConfigManager
    • ConfigManager.__init__()
    • ConfigManager.reset_session_state()
    • ConfigManager.maintain_session_state()
    • ConfigManager.get_configs()
    • ConfigManager.beginner_mode()
    • ConfigManager.expert_mode()
    • ConfigManager.generate_config()
    • ConfigManager.run_config()
• trinity.manager.state_manager module
  • StateManager
    • StateManager.__init__()
    • StateManager.save_explorer()
    • StateManager.load_explorer()
    • StateManager.save_explorer_server_url()
    • StateManager.load_explorer_server_url()
    • StateManager.save_trainer()
    • StateManager.load_trainer()
    • StateManager.save_stage()
    • StateManager.load_stage()
• trinity.manager.synchronizer module
  • Synchronizer
    • Synchronizer.trainer_status
    • Synchronizer.explorer_status_counts
    • Synchronizer._ready_condition
    • Synchronizer.model_state_dict
    • Synchronizer.model_version
    • Synchronizer.checkpoint_shard_counter
    • Synchronizer.__init__()
    • Synchronizer.add_module()
    • Synchronizer.set_trainer_status()
    • Synchronizer.get_trainer_status()
    • Synchronizer.set_explorer_status()
    • Synchronizer.get_explorer_status_counts()
    • Synchronizer.set_model_state_dict_with_step_num()
    • Synchronizer.set_model_state_dict()
    • Synchronizer.get_model_state_dict()
    • Synchronizer.get_state_dict_meta()
    • Synchronizer.setup_weight_sync_group()
    • Synchronizer.wait_new_model_state_dict()
    • Synchronizer.get_latest_model_version()
    • Synchronizer.ready_to_nccl_sync()
    • Synchronizer.get_actor()

Module contents

class trinity.manager.StateManager(path: str, trainer_name: str | None = None, explorer_name: str | None = None, config: Config | None = None, check_config: bool = False)

基类：object

A Manager class for managing the running state of Explorer and Trainer.

__init__(path: str, trainer_name: str | None = None, explorer_name: str | None = None, config: Config | None = None, check_config: bool = False)

load_explorer() → dict

load_explorer_server_url() → str | None

load_stage() → dict

load_trainer() → dict

save_explorer(current_step: int, taskset_states: List[Dict]) → None

save_explorer_server_url(url: str) → None

save_stage(current_stage: int) → None

save_trainer(current_step: int, sample_strategy_state: dict) → None

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

基类：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.

参数:

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

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

Get or create a remote Ray actor for the Synchronizer.

参数:

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

• namespace -- Optional Ray namespace for the actor.

返回:

A reference to the Synchronizer actor.

get_explorer_status_counts() → Dict[RunningStatus, int]

Return the current status counts for all explorers.

async get_latest_model_version() → int

Get the latest model version available in the synchronizer.

返回:

The current model version.

get_model_state_dict()

Return the current model state and its version.

async get_state_dict_meta()

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

返回:

(name, dtype, shape).

返回类型:

List of tuples

get_trainer_status() → RunningStatus

Get the current status of the trainer.

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.

参数:

• module -- Either 'trainer' or 'explorer'.

• trainer_step -- Optional step number from the trainer.

返回:

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

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

Update the status count for an explorer.

参数:

• status -- New status of the explorer.

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

async set_model_state_dict(model_state_dict: dict | None | str | Tuple[str, str], trainer_step: int)

Set the new model state and update the version.

参数:

• model_state_dict -- The PyTorch model state dictionary.

• trainer_step -- Step number associated with this model version.

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.

参数:

• step_num -- Training step number corresponding to the checkpoint.

• world_size -- Number of shards expected for this checkpoint.

返回:

The updated model version (step number).

async set_trainer_status(status: RunningStatus)

Update the status of the trainer.

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

Notify the explorer actor to setup weight sync group.

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

参数:

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

参数:

current_version -- Current model version known to one explorer.

返回:

The new model version after it has been updated.