许多 Haystack 组件，例如HuggingFaceLocalGenerator , AzureOpenAIGenerator，以及其他组件，允许用户选择要查询和执行的语言模型。对于与云服务交互的组件，服务提供商会自动处理配置所需硬件（如 GPU）的详细信息。但是，如果您希望在本地计算机上使用模型，则需要弄清楚如何在自己的硬件上部署它们。更复杂的是，不同的 ML 库有不同的 API 来在特定设备上启动模型。

为了尽可能简化在本地模型上运行推理的过程，Haystack 使用了框架无关的设备管理实现。通过此接口公开设备意味着您不再需要担心特定于库的调用和设备表示。

概念

Haystack 的设备管理建立在以下抽象之上：

DeviceType - 一个枚举，列出了所有支持的不同设备类型。
Device - 一个通用设备表示，由一个DeviceType 和一个唯一标识符组成。它们共同表示所有可用设备组中的单个设备。
DeviceMap - 一个将字符串映射到Device 实例的映射。字符串代表特定于模型的标识符，通常是模型参数。这使我们将模型的特定部分映射到特定设备。
ComponentDevice - 一个标记联合，由一个单独的Device 或一个DeviceMap 实例组成。支持本地推理的组件将在其构造函数中公开一个此类型的可选device 参数。

通过上述抽象，Haystack 可以完全寻址本地机器中的任何支持设备，并能够同时使用多个设备。每个支持本地推理的组件都会在内部处理将这些通用表示转换为其后端特定的表示。

📘
源代码
您可以在 Haystack GitHub 存储库中找到上述抽象的完整代码。

用法

要使用单个设备进行推理，请使用ComponentDevice.from_single 或ComponentDevice.from_str 类方法。

from haystack.utils import ComponentDevice, Device

device = ComponentDevice.from_single(Device.gpu(id=1))
# Alternatively, use a PyTorch device string
device = ComponentDevice.from_str("cuda:1")
generator = HuggingFaceLocalGenerator(model="llama2", device=device)

要使用多个设备，请使用ComponentDevice.from_multiple 类方法。

from haystack.utils import ComponentDevice, Device, DeviceMap

device_map = DeviceMap({
	"encoder.layer1": Device.gpu(id=0),
	"decoder.layer2": Device.gpu(id=1),
	"self_attention": Device.disk(),
	"lm_head": Device.cpu()
})
device = ComponentDevice.from_multiple(device_map)
generator = HuggingFaceLocalGenerator(model="llama2", device=device)

在自定义组件中集成设备

组件应公开一个类型的可选device 参数，该参数类型为ComponentDevice。公开后，它们可以决定如何处理它：

如果当 device=None 时，组件可以将该参数传递给后端。在这种情况下，后端决定将模型放置在哪个设备上。
或者，组件可以尝试在将设备传递给后端之前自动选择一个可用设备，使用ComponentDevice.resolve_device 类方法。

一旦设备被解析，组件就可以使用ComponentDevice.to_* 方法获取底层设备的后端特定表示，然后将其传递给后端。

该应在组件的to_dict 和from_dict 方法中序列化 ComponentDevice 实例。

from haystack.utils import ComponentDevice, Device, DeviceMap

class MyComponent(Component):
    def __init__(self, device: Optional[ComponentDevice] = None):
        # If device is None, automatically select a device.
        self.device = ComponentDevice.resolve_device(device)

    def warm_up(self):
        # Call the framework-specific conversion method.
        self.model = AutoModel.from_pretrained("deepset/bert-base-cased-squad2", device=self.device.to_hf())

	  def to_dict(self):
	    # Serialize the policy like any other (custom) data. 
	    return default_to_dict(self,
							 device=self.device.to_dict() if self.device else None,
							 ...)
	
	  @classmethod
	  def from_dict(cls, data):
	    # Deserialize the device data inplace before passing 
			# it to the generic from_dict function.
	    init_params = data["init_parameters"]
      init_params["device"] = ComponentDevice.from_dict(init_params["device"])
			return default_from_dict(cls, data)

# Automatically selects a device.
c = MyComponent(device=None)

# Uses the first GPU available.
c = MyComponent(device=ComponentDevice.from_str("cuda:0"))

# Uses the CPU.
c = MyComponent(device=ComponentDevice.from_single(Device.cpu()))

# Allow the component to use multiple devices using a device map.
c = MyComponent(device=ComponentDevice.from_multiple(DeviceMap({
      "layer1": Device.cpu(),
      "layer2": Device.gpu(1),
      "layer3": Device.disk()
})))

如果组件的后端提供了更专业的 API 来管理设备，它可以添加一个额外的 init 参数作为通道。例如，HuggingFaceLocalGenerator 通过huggingface_pipeline_kwargs 参数公开了一个参数，通过该参数可以传递 Hugging Face 特定的device_map 参数。

generator = HuggingFaceLocalGenerator(model="llama2", huggingface_pipeline_kwargs={
	"device_map": "balanced"
})

在这种情况下，请确保参数的优先级和选择行为已清楚记录。就HuggingFaceLocalGenerator 而言，通过huggingface_pipeline_kwargs 参数传递的设备映射会覆盖显式的device 参数，并按此记录。

概念

📘源代码

用法

在自定义组件中集成设备

📘
源代码