1M 토큰 이상 학습

에이전트는 수 시간 동안 작업을 수행할 수 있습니다. 파일을 읽고, 명령을 실행하며, 세션의 모든 턴을 계속 이어 나갑니다. 이렇게 축적된 세션은 수십만 토큰에 달하며, 이것이 프런티어 모델들이 현재 백만 토큰 컨텍스트 윈도우를 홍보하는 이유입니다.

모델이 그 정도의 컨텍스트를 잘 처리하려면 그만큼 긴 시퀀스로 학습시켜야 하며, 바로 이 지점에서 난이도가 높아집니다. 100만 토큰의 시퀀스는 단일 GPU에 담기지 않으며, 8개의 GPU로도 담기지 않습니다. 도움을 받지 않는다면요.

이 가이드의 예제는 단일 8-GPU 노드에서 각 스텝마다 정확히 하나의 이러한 시퀀스로 학습합니다.

실행

8개의 H100(또는 더 높은 사양)이 장착된 노드 1개와 transformers>=5.17.0가 필요합니다. 이 예제는 해당 릴리스에서 추가된 gradient checkpointing의 offload를 사용합니다.

accelerate launch \
    --config_file examples/sft_qwen3_8b_1m_context/context_parallel_8gpu.yaml \
    examples/sft_qwen3_8b_1m_context/sft_qwen3_8b_1m_context.py

스크립트는 PG-19의 책들을 Qwen3-8B에 대해 fine-tune합니다. 각 책이 약 100만 토큰 길이가 될 때까지 끝에서 끝으로 이어 붙입니다. 로딩과 토큰화가 약 10분 정도 걸린 후, 첫 번째 스텝이 실행됩니다:

{'loss': 4.311, 'grad_norm': 29.25, 'num_tokens': 1049000.0, 'epoch': 0.25}
  8%|▊ | 1/12 [06:20<1:09:41, 380.10s/it]

읽어볼 가치가 있는 세 가지 숫자:

• num_tokens은 1.049e6입니다. 이는 짧은 시퀀스의 배치 하나가 아니라 하나의 시퀀스입니다. 모든 토큰은 전체 100만 토큰에 걸쳐 이전의 모든 토큰에 주의를 기울입니다.
• Loss는 4.3에서 시작하며, 이는 일반적인 초기 loss 값입니다. 장문 컨텍스트에 대해 세밀하게 잘못 구성된 실행은 약 10에서 시작하며, 그 이유는 나중에 다시 다룰 것입니다.
• 스텝당 380초. 한 스텝은 하나의 시퀀스에 해당하므로, 시퀀스당 약 6분 이상 소요됩니다.

Naively, 이 시퀀스는 GPU당 288 GB가 필요합니다. 이 가이드의 나머지 부분은 이를 56으로 줄이는 과정에 대한 이야기입니다.

방금 일어난 일

아래의 모든 내용은 이미 알고 있는 설정, 즉 GPU 1개와 시퀀스당 수천 개의 토큰을 기반으로 합니다. 시퀀스가 커지면 문제가 발생하기 시작하며, 이는 일정한 순서로 발생합니다.

손실이 가장 먼저 깨지는 부분입니다

Qwen3-4B를 80 GB 카드 1개에서 실행하고 시퀀스를 늘려 보세요. 메모리가 가장 먼저 부족해지는 것은 모델이 아닙니다. 손실입니다. 메모리 프로파일을 확인해 보세요.

피크는 포워드 패스와 백워드 패스 사이에 위치하며, 여기서 손실이 계산됩니다. 피크가 그 위치에 있는 이유를 이해하려면 디코딩의 마지막 단계에서 어떤 일이 일어나는지, 그리고 손실이 그 결과로부터 어떻게 구성되는지를 살펴봐야 합니다.

디코더의 마지막 레이어는 다음과 같은 형태의 hidden state를 출력합니다$L \times H$, 여기서$L$ 은 시퀀스 길이이고$H$ 은 hidden size입니다. 이 hidden state는 language modeling head의 가중치와 곱해져 다음과 같은 형태의 logits 행렬을 생성합니다$L \times V$, 여기서$V$ 은 어휘 크기입니다. 이 행렬은 시퀀스가 길고 어휘가 크므로(10만 개 이상의 항목) 매우 큽니다. 손실은 이 행렬과 대상 라벨을 사용하여 cross entropy 함수로 계산되며, 이 행렬을 물리적으로 생성하는 것이 바로 피크의 원인입니다.

전체 logits 행렬을 물리적으로 생성하지 않고 손실을 계산할 수 있을까요? 네, 청크 단위로 진행할 수 있습니다.

매우 큰 logits 행렬: 청킹

logits를 전체 시퀀스 단위로 한 번에 계산하는 대신, 수백 개의 행을 한 번에 가져와 해당 행에 대한 손실을 계산한 후 그 조각들을 합산합니다. 이렇게 하면 logits 행렬 전체가 한 번에 메모리에 유지되지 않습니다. 행은 깔끔하게 분할되는 차원입니다. 손실은 행에 대한 합이며, 그 안의 softmax는 한 행의 어휘를 따라 실행되므로, 행의 청크는 자체 손실에 필요한 모든 것을 포함합니다.

큰 곱셈 하나가 여러 개의 작은 곱셈으로 나뉘고, 메모리 절감 효과가 충분히 커서 손실이 다시 피크가 되지 않습니다. TRL은 청크당 256행을 사용하며, 이는 각 곱셈이 GPU를 효율적으로 활용할 수 있을 만큼 충분히 크도록 유지합니다.

새로운 프로필은 다음과 같습니다.

피크가 사라졌습니다.

TRL에서 이를 활성화하는 방법은 무엇인가요? 별도로 할 일이 없습니다. chunked loss가 기본값입니다. plain loss를 사용하려면 SFTConfig에서 loss_type="nll"를 설정하세요.

이로 인해 시퀀스를 얼마나 확장할 수 있을까요? 메모리 관점에서, 이 변경 사항 하나만으로도 32k 토큰에서 160k 토큰까지 도달할 수 있습니다.

모델은 100,000번째 위치를 본 적이 없습니다: YaRN으로 위치를 재조정하세요

이제 메모리에 160k 토큰을 담을 수 있습니다. 하지만 모델이 실제로 이를 읽을 수 있을까요?

Qwen3-4B는 40,960 토큰 길이의 시퀀스로 학습되었습니다. 설정 파일에도 그렇게 명시되어 있습니다:

>>> from transformers import AutoConfig
>>> AutoConfig.from_pretrained("Qwen/Qwen3-4B").max_position_embeddings
40960

그 위치 이후의 모든 위치는 모델이 한 번도 경험해 본 적 없는 위치입니다. 크래시는 발생하지 않지만, 모델이 배치할 수 없는 토큰으로만 학습이 진행됩니다. 160k 토큰 시퀀스의 모든 토큰에 대한 손실(loss)은 다음과 같습니다:

모델은 40,960에서 무너지지 않습니다. 한동안 더 진행하다가 70k를 넘어서면서부터 비로소 맥락을 잃기 시작합니다. 이후에는 꾸준히 상승하여, 동일한 텍스트의 unigram 엔트로피(6.45)를 초과하는 수준에서 끝납니다. 이는 컨텍스트를 완전히 무시하고 토큰 빈도만으로 예측하는 모델보다 더 나쁜 결과입니다.

각 토큰은 자신의 위치를 가지며, 모델은 0에서 40,960 사이의 위치가 무엇을 의미하는지 학습했습니다. 위치 100,000에 대해 질문하면 모델은 이를 처리할 근거가 전혀 없습니다.

해결책은 모델이 모르는 위치를 절대 전달하지 않는 것입니다. 위치를 재조정합니다. 위치 100,000은 위치 25,000이 되며, 시퀀스의 모든 위치가 모델이 아는 범위로 돌아옵니다. 토큰은 변하지 않으며, 각 토큰에 부여되는 위치만 바뀝니다. YaRN가 하는 일이 바로 이것이며, 계수는 단순히 우리가 얼마나 더 멀리 가려는지를 나타냅니다. 학습은 40,960까지, 원하는 것은 163,840까지이므로 4입니다.

TRL에서 이를 활성화하는 방법은 무엇인가요? 모델이 로드될 때 RoPE 구성을 전달하세요:

training_args = SFTConfig(
    ...,
    model_init_kwargs={
        "rope_parameters": {
            "rope_type": "yarn",
            "rope_theta": 1_000_000,  # has to match what the model ships with
            "factor": 4.0,
            "original_max_position_embeddings": 40960,
        },
    },
)

재조정된 실행을 위에 배치한 동일한 측정값:

리스케일링된 실행은 160k까지 일정한 상태를 유지합니다. 마지막 20k 토큰 구간에서 평균 손실은 2.8로, 리스케일링을 적용하지 않은 경우의 7.3보다 낮습니다.

활성화 값이 남았습니다. 이를 오프로드하세요

손실이 청크로 나뉘고 위치가 재조정된 후, GPU 1개로 160k 토큰까지 처리할 수 있습니다. 하지만 다시 용량이 부족해집니다. 이번에는 카드가 무엇으로 가득 차고 있을까요?

모델과 옵티마이저는 시퀀스 길이에 관계없이 고정된 22 GB를 차지합니다. 그 위에 있는 각 빨간색 밴드는 하나의 레이어에서 저장된 활성화 값으로, 각각 0.47 GB입니다. 포워드 패스에서 작성되면 쌓이고, 백워드 패스에서 하나씩 소비될 때까지 그대로 유지됩니다. 맨 위의 회색 부분은 밀리초 내에 재계산되고 폐기되는 모든 것을 나타냅니다.

TRL은 이미 이를 자동으로 줄여 주며, 기본적으로 활성화되어 있습니다. Gradient checkpointing은 각 레이어에 입력되는 데이터만 유지하고, 레이어 내부에서 계산된 모든 것, 즉 attention scores, 넓은 MLP intermediate, norms를 폐기합니다. backward pass에서 이 값들이 필요하면 레이어를 두 번째로 실행하여 다시 계산합니다. 최종적으로 각 레이어당 sequence x hidden 크기의 저장 텐서 하나만 남게 되며, 긴 시퀀스 길이의 경우에도 이것이 여전히 GPU 카드에서 가장 큰 메모리 점유 항목입니다.

유용한 관찰은 이러한 텐서가 forward pass 동안 전혀 필요하지 않다는 점입니다. 이들은 기록된 후 방치되고, 훨씬 나중에 읽힙니다. 따라서 그 사이에 GPU에 머물러 있을 필요가 없습니다. CPU 메모리로 이동했다가 backward pass가 해당 텐서에 도달할 때 다시 가져올 수 있습니다.

training_args = SFTConfig(..., gradient_checkpointing_kwargs={"offload": True})

같은 단계, 같은 길이: 38개의 밴드가 쌓이는 대신 한 번에 4개만 상주하며, 피크가 59.9 GB에서 48.8 GB로 감소합니다. 이를 통해 확보할 수 있는 시퀀스 길이는 다음과 같습니다:

48k 미만에서는 두 방식의 성능이 동일합니다. 해당 길이에서는 이동할 데이터가 거의 없기 때문입니다. 48k를 초과하면 오프로드된 실행의 성능이 더 낮게 유지되며, 기존에 160k에서 중단되던 시퀀스가 동일한 카드에서 256k까지 도달하게 됩니다.

이 방식은 무료가 아닙니다. 저장된 각 activation은 호스트로 건너간 뒤 다시 돌아와야 하며, 이 트래픽은 인접한 compute 사이를 통과해야 합니다.

GPU 1개로는 부족합니다: 시퀀스를 여러 개로 분할

Chunked loss, offloaded activations: 이제 1개의 카드로 256k 토큰을 학습할 수 있습니다. 100만 토큰에 도달하려면 어떻게 해야 할까요? 방해 요소를 확인하려면, 백워드 패스 중 단일 레이어의 내부(위 프로파일의 회색 피크)를 살펴보세요.

Gradient checkpointing은 forward가 거의 아무것도 저장하지 않는다는 것을 의미합니다: 각 레이어는 자신의 입력만 저장하고 나머지는 버립니다. 해당 레이어는 backward가 도달할 때 다시 계산되며, 그 순간 모든 것이 한 번에 존재하게 됩니다. MLP에서 모델의 한 줄,

down_proj(act_fn(gate_proj(x)) * up_proj(x))

4개의 텐서가 전체 중간 폭으로 생성되며, 표준 구현에서는 backward가 해당 텐서들을 지나칠 때까지 4개 모두를 유지합니다. 100만 토큰, MLP가 각 토큰을 2560에서 9728로 확장하는 모델에서:

gate_proj(x)                             19 GB = 1,048,576 x 9728 x 2 bytes
act_fn(gate_proj(x))                     19 GB
up_proj(x)                               19 GB
act_fn(gate_proj(x)) * up_proj(x)        19 GB
                                  total  76 GB

76 GB, 하나의 레이어 안에서, 80 GB 카드는 이미 모델과 옵티마이저에 30 GB를 할당했습니다. 이 시점에서 답은 시퀀스를 더 지능적으로 저장하는 것이 아닙니다. 각 GPU에 이를 덜 할당하는 것입니다.

각 GPU에 시퀀스의 한 슬라이스를 할당합니다. 4개의 GPU를 사용하면 각 GPU는 1,048,576개 대신 262,144개의 토큰을 보유하게 됩니다. 4개의 텐서도 이에 따라 축소되어 76 GB 대신 19 GB가 되며, 모델의 30 GB와 함께 카드에 맞습니다.

명백한 반론은 어텐션입니다. 모든 토큰은 이전의 모든 토큰에 대해 어텐션을 수행해야 하며, 어떤 GPU에도 이전 토큰이 존재하지 않습니다. 따라서 GPU들은 자신에게 없는 정보를 교환하고, 그 결과는 하나의 거대한 GPU가 계산했을 것과 정확히 동일합니다.

exchange를 실행하는 방법은 두 가지가 있습니다. 컨텍스트 병렬성(CP)과 Ulysses 시퀀스 병렬성(SP)입니다.

FSDP2에서는 CP를 사용합니다. 각 GPU는 자신의 슬라이스를 유지하고, 다른 슬라이스가 순서대로 해당 GPU로 전달되므로, 모든 토큰이 결국 그 이전의 모든 토큰을 확인하게 됩니다.

parallelism_config:
  parallelism_config_cp_size: 4

DeepSpeed에서는 SP를 사용합니다. 슬라이스를 이동하는 대신, attention 직전에 배치만 재배열하여 각 GPU가 모든 토큰은 보유하되 attention head의 4분의 1만 담당하게 하고, 이후 다시 원래대로 복원합니다.

parallelism_config:
  parallelism_config_sp_size: 4

현재 두 설정은 서로 겹치지 않습니다: cp_size는 FSDP2를 요구하며 sp_size는 DeepSpeed 환경에서만 실행됩니다. 백엔드를 선택하면 방법이 따라옵니다. 이 섹션의 나머지 부분은 이 가이드 상단의 예제가 사용하는 FSDP2 경로를 따릅니다. 해당 예제에서 GPU는 서로 다른 배치의 독립적인 워커가 아니라, 하나의 시퀀스를 공유하는 단일 그룹입니다.

켜면 두 가지가 변경됩니다:

1. 시퀀스는 cp_size * 2의 배수로 패딩되어야 하므로, 4개의 GPU를 사용하는 경우 SFTConfig에서 pad_to_multiple_of=8을 설정해야 합니다.
2. causal-SDPA 요구 사항은 문서들이 서로 읽지 못하도록 유지하기 위해 블록 대각선 마스크에 의존하는 packing을 배제하며, TRL은 두 가지를 동시에 요청하면 예외를 발생시킵니다. 또한, 슬라이딩 윈도우 또는 청크드 어텐션을 사용하는 레이어를 가진 모델도 배제되며, accelerate는 이러한 모델을 거부합니다: OpenAI GPT-OSS, Gemma 3 및 4, Qwen3.5 및 이후 버전. Qwen3와 Qwen3-MoE는 전체적으로 full attention을 사용하므로, 이 문서에서 다루는 모델입니다.

슬라이스를 주고받는 데 드는 비용은 예상보다 적습니다. 1개 노드에서 131k 토큰 기준으로, 2개 GPU에서는 스텝당 34.6초, 4개에서는 17.8초, 8개에서는 9.5초가 소요됩니다. 그룹 크기를 두 배로 늘릴 때마다 스텝 시간이 거의 절반으로 줄어, 2개에서 8개로 확장하면 이론상 최대 4배 중 3.7배의 성능 향상을 달성합니다.

그리고 이것이 네 가지 레버 중 마지막입니다:

카드 1장은 256k를 약간 넘겨서 멈춥니다. 4장은 100만 토큰 전체를 처리하며, 여유 공간이 남도록 63.6 GB에 도달합니다. 이 가이드의 서두에 등장한 시퀀스가 여기에 맞습니다.

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True를 설정합니다. 100만 토큰에서 위 실행은 80 GB 카드에서 63.6 GB가 필요하며, 이 설정이 없으면 여전히 실패합니다: 요청하는 텐서는 크고 연속적이어서, 할당기가 이미 보유하고 있는 블록들 사이에서 이를 위한 공간을 찾을 수 없기 때문입니다.

추가 자료

• 컨텍스트 병렬성, cp_size와 이를 통해 구축되는 디바이스 메시에 대한 accelerate 가이드입니다.
• 두 교환 방식의 차이점과 각 방식의 통신 비용을 설명하는 Ulysses와 ring attention.
• RoPE 섹션에서 위치 재스케일링에 사용된 YaRN.