LLM 시대를 선도하는 최적의 딥러닝 라이브러리 JAX/Flax

JAX(잭스)는 대규모 계산의 확장성을 염두에 두고 설계된 고성능 라이브러리로, LLM 시대 애물단지로 전락한 파이토치를 빠르게 대체하고 있다. 모두의연구소 JAX/Flax LAB이 집필한 이 책은 JAX, 그리고 JAX와 함께 쓰이는 Flax(플랙스)를 본격적으로 다루는 국내 최초의 책이다. JAX 기초와 함수형 프로그래밍, 병렬처리 등의 특장점을 살펴보고, JAX와 Flax를 조합해서 CNN, ResNet, DCGAN, CLIP 모델을 실제로 구현해본다. 새로운 시대, 새로운 딥러닝의 방식을 익혀보자.

주요 내용

함수형 프로그래밍, 파이썬 라이브러리 등 JAX 사용 시 알아야 할 기초
JIT 컴파일, 자동 벡터화, pytree, 병렬처리 등 JAX의 주요 특징
CNN 튜토리얼로 알아보는 Flax 기초
ResNet, DCGAN, CLIP 모델을 구축하며 Flax에 익숙해지기
코랩, 캐글에서 TPU 환경 설정하기

책 속에서

JAX는 구글에서 개발한 고성능 수치 계산 라이브러리로, 특히 병렬 가속화 기능을 통해 대규모 모델의 효율적인 학습과 추론이 가능합니다. Flax는 JAX 기반의 심플한 신경망 라이브러리로, JAX의 장점을 살려 유연하고 확장 가능한 모델 구축을 지원합니다. 이 책은 모두의연구소 JAX/Flax LAB이 다양한 경험과 지식을 바탕으로 JAX를 어떻게 실용적으로 활용할 수 있는지에 중점을 두고 집필한 책입니다. 이론만을 설명하는 것이 아니라 실제 예제를 통해 적용 방법까지 소개합니다.
--- p.xiv

JAX에서는 해당 메서드를 데커레이터--- p.decorator)로 활용합니다. 데커레이터로 사용하면 코드 간결성이나 코드 재사용이 늘어난다는 장점을 갖고 있습니다. --- p.…) 이번 예제는 @partial에만 집중해서 살펴보겠습니다. 해당 데커레이터는 jax.jit라는 함수에서 고정시키고 싶은 인수인 n을 static_argnames로 고정시키고 컴파일됩니다. 이 방식을 취하면 n은 컴파일되어 추가적인 계산을 진행하지 않습니다. / 따라서 JAX에서 partial 데커레이터를 사용하면 굳이 선언할 필요 없이 병렬처리를 할 수 있게 도와줍니다.
--- p.11

먼저 SELU--- p.scaled exponential linear unit)를 구현한 예시를 봅시다. --- p.…) 이 출력 결과는 구글 코랩의 T4 가속기에서 실행한 결과입니다. 이제 XLA 컴파일러를 이용해보겠습니다. JAX는 jax.jit 변환을 통해 JAX와 호환되는 함수들을 JIT 컴파일합니다. 얼마나 빨라지는지 확인해보겠습니다. --- p.…) 7배 이상 빨라진 것을 확인할 수 있습니다. jax.jit이 함수 selu--- p.)를 selu_jit--- p.)으로 변환해준 덕입니다.
--- p.33

드롭아웃에서는 뉴런의 무작위적 삭제를 위한 무작위 연산이 필요합니다. 따라서 PRNG 상태를 제공해야 하며, 이를 위해 PRNG 키가 필요합니다. 즉 기존에 사용했던 키 외에 추가적인 키가 필요하게 됩니다. 이미 알고 있는 것처럼 JAX에는 PRNG 키를 제공하는 명시적인 방법이 있습니다. key, subkey = jax.random.split--- p.key)처럼 split 메서드를 사용하여 추가적인 키를 발급할 수 있습니다. --- p.…) TrainState의 초기화 방법에도 변화가 있습니다. 먼저 사용자 정의 TrainState 클래스에 key 필드를 추가해야 합니다. 이제 TrainState에서 PRNG 키 또한 상태로 관리할 수 있습니다. 이를 위해 TrainState.create--- p.) 메서드에 dropout_key를 전달합니다.
--- p.96

이번 절에서는 Flax와 TPU를 활용하여 허깅 페이스의 Transformers 라이브러리에 있는 DistilGPT2 모델을 미세조정 학습하는 과정을 다룰 것입니다. 이 예제는 고성능 하드웨어의 이점을 최대한 활용하면서도 최신 소프트웨어 도구의 유연성을 결합하여 효과적인 학습 파이프라인을 구축하는 방법을 보여줍니다. --- p.…) DistilGPT2는 GPT-2--- p.Generative Pre-trained Transformer 2)의 가장 작은 버전으로, 영어를 지원하는 사전 훈련된 모델입니다. 지식 증류--- p.knowledge distillation) 기법을 활용하여 기존 124 M개의 매개변수를 가진 모델을 82 M의 매개변수로 축소한 모델입니다. 이를 통해 모델의 크기를 줄이면서도 성능을 유지하는 효율적인 학습이 가능해집니다.
--- pp.166-167