Fix/chatglm6b 2 (#354)

This commit is contained in:
tastelikefeet
2023-07-04 01:58:57 +08:00
committed by GitHub
parent 45cf0035f4
commit 08c71f1f3d
3 changed files with 4 additions and 6 deletions

View File

@@ -196,7 +196,7 @@ tokenizer = ChatGLMTokenizer.from_pretrained(model_dir, trust_remote_code=True)
device_map_kwargs = {}
device_kwargs = {}
if args.use_lora != 0:
device_kwargs['device_map'] = 'auto'
device_map_kwargs['device_map'] = 'auto'
# No placement for model, leave the model to `device_map`
device_kwargs['device'] = 'cpu'
model = Model.from_pretrained(

View File

@@ -931,7 +931,6 @@ class ChatGLMModel(ChatGLMPreTrainedModel):
self.num_attention_heads,
self.hidden_size // self.num_attention_heads)
# seq_len, b, nh, hidden_size
print('#########################:', past_key_values.device)
past_key_values = self.dropout(past_key_values)
past_key_values = past_key_values.permute([2, 1, 0, 3, 4]).split(2)
# past_key_values = [(v[0], v[1]) for v in past_key_values]

View File

@@ -197,8 +197,8 @@ class ChatGLM6bTextGenerationPipeline(Pipeline):
if isinstance(model, str):
model_dir = snapshot_download(
model) if not os.path.exists(model) else model
config = ChatGLMConfig.from_pretrained(model_dir)
model = ChatGLMForConditionalGeneration(config).half()
model = ChatGLMForConditionalGeneration.from_pretrained(
model_dir).half()
if torch.cuda.is_available():
model = model.cuda()
if quantization_bit is not None:
@@ -235,8 +235,7 @@ class ChatGLM6bV2TextGenerationPipeline(Pipeline):
if isinstance(model, str):
model_dir = snapshot_download(
model) if not os.path.exists(model) else model
config = ChatGLM2Config.from_pretrained(model_dir)
model = ChatGLM2ForConditionalGeneration(config)
model = ChatGLM2ForConditionalGeneration.from_pretrained(model_dir)
if torch.cuda.is_available():
model = model.cuda()
if quantization_bit is not None: