TFLite로 Edge에서 LLM 추론 간소화

8월 13, 2024
Quentin Khan Software Engineer
Linkun Chen Software Engineer

XNNPack에 더 스마트한 캐시를 사용하여 첫 번째 토큰 및 피크 메모리 사용 시간 최적화


XNNPack은 모든 모델을 위한 기본 TensorFlow Lite CPU 추론 엔진입니다. 모바일, 데스크톱 및 Web 플랫폼 전반에서 게임 변경 속도를 더 높여줍니다. XNNPack에 사용되는 최적화 기법 중 하나는 Convolution, Depthwise Convolution, Transposed Convolution 및 Fully Connected 연산자의 정적 가중치를 추론 계산에 최적화된 내부 레이아웃으로 재구성하는 것입니다. 추론 중에 재구성된 가중치는 프로세서의 파이프라인에 적합한 순차적 패턴으로 액세스됩니다.

추론 지연 시간 감소에 수반되는 비용이 있는데, 재구성 시 기본적으로 XNNPack 내부 가중치의 복사본이 추가로 생성됩니다. 이전에는 XNNPack에 인메모리 캐시를 추가하여 이 비용을 절감하기 위해 노력했습니다. 이 캐시를 사용하면 동일한 모델을 독립적으로 실행하는 독립 TFLite 인터프리터 간에 구성된 가중치를 공유할 수 있습니다.

기존 캐시의 일부 단점을 해결하기 위해 TFLite XNNPack 대리자 구현을 개선했습니다.


1. 캐시는 익명 메모리에 저장되며, 메모리가 부족해질 경우 디스크로 옮겨지므로 성능 저하가 발생할 수 있습니다.

2. 프로세스를 시작할 때마다 초기 가중치를 재구성해야 합니다.

3. 재구성 시 원래의 TFLite 가중치를 읽고 새 버퍼에 쓰므로 구성 중 피크 메모리 사용량이 늘어나게 됩니다.

4. XNNPack 대리자를 통해 캐싱을 올바르게 사용하려면 여러 장황한 단계를 거쳐야 하고 세심한 수명 주기 관리도 필요합니다.

5. 프로세스 전반에서 가중치를 공유할 수 없습니다.

TFLite XNNPack delegate architecture
.

새로운 XNNPack 캐시 공급자 인터페이스

XNNPack이 업데이트되었으며 가중치 캐시 공급자를 구현할 수 있는 인터페이스를 제공합니다. 가중치 캐시 공급자는 구성된 버퍼에 액세스하기 위해 XNNPack이 채우고 쿼리하는 사전으로서 작동합니다. 주요 기능은 다음과 같습니다.

  • look_up은 구성된 버퍼 키를 조회하고 나중에 버퍼 주소를 검색하는 데 사용될 수 있는 고유 식별자(또는 NotFound에 예약된 특수 식별자)를 반환합니다.

  • reserve_space는 주어진 크기의 정보를 저장하는 데 사용될 수 있는 버퍼를 예약합니다. 그런 다음 look_up_or_insert를 사용하여 해당 버퍼를 커밋해야 합니다.

  • look_up_or_insert는 주어진 키와 일치하는 버퍼가 캐시 공급자에 존재하는지 확인합니다. 존재하지 않은 경우, 주어진 데이터는 캐시 공급자에 커밋됩니다. 이 함수는 또한 버퍼 주소 검색에 사용할 수 있는 식별자도 반환합니다.

  • offset_to_addr은 look_up과 look_up_or_insert에 의해 반환된 식별자로부터 버퍼 주소를 반환합니다.

다음 다이어그램은 XNNPack과 가중치 캐시 공급자 간의 상호작용을 보여줍니다.

The interactions between XNNPack and the weight cache provider
.

TFLite Delegate의 mmap을 사용하여 디스크에서 캐시 로드

TFLite Delegate는 이제 이 새로운 인터페이스를 사용하며 자체 가중치 캐시 공급자를 가지고 있습니다. 이 공급자는 구성된 가중치를 디스크와의 사이에서 직접 저장하고 로드할 수 있습니다. TFLite는 오랫동안 플랫버퍼 및 파일 백업 메모리 매핑을 활용해 왔습니다. 다음과 같은 이점 때문에 저희는 동일한 기법을 활용하여 그 간극을 메우고 있습니다.


재구성 오버헤드가 제거됩니다.

구성된 가중치를 디스크에 유지하면 모델 로드 시마다 비용이 많이 드는 재구성 프로세스를 우회할 수 있습니다. 이는 곧 시작 지연 시간과 피크 메모리 사용량이 모두 크게 감소한다는 뜻입니다. 초기 빌드 작업의 경우에도 이를 통해 구성된 데이터 중복을 삭제하고 같은 데이터를 재구성할 필요가 없도록 함으로써 구성 성능을 더욱 향상시킬 수 있습니다.


메모리 관리가 개선됩니다.

mmap은 운영 체제의 가상 메모리 관리를 활용하여 운영 체제에서 전체 시스템 메모리 사용량과 성능을 최적화할 수 있게 해줍니다. 저희의 경우를 예로 들자면, 신경망 연산의 일정한 가중치같이 대량의 읽기 전용 파일에 무작위로 액세스하는 데 특히 유리합니다.

구성된 데이터가 디스크에 저장되면 XNNPack 캐시는 메모리 부족에 시달리는 상황에서 성능 문제가 발생하기 쉬운 익명 메모리에 더 이상 의존하지 않습니다. 대신, XNNPack은 더 원활한 연산을 위해 운영 체제의 가상 메모리 관리 기능을 활용합니다.

mmap은 파일 시스템과 메모리 간에 데이터를 복사할 필요가 없으므로 상당한 수준으로 오버헤드를 줄이고 액세스 시간을 단축합니다.

mmap의 매뉴얼 페이지와 기타 흥미로운 읽을거리에서 파일 매핑과 메모리 사용량에 대한 자세한 정보를 직접 확인하실 수 있습니다.


프로세스 간 협업이 가능합니다.

각 프로세스의 가상 주소 공간이 동일한 물리적 메모리 페이지에 매핑됨에 따라 mmap 기반 파일 로딩이 여러 프로세스 간의 원활한 가중치 공유를 위한 문을 열어 줍니다. 이를 통해 여러 프로세스가 동일한 메모리를 공유하므로 전체 메모리 풋프린트가 감소할 뿐만 아니라 전면적으로 모델 로딩이 가속화될 수도 있습니다.

mmap-based file loading architecture
.

사용자 대상 API를 단순화합니다.

사용자에게 애플리케이션 수명 내내 캐시 객체를 설정하고 관리하도록 요구하는 대신, 캐시 파일의 경로를 간단히 제공할 수 있습니다.

std::unique_ptr<tflite::Interpreter> interpreter;
// XNNPack 대리자에 대한 옵션을 설정합니다.
TfLiteXNNPackDelegateOptions xnnpack_options = TfLiteXNNPackDelegateOptionsDefault();
xnnpack_options.weight_cache_file_path = "/tmp/cache_file.xnn_cache";
// XNNPack 대리자를 생성하여 TFLite 인터프리터에 적용합니다.
// 정적 가중치는 최초 실행 시에 구성되어 weights_cache에 기록됩니다.
// 정적 가중치는 다른 모든 실행에 대해 자동으로 로드됩니다.
TfLiteDelegate* delegate = TfLiteXNNPackDelegateCreate(&xnnpack_options);
interpreter->ModifyGraphWithDelegate(delegate);

캐시 무결성 유지

정확하고 효율적인 추론을 보장하려면 특정 조건에서 XNNPack 캐시를 무효화하는 것이 중요합니다.

모델 진화: 모델의 가중치나 구조가 변경되는 경우 캐시된 데이터가 오래되어 무효화되어야 합니다. 이는 제공된 캐시 경로에서 파일을 제거하다는 의미입니다.

XNNPack 업그레이드: XNNPack의 내부 구성 알고리즘을 업데이트하면 캐시된 가중치가 호환되지 않아 캐시를 다시 계산해야 할 수 있습니다. 다행히 XNNPack은 이를 감지할 수 있어 기존 캐시를 자동으로 바꿉니다.

본질적으로, XNNPack에서 가중치를 구성하거나 활용하는 방식에 영향을 미칠 수 있는 수정 사항은 캐시 무효화를 트리거해야 합니다.


벤치마크

세션 초기화는 가중치 구성의 지배를 받습니다. LLM의 경우 여러 하위 그래프가 동일한 가중치를 재사용하고 있습니다. 중복 삭제 기능이 동일한 가중치를 여러 번 구성하지 못하도록 하므로 캐시 작성 속도가 더 빨라집니다. 안정적인 확산 같은 더 많은 표준 모델의 경우 중복 삭제가 없습니다. 초기화 시간이 약간 더 늘어나는 것은 캐시를 디스크에 저장하기 때문입니다. 두 번째 실행부터는 캐시를 다시 로드하면 모든 경우에 초기화 시간이 전보다 훨씬 단축됩니다.

세션 초기화 개선은 당연히 LLM의 첫 번째 토큰까지의 시간에 영향을 미치는데, 벤치마크에서 대략 절반 정도로 시간이 단축됩니다.

캐시 구현으로 인한 메모리 증가도 확인할 수 있습니다. 중복 삭제 덕분에 LLM의 RSS(Resident Set Size) 값이 낮아집니다. 중복 삭제의 이점을 누리지 못하는 다른 모델의 경우에는 변화가 없습니다. 캐시를 다시 로드하면 TFLite 원본 모델이 더 이상 읽히지 않아 메모리로 가져오지 않게 되므로 피크 RSS가 훨씬 더 낮아집니다.


Pixel 8 Pro의 Gemma 2B

Benchmarks - Gemma 2B on a Pixel 8 Pro
.

Pixel 8 Pro의 Phi2

Benchmarks - Phi2 on a Pixel 8 Pro
.

Pixel 8 Pro의 안정적인 확산

Stable Diffusion on a Pixel 8 Pro
.

향후 작업

현재, 캐시는 파일 시스템 사용에 연계되어 있습니다. 기존의 할당된 메모리를 파일 백업 매핑과 교환하지 않으려는 사용 사례에 대해 데이터 중복 삭제 메커니즘을 독립적으로 사용해 그 이점을 누릴 수 있으면 좋겠습니다. mmap을 사용하면 익명 매핑을 만들 수 있어 구현을 대부분 재사용할 수 있습니다.