(hf_model_dir: str,
output_dir: str,
config: GPTConfig,
device: str = 'cuda',
calib_dataset: str = 'cnn_dailymail',
trust_remote_code: bool = True)
| 857 | |
| 858 | |
| 859 | def quantize(hf_model_dir: str, |
| 860 | output_dir: str, |
| 861 | config: GPTConfig, |
| 862 | device: str = 'cuda', |
| 863 | calib_dataset: str = 'cnn_dailymail', |
| 864 | trust_remote_code: bool = True): |
| 865 | os.makedirs(output_dir, exist_ok=True) |
| 866 | config.to_json_file(os.path.join(output_dir, 'config.json')) |
| 867 | |
| 868 | mapping = config.mapping |
| 869 | assert mapping.rank == 0, "quantize should be called at rank 0 only" |
| 870 | |
| 871 | quant_config = config.quantization |
| 872 | use_smooth_quant = quant_config._use_plugin_sq |
| 873 | int8_kv_cache = quant_config.kv_cache_quant_algo == QuantAlgo.INT8 |
| 874 | |
| 875 | assert use_smooth_quant or int8_kv_cache, "Call from_hugging_face when there is no quantization" |
| 876 | assert hf_model_dir is not None |
| 877 | ## only load and call smooth quant routine once for all ranks |
| 878 | hf_model = AutoModelForCausalLM.from_pretrained( |
| 879 | hf_model_dir, |
| 880 | device_map='auto' if device != 'cpu' else 'cpu', |
| 881 | dtype='auto' if not use_smooth_quant else torch.float16, |
| 882 | trust_remote_code=trust_remote_code) |
| 883 | |
| 884 | os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( |
| 885 | "TOKENIZERS_PARALLELISM", "false") |
| 886 | tokenizer = AutoTokenizer.from_pretrained( |
| 887 | hf_model_dir, |
| 888 | trust_remote_code=trust_remote_code, |
| 889 | use_fast=False, |
| 890 | padding_side='left') |
| 891 | |
| 892 | dataset = load_calib_dataset(calib_dataset) |
| 893 | act_range = capture_activation_range(hf_model, tokenizer, dataset) |
| 894 | if use_smooth_quant: |
| 895 | smooth_gpt_model(hf_model, act_range, quant_config.smoothquant_val) |
| 896 | |
| 897 | for rank in range(mapping.world_size): |
| 898 | # 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 |
| 899 | config = copy.deepcopy(config) |
| 900 | config.set_rank(rank) |
| 901 | weights = load_weights_from_hf_model( |
| 902 | hf_model, |
| 903 | config=config, |
| 904 | act_range=act_range, |
| 905 | ) |
| 906 | safetensors.torch.save_file( |
| 907 | weights, os.path.join(output_dir, f'rank{rank}.safetensors')) |
| 908 | del weights |
| 909 | |
| 910 | |
| 911 | def load_hf_gpt(model_dir: str, load_model_on_cpu: bool = False): |
no test coverage detected