최근 대규모 언어모델(LLM)의 발전과 함께 학습에 필요한 컴퓨팅 자원, 특히 메모리(RAM) 문제가 점점 더 중요해지고 있습니다. 특히 긴 컨텍스트(context)와 방대한 어휘집(vocabulary)을 가진 LLM을 학습할 때, 모델의 가중치(weight) 자체보다도 크로스 엔트로피 손실(Cross Entropy Loss)을 계산하는 과정에서 발생하는 중간 텐서(tensor)들이 엄청난 메모리를 잡아먹는다는 사실이 주목받고 있습니다. 이는 학습 중 메모리 부족(Out Of Memory, OOM) 현상의 주요 원인으로 지목됩니다.
구체적으로, 128K(131,072) 토큰 길이의 컨텍스트를 사용하는 경우, 로짓(logits) 텐서 하나만으로도 약 40GB에 육박하는 메모리를 소비할 수 있습니다. 이는 160억 개 매개변수(16B) 모델의 전체 가중치 크기보다도 큰 수치입니다. 기존에는 시퀀스(sequence) 길이를 단순히 청크(chunk)로 나누어 처리하는 방식이 시도되었지만, 자동 미분(autograd) 시스템이 각 청크별 연산 그래프를 역전파(backward pass) 단계까지 유지하려 하기 때문에 피크 메모리(peak memory) 사용량을 효과적으로 줄이지 못했습니다. 이러한 문제를 해결하기 위해 FLCE(Fused Loss with Chunked Embeddings)와 같은 새로운 접근 방식이 제안되었습니다. FLCE는 각 청크의 기울기(gradient)를 순전파(forward pass) 과정에서 즉시 계산하여, 대형 텐서가 연산 그래프에 오래 남아있지 않도록 함으로써 메모리 효율을 극대화합니다.
이러한 메모리 최적화 기술은 LLM 학습의 효율성을 크게 향상시키며, 더 크고 복잡한 모델을 학습할 수 있는 길을 열어줍니다. 특히 제한된 하드웨어 자원을 가진 연구실이나 스타트업에게는 고성능 GPU 없이도 대규모 모델을 실험하고 개발할 수 있는 기회를 제공합니다. 메모리 사용량과 연산 지연 시간(latency) 간의 균형점을 찾는 것은 여전히 중요한 과제이지만, FLCE와 같은 기법들은 LLM 학습의 접근성을 높이고, 궁극적으로 AI 기술의 발전을 가속화하는 데 기여할 것입니다.