LayerSkip

Cache HuggingFace Models

To avoid re-downloading models on every run, mount your host's HuggingFace cache directory to the container's cache directory:

docker run -it --rm \
    -e HUGGINGFACE_TOKEN=your_token \
    -v /path/on/host/huggingface_cache:/root/.cache/huggingface \
    layerskip:latest \
    python generate.py --help

Optimize Docker Layers

• Leverage Caching: Copy requirements.txt and install dependencies before copying the rest of the source code.
• Combine RUN Commands: Reduce layer count by combining commands:

RUN conda install pytorch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 cpuonly -c pytorch -y && \
    pip install --upgrade pip && \
    pip install --no-cache-dir -r /app/requirements.txt

To observe speedups, you must use models trained with the LayerSkip recipe. Follow these steps to access them:

1. Visit the model's Hugging Face page (e.g., facebook/layerskip-llama2-7B).
2. Submit a request form on the Hugging Face website to request access.
3. Obtain a Hugging Face user access token.
4. Log in via the CLI using huggingface-cli login and provide your token.

Available checkpoints include various Llama, CodeLlama, and Llama 3 models.

To use LayerSkip, clone the repository and set up a Conda environment with Python 3.10. Install the required dependencies using pip install -r requirements.txt.

$ git clone git@github.com:facebookresearch/LayerSkip.git
$ cd LayerSkip

$ conda create --name layer_skip python=3.10
$ conda activate layer_skip

$ pip install -r requirements.txt

For production environments (Docker Swarm), use Docker secrets instead of environment variables to manage the HUGGINGFACE_TOKEN.

1. Create the secret:

echo "your_token" | docker secret create huggingface_token -

2. Update entrypoint.sh to read the secret:

export HUGGINGFACE_TOKEN=$(cat /run/secrets/huggingface_token)

3. Deploy the service:

docker service create --name layerskip_service --secret huggingface_token layerskip:latest python generate.py --help

Run LayerSkip scripts using docker run. You must pass your HuggingFace token via the HUGGINGFACE_TOKEN environment variable using the -e flag. Use -it for interactive mode and --rm to automatically remove the container upon exit.

Basic Structure:

docker run -it --rm -e HUGGINGFACE_TOKEN=your_token layerskip:latest python your_script.py [args]

docker run -it --rm \
    -e HUGGINGFACE_TOKEN=your_huggingface_token_here \
    layerskip:latest \
    python your_script.py --help

To build the Docker image for LayerSkip, ensure you have cloned the repository and that the Dockerfile, entrypoint.sh, and .dockerignore files are in the root directory. Use the docker build command to create an image tagged as layerskip:latest.

git clone git@github.com:facebookresearch/LayerSkip.git
cd LayerSkip
docker build -t layerskip:latest .

The performance of self-speculative decoding depends on the combination of exit_layer and num_speculations. Use sweep.py to run a grid search over these hyperparameters.

Example Command:

$ torchrun sweep.py --model facebook/layerskip-llama2-7B \
    --dataset human_eval \
    --generation_strategy self_speculative \
    --num_samples 150 \
    --max_steps 256 \
    --output_dir ./logs/ \
    --sample False

Results are saved as a CSV file in the specified --output_dir.

Use correctness.py to verify the correctness of the self-speculative decoding implementation. Mount a host directory to /app/correctness to save metrics.

docker run -it --rm \
    -e HUGGINGFACE_TOKEN=your_huggingface_token_here \
    -v /path/on/host/correctness:/app/correctness \
    layerskip:latest \
    python correctness.py --model facebook/layerskip-llama2-7B \
                            --dataset human_eval \
                            --generation_strategy self_speculative \
                            --num_speculations 6 \
                            --exit_layer 4 \
                            --num_samples 10 \
                            --sample False \
                            --output_dir /app/correctness

Use generate.py to run models in interactive mode.

Regular Autoregressive Decoding:

$ torchrun generate.py --model facebook/layerskip-llama2-7B --sample True --max_steps 512

Self-Speculative Decoding (for speedup): You must specify --generation_strategy self_speculative, an --exit_layer (the layer where the draft stage exits), and --num_speculations (the number of draft tokens).

$ torchrun generate.py --model facebook/layerskip-llama2-7B \
    --sample True \
    --max_steps 512 \
    --generation_strategy self_speculative \
    --exit_layer 8 \
    --num_speculations 6

Arguments:

• --model: HuggingFace model ID.
• --sample: Enable/disable sampling (default is True).
• --temperature, --top_p, --top_k: Sampling parameters.
• --max_steps: Maximum number of steps.

Use sweep.py to perform a sweep over exit_layer and num_speculations hyperparameters. Mount a host directory to /app/sweep to persist the sweep results.

docker run -it --rm \
    -e HUGGINGFACE_TOKEN=your_huggingface_token_here \
    -v /path/on/host/sweep:/app/sweep \
    layerskip:latest \
    python sweep.py --model facebook/layerskip-llama2-7B \
                     --dataset human_eval \
                     --generation_strategy self_speculative \
                     --num_samples 150 \
                     --max_steps 256 \
                     --output_dir /app/sweep \
                     --sample False