__init__接口

接口功能

提供Generator初始化,包括Model和Sampler初始化,以及权重加载、KV Cache分配等功能。

接口实现

class GeneratorTorch(GeneratorBackend):
    cache_pool: CachePool = None
    def __init__(self, model_config):
        super().__init__(model_config)
        self.tokenizer = self.model_wrapper.tokenizer
        self.device = self.model_wrapper.device
        self.rank = self.model_wrapper.rank

参数说明

model_config为模型配置,支持字典格式输入,主要配置信息参考如下:

参数名称

是否必选

配置名称

类型

默认值

描述

model_config

必选

backend_type

string

None

模型后端类型,支持'atb'和'ms'。

world_size

int

None

TP并行数。

rank

int

None

当前进程在TP并行进程中的序号,从0开始计数,此值应小于world_size。

npu_device_id

int

None

当前进程所使用的可见npu设备序号,此值应小于环境变量ASCEND_RT_VISIBLE_DEVICES设置的可见设备数量。(未设置ASCEND_RT_VISIBLE_DEVICES时,默认所有设备可见。)

num_threads

int

8

后处理线程数。

可选

load_tokenizer

bool

True

是否加载tokenizer,设为False则不加载默认的tokenizer。

tokenizer_path

string

None

自定义加载tokenizer的路径,为None时默认使用模型权重路径。

安全说明

调用API 出现异常情况时,会直接抛出异常信息。例如: