[논문 리뷰] Sockeye 3: Fast Neural Machine Translation with PyTorch

Felix Hieber, Michael Denkowski외 12명|arXiv (Cornell University)|2022. 07. 12.

Natural Language Processing Techniques인용 수 11

한 줄 요약

Sockeye 3는 최적화된 혼합 정밀도 학습, 분산 데이터 병렬 처리 및 모델 양자화를 통해 학습 및 추론을 가속화하는 PyTorch 기반 신경 기계 번역 툴킷입니다. 다른 PyTorch 구현 대비 GPU에서 최대 126% 빠른 추론 속도와 CPU에서 최대 292% 빠른 추론 속도를 달성하며, 생산 및 연구 용도로 사용 가능한 고급 기능인 소스/타겟 접두사, 어휘 단서 목록, 공손성/어휘 빈도 제어를 지원합니다.

ABSTRACT

Sockeye 3 is the latest version of the Sockeye toolkit for Neural Machine Translation (NMT). Now based on PyTorch, Sockeye 3 provides faster model implementations and more advanced features with a further streamlined codebase. This enables broader experimentation with faster iteration, efficient training of stronger and faster models, and the flexibility to move new ideas quickly from research to production. When running comparable models, Sockeye 3 is up to 126% faster than other PyTorch implementations on GPUs and up to 292% faster on CPUs. Sockeye 3 is open source software released under the Apache 2.0 license.

연구 동기 및 목표

• PyTorch 기반으로 더 빠르고 확장성 있고, 생산 환경에 적합한 NMT 툴킷을 개발하기 위해.
• 코드베이스를 단순화하고 학습 및 추론 파이프라인을 최적화하여 연구자가 신속하게 반복할 수 있도록 하기 위해.
• 소스 및 타겟 접두사, 어휘 단서 목록, 하이브리드 디코더와 같은 고급 NMT 기능을 지원하여 세밀한 제어를 가능하게 하기 위해.
• 효율적이고, 양자화된, 분산된 학습을 통해 연구 혁신과 생산 배포 사이의 격차를 메우기 위해.
• 번역에서 공손성 및 어휘 빈도 제어를 실험할 수 있는 유연하고 오픈소스 프레임워크를 제공하기 위해.

제안 방법

• Sockeye 3는 인코더, 디코더 및 빔 서치의 정적 계산 그래프를 추적하고 최적화하기 위해 PyTorch의 JIT 컴파일러를 사용하여 동적 제어 흐름 오버헤드를 줄입니다.
• PyTorch의 DDP와 AMP를 사용하여 혼합 정밀도 학습을 분산 구현하여 FP16/FP32 혼합 정밀도 및 전체 FP16(O2) 모드를 지원함으로써 속도와 메모리 효율성을 향상시킵니다.
• 데이터 준비 과정이 병렬화되고 이진 형식으로 분할되어 고정된 메모리 사용량으로 임의로 큰 데이터셋에서 학습이 가능합니다.
• CPU 추론을 위해 PyTorch의 동적 양자화(INT8)를, GPU 추론을 위해 FP16 캐스팅을 적용하여 정확도 손실를 최소화하면서 추론 속도를 가속화합니다.
• 소스 접두사를 사용하여 모델 출력을 공손성 또는 어휘 빈도 수준에 따라 조건화함으로써 재학습 없이도 제어 가능한 생성을 가능하게 합니다.
• 비용이 덜 드는 빔 서치의 대안으로 근사 탐색을 구현하여 계산 오버헤드를 줄이면서도 번역 품질을 유지합니다.

실험 결과

연구 질문

• RQ1 PyTorch 기반 NMT 툴킷은 기존 구현 대비 어떻게 상당히 빠른 학습 및 추론 처리량을 달성할 수 있는가?
• RQ2 혼합 정밀도 및 분산 학습은 모델 품질을 손상시키지 않으면서 NMT의 확장성과 효율성을 얼마나 향상시킬 수 있는가?
• RQ3 소스 접두사 조건화는 신경 기계 번역에서 출력의 공손성과 어휘 빈도를 효과적으로 제어할 수 있는가?
• RQ4 양자화는 번역 품질을 유지하면서 추론 속도와 메모리 점유율에 어떤 영향을 미치는가?
• RQ5 최적화된 추론 컴포onent인 근사 탐색 및 모델 추적을 통해 어떤 성능 향상이 달성될 수 있는가?

주요 결과

• Sockeye 3는 다른 PyTorch NMT 구현 대비 GPU에서 최대 126% 빠른 추론 속도, CPU에서 최대 292% 빠른 추론 속도를 달성합니다.
• 8개 GPU에서의 혼합 정밀도 학습은 7.7배의 처리량 향상을 이끌어내며, GPU 효율성이 96.6%로 강력한 스케일링 성능을 보입니다.
• 비교적 빔 크기가 1인 빔 서치 대비 근사 탐색은 GPU에서 16% 빠른 번역 속도, CPU에서 6% 빠른 번역 속도를 기록합니다.
• CPU에서의 INT8 양자화는 추론 속도를 초당 2.6문장에서 4.5문장으로 향상시키며, BLEU 점수 저하가 최소한입니다.
• 공손성 제어 모델은 공손 번역에서 97.8%의 정확도를 달성하고, WMT 테스트 세트에서 45.0 BLEU 점수를 기록하여 높은 제어 정밀도를 입증합니다.
• 길이 토큰과 N-best 재랭킹을 활용한 어휘 빈도 제어는 영어-독어 번역에서 최종 시스템 순위 지표를 최대 17.4점(BERTScore × LC) 향상시킵니다.