Skip to content

script.md

Latest commit

 

History

History
565 lines (401 loc) · 10.8 KB

script.md

File metadata and controls

565 lines (401 loc) · 10.8 KB
Raw

Qwen2.5 Attention Sink Benchmark Scripts - Shreya Pagaria

To be led by Suraj Botcha & Achuth

This document explains how to reproduce the Qwen2.5 long-context benchmark runs using the provided shell scripts and plotting script.

The benchmark compares three inference modes:

  • transformers: full KV-cache baseline
  • windowed: sliding-window KV-cache baseline
  • attention_sinks: attention-sink cache with preserved sink tokens + recent window

The main metrics saved in each CSV are:

  • nll: negative log-likelihood for each decoded token
  • running_perplexity: cumulative perplexity over decoded tokens
  • latency_sec: per-token latency
  • cache_len: current KV-cache length
  • current_gpu_mem_mb: current GPU memory usage, if logged
  • peak_gpu_mem_mb: peak GPU memory usage

1. Setup

Run everything from the project root:

cd /home/spagaria/qwen2_AS

Activate the GPU environment:

conda activate qwen2_gpu

Check that CUDA is available:

python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.device_count()); print(torch.version.cuda)"

Expected output should include:

True

Check available GPUs:

nvidia-smi

For live monitoring during a run:

watch -n 1 nvidia-smi

2. Script overview

Single-run scripts

These scripts run all three modes for one model/configuration:

Script Model Purpose
scripts/run_qwen0.5B.sh Qwen/Qwen2.5-0.5B-Instruct Runs transformers, windowed, and attention_sinks for one configuration
scripts/run_qwen3B.sh Qwen/Qwen2.5-3B-Instruct Runs transformers, windowed, and attention_sinks for one configuration
scripts/run_qwen7B.sh Qwen/Qwen2.5-7B-Instruct Runs transformers, windowed, and attention_sinks for one configuration

Sweep scripts

These scripts call the single-run scripts repeatedly for ablations:

Script Model Ablation
scripts/sweep_0.5B.sh 0.5B Window-size sweep
scripts/sweep_0.5B_sink.sh 0.5B Sink-size sweep at fixed window size
scripts/sweep_3B.sh 3B Window-size sweep
scripts/sweep_3B_sink.sh 3B Sink-size sweep at fixed window size
scripts/sweep_7B.sh 7B Window-size sweep
scripts/sweep_7B_sink.sh 7B Sink-size sweep at fixed window size

Plotting script

Script Purpose
scripts/plot.py Plots modes, sink sweeps, window sweeps, or all available CSVs

3. Single configuration runs

0.5B model

Usage:

bash scripts/run_qwen0.5B.sh [MODEL] [TOKENS] [OUTDIR] [SINK_SIZE] [WINDOW_SIZE]

Example:

bash scripts/run_qwen0.5B.sh \
  Qwen/Qwen2.5-0.5B-Instruct \
  4096 \
  outputs/run_0.5B_sink4_window1020 \
  4 \
  1020

3B model

Usage:

bash scripts/run_qwen3B.sh [GPU_ID] [SINK_SIZE] [WINDOW_SIZE] [NUM_TOKENS] [OUTPUT_DIR]

Example:

bash scripts/run_qwen3B.sh \
  1 \
  4 \
  1020 \
  4096 \
  outputs/run_3B_sink4_window1020

7B model

Usage:

bash scripts/run_qwen7B.sh [GPU_ID] [SINK_SIZE] [WINDOW_SIZE] [NUM_TOKENS] [OUTPUT_DIR]

Example:

bash scripts/run_qwen7B.sh \
  0 \
  4 \
  1020 \
  4096 \
  outputs/run_7B_sink4_window1020

4. Running jobs in the background

Use nohup for long runs so the job continues after closing VSCode or disconnecting SSH.

Example:

nohup bash scripts/run_qwen7B.sh \
  0 \
  4 \
  1020 \
  4096 \
  outputs/run_7B_sink4_window1020 \
  > logs/run_7B_sink4_window1020.log 2>&1 &

Check the PID:

echo $!

Watch the log:

tail -f logs/run_7B_sink4_window1020.log

Check GPU usage:

watch -n 1 nvidia-smi

Stop a run if needed:

ps -ef | grep benchmark.perplexity
pkill -f benchmark.perplexity

5. Window-size ablation

Window ablation fixes sink size and sweeps over window sizes:

WINDOWS = 256, 512, 1020, 2048, 3072

This is used to test how increasing retained recent context affects perplexity, latency, memory, and cache length.

0.5B window sweep

nohup bash scripts/sweep_0.5B.sh \
  Qwen/Qwen2.5-0.5B-Instruct \
  4096 \
  outputs/sink_size_4_0.5B_window_all \
  4 \
  > logs/sweep_0.5B_sink4.log 2>&1 &

3B window sweep

nohup bash scripts/sweep_3B.sh \
  1 \
  4096 \
  outputs/sink_size_4_3B_window_all \
  4 \
  > logs/sweep_3B_sink4.log 2>&1 &

7B window sweep

nohup bash scripts/sweep_7B.sh \
  0 \
  4096 \
  outputs/sink_size_4_7B_window_all \
  4 \
  > logs/sweep_7B_sink4.log 2>&1 &

Argument order for window sweeps:

sweep_3B.sh / sweep_7B.sh:
  arg1 = GPU_ID
  arg2 = TOKENS
  arg3 = OUTPUT_DIR
  arg4 = SINK_SIZE

Important: the last argument is the sink size, not the window size.

6. Sink-size ablation

Sink ablation fixes the window size and sweeps over sink sizes:

SINK_SIZES = 4, 8, 20
WINDOW_SIZE = 1020

This is used to test whether more sink tokens improve long-context quality.

3B sink sweep

nohup bash scripts/sweep_3B_sink.sh \
  1 \
  4096 \
  outputs/window_1020_sinksize_all_3B \
  1020 \
  > logs/sweep_3B_sink_window1020.log 2>&1 &

7B sink sweep

nohup bash scripts/sweep_7B_sink.sh \
  0 \
  4096 \
  outputs/window_1020_sinksize_all_7B \
  1020 \
  > logs/sweep_7B_sink_window1020.log 2>&1 &

Argument order for sink sweeps:

sweep_3B_sink.sh / sweep_7B_sink.sh:
  arg1 = GPU_ID
  arg2 = TOKENS
  arg3 = OUTPUT_DIR
  arg4 = WINDOW_SIZE

Important: the last argument is the fixed window size, not the sink size.

7. Plotting results

The plotting script is:

python scripts/plot.py [OPTIONS]

Common metrics:

running_perplexity
nll
latency_sec
cache_len
current_gpu_mem_mb
peak_gpu_mem_mb

A. Same sink size + same window size + all three modes

Example for 3B, sink size 4, window size 1020:

python scripts/plot.py \
  --outputs_dir outputs \
  --model_tag Qwen_Qwen2.5-3B-Instruct \
  --compare modes \
  --sink_size 4 \
  --window_size 1020 \
  --metric running_perplexity \
  --output figures/3B_modes_sink4_window1020_ppl.png

B. Different window sizes at the same sink size

Only attention_sinks:

python scripts/plot.py \
  --outputs_dir outputs/sink_size_4_3B_window_all \
  --model_tag Qwen_Qwen2.5-3B-Instruct \
  --compare windows \
  --sink_size 4 \
  --mode attention_sinks \
  --metric running_perplexity \
  --output figures/3B_attention_sinks_window_sweep_ppl.png

All modes:

python scripts/plot.py \
  --outputs_dir outputs/sink_size_4_3B_window_all \
  --model_tag Qwen_Qwen2.5-3B-Instruct \
  --compare windows \
  --sink_size 4 \
  --metric running_perplexity \
  --output figures/3B_all_modes_window_sweep_ppl.png

C. Different sink sizes at the same window size

Only attention_sinks:

python scripts/plot.py \
  --outputs_dir outputs/window_1020_sinksize_all_3B \
  --model_tag Qwen_Qwen2.5-3B-Instruct \
  --compare sinks \
  --window_size 1020 \
  --mode attention_sinks \
  --metric running_perplexity \
  --output figures/3B_attention_sinks_sink_sweep_ppl.png

All modes:

python scripts/plot.py \
  --outputs_dir outputs/window_1020_sinksize_all_3B \
  --model_tag Qwen_Qwen2.5-3B-Instruct \
  --compare sinks \
  --window_size 1020 \
  --metric running_perplexity \
  --output figures/3B_all_modes_sink_sweep_ppl.png

D. Plot other metrics

To plot latency, change:

--metric running_perplexity

to:

--metric latency_sec

To plot cache length:

--metric cache_len

To plot peak GPU memory:

--metric peak_gpu_mem_mb

8. Suggested final experiment plan

Main benchmark

Run window sweeps for:

0.5B, 3B, 7B
sink_size = 4
window_size = 256, 512, 1020, 2048, 3072
num_tokens = 4096
prefill_tokens = 256

Sink ablation

Run sink sweeps for:

3B, 7B
window_size = 1020
sink_size = 4, 8, 20
num_tokens = 4096
prefill_tokens = 256

0.5B sink ablation can be included if already available.

9. Common issues and fixes

A. Job stops after closing VSCode

Use nohup or tmux. Do not run long jobs as plain foreground commands if disconnecting.

B. Wrong sink size appears in the log

Check argument order.

For window sweeps:

last argument = sink size

Example:

bash scripts/sweep_7B.sh 0 4096 outputs/sink_size_4_7B_window_all 4

For sink sweeps:

last argument = window size

Example:

bash scripts/sweep_7B_sink.sh 0 4096 outputs/window_1020_sinksize_all_7B 1020

C. Hugging Face dataset timeout

Cache the dataset once:

python - <<'PY'
from datasets import load_dataset

ds = load_dataset("emozilla/pg19-test", split="test", streaming=False)
print(ds)
print(ds[0].keys())
PY

Then rerun the failed sweep.

D. 7B attention_sinks NaN / blank NLL or perplexity

The 7B custom attention path may need float32 attention-score computation before softmax. The current fixed implementation should compute QK attention scores in float32 and cast the mask to the same device/dtype as the attention scores.

E. Plot script cannot find CSV files

Check that filenames contain both sink and window tags:

sink4
window1020

If outputs are spread across several folders, use a parent directory:

--outputs_dir outputs

because the plot script searches recursively.

10. Recommended paper plots

  1. Running perplexity vs step for all three modes at fixed sink=4, window=1020.
  2. GPU memory vs step for all three modes at fixed sink=4, window=1020.
  3. Cache length vs step for all three modes at fixed sink=4, window=1020.
  4. Latency vs step for all three modes at fixed sink=4, window=1020.
  5. Window-size ablation: attention_sinks and/or all modes across windows 256, 512, 1020, 2048, 3072.
  6. Sink-size ablation: attention_sinks across sink sizes 4, 8, 20 at fixed window=1020.
  7. Summary table with final perplexity, average latency, peak memory, and final cache length.

11. Quick command checklist

# Check GPU
nvidia-smi

# Monitor GPU live
watch -n 1 nvidia-smi

# Monitor logs
tail -f logs/<log_file>.log

# Stop benchmark jobs
pkill -f benchmark.perplexity

# Run 3B window sweep
nohup bash scripts/sweep_3B.sh 1 4096 outputs/sink_size_4_3B_window_all 4 > logs/sweep_3B_sink4.log 2>&1 &

# Run 7B window sweep
nohup bash scripts/sweep_7B.sh 0 4096 outputs/sink_size_4_7B_window_all 4 > logs/sweep_7B_sink4.log 2>&1 &

# Run 3B sink sweep
nohup bash scripts/sweep_3B_sink.sh 1 4096 outputs/window_1020_sinksize_all_3B 1020 > logs/sweep_3B_sink_window1020.log 2>&1 &

# Run 7B sink sweep
nohup bash scripts/sweep_7B_sink.sh 0 4096 outputs/window_1020_sinksize_all_7B 1020 > logs/sweep_7B_sink_window1020.log 2>&1 &

Check for the output directory carefully and change accordingly

############# End of the script ########################