Test: Fix Llama injection test and use a smaller model

The original Llama injection test was using a large 7B model, which is not suitable for CI.

This commit modifies the test to:
1. Use a manually created, small Llama model with a standard configuration.
2. This new approach successfully reproduces the original bug and passes after the fix.

Signed-off-by: huanyuqu <[email protected]>

Author: huanyuqu

tests/unit/inference/test_inference.py

@@ -24,6 +24,7 @@
 from transformers import pipeline
 from transformers.models.t5.modeling_t5 import T5Block
 from transformers.models.roberta.modeling_roberta import RobertaLayer
+from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer
 
 from deepspeed.accelerator import get_accelerator
 from deepspeed.git_version_info import torch_info
@@ -554,7 +555,7 @@ def test(self, model_w_task, injection_policy, query, inf_kwargs, assert_fn, dty
 
 
 @pytest.mark.seq_inference
-@pytest.mark.parametrize("model_w_task", [("meta-llama/Llama-2-7b-hf", "text-generation")], ids=["llama"])
+@pytest.mark.parametrize("model_w_task", [("hf-internal-testing/tiny-random-LlamaForCausalLM", "text-generation")], ids=["llama"])
 @pytest.mark.parametrize("dtype", [torch.half], ids=["fp16"])
 class TestLlamaInjection(DistributedTest):
     world_size = 1
@@ -571,12 +572,31 @@ def test(self, model_w_task, dtype, query, inf_kwargs, assert_fn):
             pytest.skip("This op had not been implemented on this system.", allow_module_level=True)
 
         model, task = model_w_task
+
+        tokenizer = LlamaTokenizer.from_pretrained(model)
+        config = LlamaConfig(
+            vocab_size=32000,
+            hidden_size=4096,
+            intermediate_size=11008,
+            num_hidden_layers=1,
+            num_attention_heads=32,
+            max_position_embeddings=2048,
+            initializer_range=0.02,
+            rms_norm_eps=1e-5,
+            use_cache=True,
+            pad_token_id=0,
+            bos_token_id=1,
+            eos_token_id=2,
+            tie_word_embeddings=True
+        )
+        model = LlamaForCausalLM(config)
 
         local_rank = int(os.getenv("LOCAL_RANK", "0"))
         device = torch.device(get_accelerator().device_name(local_rank))
 
         pipe = pipeline(task,
                         model=model,
+                        tokenizer=tokenizer,
                         device=torch.device("cpu"),
                         model_kwargs={"low_cpu_mem_usage": True},
                         framework="pt")
@@ -597,7 +617,7 @@ def test(self, model_w_task, dtype, query, inf_kwargs, assert_fn):
             )
             check_injection(pipe.model)
         except AttributeError as e:
-            if "'LlamaAttention' object has no attribute 'num_heads'" in e:
+            if "'LlamaAttention' object has no attribute 'num_heads'" in str(e):
                 pytest.skip("Skipping due to transformers version compatibility issue with self-attention")
             raise e