> ## Documentation Index
> Fetch the complete documentation index at: https://docs.monostate.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Flash Attention

> Speed up training with Flash Attention 2

# Flash Attention

Flash Attention 2 provides significant speedups for transformer training by optimizing memory access patterns.

## Requirements

<Warning>
  **Flash Attention 2 requires:**

  * Linux operating system
  * NVIDIA GPU with CUDA support
  * `flash-attn` package installed
</Warning>

```bash theme={null}
pip install flash-attn
```

## Quick Start

```bash theme={null}
aitraining llm --train \
  --model meta-llama/Llama-3.2-1B \
  --data-path ./data.jsonl \
  --project-name fast-model \
  --use-flash-attention-2
```

## Python API

```python theme={null}
from autotrain.trainers.clm.params import LLMTrainingParams

params = LLMTrainingParams(
    model="meta-llama/Llama-3.2-1B",
    data_path="./data.jsonl",
    project_name="fast-model",

    use_flash_attention_2=True,
)
```

## Parameters

| Parameter               | CLI Flag                  | Default | Description                                              |
| ----------------------- | ------------------------- | ------- | -------------------------------------------------------- |
| `use_flash_attention_2` | `--use-flash-attention-2` | `False` | Enable Flash Attention 2                                 |
| `attn_implementation`   | `--attn-implementation`   | `None`  | Override attention: `eager`, `sdpa`, `flash_attention_2` |

### Attention Implementation Options

| Option              | Description                                          |
| ------------------- | ---------------------------------------------------- |
| `eager`             | Standard PyTorch attention (default for some models) |
| `sdpa`              | Scaled Dot Product Attention (PyTorch 2.0+)          |
| `flash_attention_2` | Flash Attention 2 (fastest, requires flash-attn)     |

## Model Compatibility

<Warning>
  **Gemma models use eager attention by default.** Flash Attention 2 is automatically disabled for Gemma models due to compatibility issues. The `attn_implementation` is forced to `eager`.
</Warning>

### Supported Models

| Model Family | Flash Attention 2 | Notes                |
| ------------ | ----------------- | -------------------- |
| Llama        | Yes               | Full support         |
| Mistral      | Yes               | Full support         |
| Qwen         | Yes               | Full support         |
| Phi          | Yes               | Full support         |
| Gemma        | **No**            | Uses eager attention |

## With Quantization

Combine Flash Attention with quantization for maximum efficiency:

```python theme={null}
params = LLMTrainingParams(
    model="meta-llama/Llama-3.2-8B",
    data_path="./data.jsonl",
    project_name="fast-quantized",

    peft=True,
    quantization="int4",
    use_flash_attention_2=True,
)
```

```bash theme={null}
aitraining llm --train \
  --model meta-llama/Llama-3.2-8B \
  --data-path ./data.jsonl \
  --project-name fast-quantized \
  --peft \
  --quantization int4 \
  --use-flash-attention-2
```

## With Sequence Packing

Flash Attention enables efficient sequence packing:

```python theme={null}
params = LLMTrainingParams(
    model="meta-llama/Llama-3.2-1B",
    data_path="./data.jsonl",
    project_name="packed-model",

    use_flash_attention_2=True,
    packing=True,
)
```

<Note>
  Sequence packing requires Flash Attention to be enabled.
</Note>

## Performance Benefits

| Configuration      | Memory     | Speed        |
| ------------------ | ---------- | ------------ |
| Standard attention | Baseline   | Baseline     |
| SDPA               | \~15% less | \~20% faster |
| Flash Attention 2  | \~40% less | \~2x faster  |

*Results vary by model size, sequence length, and hardware.*

## Troubleshooting

### Installation Errors

If `pip install flash-attn` fails:

```bash theme={null}
# Ensure CUDA toolkit is installed
nvcc --version

# Install with specific CUDA version
pip install flash-attn --no-build-isolation
```

### Runtime Errors

**"Flash Attention is not available"**

* Verify flash-attn is installed: `python -c "import flash_attn"`
* Ensure you're on Linux with CUDA
* Check GPU compute capability (requires SM 80+, e.g., A100, H100)

**Model uses eager attention despite flag**

* Some models (like Gemma) force eager attention
* Check model documentation for compatibility

## Next Steps

<CardGroup cols={2}>
  <Card title="Quantization" href="/advanced/quantization">
    Combine with memory optimization
  </Card>

  <Card title="LoRA/PEFT" href="/advanced/lora-peft">
    Efficient fine-tuning
  </Card>
</CardGroup>