Views
No views yet
1import torch
2from datasets import load_dataset
3from transformers import AutoTokenizer
4
5from llmcompressor.transformers import SparseAutoModelForCausalLM, oneshot
6from llmcompressor.transformers.compression.helpers import (
7 calculate_offload_device_map,
8 custom_offload_device_map,
9)
10
11recipe = """
12quant_stage:
13 quant_modifiers:
14 QuantizationModifier:
15 ignore: ["lm_head"]
16 config_groups:
17 group_0:
18 weights:
19 num_bits: 8
20 type: float
21 strategy: tensor
22 dynamic: false
23 symmetric: true
24 input_activations:
25 num_bits: 8
26 type: float
27 strategy: tensor
28 dynamic: false
29 symmetric: true
30 targets: ["Linear"]
31"""
32
33model_stub = "meta-llama/Meta-Llama-3.1-405B"
34model_name = model_stub.split("/")[-1]
35
36device_map = calculate_offload_device_map(
37 model_stub, reserve_for_hessians=False, num_gpus=8, torch_dtype=torch.float16
38)
39
40model = SparseAutoModelForCausalLM.from_pretrained(
41 model_stub, torch_dtype=torch.float16, device_map=device_map
42)
43tokenizer = AutoTokenizer.from_pretrained(model_stub)
44
45output_dir = f"./{model_name}-FP8"
46
47DATASET_ID = "HuggingFaceH4/ultrachat_200k"
48DATASET_SPLIT = "train_sft"
49NUM_CALIBRATION_SAMPLES = 512
50MAX_SEQUENCE_LENGTH = 4096
51
52ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
53ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
54
55def preprocess(example):
56 return {
57 "text": tokenizer.apply_chat_template(
58 example["messages"],
59 tokenize=False,
60 )
61 }
62
63ds = ds.map(preprocess)
64
65def tokenize(sample):
66 return tokenizer(
67 sample["text"],
68 padding=False,
69 max_length=MAX_SEQUENCE_LENGTH,
70 truncation=True,
71 add_special_tokens=False,
72 )
73
74ds = ds.map(tokenize, remove_columns=ds.column_names)
75
76oneshot(
77 model=model,
78 output_dir=output_dir,
79 dataset=ds,
80 recipe=recipe,
81 max_seq_length=MAX_SEQUENCE_LENGTH,
82 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
83 save_compressed=True,
84)| Benchmark | Meta-Llama-3.1-405B | Meta-Llama-3.1-405B-FP8(this model) | Recovery |
| MMLU (5-shot) | * | 84.72 | * |
| ARC Challenge (0-shot) | 95.99 | 95.82 | 99.82% |
| GSM-8K (5-shot, strict-match) | 88.10 | 87.94 | 99.82% |
| Hellaswag (10-shot) | 90.02 | 89.14 | 99.02% |
| Winogrande (5-shot) | 87.61 | 86.42 | 98.64% |
| TruthfulQA (0-shot) | 49.83 | 47.93 | 96.19% |
| Average | * | 82.00 | 98.70% |
lm_eval \
--model vllm \
--model_args pretrained="neuralmagic/Meta-Llama-3.1-405B-FP8",dtype=auto,add_bos_token=True,max_model_len=4096,tensor_parallel_size=8 \
--tasks mmlu \
--num_fewshot 5 \
--batch_size autolm_eval \
--model vllm \
--model_args pretrained="neuralmagic/Meta-Llama-3.1-405B-FP8",dtype=auto,add_bos_token=True,max_model_len=4096,tensor_parallel_size=8 \
--tasks arc_challenge_llama_3.1_instruct \
--num_fewshot 25 \
--batch_size autolm_eval \
--model vllm \
--model_args pretrained="neuralmagic/Meta-Llama-3.1-405B-FP8",dtype=auto,add_bos_token=True,max_model_len=4096,tensor_parallel_size=8 \
--tasks gsm8k \
--num_fewshot 5 \
--batch_size autolm_eval \
--model vllm \
--model_args pretrained="neuralmagic/Meta-Llama-3.1-405B-FP8",dtype=auto,add_bos_token=True,max_model_len=4096,tensor_parallel_size=8 \
--tasks hellaswag \
--num_fewshot 10 \
--batch_size autolm_eval \
--model vllm \
--model_args pretrained="neuralmagic/Meta-Llama-3.1-405B-FP8",dtype=auto,add_bos_token=True,max_model_len=4096,tensor_parallel_size=8 \
--tasks winogrande \
--num_fewshot 5 \
--batch_size autolm_eval \
--model vllm \
--model_args pretrained="neuralmagic/Meta-Llama-3.1-405B-FP8",dtype=auto,add_bos_token=True,max_model_len=4096,tensor_parallel_size=8 \
--tasks truthfulqa_mc \
--num_fewshot 0 \
--batch_size auto