약 1년 전, AMD는 AI Max+ 시리즈 CPU (일명 Strix Halo)를 출시했습니다. 제 YouTube 피드는 Mac이 아닌 하드웨어에서 통합 메모리 (Unified Memory)를 확보하기로 한 아키텍처 결정에 대한 찬사로 가득 찬 것 같았습니다. 저는 결국 작년 11월에 128GB RAM을 탑재한 GMKTec EVO x2를 구매했습니다. 저는 다양한 튜토리얼을 따라 해보고, 비공식 릴리스를 통해 rocm을 설치하려고 시도하며, 로컬 LLM (Large Language Models)을 실행해 보기 시작했습니다!

저는 Strix Halo에서 1B 파라미터 LLM을 학습하려고 시도했던 저의 여정에 대해 쓰고 싶었습니다! 또한 그 과정에서 배운 모든 것들을 공유할 것입니다. 저에게 도움이 되었던 모든 주제에 대한 링크도 제공하겠습니다.

요약 (TL;DR):

이곳은 제가 Antigravity를 사용하여 만든 AI 학습 프레임워크를 위한 리포지토리입니다: https://github.com/genuinelucifer/onix
이 리포지토리에는 데이터셋 다운로드부터 사전 학습 (Pretraining), 미세 조정 (Finetuning), 로컬 실행에 이르기까지 필요한 모든 도구가 포함되어 있습니다.

또한 이 리포지토리에는 tiny-stories 데이터셋을 사용하여 모델을 학습하고 미세 조정하는 전체 프로세스를 설명하는 문서가 포함되어 있습니다. 프레임워크를 사용하여 로컬에서 실행하는 방법도 언급되어 있습니다: https://github.com/genuinelucifer/onix/blob/main/docs/training_llama1b_on_tinystories.md

참고: 프레임워크를 위한 코드만 AI로 작성되었습니다. 이 블로그는 제가 완전히 직접 작성했습니다.

시작 단계:

게임도 즐기고 싶었기 때문에 Windows 11이 설치된 미니 PC를 구매했습니다. 당시에는 Steam 게임에 대한 Proton 지원이 얼마나 발전했는지 알지 못했습니다. ROCm 드라이버를 설치하고 Windows에서 모델을 실행하려고 시도했는데, 그것은 실수였습니다. Windows 지원은 매우 부족하며, 유튜버들 중 누구도 (적어도 Strix Halo에서는) AI 모델을 실행하기 위해 Windows를 사용하는 것 같지 않았습니다.

배운 점: 로컬 AI 학습/실행을 위해서는 Linux (최신 릴리스)를 사용하세요.

Windows에서 ComfyUI나 Qwen-image-edit를 제대로 작동시키기 위해 몇 주 동안 시도하고 실패한 끝에, 마침내 Ubuntu를 듀얼 부팅(dual boot)했습니다. 덕분에 사람들이 Linux에서 모델을 실행하고 학습시키기 위해 만들어 놓은 방대한 튜토리얼에 접근할 수 있었습니다. Strix Halo에서 모델을 실행하기 위한 튜토리얼을 검색하는 대부분의 사람들은 Donato Capitella가 만든 멋진 strix-halo용 툴박스(toolboxes for strix-halo)를 거의 확실하게 발견하게 될 것입니다. 이것들은 제가 PC에서 모델을 실행하는 여정을 시작하는 데 큰 도움이 되었습니다.

하지만, 해당 툴박스의 문서는 BIOS 설정에서 전용 VRAM(dedicated VRAM)을 최소값으로 설정한 다음, grub 설정을 편집하여 RAM과 VRAM 사이의 공유 메모리(shared memory) 크기를 늘리기 위해 다음 내용을 추가하라고 요청합니다:

amd_iommu=off amdgpu.gttsize=126976 ttm.pages_limit=32505856

이 변경 사항을 성급하게 적용한 것이 저의 두 번째 실수였습니다. 이는 100GB 이상의 VRAM이 필요한 모델을 실행하는 데는 훌륭한 해킹(hack)일 수 있지만, 저는 그렇게 큰 VRAM이 필요한 모델을 실행해 본 적이 없었습니다. 저는 Strix Halo가 메모리 제한(memory limited)보다 연산 제한(compute limited)이 훨씬 더 심하다는 것을 알게 되었습니다. 훨씬 나중에, 전용 VRAM을 96GB로 설정하고 RAM을 32GB로 설정하는 것이 저에게 훨씬 더 나은 해결책이라는 것을 발견했습니다 (이에 대해서는 나중에 더 자세히 다루겠습니다).

배운 점: 다른 사람들에게 효과적이었던 설정만 보고 설정을 성급하게 최적화하지 마세요.

ComfyUI와 LM-Studio를 사용하여 몇 가지 모델을 실행해 보았는데, 30B 파라미터 (30B parameter) 모델에서 초당 약 40 토큰 (40 tokens/second) 정도의 속도가 나왔으며 이는 꽤 만족스러운 수준이었습니다. ComfyUI에서 FLUX 모델을 통해 1080p 이미지를 생성하는 데는 약 50초가 걸립니다. 따라서 이 PC는 로컬 코딩 에이전트 (local coding agents) 및 이미지 생성에 적합합니다. 하지만 비디오 생성은 여전히 역부족입니다. ComfyUI에서 Hunyuan video 모델을 사용하여 4초짜리 영상을 만드는 데 약 6시간이 걸렸습니다!

모델 실행에 대한 모든 실험을 마친 후, 저는 저만의 모델을 직접 학습시키고 싶었습니다. 단순히 기존 모델을 실행하는 것만이 제가 이 PC를 구매한 이유는 아니었기 때문입니다.

2단계: LLM에 대해 배우고 목표 설정하기

저는 이미 CNN, RNN 등에 대해 학습했었기에 모델이 어떻게 작동하는지에 대한 모든 기초 지식은 갖추고 있었습니다. 관련된 수학적 원리와 기본적인 모델 아키텍처 (architecture)가 어떻게 생겼는지도 알고 있었습니다.
제가 배우고 싶었던 것은 수학적인 세부 사항에 너무 매몰되지 않으면서, 실제 LLM을 구성하는 모든 구성 요소 (components)를 파악하는 것이었습니다. 또한 연구 (research) 관점보다는 소프트웨어 (software) 관점에서 학습하는 데 집중했습니다.

LLM이 어떻게 작동하는지에 대한 모든 기본 원리와 LLM 내부의 구성 요소가 무엇인지 배우기 위해 Sebastian Raschka의 저서인 Build a Large Language Model from Scratch를 구매했습니다. 이는 제 사용 사례에 완벽한 책이었습니다! 전통적인 방식(텍스트를 읽고 모든 코드 라인을 수동으로 작성하는 방식)으로 이 책을 끝내는 데 약 1.5개월이 걸렸습니다. 제가 작성한 코드는 my yallm repo에 업로드해 두었습니다.

그 이후, 저는 워크플로우 (workflow)를 가속화하고 싶었습니다. 또한 개인 프로젝트를 통해 AI 도구들에 대한 모든 화제성을 직접 확인해보고 싶었습니다. 그래서 Gemini Pro 구독을 시작하고, Antigravity를 설치한 뒤, 제 yallm 리포지토리의 코드를 에이전트 (agent)에게 전달하여 이제 모든 코드를 작성하도록 시키기로 결정했습니다. 또한 (제 생각에는) 앞으로 다양한 유형의 거대 언어 모델 (LLM)들을 학습시킬 예정이었기에, 제 PC에서 LLM을 처음부터 학습시킬 수 있는 범용 프레임워크 (generic framework)를 만들고 싶었습니다.

저는 이 프레임워크를 위해 코드를 전혀 작성하지 않는 것을 목표로 삼았습니다. 저는 (때때로) 코드를 검토하면서, 모델을 학습 (training), 미세 조정 (finetuning) 및 실행 (running)할 때 어떤 결과가 나오는지 확인하기만 할 것입니다.

그래서 저는 전속력으로 달려 Onix 프레임워크 (ONNX 모델 포맷을 활용한 가벼운 언어유희)를 만들었고, 이를 사용하여 제 Strix Halo에서 1B 파라미터 모델을 완전히 학습, 미세 조정 및 실행했습니다. 저는 이 프레임워크를 계속 발전시켜 나갈 계획이며 (앞으로는 100% AI를 통해서만 하지는 않겠지만), 더 많은 최적화와 함께 더 다양한 유형의 모델 아키텍처 (model architectures) 학습 지원을 추가하기를 희망합니다. 현재는 텍스트 기반 LLM, 이미지를 위한 VQ-VAE 임베딩 (embedding), 그리고 학습된 VQ-VAE를 사용하여 텍스트를 기반으로 이미지를 생성하는 "멀티모달 (multimodal)" 모델 학습을 지원합니다.

모델을 사전 학습 (pretrain)할 수 있는 영어 데이터셋을 찾아보았고, 첫 번째 데이터셋으로 tiny-stories 데이터셋을 선택했습니다. 이 데이터셋은 약 4억 7천만 개의 토큰 (tokens)을 가지고 있습니다. (1B 파라미터 모델을 학습시키려 한다는 점을 고려했을 때) 더 큰 데이터셋으로 넘어가기 전 유용한 테스트를 수행하기에 충분히 작은 규모로 보였습니다.

3단계: 모델 사전 학습을 위한 메모리 점유율 (Memory footprint)

그 후, 저는 시작 모델 아키텍처로 Llama를 선택했습니다. Llama의 아키텍처는 문서화가 매우 잘 되어 있고 재사용하기 쉽습니다. 저는 새로운 모델 아키텍처에서도 대부분의 코드를 재사용할 수 있도록 Sebastian Raschka의 저서에서 읽었던 GPT2 토크나이저 (tokenizer)를 고수했습니다. 그 후 사전 학습 과정을 시작했습니다.

1B Llama 모델로 학습을 시도할 때, 저는 **1024 토큰 컨텍스트 윈도우 (context window)**로 시작했습니다. 그랬더니 약 26 GB의 VRAM을 사용하는 것을 확인했습니다! 이는 안도감이 들면서도 동시에 충격적이었습니다. 1B 파라미터 모델이 왜 26 GB의 VRAM을 사용하는지 몰랐기에 충격이었고, 당시 저에게 110 GB 이상의 공유 메모리 (shared memory)가 있었기에 안도감이 들었습니다. 덕분에 컨텍스트 윈도우와 배치 크기 (batch size)를 모두 늘릴 수 있었습니다. 심지어 모델 크기를 키워 훨씬 더 큰 모델을 학습시킬 수도 있었습니다. 아, 그때는 참 순진했던 시절이었죠!

그래서 메모리 사용량을 살펴보았고, 결과는 다음과 같았습니다. 저는 기본 fp32 정밀도 (accuracy)로 모델을 학습시키고 있습니다. fp32 모델 파라미터를 저장하는 데는 약 4바이트 (32비트)가 필요합니다.

다음은 학습 중 각 항목이 어느 정도의 VRAM을 필요로 하는지에 대한 대략적인 추정치입니다 (정확한 값을 찾는 방법은 확실하지 않습니다):

1 - 총 정적 오버헤드 (Total Static overhead): (~16 GB)

• ~1B 모델 파라미터 (model parameters)
• ~1B 모델 그래디언트 (gradients)
• AdamW 옵티마이저 (optimizer)를 위한 ~2B 파라미터. 모델 파라미터당 모멘텀 (momentum)과 분산 (variance)이라는 2개의 값을 가집니다.

2 - 총 활성화 메모리 (Total Activation Memory): (~5GB)

• Self-attention matrix (~context size^2 * n_heads * n_layers)
• 피드 포워드 블록 (Feed forward block) 출력값

3 - 모델의 평가 (eval) 단계를 위한 캐시된 파라미터 (학습의 첫 번째 반복 이후에 평가를 수행하기 때문입니다).

4. Pytorch/rocm 캐싱 (Caching)에도 약간의 오버헤드가 있습니다.

배운 점: 100GB VRAM으로는 100B 모델을 학습시킬 수 없습니다! :(

그래서 저는 배치 크기를 8로 늘렸습니다. 그랬더니 이제 약 65 GB의 VRAM을 사용하는 것을 확인했습니다. VRAM을 약 8배 더 사용하며 OOM (Out of Memory)이 발생할 것이라 예상했기에 이 또한 충격이었습니다. 하지만 발생하지 않았습니다!

모델 인스턴스가 하나뿐이기 때문에 정적 오버헤드는 동일하게 유지된다는 것이 밝혀졌습니다. 따라서 모든 모델 및 옵티마이저 파라미터는 여전히 약 16GB의 VRAM만 필요합니다. 하지만 다른 활성화 값들 (activations)은 서로 다른 데이터 세트에서 병렬로 작동하기 위해 배치 전체에 걸쳐 복제되어야 합니다.

학습 내용 (Learning): 배치 크기 (batch size)가 커짐에 따라 활성화 메모리 (activation memory)는 선형적으로 증가합니다. 정적 오버헤드 (static overhead)는 동일하게 유지됩니다.

이후, 저는 더 큰 컨텍스트 윈도우 (context window)로 학습하기로 결정했습니다. 왜냐하면 모델을 최종적으로 어떤 데이터로 학습시키든, 일종의 챗봇 (chatbot)처럼 작동할 것이라고 가정했기 때문입니다. 1024는 매우 작은 컨텍스트 윈도우입니다. 그래서 컨텍스트 윈도우를 8192 토큰으로 늘렸습니다 (이전보다 8배 증가했으며, 저의 미숙한 추측으로는 필요한 최소한의 컨텍스트 양입니다).
그리고 놀랍게도, 배치 크기가 1임에도 불구하고 8192 컨텍스트 크기의 Llama 모델을 학습시키려 할 때 OOM (Out of Memory) 문제에 직면했습니다.

이 경우에도 정적 메모리 오버헤드 (static memory overhead)는 여전히 동일합니다. 셀프 어텐션 (self-attention)은 본질적으로 이차 함수적 (quadratic) 성질을 가지므로, 컨텍스트 크기의 제곱 ($ ext{context size}^2$)에 비례합니다. 만약 컨텍스트 크기를 8배 늘린다면, 셀프 어텐션 파라미터는 64배 증가합니다.
따라서, 필요한 셀프 어텐션 메모리 자체만으로도 150GB를 초과하게 됩니다!

학습 내용 (Learning): 컨텍스트 윈도우 크기에 따라 셀프 어텐션 메모리는 이차 함수적으로 (quadratically) 증가합니다.

이제 저는 이 메모리 사용량을 최적화해야 했습니다. 셀프 어텐션은 더 이상 유용해 보이지 않았습니다. 그래서 저는 Flash Attention 메커니즘으로 전환했습니다. 이를 통해 모든 어텐션 파라미터가 이차 함수적이 아닌 선형적으로 (linearly) 증가하게 되었습니다.
Flash Attention (SDPA)를 사용했을 때, 배치 크기 1 기준으로 제 모델은 58GB를 사용했습니다.

학습 내용 (Learning): Flash Attention을 사용하면 모든 어텐션 파라미터는 컨텍스트 크기에 따라 선형적으로 증가합니다.

그 다음 저는 그래디언트 체크포인팅 (gradient checkpointing)을 시도했습니다. 이는 학습 과정에서 순전파 (forward pass)의 모든 값을 버리고 역전파 (backward pass) 중에 이를 다시 계산할 수 있게 해주는 방식입니다. 이를 통해 메모리를 더 절약할 수 있습니다. 이 시점에서 VRAM은 약 29GB만 사용되었습니다. 그리고 매 반복 (iteration)에는 약 47초가 소요되었습니다.

지금까지의 모든 데이터는 이 표에 요약되어 있습니다:

단계 (Stage)	컨텍스트 크기 (Context size)	배치 크기 (Batch Size)	최적화 (Optimization)	VRAM (GB)	스텝당 시간 (Time s/step)
0	1024	1	-	26	3.5
...
이 표의 데이터는 이 포스트를 작성하는 동안 재계산되었음을 유의하십시오. 따라서 다음 표와 100% 일치하지 않을 수 있습니다. 이 시점에서는 이미 전용 VRAM (dedicated VRAM)을 활성화한 상태였습니다. 또한 PyTorch, ROCm, Ubuntu 등을 업그레이드했습니다.

이제 주요 문제가 드러나기 시작했습니다. 배치 크기 (batch-size)를 2배로 늘렸을 때, 스텝당 학습 시간이 약 2배 더 소요되었습니다 (그래디언트 체크포인팅 (grad checkpointing) 사용 여부와 관계없이). 따라서 GPU가 메모리가 아닌 연산 능력 (compute)에 의해 제한되고 있다는 점이 매우 명확해졌습니다.

이 시점에서 저는 더 많은 VRAM을 사용하려는 시도를 중단했습니다. 이제는 학습 시간 (training time)을 개선하는 데 집중해야 했습니다. 저는 그래디언트 체크포인팅 (grad checkpointing)을 사용하지 않기로 결정했습니다. 이는 약 25% 더 느렸으며, 이 정도 규모에서는 시간이 더 중요하기 때문입니다.

이제 학습 과정의 각 스텝에는 약 80초가 소요되었습니다. tiny-stories 데이터셋의 경우, 배치 크기 (batch-size) 2를 사용할 때 1 에포크 (epoch)의 학습을 완료하기 위해 여전히 24,414번의 학습 스텝이 필요했습니다.

즉, (매 100번의 학습 스텝마다 평가 스텝을 추가한 후에도) 단 1 에포크의 학습을 완료하는 데 약 한 달간의 연속적인 학습이 필요하다는 뜻이었습니다! 이는 실행 불가능했습니다.

또한 학습 데이터의 크기를 늘리는 것도 완전히 비현실적이었기에 더 이상 의미가 없었습니다. 그래서 저는 tiny-stories를 모델 학습에 사용할 유일한 데이터셋으로 결정했습니다.