LLM from scratch, 파트 28 – RTX 3090에서 베이스 모델을 처음부터 훈련하기

Source: Dev.to

설정: 왜 RTX 3090을 선택했는가

먼저 하드웨어부터 얘기해 보겠습니다. 저는 RTX 3090을 한동안 사용해 왔으며, 24 GB VRAM 덕분에 딥러닝에 많이 쓰이는 선택지입니다. 성능이 뛰어난 만큼 비용도 만만치 않으니, 과열 없이 작업을 처리할 수 있도록 작업 공간을 효율적으로 구성해야 했습니다.

첫 번째 단계는 필요한 라이브러리를 설치하는 것이었습니다. 저는 성능과 사용 편의성의 균형이 좋은 PyTorch를 선택했습니다.

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
pip install transformers datasets

이제 모델을 파인‑튜닝하기 시작하면 본격적인 작업이 시작됩니다.

아키텍처: 처음부터 직접 설계하기

재미있는 부분, 즉 베이스 모델의 아키텍처를 정의하는 단계입니다. 저는 표준 트랜스포머 아키텍처를 선택했습니다.

import torch
import torch.nn as nn

class SimpleTransformer(nn.Module):
    def __init__(self, input_dim, num_heads, num_layers):
        super(SimpleTransformer, self).__init__()
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=input_dim, nhead=num_heads),
            num_layers=num_layers
        )

    def forward(self, x):
        return self.encoder(x)

model = SimpleTransformer(input_dim=512, num_heads=8, num_layers=6)

하이퍼파라미터 튜닝—레이어와 헤드 수를 적절히 선택하는 것이 핵심이었습니다.

학습: 좋은 점, 나쁜 점, 그리고 계산량이 많은 점

모델을 준비한 뒤, 다양한 출처에서 수집한 데이터를 모델에 공급했습니다. 모든 데이터가 동일하게 좋은 것은 아니므로, 데이터 정제가 필수적이었습니다.

from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
train_data = dataset["train"]

몇 에포크가 지나자 모델이 일관된 텍스트를 생성하기 시작했지만, 가끔씩 의미 없는 출력이 나타나면서 스크래치 학습의 한계와 데이터 품질의 중요성을 다시금 깨달았습니다.

문제 해결: 문제가 발생했을 때

제가 겪은 주요 문제는 그래디언트 폭발이었습니다. 그래디언트 클리핑을 적용해 학습을 안정화시켰습니다.

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

이 한 줄 코드로 손실 값이 급격히 상승하는 것을 방지하고, 학습 과정을 보다 관리하기 쉬운 수준으로 만들 수 있었습니다.

실제 적용 사례: 이 모든 것이 어디에 쓰이는가

학습된 모델은 작은 프로젝트용 코드 스니펫을 생성해 주는, 언제든지 작동하는 어시스턴트처럼 활용될 수 있습니다. 하지만 동시에 윤리적 고려사항도 떠오릅니다. 문제가 있는 데이터로 학습된 모델은 편향을 지속시킬 수 있기 때문에 개발자는 이러한 위험을 인식하고 조심해야 합니다.

앞으로의 생각: 향후 계획

앞으로는 특정 작업에 대한 파인‑튜닝과 효율성을 위한 혼합 정밀도 학습 같은 고급 기법을 탐구할 예정입니다. 비슷한 길을 걷고 계신 분들은 사용 중인 도구와 겪은 어려움을 자유롭게 공유해 주세요.

스크래치에서 LLM을 구축하는 과정은 실패를 받아들이고 작은 성공을 축하하며 지속적으로 한계를 뛰어넘는 보람찬 경험이었습니다. RTX 3090을 잡고 함께 멋진 무언가를 만들어 봅시다!