PyTorch 한 파일에서의 Mamba 구현

 (github.com/johnma2006)
1P by GN⁺ | ★ favorite | 댓글 1개
  • mamba-minimal은 Mamba를 PyTorch 한 파일로 단순하고 최소한으로 구현한 프로젝트임
  • 공식 구현과 forward/backward pass에서 같은 수치 출력을 내는 것을 목표로 함
  • 코드는 단순화되어 있고, 읽기 쉽도록 주석이 포함된 형태로 구성됨
  • 공식 구현의 핵심 최적화는 포함하지 않아 속도를 제공하지 않으며, 적절한 파라미터 초기화도 포함하지 않음
  • 데모는 state-spaces/mamba-370m과 EleutherAI/gpt-neox-20b 토크나이저를 사용해 프롬프트 완성 예시를 실행함

프로젝트 개요

  • mamba-minimal은 Mamba를 PyTorch 한 파일로 구현한 단순한 최소 구현체임
  • 목표는 공식 구현과 같은 동작을 더 읽기 쉬운 코드로 보여주는 것임
  • 주요 특징:
    • 공식 구현과 forward pass 및 backward pass에서 동등한 수치 출력
    • 단순화된 코드
    • 읽기 쉽고 주석이 달린 구현

포함하지 않는 것

  • 속도는 목표가 아님
    • 공식 구현은 강하게 최적화되어 있음
    • 해당 최적화는 Mamba 논문의 핵심 기여에 포함됨
    • 이 구현은 가독성을 위해 대부분의 구현을 단순하게 유지함
  • 적절한 파라미터 초기화는 포함하지 않음
    • 가독성을 희생하지 않고 추가할 수 있는 항목으로 제시됨

데모 사용 예

  • demo.ipynb에서 프롬프트 완성 예시를 볼 수 있음
  • 예시는 model.Mamba와 Hugging Face transformers의 AutoTokenizer를 사용함
  • 사용 모델과 토크나이저:
    • state-spaces/mamba-370m
    • EleutherAI/gpt-neox-20b
  • 예시 프롬프트는 Mamba is the이며, 생성 결과에는 Mamba를 독사로 설명하는 문장이 포함됨

참고 자료

함께 보면 좋은 글 β

GN⁺   [-]
Hacker News 의견들

  • 예전에 동료와 함께 공유되는 모델 코드를 대부분 분리한 라이브러리를 만들었고, 이를 쓰면 많은 모델을 Python import와 주석을 제외하고 약 100줄로 구현할 수 있음
    BERT: https://github.com/explosion/curated-transformers/blob/main/...
    Llama 1/2: https://github.com/explosion/curated-transformers/blob/main/...
    MPT: https://github.com/explosion/curated-transformers/blob/main/...
    TorchScript JIT, PyTorch flash attention 같은 기능도 지원함

    • 이 라이브러리를 꼭 살펴볼 예정임. xformers도 봤는지 궁금함
      xformers도 비슷한 문제를 다루지만, Triton을 이용해 성능 좋은 Transformer 모듈을 제공하는 데 더 초점이 있음. 다만 라이브러리의 특정 구성요소만 가져다 쓰기는 쉽지 않았고, 런타임 오류가 계속 나서 일단 미뤄둠. BERT 아키텍처 기반으로 뭔가 만들고 있어서 참고해보겠음
    • 이 라이브러리에 감탄함. Hugging Face 구현은 별로 좋아하지 않았는데, 이건 추상화 수준이 딱 맞는 아름다운 API처럼 보임
      다음 프로젝트에서 써볼 생각임

  • 원래 Mamba 코드는 속도 최적화와 기타 요소가 많아서 바로 이해하기 어려운데, 이 구현은 학습에 도움이 될 듯함
    토큰을 하나씩 추론할 때는 모든 것이 훨씬 단순해짐. 직접 만든 Mamba 추론 구현도 있음: https://github.com/rbitr/llm.f90/tree/master/ssm

    • Fortran이라니. 왜 Fortran을 썼는지 궁금함
      오래 검증된 과학 계산 코드의 기반이고 PyTorch나 Numpy 같은 라이브러리로 감싸 쓰이는 경우가 많다는 건 알지만, 요즘 인기 있는 언어는 아님. 선택한 이유가 궁금함

  • Mamba에 대해 머신러닝 연구자가 아닌 사람도 이해할 수 있게 설명해줬으면 하는 부분이 있음

    1. Transformer 너머의 상태공간 모델이 가진 전체적인 통찰이 무엇인지
    2. Mamba가 S4, H3, Monarch 같은 전작보다 더 성공적이거나 흥미로워진 점진적 혁신이 무엇인지
    3. 문맥 길이의 아제곱 미만 확장성 말고 어떤 의미가 있는지. 예를 들어 100k 토큰 이상의 문맥 길이에 관심이 없다면, 비슷한 크기의 모델과 데이터셋에서 Mamba가 학습 계산 효율이 더 좋을 가능성이 있는지 궁금함
    • 논문 저자들보다 내 지능은 한참 낮지만 그래도 이해해보려고 노력했음. 컴퓨터공학을 공부했고 기본적인 제어이론과 학부 수준의 이산시간 시스템 직관은 있지만, 이 논문을 제대로 이해하려면 상태공간 모델을 훨씬 더 공부해야 할 듯함
      Mamba의 핵심 통찰은 상태공간 모델의 오래된 문제를 푸는 데 있음. 상태공간 모델은 입력 문맥을 압축하는 데는 좋지만, 입력을 은닉 상태로 압축하는 과정에서 Transformer처럼 문맥을 효과적으로 활용하는 데 필요한 정보가 지워짐
      해결책은 논문에서 선택 메커니즘이라고 부르는 것을 만드는 것임. 이 메커니즘은 입력 의존적이라서, 입력이 바뀔 때마다 모델이 각 단계의 출력을 조정할 수 있음. 이를 위해 몇몇 상태공간 변수를 입력 불변이 아니라 입력 의존적으로 만들고, 각 시점의 입력을 상태공간 변수로 투영하도록 선형 계층 등을 붙임
      하지만 상태공간 변수를 입력 의존적으로 만들면 계산 오버헤드가 생김. 이를 현대 GPU 메모리 구조를 최대한 활용하는 하드웨어 인식 알고리즘으로 해결해서, HBM 안팎으로 데이터를 옮기는 일을 가능한 한 피함
      Tri Dao는 Flash Attention을 만든 사람이고, 이것도 Transformer에서 하드웨어를 더 효율적으로 쓰는 방법이었음. 이런 쪽이 정말 그의 전문 분야임
    • Attention은 문맥 길이에 대해 이차적으로 커지고, 게이팅이 있는 순환 신경망(LSTM, GRU 등)은 선형이며, 이 새 아키텍처들도 선형임. 초기 순환 신경망은 폭주하는 기울기를 피하려고 게이팅을 썼지만, 새 접근들은 안정성을 보장하는 동역학 시스템 이론을 써서 게이팅이 두 문제를 동시에 풀지 않고 기억에 집중할 수 있게 함
      NeurIPS 2023 직전에 나온 Mamba와 Based는 다중 질의 연관 회상(MQAR)과, 다중 헤드 Attention에서 영감을 받은 게이팅/선택의 데이터 의존성을 포함했음. 이 둘이 Hyena 및 이전 상태공간 아키텍처에 빠져 있던 핵심 요소였고, 새 모델들이 연관 회상 과제에서 Attention만큼 좋아지며, 조회가 아닌 다른 과제에서는 어쩌면 Attention보다 약간 더 나아질 가능성도 보임
      물론 Mamba의 큰 디테일은 효율적인 CUDA 구현임. 그게 없으면 Transformer가 이미 적합한 작업에서는 이 아키텍처의 의미가 줄어들 수 있음
      문맥 길이를 크게 걱정하지 않아도 새 영역은 많이 열림. DNA 서열 분석은 긴 의존성을 가진 선형 작업이고, 이미지·비디오·고차원 정보를 토큰 스트림으로 보는 방식도 생각해볼 수 있음. 예전 CRT 모니터처럼 픽셀을 훑는 식임
      AI의 초기 꿈 중 하나는 환경과 계속 상호작용하는 에이전트의 단일 학습 궤적이 지속적으로 진화하는 것이었고, 이런 무한 문맥 길이 모델이 그런 꿈을 더 쉽게 만들 수도 있음
      다만 현재로서는 이런 모델의 중요한 실제 작업에 대한 다운스트림 응용은 Attention 기반의 성숙한 응용에 비해 대체로 검증과 튜닝이 덜 됐음. 오래된 순환 신경망과의 비유가 어느 정도 도움은 되지만, 지난 5년간 사람들은 Attention과 Transformer에 지나치게 특화돼 있어서 Transformer 쪽 관성이 큼
    • 비슷한 크기의 모델과 데이터셋에서 Mamba가 더 계산 효율적으로 학습될 수 있는지는 나도 알고 싶음
      원 논문에서는 파라미터가 변환된 뒤 모델을 선형 점화식이나 전역 합성곱 두 방식으로 계산할 수 있다고 설명함. 일반적으로 전체 입력 시퀀스를 미리 볼 수 있는 학습에서는 병렬화가 쉬운 합성곱 모드를 쓰고, 입력을 한 시점씩 보는 자기회귀 추론에서는 효율적인 순환 모드로 전환함
      그래서 학습은 RetNet의 병렬 순전파 모드처럼 병렬화 가능함. 기본 추론은 가능한 한 긴 문맥을 얻기 위해 순환 모드로 수행되며, 청킹이 없어서 추론 중 RAM과 VRAM을 얼마나 먹을지는 판단하기 어려움
    • 이 영상이 찾던 것에 정확히 맞을 것 같음
      논문을 설명하면서도 큰 그림에서 어디에 들어맞는지 맥락을 많이 줌. 전개를 듣는 게 꽤 흥미로움
      https://youtu.be/ouF-H35atOY?si=y2Ckp9MCFd7ulLL3
    • 내가 알기로 Mamba는 기본적으로 긴 합성곱이라고 부를 수 있는 상태공간 모델 연구의 연장선임
      각 토큰이 모든 다른 토큰에 얼마나 주의를 기울이는지 계산하는 이차 Attention 대신, 입력과 같은 길이의 긴 합성곱 커널을 어떻게든 계산한 뒤 conv1d를 적용함
      제한적으로 이해한 바로는 FFT를 적용하고, 행렬곱을 하고, 다시 IFFT로 되돌리는 것과 약간 관련 있음. 작동한다는 건 알지만 느림. FFT를 계산하는 방법은 여러 가지가 있고, 그중 하나가 나비 행렬임. 아마 근사일 뿐이지만 충분히 좋고 현재 하드웨어에서 매우 빠르고 효율적인 듯함
      이차 복잡도는 나쁘게 들리지만 실제로는 하드웨어 제약 때문에 아제곱 미만 알고리즘이 더 느린 경우가 많음. 그래서 상태공간 모델에 대한 기대가 크더라도 Llama가 끝났다고 말하기는 쉽지 않음. Mamba가 규모를 키워도 잘 되는지도 아직 모르고, 그걸 알려면 실제로 학습에 수백만 달러를 써봐야 함. 그래도 낙관적임
      아제곱 미만 계열의 또 다른 흥미로운 모델은 RWKV임. 살펴볼 만하지만, 아마 이미 팟캐스트에서 다뤘던 것 같음
      독학했고 논문도 예전에 대충 훑은 정도라 많이 틀렸을 수 있음. 또 Attention에는 보통 KV 캐시가 있어서 성능에 큰 도움이 되는데, Mamba에서는 그걸 할 수 없다고 봄

  • “Mamba는 추정 길이가 150m가 넘는 세계에서 가장 긴 독사”라는 문장에서 웃음이 났음
    그래도 정말 훌륭하고, arXiv 논문을 참조해줘서 논문을 직접 해석하기보다 이런 글을 소비하는 나 같은 사람도 안쪽을 살짝 들여다볼 수 있어서 좋았음

    • Mamba라는 이름은 좋음. [S]elective [S]tructured [S]tate [S]pace [S]equence models라서 sSSSS가 되고, 뱀 소리 같음
    • 가장 긴 독사는 킹코브라인 줄 알았음. 간단히 Google 검색해도 그렇게 나왔음
      나중에 그 문장에 정정문을 내야 한다면 재미있을 듯함

  • 알고리즘의 핵심은 병렬 접두 스캔일 거라고 예상했음. 그게 Mamba의 요점 아닌가 싶음
    for i in range(l):
    x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
    y = einsum(x, C[:, i, :], 'b d_in n , b n -> b d_in')
    ys.append(y)

  • 멍청한 질문일 수 있지만, Hugging Face에 올라온 Mamba 모델을 학습시키는 난이도가 궁금함
    가장 큰 모델이 2.8B로 보이는데, The Pile 같은 데이터셋으로 학습하려면 GPU가 몇 개 필요하고 얼마나 오래 걸릴까?

    • 나도 알고 싶은 훌륭한 질문임. 답은 같은 크기의 Transformer보다 상당히 빠른 것 같고, 최종 결과도 거의 모든 벤치마크에서 Transformer보다 더 좋은 점수를 낼 듯함
      추론도 RAM을 절반만 쓰면서 3~5배 빠르게 수행될 것으로 보임

  • 공식 CUDA 버전을 풀어보려다 첫 시도가 실패한 뒤 결국 손대지 못했는데, 이 구현이 훨씬 좋아 보임

  • 또 하나의 단일 파일 PyTorch 구현이라니 정말 훌륭함. 이전에 했던 hlb-CIFAR10 및 관련 프로젝트, 그리고 minGPT나 DawnBench 같은 선행 영향들이 단순한 단일 파일 형식을 조금이라도 밀어붙이는 데 도움이 되었기를 바람
    이런 작업은 효율적인 머신러닝 연구에 중요하며, 지금 이 분야를 위해 할 수 있는 가장 중요한 일 중 하나일 수 있음
    연구는 혁신의 속도로 전진하고, 혁신은 실험 실행 시간의 역수에 따라 빨라지며, 이는 연구나 간단한 해킹 목적에 대한 코드의 Kolmogorov 복잡도와 분명히 관련 있음
    이런 도구가 연구에 얼마나 중요한지, 개인적으로 지식 발견 과정을 얼마나 빠르게 해줬는지는 아무리 강조해도 부족함. 아이디어를 몇 분 만에 빠르게 스케치하고 즉각적으로 신호대잡음비가 높은 결과를 받는 능력은 연구 진행에 필수 요소가 됐음
    지식 증류와 MDL(https://en.wikipedia.org/wiki/Minimum_description_length)은 현재 논문 제출·리뷰 과정이 부추기는 듯한 불필요한 장식, 잡동사니, 과도하게 조밀한 “뒤처지지 않으려는” 저가치 주제 경쟁을 되돌리는 데 매우 중요하다고 봄
    최근에는 이 문제를 피하고 조금 더 나은 확장 해법으로 가고 싶어서, 코드를 1파일짜리 자체 포함 짧은 gist인 “코드 스케치”로 배포하기 시작했음. 개발 시간을 줄이고, 개념을 담은 거칠고 다듬지 않은 동작 코드를 사람들에게 바로 전달할 수 있음. 지금까지는 꽤 잘 작동하는 듯하고 계속하고 싶음
    이런 코드를 더 많이 보고 싶음. 대규모로 데이터를 학습하는 연구자들이라면, 정보 전파 방식에서도 데이터 효율적이어야 함

    • 2023년은 AI 연구가 말도 안 되는 속도로 전개되는 걸 보는 것만으로도 흥미로운 해였음. ArXiV, PyTorch, GitHub, Hugging Face, 간결한 오픈소스 Python 코드 같은 기반 요소들이 이 새 분야의 발전을 극적으로 가속하고 있음
      인류가 상당한 복잡도를 가진 무언가를 이렇게 빠르게 발전시킨 적은 아마 없을 것 같음
      비슷한 속도가 보이는 곳은 SpaceX 정도인데, 올해도 최첨단 로켓 두 대를 발사했음. 2024년에는 무엇이 나올지 궁금함
    • 작은 성능 개선 가능성이 있음. 여기서는 x_proj에 바이어스가 없으니 x_proj와 dt_proj 가중치를 합칠 수 있을 것처럼 보임
      가중치 조정 요구사항이 있다면 런타임에 단순히 할 수 있을 수도 있고, 단일 커널과 바이어스가 결국 더 빠를 것 같음. 확실하진 않음

  • 원 논문 토론이 있었는지 궁금함. 놓친 것 같은데 꽤 흥미로움
    “효율적인 구현 부족으로 메모리 부족이나 비현실적인 계산 요구가 발생해, SSM으로도 해석할 수 있는 이전의 강력한 순환 모델인 RWKV와 RetNet 기준선의 문맥 길이 8k 전체 결과가 빠져 있다”는 부분이 잘 이해되지 않았음
    RetNet은 메모리를 많이 쓰지 않고, 청크 단위 순전파 구현을 쓰면 VRAM 사용량이 청크 크기로 제한됨. 이 부분이 문맥 길이를 테스트하는 핵심임
    원래 Mamba 모델을 테스트해본 사람이 있는지 궁금함. 병렬 순전파 모드의 RetNet과 비교해 학습 속도가 어느 정도일까?

  • 복잡한 것을 핵심만 남기고 정리한 구현은 늘 좋음

답변달기