> ## 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

> 使用 Flash Attention 2 加速训练

# Flash Attention

Flash Attention 2 通过优化内存访问模式，为 transformer 训练提供显著的加速。

## 要求

<Warning>
  **Flash Attention 2 需要：**

  * Linux 操作系统
  * 支持 CUDA 的 NVIDIA GPU
  * 已安装 `flash-attn` 包
</Warning>

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

## 快速开始

```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,
)
```

## 参数

| 参数                      | CLI 标志                    | 默认值     | 描述                                       |
| ----------------------- | ------------------------- | ------- | ---------------------------------------- |
| `use_flash_attention_2` | `--use-flash-attention-2` | `False` | 启用 Flash Attention 2                     |
| `attn_implementation`   | `--attn-implementation`   | `None`  | 覆盖注意力：`eager`、`sdpa`、`flash_attention_2` |

### 注意力实现选项

| 选项                  | 描述                                  |
| ------------------- | ----------------------------------- |
| `eager`             | 标准 PyTorch 注意力（某些模型的默认值）            |
| `sdpa`              | 缩放点积注意力（PyTorch 2.0+）               |
| `flash_attention_2` | Flash Attention 2（最快，需要 flash-attn） |

## 模型兼容性

<Warning>
  **Gemma 模型默认使用 eager 注意力。** 由于兼容性问题，Gemma 模型会自动禁用 Flash Attention 2。`attn_implementation` 被强制为 `eager`。
</Warning>

### 支持的模型

| 模型系列    | Flash Attention 2 | 备注           |
| ------- | ----------------- | ------------ |
| Llama   | 是                 | 完全支持         |
| Mistral | 是                 | 完全支持         |
| Qwen    | 是                 | 完全支持         |
| Phi     | 是                 | 完全支持         |
| Gemma   | **否**             | 使用 eager 注意力 |

## 与量化结合

将 Flash Attention 与量化结合以获得最大效率：

```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
```

## 与序列打包结合

Flash Attention 支持高效的序列打包：

```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>
  序列打包需要启用 Flash Attention。
</Note>

## 性能优势

| 配置                | 内存       | 速度       |
| ----------------- | -------- | -------- |
| 标准注意力             | 基线       | 基线       |
| SDPA              | \~15% 更少 | \~20% 更快 |
| Flash Attention 2 | \~40% 更少 | \~2x 更快  |

*结果因模型大小、序列长度和硬件而异。*

## 故障排除

### 安装错误

如果 `pip install flash-attn` 失败：

```bash theme={null}
# 确保安装了 CUDA 工具包
nvcc --version

# 使用特定 CUDA 版本安装
pip install flash-attn --no-build-isolation
```

### 运行时错误

**"Flash Attention 不可用"**

* 验证 flash-attn 已安装：`python -c "import flash_attn"`
* 确保您在 Linux 上且支持 CUDA
* 检查 GPU 计算能力（需要 SM 80+，例如 A100、H100）

**模型尽管设置了标志仍使用 eager 注意力**

* 某些模型（如 Gemma）强制使用 eager 注意力
* 检查模型文档以了解兼容性

## 下一步

<CardGroup cols={2}>
  <Card title="Quantization" href="/advanced/quantization">
    与内存优化结合
  </Card>

  <Card title="LoRA/PEFT" href="/advanced/lora-peft">
    高效微调
  </Card>
</CardGroup>