minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NPU] RL update_weights_from_disk/ tensor /distributed (#26717)

Browse Source
This commit is contained in:
McZyWu
2026-06-09 16:52:36 +08:00
committed by GitHub
parent cd6efcb947
commit c6be251c5b
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+9
View File
@@ -70,6 +70,7 @@ from sglang.srt.utils import (
get_bool_env_var, get_bool_env_var,
is_cpu, is_cpu,
is_hip, is_hip,
is_npu,
print_info_once, print_info_once,
round_up, round_up,
) )
@@ -78,6 +79,7 @@ from sglang.srt.utils.custom_op import register_custom_op
_is_hip = is_hip() _is_hip = is_hip()
_is_cpu_amx_available = cpu_has_amx_support() _is_cpu_amx_available = cpu_has_amx_support()
_is_cpu = is_cpu() _is_cpu = is_cpu()
_is_npu = is_npu()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -778,6 +780,13 @@ class FusedMoE(torch.nn.Module):
# expert weights into block layout. During weight update, we must restore # expert weights into block layout. During weight update, we must restore
# canonical load-time shapes before copying checkpoint tensors. # canonical load-time shapes before copying checkpoint tensors.
if isinstance(method, UnquantizedFusedMoEMethod): if isinstance(method, UnquantizedFusedMoEMethod):
if _is_npu:
if weight_name.endswith(".experts.w2_weight"):
if param.data.shape[1] != loaded_weight.shape[0]:
param.data = param.data.transpose(1, 2).contiguous()
if weight_name.endswith(".experts.w13_weight"):
if param.data.shape[2] != loaded_weight.shape[1]:
param.data = param.data.transpose(1, 2).contiguous()
method.maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load( method.maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load(
layer=self, layer=self,
param=param, param=param,