Buckets:
| #!/usr/bin/env python3 | |
| import argparse | |
| import os | |
| import sys | |
| import time | |
| from pathlib import Path | |
| sys.path.append(str(Path(__file__).parent.parent / "src")) | |
| from src.main import main as run_experiment | |
| from src.config import Config | |
| from src.utils.logging_utils import setup_logging | |
| from loguru import logger | |
| def parse_args(): | |
| parser = argparse.ArgumentParser(description="Train RL-LLM models") | |
| parser.add_argument( | |
| "--experiment_type", | |
| type=str, | |
| default="rlvr", | |
| choices=["rlvr", "rlhf", "comparison"], | |
| help="Type of experiment to run", | |
| ) | |
| parser.add_argument( | |
| "--config", type=str, default="config.yaml", help="Path to configuration file" | |
| ) | |
| parser.add_argument( | |
| "--num_epochs", type=int, default=10, help="Number of training epochs" | |
| ) | |
| parser.add_argument( | |
| "--batch_size", type=int, default=8, help="Batch size for training" | |
| ) | |
| parser.add_argument( | |
| "--learning_rate", type=float, default=1e-4, help="Learning rate" | |
| ) | |
| parser.add_argument( | |
| "--model_name", type=str, default="gpt2", help="Model name to use" | |
| ) | |
| parser.add_argument( | |
| "--output_dir", | |
| type=str, | |
| default="./results", | |
| help="Output directory for results", | |
| ) | |
| parser.add_argument( | |
| "--log_level", | |
| type=str, | |
| default="INFO", | |
| choices=["DEBUG", "INFO", "WARNING", "ERROR"], | |
| help="Logging level", | |
| ) | |
| parser.add_argument("--seed", type=int, default=42, help="Random seed") | |
| parser.add_argument( | |
| "--device", type=str, default="auto", help="Device to use (cuda, cpu, auto)" | |
| ) | |
| parser.add_argument( | |
| "--save_checkpoints", | |
| action="store_true", | |
| help="Save model checkpoints during training", | |
| ) | |
| parser.add_argument( | |
| "--resume_from", type=str, default=None, help="Resume training from checkpoint" | |
| ) | |
| return parser.parse_args() | |
| def create_config_from_args(args): | |
| config = Config() | |
| config.model.model_name = args.model_name | |
| config.model.device = args.device | |
| config.rl.num_epochs = args.num_epochs | |
| config.rl.batch_size = args.batch_size | |
| config.rl.learning_rate = args.learning_rate | |
| config.experiment.experiment_type = args.experiment_type | |
| config.experiment.output_dir = args.output_dir | |
| config.experiment.seed = args.seed | |
| config.experiment.save_checkpoints = args.save_checkpoints | |
| config.experiment.resume_from = args.resume_from | |
| return config | |
| def main(): | |
| args = parse_args() | |
| log_file = os.path.join(args.output_dir, "training.log") | |
| os.makedirs(args.output_dir, exist_ok=True) | |
| setup_logging(log_file, args.log_level) | |
| logger.info("Starting RL-LLM training") | |
| logger.info(f"Arguments: {args}") | |
| try: | |
| if os.path.exists(args.config): | |
| config = Config.from_yaml(args.config) | |
| logger.info(f"Loaded configuration from {args.config}") | |
| else: | |
| config = create_config_from_args(args) | |
| logger.info("Created configuration from command line arguments") | |
| config = create_config_from_args(args) | |
| os.makedirs(args.output_dir, exist_ok=True) | |
| config_path = os.path.join(args.output_dir, "config.yaml") | |
| config.to_yaml(config_path) | |
| logger.info(f"Saved configuration to {config_path}") | |
| start_time = time.time() | |
| if args.experiment_type == "rlvr": | |
| result = run_experiment("rlvr", config) | |
| elif args.experiment_type == "rlhf": | |
| result = run_experiment("rlhf", config) | |
| elif args.experiment_type == "comparison": | |
| result = run_experiment("comparison", config) | |
| else: | |
| raise ValueError(f"Unknown experiment type: {args.experiment_type}") | |
| end_time = time.time() | |
| training_time = end_time - start_time | |
| logger.info(f"Training completed in {training_time:.2f} seconds") | |
| logger.info(f"Results saved to: {args.output_dir}") | |
| if result: | |
| logger.info("Training Summary:") | |
| if "training_history" in result: | |
| final_epoch = result["training_history"][-1] | |
| logger.info( | |
| f"Final average reward: {final_epoch.get('avg_reward', 0):.3f}" | |
| ) | |
| logger.info( | |
| f"Final success rate: {final_epoch.get('success_rate', 0):.3f}" | |
| ) | |
| if "final_metrics" in result: | |
| metrics = result["final_metrics"] | |
| logger.info(f"Final evaluation metrics: {metrics}") | |
| except Exception as e: | |
| logger.error(f"Training failed: {e}") | |
| raise | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 4.76 kB
- Xet hash:
- 85f43766f60daa83f929893665d1a13dfa32a7caf1137198ca38d5d8a1b141af
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.