meta-pytorch/attention-gym
Helpful tools and examples for working with flex-attention
해결하는 문제
Attention Gym은 새로운 어텐션 메커니즘을 구현하고 실험할 수 있는 플레이그라운드와 빌딩 블록 라이브러리를 제공합니다. 사용자 정의 어텐션 마스크 및 스코어 수정을 간편하게 만들며, 일반적으로 처음부터 구현하기 어려운 스파스 및 선형 어텐션 변형에 대한 효율적인 구현을 제공합니다.
작동 방식
이 프로젝트는 세 가지 주요 워크플로우를 통해 작동합니다:
- FlexAttention 통합: PyTorch의 FlexAttention API에 직접 전달할 수 있는
mask_mod 및 score_mod 함수를 제공하여 어텐션 계산 방식을 사용자 정의할 수 있습니다.
- 전용 커널:
gather_attn을 통해 스파스 어텐션을 위한 전용 API와 커널, GDN 및 KDA 경로를 포함한 선형 어텐션을 위한 커널을 제공합니다.
- 실용적인 예제: 페이지 기반 어텐션, 링 어텐션, CUDA Graphs를 포함한 벤치마크 및 실제 워크로드 세트를 제공하여 이러한 빌딩 블록을 실제 운영 환경과 유사한 시나리오에서 어떻게 적용할 수 있는지 보여줍니다.
대상 사용자
PyTorch 2.5 이상을 사용하여 새로운 어텐션 변형을 개발하거나 어텐션 성능을 최적화하고자 하는 AI 연구자 및 개발자에게 적합합니다.
주요 특징
- PyTorch FlexAttention과의 통합으로 마스크 및 스코어 수정이 간편합니다.
- 스파스 어텐션 패턴과 효율적인 선형 어텐션 커널 지원.
- GDN 및 KDA(체크, 반복, 디코딩 경로)를 위한 전용 프리미티브.
- 변동 스파스 어텐션, 컴파일, 프로파일링을 포함한 포괄적인 예제.
관련
- 프로젝트
lucidrains/perceiver-pytorchPerceiver 및 Perceiver IO 아키텍처(및 소규모 실험적 변형)를 구현한 PyTorch 라이브러리입니다. 즉시 사용 가능한 Perceiver, PerceiverIO 및 PerceiverLM 클래스를 제공하며, 푸리에 위치 인코딩을 지원하고 pip install perceiver-pytorch를 통해 설치할 수 있습니다. 이 저장소는 대규모 멀티모달 입력을 처리하기 위한 최첨단 Transformer 모델에 중점을 둔 AI/ML 프로젝트입니다.
- 프로젝트
Shmuma/ptanPTAN (PyTorch AgentNet)은 PyTorch 기반 에이전트를 위한 재사용 가능한 강화학습 빌딩 블록(경험 버퍼, 학습 루프, Ignite 통합 등)을 제공하는 소규모 오픈소스 라이브러리로, OpenAI-Gym 환경에서 작동합니다.
- 프로젝트
ixaxaar/pytorch-dncPyTorch 라이브러리로 Differentiable Neural Computers (DNC), Sparse DNC (SDNC), Sparse Access Memory (SAM)를 구현. 설치 지침, API 사용법, 디버깅 지원, 예제 학습 작업(복사, 덧셈, argmax) 포함.
- 프로젝트
meta-pytorch/torchtunetorchtune은 PyTorch 기반의 라이브러리로, Llama, Gemma, Qwen, Phi 등 여러 현대적 LLM(Large Language Models)의 파인튜닝, RL-HF, 증류, 양자화 대응 학습, 추론을 위한 즉시 실행 가능한 레시피와 설정을 제공합니다. 분산 학습, 메모리 절약 기술, 디바이스 지원을 간단한 CLI(`tune`)와 YAML 파일 뒤에 추상화하여 Hugging Face, torch-ao, FSDP2, 로깅 도구와 통합됩니다. 2025년 개발이 중단되었지만, LLM 실험을 위한 실용적이고 오픈소스 툴킷으로 여전히 유용합니다.
- 프로젝트
webis-de/small-textscikit-learn, PyTorch, Hugging Face transformers를 지원하며, 다양한 쿼리 전략, 초기화 방법, 정지 기준을 제공하는 텍스트 분류용 즉시 사용 가능한 액티브 러닝 루프를 제공하는 파이썬 라이브러리입니다.