CANN/hixl LLM数据分发API

LLMDataDist

【免费下载链接】hixlHIXL（Huawei Xfer Library）是一个灵活、高效的昇腾单边通信库，面向集群场景提供简单、可靠、高效的点对点数据传输能力。项目地址: https://gitcode.com/cann/hixl

产品支持情况

说明：针对Atlas A2 训练系列产品/Atlas A2 推理系列产品，仅支持Atlas 800I A2 推理服务器、A200I A2 Box 异构组件。针对Ascend 950PR/Ascend 950DT，不支持link、unlink和query_register_mem_status。

LLMDataDist构造函数

函数功能

构造LLMDataDist。

函数原型

__init__(role: LLMRole, cluster_id: int)

参数说明

调用示例

from llm_datadist import LLMDataDist, LLMRole llm_datadist = LLMDataDist(LLMRole.DECODER, 0)

返回值

正常情况下返回LLMDataDist的实例。

参数错误可能抛出TypeError或ValueError。

约束说明

无

init

函数功能

初始化LLMDataDist，需要在init的options中配置CacheManager模式，即enable_cache_manager设置为True或者指定local_comm_res。

该模式下，Decode和Prompt可以双向拉取Cache。

函数原型

init(options: Dict[str, str])

参数说明

调用示例

from llm_datadist import LLMDataDist, LLMRole, LLMConfig llm_datadist = LLMDataDist(LLMRole.PROMPT, 0) llm_config = LLMConfig() llm_config.enable_cache_manager = True llm_config.device_id = 0 engine_options = llm_config.generate_options() llm_datadist.init(engine_options)

返回值

正常情况下无返回值。

异常情况会抛出LLMException。

参数错误可能抛出TypeError或ValueError。

约束说明

初始化成功后，系统退出前需要调用finalize。

finalize

函数功能

释放LLMDataDist。

函数原型

finalize()

参数说明

无

调用示例

from llm_datadist import LLMDataDist, LLMRole, LLMConfig llm_datadist = LLMDataDist(LLMRole.DECODER, 0) llm_config = LLMConfig() llm_config.enable_cache_manager = True llm_config.device_id = 0 engine_options = llm_config.generate_options() llm_datadist.init(engine_options) llm_datadist.finalize()

返回值

无

约束说明

• 初始化成功后，系统退出前需要调用finalize。

• finalize不能和其他接口并发调用。

link_clusters

函数功能

单边建链，由Client单侧发起建链。是Client还是Server与角色prompt或者decoder无关。设置listen_ip_info标识端口监听，即为Server端。

函数原型

link_clusters(clusters: Union[List[LLMClusterInfo], Tuple[LLMClusterInfo]], timeout=3000)

参数说明

调用示例

请参考样例运行。

返回值

正常情况下返回两个值的元组，第一个值是接口的返回值，类型是LLMStatusCode，第二个是每个集群建链结果的列表，类型是LLMStatusCode。

参数错误可能抛出TypeError或ValueError。

约束说明

• 建链的要求如下。建链数量过多存在内存OOM及KV Cache传输的性能风险。

  • 允许创建的最大通信数量=512。

• 建议超时时间配置200ms以上。如果TLS处于开启状态，建议超时时间配置为2000ms以上。查询TLS状态可以使用如下命令：

  hccn_tool [-i %d] -tls -g [host]

• 调用该接口前需提前注册所有内存，否则建链后注册不支持远端访问。

• 容器场景若未配置local_comm_res或配置为空，需在容器内映射“/etc/hccn.conf”文件或者确保默认路径“/usr/local/Ascend/driver/tools”下存在hccn_tool，如果两者都不能满足，则需要用户将hccn_tool所在路径配置到PATH中。配置实例如下，hccn_tool_install_path表示hccn_tool所在路径。

  export PATH=$PATH:{hccn_tool_install_path}

unlink_clusters

函数功能

单边断链，可以由Client单侧发起，通知Server进行断链；也可以Client和Server均发起强制断链，只清理本地链路。

函数原型

unlink_clusters(clusters: Union[List[LLMClusterInfo], Tuple[LLMClusterInfo]], timeout=3000, force=False)

参数说明

调用示例

请参考样例运行。

返回值

正常情况下返回两个值的元组，第一个值是接口的返回值，类型是LLMStatusCode，第二个是每个集群建链结果的列表，类型是LLMStatusCode。

参数错误可能抛出TypeError或ValueError。

约束说明

无

switch_role

函数功能

切换当前LLMDataDist的角色，同时可通过配置switch_options切换Client或者Server。

函数原型

switch_role(self, role: LLMRole, switch_options: Optional[Dict[str, str]] = None)

参数说明

调用示例

请参考样例运行。

返回值

• 正常情况下无返回值。
• 传入数据类型错误情况下会抛出TypeError或ValueError异常。
• 如果switch_role时存在残留链路资源，则会抛出LLMException，status_code为LLM_EXIST_LINK。

约束说明

无

link

函数功能

双边建链，通过建立通信域方式建链。（推荐使用更易用的单边建链的link_clusters接口，无需指定ranktable信息）

如果要在单机内且想使用RDMA网卡通信，需要设置HCCL_INTRA_ROCE_ENABLE=1环境变量。

函数原型

link(comm_name: str, cluster_rank_info: Dict[int, int], rank_table: str) -> int

参数说明

如上表格中ranktable具体信息请参见《HCCL集合通信库用户指南》。
如上表格中rank_table的配置示例如下:

{ "server_count": "2", "server_list": [ { "device": [ { "device_id": "0", "device_ip": "x.x.x.x", "rank_id": "0" }, ], "server_id": "xxxx" },{ "device": [ { "device_id": "0", "device_ip": "x.x.x.x", "rank_id": "1" }, ], "server_id": "xxxx" } ], "status": "completed", "version": "1.0" }

调用示例

请参考样例运行。

返回值

正常情况下返回标识通信域的ID。

参数错误可能抛出TypeError或ValueError。

约束说明

• 需要通信域内所有节点同时发起。
• 通信域内节点数量最大支持4。
• 通信域数量建议不超过16，最大支持512。建链数量过多存在内存OOM及cache传输的性能下降风险。
• link底层会在通信域内部交换内存描述符，对于用来传输的cache，需要在link前先配置内存池或者先注册内存。
• 不支持fork子进程方式调用。
• 最多支持16条链路并发建链，超过16条底层会排队。
• 需保证多通信域建链不出现循环依赖。
• Ascend 950PR/Ascend 950DT不支持该接口。

unlink

函数功能

双边建链对应的断链接口，调用此接口进行断链。

函数原型

unlink(comm_id: int)

参数说明

调用示例

请参考样例运行。

返回值

正常情况下无返回值。

异常场景会抛出LLMException异常。

参数错误可能抛出TypeError或ValueError。

Ascend 950PR/Ascend 950DT不支持该接口。

约束说明

无

query_register_mem_status

函数功能

调用此接口查询注册内存状态。

函数原型

query_register_mem_status(comm_id: int) -> RegisterMemStatus

参数说明

调用示例

请参考样例运行。

返回值

正常场景下无返回值。

异常场景会抛出LLMException异常。

参数错误可能抛出TypeError或ValueError。

Ascend 950PR/Ascend 950DT不支持该接口。

约束说明

配合link使用。

cache_manager

函数功能

获取CacheManager实例。

函数原型

cache_manager()

参数说明

无

调用示例

from llm_datadist import LLMDataDist, LLMRole llm_datadist = LLMDataDist(LLMRole.DECODER, 0) ... llm_datadist.init(engine_options) cache_manager = llm_datadist.cache_manager

返回值

返回CacheManager实例。

约束说明

无

【免费下载链接】hixlHIXL（Huawei Xfer Library）是一个灵活、高效的昇腾单边通信库，面向集群场景提供简单、可靠、高效的点对点数据传输能力。项目地址: https://gitcode.com/cann/hixl