DOCSEARCHbyAlgolia
Submit Docs

StreamingLLM Documentation

repository·main·Indexed 27 days ago

https://github.com/mit-han-lab/streaming-llm

A framework enabling Large Language Models to handle infinite-length input sequences using attention sinks and a sliding window of recent tokens. It manages memory by retaining KV states of initial tokens and recent tokens without requiring fine-tuning, while maintaining the model's original pre-trained context window capacity.

Tokens
505
Snippets
2
Records
3
Agent score
43%

What's inside StreamingLLM

  1. Understand StreamingLLM capabilities and limitations

    main

    Key Concepts

    • Infinite-length inputs: StreamingLLM enables LLMs to handle continuous input streams (like multi-round dialogues) by retaining only the most recent tokens and 'attention sinks' (the KV states of initial tokens), discarding intermediate tokens to manage memory.
    • Attention Sinks: By keeping the KV states of the very first tokens, the model maintains performance even when the sequence length exceeds the original training window.

    Limitations

    • No Context Expansion: StreamingLLM does not expand the model's original pre-trained context window. For example, if Llama-2 was trained with a 4096 window, the maximum cache size remains 4096. It only manages the sliding window of that capacity.
    • No Long-term Memory: The model does not have enhanced long-term memory. If you input a whole book, the model will only be able to process and summarize the most recent tokens (e.g., the concluding paragraphs).

  2. Install StreamingLLM

    main

    To set up the environment for StreamingLLM, create a new Conda environment with Python 3.8, install PyTorch, and then install the required dependencies including transformers==4.33.0. Finally, run the setup script in development mode.

    conda create -yn streaming python=3.8
conda activate streaming

pip install torch torchvision torchaudio
pip install transformers==4.33.0 accelerate datasets evaluate wandb scikit-learn scipy sentencepiece

python setup.py develop

  3. Run the Streaming Llama Chatbot

    main

    You can run the Streaming Llama Chatbot demo using the provided example script. Use the --enable_streaming flag to activate the StreamingLLM functionality.

    CUDA_VISIBLE_DEVICES=0 python examples/run_streaming_llama.py  --enable_streaming