Views
No views yet
varlen attention 和 varlen RoPE算子,同时也探讨了批量推理时引入的填充token对输出的影响,以及如何通过设计兼容varlen的KV Cache类直接基于Huggingface GenerationMixin接口无缝切块填充推理和无填充变长序列推理1import os
2import torch
3
4from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
5
6model: AutoModelForCausalLM = AutoModelForCausalLM.from_pretrained("SophieA17/Sophie0-DPO", trust_remote_code=True)
7tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained("SophieA17/Sophie0-DPO", trust_remote_code=True)
8
9model = model.to(device="cuda:0", dtype=torch.bfloat16)
10inputs = [
11 "<s><user>Could you please introduce youself?</s>\n",
12 "<s><user>Where is the best place for traveling in summer?</s>\n"
13]
14
15input_ids = tokenizer(inputs, return_tensors="pt", padding=True, padding_side="left", return_token_type_ids=False).to(model.device)
16
17generation_config = GenerationConfig(
18 bos_token_id=tokenizer.bos_token_id,
19 eos_token_id=tokenizer.eos_token_id,
20 pad_token_id=tokenizer.pad_token_id,
21 max_new_tokens=1024,
22 do_sample=True,
23 top_k=20,
24 top_p=0.8,
25 temperature=0.8,
26 repeat_penalty=1.1,
27 use_cache=True
28)
29
30outputs = model.generate(**input_ids, use_varlen_inference=True, generation_config=generation_config)
31outputs = tokenizer.batch_decode(outputs, skip_special_tokens=False)