(ckpt_prefix, save_prefix, opt_model_path)
| 679 | |
| 680 | |
| 681 | def convert_ckpt(ckpt_prefix, save_prefix, opt_model_path): |
| 682 | opt_variables_path = f'{opt_model_path}/variables/variables' |
| 683 | opt_variables = [name for name, _ in tf.train.list_variables(opt_variables_path)] |
| 684 | input_variables = [name for name, _ in tf.train.list_variables(ckpt_prefix)] |
| 685 | # Find embedding variables |
| 686 | v2_quant_dict = {INT8: [], BF16: [], FP16: []} |
| 687 | ev_quant_dict = {INT8: [], BF16: [], FP16: []} |
| 688 | updated_variables = list(set(opt_variables).difference(set(input_variables))) |
| 689 | for var_name in updated_variables: |
| 690 | suffix = var_name.split('-')[-1] |
| 691 | suffixes = [s for s in ev_suffix if s != 'values'] |
| 692 | if suffix in suffixes or var_name.endswith('/int8_scale'): |
| 693 | continue |
| 694 | is_qvar = False |
| 695 | for opt_dtype in [INT8, BF16, FP16]: |
| 696 | if var_name.endswith(f'/{opt_dtype.lower()}_data'): |
| 697 | v2_quant_dict[opt_dtype].append(var_name) |
| 698 | is_qvar = True |
| 699 | if var_name.endswith(f'/{opt_dtype.lower()}_data-values'): |
| 700 | ev_quant_dict[opt_dtype].append(var_name[:-7]) |
| 701 | is_qvar = True |
| 702 | if not is_qvar: |
| 703 | print(f'Not found variable: {var_name}') |
| 704 | # Update checkpoint |
| 705 | variable_path = ckpt_prefix |
| 706 | for quant_dict, is_ev in [(v2_quant_dict, False), (ev_quant_dict, True)]: |
| 707 | for opt_dtype, quant_names in quant_dict.items(): |
| 708 | if len(quant_names) == 0: |
| 709 | continue |
| 710 | len_suffix = len(f'/{opt_dtype.lower()}_data') |
| 711 | names = [name[:-len_suffix] for name in quant_names] |
| 712 | if opt_dtype == INT8: |
| 713 | scale_names = [f'{name}/int8_scale' for name in names] |
| 714 | else: |
| 715 | scale_names = [] |
| 716 | tmp_path = tempfile.mkdtemp(dir='.') |
| 717 | opt_path = f'{tmp_path}/variables' |
| 718 | dtype = get_tf_dtype(opt_dtype) |
| 719 | quantize_embedding_variable.quantize_by_name( |
| 720 | variable_path, opt_path, names, quant_names, scale_names, dtype, is_ev |
| 721 | ) |
| 722 | if variable_path != ckpt_prefix: |
| 723 | tf.io.gfile.rmtree(os.path.dirname(variable_path)) |
| 724 | variable_path = opt_path |
| 725 | # Remove unused variables |
| 726 | cur_variables = [name for name, _ in tf.train.list_variables(variable_path)] |
| 727 | removed_variables = list(set(cur_variables).difference(set(opt_variables))) |
| 728 | save_dir = os.path.dirname(save_prefix) |
| 729 | if not tf.io.gfile.exists(save_dir): |
| 730 | tf.io.gfile.makedirs(save_dir) |
| 731 | quantize_embedding_variable.remove_variables_by_name( |
| 732 | variable_path, save_prefix, removed_variables |
| 733 | ) |
| 734 | if variable_path != ckpt_prefix: |
| 735 | tf.io.gfile.rmtree(os.path.dirname(variable_path)) |
| 736 | print('Convert Result:') |
| 737 | for opt_dtype, quant_names in v2_quant_dict.items(): |
| 738 | for name in quant_names: |
nothing calls this directly
no test coverage detected