xtuner.v1.train.rl_trainer.RLTrainer

xtuner.v1.train.rl_trainer.RLTrainer#

class xtuner.v1.train.rl_trainer.RLTrainer(*, load_from: str | Path, resources: AcceleratorResourcesConfig, cpu_resources: CPUResourcesConfig | None = None, rollout_config: RolloutConfig, dataflow_config: DataFlowConfig, judger_config: JudgerConfig, replay_buffer_config: ReplayBufferConfig, train_worker_cfg: WorkerConfig, evaluator_config: EvaluatorConfig | None = None, tokenizer_path: str | Path, work_dir: Path | str | None = None, log_dir: Path | str | None = None, total_epochs: int, auto_resume: bool = False, load_checkpoint_cfg: LoadCheckpointConfig = pydantic.BaseModel, strict_load: bool = True, checkpoint_interval: int | None = -1, checkpoint_maxkeep: int | None = -1, checkpoint_no_save_optimizer: bool = False, skip_checkpoint_validation: bool = False, hf_interval: int | None = None, hf_max_keep: int | None = None, seed: int = 42, debug: bool = False, debug_rollout: bool = False, rollout_steps: int | None = None, exp_tracker: Literal['tensorboard', 'jsonl'] = 'tensorboard', display_all_workers_log: bool = False, trainer_cfg: RLTrainerConfig | None = None, advantage_estimator_config: BaseAdvantageConfig = pydantic.BaseModel)[source]#

Universal Reinforcement Learning Trainer for XTuner.

A flexible RL training orchestrator that supports multiple RL algorithms through pluggable training workers and controllers. Manages the complete RL training workflow including rollout generation, policy updates, evaluation, and checkpoint management.

Training Workflow:

  1. Initialize distributed workers and rollout environment

  2. Generate experiences using current policy

  3. Update policy using algorithm-specific training logic

  4. Synchronize weights between training and rollout workers

  5. Evaluate model performance and save checkpoints

Parameters:

  • load_from (str | Path) – Path to the base model to load. Should be a HuggingFace model path (e.g., “meta-llama/Llama-2-7b-hf”) or local model directory.

  • resources (AcceleratorResourcesConfig) – Configuration for distributed computing resources including number of workers, GPU allocation, and placement groups.

  • rollout_config (RolloutConfig) – Configuration for rollout workers that generate experiences by interacting with the environment.

  • dataflow_config (DataFlowConfig) – Data orchestration configuration controlling experience collection, batch formation, and data distribution across workers.

  • judger_config (JudgerConfig) – Configuration for the reward model or scoring system that evaluates generated responses and provides training signals.

  • replay_buffer_config (ReplayBufferConfig) – Settings for experience replay buffer including capacity, sampling strategy, and data retention policies.

  • evaluator_config (EvaluatorConfig | None) – Evaluation configuration specifying metrics, evaluation datasets, and assessment frequency for monitoring training progress. Defaults to None.

  • train_worker_cfg (WorkerConfig) – Configuration for distributed training workers including model architecture, optimizer settings, loss functions, and parallelism.

  • tokenizer_path (str | Path) – Path to the tokenizer for text preprocessing. Should be compatible with the base model specified in load_from.

  • work_dir (Path | str | None) – Working directory for experiment outputs, checkpoints, and logs. Defaults to None.

  • log_dir (Path | str | None) – Directory for training logs and monitoring outputs. Defaults to None.

  • total_epochs (int) – Total number of training epochs to execute.

  • enable_evaluate (bool) – Whether to perform periodic evaluation during training.

  • resume_config (ResumeConfig | None) – Configuration for resuming training from a previous checkpoint. Defaults to None.

  • auto_resume (bool) – Whether to automatically resume training. Defaults to False.

  • load_checkpoint_cfg (LoadCheckpointConfig) – Configuration for loading checkpoints.

  • strict_load (bool) – Whether to strictly enforce checkpoint loading compatibility. Defaults to True.

  • hf_interval (int | None) – Interval (in epochs) for saving HuggingFace format checkpoints. Defaults to None.

  • hf_max_keep (int | None) – Maximum number of HuggingFace checkpoints to retain. Defaults to None.

  • seed (int) – Random seed for reproducible training. Defaults to 42.

  • debug (bool) – Enable debug mode with additional logging. Defaults to False.

  • debug_rollout (bool) – Enable debug mode for rollout workers. Defaults to False.

  • rollout_steps (int | None) – Total number of rollout steps to perform. If specified, overrides total_epochs. Defaults to None.

  • display_all_workers_log (bool) – Whether to display logs from all workers. Defaults to False.

  • exp_tracker (Literal["tensorboard", "jsonl"]) – Type of experiment tracker to use. Options are “tensorboard” or “jsonl”. Defaults to “tensorboard”.

Examples:

Example configuration for GRPO RL training setup:

trainer = RLTrainer(
    load_from="Qwen3-8B",
    resources=resources_config,
    rollout_config=rollout_cfg,
    dataflow_config=dataflow_cfg,
    judger_config=judger_cfg,
    replay_buffer_config=buffer_cfg,
    evaluator_config=eval_cfg,
    train_worker_cfg=worker_cfg,
    tokenizer_path="Qwen3-8B",
    total_epochs=10,
    enable_evaluate=True
)
trainer.fit()

Methods

fit()

Run the RL training loop.

from_config(config)

Create a Trainer instance from a TrainerConfig.
fit()[source]#

Run the RL training loop.

This method executes the main rl training loop, iterating generating through the dataset and performing training steps. It handles rollout, prepare training data, update policy , synchronize model weights, and evaluation.

classmethod from_config(config: RLTrainerConfig) → Self[source]#

Create a Trainer instance from a TrainerConfig.

Parameters:

config (TrainerConfig) – TrainerConfig instance containing all configuration parameters.

Returns:

Trainer instance initialized with the provided config.

Return type:

Self

Contents