From 31555506e22a33c19812c8b694aef9e92505cb8b Mon Sep 17 00:00:00 2001 From: helloyongyang Date: Fri, 11 Oct 2024 17:13:28 +0800 Subject: [PATCH] support llama3.1 llama3.2 & transformers 4.45.2 version & fix two devices bug --- llmc/models/falcon.py | 5 +++-- llmc/models/llama.py | 5 +++-- llmc/models/phi.py | 5 +++-- llmc/models/qwen2.py | 5 +++-- llmc/models/stablelm.py | 5 +++-- requirements/runtime.txt | 2 +- 6 files changed, 16 insertions(+), 11 deletions(-) diff --git a/llmc/models/falcon.py b/llmc/models/falcon.py index 5bac37772..e13019a3f 100644 --- a/llmc/models/falcon.py +++ b/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 diff --git a/llmc/models/llama.py b/llmc/models/llama.py index 665bbf6d0..04a7e1d3f 100644 --- a/llmc/models/llama.py +++ b/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'] diff --git a/llmc/models/phi.py b/llmc/models/phi.py index 95a2232fa..a7e17144b 100644 --- a/llmc/models/phi.py +++ b/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'] diff --git a/llmc/models/qwen2.py b/llmc/models/qwen2.py index 42a5dadb4..6c3271fac 100644 --- a/llmc/models/qwen2.py +++ b/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'] diff --git a/llmc/models/stablelm.py b/llmc/models/stablelm.py index a39bc1b14..125c870a4 100644 --- a/llmc/models/stablelm.py +++ b/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'] diff --git a/requirements/runtime.txt b/requirements/runtime.txt index 43244d018..4a5fd655c 100644 --- a/requirements/runtime.txt +++ b/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