AviSoori1x/makeMoE

Andrej Karpathy의 makemore (https://github.com/karpathy/makemore)에서 영감을 얻고 (그리고 상당 부분 기반을 둔) 밑바닥부터 구현한 희소 전문가 혼합 (Sparse Mixture of Experts, MoE) 언어 모델 :)

이 과정을 설명하는 HuggingFace Community Blog: https://huggingface.co/blog/AviSoori1x/makemoe-from-scratch

전문가 용량 (expert capacity)을 상세히 다룬 Part #2: https://huggingface.co/blog/AviSoori1x/makemoe2

이 프로젝트는 밑바닥부터 구현한 희소 전문가 혼합 (Sparse Mixture of Experts, MoE) 언어 모델의 구현체입니다. 이는 Andrej Karpathy의 'makemore' 프로젝트에서 영감을 얻었으며 상당 부분 그에 기반을 두고 있으며, 해당 구현체로부터 재사용 가능한 컴포넌트들을 가져왔습니다. makemore과 마찬가지로, makeMoE 또한 자기회귀적 (autoregressive) 문자 단위 (character-level) 언어 모델이지만, 앞서 언급한 희소 전문가 혼합 (Sparse Mixture of Experts) 아키텍처를 사용합니다.

makemore과 마찬가지로, pytorch가 유일한 요구 사항입니다 (따라서 '밑바닥부터 (from scratch)'라는 주장이 정당하기를 바랍니다).

makemore 아키텍처로부터의 주요 변경 사항

• 단일 피드 포워드 신경망 (feed forward neural net) 대신 희소 전문가 혼합 (Sparse Mixture of Experts) 사용.
• Top-k 게이팅 (Top-k gating) 및 노이지 Top-k 게이팅 (noisy top-k gating) 구현.
• 초기화 (initialization) - 여기서는 Kaiming He 초기화를 사용하지만, 이 노트북의 목적은 수정 가능하도록(hackable) 만드는 것이므로 Xavier Glorot 등을 교체하여 직접 테스트해 볼 수 있습니다.
• 전문가 용량 (Expert Capacity) -- 가장 최근 업데이트 (2024/03/18)

makemore으로부터 변경되지 않은 사항

• 데이터셋, 전처리 (tokenization), 그리고 Andrej가 처음에 선택한 언어 모델링 작업 - 셰익스피어 스타일의 텍스트 생성
• 인과적 셀프 어텐션 (Causal self attention) 구현
• 훈련 루프 (Training loop)
• 추론 로직 (Inference logic)

이 구현을 위해 집중적으로 참조된 논문들:

• Outrageously Large Neural Networks: The Sparsely-Gated Mixture-Of-Experts layer: https://arxiv.org/pdf/1701.06538.pdf
• Mixtral of experts: https://arxiv.org/pdf/2401.04088.pdf

makeMoE.py는 pytorch를 사용한 전체 구현이 단일 파일에 담겨 있습니다.

makMoE_from_Scratch.ipynb는 전체 모델 아키텍처에 대한 직관과 모든 것이 어떻게 결합되는지를 단계별로 설명합니다. 여기서부터 시작하는 것을 권장합니다.

makeMoE_from_Scratch_with_Expert_Capacity.ipynb는 위의 단계별 설명을 바탕으로 구축되었으며, 더 효율적인 학습을 위해 전문가 용량 (Expert Capacity) 기능을 추가했습니다.

makeMoE_Concise.ipynb는 통합된 해킹 가능한 (hackable) 구현체로, 여러분이 직접 수정하고, 이해하고, 개선하여 자신만의 것으로 만들 것을 권장합니다.

이 코드는 컴퓨팅을 위해 단일 A100을 사용하여 Databricks에서 전적으로 개발되었습니다. 만약 Databricks에서 이 코드를 실행한다면, 원하는 클라우드 제공업체에서 임의의 크기를 가진 대규모 GPU 클러스터로 문제없이 확장할 수 있습니다.

저는 필요한 모든 지표 (Metrics)를 추적하고 기록하는 데 도움이 된다고 판단하여 MLFlow (Databricks에 사전 설치되어 있으며, 완전한 오픈 소스이므로 다른 곳에서도 쉽게 pip install 할 수 있습니다)를 사용하기로 선택했습니다. 이는 전적으로 선택 사항이지만 권장됩니다.

이 구현은 성능보다는 가독성과 해킹 가능성 (Hackability)에 중점을 두고 있으므로, 이를 개선할 수 있는 방법이 많이 있다는 점을 유의해 주세요. 시도해 보시고 저에게 알려주세요!

이 내용이 유용하기를 바랍니다. 즐거운 해킹 되세요!!