Update finetune.py

finetune.py

@@ -7,7 +7,7 @@ from typing import List
 import torch
 import transformers
 
-from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig, AutoTokenizer, AutoModelForCausalLM, AutoConfig
 
 from utils import *
 from collator import Collator
@@ -27,13 +27,13 @@ def train(args):
     if ddp:
         device_map = {"": local_rank}
 
-    config = LlamaConfig.from_pretrained(args.base_model)
-    tokenizer = LlamaTokenizer.from_pretrained(
+    config = AutoConfig.from_pretrained(args.base_model)
+    tokenizer = AutoTokenizer.from_pretrained(
         args.base_model,
         model_max_length = args.model_max_length,
         padding_side="right",
     )
-    tokenizer.pad_token_id = 0
+    tokenizer.pad_token_id = tokenizer.eos_token_id
     gradient_checkpointing = True
 
     train_data, valid_data = load_datasets(args)
@@ -48,7 +48,7 @@ def train(args):
     collator = Collator(args, tokenizer)
 
 
-    model = LlamaForCausalLM.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         args.base_model,
         # torch_dtype=torch.float16,
         device_map=device_map,
@@ -85,7 +85,7 @@ def train(args):
             eval_steps=args.save_and_eval_steps,
             save_steps=args.save_and_eval_steps,
             output_dir=args.output_dir,
-            save_total_limit=5,
+            save_total_limit=20,
             load_best_model_at_end=True,
             deepspeed=args.deepspeed,
             ddp_find_unused_parameters=False if ddp else None,