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

50
python/models/qwen3_5.py
View File

@@ -141,15 +141,12 @@ def _dp_all_gather(
) -> torch.Tensor:
if world_size <= 1:
return tensor
try:
from vllm.distributed import get_dp_group
from vllm.distributed import get_dp_group
group = get_dp_group()
gathered = [torch.empty_like(tensor) for _ in range(world_size)]
torch.distributed.all_gather(gathered, tensor, group=group)
return torch.cat(gathered, dim=dim)
except (ImportError, RuntimeError):
return tensor.repeat(world_size, *([1] * (tensor.dim() - 1)))
group = get_dp_group()
gathered = [torch.empty_like(tensor) for _ in range(world_size)]
torch.distributed.all_gather(gathered, tensor, group=group)
return torch.cat(gathered, dim=dim)
def _dp_all_gather_variable(
@@ -158,25 +155,22 @@ def _dp_all_gather_variable(
dp_rank: int,
group_name: str = "dp",
) -> torch.Tensor:
try:
from vllm.distributed import get_dp_group
from vllm.distributed import get_dp_group
group = get_dp_group()
world_size = len(token_counts)
hidden_dim = tensor.shape[1] if tensor.dim() > 1 else 1
recv_tensors = []
for i, count in enumerate(token_counts):
if i == dp_rank:
recv_tensors.append(tensor[:count])
else:
recv_tensors.append(
torch.empty(
count, hidden_dim, dtype=tensor.dtype, device=tensor.device
)
group = get_dp_group()
world_size = len(token_counts)
hidden_dim = tensor.shape[1] if tensor.dim() > 1 else 1
recv_tensors = []
for i, count in enumerate(token_counts):
if i == dp_rank:
recv_tensors.append(tensor[:count])
else:
recv_tensors.append(
torch.empty(
count, hidden_dim, dtype=tensor.dtype, device=tensor.device
)
torch.distributed.all_gather(
recv_tensors, tensor[: token_counts[dp_rank]], group=group
)
return torch.cat(recv_tensors, dim=0)
except (ImportError, RuntimeError):
return tensor
)
torch.distributed.all_gather(
recv_tensors, tensor[: token_counts[dp_rank]], group=group
)
return torch.cat(recv_tensors, dim=0)