From 9eec3f9fc98c4bafe2b0a07bef593b4259107c6b Mon Sep 17 00:00:00 2001 From: sufubao Date: Thu, 11 Sep 2025 11:40:19 +0800 Subject: [PATCH] fix tl.where and set the default of loader worker number --- lightllm/common/basemodel/layer_weights/hf_load_utils.py | 3 ++- lightllm/common/fused_moe/grouped_topk.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/lightllm/common/basemodel/layer_weights/hf_load_utils.py b/lightllm/common/basemodel/layer_weights/hf_load_utils.py index 8cf66a5ad..897d6e949 100755 --- a/lightllm/common/basemodel/layer_weights/hf_load_utils.py +++ b/lightllm/common/basemodel/layer_weights/hf_load_utils.py @@ -60,7 +60,8 @@ 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)) + cpu_count = min(16, os.cpu_count() or 4) + worker = min(cpu_count, len(candidate_files)) 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" diff --git a/lightllm/common/fused_moe/grouped_topk.py b/lightllm/common/fused_moe/grouped_topk.py index d380e091e..058dd150e 100644 --- a/lightllm/common/fused_moe/grouped_topk.py +++ b/lightllm/common/fused_moe/grouped_topk.py @@ -35,7 +35,7 @@ def _compare_and_swap(x, x_1, ids, flip, i: tl.core.constexpr, n_dims: tl.core.c iright = right.to(idtype, bitcast=True) ix = x.to(idtype, bitcast=True) - cond = (left > right) ^ flip + cond = ((left > right) ^ flip) > 0 ret = ix ^ tl.core.where(cond, ileft ^ iright, zeros_like(ix))