Keras에서 Llama 3.2

Llama 3.2는 Keras에서 즉시 사용할 수 있으며, 필요에 따라 실시간 변환을 통해 표준 Hugging Face 체크포인트(​safetensors​ 포함)를 로드하는 것을 지원합니다. 이 통합을 통해 개발자는 Keras 생태계 내에서 Llama 3.2를 활용할 수 있으며, 다중 백엔드 유연성과 통합 훈련 도구의 이점을 누릴 수 있습니다.

JAX, PyTorch, TensorFlow를 위한 다중 백엔드 지원

Keras는 다중 백엔드 모델링 라이브러리로, 동일한 모델을 JAX, PyTorch 또는 TensorFlow에서 실행할 수 있게 합니다. 백엔드는 Keras를 임포트하기 전에 환경 변수로 지정합니다:

import os
os.environ["KERAS_BACKEND"] = "jax" # 옵션: "jax", "torch", 또는 "tensorflow"

이 유연성을 통해 사용자는 최적화된 성능을 위해 XLA 컴파일이 적용된 JAX를 활용할 수 있습니다.

Keras‑Hub와 모델 통합

keras-hub(이전 명칭: KerasNLP 및 KerasCV)은 Keras용 사전 학습 모델 모음입니다. 여기에는 Llama 3, Gemma, StableDiffusion, Segment Anything 등 인기 모델에 대한 표준 Keras 구현이 포함됩니다.

Llama 3.2는 keras_hub의 Llama3CausalLM 클래스를 사용해 로드할 수 있습니다:

from keras_hub import models.Llama3CausalLM
model = Llama3CausalLM.from_preset("hf://meta-llama/Llama-3.2-1B-Instruct", dtype="bfloat16")

"배터리 포함" LLM 기능

Keras LLM은 토크나이저를 모델 객체에 직접 통합하여 사용성을 높였습니다. 이를 통해 원시 문자열에 대한 고수준 연산이 가능합니다:

  • 생성: model.generate("Hi there!")는 문자열 입력으로부터 직접 텍스트 출력을 생성합니다.
  • 훈련: model.fit(strings)는 문자열 리스트 또는 데이터셋에 바로 훈련을 수행합니다.

채팅 및 지시 튜닝

Llama‑3.2‑1B‑Instruct와 같은 지시 튜닝된 변형은 특정 태그 형식을 사용해 턴‑바이‑턴 대화를 지원합니다. Llama 3.2에 필요한 형식에는 <|start_header_id|>system<|end_header_id|>, <|start_header_id|>user<|end_header_id|> 및 <|eot_id|>와 같은 태그가 포함됩니다. 형식이 맞춰진 문자열은 바로 model.generate()에 전달할 수 있습니다.

저수준 모델 접근

보다 세밀한 제어가 필요한 사용자를 위해 Keras는 기본 구성 요소에 대한 접근을 제공합니다:

  • 토크나이저: model.preprocessor.tokenizer를 통해 접근할 수 있습니다. 이는 텍스트를 정수 벡터로 변환합니다.
  • 백본: 핵심 모델 아키텍처는 model.backbone을 통해 접근할 수 있습니다.

전처리기(Preprocessor) 개념

Keras의 Preprocessor는 데이터 변환을 위한 포괄적인 도구입니다. CausalLM 작업에 대해 전처리기는 다음을 수행합니다:

  1. 시작 및 종료 텍스트 토큰 추가.
  2. 토큰 시퀀스 패딩 및 마스크 생성.
  3. 훈련을 위한 "예상 출력" 생성(입력 문자열을 한 토큰씩 오른쪽으로 이동).

훈련 및 Hub 통합

Keras는 model.fit(ds)를 통해 사용할 수 있는 내장 트레이너를 제공합니다. 이 트레이너는 분산 훈련, 혼합 정밀도, 양자화, LoRA 및 QLoRA와 같은 파라미터 효율적 미세조정 방법과 같은 Keras 기능과 호환됩니다.

미세조정된 모델은 model.save_to_preset()으로 로컬에 저장한 뒤 keras_hub.upload_preset()을 사용해 Hugging Face Hub에 직접 업로드할 수 있습니다.

분산 모델 병렬화

Keras는 JAX와 XLA 컴파일러를 통한 고급 모델 병렬화 경로를 간소화합니다. 이는 단일 가속기(예: Llama 3.1 8B)로는 다루기 힘든 대형 모델에 특히 유용합니다.

사용자는 DeviceMesh와 LayoutMap을 정의해 모델을 여러 GPU 또는 TPU에 샤딩할 수 있습니다. 대부분의 모델은 get_layout_map(device_mesh)를 통해 합리적인 기본값을 제공하지만, 사용자는 성능 최적화를 위해 맞춤형 레이아웃 맵을 정의할 수도 있습니다. 예를 들어, TPU v5e에서 맞춤형 레이아웃 맵을 사용하면 에포크 시간이 62초에서 54초로 단축됩니다.

Sources