| """Launch the trainer""" |
| import argparse |
| import asyncio |
| import os |
| import sys |
| import traceback |
| from pathlib import Path |
| from pprint import pprint |
| from typing import Optional |
|
|
| import ray |
|
|
| from trinity.buffer.pipelines.task_pipeline import check_and_run_task_pipeline |
| from trinity.common.config import Config, load_config |
| from trinity.common.constants import DEBUG_NAMESPACE, PLUGIN_DIRS_ENV_VAR |
| from trinity.explorer.explorer import Explorer |
| from trinity.manager.checkpoint_converter import Converter |
| from trinity.manager.state_manager import StateManager |
| from trinity.trainer.trainer import Trainer |
| from trinity.utils.dlc_utils import is_running, setup_ray_cluster, stop_ray_cluster |
| from trinity.utils.log import get_logger |
| from trinity.utils.plugin_loader import load_plugins |
|
|
| logger = get_logger(__name__) |
|
|
|
|
| def bench(config: Config) -> None: |
| """Evaluate model.""" |
| config.explorer.name = "benchmark" |
| try: |
| explorer = Explorer.get_actor(config) |
| ray.get(explorer.prepare.remote()) |
| ray.get(explorer.benchmark.remote()) |
| logger.info("Benchmark finished.") |
| ray.get(explorer.shutdown.remote()) |
| except Exception: |
| logger.error(f"Benchmark failed:\n{traceback.format_exc()}") |
|
|
|
|
| def explore(config: Config) -> None: |
| """Run explorer.""" |
| try: |
| explorer = Explorer.get_actor(config) |
| ray.get(explorer.prepare.remote()) |
| ray.get(explorer.sync_weight.remote()) |
| ray.get(explorer.explore.remote()) |
| ray.get(explorer.shutdown.remote()) |
| except Exception: |
| logger.error(f"Explorer failed:\n{traceback.format_exc()}") |
|
|
|
|
| def train(config: Config) -> None: |
| """Run trainer.""" |
| try: |
| trainer = Trainer.get_actor(config) |
| ray.get(trainer.prepare.remote()) |
| ray.get(trainer.sync_weight.remote()) |
| ray.get(trainer.train.remote()) |
| ray.get(trainer.shutdown.remote()) |
| except Exception: |
| logger.error(f"Trainer failed:\n{traceback.format_exc()}") |
|
|
|
|
| def serve(config: Config) -> None: |
| """Run explorer in server mode.""" |
| try: |
| explorer = Explorer.get_actor(config) |
| ray.get(explorer.prepare.remote()) |
| ray.get(explorer.sync_weight.remote()) |
| ray.get(explorer.serve.remote()) |
| ray.get(explorer.shutdown.remote()) |
| except Exception: |
| logger.error(f"Explorer failed:\n{traceback.format_exc()}") |
|
|
|
|
| def both(config: Config) -> None: |
| """Setup both explorer and trainer. |
| |
| For the explorer, a step contains `batch_size * sync_interval` number |
| of rollout tasks. |
| |
| For the trainer, it has to consume all experiences generated by the explorer in |
| the latest step. The specific number of experiences may vary for different |
| algorithms and tasks. |
| """ |
| try: |
| explorer = Explorer.get_actor(config) |
| trainer = Trainer.get_actor(config) |
| ray.get([explorer.__ray_ready__.remote(), trainer.__ray_ready__.remote()]) |
| ray.get( |
| [ |
| explorer.prepare.remote(), |
| trainer.prepare.remote(), |
| ] |
| ) |
| ray.get( |
| [ |
| explorer.sync_weight.remote(), |
| trainer.sync_weight.remote(), |
| ] |
| ) |
| ready_ref, wait_ref = ray.wait( |
| [ |
| explorer.explore.remote(), |
| trainer.train.remote(), |
| ], |
| num_returns=1, |
| ) |
|
|
| ready = ray.get(ready_ref[0]) |
| if ready == config.trainer.name: |
| logger.info( |
| "===========================================================\n" |
| "> Launcher detected that the `Trainer` process has finished.\n" |
| "> Stopping the explorer process immediately.\n" |
| "===========================================================" |
| ) |
| ray.wait(wait_ref, timeout=5) |
| elif ready == config.explorer.name: |
| logger.info( |
| "===============================================================\n" |
| "> Launcher detected that the `Explorer` process has finished.\n" |
| "> `Trainer` process may need to save the model checkpoint.\n" |
| f"> Waiting {config.synchronizer.sync_timeout} s for the trainer process...\n" |
| "> You can force stop the `Trainer` process by pressing Ctrl+C.\n" |
| "===============================================================" |
| ) |
| ray.wait(wait_ref, timeout=config.synchronizer.sync_timeout) |
| ray.wait( |
| [explorer.shutdown.remote(), trainer.shutdown.remote()], |
| timeout=config.synchronizer.sync_timeout, |
| num_returns=2, |
| ) |
| except Exception: |
| logger.error(f"Explorer or Trainer failed:\n{traceback.format_exc()}") |
|
|
|
|
| MODE_MAP = { |
| "explore": explore, |
| "train": train, |
| "both": both, |
| "bench": bench, |
| "serve": serve, |
| } |
|
|
|
|
| def run_stage(config: Config) -> None: |
| ray.init( |
| address=config.cluster.ray_address, |
| ignore_reinit_error=True, |
| namespace=config.ray_namespace, |
| runtime_env={"env_vars": config.get_envs()}, |
| ) |
| pprint(config) |
| try: |
| check_and_run_task_pipeline(config) |
| MODE_MAP[config.mode](config) |
| finally: |
| if config.monitor.enable_ray_timeline: |
| timeline_file = os.path.join(config.monitor.cache_dir, "timeline.json") |
| logger.info(f"Exporting Ray timeline to {timeline_file}...") |
| ray.timeline(filename=timeline_file) |
| logger.info("Done. You can open the timeline file in `chrome://tracing`") |
| ray.shutdown() |
|
|
|
|
| def run(config_path: str, dlc: bool = False, plugin_dir: str = None): |
| if plugin_dir: |
| os.environ[PLUGIN_DIRS_ENV_VAR] = plugin_dir |
| load_plugins() |
| config = load_config(config_path) |
|
|
| if dlc: |
| cluster_namespace = f"{config.project}-{config.name}" |
| config.cluster.ray_address = setup_ray_cluster(namespace=cluster_namespace) |
|
|
| if not is_running(): |
| raise RuntimeError("Ray is not running, please start it by `ray start --head`.") |
|
|
| try: |
| if config.stages: |
| from trinity.trainer.verl.utils import get_latest_hf_checkpoint_path |
|
|
| state_manager = StateManager( |
| path=os.path.join(config.checkpoint_root_dir, config.project, config.name) |
| ) |
| latest_stage = state_manager.load_stage().get("latest_stage", 0) |
| prev_stage_checkpoint = None |
| for i, stage_config in enumerate(config): |
| if i < latest_stage: |
| logger.info( |
| "===========================================================\n" |
| f"> Skipping completed stage {i + 1}/{len(config.stages)}...\n" |
| "===========================================================" |
| ) |
| stage_config.check_and_update() |
| else: |
| logger.info( |
| "===========================================================\n" |
| f"> Starting stage {i + 1}/{len(config.stages)}...\n" |
| "===========================================================" |
| ) |
| state_manager.save_stage(i) |
| if prev_stage_checkpoint is not None: |
| stage_config.model.model_path = prev_stage_checkpoint |
| stage_config.check_and_update() |
| run_stage(stage_config) |
| logger.info( |
| "===========================================================\n" |
| f"> Stage {i + 1}/{len(config.stages)} finished.\n" |
| "===========================================================" |
| ) |
| prev_stage_checkpoint = get_latest_hf_checkpoint_path(stage_config) |
| else: |
| config.check_and_update() |
| run_stage(config) |
|
|
| finally: |
| if dlc: |
| stop_ray_cluster(namespace=cluster_namespace) |
|
|
|
|
| def studio(port: int = 8501): |
| from streamlit.web import cli as stcli |
|
|
| current_dir = Path(__file__).resolve().parent.parent |
| config_manager_path = os.path.join(current_dir, "manager", "config_manager.py") |
|
|
| sys.argv = [ |
| "streamlit", |
| "run", |
| config_manager_path, |
| "--server.port", |
| str(port), |
| "--server.fileWatcherType", |
| "none", |
| ] |
| sys.exit(stcli.main()) |
|
|
|
|
| def debug( |
| config_path: str, |
| module: str, |
| output_dir: str = "debug_output", |
| disable_overwrite: bool = False, |
| enable_profiling: bool = False, |
| port: int = 8502, |
| plugin_dir: str = None, |
| ): |
| """Debug a module.""" |
| if plugin_dir: |
| os.environ[PLUGIN_DIRS_ENV_VAR] = plugin_dir |
| load_plugins() |
| config = load_config(config_path) |
| config.mode = "explore" |
| config.check_and_update() |
| sys.path.insert(0, os.getcwd()) |
| config.ray_namespace = DEBUG_NAMESPACE |
| ray.init( |
| namespace=config.ray_namespace, |
| runtime_env={"env_vars": config.get_envs()}, |
| ignore_reinit_error=True, |
| ) |
| from trinity.common.models import create_debug_inference_model |
|
|
| if module == "inference_model": |
| asyncio.run(create_debug_inference_model(config)) |
|
|
| elif module == "workflow": |
| from trinity.explorer.workflow_runner import DebugWorkflowRunner |
|
|
| runner = DebugWorkflowRunner(config, output_dir, enable_profiling, disable_overwrite) |
| asyncio.run(runner.debug()) |
| elif module == "viewer": |
| from streamlit.web import cli as stcli |
|
|
| current_dir = Path(__file__).resolve().parent.parent |
| viewer_path = os.path.join(current_dir, "buffer", "viewer.py") |
| output_dir_abs = os.path.abspath(output_dir) |
| if output_dir_abs.endswith("/"): |
| output_dir_abs = output_dir_abs[:-1] |
| print(f"sqlite:///{output_dir_abs}/experiences.db") |
| sys.argv = [ |
| "streamlit", |
| "run", |
| viewer_path, |
| "--server.port", |
| str(port), |
| "--server.fileWatcherType", |
| "none", |
| "--", |
| "--db-url", |
| f"sqlite:///{output_dir_abs}/experiences.db", |
| "--table", |
| "debug_buffer", |
| "--tokenizer", |
| config.model.model_path, |
| ] |
| sys.exit(stcli.main()) |
| else: |
| raise ValueError( |
| f"Only support 'inference_model' and 'workflow' for debugging, got {module}" |
| ) |
|
|
|
|
| def convert(checkpoint_dir: str, base_model_dir: Optional[str] = None) -> None: |
| if "global_step_" in checkpoint_dir: |
| while not os.path.basename(checkpoint_dir).startswith("global_step_"): |
| checkpoint_dir = os.path.dirname(checkpoint_dir) |
| converter = Converter(base_model_dir) |
| converter.convert(checkpoint_dir) |
|
|
|
|
| def main() -> None: |
| """The main entrypoint.""" |
| parser = argparse.ArgumentParser() |
| subparsers = parser.add_subparsers(dest="command", required=True) |
|
|
| |
| run_parser = subparsers.add_parser("run", help="Run RFT process.") |
| run_parser.add_argument("--config", type=str, required=True, help="Path to the config file.") |
| run_parser.add_argument( |
| "--plugin-dir", |
| type=str, |
| default=None, |
| help="Path to the directory containing plugin modules.", |
| ) |
| run_parser.add_argument( |
| "--dlc", action="store_true", help="Specify when running in Aliyun PAI DLC." |
| ) |
|
|
| |
| studio_parser = subparsers.add_parser("studio", help="Run studio.") |
| studio_parser.add_argument( |
| "--port", type=int, default=8501, help="The port for Trinity-Studio." |
| ) |
|
|
| |
| debug_parser = subparsers.add_parser("debug", help="Debug the code.") |
| debug_parser.add_argument("--config", type=str, help="Path to the config file.") |
| debug_parser.add_argument( |
| "--module", |
| type=str, |
| choices=["inference_model", "workflow", "viewer"], |
| help="The module to start debugging, only support 'inference_model', 'workflow' and 'viewer' for now.", |
| ) |
| debug_parser.add_argument( |
| "--plugin-dir", |
| type=str, |
| default=None, |
| help="Path to the directory containing plugin modules.", |
| ) |
| debug_parser.add_argument( |
| "--output-dir", |
| type=str, |
| default="debug_output", |
| help="The output directory for debug files.", |
| ) |
| debug_parser.add_argument( |
| "--disable-overwrite", action="store_true", help="Disable overwriting the output directory." |
| ) |
| debug_parser.add_argument( |
| "--enable-profiling", |
| action="store_true", |
| help="Whether to use viztracer for workflow profiling.", |
| ) |
| debug_parser.add_argument( |
| "--output-file", |
| type=str, |
| default=None, |
| help="[DEPRECATED] Please use --output-dir instead.", |
| ) |
| debug_parser.add_argument( |
| "--port", |
| type=int, |
| default=8502, |
| help="The port for Experience Viewer.", |
| ) |
|
|
| convert_parser = subparsers.add_parser( |
| "convert", help="Convert checkpoint to huggingface format." |
| ) |
| convert_parser.add_argument( |
| "--checkpoint-dir", |
| type=str, |
| required=True, |
| help="The path to the checkpoint directory.", |
| ) |
| convert_parser.add_argument( |
| "--base-model-dir", |
| type=str, |
| default=None, |
| help="The path to the base model.", |
| ) |
|
|
| args = parser.parse_args() |
| if args.command == "run": |
| |
| run(args.config, args.dlc, args.plugin_dir) |
| elif args.command == "studio": |
| studio(args.port) |
| elif args.command == "debug": |
| debug( |
| args.config, |
| args.module, |
| args.output_dir, |
| args.disable_overwrite, |
| args.enable_profiling, |
| args.port, |
| args.plugin_dir, |
| ) |
| elif args.command == "convert": |
| convert(args.checkpoint_dir, args.base_model_dir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
