[Fix] Support origin HF weight backfilling in parallel converter

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.")