Gradient checkpointing

forward pass는 일반적으로 backward pass에서 재사용하기 위해 모든 중간 activations를 캐시합니다. 그러나 activations는 배치 크기와 시퀀스 길이에 따라 비례하여 증가합니다. Gradient checkpointing은 특정 activations만 저장하고 나머지는 폐기합니다. 이를 통해 backward pass는 필요할 때 일부 activations를 실시간으로 다시 계산하도록 강제됩니다.

Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it

backward pass에서 폐기된 activations를 다시 계산하기 때문에 학습은 일반적으로 더 느리지만, checkpointing은 activation memory를 줄입니다.

gradient_checkpointing=True로 설정하여 활성화합니다.

메모리 사용량을 추가로 줄이려면 gradient accumulation과 함께 사용하세요.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)

부분 체크포인트

Full gradient checkpointing은 모든 체크포인트 가능한 레이어를 다시 계산합니다. 런에 메모리 여유가 있다면, 레이어를 덜 체크포인트하여 메모리 절감의 일부를 속도로 교환할 수 있습니다.

every_n_layers를 gradient_checkpointing_enable()에 전달하여 체크포인트 간격을 선택합니다.

every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun

every_n_layers=2를 사용하면 첫 번째 레이어와 그 이후의 모든 두 번째 레이어가 체크포인트됩니다. 체크포인트된 레이어는 포워드 패스 동안 활성화 값을 폐기하고 백워드 패스 동안 다시 계산하는 반면, 다른 레이어는 활성화 값을 메모리에 유지합니다.

model.gradient_checkpointing_enable(every_n_layers=2)

기본값인 every_n_layers=1은 모든 레이어에 대해 체크포인트를 생성합니다. 값이 더 크면 첫 번째 레이어에 체크포인트를 생성한 후, 이후 n개 레이어마다 체크포인트를 생성하며, 나머지 레이어의 활성화 값은 메모리에 유지됩니다. 예를 들어, every_n_layers=2인 경우 1, 3, 5번 레이어 등에 체크포인트를 생성합니다. GradientCheckpointingLayer를 상속하는 모듈만 카운트에 포함됩니다. 그레디언트 체크포인트를 지원하는 다른 모듈은 계속 활성화된 상태로 유지됩니다.

To use partial gradient checkpointing with Trainer, set every_n_layers in gradient_checkpointing_kwargs.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)

저장된 활성화 값 오프로딩

Gradient checkpointing은 체크포인트가 적용된 각 레이어당 하나의 activation을 GPU에 유지합니다. 긴 시퀀스 길이의 경우, 상당한 메모리를 소모할 수 있습니다(대략 layers x sequence x hidden x bytes_per_element). offload을 설정하면 이러한 activation을 고정된 호스트 메모리에 유지할 수 있으며, 이 경우 forward pass 중에 device-to-host 복사, backward pass 중에 host-to-device 복사가 발생하는 비용이 발생합니다.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

두 복사본 모두 compute stream에서 실행되므로, 이는 느린 스텝을 메모리와 교환하는 것입니다. 실행이 다른 방법으로 적합하지 않을 때 사용하며, 속도를 높이기 위한 용도로는 사용하지 마세요.

다음 단계

GPU 메모리 사용량이 학습 중 GPU에서 어떻게 발생하는지 이해하려면 GPU memory usage 문서를 읽어보세요.
메모리를 추가로 줄이고 학습 속도를 높이기 위해 더 낮은 정밀도의 데이터 타입을 사용하는 방법을 배우려면 Mixed precision training 가이드를 참조하세요.
커스텀 퓨즈드 커널을 사용하여 학습을 가속화하는 방법을 배우려면 Kernels 가이드를 참조하세요.

Update on GitHub

Transformers

Gradient checkpointing

부분 체크포인트

저장된 활성화 값 오프로딩

다음 단계