디바이스 관리

디바이스 관리 (Device Management)

이 페이지에서는 Haystack 관점에서 디바이스(device) 관리가 무엇인지 다뤄볼게요. 로컬 머신에서 모델 추론을 돌리는 일이 얼마나 번거로울 수 있는지부터, Haystack이 그 문제를 어떻게 풀었는지 순서대로 설명합니다.

출처: 공식문서

왜 디바이스 관리를 신경 써야 할까요?

HuggingFaceLocalGenerator, AzureOpenAIGenerator 등 많은 Haystack 컴포넌트는 어떤 언어 모델을 호출하고 실행할지를 사용자가 고를 수 있게 해 줘요. 클라우드 기반 서비스와 연동하는 컴포넌트라면, 필요한 하드웨어(GPU 같은 것)를 프로비저닝하는 세부 사항은 서비스 제공자가 알아서 처리합니다. 하지만 로컬 머신에서 모델을 쓰려면 내 하드웨어에 어떻게 배포할지 스스로 해결해야 해요. 게다가 ML 라이브러리마다 특정 디바이스에서 모델을 실행하는 API가 제각각이라 더 복잡해지죠.

로컬 모델에서 추론을 실행하는 과정을 최대한 단순하게 만들기 위해, Haystack은 프레임워크에 종속되지 않는(framework-agnostic) 디바이스 관리 구현을 사용합니다. 이 인터페이스를 통해 디바이스를 노출하면, 더 이상 라이브러리별 호출 방식이나 디바이스 표현을 신경 쓸 필요가 없어요.

개념 (Concepts)

Haystack의 디바이스 관리는 다음과 같은 추상화 위에 세워져 있어요:

• DeviceType — 지원되는 모든 디바이스 유형을 나열하는 열거형(enumeration)이에요.
• Device — DeviceType과 고유 식별자로 구성된 디바이스의 일반적인 표현이에요. 이 둘이 합쳐져 사용 가능한 전체 디바이스 중 하나의 디바이스를 나타냅니다.
• DeviceMap — 문자열을 Device 인스턴스에 매핑하는 구조예요. 문자열은 보통 모델 파라미터 같은 모델별 식별자를 나타내며, 모델의 특정 부분을 특정 디바이스에 배치할 수 있게 해 줍니다.
• ComponentDevice — 하나의 Device 또는 DeviceMap 인스턴스를 담는 태그된 유니언(tagged union)이에요. 로컬 추론을 지원하는 컴포넌트는 생성자에서 이 타입의 선택적 device 파라미터를 노출합니다.

이 추상화들 덕분에 Haystack은 내 로컬 머신에 있는 어떤 지원 디바이스든 온전히 다룰 수 있고, 동시에 여러 디바이스를 쓰는 것도 지원해요. 로컬 추론을 지원하는 모든 컴포넌트는 이 일반적인 표현을 백엔드별 표현으로 변환하는 일을 내부에서 처리합니다.

이 추상화들의 전체 코드는 Haystack GitHub repo에서 확인할 수 있어요.

사용법 (Usage)

아래 예제는 transformers-haystack 통합에 포함된 HuggingFaceLocalGenerator를 사용합니다. 먼저 설치해 주세요:

pip install transformers-haystack

추론에 단일 디바이스를 쓰려면 ComponentDevice.from_single 또는 ComponentDevice.from_str 클래스 메서드를 사용합니다:

from haystack.utils import ComponentDevice, Device

device = ComponentDevice.from_single(Device.gpu(id=1))
# 또는 PyTorch 디바이스 문자열 사용
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)

컴포넌트는 ComponentDevice 타입의 선택적 device 파라미터를 노출해야 합니다. 일단 노출되고 나면 컴포넌트는 이 값을 어떻게 처리할지 결정할 수 있어요:

• device=None이면 그 값을 그대로 백엔드에 넘깁니다. 이 경우 백엔드가 모델을 어느 디바이스에 배치할지 결정해요.
• 또는 컴포넌트가 ComponentDevice.resolve_device 클래스 메서드로 백엔드에 넘기기 전에 사용 가능한 디바이스를 자동으로 고를 수도 있어요.

디바이스가 결정되고 나면 컴포넌트는 ComponentDevice.to_* 메서드로 백엔드별 표현을 얻어 백엔드에 전달합니다.

ComponentDevice 인스턴스는 컴포넌트의 to_dict와 from_dict 메서드에서 직렬화되어야 해요.

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를 제공한다면, 그 통로(conduit) 역할을 하는 추가 init 파라미터를 넣을 수도 있어요. 예를 들어 HuggingFaceLocalGenerator는 Hugging Face 전용 device_map 인자를 전달할 수 있는 huggingface_pipeline_kwargs 파라미터를 노출합니다:

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

이런 경우 파라미터 우선순위와 선택 동작이 명확히 문서화되어야 해요. HuggingFaceLocalGenerator의 경우 huggingface_pipeline_kwargs로 전달된 device map이 명시적인 device 파라미터를 덮어쓰며, 그렇게 동작한다고 문서화되어 있습니다.

더 알아보기 (Learn more)

• 공식문서: Device Management