support llama3.1 llama3.2 & transformers 4.45.2 version & fix two devices bug (#120)

Author: helloyongyang

llmc/models/falcon.py

@@ -13,15 +13,16 @@ def find_blocks(self):
 
     def find_embed_layers(self):
         self.word_embeddings = self.model.transformer.word_embeddings
+        self.rotary_emb = self.model.model.rotary_emb
 
     def find_block_name(self):
         self.block_name_prefix = 'model.transformer.h'
 
     def get_embed_layers(self):
-        return [self.word_embeddings]
+        return [self.word_embeddings, self.rotary_emb]
 
     def get_layers_except_blocks(self):
-        return [self.word_embeddings, self.model.transformer.ln_f]
+        return [self.word_embeddings, self.rotary_emb, self.model.transformer.ln_f]
 
     def has_bias(self):
         return False

llmc/models/llama.py

@@ -13,13 +13,14 @@ def find_blocks(self):
 
     def find_embed_layers(self):
         self.embed_tokens = self.model.model.embed_tokens
+        self.rotary_emb = self.model.model.rotary_emb
 
     def find_block_name(self):
         self.block_name_prefix = 'model.layers'
         self.pairs = {'q_proj': 'qkv', 'o_proj': 'out', 'up_proj': 'fc1'}
 
     def get_embed_layers(self):
-        return [self.embed_tokens]
+        return [self.embed_tokens, self.rotary_emb]
 
     def get_head_layers(self):
         return [self.model.lm_head]
@@ -28,7 +29,7 @@ def get_pre_head_layernorm_layers(self):
         return [self.model.model.norm]
 
     def get_layers_except_blocks(self):
-        return [self.embed_tokens, self.model.model.norm, self.model.lm_head]
+        return [self.embed_tokens, self.rotary_emb, self.model.model.norm, self.model.lm_head] # noqa
 
     def skip_layer_name(self):
         return ['lm_head']

llmc/models/phi.py

@@ -13,13 +13,14 @@ def find_blocks(self):
 
     def find_embed_layers(self):
         self.embed_tokens = self.model.model.embed_tokens
+        self.rotary_emb = self.model.model.rotary_emb
 
     def find_block_name(self):
         self.block_name_prefix = 'model.layers'
         self.pairs = {'q_proj': 'qkv', 'o_proj': 'out', 'up_proj': 'fc1'}
 
     def get_embed_layers(self):
-        return [self.embed_tokens]
+        return [self.embed_tokens, self.rotary_emb]
 
     def get_head_layers(self):
         return [self.model.lm_head]
@@ -28,7 +29,7 @@ def get_pre_head_layernorm_layers(self):
         return [self.model.model.final_layernorm]
 
     def get_layers_except_blocks(self):
-        return [self.embed_tokens, self.model.model.final_layernorm, self.model.lm_head]
+        return [self.embed_tokens, self.rotary_emb, self.model.model.final_layernorm, self.model.lm_head] # noqa
 
     def skip_layer_name(self):
         return ['lm_head']

llmc/models/qwen2.py

@@ -13,13 +13,14 @@ def find_blocks(self):
 
     def find_embed_layers(self):
         self.embed_tokens = self.model.model.embed_tokens
+        self.rotary_emb = self.model.model.rotary_emb
 
     def find_block_name(self):
         self.block_name_prefix = 'model.layers'
         self.pairs = {'q_proj': 'qkv', 'o_proj': 'out', 'up_proj': 'fc1'}
 
     def get_embed_layers(self):
-        return [self.embed_tokens]
+        return [self.embed_tokens, self.rotary_emb]
 
     def get_head_layers(self):
         return [self.model.lm_head]
@@ -28,7 +29,7 @@ def get_pre_head_layernorm_layers(self):
         return [self.model.model.norm]
 
     def get_layers_except_blocks(self):
-        return [self.embed_tokens, self.model.model.norm, self.model.lm_head]
+        return [self.embed_tokens, self.rotary_emb, self.model.model.norm, self.model.lm_head] # noqa
 
     def skip_layer_name(self):
         return ['lm_head']

llmc/models/stablelm.py

@@ -10,6 +10,7 @@ def __init__(self, model_path, torch_dtype, device_map=None, use_cache=False):
 
     def find_blocks(self):
         self.blocks = self.model.model.layers
+        self.rotary_emb = self.model.model.rotary_emb
 
     def find_embed_layers(self):
         self.embed_tokens = self.model.model.embed_tokens
@@ -19,7 +20,7 @@ def find_block_name(self):
         self.pairs = {'q_proj': 'qkv', 'o_proj': 'out', 'up_proj': 'fc1'}
 
     def get_embed_layers(self):
-        return [self.embed_tokens]
+        return [self.embed_tokens, self.rotary_emb]
 
     def get_head_layers(self):
         return [self.model.lm_head]
@@ -28,7 +29,7 @@ def get_pre_head_layernorm_layers(self):
         return [self.model.model.norm]
 
     def get_layers_except_blocks(self):
-        return [self.embed_tokens, self.model.model.norm, self.model.lm_head]
+        return [self.embed_tokens, self.rotary_emb, self.model.model.norm, self.model.lm_head] # noqa
 
     def skip_layer_name(self):
         return ['lm_head']

requirements/runtime.txt

@@ -1,7 +1,7 @@
 torch>=2.1.0
 pillow
 loguru
-transformers>=4.41.2
+transformers==4.45.2
 huggingface-hub
 sentencepiece
 protobuf