add test code

test_modeling.py

@@ -0,0 +1,31 @@
+# torchrun --nproc_per_node=2 test_modeling.py --local_path {local_path}
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_path", type=str, default="petil777/srv1_parallel")
+parser.add_argument("--revision", type=str, default=None)
+parser.add_argument("--local_path", type=str, required=True)
+args = parser.parse_args()
+
+model_path = "petil777/srv1_parallel"
+
+# Distributed.launch will occur inside.
+model = AutoModelForCausalLM.from_pretrained(args.model_path,local_path=args.local_path, 
+                                             revision=args.revision, trust_remote_code=True)
+
+model.eval()
+tokenizer = model.tokenizer
+
+rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+model = model.to(f"cuda:{rank}")
+
+
+input_str="apple is red and banana is"
+input_dict = tokenizer(input_str, return_tensors="pt") #input_ids, attention_mask
+input_ids= input_dict.input_ids
+input_ids=input_ids.to(f"cuda:{rank}")
+
+out_tensor = model.generate(input_ids, top_k=0,return_dict_in_generate=True,output_scores=True,output_hidden_states=True)
+if rank == 0:
+    print(tokenizer.decode(out_tensor.sequences[0]))