머신 러닝을 위한 머신 러닝 기술의 발전

 (blog.research.google)
1P by GN⁺ | ★ favorite | 댓글 1개
  • ML 모델이 커질수록 실행 효율은 하드웨어뿐 아니라 컴파일러 최적화에 크게 좌우되며, Google Research와 Google DeepMind는 이 영역을 다시 ML로 개선하는 접근을 공개함
  • TensorFlow, JAX, PyTorch 같은 프레임워크 아래에서 내려지는 컴파일러 결정은 같은 모델이라도 실행 시간과 자원 사용을 크게 바꿀 수 있음
  • TpuGraphs는 TPU용 ML 프로그램의 계산 그래프, 컴파일 설정, 실행 시간을 담아 학습 기반 비용 모델 연구에 쓸 수 있는 데이터셋임
  • 큰 계산 그래프를 한 번에 학습하기 어려운 제약을 줄이기 위해 Graph Segment Training은 그래프를 세그먼트로 나눠 메모리 사용을 낮추고 학습 시간을 3배 단축함
  • Kaggle 대회에서는 66개국 616팀 792명이 참여해 그래프 압축, 패딩 값 조정, 노드 특성 추가, 설정 간 attention 등 실제 비용 예측 모델 개선 기법을 검증함

ML 컴파일러가 실행 성능을 좌우하는 이유

  • 최신 ML 모델은 자연어 이해, 대화, 이미지 생성, 동영상 생성 같은 작업을 수행하며, TensorFlow, JAX, PyTorch 같은 ML 프로그래밍 프레임워크로 작성·학습됨
  • 프레임워크는 행렬 곱, 컨볼루션 같은 선형대수 연산과 2D 컨볼루션 레이어, 트랜스포머 레이어 같은 신경망 계층을 제공함
  • 사용자는 모델을 하드웨어에서 효율적으로 실행하는 세부 방식을 직접 다루지 않아도 되며, 프레임워크 아래의 컴파일러가 모델을 자동 최적화함
  • 다만 컴파일러는 복잡한 최적화 문제를 휴리스틱으로 푸는 경우가 많아 항상 최적 성능을 내기는 어려움

계산 그래프와 두 단계 최적화

  • ML 컴파일러는 사용자가 작성한 수학적 명령을 실제 하드웨어에서 실행할 수 있는 명령으로 변환함
  • ML 프로그램은 계산 그래프로 표현될 수 있음
    • 노드는 matrix multiplication 같은 텐서 연산을 나타냄
    • 엣지는 한 노드에서 다른 노드로 흐르는 텐서를 나타냄
  • 컴파일러 최적화는 크게 두 종류로 나뉨
    • 그래프 수준 최적화: 전체 그래프의 맥락을 고려해 결정을 내리고 그래프 전체를 변환함
    • 커널 수준 최적화: fused subgraph인 하나의 커널을 다른 커널과 독립적으로 변환함

메모리 레이아웃의 성능 트레이드오프

  • 행렬 같은 2D 텐서는 메모리에 [A B C a b c] 또는 [A a B b C c] 형태로 저장될 수 있으며, 이는 각각 row-major와 column-major 레이아웃에 해당함
  • ML 컴파일러의 중요한 최적화 중 하나는 프로그램의 모든 중간 텐서에 메모리 레이아웃을 배정하는 작업임
  • 특정 레이아웃이 개별 연산에는 가장 효율적일 수 있지만, add와 convolution 사이에서 레이아웃이 맞지 않으면 컴파일러가 추가 copy 연산을 삽입해야 함
  • 반대로 개별 연산 성능은 조금 낮아도 레이아웃 변환이 필요 없는 구성이 전체 실행에서는 더 나을 수 있음
  • XLA 벤치마크 suite에서 기본 컴파일러 설정 대신 최적 레이아웃 구성을 선택했을 때 최대 32% 속도 향상이 관찰됨

TpuGraphs 데이터셋

  • TpuGraphs는 Google의 커스텀 TPU에서 실행되는 프로그램을 위한 학습 기반 비용 모델 데이터셋임
  • 목표는 입력 프로그램과 컴파일러 설정을 받아 프로그램의 실행 시간을 예측하는 비용 모델을 학습하는 것임
  • 데이터셋은 두 가지 XLA 컴파일러 설정을 대상으로 함
    • layout: 행렬의 row-major·column-major 개념을 고차원 텐서로 일반화한 설정
    • tiling: 타일 크기 설정
  • 각 예제는 ML 워크로드의 계산 그래프, 컴파일 설정, 해당 설정으로 컴파일했을 때의 실행 시간을 포함함
  • 그래프는 오픈소스 ML 프로그램에서 수집됐으며 ResNet, EfficientNet, Mask R-CNN, Transformer 같은 모델 아키텍처를 포함함
  • 다운로드 방법과 시작 코드는 TpuGraphs GitHub에 제공됨
  • TpuGraphs는 비슷한 그래프 크기의 기존 최대 graph property prediction 데이터셋보다 그래프 수가 25배 많고, 기존 ML 프로그램 성능 예측 데이터셋보다 평균 그래프 크기가 770배 큼

기준 비용 모델과 GNN 구조

  • TpuGraphs는 기준 학습 기반 비용 모델을 함께 제공하며, 입력 프로그램이 그래프로 표현되기 때문에 GNN을 사용함
  • 노드 특성은 두 부분으로 구성됨
    • opcode id: 텐서 연산 종류를 나타내는 가장 중요한 노드 정보
    • 그 외 노드 특성
  • 기준 모델은 opcode id를 embedding lookup table을 통해 opcode embedding으로 변환함
  • opcode embedding과 나머지 노드 특성을 결합해 GNN 입력으로 사용함
  • GNN이 만든 노드 embedding은 sum과 mean 같은 단순 graph pooling reduction으로 고정 크기 그래프 embedding으로 결합됨
  • 최종 그래프 embedding은 feedforward layer를 거쳐 하나의 scalar 출력으로 변환됨

Graph Segment Training으로 큰 그래프 학습

  • Graph Segment Training은 제한된 메모리 용량의 장치에서 큰 그래프를 다루기 위한 GNN 학습 확장 기법임
  • 이 방법은 예측 대상이 노드나 엣지가 아니라 전체 그래프인 graph-level prediction 상황을 겨냥함
  • 계산 그래프는 수십만 개 노드를 포함할 수 있어, 전체 그래프를 한 번에 사용하는 Full Graph Training은 계산적으로 불가능할 수 있음
  • GST는 큰 그래프를 작은 세그먼트로 나누고, 세그먼트의 무작위 부분집합만 선택해 모델을 업데이트함
  • 나머지 세그먼트는 중간 activation을 저장하지 않고 embedding을 생성해 메모리 사용을 줄임
  • 모든 세그먼트 embedding을 결합해 원래 큰 그래프의 embedding을 만들고, 이를 예측에 사용함
  • historical embedding table과 segment dropout을 함께 도입해 historical embedding의 staleness를 완화함
  • 전체 방법은 end-to-end 학습 시간을 3배 단축함

Kaggle 대회에서 검증된 개선 기법

  • Fast or Slow? Predict AI Model Runtime Kaggle 대회는 TpuGraphs 데이터셋을 기반으로 열렸고, 66개국 616팀 792명이 참여함
  • 제출 수는 10,507건이었으며, 153명에게는 첫 Kaggle 대회였고 이 중 47명은 상위 100명 안에 들었음
  • 참여 팀들은 여러 기법을 실험함
    • 그래프 pruning·compression: GST 대신 큰 그래프를 압축하는 방법을 실험했으며, 설정 가능한 노드와 그 즉시 이웃을 포함하는 subgraph만 유지하는 방식이 사용됨
    • 패딩 값 변경: 기본 padding 값 0이 유효한 feature 값과 충돌하므로 -1을 쓰면 모델 정확도를 크게 개선할 수 있었음
    • 노드 특성 추가·인코딩 변경: dot general의 contracting dimensions 같은 추가 노드 특성이 중요하며, 노드 특성 인코딩 방식도 결과에 영향을 줄 수 있음
    • cross-configuration attention: 우승 팀은 모델이 설정들을 명시적으로 비교할 수 있는 단순 layer를 설계했고, 각 설정을 개별적으로 추론하게 하는 방식보다 훨씬 나은 결과를 냄
  • 대회 결과와 우승 솔루션은 2023년 12월 16일 NeurIPS의 ML for Systems workshop 대회 세션에서 다뤄질 예정임

NeurIPS Expo 관련 세션

  • 구조화 데이터와 인공지능 연구에 관심 있는 독자를 위해 NeurIPS Expo 패널 Graph Learning Meets Artificial Intelligence가 2023년 12월 9일 열림
  • 이 패널은 학습 기반 비용 모델의 발전 등을 다룸

함께 보면 좋은 글 β

GN⁺   [-]
Hacker News 의견들

  • ML 컴파일러는 과장되고 있음. 전통적인 컴파일러와 같은 절충임: 전문 성능 프로그래머를 고용하는 것보다 처리량은 훨씬 많이 얻지만, 후자는 보통 훨씬 더 빠르고 경우에 따라 몇 자릿수 차이로 앞설 수 있음
    여러 층위에서 부족함. 알고리즘 차원에서는 네트워크를 더 빠르게 바꾸는 요령을 사람에게 되먹임해 주지 못하고, 아주 기본적인 신호 정도만 줌. 의도도 손실됨. ML 네트워크 설계자는 Python으로 구조를 지정하지만 여러 단계의 하향 변환을 거치면 완전히 엉뚱한 결과가 나올 수 있음. 최근에는 어떤 컴파일러가 slice update를 하면서 배열의 가능한 모든 인덱스 범위를 만들고, 그걸 잘라 업데이트할 인덱스를 얻은 뒤 scatter를 하는 걸 봤고, 단일 memcpy 호출로 대체함. 커널도 비효율적임. 이런 컴파일러 출력과 숙련된 어셈블리 프로그래머가 맞붙을 때마다 컴파일러가 지고, 보통 30% 이상 차이남. 쉽게 해결될 것처럼 보이지만 지난 50년 동안 아무도 제대로 못 풀었다면 말처럼 단순하지 않은 게 분명함

    • 체스 엔진 Stockfish를 보면, 보드 평가에서 수년간 사람이 작성한 휴리스틱을 버리고 작은 신경망으로 바꿨는데 더 잘함
      컴파일러에도 인라인화, 루프 펼치기, 벡터화 같은 휴리스틱이 많으니 신경망이 도움이 될 수 있고, 사람이 작성한 수많은 휴리스틱보다 유지보수도 쉬울 수 있음
    • 전통적인 컴파일러와 같은 절충이라고 했는데, 그 전통적인 컴파일러가 엄청나게 유용했다는 점이 재미있음
    • 너무 단정적이고 닫힌 태도로 보임
    • 맞음. 손으로 어셈블리를 작성할 사람을 고용하면 되는데 왜 누가 gcc/clang을 쓰겠음?
    • 그 처리량이 핵심임. 모든 ML 작업마다 성능 전문가를 붙일 수는 없음
      이런 최적화가 아예 없는 것보다는 여전히 훨씬 낫다

  • 이걸 좀 현실적으로 설명해 줄 수 있음? 현재 ML 컴파일러의 실제 상태가 어떻고, 가까운 시일 내에 기대할 수 있는 게 뭔지 궁금함

    • 가장 쉬운 접근 중 하나는 torch.compile임. PyTorch 컴파일러의 최신 반복판이고, 이전 방식으로는 TorchScript와 FX Tracing이 있었음
      그냥 model = torch.compile(model)이라고 쓰면 됨. “이 163개 오픈소스 모델 전반에서 torch.compile은 93%의 경우 동작했고, NVIDIA A100 GPU에서 학습이 43% 빨라졌다. Float32 정밀도에서는 평균 21%, AMP 정밀도에서는 평균 51% 빨라졌다.”[1] Google은 이런 방법의 연구개발에 더 많은 사람이 참여하게 하려는 것으로 보임
      [1] https://pytorch.org/get-started/pytorch-2.0/
    • 가까운 시일 내 기대는 모델이 개발된 프레임워크에 대해 명시적인 벤더 지원이 없어도 AMD, CUDA, TPU, CPU 등을 쓸 수 있게 되는 것임
      현실은 복잡하니 꽤 단순화해서 말하면, 계산 그래프를 어떤 중간 표현으로 컴파일하고 거기에 맞는 백엔드를 구현하는 방식임. 관련 프로젝트로는 stableHLO, IREE, openXLA를 보면 됨. Jax의 jit 컴파일러도 이런 컴파일러의 한 형태라고 볼 수 있음. 추적된 연산을 XLA로 내리고, XLA가 다시 백엔드에서 동작하도록 여러 마법을 부림. 결국 아래로 내려갈수록 계속 변환과 추상화임
    • torch.compile을 보면 됨

  • 요약하면, 계산 그래프의 실행 시간 성능 예측을 그래프 신경망(GNN) 으로 개선하는 작업임. 각 노드의 opcode에 대한 임베딩 사전과 shape, bits, window size 같은 다른 노드 특징을 함께 사용함([1])
    다양한 XLA 컴파일 설정과 TPU에서의 결과 성능을 담은 큰 그래프 데이터셋을 [2]에 공개했고, 그래프를 분할하는 방식(METIS graph partition은 처음 봄)과 여러 학습 기법으로 [3]에서 이전보다 큰 그래프의 예측을 개선함. 이건 주어진 그래프의 성능을 예측하는 이야기지, 동등한 새 그래프를 개선·제안·수정하는 이야기는 아님. FunSearch처럼 예측력이 괜찮은 모델은 진화 탐색과 함께 쓸 수 있음
    [1] https://github.com/google-research-datasets/tpu_graphs#featu...
    [2] TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs https://arxiv.org/abs/2308.13490
    [3] Learning Large Graph Property Prediction via Graph Segment Training https://arxiv.org/abs/2305.12322

  • 저 그래프에서 합성곱이 어떻게 동작하는지 설명해 줄 수 있음? shape가 [2,4,16]인 텐서에 shape [4,16,8]인 커널을 합성곱했는데 [2,8] 텐서가 나온다니, 어떻게 가능한가?

    • 도움이 될지 모르겠지만, 입력 텐서 [2,4,16]에서 2는 배치 크기, 4는 입력 특징 차원, 16은 입력 채널 차원으로 볼 수 있음
      커널 [4,16,8]에서 4는 필터 창의 크기, 16은 입력 채널 차원과 맞는 값, 8은 출력 채널 차원임. 출력 [2,8]에서 2는 배치 크기로 유지되고, 8은 커널의 출력 채널 차원과 맞음. 겉보기에는 차원이 안 맞아 보이지만, 그래프에서의 합성곱은 이웃 구조를 활용함. 커널이 그래프 위를 이동하며 현재 노드와 특정 반경 안의 이웃 특징에 가중치를 적용하고, 그 가중합을 모아 각 출력 채널의 새 특징을 만듦. 그래프 구조와 간선 가중치, 패딩과 보폭 같은 구현 세부사항도 출력 shape에 영향을 줄 수 있음

  • Gemini는 어떤 상태임?

    • GPT-4가 계속 우세하다는 점이 흥미로움: https://twitter.com/lmsysorg/status/1735729398672716114
      떠오르는 것만 해도 Llama, Claude, Gemini, Falcon, Mistral 같은 기반 모델이 최소 다섯 개 있고 서로 엎치락뒤치락하지만, GPT는 여전히 한 단계 위에 있고 1년째 그러고 있음. Transformer 기반 대규모 언어 모델은 충분히 단순해서 GPU 시간에 백만 달러쯤 쓸 수 있으면 누구나 만들 수 있다는 게 드러났지만, OpenAI를 완전히 따라잡지는 못함. 그들의 특별한 비결이 뭘까?

  • Transformer 자체는 어떤가? 어떤 의미에서 최적이라는 단서가 있나?

  • 첫 문단에서 핵심을 묻어버린 느낌이지만, 나머지는 멋진 내용임

  • 지금 ML 발전 속도는 놀라움. 특이점은 믿지 않지만, 소프트웨어와 사회를 아무도 예측 못 하는 방식으로 바꾸고 있음

    • 이것과 FunSearch를 보면 특이점이 임박한 것처럼 보임
      https://deepmind.google/discover/blog/funsearch-making-new-d...
    • 내게는 닷컴, 모바일, 클라우드, VR 다음에 온 또 하나의 골드러시로 보임
    • 5년 뒤에는 사람들이 지금처럼 프로그래밍하지 않을 것 같음
    • 먼저 치료하기 어려운 질병의 치료법을 내놓는 걸 보고 싶음. 특이점 자체는 인간에게 이득이 없으면 의미가 없고, 그 이득은 주로 건강 개선과 고통 감소에 있어야 함
답변달기