이 문서는 보강 - 토크나이저, forward, 생성 루프에서 갈라져 나온 곁가지입니다. 그 글이 두 메서드를 나란히 읽는다면, 여기서는 한쪽만 그림 위주로 더 파고듭니다. 다른 한쪽은 generate()의 세 단계에서 다룹니다.

model.generate() 는 자기회귀 루프를 함수 안에서 최대 50바퀴 돌아줍니다. generate_forward_batch() 는 그 루프의 한 바퀴만 실행하고, 루프 자체는 바깥의 서빙 코드가 돌립니다. 추상화를 벗겨내면 무엇을 손으로 해야 하는지, 각각 무엇을 맡고 있는지 따라가 봅니다.

읽는 위치

llm/model_worker.py:65 의 generate_forward_batch() 는 프롬프트 배치 하나를 받아 각자에게 토큰 한 개씩을 돌려줍니다. 바로 위에 있는 generate() 는 같은 배치로 50토큰을 만들어 주는데 실질 코드가 30줄입니다. 토큰을 하나만 만드는 이쪽은 35줄입니다. 덜 만드는 쪽이 더 깁니다.

길어진 그 차이가 generate() 안에 접혀 있던 부분입니다. 마지막 자리의 점수를 골라내고, 확률로 바꾸고, 주사위를 굴리는 일이 밖으로 나왔습니다.

# transformers   model_worker.py:77 · 배치를 한 텐서로 묶는다
encoded = self.tokenizer(
    [p['prompt'] for p in prompts],
    return_tensors="pt", padding=True,
    truncation=True, max_length=512,
).to(self.device)

# transformers   model_worker.py:89 · forward 딱 한 번
with torch.no_grad():
    outputs = self.model(
        input_ids=encoded.input_ids,
        attention_mask=encoded.attention_mask,
        use_cache=False,
    )

    # PyTorch      model_worker.py:96 · 여기부터가 손으로 하는 구간
    next_token_logits = outputs.logits[:, -1, :]
    next_token = torch.multinomial(
        torch.softmax(next_token_logits / 0.7, dim=-1),
        num_samples=1,
    ).squeeze(-1)

    # transformers   model_worker.py:104 · 배치를 풀어 하나씩 담는다
    for i, prompt_data in enumerate(prompts):
        token = self.tokenizer.decode(
            next_token[i].unsqueeze(0), skip_special_tokens=True
        )
        results.append({'request_id': ..., 'token': token, ...})

가운데 네 줄이 이 문서가 다룰 전부입니다. [:, -1, :] 가 어떤 낭비를 감수하는 슬라이스인지, / 0.7 이 분포를 어떻게 바꾸는지, multinomial 이 argmax 와 무엇이 다른지를 차례로 봅니다.

마지막 for 루프 안에 모델 호출이 없다는 점도 눈여겨볼 만합니다. 연산은 이미 위에서 배치 전체에 대해 끝났고, 저 루프는 [4] 짜리 텐서를 인덱스로 풀어 담기만 합니다. 무거운 일은 배치당 한 번이고, 파이썬 루프는 요청당 한 번입니다.

아래 그림에 적힌 숫자는 전부 facebook/opt-125m 을 실제로 돌려서 나온 값입니다.

forward 한 번: 모든 자리에 대해 동시에 답한다

model.generate() 와 model() 은 이름이 비슷하지만 규모가 다릅니다. 뒤엣것은 forward 딱 한 번이고, 돌아오는 outputs.logits 의 모양은 [B, T, vocab] 입니다. B 는 배치 크기, T 는 입력 토큰 수, vocab 은 어휘 사전 크기입니다. facebook/opt-125m 이면 vocab 은 50272 입니다.

중요한 것은 가운데 축입니다. 모델은 마지막 자리만 보고 답하지 않습니다. 입력의 모든 위치 t 마다 "이 자리 다음에 올 토큰의 점수 50272개"를 한꺼번에 내놓습니다. 프롬프트가 4토큰이면 예측이 4벌 나옵니다.

opt-125m 은 "brown" 다음에 'ie' 를 가장 높게 봅니다. 125M 파라미터짜리 모델이라 "brownie" 로 이어가려는 것이고, 이것이 실제 측정값입니다. 앞의 세 자리 예측도 똑같이 계산되었지만 logits[:, -1, :] 한 줄에 걸러집니다.

T개를 계산해놓고 1개만 쓴다

그러면 앞의 세 벌은 왜 계산할까요. 버리려고 계산하는 것이 아니라, 학습 때 쓰라고 만들어진 출력 형태를 추론이 그대로 물려받았기 때문입니다. 학습에서는 정답 문장을 통째로 넣고 T 개 위치의 예측을 전부 정답과 비교해 손실 하나를 만듭니다. 이것이 teacher forcing 이고, forward 한 번으로 T 자리를 동시에 배울 수 있어서 학습이 빠릅니다.

추론에서는 t 가 0부터 T-2 까지의 예측이 쓸모없습니다. 그 자리 다음에 무엇이 오는지 이미 프롬프트에 적혀 있기 때문입니다. 아직 모르는 자리는 마지막 하나뿐입니다.

"forward 1회 = 토큰 1개"는 모델의 성질이 아니라 소비 방식의 결과입니다. model_worker.py:96 의 슬라이스 한 줄이 그 버리는 동작을 담당합니다.

샘플링: 점수를 확률로 바꾸고 주사위를 굴린다

generate() 를 썼다면 여기서 끝났습니다. 그 안에 샘플링이 들어 있기 때문입니다. 직접 부르면 점수 50272개를 받아 토큰 하나를 고르는 일을 손으로 해야 합니다. model_worker.py:96-100 의 네 줄이 그 일 전부입니다.

dim=-1 은 "마지막 축을 따라"라는 뜻입니다. 여기서는 50272짜리 어휘 축이므로, 4개 행 각각이 독립적으로 합 1인 확률 분포가 됩니다.

temperature로 나눈다는 것

0.7 로 나누는 순간 무슨 일이 벌어지는지는 숫자로 보는 편이 빠릅니다. 아래는 "The quick brown" 다음 토큰에 대한 실제 확률입니다. 나누기 전과 후를 나란히 놓으면, 1등이 가져가는 몫이 두 배 이상 늘고 나머지 5만여 개의 몫이 3분의 1 아래로 줄어듭니다.

나눗셈 자체는 순서를 바꾸지 않습니다. 1등은 나눠도 1등입니다. 바뀌는 것은 1등과 나머지의 격차뿐이고, 그 격차가 곧 "다른 후보가 뽑힐 여지"입니다. T 를 0에 가깝게 보내면 사실상 greedy 가 되고, 크게 키우면 분포가 평평해져 아무 토큰이나 나옵니다.

multinomial과 argmax

argmax 는 가장 높은 칸을 무조건 고릅니다. 이것이 greedy 디코딩이고, 같은 프롬프트에는 언제나 같은 답이 나옵니다. multinomial 은 확률을 그대로 주사위 눈금으로 삼습니다. 0.641 짜리 토큰은 열 번 중 여섯 번쯤 나오고, 0.121 짜리도 가끔 나옵니다. 실행할 때마다 답이 달라지는 이유가 이것입니다.

배치 4개로 실제로 비교해 보면 차이가 바로 드러납니다.

넷 중 셋이 다릅니다.

그리고 위 다섯 단계 어디에도 파이썬 for 가 없습니다. 나눗셈도 softmax 도 multinomial 도 [4, 50272] 행렬을 통째로 받아 [4] 를 돌려줍니다. 4개 요청은 서로를 모른 채 같은 커널 안에서 동시에 처리되고, 파이썬 루프는 결과를 dict 로 포장할 때 한 번만 돕니다.

루프의 소유권: 같은 루프를 누가 돌리는가

여기까지가 한 바퀴입니다. 토큰 하나가 나왔고, 다음 토큰을 만들려면 이 바퀴를 또 돌아야 합니다. "실행, 토큰 선택, 입력 끝에 붙이기, 재실행"이라는 자기회귀 루프는 두 경로 모두에 있습니다. 다른 것은 그 루프가 어느 함수 안에서 도느냐뿐이고, 이 차이 하나가 서빙의 성격을 전부 바꿉니다.

왼쪽에서 ×50 화살표는 transformers 함수 몸통 안에 있습니다. 밖에서는 손댈 수 없습니다. 오른쪽에서 같은 화살표는 llm.py:50 의 while True 입니다. 루프를 밖으로 꺼냈기 때문에 매 바퀴 두 가지가 가능해집니다. 토큰이 나올 때마다 클라이언트에게 보내는 것, 그리고 다음 바퀴의 배치 구성원을 새로 정하는 것입니다.

오른쪽 그림의 get_next_batch() 가 루프 안에 있다는 점을 눈여겨볼 만합니다. 끝난 시퀀스는 빠지고 대기 중이던 요청은 다음 바퀴에 들어옵니다. 배치를 한 번 짜서 끝까지 끌고 가는 것이 아니라 매 스텝 다시 짜는 셈인데, 이 아이디어를 제대로 밀고 나간 것이 vLLM 의 continuous batching 입니다.

메소드의 역할 분담

generate_forward_batch() 한 함수 안에 다양한 부분을 직접 나열하여 작성하였습니다.

• PyTorch의 요소
  • .to(self.device)
  • torch.no_grad()
  • multinomial
• HuggingFace의 transformers
  • tokenizer
  • self.model

generate_vllm() 에는 토크나이저도, .to(device) 도, no_grad() 도, multinomial 도 없습니다. 사라진 것이 아니라 vLLM 안으로 들어간 것입니다. 위 목록이 두 줄이 되는 대가로, 그 안에서 무슨 일이 벌어지는지가 보이지 않게 됩니다.