Fix: Adapt Llama injection policy for newer transformers versions

The LlamaAttention module in recent transformers versions stores configuration
like 'num_heads' and 'rope_theta' in a 'config' object, rather than as
direct attributes.

This change updates the LlamaLayerPolicy to fetch these attributes from
the 'config' object first, falling back to the direct attribute to maintain
backward compatibility. This resolves the AttributeError during kernel
injection with newer transformers versions.

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

Author: huanyuqu

deepspeed/module_inject/containers/llama.py

@@ -34,7 +34,10 @@ def create_module(self, config=None):
         _config.rotate_half = True
         _config.rotate_every_two = False
         _config.rotary_dim = self.hidden_size // self.num_attention_heads
-        _config.rope_theta = self.policy.client_module.self_attn.rope_theta
+        if hasattr(self.policy.client_module.self_attn, 'config'):
+            _config.rope_theta = self.policy.client_module.self_attn.config.rope_theta
+        else:
+            _config.rope_theta = self.policy.client_module.self_attn.rope_theta
         self.module = DeepSpeedGPTInference(_config, mp_group=self.mp_group)
 
         return self.module
@@ -128,9 +131,13 @@ def __init__(self, client_module, inference=True):
             LLAMALayerPolicy._orig_layer_class = None
 
     def get_hidden_heads(self):
+        if hasattr(self.client_module.self_attn, 'config'):
+            num_heads = self.client_module.self_attn.config.num_attention_heads
+        else:
+            num_heads = self.client_module.self_attn.num_heads
         hidden_heads = (
             self.client_module.self_attn.q_proj.in_features,
-            self.client_module.self_attn.num_heads,
+            num_heads,
             self.client_module.input_layernorm.variance_epsilon,
             self.client_module.mlp.gate_proj.out_features,
         )