1import multiprocessing
2
3from datasets import load_dataset
4from tqdm.rich import tqdm
5from transformers import AutoTokenizer, AutoModelForCausalLM
6from trl import (
7 ModelConfig,
8 SFTTrainer,
9 get_peft_config,
10 get_quantization_config,
11 get_kbit_device_map,
12 SFTConfig,
13 ScriptArguments
14)
15from trl.commands.cli_utils import TrlParser
16
17tqdm.pandas()
18
19if __name__ == "__main__":
20 parser = TrlParser((ScriptArguments, SFTConfig, ModelConfig))
21 args, training_args, model_config = parser.parse_args_and_config()
22
23 quantization_config = get_quantization_config(model_config)
24 model_kwargs = dict(
25 revision=model_config.model_revision,
26 trust_remote_code=model_config.trust_remote_code,
27 attn_implementation=model_config.attn_implementation,
28 torch_dtype=model_config.torch_dtype,
29 use_cache=False if training_args.gradient_checkpointing else True,
30 device_map=get_kbit_device_map() if quantization_config is not None else None,
31 quantization_config=quantization_config,
32 )
33
34 model = AutoModelForCausalLM.from_pretrained(model_config.model_name_or_path,
35 **model_kwargs)
36 tokenizer = AutoTokenizer.from_pretrained(
37 model_config.model_name_or_path, trust_remote_code=model_config.trust_remote_code, use_fast=True
38 )
39 tokenizer.pad_token = tokenizer.eos_token
40
41 train_dataset = load_dataset(args.dataset_name,
42 split=args.dataset_train_split,
43 num_proc=multiprocessing.cpu_count())
44
45 trainer = SFTTrainer(
46 model=model,
47 args=training_args,
48 train_dataset=train_dataset,
49 processing_class=tokenizer,
50 peft_config=get_peft_config(model_config),
51 )
52
53 trainer.train()
54
55 trainer.save_model(training_args.output_dir)
1from vllm import LLM
2from datasets import load_dataset
3from vllm.sampling_params import SamplingParams
4from transformers import AutoTokenizer
5
6MODEL_PATH = "autodl-tmp/saves/Qwen2.5-1.5B-ultrachat200k"
7
8model = LLM(MODEL_PATH,
9 tensor_parallel_size=1,
10 dtype='bfloat16')
11tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
12
13input = tokenizer.apply_chat_template([{"role": "user", "content": "Which province is Shenyang in?"}],
14 tokenize=False,
15 add_generation_prompt=True)
16sampling_params = SamplingParams(max_tokens=1024,
17 temperature=0.7,
18 logprobs=1,
19 stop_token_ids=[tokenizer.eos_token_id])
20
21vllm_generations = model.generate(input,
22 sampling_params)
23
24print(vllm_generations[0].outputs[0].text)
25# print result: Shenyang is in Liaoning province, China.
