파이토치 버퍼(register_buffer) 정리: 모델을 GPU로 옮길 때 마스크가 따라가지 않는 오류 해결법

핵심 메시지

• 파이토치 모듈을 .to('cuda')로 옮기면 파라미터는 자동으로 GPU에 올라가지만, 평범한 텐서 속성은 CPU에 그대로 남는다.
• 인과 어텐션 마스크처럼 학습하지 않는 고정 텐서가 CPU에 남으면 '마스크는 CPU, 모듈은 CUDA' 형태의 런타임 오류가 난다.
• 해결책은 그 텐서를 nn.Parameter로 만드는 것이 아니라 register_buffer로 등록하는 것이다. 파라미터는 옵티마이저로 넘어가 학습 대상이 되기 때문이다.
• 버퍼로 등록하면 모듈을 GPU로 보낼 때 함께 이동하고, 나중에 다시 CPU로 되돌리는 것도 모듈 단위로 한 번에 처리된다.
• 마스크를 forward에서 매번 새로 만들지 않고 초기화 때 한 번 만들어 두는 것은 효율 때문이며, 발표자는 매번 생성할 경우 상당한 성능 손해를 본다고 설명한다.

쉽게 이해하기

대규모 언어 모델을 직접 구현해 본 사람이라면 파이토치 모듈 안에서 self.register_buffer(...) 같은 줄을 본 적이 있을 것이다. 이 영상은 그 '버퍼'가 정확히 무엇이고 언제 필요한지를, 인과(causal) 셀프 어텐션 모듈을 직접 고쳐 가며 보여 준다. 어텐션 계산 자체를 이해할 필요는 없고, 모듈 안에 어떤 값이 들어 있는지만 따라가면 된다.

출발점은 버퍼를 쓰지 않은 어텐션 모듈이다. 입력 배치와 모듈을 각각 GPU로 옮긴 뒤 순전파를 돌리면, 곧바로 '모듈은 CUDA에 있는데 마스크는 CPU에 있다'는 런타임 오류가 발생한다. 코드에는 잘못된 부분이 없어 보이지만, 모듈이 들고 있는 값들이 서로 다른 장치에 흩어져 있는 것이다.

원인을 확인하는 방법은 단순하다. 선형 계층의 가중치와 마스크의 device를 각각 출력해 보면, 가중치는 CUDA로 옮겨졌지만 마스크는 여전히 CPU에 남아 있다. 가중치는 nn.Parameter라서 파이토치가 모듈 이동 시 자동으로 챙기지만, 상삼각 행렬로 만든 마스크는 그냥 일반 텐서라 아무도 챙겨 주지 않기 때문이다.

그렇다면 마스크도 nn.Parameter로 만들면 되지 않을까. 발표자는 그것이 답이 아니라고 말한다. 파라미터는 관례상 '학습으로 바뀌어야 하는 값'이고 옵티마이저에 전달되는데, 고정된 마스크는 학습할 이유가 없기 때문이다. 대신 register_buffer로 이름 문자열과 텐서를 등록하면, 학습 대상은 아니면서 모듈의 일부로 취급되어 장치 이동을 함께 따라간다.

버퍼로 바꾼 뒤 같은 코드를 다시 돌리면 오류 없이 GPU에서 순전파가 끝나고, 마스크의 device도 CUDA로 찍힌다. 모듈을 다시 CPU로 돌려도 마스크가 같이 돌아온다. 발표자는 이 짧은 편의성이 큰 모델에서 특히 중요하다고 정리한다. 모듈 안 일부만 GPU에 올라가고 나머지가 남으면 코드가 의도대로 동작하지 않기 때문이다.

주요 인사이트

• 파이토치에서 '모듈을 GPU로 옮긴다'는 말은 사실 파라미터와 버퍼를 옮긴다는 뜻이다. 그 밖의 속성은 개발자가 직접 관리해야 한다.
• 학습 여부가 파라미터와 버퍼를 가르는 기준이다. 값이 고정돼 있고 옵티마이저가 건드릴 이유가 없다면 파라미터가 아니라 버퍼가 맞다.
• 장치 불일치 오류는 보통 코드 로직이 아니라 '무엇이 모듈의 일부로 등록됐는가'의 문제다. 오류 메시지를 보고 텐서의 device를 하나씩 찍어 보는 습관이 진단을 빠르게 한다.
• 마스크를 forward마다 새로 만드는 대신 초기화 때 한 번 만들어 재사용하는 것은 성능 최적화이자, 버퍼가 필요해지는 직접적인 이유이기도 하다.
• 버퍼는 텐서의 종류가 바뀌는 것이 아니다. 여전히 평범한 텐서지만 모듈에 '등록'되어 함께 관리된다는 점이 달라진다.

자주 묻는 질문

관련 AI 소식