Skip to content

Derive CAM HCCL buffer sizes per role #234

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ShwStone wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Derive CAM HCCL buffer sizes per role #234

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 33 additions & 10 deletions afd_plugin/connectors/npu/async_cam.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,13 @@
AFDTransferMetadata,
AFDTransferState,
)
from afd_plugin.distributed import init_afd_process_group
from afd_plugin.distributed import (
create_hccl_process_group_options,
init_afd_process_group,
)
from afd_plugin.distributed.cam_hccl_buffer import (
derive_cam_hccl_buffer_plan_from_config,
)

if TYPE_CHECKING:
from torch.distributed.distributed_c10d import ProcessGroup
Expand Down Expand Up @@ -238,8 +244,8 @@ def __init__(

Communication resources are created collectively by
``init_afd_connector``. ``role_rank`` is resolved before connector
construction; ``attn_ranks_per_dp`` is used as the CAM Attention TP
width.
construction; ``attn_ranks_per_dp`` supplies the number of NPUs in one
Attention data-parallel group.
"""
super().__init__(rank, local_rank, vllm_config, afd_config, role_rank)
self._initialized = False
Expand All @@ -251,7 +257,21 @@ def __init__(
self.group_name = ""
self.max_seq_len = vllm_config.scheduler_config.max_num_batched_tokens
self.comm_id = CAM_COMM_ID
self.tp_size = self.extra_info.attn_ranks_per_dp
self.num_npus_per_dp_group = self.extra_info.attn_ranks_per_dp
self.hccl_buffer_plan = derive_cam_hccl_buffer_plan_from_config(
vllm_config,
afd_config,
)
self.hccl_buffer_size_mb = self.hccl_buffer_plan.buffer_size_mb_for_role(
afd_config.role,
)
logger.info(
"CAM async %s HCCL buffer size is %d MB (auto-derived with "
"1.1x headroom from %d required bytes)",
afd_config.role,
self.hccl_buffer_size_mb,
self.hccl_buffer_plan.required_bytes_for_role(afd_config.role),
)
self.cam_pg: ProcessGroup | None = None
self.topology = build_async_topology(
afd_config,
Expand Down Expand Up @@ -292,6 +312,9 @@ def init_afd_connector(self) -> None:
rank=self.world_rank,
group_name=AFD_ASYNC_CAM_GROUP_NAME,
timeout=timedelta(minutes=30),
pg_options=create_hccl_process_group_options(
self.hccl_buffer_size_mb,
),
)
backend = self.cam_pg._get_backend(torch.device("npu"))
self.group_name = str(backend.get_hccl_comm_name(self.world_rank))
Expand Down Expand Up @@ -523,7 +546,7 @@ def send_attn_output(
rank=self.world_rank,
world_size=self.topology.world_size,
layer_idx=states.layer_idx,
tp_size=self.tp_size,
num_npus_per_dp_group=self.num_npus_per_dp_group,
dynamic_quant=self.dynamic_quant,
group_name=self.group_name,
)
Expand All @@ -542,7 +565,7 @@ def send_attn_output(
self.world_rank,
self.topology.world_size,
states.layer_idx,
self.tp_size,
self.num_npus_per_dp_group,
self.dynamic_quant,
self.group_name,
)
Expand Down Expand Up @@ -691,7 +714,7 @@ def recv_attn_output(
expert_per_rank=self.expert_per_rank,
rank=self.world_rank,
world_size=self.topology.world_size,
tp_size=self.tp_size,
num_npus_per_dp_group=self.num_npus_per_dp_group,
dynamic_quant=self.dynamic_quant,
group_name=self.group_name,
)
Expand All @@ -707,7 +730,7 @@ def recv_attn_output(
self.expert_per_rank,
self.world_rank,
self.topology.world_size,
self.tp_size,
self.num_npus_per_dp_group,
self.dynamic_quant,
self.group_name,
)
Expand Down Expand Up @@ -782,7 +805,7 @@ def send_ffn_output(
expert_per_rank=self.expert_per_rank,
rank=self.world_rank,
world_size=self.topology.world_size,
tp_size=self.tp_size,
num_npus_per_dp_group=self.num_npus_per_dp_group,
group_name=self.group_name,
)
torch.ops.umdk_cam_op_lib.async_combine_send(
Expand All @@ -799,7 +822,7 @@ def send_ffn_output(
self.expert_per_rank,
self.world_rank,
self.topology.world_size,
self.tp_size,
self.num_npus_per_dp_group,
self.group_name,
)

Expand Down
29 changes: 28 additions & 1 deletion afd_plugin/connectors/npu/camp2p.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,14 @@
recv_control_payload,
send_control_payload,
)
from afd_plugin.distributed import init_afd_process_group, topology_from_config
from afd_plugin.distributed import (
create_hccl_process_group_options,
init_afd_process_group,
topology_from_config,
)
from afd_plugin.distributed.cam_hccl_buffer import (
derive_cam_hccl_buffer_plan_from_config,
)

if TYPE_CHECKING:
from vllm.config import VllmConfig
Expand Down Expand Up @@ -283,6 +290,20 @@ def __init__(
self.hidden_size = hf_config.hidden_size
self.num_experts_per_tok = hf_config.num_experts_per_tok
self.num_routed_experts = hf_config.n_routed_experts
self.hccl_buffer_plan = derive_cam_hccl_buffer_plan_from_config(
vllm_config,
afd_config,
)
self.hccl_buffer_size_mb = self.hccl_buffer_plan.buffer_size_mb_for_role(
afd_config.role,
)
logger.info(
"CAM P2P %s HCCL buffer size is %d MB (auto-derived with "
"1.1x headroom from %d required bytes)",
afd_config.role,
self.hccl_buffer_size_mb,
self.hccl_buffer_plan.required_bytes_for_role(afd_config.role),
)
self.control_plane = CAMP2pAFDControlPlane(self)

@property
Expand Down Expand Up @@ -325,6 +346,9 @@ def init_afd_connector(self) -> None:
rank=self.world_rank,
group_name=group_name,
timeout=timedelta(minutes=30),
pg_options=create_hccl_process_group_options(
self.hccl_buffer_size_mb,
),
)
self.afd_pg_list.append(afd_pg)
backend = afd_pg._get_backend(torch.device("npu"))
Expand All @@ -346,6 +370,9 @@ def init_afd_connector(self) -> None:
rank=self.world_rank,
group_name="afd_moe",
timeout=timedelta(minutes=30),
pg_options=create_hccl_process_group_options(
self.hccl_buffer_size_mb,
),
)
backend = self.ffn_pg._get_backend(torch.device("npu"))
self.hccl_comm_name1 = str(
Expand Down
7 changes: 6 additions & 1 deletion afd_plugin/distributed/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,11 @@


def __getattr__(name: str):
if name in {"DefaultProcessGroupSwitcher", "init_afd_process_group"}:
if name in {
"DefaultProcessGroupSwitcher",
"create_hccl_process_group_options",
"init_afd_process_group",
}:
from afd_plugin.distributed import afd_process_group

value = getattr(afd_process_group, name)
Expand All @@ -25,6 +29,7 @@ def __getattr__(name: str):
"AFDRankMapping",
"DefaultProcessGroupSwitcher",
"build_rank_mapping",
"create_hccl_process_group_options",
"init_afd_process_group",
"resolve_role_rank",
"topology_from_config",
Expand Down
14 changes: 14 additions & 0 deletions afd_plugin/distributed/afd_process_group.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,19 @@ def __exit__(self, exc_type: object, exc_value: object, tb: object) -> None:
_update_default_pg(self.default_group)


def create_hccl_process_group_options(hccl_buffer_size: int) -> Any:
"""Create fresh HCCL options for one plugin-owned process group.

``hccl_buffer_size`` is expressed in MB. Per-process-group options keep CAM
buffer sizing independent from process-wide ``HCCL_BUFFSIZE``.
"""
import torch_npu

options = torch_npu._C._distributed_c10d.ProcessGroupHCCL.Options()
options.hccl_config = {"hccl_buffer_size": hccl_buffer_size}
return options


def init_afd_process_group(
*,
backend: str,
Expand Down Expand Up @@ -99,5 +112,6 @@ def init_afd_process_group(

__all__ = [
"DefaultProcessGroupSwitcher",
"create_hccl_process_group_options",
"init_afd_process_group",
]
Loading
Loading