[TIL] PyTorch + DeepSpeed 학습 중 !block->expandable_segment_ 에러 해결

문제 상황

• 환경: AWS g5.12xlarge, CUDA 12.4, PyTorch 2.4 + DeepSpeed ZeRO-3 + Flash-Attn + InternVL3-2B
• 에러 증상: 학습 도중 !block->expandable_segment_ 발생

RuntimeError: !block->expandable_segment_ INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":2586

• 특징:
  • 특정 step에서 GPU OOM으로 터짐
  • expandable_segments=True 설정 시 주로 재현
  • DeepSpeed ZeRO-3 + gradient checkpointing이 켜진 상태에서 자주 발생

원인 분석

• PyTorch CUDA allocator의 expandable segment 모드( PYTORCH_CUDA_ALLOC_CONF="expandable_segments=True" )와
  DeepSpeed ZeRO-3의 parameter offloading / all_gather hooks가 충돌하면서 발생.
• 특히 Vision Transformer(ViT) 쪽 torch.utils.checkpoint.checkpoint가 재진입(reentrant) 모드일 때,
  backward 시 동일 레이어 재호출로 allocator 내부 블록 확장 로직이 꼬이는 현상이 주요 트리거.

해결 방법

1. Allocator 설정 제거 : export PYTORCH_CUDA_ALLOC_CONF="max_split_size_mb=32,garbage_collection_threshold=0.8"  제거
2. ViT 체크포인트 비재진입 설정
   • 기존: torch.utils.checkpoint.checkpoint(self.layer_fn, hidden_states)
   • 수정: torch.utils.checkpoint.checkpoint( self.layer_fn, hidden_states, use_reentrant=False)
   • 비재진입 모드로 바꾸면, backward 시 함수 재호출을 피하고 메모리 충돌 위험이 줄어듦.

• PYTORCH_CUDA_ALLOC_CONF=expandable_segments=True는 CUDA 메모리 단편화를 줄여주는 대신, 큰 연속 메모리 블록을 재할당할 수 없게 강제하는 옵션으로, expandable_segments=True로 설정하면 여러 작은 청크로 나눠 쓰는 대신, 어떤 연산이 갑자기 큰 텐서를 요구하면 !block->expandable_segment_ 오류가 발생
• unset PYTORCH_CUDA_ALLOC_CONF하면, CUDA allocator가 기본 동작(필요시 큰 블록을 할당/재할당)으로 돌아가서, 큰 텐서 요청도 받아들이지만, 대신 메모리 단편화가 생길 수 있고 OOM 위험은 올라감
• 그런데 실제로는 expandable_segments=True로 터지는 원인은 ZeRO-3 all-gather 버퍼 + ViT flash-attn이 동시에 큰 contiguous 블록을 요구하기 때문으로, unset하였을 때 PyTorch allocator가 큰 블록을 새로 잡아줘서 OOM 없이 넘어감