dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

under test, not sure no errors

Browse Source
This commit is contained in:
DP Migration
2026-09-01 10:28:11 +00:00
parent 94d77cf0b4
commit 96f4bafcef
3 changed files with 43 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/layers/fused_moe.py
View File

@@ -104,15 +104,11 @@ def _dp_all_gather(
"""All-gather ``tensor`` along ``dim`` across the DP process group.""" """All-gather ``tensor`` along ``dim`` across the DP process group."""
if world_size <= 1: if world_size <= 1:
return tensor return tensor
try:
from vllm.distributed import get_dp_group from vllm.distributed import get_dp_group
group = get_dp_group() group = get_dp_group()
gathered = [torch.empty_like(tensor) for _ in range(world_size)] gathered = [torch.empty_like(tensor) for _ in range(world_size)]
torch.distributed.all_gather(gathered, tensor, group=group) torch.distributed.all_gather(gathered, tensor, group=group)
return torch.cat(gathered, dim=dim) return torch.cat(gathered, dim=dim)
except (ImportError, RuntimeError):
# Fallback: repeat for testing without actual distributed backend
return tensor.repeat(world_size, *([1] * (tensor.dim() - 1)))
def _dp_all_gather_variable( def _dp_all_gather_variable(
@@ -123,7 +119,6 @@ def _dp_all_gather_variable(
) -> torch.Tensor: ) -> torch.Tensor:
"""Variable-length all-gather: each rank contributes a different number """Variable-length all-gather: each rank contributes a different number
of tokens. Returns a compact concatenation without padding.""" of tokens. Returns a compact concatenation without padding."""
try:
from vllm.distributed import get_dp_group from vllm.distributed import get_dp_group
group = get_dp_group() group = get_dp_group()
world_size = len(token_counts) world_size = len(token_counts)
@@ -138,5 +133,3 @@ def _dp_all_gather_variable(
) )
torch.distributed.all_gather(recv_tensors, tensor[:token_counts[dp_rank]], group=group) torch.distributed.all_gather(recv_tensors, tensor[:token_counts[dp_rank]], group=group)
return torch.cat(recv_tensors, dim=0) return torch.cat(recv_tensors, dim=0)
except (ImportError, RuntimeError):
return tensor

5
python/model_executor/executor.py
View File

@@ -120,13 +120,8 @@ class ModelExecutor:
).lower() ).lower()
graph_disabled = graph_backend in ("", "off", "none", "0") graph_disabled = graph_backend in ("", "off", "none", "0")
if graph_disabled and config.get("enable_graph", False): if graph_disabled and config.get("enable_graph", False):
# Default to ACL graph on NPU platforms
try:
import torch_npu # noqa: F401 import torch_npu # noqa: F401
return "aclgraph" return "aclgraph"
except ImportError:
pass
return graph_backend return graph_backend
@torch.inference_mode() @torch.inference_mode()

6
python/models/qwen3_5.py
View File

@@ -141,15 +141,12 @@ def _dp_all_gather(
) -> torch.Tensor: ) -> torch.Tensor:
if world_size <= 1: if world_size <= 1:
return tensor return tensor
try:
from vllm.distributed import get_dp_group from vllm.distributed import get_dp_group
group = get_dp_group() group = get_dp_group()
gathered = [torch.empty_like(tensor) for _ in range(world_size)] gathered = [torch.empty_like(tensor) for _ in range(world_size)]
torch.distributed.all_gather(gathered, tensor, group=group) torch.distributed.all_gather(gathered, tensor, group=group)
return torch.cat(gathered, dim=dim) return torch.cat(gathered, dim=dim)
except (ImportError, RuntimeError):
return tensor.repeat(world_size, *([1] * (tensor.dim() - 1)))
def _dp_all_gather_variable( def _dp_all_gather_variable(
@@ -158,7 +155,6 @@ def _dp_all_gather_variable(
dp_rank: int, dp_rank: int,
group_name: str = "dp", group_name: str = "dp",
) -> torch.Tensor: ) -> torch.Tensor:
try:
from vllm.distributed import get_dp_group from vllm.distributed import get_dp_group
group = get_dp_group() group = get_dp_group()
@@ -178,5 +174,3 @@ def _dp_all_gather_variable(
recv_tensors, tensor[: token_counts[dp_rank]], group=group recv_tensors, tensor[: token_counts[dp_rank]], group=group
) )
return torch.cat(recv_tensors, dim=0) return torch.cat(recv_tensors, dim=0)
except (ImportError, RuntimeError):
return tensor