diff --git a/tools/convert_torch_dist_to_hf_parallel.py b/tools/convert_torch_dist_to_hf_parallel.py index 763254d42..bcd049da9 100644 --- a/tools/convert_torch_dist_to_hf_parallel.py +++ b/tools/convert_torch_dist_to_hf_parallel.py @@ -281,24 +281,27 @@ def save_tensors(args, model_name, state_dict, output_dir, chunk_size, vocab_siz print(f"Total parameters to process: {len(param_list)}") all_converted_tensors = [] - lock = threading.Lock() def process_and_collect(name_param_pair): name, param = name_param_pair - try: - converted = process_param(args, model_name, name, param, vocab_size) - return converted - except Exception as e: - print(f"Error processing {name}: {e}") - return [] + return process_param(args, model_name, name, param, vocab_size) + conversion_errors = [] print(f"Processing with {max_workers} workers") with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = {executor.submit(process_and_collect, (name, param)): name for name, param in param_list} for future in tqdm(as_completed(futures), total=len(futures), desc="Converting parameters"): - converted = future.result() - with lock: - all_converted_tensors.extend(converted) + name = futures[future] + try: + converted = future.result() + except Exception as error: + conversion_errors.append(f"- {name}: {type(error).__name__}: {error}") + continue + all_converted_tensors.extend(converted) + + if conversion_errors: + details = "\n".join(conversion_errors) + raise RuntimeError(f"Failed to convert {len(conversion_errors)} parameter(s):\n{details}") current_size = 0 total_size = 0 @@ -353,6 +356,53 @@ def copy_assets(origin_hf_dir, output_dir): shutil.copy(src, dst) +def save_missing_tensors(origin_hf_dir, converted_names, output_dir, chunk_size, start_file_index): + safetensors_files = sorted(f for f in os.listdir(origin_hf_dir) if f.endswith(".safetensors")) + missing_weight_map = {} + current_tensors = {} + current_size = 0 + total_size = 0 + file_index = start_file_index + + def flush_current_tensors(): + nonlocal current_tensors, current_size, file_index + if not current_tensors: + return + + filename = f"model-{file_index:05d}.safetensors" + filepath = os.path.join(output_dir, filename) + print(f"saving {len(current_tensors)} missing tensors to {filepath}") + safetensors.torch.save_file(current_tensors, filepath) + for name in current_tensors: + missing_weight_map[name] = filename + current_tensors = {} + current_size = 0 + file_index += 1 + + for filename in safetensors_files: + filepath = os.path.join(origin_hf_dir, filename) + with safetensors.safe_open(filepath, framework="pt", device="cpu") as f: + for name in f.keys(): + if name in converted_names: + continue + if name in missing_weight_map or name in current_tensors: + raise ValueError(f"Duplicate tensor {name} found in origin HF checkpoint") + + tensor = f.get_tensor(name) + tensor_size = tensor.numel() * tensor.element_size() + if current_tensors and tensor_size + current_size > chunk_size: + flush_current_tensors() + + print(f"add {name} from origin hf checkpoint") + current_tensors[name] = tensor + current_size += tensor_size + total_size += tensor_size + + flush_current_tensors() + print(f"Added {len(missing_weight_map)} missing tensors from origin HF checkpoint") + return missing_weight_map, total_size, file_index + + def conversion_worker( worker_id, keys, @@ -416,6 +466,9 @@ def conversion_worker( parser.add_argument( "-f", "--force", action="store_true", help="Force overwrite the output directory if it exists." ) + parser.add_argument( + "-a", "--add-missing-from-origin-hf", action="store_true", help="Add missing weights from origin hf checkpoint" + ) parser.add_argument( "--chunk-size", type=int, @@ -449,6 +502,8 @@ def conversion_worker( raise ValueError( "Either --model-name or --origin-hf-dir must be provided, so that we can know the name of the params." ) + if args.add_missing_from_origin_hf and args.origin_hf_dir is None: + raise ValueError("--add-missing-from-origin-hf requires --origin-hf-dir") if args.model_name is None: hf_config = AutoConfig.from_pretrained(args.origin_hf_dir, trust_remote_code=True) @@ -557,6 +612,17 @@ def conversion_worker( if os.path.exists(temp_index): os.remove(temp_index) + if args.add_missing_from_origin_hf: + missing_weight_map, missing_size, final_file_index = save_missing_tensors( + args.origin_hf_dir, + set(final_weight_map), + args.output_dir, + args.chunk_size, + final_file_index, + ) + final_weight_map.update(missing_weight_map) + total_size += missing_size + total_files = final_file_index - 1 final_weight_map_fixed = {} for i in range(1, total_files + 1): @@ -570,6 +636,12 @@ def conversion_worker( if v == old_name: final_weight_map_fixed[k] = new_name + if len(final_weight_map_fixed) != len(final_weight_map): + raise RuntimeError( + f"Final weight map is incomplete: expected {len(final_weight_map)} tensors, " + f"found {len(final_weight_map_fixed)}" + ) + index_data = {"metadata": {"total_size": total_size}, "weight_map": final_weight_map_fixed} json.dump(index_data, open(os.path.join(args.output_dir, "model.safetensors.index.json"), "w"), indent=2) print("Model converted and saved.")