Environment:FlagOpen FlagEmbedding Finetuning Environment
| Knowledge Sources | |
|---|---|
| Domains | Infrastructure, Deep_Learning, Distributed_Training |
| Last Updated | 2026-02-09 21:00 GMT |
Overview
Extended GPU environment with DeepSpeed and Flash Attention for distributed fine-tuning of embedders and rerankers.
Description
This environment extends the base Python/PyTorch environment with additional packages required for fine-tuning workflows. It includes DeepSpeed for distributed training (ZeRO Stage 0 and Stage 1 configurations are provided) and Flash Attention 2 for efficient attention computation on supported GPUs. These are declared as optional extras in `setup.py` under the `[finetune]` extra. Flash Attention 2.10+ is required for certain MiniCPM reranker models.
Usage
Use this environment when running Embedder Finetuning or Reranker Finetuning workflows. It is required by the `EmbedderRunner.run()` and `RerankerRunner.run()` implementations which use `torchrun` for distributed training with DeepSpeed ZeRO.
System Requirements
| Category | Requirement | Notes |
|---|---|---|
| OS | Linux | DeepSpeed has limited Windows/macOS support |
| Hardware | NVIDIA GPU with CUDA | Flash Attention requires Ampere (A100) or newer |
| VRAM | 16GB+ recommended | For 7B+ parameter models with gradient checkpointing |
| Multi-GPU | 1+ GPUs | DeepSpeed supports multi-node multi-GPU |
Dependencies
System Packages
- NVIDIA CUDA toolkit >= 11.0
- NVIDIA cuDNN
- `ninja` (build tool for Flash Attention compilation)
Python Packages
- All packages from the base Python_PyTorch_Environment
- `deepspeed` (distributed training)
- `flash-attn` (Flash Attention 2 for efficient attention; >= 2.10 for MiniCPM models)
Credentials
- `HF_TOKEN`: Required if fine-tuning gated models (e.g., Llama-based). Set via environment variable.
Quick Install
# Install FlagEmbedding with fine-tuning extras
pip install FlagEmbedding[finetune]
# Or install components individually
pip install FlagEmbedding deepspeed flash-attn
Code Evidence
Finetune extras from `setup.py:27-29`:
extras_require={
'finetune': ['deepspeed', 'flash-attn'],
},
Flash Attention detection from `FlagEmbedding/inference/reranker/decoder_only/models/modeling_minicpm_reranker.py:48-62`:
from transformers.utils import (
is_flash_attn_2_available,
is_flash_attn_greater_or_equal_2_10,
)
# ...
try:
from flash_attn import flash_attn_func, flash_attn_varlen_func
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
except:
pass
Flash Attention model argument from `FlagEmbedding/finetune/embedder/decoder_only/base/arguments.py:39`:
use_flash_attn: bool = field(
default=False,
metadata={"help": "Use flash attention or not."}
)
Auto-configuration of Flash Attention from `FlagEmbedding/finetune/reranker/decoder_only/layerwise/configuration_minicpm_reranker.py:183-187`:
try:
import flash_attn
self._attn_implementation = "flash_attention_2"
except:
pass
Common Errors
| Error Message | Cause | Solution |
|---|---|---|
| `ImportError: No module named 'deepspeed'` | DeepSpeed not installed | `pip install deepspeed` |
| `ImportError: flash_attn` | Flash Attention not installed | `pip install flash-attn` (requires CUDA and ninja) |
| Flash Attention compilation error | Missing CUDA toolkit or ninja | `pip install ninja` and verify CUDA is in PATH |
| `is_flash_attn_greater_or_equal_2_10` fails | Flash Attention < 2.10 | `pip install flash-attn>=2.10` for MiniCPM models |
Compatibility Notes
- Flash Attention: Requires NVIDIA Ampere (A100) or newer GPUs. Does not work on older architectures (V100, P100).
- DeepSpeed ZeRO: Stage 0 and Stage 1 configs are provided in `examples/finetune/`. Stage 0 disables sharding; Stage 1 shards optimizer states.
- Linux only: DeepSpeed has limited support on Windows and macOS. Linux is strongly recommended for fine-tuning.
- MiniCPM models: The layerwise reranker based on MiniCPM auto-enables Flash Attention 2 if the package is detected.