Shardy는 모든 방언을 위한 MLIR 기반 텐서 파티션 시스템입니다. GSPMD 및 PartIR팀의 공동작업으로 구축된 이 시스템은 두 시스템의 장점과 팀 및 사용자의 공유 경험을 통합합니다.
이점
- GSPMD의 전파와 PartIR의 증분 분할을 결합하여 사용자의 제어 및 예측 가능성을 개선했습니다.
- 공유된 경험에 기반한 새로운 기능(예: 사용자가 해결 방법을 모르면 추가 커뮤니케이션이 발생하는 것으로 악명 높은 모양 변경을 위한 새로운 지원)
- 축 기반 샤딩 표현을 사용하여 최종 사용자 속도를 높이는 등 사용성 및 디버그 가능성 개선
- MLIR를 사용하는 간단한 오픈소스 코드베이스로, 사용자를 지원하는 더 광범위한 활성 참여자 (내부, 외부, 여러 시간대)가 있습니다.
구성요소
- 샤딩 표현: 특정 논리 메시 (가능한 경우 여러 메시 중 하나)에 바인딩되고 측정기준 샤딩 및 축 제약, 모양 변경과 같은 작업을 위한 축 분할, 증분 파티셔닝 우선순위 등을 지원하는 축 기반 샤딩 표현입니다.
- 컴파일러 API: 샤딩 표현과 함께 사용하여 샤딩 전파에 영향을 줄 수 있는 컴파일러 구성요소 집합입니다.
- 입력/출력 샤딩 - 샤딩을 기본 함수의 입력 또는 출력에 연결하여 함수에 제공될 때/함수에서 반환될 때 입력/출력 텐서를 샤딩해야 하는 방식을 나타냅니다.
- 샤딩 제약 조건 - 중간 텐서 (예: matmul의 결과)에 샤딩을 연결하여 이 텐서 또는 사용의 하위 집합을 샤딩해야 함을 나타냅니다.
- Shard As/Like - 여러 텐서를 ID별로 그룹화하여 동일한 방식으로 샤딩해야 함을 나타냅니다.
- 수동 계산 - 메시 축의 하위 집합을 사용하여 수동으로 분할된 하위 계산을 묶습니다. 여기서 이러한 수동 축을 따라 샤딩이 모든 입력과 출력에 지정되고 하위 계산 내에서 텐서 유형은 이러한 샤딩과 관련하여 로컬입니다.
- 샤딩 전파: 사용자 우선순위 및 샤딩 제약 조건을 컴파일러 비용 모델 및 휴리스틱과 결합하는 전파 알고리즘입니다.
- 사용자 정의 우선순위(예: 일괄 병렬 처리 후 ZeRO 실행)
- 연산 기반 우선순위(예: 요소별 연산이 먼저, matmul 등이 나중에 실행됨)
- 더 세분화된 휴리스틱(예: 배치 측정기준 선호)
- SPMD 파티셔너: 프로그램을 SPMD 프로그램으로 분할하고 프로세스에 필요한 데이터 이동/형식 지정 및 집합 연산을 추가하여 샤딩 전파 결정을 낮추는 구성요소입니다.
- 단기적으로는 초기 구현에서 현재 GSPMD SPMD 파티셔너를 사용합니다.
- 장기적으로는 새로운 MLIR 기반 SPMD 파티셔너를 만들 계획입니다.
코드 저장소
Shardy 프로젝트는 현재 개발 중이며 오픈소스 커뮤니티의 의견을 기다리고 있습니다. Shardy 코드는 https://github.com/openxla/shardy에서 확인할 수 있습니다.