From 0757f2353a3e7566b9d7abea28afc420ab00619e Mon Sep 17 00:00:00 2001 From: qrskannbara <94727257+albaNnaksqr@users.noreply.github.com> Date: Thu, 13 Aug 2026 10:55:10 +0800 Subject: [PATCH 1/3] fix: add missing HF weights in parallel converter --- .../test_convert_torch_dist_to_hf_parallel.py | 78 +++++++++++++++++++ tools/convert_torch_dist_to_hf_parallel.py | 75 ++++++++++++++++++ 2 files changed, 153 insertions(+) create mode 100644 tests/test_convert_torch_dist_to_hf_parallel.py diff --git a/tests/test_convert_torch_dist_to_hf_parallel.py b/tests/test_convert_torch_dist_to_hf_parallel.py new file mode 100644 index 0000000000..b4853fb581 --- /dev/null +++ b/tests/test_convert_torch_dist_to_hf_parallel.py @@ -0,0 +1,78 @@ +import importlib.util +from pathlib import Path + +import safetensors.torch +import torch + + +CONVERTER_PATH = Path(__file__).parents[1] / "tools" / "convert_torch_dist_to_hf_parallel.py" +SPEC = importlib.util.spec_from_file_location("convert_torch_dist_to_hf_parallel", CONVERTER_PATH) +assert SPEC is not None and SPEC.loader is not None +CONVERTER = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(CONVERTER) +save_missing_tensors = CONVERTER.save_missing_tensors + + +def _load_output_tensor(output_dir, weight_map, name): + return safetensors.torch.load_file(output_dir / weight_map[name])[name] + + +def test_save_missing_tensors_only_copies_unconverted_weights(tmp_path): + origin_dir = tmp_path / "origin" + output_dir = tmp_path / "output" + origin_dir.mkdir() + output_dir.mkdir() + + converted = torch.tensor([1.0, 2.0]) + missing_visual = torch.tensor([3.0, 4.0]) + missing_mtp = torch.tensor([5.0, 6.0, 7.0]) + safetensors.torch.save_file( + {"model.converted.weight": converted, "model.visual.weight": missing_visual}, + origin_dir / "model-00001-of-00002.safetensors", + ) + safetensors.torch.save_file( + {"mtp.weight": missing_mtp}, + origin_dir / "model-00002-of-00002.safetensors", + ) + + weight_map, total_size, next_file_index = save_missing_tensors( + origin_dir, + {"model.converted.weight"}, + output_dir, + chunk_size=12, + start_file_index=3, + ) + + assert set(weight_map) == {"model.visual.weight", "mtp.weight"} + assert total_size == ( + missing_visual.numel() * missing_visual.element_size() + missing_mtp.numel() * missing_mtp.element_size() + ) + assert next_file_index == 5 + assert sorted(path.name for path in output_dir.glob("*.safetensors")) == [ + "model-00003.safetensors", + "model-00004.safetensors", + ] + torch.testing.assert_close(_load_output_tensor(output_dir, weight_map, "model.visual.weight"), missing_visual) + torch.testing.assert_close(_load_output_tensor(output_dir, weight_map, "mtp.weight"), missing_mtp) + + +def test_save_missing_tensors_writes_nothing_when_checkpoint_is_complete(tmp_path): + origin_dir = tmp_path / "origin" + output_dir = tmp_path / "output" + origin_dir.mkdir() + output_dir.mkdir() + tensor = torch.tensor([1.0, 2.0]) + safetensors.torch.save_file({"model.weight": tensor}, origin_dir / "model.safetensors") + + weight_map, total_size, next_file_index = save_missing_tensors( + origin_dir, + {"model.weight"}, + output_dir, + chunk_size=1024, + start_file_index=2, + ) + + assert weight_map == {} + assert total_size == 0 + assert next_file_index == 2 + assert list(output_dir.iterdir()) == [] diff --git a/tools/convert_torch_dist_to_hf_parallel.py b/tools/convert_torch_dist_to_hf_parallel.py index 763254d42c..e5dcde128c 100644 --- a/tools/convert_torch_dist_to_hf_parallel.py +++ b/tools/convert_torch_dist_to_hf_parallel.py @@ -353,6 +353,59 @@ 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): + """Save tensors that exist only in the original HF checkpoint. + + The parallel converter writes temporary ``model-NNNNN.safetensors`` files + before adding the final ``-of-NNNNN`` suffix. Missing tensors follow the + same convention so they can participate in the existing final rename. + """ + 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 +469,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 +505,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 +615,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 +639,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.") From d0d368efe778d85040de3c3d65bbf11f98c9f1dd Mon Sep 17 00:00:00 2001 From: qrskannbara <94727257+albaNnaksqr@users.noreply.github.com> Date: Thu, 13 Aug 2026 19:58:45 +0800 Subject: [PATCH 2/3] chore: trim parallel converter test scaffolding --- .../test_convert_torch_dist_to_hf_parallel.py | 78 ------------------- tools/convert_torch_dist_to_hf_parallel.py | 6 -- 2 files changed, 84 deletions(-) delete mode 100644 tests/test_convert_torch_dist_to_hf_parallel.py diff --git a/tests/test_convert_torch_dist_to_hf_parallel.py b/tests/test_convert_torch_dist_to_hf_parallel.py deleted file mode 100644 index b4853fb581..0000000000 --- a/tests/test_convert_torch_dist_to_hf_parallel.py +++ /dev/null @@ -1,78 +0,0 @@ -import importlib.util -from pathlib import Path - -import safetensors.torch -import torch - - -CONVERTER_PATH = Path(__file__).parents[1] / "tools" / "convert_torch_dist_to_hf_parallel.py" -SPEC = importlib.util.spec_from_file_location("convert_torch_dist_to_hf_parallel", CONVERTER_PATH) -assert SPEC is not None and SPEC.loader is not None -CONVERTER = importlib.util.module_from_spec(SPEC) -SPEC.loader.exec_module(CONVERTER) -save_missing_tensors = CONVERTER.save_missing_tensors - - -def _load_output_tensor(output_dir, weight_map, name): - return safetensors.torch.load_file(output_dir / weight_map[name])[name] - - -def test_save_missing_tensors_only_copies_unconverted_weights(tmp_path): - origin_dir = tmp_path / "origin" - output_dir = tmp_path / "output" - origin_dir.mkdir() - output_dir.mkdir() - - converted = torch.tensor([1.0, 2.0]) - missing_visual = torch.tensor([3.0, 4.0]) - missing_mtp = torch.tensor([5.0, 6.0, 7.0]) - safetensors.torch.save_file( - {"model.converted.weight": converted, "model.visual.weight": missing_visual}, - origin_dir / "model-00001-of-00002.safetensors", - ) - safetensors.torch.save_file( - {"mtp.weight": missing_mtp}, - origin_dir / "model-00002-of-00002.safetensors", - ) - - weight_map, total_size, next_file_index = save_missing_tensors( - origin_dir, - {"model.converted.weight"}, - output_dir, - chunk_size=12, - start_file_index=3, - ) - - assert set(weight_map) == {"model.visual.weight", "mtp.weight"} - assert total_size == ( - missing_visual.numel() * missing_visual.element_size() + missing_mtp.numel() * missing_mtp.element_size() - ) - assert next_file_index == 5 - assert sorted(path.name for path in output_dir.glob("*.safetensors")) == [ - "model-00003.safetensors", - "model-00004.safetensors", - ] - torch.testing.assert_close(_load_output_tensor(output_dir, weight_map, "model.visual.weight"), missing_visual) - torch.testing.assert_close(_load_output_tensor(output_dir, weight_map, "mtp.weight"), missing_mtp) - - -def test_save_missing_tensors_writes_nothing_when_checkpoint_is_complete(tmp_path): - origin_dir = tmp_path / "origin" - output_dir = tmp_path / "output" - origin_dir.mkdir() - output_dir.mkdir() - tensor = torch.tensor([1.0, 2.0]) - safetensors.torch.save_file({"model.weight": tensor}, origin_dir / "model.safetensors") - - weight_map, total_size, next_file_index = save_missing_tensors( - origin_dir, - {"model.weight"}, - output_dir, - chunk_size=1024, - start_file_index=2, - ) - - assert weight_map == {} - assert total_size == 0 - assert next_file_index == 2 - assert list(output_dir.iterdir()) == [] diff --git a/tools/convert_torch_dist_to_hf_parallel.py b/tools/convert_torch_dist_to_hf_parallel.py index e5dcde128c..8d6e35c1be 100644 --- a/tools/convert_torch_dist_to_hf_parallel.py +++ b/tools/convert_torch_dist_to_hf_parallel.py @@ -354,12 +354,6 @@ def copy_assets(origin_hf_dir, output_dir): def save_missing_tensors(origin_hf_dir, converted_names, output_dir, chunk_size, start_file_index): - """Save tensors that exist only in the original HF checkpoint. - - The parallel converter writes temporary ``model-NNNNN.safetensors`` files - before adding the final ``-of-NNNNN`` suffix. Missing tensors follow the - same convention so they can participate in the existing final rename. - """ safetensors_files = sorted(f for f in os.listdir(origin_hf_dir) if f.endswith(".safetensors")) missing_weight_map = {} current_tensors = {} From 539c0eb84dd4945e9c03da9dd283501279e6f1ec Mon Sep 17 00:00:00 2001 From: qrskannbara <94727257+albaNnaksqr@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:55:44 +0800 Subject: [PATCH 3/3] fix: fail on parallel conversion errors --- tools/convert_torch_dist_to_hf_parallel.py | 23 ++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/tools/convert_torch_dist_to_hf_parallel.py b/tools/convert_torch_dist_to_hf_parallel.py index 8d6e35c1be..bcd049da9b 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