LabHub
시작하기
배우기 러닝패스 코스

MiniMind — 작은 언어 모델을 처음부터 끝까지 직접 학습한다

MiniMind 의 generate 로 KV 캐시와 샘플링을 숫자로 잰다

LabHub 에서 이어서 보기

목표

MiniMind 의 generate 로 KV 캐시가 있을 때와 없을 때 욕심쟁이 생성이 같은 토큰을 내는지 확인하고, forward 에 들어간 토큰 수를 식과 맞추고, 시간을 잰다. 캐시의 실제 바이트를 식과 맞추고, 온도와 top-p 가 무엇을 바꾸는지 SFT 모델로 센다.

왜 중요한가

생성 비용은 캐시가 정한다. 캐시가 없으면 새 토큰 하나를 만들 때마다 지금까지의 전체를 다시 계산하고, 있으면 한 토큰만 계산한다. 그 차이는 길이가 길수록 제곱으로 벌어진다. 대신 캐시는 메모리를 먹고, 그 크기는 설정(층·KV 헤드·head_dim)으로 정확히 계산된다 — 서빙에서 한 카드에 몇 명을 받을지가 이 숫자다. 뽑는 방법은 캐시와 별개다. 캐시는 '같은 결과를 싸게' 이고, 샘플링은 '어떤 결과를' 이다. 둘을 섞어 생각하면 '캐시를 켰더니 답이 바뀌었다' 같은 잘못된 진단을 하게 된다.

단계

  1. 기준 사전학습 모델로, 검증 문서 앞 6개를 이어 200토큰까지 자른 프롬프트(앞에 bos) 뒤를 64토큰 욕심쟁이로 캐시 있이·없이 생성해 같은지 /root/mm/infer/same.json 에 적으세요.
  2. 같은 두 생성에서 forward 에 들어간 토큰 수를 세어 /root/mm/infer/count.json 에 적으세요.
  3. 두 방식의 생성 시간을(한 번 데운 뒤 세 번 재서 가운데 값) /root/mm/infer/time.json 에 적으세요.
  4. 토큰 100개를 넣은 뒤의 KV 캐시 크기를 재고 식과 맞춰 /root/mm/infer/kvbytes.json 에 적으세요.
  5. SFT 모델로 '가람 마을의 특산물은 뭐야?' 에 온도 0.3·1.0·1.5 로 30번씩 답하게 해 서로 다른 답의 수를 /root/mm/infer/temp.json 에 적으세요.
  6. 사전학습 모델에 [bos] 민수는 을 넣은 다음 토큰 분포에서, MiniMind 의 top-p 규칙으로 p=0.5·0.9·0.99 에 살아남는 후보 수를 /root/mm/infer/topp.json 에 적으세요.
  7. /root/mm/infer/report.md 에 ## KV 캐시 ## 온도 ## top-p 세 절을 쓰고, 3단계의 speedup 과 4단계의 캐시 바이트를 넣으세요.

참고

캐시가 있든 없든 같은 답

스크립트 /root/mm/infer/cache.py 를 써서 /opt/mm/ref/pretrain.pth 를 불러, /opt/mm/data/pretrain_val.jsonl 앞 6개 문서의 text 를 공백으로 이어 자른 토큰의 앞 200개에 bos(1)를 붙인 프롬프트 뒤를 64토큰 욕심쟁이(do_sample=False, top_k=0, top_p=1.0, eos_token_id=None)로 use_cache=True·False 두 번 생성하세요. 새 토큰이 같은지(same_ids)와 캐시 쪽 새 토큰(ids)을 /root/mm/infer/same.json 에 prompt_len·new_tokens 와 함께 적으세요.

인과 마스크 덕분에 앞 토큰의 K·V 는 뒤에 토큰이 붙어도 바뀌지 않습니다. 그래서 꺼내 써도 결과가 같습니다. 여기서 다르다면 캐시가 아니라 조건(샘플링·난수)이 다른 것입니다.

forward 에 들어간 토큰을 센다

1단계의 두 생성에서 model.forward 에 들어간 input_ids 의 길이를 모두 더해 /root/mm/infer/count.json 에 prompt_len·new_tokens·fed_with_cache·fed_without_cache 로 적으세요.

캐시가 있으면 첫 걸음에 프롬프트 전체, 그 뒤로는 한 토큰씩입니다. 없으면 걸음마다 지금까지의 전체를 다시 넣습니다. 채점기는 여러분이 적은 두 값을 이 두 식으로 계산한 값과 견줍니다.

시간으로 보면

캐시 있이·없이 64토큰 생성을 한 번 데운 뒤 각각 세 번 재서 가운데 값을 /root/mm/infer/time.json 에 with_cache_s·without_cache_s·speedup(없음÷있음)으로 적으세요.

time.perf_counter() 로 잽니다. 첫 실행은 메모리 할당과 준비 때문에 느려 제외합니다. 토큰 수로는 50배 넘게 차이 나는데 시간은 그만큼 차이 나지 않습니다 — 작은 모델에서는 걸음마다의 고정 비용(파이썬·커널 호출)이 크기 때문입니다.

캐시는 몇 바이트인가

기준 사전학습 모델에 무작위 토큰 100개(torch.randint(3, 어휘, (1, 100), generator=시드 0))를 use_cache=True 로 넣고, 돌려받은 past_key_values 의 모든 텐서 바이트 합(measured_bytes)과 식으로 계산한 값(formula_bytes = 2 × 층 × 토큰 × KV 헤드 × head_dim × 4), 첫 층 K 의 모양(k_shape)을 /root/mm/infer/kvbytes.json 에 tokens 와 함께 적으세요.

MiniMind 의 Attention 은 past_kv = (xk, xv) 를 repeat_kv 전에 만듭니다. 그래서 K 의 모양은 (배치, 길이, KV 헤드 2, 32) 이고, GQA 로 줄인 만큼 캐시도 작습니다.

온도를 올리면

SFT 모델(/opt/mm/ref/sft.pth)에 '가람 마을의 특산물은 뭐야?' 를 채팅 형식(mmkit.chat_text(…, add_generation_prompt=True))으로 넣고, 온도 0.3·1.0·1.5 마다 torch.manual_seed(0) 을 건 뒤 do_sample=True, top_k=0, top_p=1.0, max_new_tokens=24, num_return_sequences=30 으로 생성해, <|im_end|> 앞까지 자른 서로 다른 답의 수를 /root/mm/infer/temp.json 에 {"0.3": {"distinct": n}, "1.0": …, "1.5": …} 로 적으세요.

온도마다 같은 난수에서 출발해야 온도의 효과만 봅니다. 온도 1.5 에서는 서른 개가 거의 모두 다르고 글자가 깨진 답도 나옵니다 — 드문 토큰이 뽑히면 작은 모델은 제자리로 돌아오지 못합니다.

top-p 가 남기는 후보

기준 사전학습 모델에 [bos] + '민수는' 을 넣은 마지막 자리의 로짓에서, MiniMind generate 의 top-p 규칙(확률 내림차순 누적합이 p 를 넘는 자리부터 자르되 한 칸 밀어 p 를 처음 넘기는 토큰은 살리고, 맨 위 하나는 늘 살림)으로 p=0.5·0.9·0.99 에 살아남는 후보 수를 /root/mm/infer/topp.json 에 {"0.5": n, "0.9": n, "0.99": n} 으로 적으세요.

/opt/minimind/model/model_minimind.py 의 generate 에서 top_p 부분 세 줄을 그대로 옮기면 됩니다. '민수는' 뒤에는 동사·장소 등 여러 갈래가 있어 후보가 여럿 남습니다. 뾰족한 자리(예: '마을의 특산물은' 뒤)와 견줘 보세요.

추론 비용과 뽑기의 기록

/root/mm/infer/report.md 에 ## KV 캐시 ## 온도 ## top-p 세 절을 쓰고, 3단계의 speedup 과 4단계의 measured_bytes 를 숫자로 넣으세요.

캐시가 토큰 수로는 몇 배, 시간으로는 몇 배였는지 나란히 적고 그 차이의 이유를 한 줄 덧붙이세요.