Conversational AI / NLP

NVIDIA cuDNN 9로 트랜스포머 가속화

Reading Time: 7 minutes

NVIDIA CUDA 딥 뉴럴 네트워크 라이브러리(cuDNN)는 최첨단 성능으로 딥 러닝 기본 요소를 가속화하기 위한 GPU 가속 라이브러리입니다.

cuDNN은 PyTorch, TensorFlow 및 XLA(가속 선형 대수)와 같은 인기 있는 딥 러닝 프레임워크와 통합됩니다. 이러한 프레임워크는 직접 GPU 프로그래밍의 복잡성을 추상화하므로 기본 하드웨어에 신경을 쓰는 대신 모델을 설계하고 트레이닝하는 데 집중할 수 있습니다. cuDNN은 내부에서 성능 엔진 역할을 하여 이러한 프레임워크에서 작업을 최대한 효율적으로 실행할 수 있게 해줍니다.

최근에는 SDPA(Scaled Dot Product Attention)가 거대 언어 모델(LLM)과 같은 주요 워크로드에서 성능을 좌우하는 결정적 기본 요소로 부상했습니다. cuDNN은 이 기본 요소에 대한 지원을 추가했으며 플래시 어텐션 및 그 외 최적화를 사용하여 릴리스마다 성능을 개선하는 동시에 기능 지원 영역을 확장하여 다양한 어텐션 사용 사례를 지원하고 있습니다. 

NVIDIA H200 Tensor 코어 GPU에서 cuDNN은 FP8에서 최대 1.2PFLOPS를 달성할 수 있습니다. 엔드 투 엔드 예시로, 저희 팀은 Llama2 70B LoRA 미세 조정을 위해 cuDNN FP8 SDPA를 활성화한 후 1.15배의 속도 향상을 측정했습니다. 이 실험에서는 8GPU H200 노드에서 NVIDIA NeMoNVIDIA TE(Transformer Engine)를 사용했습니다.

소프트웨어 스택e2e 가속
cuDNN이 비활성화된 NeMo + TE1x(기준)
NeMo + TE(BF16의 cuDNN SDPA 포함)1.11x
NeMo + TE(FP8의 cuDNN SDPA 포함)1.15x
표 1. 8-GPU H200 노드에서 엔드 투 엔드 트레이닝 실행(Llama2 70B LoRA 미세 조정)의 일부로 SDPA에 cuDNN을 사용 시 그에 따른 영향

이 게시물에는 cuDNN SDPA로 구현 가능한 성능에 관한 자세한 내용과 사용 방법을 살펴보며, cuDNN 9에서 주목할만한 그 밖의 새로운 기능을 간략히 요약하고 있습니다.

SDPA(Scaled Dot Product Attention) 성능

NVIDIA는 어텐션 알고리즘을 단일 커널로 융합하는 APEX 라이브러리의 융합 멀티헤드 어텐션(fMHA) 커널을 오픈 소스로 공개하여 어텐션을 지원하기 시작했습니다. Tri Dao의 혁신적인 작업은 이 커널을 출발점으로 사용하여 플래시 어텐션의 형태로 대규모 성능 개선과 기능성을 구현했습니다. 자세한 내용은 GitHub의 /Dao-AILab/flash-attentionFlashAttention-2: 더 나은 병렬 처리 및 작업 분할로 더 빨라진 어텐션 문서를 참조하세요.

그런 다음, NVIDIA는 더 빠르고 유연한 구현으로 융합된 어텐션의 최첨단 기술을 발전시켰습니다. 이 구현은 이제 NVIDIA Hopper 아키텍처 GPU를 위한 NVIDIA TE의 기본 융합 어텐션 백엔드입니다. 

XLA는 현재 cuDNN SDPA로의 경로를 제공하며 JAX SDPA API를 통해 액세스하거나 XLA 컴파일러에 의존하여 JAX/PyTorch의 맞춤형 구현에서 cuDNN SDPA로 낮출 수 있습니다. 

PyTorch 즉시 모드 SDPA는 현재 cuDNN을 사용하지 않는 상태지만, cuDNN 기반 구현이 진행 중입니다. 자세한 내용은 FpropBprop용 PyTorch PR을 참조하세요.  

cuDNN SDPA 구현에서는 다음을 캡슐화합니다.

  • 기본 NVIDIA 하드웨어 아키텍처에 대한 심층적인 이해
  • 비 플래시에서부터 플래시 어텐션 v2 및 그 사이의 모든 것에 이르기까지 모든 최첨단 SDPA 알고리즘의 구현
  • 문제 크기 및 대상 GPU 아키텍처에 따라 성능 조절(예: 타일 크기)을 자동으로 설정하는 휴리스틱

이는 NVIDIA GPU에서 최고의 가용 성능을 구현합니다. 그림 1과 2는 다양한 사용 사례에서 cuDNN 9 BF16이 BF16에서도 사용 가능한 최고의 PyTorch eager 구현보다 최대 2배 더 빠름을 보여줍니다. cuDNN FP8 구현은 최대 3배 더 빠른 속도로 진행됩니다.

성능이 향상되면 시퀀스 길이가 더 길어지고 모델의 사전 트레이닝 및 미세 조정 시간이 단축됩니다. 다른 공개 벤치마크(예: 플래시 어텐션 혹은 벤치마크)에 따라 이 게시물은 GPU 시간만 보고하며 호스트 오버헤드는 포함하지 않습니다.

그림 1. 인과 마스크, 헤드 차원 128이 있는 SDPA(순방향 전용)
그림 2. 인과 마스크, 헤드 차원 128이 있는 SDPA(순방향 + 역방향)
그림 3. 인과 마스크, 헤드 차원 256이 없는 SDPA(순방향 전용)

cuDNN 그래프로서의 SDPA

cuDNN의 SDPA는 Tensor 연산의 cuDNN 그래프로 지정할 수 있습니다. 주어진 그래프에서 cuDNN 라이브러리에는 이를 실행할 수 있는 일부 엔진 세트가 있습니다. 

일부 그래프에는 적합한 엔진이 없을 수도 있겠지만 이 엔진 세트는 GPU에서 원자적으로 실행되는 실용적인 cuDNN 그래프에 대해 적어도 하나의 엔진을 제공하는 데 그 의의가 있습니다. 엔진은 일반적으로 하나의 융합된 커널이지만 때로는 작은 협력 커널 세트에서 제공됩니다. 이는 전체 프레임워크 그래프의 “부분 그래프”로 볼 수 있습니다.

SDPA(다양한 형태)는 그러한 부분 그래프의 전형적인 예에 속합니다. 이러한 패턴을 지원하는 엔진은 플래시 어텐션과 같은 사용 가능한 최상의 알고리즘을 사용하여 성능을 크게 저하시키지 않고도 최대한 유연하게 설계되었습니다. 즉, 어텐션 연산을 변형하면서도 이를 GPU에서 효율적으로 실행할 수 있습니다. 이 섹션은 현재 가용한 유연성과 지원을 설명하는 데 그 목적이 있습니다.

그림 4는 순방향 전파(fprop) 사용 사례에 대한 cuDNN 그래프를 보여줍니다. 이는 뉴럴 네트워크 내에서 어텐션 메커니즘인 SDPA의 순방향 연산과 관련된 일련의 연산을 나타냅니다.

그림 4. SDPA를 위해 세분화된 fprop 연산 cuDNN 그래프 

인과 마스크 로직은 개별 gen 인덱스, 점별 연산 및 선택 연산으로 낮아집니다. softmax 로직은 6개의 cuDNN 연산으로 낮아집니다(그림 4).  cuDNN 연산의 이러한 세분화는 맞춤형 모델을 표현할 수 있는 유연성을 제공합니다. 그림 5는 가능한 조합을 보여줍니다. 

그림 5. cuDNN 어텐션 지원

그림 5는 헤드 차원이 더 크고(256) 인과 마스크가 없을 때 cuDNN FP8 포워드 플래시 어텐션이 최대 1.2PFLOP를 구현할 수 있음을 보여줍니다. 

BMM1과 Softmax 간에 임의의 점별 연산으로 맞춤형 그래프를 구성할 수도 있습니다. 이러한 유연성 덕분에 아직 발견되지 않은 새로운 변이를 지원할 수 있습니다. 

크리에이티브 연구원이 표준 어텐션 패턴을 수정하면 덜 최적화된 프레임워크 구현으로 되돌아가 성능 절벽에 부딪힐 가능성이 줄어듭니다. 그림 2는 전체 트레이닝 실행 결과를 나타냅니다. 역전파 어텐션 그래프에도 이와 유사한 유연성을 제공할 수 있습니다.

SDPA 사용 연습

cuDNN 그래프를 생성하고 실행하는 데 사용할 수 있는 여러 API 진입점을 소개하면 다음과 같습니다.

  • 프런트엔드 API(C++ 및 Python 변형 모두 포함)
  • 백엔드 API(C만 해당)

이전 섹션에서 설명한 모든 cuDNN 그래프 개념은 두 API 수준에 모두 적용됩니다. 하지만 C 인터페이스가 필요하지 않다면 cuDNN 팀은 Python 또는 C++에서 cuDNN 프런트엔드 API와 인터페이스로 연결할 것을 권장합니다. 프런트엔드 API는 훨씬 더 간결하며 몇 가지 편의성을 제공합니다.

예를 들어, 프런트엔드 API는 작업 노드의 개념을 확장하여 여러 작업과 작업 간의 데이터 흐름을 캡슐화하는 노드를 지원합니다. 즉, SDPA와 같은 일반적인 그래프 패턴의 세부 사항을 추상화하는 편리한 노드입니다. 노드는 잘 알려진 변형을 구성할 수 있도록 유연성을 계속 지원합니다. 이 SDPA 사용 연습은 가장 간단한 사례인 Python에서 생성된 SDPA 노드로 시작합니다.

프런트엔드 리포지토리의 SDPA Python 예제는 구성 옵션과 기본 사용 흐름을 나타냅니다.

  1. 적절한 데이터 유형으로 cudnn.pygraph 개체를 초기화합니다. 
  2. Tensor의 차원, 레이아웃 및 데이터 유형을 포함하는 Tensor 개체를 만듭니다. 
  3. 확장된 내적 플래시 어텐션 노드를 생성하고 필요한 구성을 제공합니다. 
  4. 그래프를 작성하고 디바이스 포인터를 제공합니다. 
  5. 그래프를 실행합니다.  
# 1. cuDNN graph
graph = cudnn.pygraph(
          io_data_type=cudnn.data_type.BFLOAT16,
          intermediate_data_type=cudnn.data_type.FLOAT,
          compute_data_type=cudnn.data_type.FLOAT,
        )
 
# 2. tensor attributes
q = make_tensor_attr(graph, q_gpu, "q")
k = make_tensor_attr(graph, k_gpu, "k")
v = make_tensor_attr(graph, v_gpu, "v")
attn_scale = make_tensor_attr(graph, attn_scale_cpu, "attn_scale", is_pass_by_value=True)
 
# 3. sdpa node with the configurations
o, stats = graph.scaled_dot_product_flash_attention(
                    name="scaled_dot_product_flash_attention",
                    q=q,
                    k=k,
                    v=v,
                    is_inference=false,
                    attn_scale=attn_scal,
                    bias=None,
                    use_alibi_mask=false,
                    use_padding_mask=false,
                    seq_len_q=None,
                    seq_len_kv=None,
                    use_causal_mask=true,
                    dropout=None,
                )
 
# 4. build the graph and provide the device pointers
graph.build()
variant_pack = {
                 q: q_gpu,
                 k: k_gpu,
                 v: v_gpu,
                 o: o_gpu
               }
 
# 5. execute the graph
graph.execute()

SDPA 노드가 기본적으로 제공하는 것보다 더 많은 유연성을 원한다면 기본 SDPA 그래프를 구성하는 코드는 오픈 소스이므로 맞춤화할 수 있습니다. 

예를 들어, 다음 코드 예제에서는 SDPA 노드 내의 확장 노드를 맞춤화합니다. 자세한 내용은 scaled_dot_product_flash_attention.h를 참조하세요. 

// Optional scale
if (options.inputs.Attn_scale) {
      // Lower options to scale options
      auto attn_scale_output = std::make_shared<Tensor_attributes>();
      attn_scale_output->set_is_virtual(true);
 
      Pointwise_attributes scale_attributes;
      scale_attributes.set_name("attn_scale");
      scale_attributes.set_mode(PointwiseMode_t::MUL);
      scale_attributes.inputs.IN_0 = last_output;
      scale_attributes.inputs.IN_1 = options.inputs.Attn_scale;
      last_output = scale_attributes.outputs.OUT_0 = attn_scale_output;
      auto scale_node = std::make_unique<PointwiseNode>(std::move(scale_attributes), context);
       
      sub_nodes.emplace_back(std::move(scale_node));
}

맞춤형 트랜스포머를 가속화하는 방법에 관한 자세한 내용은 NVIDIA cuDNN 설명서를 참조하세요.

그 외 주목할만한 cuDNN 9 기능

SDPA 개선 사항 외에도 cuDNN 9에는 다음과 같이 몇 가지 중요한 개선 사항이 추가되었습니다.

  • 행렬곱 및 합성곱에 대한 혼합 입력 정밀도 지원
  • 향상된 오류 보고
  • 하드웨어 이전 버전과의 호환성
  • 간소화된 설치

행렬곱 및 합성곱에 대한 혼합 입력 정밀도 지원 

입력 피연산자의 데이터 유형이 동일한 유형(FP16, FP32)이어야 하는 행렬곱(matmul) 및 합성곱 API은 활성화가 FP16에 있고 가중치가 INT8에 있을 수 있는 AWQ(Activation-aware Weight Quantization)와 같은 사례에 적합하지 않습니다. 더 높은 정밀도로 연산하려면 데이터를 캐스팅해야 합니다. 온라인이 아닌 경우, 전환 비용과 함께 추가 메모리가 필요합니다. 

cuDNN은 이제 성능 및 메모리 최적화를 위한 온라인 융합 유형 변환을 통해 A 및 B 피연산자가 서로 다른 데이터 유형일 수 있는 혼합 입력 정밀도 행렬곱과 합성곱을 지원합니다. 피연산자 중 하나를 다른 유형으로 캐스팅할 수 있습니다. cuDNN은 최적화된 커널에서 필요한 변환을 처리합니다.

그림 6은 cuDNN 혼합 입력 정밀도 행렬곱과 융합되지 않은 워크플로우 간의 속도 향상을 보여줍니다. 회색 막대는 입력 A와 B가 각각 FP16 및 INT8 정밀도이고 A가 INT8로 변환된 후 INT32 누산을 통해 INT8xINT8 행렬곱이 뒤따르는 사례를 나타냅니다. 녹색 막대는 A가 INT8에서 FP16으로 업컨버팅된 후 FP32 누산을 통한 FP16xFP16 행렬곱이 뒤따르는 사례를 나타냅니다.

그림 6. cuDNN에서 혼합 입력 정밀도 행렬곱 및 합성곱 지원

향상된 오류 보고

로깅은 소프트웨어 개발, 특히 cuDNN을 사용하는 딥 러닝 프레임워크와 같은 복잡한 시스템에서 필수적입니다. 과거에 cuDNN에서 흔히 발생한 문제점은 오류 및 경고 디버깅의 어려움이었습니다. NVIDIA는 이 문제를 해결하기 위해 오류 보고를 지속적으로 개선해 왔습니다. 

cuDNN 9는 다음과 함께 이 작업에 추가됩니다.

  • 보다 구체적인 오류 코드
  • 증가하는 오류 코드 수를 정리하는 데 도움이 되는 분류
  • 로깅 규칙에 부합하는 중첩된 로깅 수준
  • cuDNNGetLastErrorstring(프로그래밍 방식으로 마지막 오류 메시지를 가져오는 새 함수)

자세한 내용은 개발자 가이드의 오류 및 API 로깅 섹션을 참조하세요.

하드웨어 이전 버전과의 호환성

버전 9.0.0 이전의 cuDNN 라이브러리는 라이브러리 출시일에 공개적으로 사용 가능한 최신 GPU 아키텍처까지 지원했습니다. 예를 들어, cuDNN 8.9.x는 NVIDIA Hopper까지 지원했습니다(즉, 컴퓨팅 기능 9.0). 향후 GPU 아키텍처에서 8.9.x를 실행하는 것은 지원되지 않습니다. 

그러나 cuDNN 9은 API의 대규모 부분 집합에 대한 하드웨어 이전 버전과의 호환성을 제공합니다. 즉, cuDNN v9를 통해 이 API 부분 집합만 사용하는 프로그램은 향후 출시될 GPU에서도 작동하며 이러한 프로그램의 사용자는 향후 출시될 GPU를 사용하기 위해 cuDNN 설치를 억지로 업그레이드하지 않아도 됩니다. 

9.0 이상의 컴퓨팅 성능을 갖춘 일부 GPU에서 실행할 때 오류가 발생하는 대신 이전 버전과의 호환성은 라이브러리가 기능적으로 동등한 구현을 찾고 PTX JIT를 사용하여 새 아키텍처를 대상으로 한다는 것을 의미합니다.

지원 제한 및 모범 사례에 관한 자세한 내용은 개발자 가이드의 하드웨어 이전 버전과의 호환성 섹션을 참조하세요.

간소화된 설치

Python 환경에서는 이제 라이브러리 외에도 pip를 사용하여 새 Python 프런트엔드를 설치할 수 있습니다.

# cuDNN library for CUDA 12
pip install nvidia-cudnn-cu12 
 
# cuDNN frontend
pip install git+https://github.com/NVIDIA/cudnn-frontend.git

또한 cuDNN 9은 RPM 및 Debian 메타 패키지의 설치 프로세스를 간소화했습니다. 예를 들어, 다음 명령을 실행하면 Ubuntu 22에 cuDNN을 설치합니다.

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get install -y cudnn

이 새 흐름은 Debian 패키지 설치로 키링 설정을 간소화하고 cudnn 메타 패키지를 사용하여 정확한 cuDNN 라이브러리 버전을 추상화합니다.

자세한 내용 및 전체 지침은 cuDNN 설치 가이드를 참조하십시오. 

다음 단계

피드백, 질문 또는 의견이 있는 경우 cuDNN NVIDIA 개발자 포럼(cuDNN 라이브러리 주제 관련) 또는 GitHub의 NVIDIA/cudnn-frontend(프런트엔드 주제 관련)에 게시할 수 있습니다. cuDNN 다운로드 시작하기.

앞으로 더 많은 새 cuDNN 기능을 기대해 주세요. 이 게시물에서 간략하게 설명한 진행 상황은 중요한 이정표에 속하지만 앞으로 더 많은 일들이 일어날 것입니다. 

AI의 지속적인 확장으로 업계가 하드웨어와 소프트웨어 통합의 한계에 도달함에 따라 NVIDIA는 딥 러닝 프레임워크와 그래프 컴파일러 전반에서 cuDNN을 보다 효과적으로 광범위하게 사용할 수 있도록 성능을 최적화하고 사용자 경험을 지속적으로 개선하고 있습니다.

도움 주신 분들

NVIDIA cuDNN 팀은 회사 전반의 다른 많은 팀과 협업하여 이 게시물의 기술 콘텐츠를 제공했습니다.

관련 리소스

Discuss (0)

Tags