Quantize the save the model as TRT-LLM checkpoint to output_dir
(hf_model_dir: str,
output_dir: str,
config: LLaMAConfig,
device: str = 'cuda',
calib_dataset: str = 'cnn_dailymail',
trust_remote_code: bool = True,
calib_batches: int = 512,
calib_max_seq_length: int = 512)
| 1100 | |
| 1101 | |
| 1102 | def quantize(hf_model_dir: str, |
| 1103 | output_dir: str, |
| 1104 | config: LLaMAConfig, |
| 1105 | device: str = 'cuda', |
| 1106 | calib_dataset: str = 'cnn_dailymail', |
| 1107 | trust_remote_code: bool = True, |
| 1108 | calib_batches: int = 512, |
| 1109 | calib_max_seq_length: int = 512): |
| 1110 | ''' |
| 1111 | Quantize the save the model as TRT-LLM checkpoint to output_dir |
| 1112 | ''' |
| 1113 | os.makedirs(output_dir, exist_ok=True) |
| 1114 | config.to_json_file(os.path.join(output_dir, 'config.json')) |
| 1115 | |
| 1116 | mapping = config.mapping |
| 1117 | assert mapping.rank == 0, "quantize should be called at rank 0 only" |
| 1118 | |
| 1119 | quant_config = config.quantization |
| 1120 | use_smooth_quant = quant_config._use_plugin_sq |
| 1121 | int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 |
| 1122 | |
| 1123 | assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" |
| 1124 | assert hf_model_dir is not None |
| 1125 | ## only load and call smooth quant routine once for all ranks |
| 1126 | hf_config = AutoConfig.from_pretrained(hf_model_dir, |
| 1127 | trust_remote_code=trust_remote_code) |
| 1128 | assert "llava" not in hf_config.model_type, "Smooth quant llava/vila/llava_next is not supported yet" |
| 1129 | hf_model = AutoModelForCausalLM.from_pretrained( |
| 1130 | hf_model_dir, |
| 1131 | device_map='auto' if device != 'cpu' else 'cpu', |
| 1132 | dtype='auto' if not use_smooth_quant else torch.float16, |
| 1133 | trust_remote_code=trust_remote_code) |
| 1134 | |
| 1135 | os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( |
| 1136 | "TOKENIZERS_PARALLELISM", "false") |
| 1137 | tokenizer = AutoTokenizer.from_pretrained( |
| 1138 | hf_model_dir, |
| 1139 | trust_remote_code=trust_remote_code, |
| 1140 | use_fast=False, |
| 1141 | padding_side='left') |
| 1142 | |
| 1143 | dataset = load_calib_dataset(calib_dataset) |
| 1144 | |
| 1145 | if calib_batches == -1: # use the whole dataset if calib_batches is -1 |
| 1146 | calib_batches = len(dataset) |
| 1147 | |
| 1148 | act_range = capture_activation_range(hf_model, |
| 1149 | tokenizer, |
| 1150 | dataset, |
| 1151 | num_samples=calib_batches, |
| 1152 | seq_len=calib_max_seq_length) |
| 1153 | qkv_para, smoother = {}, {} |
| 1154 | if use_smooth_quant: |
| 1155 | smooth_llama_model(hf_model, act_range, quant_config.smoothquant_val, |
| 1156 | qkv_para, smoother) |
| 1157 | |
| 1158 | for rank in range(mapping.world_size): |
| 1159 | # To avoid changing the mapping arg in-place, also the given mapping from caller is rank agnostic, since quantize is called from only one rank |
nothing calls this directly
no test coverage detected