Update test_ddp.py

test_ddp.py

@@ -11,7 +11,8 @@ from torch.nn.parallel import DistributedDataParallel
 from peft import PeftModel
 from torch.utils.data import DataLoader
 from tqdm import tqdm
-from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+# from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 
 from utils import *
 from collator import TestCollator
@@ -33,12 +34,12 @@ def test_ddp(args):
     device_map = {"": local_rank}
     device = torch.device("cuda",local_rank)
 
-    tokenizer = LlamaTokenizer.from_pretrained(args.ckpt_path)
-    args.lora=True
+    tokenizer = AutoTokenizer.from_pretrained(args.ckpt_path)
+    #args.lora=True
     if args.lora:
-        model = LlamaForCausalLM.from_pretrained(
+        model = AutoModelForCausalLM.from_pretrained(
             args.base_model,
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             low_cpu_mem_usage=True,
             device_map=device_map,
         )
@@ -46,13 +47,13 @@ def test_ddp(args):
         model = PeftModel.from_pretrained(
             model,
             args.ckpt_path,
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             device_map=device_map,
         )
     else:
-        model = LlamaForCausalLM.from_pretrained(
+        model = AutoModelForCausalLM.from_pretrained(
             args.ckpt_path,
-            torch_dtype=torch.float16,
+            torch_dtype=torch.bfloat16,
             low_cpu_mem_usage=True,
             device_map=device_map,
         )