[Fix] NPU deepep hccl buffer and fix IPC safe check (#17804)
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -2359,6 +2359,7 @@ class SafeUnpickler(pickle.Unpickler):
|
||||
"sglang.srt.model_executor.model_runner.",
|
||||
"sglang.srt.layers.",
|
||||
"sglang.srt.utils.",
|
||||
"torch_npu.",
|
||||
}
|
||||
|
||||
DENY_CLASSES = {
|
||||
|
||||
Reference in New Issue
Block a user