Padding wastes compute; packing blurs document boundaries

한 줄 요약

사전학습 자료는 '토큰의 긴 줄' 로 만들어 배치로 잘라 먹인다. MiniMind 는 문서 하나를 한 줄로 두고 최대 길이까지 패딩한다 — 단순하고 문서가 섞이지 않지만, 문서가 짧으면 계산 대부분이 패딩에 버려진다. 여러 문서를 이어 붙여 꽉 채우는 패킹은 버리는 계산이 없는 대신 한 창 안에서 앞 문서가 뒤 문서에 보인다. 이 모듈은 우리 말뭉치에서 두 방식의 비용을 숫자로 잰다.

왜 이게 필요했나

모델은 (배치, 길이) 모양의 정수 텐서를 받는다. 문서 길이는 제각각인데 텐서는 네모여야 하므로 어딘가에서 맞춰야 한다. 방법은 둘이다.

패딩. 문서 하나를 한 줄에 두고 모자란 자리를 패딩 토큰으로 채운다. MiniMind 의 PretrainDataset 이 이렇게 한다.

tokens = tokenizer(text, max_length=self.max_length - 2, truncation=True).input_ids
tokens = [bos_token_id] + tokens + [eos_token_id]
input_ids = tokens + [pad_token_id] * (self.max_length - len(tokens))
labels = input_ids.clone(); labels[input_ids == pad_token_id] = -100

패딩 자리는 라벨이 -100 이라 손실에는 안 들어가지만 계산은 똑같이 한다. 행렬 곱은 패딩인지 모른다. 우리 말뭉치의 문서는 평균 33토큰인데 최대 길이를 128로 두면, 계산의 70% 넘게가 패딩에 쓰인다. 거꾸로 최대 길이를 짧게 잡으면 긴 문서가 잘린다. MiniMind README 도 이 맞바꿈을 적어 둔다 — 짧은 표본은 패딩으로 계산을 버리고, 긴 표본은 잘려 정보를 잃는다. 그래서 자료마다 권장 max_seq_len 을 따로 적는다.

패킹. 문서마다 [bos] … [eos] 로 감싸 한 줄로 길게 잇고, 그 줄을 고정 길이로 자른다. 버리는 자리가 없다. 대신 한 창에 문서 서너 개가 들어가고, 인과 마스크는 '앞의 모든 토큰' 을 보게 하므로 뒤 문서의 토큰이 앞 문서를 본다. 대부분의 사전학습은 이것을 감수한다 — [eos]·[bos] 가 경계를 알려 주고, 모델이 경계 너머를 무시하는 법을 배우기 때문이다. 경계를 넘지 못하게 하려면 문서별 어텐션 마스크를 따로 만들어야 하는데, 그만큼 구현이 무거워진다.

어떻게 동작하나

패킹한 결과는 정수 배열 하나다. 어휘가 65,536 보다 작으면 uint16 으로 충분해 int64 의 4분의 1 공간이면 된다. 학습 루프는 이 배열에서 임의 위치를 골라 길이 만큼 떼어 배치를 만든다.

ix = torch.randint(0, len(train) - seq, (batch,), generator=g)
x = torch.stack([train[i:i + seq] for i in ix])
loss = model(x, labels=x).loss      # 한 칸 미는 일은 모델 안에서

MiniMind 모델은 labels 로 입력과 같은 텐서를 받고, 안에서 logits[..., :-1] 과 labels[..., 1:] 을 맞춰 다음 토큰 예측 손실을 낸다. 그래서 자료 쪽에서 입력과 정답을 따로 밀어 둘 필요가 없다.

검증 자료는 문서 단위로 떼어야 한다. 한 배열을 앞 95%·뒤 5% 로 자르면 한 문서가 둘로 갈려 양쪽에 걸칠 수 있고, 같은 문서가 학습과 검증에 모두 있으면 검증 손실은 외운 것을 잰다. 떼어 낸 뒤에도 같은 글이 학습 쪽에 글자 그대로 있는지 한 번 확인한다.

현장에서 만나는 모습

짧은 질의응답·채팅 기록처럼 문서가 짧은 자료를 패딩으로 학습하면, GPU 가 바쁘게 도는데 손실은 느리게 내려간다. 사용률은 높지만 실제로 배우는 토큰이 적기 때문이다. 이때 한 걸음에 들어간 진짜 토큰 수를 세어 보면 원인이 바로 보인다. 반대로 패킹으로 바꾼 뒤 모델이 문서 경계를 넘어 엉뚱한 내용을 이어 쓰기 시작했다면 [eos] 를 제대로 붙였는지부터 본다.

학습 예산도 토큰으로 센다. '몇 걸음' 은 배치 크기와 길이에 따라 뜻이 달라지므로, 한 걸음이 먹는 토큰과 말뭉치 전체 토큰으로 몇 걸음이 한 바퀴(에폭)인지 계산해 두어야 손실 곡선의 굴곡을 읽을 수 있다.

MiniMind 원본과 이 코스가 다른 점

MiniMind 는 HuggingFace datasets 로 jsonl 을 읽고, 표본마다 __getitem__ 에서 그때그때 토크나이저를 돌린다. 1.2GB 말뭉치를 미리 토큰으로 바꿔 두지 않아도 되는 대신, 학습하는 동안 CPU 가 계속 토큰을 자른다(그래서 기본 num_workers 가 8 이다). 이 코스는 말뭉치가 2MB 도 안 되므로 한 번에 모두 잘라 uint16 배열 하나로 저장하고, 학습 루프는 그 배열에서 창을 떼어 쓰기만 한다. 큰 사전학습에서도 이 방식(미리 토큰화해 이진 파일로)이 흔하다 — 자료를 여러 번 도는 동안 같은 토큰화를 되풀이하지 않고, 파일을 메모리 매핑해 필요한 부분만 읽을 수 있기 때문이다.

다음 실습에서 할 것

기준 토크나이저로 문서 길이를 재고, MiniMind 식 패딩이 얼마를 버리는지 센다. 학습·검증 말뭉치를 문서 단위로 패킹해 uint16 배열로 저장하고, 검증 문서가 학습 쪽에 새지 않았는지, 패킹한 창에서 토큰 몇 %가 앞 문서를 볼 수 있는지, 한 바퀴가 몇 걸음인지 계산한다.