GPT-2 실행하기
이제 Python 텐서 언어, MLIR 빌더, 변환 파이프라인, 네이티브 실행 경계가 모두 갖추어졌습니다. 남은 작업은 이 구성 요소들을 사전 학습된 모델에 연결하고, 그 결과가 원래의 NumPy 구현과 일치하는지 확인하는 것입니다.
이 예제는 소형 GPT-2 체크포인트를 사용합니다. 이 체크포인트는 12개의 트랜스포머 블록, 768의 임베딩 너비, 12개의 어텐션 헤드, 50,257개 토큰의 어휘로 구성됩니다. 동반 프로젝트의 다운로드 스크립트는 특정 체크포인트 리비전을 고정하므로, 나중에 원격 파일이 변경되더라도 예제가 모르는 사이에 달라지지 않습니다.
tiny-mlir 체크아웃에서 다음을 실행하십시오.
uv sync
uv run python fetch_model.py model
uv run python example.py "Alan Turing theorized that computers would one day become" --tokens 10
모델 다운로드 용량은 약 550 MB입니다. 컴파일과 실행에는 uv sync로 설치되는 MLIR Python 바인딩을 사용하므로, 예제를 실행하는 데 별도의 명령줄 LLVM 설치가 필요하지 않습니다.
최상위 예제
진입점은 프롬프트, 생성 길이, 그리고 선택적인 체크포인트 디렉터리를 인자로 받습니다. 이 인자들을 파싱한 뒤 모델을 불러오고 생성된 텍스트를 출력합니다.
model, tokenizer = load_model(args.model)
print(generate(model, tokenizer, args.prompt, args.tokens))
참조 예제를 실행하려면 프롬프트를 명령줄에 직접 전달하십시오.
uv run python example.py "Alan Turing theorized that computers would one day become" --tokens 10
컴파일된 모델과 NumPy 참조 구현이 내놓는 완성 결과는 다음과 같습니다.
Alan Turing theorized that computers would one day become the most powerful machines on the planet.
마침표 뒤에는 두 개의 개행 토큰이 있습니다. 이 토큰들은 일반적으로 렌더링된 문장에서는 보이지 않지만, 10개 토큰으로 이루어진 출력의 일부입니다. 디코딩된 텍스트뿐만 아니라 토큰 ID도 확인하겠습니다.
매개변수 불러오기
load_model은 config.json, model.safetensors, tokenizer.json을 읽습니다. 모델 생성자는 필요한 모든 매개변수의 형상을 검사하며, 체크포인트 배열이 float32일 것을 요구합니다. 또한 임베딩 너비가 헤드 수로 나누어떨어지는지도 검사합니다. 로더는 활성화 함수나 레이어 정규화 엡실론이 다르면 잘못된 공식을 조용히 적용하는 대신 이를 거부합니다.
임베딩 너비가 인 블록에서 주요 프로젝션의 형태는 다음과 같습니다.
| 매개변수 | 형태 |
|---|---|
| 쿼리/키/값 프로젝션 | |
| 어텐션 출력 프로젝션 | |
| 피드포워드 확장 | |
| 피드포워드 출력 프로젝션 |
각 프로젝션에는 출력 너비에 맞는 바이어스 벡터도 있습니다. 두 레이어 정규화에는 길이가 인 게인 벡터와 바이어스 벡터가 있습니다.
배열은 여전히 컴파일된 함수의 런타임 입력으로 남습니다. 가중치 값의 집합마다 같은 블록을 별도로 복사해 컴파일하지는 않습니다. 호환되는 형상과 dtype은 하나의 특수화를 공유합니다.
임베딩
초기 잔차 스트림은 토큰 임베딩과 학습된 위치 임베딩을 더합니다:
이 역시 작은 컴파일된 함수입니다:
@jit
def embed(ids, words, positions):
return gather(words, ids) + positions[:ids.shape[0]]
gather 연산의 열 인덱싱은 아핀이며, 행 인덱스는 토큰 ID 텐서에서 가져옵니다. 이 연산의 linalg.generic 영역은 tensor.extract로 선택된 테이블 항목을 로드합니다. 모델은 네이티브 코드에 진입하기 전에 토큰 ID가 어휘 범위 안에 있는지 검사하며, JIT의 gather 경계에서도 인덱스를 검사합니다.
위치 임베딩은 현재 시퀀스 길이에 해당하는 정적 슬라이스입니다. 토큰 시퀀스는 비어 있지 않아야 하며 설정된 컨텍스트 윈도에 들어가야 합니다.
트랜스포머 스택
각 블록은 Part 11의 두 컴파일된 함수를 실행합니다:
x = x + attention(layer_norm(x))
x = x + feedforward(layer_norm(x))
실제 호출에서는 각 블록의 정규화 매개변수와 투영 배열을 전달합니다. attend와 feedforward는 평범한 Python 정의이며, 그 텐서 표현식은 함께 추적되고 로워링됩니다. 모델 루프 자체는 Python에 남아 있으므로 블록의 순서와 각 블록의 매개변수를 쉽게 살펴볼 수 있습니다.
첫 번째 어텐션 프로젝션은 float32 임베딩에서 시작합니다. 레퍼런스의 산술 연산은 이후의 값을 float64로 승격하며, 우리의 타입 그래프도 그 동작을 따릅니다. 따라서 첫 번째 블록은 차원이 같더라도 이후 블록과 다른 특수화가 필요할 수 있습니다.
마지막 블록 다음에는 정규화를 수행하고, 같은 토큰 임베딩 행렬을 사용해 어휘로 프로젝션합니다:
@jit
def project(x, gain, bias, words):
return layer_norm(x[-1:], gain, bias) @ words.T
다음 토큰을 선택하는 데에는 마지막 위치의 로짓만 필요합니다. 그래서 이 함수는 (1, vocab_size) 형상을 반환합니다. 레퍼런스는 모든 위치의 로짓을 계산하며, 검증 과정에서는 그 마지막 행을 우리의 결과와 비교합니다.
디코딩
이 예제는 탐욕적 디코딩을 사용합니다. 각 단계에서 가장 큰 로짓의 인덱스를 고르고, 그 토큰을 덧붙인 다음, 더 길어진 접두사를 평가합니다:
for _ in range(max_tokens):
ids.append(int(np.argmax(model(ids)[0])))
토큰화, 전체 흐름 조율, 그리고 마지막 argmax는 Python 쪽에 남습니다. 모델의 임베딩 조회, 정규화, 어텐션, 프로젝션, 활성화, 잔차 산술은 컴파일된 MLIR을 통해 실행됩니다.
NumPy 레퍼런스의 generate_tokens와 마찬가지로, 이 루프는 텍스트 끝 토큰에서 일찍 멈추지 않고 요청된 개수만큼 토큰을 생성합니다. 모델은 프롬프트와 요청된 생성 길이가 컨텍스트 한도에 들어가는지 검증합니다.
키/값 캐시는 구현하지 않았습니다. 각 단계에서는 접두부를 다시 계산하며, 길이가 새로 바뀔 때마다 새로운 정적 텐서 형태가 선택됩니다. 이는 참조 알고리즘을 그대로 구현한 것으로, 컴파일러의 동작을 따라가기에 유용합니다. 디코딩 캐시를 도입하면 모델 인터페이스와 어텐션 형태가 바뀌므로, 별도의 변경으로 검증해야 합니다.
모델 전체를 하나의 커널로
위에서 다룬 서브레이어별 경로는 이 시리즈가 지금까지 구축해 온 경로입니다. 이 프로젝트에는 두 번째 모드인 --fused가 있으며, 이 모드는 NumPy 참조 구현의 순방향 패스 자체를 컴파일합니다. tiny-gpt2의 gpt2 함수는 9부에서 다룬 배열 프로토콜을 통해 작성된 그대로 트레이스되며, np.mean 정규화, np.split 프로젝션, np.triu 마스크, 중첩된 매개변수 튜플도 모두 그대로 포함됩니다. 그 결과 열두 개의 블록을 모두 담은 하나의 모듈이 만들어지므로, 융합과 matmul 치환이 모델 전체를 한꺼번에 볼 수 있습니다.
시퀀스 길이마다 커널을 따로 만든다면 생성되는 토큰마다 컴파일을 한 번씩 해야 합니다. 그 대신 융합 모드는 토큰 시퀀스를 다음 2의 거듭제곱 길이로 패딩합니다. 인과 마스크가 실제 위치들이 패딩에 어텐션하지 못하도록 막기 때문에 마지막 실제 토큰에 대한 로짓 행은 변하지 않으며, 토큰 열 개를 생성할 때 모듈을 열 개가 아니라 두어 개만 컴파일합니다. 14부에서 이 모드를 더 자세히 살펴봅니다. 여기서 이 모드가 중요한 이유는 아래의 검증이 볼 수 있는 대상이 달라지기 때문입니다.
실제 참조 구현과 대조하기
문장이 일치한다는 것은 유용한 결과이지만, 수치 오류가 어디에서 상쇄되었거나 눈에 띄지 않고 지나갔는지는 알려 주지 않습니다. 검증 스크립트는 원본 tiny-gpt2 저장소를 그 자체의 Python 환경에서 실행하고, 디코딩 단계마다 모델을 캡처합니다.
uv run python verify.py
이 스크립트는 공개된 NumPy 참조 구현의 고정된 리비전을 내려받고, uv로 잠금된 환경을 설치한 뒤, 이에 맞는 토크나이저 파일을 가져옵니다. Git과 네트워크 접근이 필요하지만, 형제 체크아웃이나 수동으로 참조 구현을 설정할 필요는 없습니다. --reference /path/to/tiny-gpt2를 지정하면 기존 체크아웃을 선택적으로 사용할 수 있습니다. 두 구현 모두 동일한 고정 체크포인트를 사용합니다. 스크립트는 자신의 배열을 참조 구현의 매개변수 데이터클래스에 맞게 변환한 다음, 참조 구현에 이미 있는 수치 연산을 변경 없이 호출합니다. 또한 캡처한 트레이스를 참조 구현의 완전한 gpt2 함수와 대조하여 확인합니다.
열 단계 각각에 대해 비교 대상은 다음과 같습니다.
- 토큰 임베딩과 위치 임베딩.
- 모든 트랜스포머 블록의 두 잔차 출력.
- 마지막 위치의 어휘 로짓.
- 선택된 토큰 ID.
융합 모드에서 순방향 패스는 포착할 중간값이 없는 하나의 컴파일된 함수이므로, 최종 로짓과 선택된 토큰만 비교합니다. 블록별 검사는 서브레이어별 경로에 해당합니다. 지속적 통합에서는 두 모드를 모두 실행하며, GPU 타깃에서는 같은 스크립트를 --gpu와 함께 실행합니다.
동일한 dtype을 요구하며 배열을 rtol=3e-4와 atol=3e-4로 비교합니다. 토큰 ID와 디코딩된 완성 결과는 정확히 일치해야 합니다. 원본 저장소의 NumPy 2.4.4 환경을 대상으로 실행한 테스트에서, 캡처한 모든 배열에 걸친 최대 절대 차이는 약 였습니다.
생성된 ID는 정확히 다음과 같았습니다:
[262, 749, 3665, 8217, 319, 262, 5440, 13, 198, 198]
그리고 공백을 명시적으로 표시한 완성 결과는 다음과 같았습니다:
' the most powerful machines on the planet.\n\n'
이는 중간 부동소수점 계산에 대해서는 수치적 일치를, 요구되는 관찰 가능한 출력에 대해서는 정확한 일치를 의미합니다. 모든 부동소수점 값이 비트 단위까지 동일하다거나 가능한 모든 프롬프트를 테스트했다고 주장하는 것은 아닙니다.
우리가 만든 것
전체 패키지는 몇 개의 모듈에 담깁니다. expr.py에는 타입이 지정된 텐서 그래프가 있고, builder.py는 MLIR 연산을 구성하며, jit.py는 특수화와 컴파일러 패스를 처리하고, ops.py에는 트랜스포머 공식이 담겨 있으며, model.py는 이들을 체크포인트 및 디코딩 루프와 연결합니다. 세 개의 모듈이 이 경로를 변경하지 않고 그 위에 추가로 구축됩니다. blas.py는 matmul을 시스템 BLAS에 넘기고, fused.py는 참조 구현을 단일 커널로 추적하며, gpu.py는 NVIDIA 하드웨어를 대상으로 합니다.
중요한 특성은 컴파일러가 계산의 구성 요소를 이해한다는 점입니다. 행렬 곱셈은 행렬 연산으로 남고, 리덕션은 명시적인 축과 영역을 가지며, 브로드캐스트는 인덱싱 맵을 가집니다. GELU와 어텐션은 이러한 기본 요소들의 조합입니다. MLIR의 패스는 저장 방식과 실행 세부 사항이 확정되기 전에 그 구조를 변환할 수 있습니다.
이 장의 모든 내용은 CPU에서 실행되었습니다. 실행 가능한 최종 결과물은 NumPy 모델의 요구 출력을 재현하는 작은 네이티브 컴파일러이며, 레퍼런스 검사는 컴파일러에 대한 이후 변경을 판단할 구체적인 방법을 제공합니다. 다음 장에서는 그러한 변경 하나를 수행합니다: 동일한 텐서 프로그램, 동일한 공식, 동일한 검사를 GPU에서 실행합니다.
이로써 Python 텐서 표현식에서 실제로 실행되는 사전 학습된 트랜스포머에 이르는 경로가 완성됩니다. 동작하는 MLIR 컴파일러를 만들고 그것으로 GPT-2를 실행하신 것을 축하합니다!