[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)
|
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
|
@dataclass
|
||||||
class GraphCaptureContext:
|
class GraphCaptureContext:
|
||||||
stream: torch.get_device_module().Stream
|
stream: torch.get_device_module().Stream
|
||||||
@@ -274,8 +292,9 @@ class GroupCoordinator:
|
|||||||
pg_options=MooncakeBackendOptions(active_ranks_cpu),
|
pg_options=MooncakeBackendOptions(active_ranks_cpu),
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
pg_options = get_torch_distributed_pg_options(group_name)
|
||||||
device_group = torch.distributed.new_group(
|
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
|
# a group with `gloo` backend, to allow direct coordination
|
||||||
# between processes through the CPU.
|
# between processes through the CPU.
|
||||||
@@ -1629,6 +1648,8 @@ def init_distributed_environment(
|
|||||||
assert timeout > 0, "timeout must be positive"
|
assert timeout > 0, "timeout must be positive"
|
||||||
timeout = timedelta(seconds=timeout)
|
timeout = timedelta(seconds=timeout)
|
||||||
|
|
||||||
|
pg_options = get_torch_distributed_pg_options()
|
||||||
|
|
||||||
# this backend is used for WORLD
|
# this backend is used for WORLD
|
||||||
torch.distributed.init_process_group(
|
torch.distributed.init_process_group(
|
||||||
backend=backend,
|
backend=backend,
|
||||||
@@ -1636,6 +1657,7 @@ def init_distributed_environment(
|
|||||||
world_size=world_size,
|
world_size=world_size,
|
||||||
rank=rank,
|
rank=rank,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
|
pg_options=pg_options,
|
||||||
)
|
)
|
||||||
|
|
||||||
# set the local rank
|
# set the local rank
|
||||||
|
|||||||
@@ -2359,6 +2359,7 @@ class SafeUnpickler(pickle.Unpickler):
|
|||||||
"sglang.srt.model_executor.model_runner.",
|
"sglang.srt.model_executor.model_runner.",
|
||||||
"sglang.srt.layers.",
|
"sglang.srt.layers.",
|
||||||
"sglang.srt.utils.",
|
"sglang.srt.utils.",
|
||||||
|
"torch_npu.",
|
||||||
}
|
}
|
||||||
|
|
||||||
DENY_CLASSES = {
|
DENY_CLASSES = {
|
||||||
|
|||||||
Reference in New Issue
Block a user