지난 이틀 동안 여러 GPU 메모리 프로파일 플롯을 들여다보았습니다. GPU 메모리 프로파일을 어떻게 활용해야 할지에 대한 정보가 많지 않아서 한 줄로 처리하는 스크립트를 작성하고 블로그 게시물을 작성했습니다. 그러면 이 모든 이상한 모양은 무엇을 의미할까요? 빠른 튜토리얼🧵
먼저, 코드 블록을 프로파일링하려면 프로젝트에 이 Python 파일을 끌어다 놓고 아래와 같은 코드를 사용하여 프로파일링 데이터를 가져옵니다. https://t.co/iYTzEFksVR
아래와 같은 시각화가 나타납니다. 3개의 학습 단계에 대해 3개의 반복 패턴을 즉시 확인할 수 있습니다. 좋은 시작입니다. 이 줄무늬는 계층적으로 발생하는 메모리 연산을 나타냅니다. 아래의 상수 대역은 모델 매개변수 + 옵티마이저 상태입니다. 하지만 그 모든 것들이 무엇일까요?
첫 번째 학습 단계를 살펴보고 분석해 보겠습니다. 정방향 패스인 반삼각형으로 시작합니다. 그다음 거대한 빨간색 평행사변형인 로짓(logits) 생성으로 끝납니다. 처음에 뾰족하게 튀어나온 부분은 소프트맥스(softmax)를 위한 임시 메모리입니다. 로짓(logits)이 큰 이유는 어휘가 15만 1천 개나 되기 때문입니다.
다음으로 zero_grad()를 호출하여 이전 .grad 할당을 지웁니다. zero_grad()가 끝나면 작은 스파이크가 나타나고 그 후 큰 클리프가 나타납니다. 스파이크는 .backward()가 소프트맥스에서 시작할 때 발생하고, 클리프는 로짓 텐서가 더 이상 필요하지 않기 때문에 발생합니다. 이로 인해 LM 헤드 그라디언트(녹색 띠)가 생성됩니다.
위 그림에서 오른쪽의 물결 모양 패턴을 보세요. 이는 활성화 체크포인팅을 활성화했기 때문입니다. 활성화 체크포인팅으로 인해 각 계층에서 활성화를 다시 계산하면서 작은 스파이크가 발생하고, .backward() 함수가 계층별로 진행됨에 따라 입력 텐서가 해제됩니다.
마지막으로 제가 가장 좋아하는 부분입니다. optimizer.step()입니다. 깃대가 있는 사원처럼 정말 아름답습니다. 두 개의 기초 위에 세워져 있죠. 주황색 줄무늬는 실제로 .backward() 함수의 마지막 부분에서 생성된 임베딩 레이어 grad입니다.
위 그림에서 녹색 줄무늬는 Adam의 grad^2의 제곱근을 저장하는 데 필요한 저장 공간입니다. 융합된 Adam은 레이어별로 업데이트를 계산하여 줄무늬를 생성합니다. 맨 위에서는 업데이트를 한꺼번에 적용합니다. 스파이크는 업데이트를 적용하기 전에 나누기에 필요한 임시 버퍼입니다.
이 아름다운 시각화에 대해 더 이야기할 내용이 많습니다. 오류나 다른 유용한 정보가 있으면 알려주세요. 자세한 내용은 내 블로그에서 확인하세요: GPU 메모리 프로파일링을 가져오고 해석하는 방법 https://t.co/3uPt0S6RIp