fix

Author: sufubao

lightllm/common/basemodel/layer_weights/hf_load_utils.py

@@ -51,6 +51,7 @@ def load_hf_weights(data_type, weight_dir, pre_post_layer=None, transformer_laye
         candidate_files = list(filter(lambda x: x.endswith(".bin"), files))
     assert len(candidate_files) != 0, "can only support pytorch tensor and safetensors format for weights."
     from functools import partial
+    from multiprocessing import cpu_count
     from multiprocessing.pool import ThreadPool as Pool
 
     partial_func = partial(
@@ -60,7 +61,7 @@ def load_hf_weights(data_type, weight_dir, pre_post_layer=None, transformer_laye
         transformer_layer_list=transformer_layer_list,
         weight_dir=weight_dir,
     )  # noqa
-    worker = int(os.environ.get("LOADWORKER", 1))
+    worker = min(24, cpu_count())
     with Pool(worker) as p:
         iterator = p.imap_unordered(partial_func, candidate_files, chunksize=1)
         desc_str = f"pid {os.getpid()} Loading model weights with {worker} workers"