File size: 3,319 Bytes
0c51b93 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import argparse
import os
os.environ["TRANSFORMERS_NO_COMPILE"] = "1"
from accelerate import Accelerator
from sotopia_rl import SotopiaGRPOTrainer
if __name__ == '__main__':
parser = argparse.ArgumentParser(description="Train a model with GRPO using a reward model.")
parser.add_argument("--model_name", type=str, default="/data/models/gemma-2-2b-it",
help="Path to the model")
parser.add_argument("--per_device_train_batch_size", type=int, default=1,
help="Batch size per device for training")
parser.add_argument("--per_device_eval_batch_size", type=int, default=1,
help="Batch size per device for evaluation")
parser.add_argument("--num_train_epochs", type=int, default=3,
help="Number of training epochs")
parser.add_argument("--num_grpo_epochs", type=int, default=4,
help="Number of GRPO epochs per update")
parser.add_argument("--learning_rate", type=float, default=5e-6,
help="Learning rate for optimizer")
parser.add_argument("--max_length", type=int, default=4096,
help="Maximum length of input sequences")
parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
help="Number of steps to accumulate gradients before performing an update")
parser.add_argument("--val_ratio", type=float, default=0.05,
help="Ratio of validation data")
parser.add_argument("--response_length", type=int, default=128,
help="Maximum length of generated responses")
parser.add_argument("--num_generations", type=int, default=4,
help="Number of generations for GRPO")
parser.add_argument("--beta", type=float, default=0.04,
help="Beta parameter for GRPO")
# Adapter parameters
parser.add_argument("--policy_adapter_path", type=str, default=None,
help="Path to policy model adapter")
parser.add_argument("--reward_adapter_path", type=str, default=None,
help="Path to reward model adapter")
# Data and checkpoint paths
parser.add_argument("--grpo_data_path", type=str, required=True,
help="Path to the reward data file")
parser.add_argument("--template_path", type=str, required=True,
help="Path to the Jinja template file")
parser.add_argument("--output_dir", type=str, default="checkpoints",
help="Directory to save the best LoRA checkpoint")
parser.add_argument("--save_steps", type=int, default=50,
help="Number of steps between saving checkpoints")
# Logging parameters
parser.add_argument("--wandb_project", type=str, default="grpo-model-training",
help="Wandb project name")
parser.add_argument("--wandb_run_name", type=str, default=None,
help="Wandb run name")
parser.add_argument("--use_lora_train_grpo", action="store_true",
help="Use LoRA for training GRPO")
args = parser.parse_args()
accelerator = Accelerator()
trainer = SotopiaGRPOTrainer(args, accelerator)
trainer.train()
|