[Fix] NPU deepep hccl buffer and fix IPC safe check (#17804)

This commit is contained in:
1StepForever
2026-03-03 14:56:06 +08:00
committed by GitHub
parent 43a9249e7c
commit 3c01b44700
2 changed files with 24 additions and 1 deletions

View File

@@ -69,6 +69,24 @@ TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
REDUCE_OP_SUM = int(torch.distributed.ReduceOp.SUM)
def get_torch_distributed_pg_options(group_name=None):
if not _is_npu:
return None
# Only create HCCL options for default group or MoE-related groups
if group_name is not None and "moe" not in group_name:
return None
import torch_npu
options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()
hccl_buffer_size = int(
os.environ.get("DEEPEP_HCCL_BUFFSIZE") or os.environ.get("HCCL_BUFFSIZE") or 200
)
options.hccl_config = {"hccl_buffer_size": hccl_buffer_size}
return options
@dataclass
class GraphCaptureContext:
stream: torch.get_device_module().Stream
@@ -274,8 +292,9 @@ class GroupCoordinator:
pg_options=MooncakeBackendOptions(active_ranks_cpu),
)
else:
pg_options = get_torch_distributed_pg_options(group_name)
device_group = torch.distributed.new_group(
ranks, backend=torch_distributed_backend
ranks, backend=torch_distributed_backend, pg_options=pg_options
)
# a group with `gloo` backend, to allow direct coordination
# between processes through the CPU.
@@ -1629,6 +1648,8 @@ def init_distributed_environment(
assert timeout > 0, "timeout must be positive"
timeout = timedelta(seconds=timeout)
pg_options = get_torch_distributed_pg_options()
# this backend is used for WORLD
torch.distributed.init_process_group(
backend=backend,
@@ -1636,6 +1657,7 @@ def init_distributed_environment(
world_size=world_size,
rank=rank,
timeout=timeout,
pg_options=pg_options,
)
# set the local rank

View File

@@ -2359,6 +2359,7 @@ class SafeUnpickler(pickle.Unpickler):
"sglang.srt.model_executor.model_runner.",
"sglang.srt.layers.",
"sglang.srt.utils.",
"torch_npu.",
}
DENY_CLASSES = {