[Open source contribution] langchain4j 오픈소스 기여 경험기 - 임베딩 토큰 분할에 숨어 있던 무한 루프를 고치다

프로그램이 예외를 던지며 멈추면 적어도 어디가 잘못됐는지 단서가 남습니다. 그런데 프로그램이 멈추지도, 끝나지도 않고 같은 일을 영원히 반복한다면 단서조차 없습니다. 화면은 응답이 없고, 로그는 더 이상 늘지 않으며, 무엇을 기다리는지도 알 수 없습니다. 무한 루프는 이렇게 침묵하는 고장입니다.

 

이번 글에서 다룰 LangChain4j의 임베딩 모듈 수정은 바로 그 무한 루프에 관한 것이었습니다. 평범한 문장을 임베딩할 때는 멀쩡하던 코드가, 아주 특별한 형태의 입력 하나를 만나면 영원히 끝나지 않는 반복에 빠졌습니다. 결과만 보면 반복문에 조건 몇 줄을 더하고 테스트 하나를 추가한 작은 변경이지만, 왜 그 반복문이 멈추지 못했는지, 그리고 멈추는 반복문이란 무엇을 보장해야 하는지를 따져 보는 과정이 저에게는 오래 남았습니다. 이 글은 그 한 건의 기여를 처음부터 끝까지 따라가며 정리한 기록입니다.

---

긴 텍스트를 임베딩하려면 잘라야 합니다

먼저 이 코드가 무슨 일을 하는지부터 짚어 두겠습니다. LangChain4j에는 외부 서비스를 거치지 않고 컴퓨터 안에서 직접 문장을 벡터로 바꾸는 임베딩 모델들이 있습니다. 이런 모델의 속을 들여다보면 OnnxBertBiEncoder라는 클래스가 실제 변환을 맡고 있습니다. 텍스트를 받아 토큰으로 쪼개고, 그 토큰들을 모델에 넣어 벡터를 얻는 일을 합니다.

 

여기서 두 가지 사실을 알아 두면 이후 이야기가 분명해집니다. 하나는 토큰화 방식이고, 다른 하나는 길이 제한입니다.

 

BERT 계열 모델은 단어를 통째로 다루지 않고 더 작은 조각으로 쪼개는 방식을 씁니다. 흔히 WordPiece라고 부르는 이 방식에서는, 모델이 아는 어휘에 없는 긴 단어를 여러 조각으로 나눕니다. 이때 한 단어의 첫 조각이 아닌 이어지는 조각들에는 앞에 ##이라는 표시가 붙습니다. 예를 들어 아주 긴 단어 하나가 "super", "##cali", "##fragi", "##listic"처럼 나뉘는 식입니다. 여기서 ##이 붙은 조각은 "앞 조각에 이어지는 같은 단어의 일부"라는 뜻입니다. 즉 ##이 붙은 토큰은 홀로 떨어져 나오면 안 되고, 앞 토큰과 한 덩어리로 묶여 있어야 자연스럽습니다.

 

또 하나, 이런 모델은 한 번에 처리할 수 있는 토큰의 개수에 상한이 있습니다. 그래서 토큰이 그 상한보다 많아지면, 텍스트를 여러 조각으로 나눠 각각 임베딩한 뒤 그 결과를 가중 평균으로 합칩니다. 이 "토큰 목록을 일정 크기씩 나누는" 일을 맡은 것이 partition이라는 메서드입니다. 이번 버그는 바로 이 메서드 안에 있었습니다.

 

---

단어를 쪼개지 않으려는 배려가 함정이 되었습니다

partition 메서드의 뼈대는 단순합니다. 토큰 목록을 앞에서부터 일정 크기(partitionSize)씩 끊어 나갑니다. 시작 위치를 from, 끝 위치를 to라 하면, 매번 to를 from에서 한 묶음만큼 떨어진 곳으로 잡고, from부터 to까지를 한 조각으로 잘라낸 다음, 다음 차례를 위해 from을 방금의 to 자리로 옮깁니다. 이렇게 from이 토큰 목록 끝에 닿을 때까지 반복합니다.

 

여기에 한 가지 배려가 더해져 있었습니다. 조각을 끊는 자리가 하필 한 단어의 중간이면, 같은 단어가 두 조각으로 갈라지게 됩니다. 이를 피하려고, 끊는 자리인 to가 ##이 붙은 토큰(즉 단어의 이어지는 조각)을 가리키고 있으면 to를 한 칸씩 앞으로 당겨, 단어가 시작되는 자리까지 물러나도록 했습니다. 단어의 경계에서 깔끔하게 끊으려는 의도였습니다. 수정 전 코드는 이 부분이 다음과 같았습니다.

while (tokens.get(to).startsWith("##")) {
    to--;
}

 

평범한 텍스트에서는 이 코드가 잘 동작합니다. 한 단어가 몇 조각으로 나뉘더라도, 그 조각의 개수는 한 묶음 크기에 비하면 훨씬 작기 때문입니다. to를 조금 당기면 곧 단어의 시작에 닿고, 그 자리에서 깔끔하게 끊긴 뒤 from이 그만큼 앞으로 나아갑니다. 문제가 드러나지 않습니다.

 

함정은 "한 단어의 조각 수가 한 묶음 크기보다 큰" 극단적인 경우에 있었습니다. 이런 일은 공백 없이 아주 길게 이어진 문자열처럼, 하나의 단어가 수백 개의 조각으로 쪼개질 때 일어납니다. 이때 from에서 한 묶음만큼 떨어진 to 자리부터 from 바로 뒤까지가 전부 ##이 붙은 같은 단어의 조각들로 채워집니다. 그러면 to를 당기는 위 반복문이 멈출 줄을 모릅니다. ##이 아닌 토큰을 만나야 멈추는데, 그런 토큰이 나오지 않으니 to는 계속 내려가 결국 from까지 도달합니다.

---

매번 한 걸음도 못 나아가면 반복문은 끝나지 않습니다

to가 from까지 내려오면 두 가지 일이 한꺼번에 잘못됩니다.

 

먼저 from부터 to까지를 잘라낸 조각이 빈 조각이 됩니다. 시작과 끝이 같은 자리이니 그 사이에는 아무 토큰도 없습니다. 의미 있는 한 조각을 만들어야 할 자리에서 텅 빈 조각이 나오는 것입니다.

 

더 큰 문제는 그다음입니다. 반복을 한 바퀴 돌고 나면 from을 to 자리로 옮기는데, to가 from과 같으니 from은 제자리에 머뭅니다. 다음 바퀴에서도 똑같은 계산이 반복되고, to는 다시 from까지 내려오고, from은 또 제자리입니다. 시작 위치가 한 걸음도 앞으로 나아가지 못한 채 같은 자리를 영원히 맴돕니다. 이것이 무한 루프의 정체였습니다.

 

여기서 분명해지는 원칙이 하나 있습니다. 어떤 일을 반복해서 끝내려면, 매 반복이 반드시 목표를 향해 조금이라도 나아가야 한다는 것입니다. 이 코드에서 "나아간다"는 것은 시작 위치 from이 토큰 목록의 끝을 향해 전진한다는 뜻입니다. 그런데 단어 경계를 맞추려고 끝 위치를 당기는 배려가, 특정 입력에서는 그 전진을 통째로 삼켜 버렸습니다. 한 걸음 나아가려다 제자리로 끌려 내려온 것입니다. 좋은 의도를 가진 코드라도 경계 조건을 놓치면 이렇게 멈추지 못하는 함정이 된다는 것을, 이 버그가 분명하게 보여 주었습니다.

---

적어도 한 걸음은 남겨 두도록 고쳤습니다

원인이 분명해지니 고칠 방향도 두 갈래로 또렷해졌습니다. 첫째, 끝 위치를 당기는 일이 시작 위치 아래로는 절대 내려가지 않게 막아야 했습니다. 둘째, 한 단어가 한 묶음 크기보다 길어서 깔끔하게 끊을 자리가 없다면, 단어를 쪼개는 일을 감수하고서라도 어떻게든 한 묶음만큼은 잘라 전진해야 했습니다. 수정 후 코드는 다음과 같습니다.

while (to > from && tokens.get(to).startsWith("##")) {
    to--;
}
if (to == from) {
    to = from + partitionSize;
}

 

첫 번째 줄의 to > from 조건이 끝 위치를 당기는 일에 바닥을 만들어 줍니다. 이제 to는 시작 위치 바로 위까지만 내려가고, 그 아래로는 내려가지 않습니다. 적어도 토큰 한 개만큼의 전진은 남겨 두는 것입니다.

 

두 번째 조건은 그래도 끝 위치가 시작 위치까지 내려와 버린 경우, 즉 한 묶음 전체가 통째로 한 단어의 조각들로 채워진 경우를 다룹니다. 이때는 단어 경계를 지킬 방법이 애초에 없습니다. 단어 하나가 한 묶음보다 길기 때문입니다. 그래서 이 경우에는 단어가 갈라지는 것을 받아들이고, 끝 위치를 한 묶음만큼 떨어진 자리로 도로 밀어 한 조각을 온전히 잘라냅니다. 단어를 쪼개는 것은 아쉽지만, 멈추지 않고 나아가는 쪽이 영원히 멈춰 있는 것보다 낫습니다. 어차피 한 묶음에 담기지 않는 단어는 어떤 방법으로도 한 조각에 다 담을 수 없습니다.

 

여기서 마음을 놓을 수 있었던 점은, 이 두 조건이 평범한 입력에서는 아무 일도 하지 않는다는 것입니다. 한 단어의 조각 수가 한 묶음 크기보다 작은 보통의 경우, 끝 위치를 당기는 일은 단어의 시작에 닿으면 시작 위치보다 위에서 멈춥니다. 그러면 to는 from보다 큰 값이고, 두 번째 조건도 성립하지 않습니다. 즉 새로 더한 두 조건은 극단적인 입력에서만 깨어나고, 기존에 잘 동작하던 경우의 결과는 조금도 바꾸지 않습니다. 동작을 바꾸지 않으면서 깨지던 경우만 바로잡는 변경이라, 호환성을 걱정할 부분이 없었습니다.

---

끝나는가를 검증하는 테스트

코드를 고치는 것만큼 신경 쓴 것이 테스트였습니다. LangChain4j에는 "테스트가 없으면 리뷰하지 않는다"는 원칙이 있어, 변경을 증명하는 테스트가 반드시 필요했습니다. 그런데 무한 루프는 테스트하기에 묘한 구석이 있습니다. 보통의 테스트는 "결과가 기대한 값과 같은가"를 확인하지만, 무한 루프는 애초에 결과가 나오지 않기 때문입니다. 고쳐지지 않은 코드에 이 입력을 넣으면 테스트 자체가 끝나지 않고 매달려 버립니다.

 

그래서 테스트는 두 가지를 동시에 확인하도록 작성했습니다. 하나는 "끝나는가"이고, 다른 하나는 "올바른 결과를 내는가"입니다. 끝나는지를 확인하기 위해 테스트에 제한 시간을 걸었습니다. 정해진 시간 안에 끝나지 않으면 테스트가 실패하도록 한 것입니다. 이렇게 하면 혹시 누군가 나중에 이 반복문의 종료 조건을 되돌려 버려도, 테스트가 무한정 매달리는 대신 시간 초과로 분명하게 실패를 알려 줍니다.

 

그 위에서, 결과가 올바른지도 함께 검증했습니다. 문제가 되는 입력으로 한 단어가 여러 조각으로 길게 이어진 토큰 목록을 만들고, 한 묶음 크기를 일부러 작게 잡아 극단적인 상황을 재현했습니다.

@Test
@Timeout(5)
public void testSingleWordSubwordRunLongerThanPartitionSize() {

    List<String> tokens = asList("[CLS]", "super", "##cali", "##fragi", "##listic", "[SEP]");
    int partitionSize = 1;

    List<List<String>> partitions = partition(tokens, partitionSize);

    assertThat(partitions).isNotEmpty();
    assertThat(partitions).allSatisfy(partition -> assertThat(partition).isNotEmpty());

    List<String> flattened = new ArrayList<>();
    partitions.forEach(flattened::addAll);
    assertThat(flattened).containsExactly("super", "##cali", "##fragi", "##listic");
}

 

이 테스트는 세 가지를 한꺼번에 못 박습니다. 제한 시간 안에 끝나는가, 빈 조각이 하나도 없는가, 그리고 잘라낸 조각들을 다시 이어 붙이면 원래 토큰들이 빠짐도 중복도 없이 순서 그대로 담기는가입니다. 마지막 검증이 특히 중요했습니다. 무한 루프만 막고 끝내는 것이 아니라, 단어를 쪼개더라도 토큰을 하나도 잃지 않고 순서대로 모두 담아야 임베딩 결과가 온전하기 때문입니다. 끝나기만 하는 것과 올바르게 끝나는 것은 다른 일이고, 테스트는 그 둘을 모두 지켜 주어야 했습니다.

---

저장소의 절차를 따라가며 배운 것

코드와 테스트를 마친 뒤에는 저장소가 요구하는 절차를 익히는 일이 남았습니다. LangChain4j는 버그 수정의 경우 곧장 PR을 올리기보다, 먼저 이슈를 등록해 문제를 공유하고 그 이슈 번호를 연결한 PR을 올리는 흐름을 따릅니다. 특히 이런 무한 루프는 "어떤 입력이, 왜, 어떻게 코드를 멈춰 세우는가"를 글로 차근차근 설명해 두는 것이 중요했습니다. 재현 조건과 원인을 먼저 정리해 두면, 검토하는 사람도 코드를 보기 전에 문제의 그림을 그릴 수 있기 때문입니다.

 

변경을 작게 유지하라는 원칙도 내내 의식했습니다. 이번 수정의 본질은 반복문에 종료를 보장하는 두 조건을 더한 것과, 그것을 지켜 주는 테스트 하나를 추가한 것입니다. 이 한 가지 목적에서 벗어나는 다른 손질을 섞지 않으려 했습니다. 바뀐 줄 하나하나가 "반복문이 반드시 끝나도록 만든다"는 목적으로 곧장 설명될 수 있어야 한다는 기준을 세워 두고 변경 내용을 다시 살폈습니다. 이 저장소는 여러 커밋을 하나로 합쳐 병합하기 때문에 PR 제목이 그대로 기록에 남는데, 그 제목 역시 어떤 메서드의 어떤 문제를 고쳤는지가 한눈에 드러나도록 다듬었습니다.

 

한 가지 더 신경 쓴 것은, 이 버그가 겉으로 드러나는 증상만 가려 두는 식으로 처리되어 있었다는 점입니다. 이전에 어떤 방어 장치가 더해져 있어서 빈 조각으로 인한 일부 증상은 덜 보이게 되어 있었지만, 정작 무한 루프를 일으키는 근본 원인인 반복문 자체는 그대로였습니다. 증상을 가리는 것과 원인을 고치는 것은 다른 일입니다. 이번 수정은 증상이 시작되는 자리가 아니라, 멈추지 못하는 반복문 그 자체를 바로잡는 데 초점을 맞췄습니다.

---

작은 수정에서 배운 것

이 기여를 마치고 돌아보면, 바뀐 코드의 양은 많지 않습니다. 반복문에 조건 한 줄을 보태고, 예외적인 경우를 다루는 짧은 분기를 더하고, 그것을 지켜 주는 테스트 하나를 추가한 것이 전부입니다. 하지만 그 적은 변경에 도달하기까지 거친 과정에서 얻은 것은 줄 수와 비례하지 않았습니다.

 

가장 크게 남은 것은, 반복문을 짤 때 "이 반복이 반드시 끝나는가"를 스스로 물어야 한다는 감각입니다. 반복문이 끝나려면 매 반복이 종료를 향해 조금이라도 나아가야 합니다. 이 코드에서는 시작 위치가 매번 전진해야 했는데, 단어 경계를 맞추려는 배려가 특정 입력에서 그 전진을 삼켜 버렸습니다. 좋은 의도로 더한 처리가 진행을 가로막을 수 있다는 것, 그래서 반복문 안에서 위치를 조정하는 코드는 항상 "그래도 한 걸음은 나아가는가"를 함께 따져야 한다는 것을 이번에 분명히 배웠습니다.

 

또 하나는 경계 조건의 무게입니다. 이 버그는 평범한 입력에서는 전혀 드러나지 않고, 한 단어가 한 묶음보다 긴 극단적인 경우에만 나타났습니다. 대부분의 입력에서 멀쩡하다는 것이 곧 안전하다는 뜻은 아니었습니다. 가장 길거나 가장 짧은 입력, 한 종류로만 채워진 입력 같은 극단을 일부러 떠올려 보는 일이 왜 중요한지를, 이 한 단어짜리 무한 루프가 일러 주었습니다.

 

마지막으로, 버그의 성질에 따라 테스트의 방식도 달라져야 한다는 점입니다. 결과를 비교하는 평범한 테스트로는 무한 루프를 잡을 수 없었습니다. 끝나지 않는 고장이니, 끝나는지 자체를 시간 제한으로 검증해야 했습니다. 무엇을 검증하려는지에 따라 테스트의 도구가 달라진다는 것, 그리고 끝나기만 하는 것과 올바르게 끝나는 것을 모두 확인해야 비로소 수정이 완성된다는 것을 이번에 배웠습니다.

 

여전히 저는 오픈소스 기여를 익혀 가는 사람이고, 한 건의 작은 무한 루프 수정으로 무언가를 다 안다고 말할 수는 없습니다. 다만 멈추지 못하는 반복문의 원인을 차분히 따라가고, 평범한 동작을 건드리지 않으면서 극단적인 경우만 바로잡고, 그 동작을 지켜 줄 테스트를 남겨 두는 이 과정 자체가, 기능 하나를 더 만드는 것보다 저를 더 단단하게 만들어 준다는 것은 분명히 느꼈습니다. 다음에 또 어떤 멈추지 않는 반복문을 만나게 될지는 모르지만, 그때도 같은 방식으로 차분히 따라가 볼 생각입니다.

 

---

 

이 글에서 다룬 기여: langchain4j/langchain4j#5454