Llama 3.2 in Keras

Llama 3.2 is available for use in Keras immediately, supporting the loading of standard Hugging Face checkpoints (including safetensors) with on-the-fly conversion if necessary. This integration allows developers to leverage Llama 3.2 within the Keras ecosystem, benefiting from multi-backend flexibility and integrated training tools.

Multi-Backend Support for JAX, PyTorch, and TensorFlow

Keras operates as a multi-backend modeling library, enabling the same model to run on JAX, PyTorch, or TensorFlow. The backend is specified via an environment variable before importing Keras:

import os
os.environ["KERAS_BACKEND"] = "jax" # Options: "jax", "torch", or "tensorflow"

This flexibility allows users to utilize JAX with XLA compilation for optimized performance.

Keras-Hub and Model Integration

keras-hub (formerly known as KerasNLP and KerasCV) is the collection of pre-trained models for Keras. It provides canonical Keras implementations of popular models, including Llama 3, Gemma, StableDiffusion, and Segment Anything.

Llama 3.2 can be loaded using the Llama3CausalLM class from keras_hub:

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

"Batteries Included" LLM Capabilities

Keras LLMs are designed for ease of use by integrating the tokenizer directly into the model object. This allows for high-level operations on raw strings:

  • Generation: model.generate("Hi there!") produces text outputs directly from string inputs.
  • Training: model.fit(strings) allows training directly on a list or dataset of input strings.

Chatting and Instruction Tuning

Instruction-tuned variants, such as Llama-3.2-1B-Instruct, support turn-by-turn conversations using specific tagging formats. The required format for Llama 3.2 includes tags such as <|start_header_id|>system<|end_header_id|>, <|start_header_id|>user<|end_header_id|>, and <|eot_id|>. Once formatted, these strings can be passed directly to model.generate().

Low-Level Model Access

For users requiring more control, Keras provides access to the underlying components:

  • Tokenizer: Accessible via model.preprocessor.tokenizer. This transforms text into integer vectors.
  • Backbone: The core model architecture is accessible via model.backbone.

The Preprocessor Concept

The Preprocessor in Keras is a comprehensive tool for data transformation. For CausalLM tasks, the preprocessor handles:

  1. Adding start and end text tokens.
  2. Padding token sequences and generating masks.
  3. Generating "expected outputs" for training (the input string shifted by one).

Training and Hub Integration

Keras includes a built-in trainer accessible via model.fit(ds). This trainer is compatible with Keras features including distributed training, mixed precision, quantization, and parameter-efficient fine-tuning methods like LoRA and QLoRA.

Fine-tuned models can be uploaded directly to the Hugging Face Hub using keras_hub.upload_preset() after saving the model locally with model.save_to_preset().

Distributed Model Parallelism

Keras provides a streamlined path to advanced model parallelism through JAX and the XLA compiler, which is particularly useful for models too large for a single accelerator (e.g., Llama 3.1 8B).

Users can shard models across multiple GPUs or TPUs by defining a DeviceMesh and a LayoutMap. While most models provide sensible defaults via get_layout_map(device_mesh), users can also define custom layout maps to optimize performance. For example, a custom layout map on a TPU v5e can reduce epoch time from 62 seconds to 54 seconds.

Sources