Text-Generation Pipeline on Intel Gaudi 2 AI Accelerator

Hugging Face has released a custom text-generation pipeline designed for the Intel Gaudi 2 AI accelerator, allowing developers to run Llama 2 models (7b, 13b, and 70b) with minimal code. This pipeline provides a high level of abstraction by handling end-to-end text generation, including both pre-processing and post-processing.

High-Level Pipeline Capabilities

The GaudiTextGenerationPipeline offers a flexible interface for generating text from single or multiple prompts. It is designed to be integrated into various workflows through three primary methods:

  • Standalone Scripts: Using the run_pipeline.py script provided in the Optimum Habana repository.
  • Custom Python Integration: Importing the pipeline class directly into Python scripts for programmatic control.
  • Framework Integration: Initializing LangChain classes using the pipeline for complex LLM orchestration.

Technical Implementation and Requirements

Hardware and Software Prerequisites

To utilize the pipeline, users must have access to gated Llama 2 models via Meta and Hugging Face, and the following software environment:

  • Optimum Habana: Version 1.10.4.
  • DeepSpeed: Required for distributed inference (e.g., for Llama-2-70b). The version should match the SynapseAI version (e.g., DeepSpeed 1.14.0 for SynapseAI 1.14.0).
  • Authentication: A Hugging Face access token is required to download gated models via huggingface-cli login.

Execution and Configuration

The pipeline supports several performance-optimizing flags and generation parameters:

  • Performance Optimizations: Users can enable --use_hpu_graphs and --use_kv_cache to improve inference efficiency.
  • Generation Control: The pipeline supports standard sampling parameters including --do_sample, --temperature, and --top_p.
  • Distributed Inference: For large models like Llama-2-70b, the pipeline is launched using gaudi_spawn.py with the --use_deepspeed flag and a specified --world_size (e.g., 8).

LangChain Compatibility

The text-generation pipeline is compatible with LangChain, allowing it to be used as an LLM backend. By setting the use_with_langchain=True constructor argument in the GaudiTextGenerationPipeline, the pipeline can be passed into the HuggingFacePipeline class in LangChain.

The pipeline class has been validated specifically for LangChain version 0.0.191 and may not work with other versions.

Summary of Model Support

Model Deployment Method Key Requirements
Llama-2-7b Single HPU / run_pipeline.py Optimum Habana 1.10.4
Llama-2-13b Single HPU / run_pipeline.py Optimum Habana 1.10.4
Llama-2-70b Distributed / gaudi_spawn.py DeepSpeed + Optimum Habana 1.10.4

Sources