[shardformer] adapted T5 and LLaMa test to use kit (#4049)

* [shardformer] adapted T5 and LLaMa test to use kit

* polish code

tests/kit/model_zoo/transformers/llama.py

@@ -0,0 +1,76 @@
+import torch
+import transformers
+
+from ..registry import ModelAttribute, model_zoo
+
+try:
+    from transformers import LlamaConfig, LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel
+    HAS_LLAMA = True
+except ImportError:
+    HAS_LLAMA = False
+
+if HAS_LLAMA:
+    # ===============================
+    # Register LLaMA
+    # ===============================
+
+    def data_gen():
+        # the input ids are corresponding to the sentence
+        # 'Hello, my dog is cute'
+        #
+        # the code is give below:
+        # -----------------------------------
+        # from transformers import LlamaTokenizerFast
+        # tokenizer = LlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
+        # input = 'Hello, my dog is cute'
+        # tokenized_input = tokenizer(input, return_tensors='pt').to('cuda')
+        # -----------------------------------
+
+        input_ids = torch.Tensor([[1, 15043, 29892, 590, 11203, 338, 274, 1082]]).long()
+        attention_mask = torch.Tensor([[1, 1, 1, 1, 1, 1, 1, 1]]).long()
+        return dict(input_ids=input_ids, attention_mask=attention_mask)
+
+    # label is needed for casual lm
+    def data_gen_for_casual_lm():
+        data = data_gen()
+        labels = data['input_ids'].clone()
+        data['labels'] = labels
+        return data
+
+    # transform the output to a dict
+    output_transform_fn = lambda x: x
+
+    # function to get the loss
+    loss_fn = lambda output: output.last_hidden_state.mean()
+    loss_fn_for_casual_lm = lambda output: output.loss
+    loss_fn_for_seq_classification = lambda output: output.logits.mean()
+
+    config = LlamaConfig(num_hidden_layers=4,
+                         hidden_size=128,
+                         intermediate_size=256,
+                         num_attention_heads=4,
+                         max_position_embeddings=128,
+                         num_labels=16)
+
+    # register the following models
+    # transformers.LlamaModel,
+    # transformers.LlamaForCausalLM,
+    # transformers.LlamaForSequenceClassification,
+    model_zoo.register(name='transformers_llama',
+                       model_fn=lambda: transformers.LlamaModel(config),
+                       data_gen_fn=data_gen,
+                       output_transform_fn=output_transform_fn,
+                       loss_fn=loss_fn,
+                       model_attribute=ModelAttribute(has_control_flow=True))
+    model_zoo.register(name='transformers_llama_for_casual_lm',
+                       model_fn=lambda: transformers.LlamaForCausalLM(config),
+                       data_gen_fn=data_gen_for_casual_lm,
+                       output_transform_fn=output_transform_fn,
+                       loss_fn=loss_fn_for_casual_lm,
+                       model_attribute=ModelAttribute(has_control_flow=True))
+    model_zoo.register(name='transformers_llama_for_sequence_classification',
+                       model_fn=lambda: transformers.LlamaForSequenceClassification(config),
+                       data_gen_fn=data_gen,
+                       output_transform_fn=output_transform_fn,
+                       loss_fn=loss_fn_for_seq_classification,
+                       model_attribute=ModelAttribute(has_control_flow=True))