Neuron Kernel Interface(NKI), NxD 훈련 및 훈련용 JAX를 새롭게 지원하는 AWS Neuron
AWS Trainium 및 Inferentia용 프로그래밍 인터페이스인 Neuron Kernel Interface(NKI)(베타)가 추가된 Neuron 2.20의 출시가 오늘 발표되었습니다. 개발자는 해당 인터페이스를 사용하여 새 기능, 최적화, 과학 혁신을 위해 최적화된 컴퓨팅 커널을 구축할 수 있습니다. 또한 이번 릴리스에는 nXD Training(베타)도 추가되었습니다. 효율적인 분산 훈련을 지원하는 PyTorch 기반 라이브러리인 nXD Training에는 NeMo와 호환되는 사용자에게 친숙한 인터페이스가 포함되어 있습니다. 이번 릴리스부터는 JAX 프레임워크(베타)도 지원됩니다.
AWS Neuron은 생성형 AI용으로 특별히 구축된 AWS Inferentia 및 Trainium 기반 인스턴스용 SDK입니다. Neuron은 PyTorch 같은 인기 있는 ML 프레임워크와 통합됩니다. 여기에는 Trn1 및 Inf2 인스턴스에서 AI 모델의 고성능 훈련 및 추론을 지원하는 컴파일러, 런타임, 도구, 라이브러리가 포함되어 있습니다.
이번 릴리스에는 모델 훈련 및 추론용 지원 기능 및 성능 개선 사항이 추가되었습니다. 훈련의 경우에는 시퀀스 길이가 최대 32K인 Llama 3.1 8B 및 70B 모델이 추가로 지원되며, PyTorch 혼합 정밀도 기본 지원을 위한 torch.autocast()와 PEFT LoRA 기법도 지원됩니다. 추론의 경우에는 Neuron 2.20은 Llama 3.1(405b, 70b, 8b)을 추가로 지원하며 Pixart-alpha 및 Pixart-sigma와 같은 Diffusion-Transformer(DiT) 모델도 지원합니다. 또한 이번 릴리스에서는 디바이스의 상위 p개 샘플 추출을 통한 추론, 그리고 Flash Decoding을 통한 128K 컨텍스트 길이도 추가로 지원됩니다. 그와 동시에 Neuron Compiler의 RMSNorm/RMSNormDx 연산자 및 Rocky 9.0 운영 체제도 추가 지원됩니다.
자세한 내용은 Neuron 출시 정보를 참조하세요.