Quantize the save the model as TRT-LLM checkpoint to output_dir
(hf_model_dir: str,
output_dir: str,
config: ChatGLMConfig,
calib_dataset: str = 'cnn_dailymail',
device: str = 'auto',
trust_remote_code: bool = True)
| 663 | |
| 664 | |
| 665 | def quantize(hf_model_dir: str, |
| 666 | output_dir: str, |
| 667 | config: ChatGLMConfig, |
| 668 | calib_dataset: str = 'cnn_dailymail', |
| 669 | device: str = 'auto', |
| 670 | trust_remote_code: bool = True): |
| 671 | ''' |
| 672 | Quantize the save the model as TRT-LLM checkpoint to output_dir |
| 673 | ''' |
| 674 | os.makedirs(output_dir, exist_ok=True) |
| 675 | config.to_json_file(os.path.join(output_dir, 'config.json')) |
| 676 | |
| 677 | mapping = config.mapping |
| 678 | assert mapping.rank == 0, "quantize should be called at rank 0 only" |
| 679 | |
| 680 | quant_config = config.quantization |
| 681 | use_smooth_quant = quant_config._use_plugin_sq |
| 682 | int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 |
| 683 | |
| 684 | assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" |
| 685 | assert hf_model_dir is not None |
| 686 | ## only load and call smooth quant routine once for all ranks |
| 687 | if config.chatglm_version == 'glm': |
| 688 | device_map = 'cuda' if device != "cpu" else 'cpu' |
| 689 | else: |
| 690 | device_map = 'auto' if device != "cpu" else 'cpu' |
| 691 | hf_model = AutoModel.from_pretrained( |
| 692 | hf_model_dir, |
| 693 | trust_remote_code=trust_remote_code, |
| 694 | dtype='auto' if config.chatglm_version != 'glm' else getattr( |
| 695 | torch, config.dtype), |
| 696 | device_map=device_map) |
| 697 | |
| 698 | os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( |
| 699 | "TOKENIZERS_PARALLELISM", "false") |
| 700 | tokenizer = AutoTokenizer.from_pretrained( |
| 701 | hf_model_dir, |
| 702 | trust_remote_code=trust_remote_code, |
| 703 | ) |
| 704 | dataset = load_calib_dataset(calib_dataset) |
| 705 | |
| 706 | act_range = capture_activation_range(hf_model, |
| 707 | tokenizer, |
| 708 | dataset, |
| 709 | num_samples=64) |
| 710 | smoother = {} |
| 711 | if use_smooth_quant: |
| 712 | smooth_chatglm_model(hf_model, act_range, quant_config.smoothquant_val, |
| 713 | smoother) |
| 714 | |
| 715 | for rank in range(mapping.world_size): |
| 716 | # 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 |
| 717 | config = copy.deepcopy(config) |
| 718 | config.set_rank(rank) |
| 719 | weights = load_weights_from_hf_model( |
| 720 | hf_model, |
| 721 | config=config, |
| 722 | act_range=act_range, |
nothing calls this directly
no test coverage detected