Martin Mass, David G. Andersen, Michael Isard, Mohammad Mahdi Javanmard, Kathryn S. McKinley, Colin Raffel ACM ASPLOS 2025
개요
Huge page에서 Fragmentation을 줄이는 것이 매우 중요하다. 본 논문은 AI기술을 사용해서 Object lifetime을 추적, Fragmentation을 줄이는 기법을 제시하였다.
Motivation & Importance
기존 OS에서 사용하는 Huge page management는 life-time에 대한 고려가 없기 때문에 Fragmentation에 취약하다. C/C++프로그램 같은 경우에는 Pointer migration이 불가능하기 때문에, 한번 할당된 메모리를 Life-time에 따라서 재배치하는 기능또한 사용할 수 없기 때문에 문제가 심각해진다. 따라서 오브젝트 Lifetime을 Allocation전에 예측하고, 이에 따라서 Long-lived object와 Short-lived object을 묶어서 배치하여, Huge page를 사용하지만, Fragmentation을 줄일 수 있는 연구가 필요하다.
Main Idea
LSTM을 이용한 AI모델로 Object의 Life-time을 정확히 예측하고, 같은 Life-time을 가질 것으로 예상되는 오브젝트끼리 묶어서 Allocation을 하면 Fragmentation이 줄어들 것이다.
Challenge
- Lifetime Prediction Challenges
- 기존 방식들은 Allocation stie나 Matching calling context를 기준으로 Lifetime을 예측하였다. 이를 위해서 Runtime에 정보를 수집하거나, Offline analysis를 이용하였다. 그러나 Stack tracking과 같은 방식을 통해서 Runtime에 수집하는 기법은 Overhead가 대략 14%에 이를 정도로 크다는 단점이 있었다.
- Coverage and Accuracy
- 대부분의 경우 Offline analysis를 이용하는 방식은, Coverage가 낮기 때문에 Accuracy가 떨어진다는 문제가 있었다. Accuracy가 떨어지면, Lifetime 예측의 정확도가 떨어져, Fragmentation이 커진다는 단점이 있었다.
- Instability
- Stack tracing과 같은 방식은, Address space layout randomization과 같은 이유 때문에, 반복되는 수행에서 정확도를 보장하지 못한다. 즉, Application의 여러 수행에서, Prediction의 정확도가 계속 바뀌는 문제가 생긴다.
Design
여러개의 Execution을 Sampling하여서, Tokenized된 Calling context를 바탕으로 Object의 Lifetime을 예측하는 모델을 학습시키고, 학습된 모델을 바탕으로 발견하지 못한 새로운 Context에서도 Object의 Lifetime을 예측하도록 하였다. 또한 예측시에 발생하는 Overhead를 줄이기 위해서, Hasing-based mechanisms을 통해, 기존의 이미 예측된 Context의 경우에는 Fast-path로 예측하도록 하여서 오버헤드를 줄였다.
- Sampling-based Data Collection
- 5분정도 Stack trace, Object size, Address at allocation and deallocation time을 이용하여 LSTM모델에 Lifetime을 예측할 수 있도록 학습시켰다. Sampling은 TCMalloc을 Extend시켜서 진행하였다.
- Lifetime Prediction Model
- Input을 8개의 클래스로 구성된 Lifetime으로 예측하는 모델을 개발하였다. LSTM의 Input은 Stack trace로 나온 String이 들어가도록 하였다. LLM과 유사한 특징을 가지는 LSTM을 이용함으로써, 문자열을 보다 효율적으로 처리하도록 하였다.
- Lifetime Aware Allocator Design
- Huge page만을 이용하여서, 직접 Object의 Life-time에 따라서 Allocation을 수행하는 Allocator을 디자인 하였다. Fragmentation을 방지하기 위해서 Huge page를 8KB의 Object로 쪼개고, 그들의 Liveness를 추적하였다. LLAMA는 최대한 같은 Lifetime을 가지는 Object끼리 같은 Huge page를 공유하도록 Allocation을 수행하였다. 한 Huge page의 모든 Object가 Free되면, Huge page는 OS로 Free시켰다. 만약 Huge page가 Fragmented가 되면, 최대한 Short-life time을 가지는 Object들이 그 Fragmented된 Huge page에 Allocated되도록 Recycle하여서 Fragmentation을 줄였다. Lifetime예측이 틀릴 수도 있기 때문에, 틀린 Lifetime에 대한 처리를 하는 것도 중요하다. 이러한 잘못된 예측을 LLAMA에서는 Deadline을 이용하여 해결하였다. 만약 데드라인 예측이 틀리면, LLAMA는 Huge page의 Lifetime class를 보다 상위의 Class로 승격시키고, 새로운 데드라이을 추가하였다. 이를 통해서 잘못된 예측이 있더라도 동적으로 복구되도록 하였다. (약간 FFMalloc스타일의 Allocation인듯?)
- Low-Latency and Accurate Prediction
- TCMalloc은 Allocation은 100cycle도 안돼서 수행한다. 그러나 AI모델을 빨라도 몇 Microsecond가 걸린다. 이 문제를 해결하기 위해서 Allocation단계에서 Lifetime예측 모델의 Output을 Hash를 통해서 체크할 수 있도록 하여서, 같은 조건에서는 hash hit이 나고, Fast-path로 Allocate가 가능하도록 하였다. 또한 Cache된 값이 나중에 바뀐 환경에서도 잘못된 값을 캐싱하는 문제를 해결하기 위해서 (대략 14%정도의 Allocation이 Cache된 값이 바뀌는 것으로 조사됨), 1000개의 Hit마다 새로 Cache값을 계산하도록 하였다.
Conclusion
본 논문을 읽고 AI모델과 Allocator의 관계에 대해서 새롭게 생각할 수 있는 계기가 되었다. AI모델을 System의 분야라고 생각했던 많은 문야에 적용하는 연구는 매우 Promising할 것으로 생각된다.