DGX Spark에 모델 추론을 위한 FlashAttention 2 설치 후기

설치 배경

MOSS-TTS 설치 후 추론 속도에 큰 불편함은 없었습니다. 다만 FlashAttention 2를 설치하면 추론 속도가 빨라지고 GPU 메모리 사용량이 절감될 수 있다고 하여, 실제로 얼마나 개선되는지 궁금했습니다. 이에 FlashAttention 2 설치를 진행했습니다.

FlashAttention 2란 무엇이고 추론 효율이 증가하는 원리는?

FlashAttention 2는 Transformer 계열에서 Attention 연산을 더 효율적으로 처리하기 위한 구현체로 이해하고 있습니다. 제 추측으로는, Attention 계산 과정에서 메모리 접근과 중간 텐서 생성을 줄이거나 연산 흐름을 최적화하여 속도 및 메모리 효율이 개선되는 방식이라고 생각했습니다. 다만 이 효과는 모델 구조, 입력 길이, dtype(torch.float16 / torch.bfloat16), GPU 아키텍처에 따라 차이가 날 수 있다고 판단했습니다.

DGX Spark에 FlashAttention 2 설치하기

MOSS-TTS README를 참조하면 다음과 같은 문구가 있습니다.

FlashAttention 2 is only available on supported GPUs and is typically used with torch.float16 or torch.bfloat16

torch.float16에서 사용할 수 있다면 DGX Spark에서도 사용할 수 있을 것으로 판단하고 설치를 시도했습니다.

flash-attn2-on-spark

1. 사전 확인 사항

  • DGX Spark는 CUDA 13.0을 사용하므로 --extra-index-url https://download.pytorch.org/whl/cu130에서 종속성 패키지를 찾아 설치했습니다. 그래도 확인을 위해 nvidia-smi로 CUDA 버전을 확인했습니다.

  • 설치 도중 PyTorch를 찾는 과정에서 기존 venv 환경에 이미 설치된 PyTorch를 그대로 사용해야 했습니다. venv 환경에서 빌드 시 임시 격리 환경이 생성되지 않도록 --no-build-isolation 옵션을 넣었습니다.

  • Spark 환경에서는 wheel 설치에 실패했고, 설치 로그에 다음이 출력되었습니다. aarch64 아키텍쳐 때문입니다. Spark를 쓰면서 자주 겪는 일이라 이제 놀랍지도 않습니다. 짜증도 안납니다. 익숙한 메시지입니다.

Precompiled wheel not found. Building from source...

선택의 여지가 없이 소스 빌드로 진행해야 합니다. 소스 빌드 과정에서 필요하므로 ninja를 venv에 설치합니다.

bash pip install ninja

  • 호스트 시스템에 Python 3.12 개발용 라이브러리가 필요하므로 없으신 분은 설치하세요.
sudo apt update
sudo apt install python3.12-dev -y

flash-attn은 C++와 CUDA 코드를 컴파일하여 Python과 연결하는 과정을 거치는데, 이때 Python의 내부 구조가 정의된 Python.h 파일이 필요했습니다. 일반적인 Python 실행 환경에는 이 파일이 포함되어 있지 않아 개발자용 패키지를 별도로 설치해야합니다.

2. 설치 명령

이번 포스트의 핵심입니다. 위 사항들을 모두 고려하여 다음 명령어로 설치했습니다.

TORCH_CUDA_ARCH_LIST="12.0" MAX_JOBS=1 pip install --no-build-isolation --extra-index-url https://download.pytorch.org/whl/cu130 -e ".[flash-attn]"

3. 명령어 조합 이유 (시행착오를 통한 결정)

처음에는 머신에서 gpt-oss-120b도 돌아가고 있는 상태였습니다. 그 상태에서 pip install ... -e ".[flash-attn]"를 실행했다가 CPU 사용량이 급격히 치솟으면서 시스템이 멈췄고, 터미널이 먹통이 되어 물리 스위치로 강제 재부팅했습니다. 이후에는 리소스를 쓰는 작업들을 모두 내리고 설치에 집중했습니다.

여러 번의 시행착오 끝에 위 명령어로 설치를 완료했습니다. 총 소요 시간은 1~2시간 정도 걸렸습니다. 정확히 측정하지는 않았지만, 1시간 이상 설치가 진행되어 자리를 비웠고, 식사 후 돌아왔을 때 설치가 완료되어 있었습니다.

설치 중 메모리는 지속적으로 약 24GB 정도 점유하는 느낌이었습니다. 문제는 CPU였고, 설치에 집중할 수 있도록 다른 작업을 내려두는 편이 안정적이었습니다.

옵션을 넣은 이유는 다음과 같습니다.

  • TORCH_CUDA_ARCH_LIST="12.0": Blackwell 아키텍처만 명시적으로 타겟팅하여 설치 시간을 단축하려는 목적이었습니다.
  • MAX_JOBS=1: 직전에 시스템이 멈춘 경험이 있어서 보수적으로 1로 설정했습니다. 그 결과 설치 시간이 60분 이상 소요되었습니다.

설치 후 추론의 향상 효과

1. 속도 측면

속도는 솔직히 크게 체감되지는 않았습니다. flash-attn 미설치 시에도 충분히 빨랐기 때문에, 수 초 단축이 있었다고 해도 체감상 “확실히 빨라졌다”는 느낌은 크지 않았습니다.

  • 7초 정도의 결과물을 생성하는데 8~9초 정도 소요되었습니다.
  • 25초 길이의 결과물을 생성하는데 32초 정도 소요되었습니다.
  • 16초 결과물은 21초 소요되었습니다.

즉, 대략 생성된 결과물 길이 대비 약 1.3배 정도의 추론 시간이 소요되는 느낌이었습니다.

2. 메모리 측면

메모리도 변화를 느끼기 어려웠습니다. 추론 중 nvidia-smi 수치를 지속 관찰했지만, 메모리 사용량이 눈에 띄게 늘거나 줄지는 않았습니다. 추론 도중 전력은 36W 내외였고, 온도는 46도에서 53도 정도로 상승하는 수준이었습니다.

정리

  • DGX Spark 환경에서 FlashAttention 2는 wheel 설치에 실패하여 소스 빌드로 설치했습니다.
  • 설치 자체는 성공했지만, 빌드 시간이 오래 걸렸고 CPU 부하가 상당했습니다.
  • 설치 후 속도와 메모리 측면에서 기대했던 만큼의 체감 개선은 크지 않았습니다.
  • 다음에 빌드를 다시 하게 된다면 MAX_JOBS=4 정도로 올려도 될 듯합니다. 빌드 시간이 이론상 1/4 수준으로 줄어들 것 같습니다.

관련글