Activation-Beacon 长上下文训练指南:从预训练到微调(Mistral / Llama-3 / Qwen-2)
Activation-Beacon 长上下文训练指南从预训练到微调Mistral / Llama-3 / Qwen-2【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbeddingActivation-Beacon 是 FlagEmbedding 仓库research/Long_LLM/activation_beacon目录下的长上下文插件模块通过引入可压缩的 beacon token 将大语言模型的上下文窗口扩展到数万甚至数十万 token。本文基于 examples/training.md 展开完整讲解其预训练 微调两阶段训练方案如何准备数据、如何逐参数读懂并运行六条torchrun训练命令、三个官方支持模型Mistral、Llama-3、Qwen-2的差异化配置以及 beacon 各超参数背后的源码机制。读完本文你可以直接从零复现一个长上下文模型的全流程训练。两阶段训练方案总览Activation-Beacon 的训练分为两个阶段这与传统预训练 指令微调范式一致但数据构成和序列长度策略有明显差异预训练Pretrain数据来源redpajama 采样约1B token任务为自回归语言建模处理方式每个文档末尾追加eos标记不做 packing即不把多个短文档拼接进同一条样本文档边界被完整保留最大上下文20K token。微调Finetune数据构成混合 4 类来源总量约 28K 条LongAlpaca 12k 中采样5K条长指令样本Booksum 中采样2K条书籍摘要样本由 GPT-3.5 生成的16K条合成长上下文问答数据one_detail_book与one_detail_paper两类从预训练数据中回采5K条最大上下文同样为20K token但--min_length从预训练的 2400 提高到 7200保证每条样本都是真正的长序列。从源码角度看这两个阶段共用同一套训练入口 main/train.py差异完全由train_data、min_length、beacon_ratio、chat_template等命令行参数体现。微调阶段--beacon_ratio只使用2 4 8预训练为2 4 8 16 32意味着微调对上下文压缩强度的要求更温和重点在于让模型学会遵循长指令、回答长文问答。环境准备与数据下载训练前需要按 README.md 完成两件事创建运行环境、下载数据。环境安装以 Python 3.10 CUDA 12.1 为例conda create beacon python3.10.14 conda activate beacon # 可根据实际 CUDA 版本调整 conda install pytorch pytorch-cuda12.1 -c pytorch -c nvidia pip install transformers deepspeed accelerate datasets peft pandas seaborn rouge fuzzywuzzy jieba python-Levenshtein pip install flash-attn --no-build-isolation下载训练与评测数据long-llm.tar.gz解压位置可任选官方示例放在/datawget https://huggingface.co/datasets/namespace-Pt/projects/resolve/main/long-llm.tar.gz?downloadtrue -O /data/long-llm.tar.gz cd /data tar -xzvf long-llm.tar.gz这里有一个贯穿全部训练/评测命令的关键机制所有以long-llm:前缀开头的路径都会被解析为相对于data_root的相对路径。例如long-llm:redpajama/train.json实际指向${data_root}/redpajama/train.json。data_root的默认值是/data/long-llm可以在 src/args.py 中修改从而避免每条命令都写绝对路径。该解析逻辑实现在ModelArgs.resolve_path中src/args.py并在__post_init__中自动对train_data、eval_data生效。训练入口与核心参数体系所有训练命令都通过torchrun --nproc_per_node 8 $DDP -m main.train ...启动其中$DDP是环境变量形式的附加参数典型内容如--master_addr... --master_port...。入口脚本 main/train.py 的调用链如下用HfArgumentParser同时解析 ModelArgs 与 TrainingArgs 两类参数通过get_model_and_tokenizer加载基座模型若开启 beacon 且only_train_beaconTrue则冻结所有非 beacon 参数源码中通过检查参数名是否包含beacon实现见 main/train.py加载训练集与评测集构造ActivationBeaconTrainer并调用train()。TrainingArgs继承自 HuggingFaceTrainingArguments因此learning_rate、num_train_epochs、save_strategy、evaluation_strategy、logging_steps、bf16等均是标准语义。下面按功能分组说明训练相关核心参数默认值取自 src/args.py参数默认值说明--enable_beaconFalse是否启用 Activation-Beacon 插件模块训练长上下文模型必须开启--beacon_windowNone初始滑窗大小即每个 window 保留的原始 token 数Mistral/Qwen-2 用 2048Llama-3 用 1024--beacon_strideNone滑窗步长。源码断言beacon_window beacon_stride且交错interleave模式要求二者相等见 src/modeling_beacon.py--beacon_attnNonebeacon token 的注意力掩码分配方式可选segmentation、step-expansion、full-coverage训练命令统一用full-coverage--beacon_attend_prevNone前序 beacon token 之间是否允许互相 attendTrue时历史压缩信息可跨窗口传播--beacon_sink_sizeNone序列头部始终保留的原始激活数量StreamingLLM 的 attention sink 思想。Llama-3 必须设为 1Mistral/Qwen-2 设为 0--beacon_ratioNone压缩比率列表预训练用2 4 8 16 32微调用2 4 8训练时按--beacon_ratio_mix策略从列表中挑选--beacon_ratio_mixNone每条样本压缩比的选取策略可选step-random、instance-random、sequence或adapt-x系列--beacon_paramNone为 beacon 引入可训练参数投影的注意力子模块可选q、k、v、o训练命令统一为q k v--beacon_embed_initeosbeacon token 的 embedding 初始化来源可从eos或bosembedding 初始化--beacon_posNonebeacon token 的摆放位置可选append追加到窗口末尾或interleave交错插入训练命令统一用interleave--beacon_parallel_windowNone并行处理的窗口数量用于吞吐优化--min_length0训练样本的最小 token 数短于此长度的样本被过滤预训练 2400微调 7200--max_length4096每条输入的最大 token 数训练统一为 20000--group_by_strideNone按 beacon 滑窗步数对样本分组relaxed/strict训练统一用strict--only_train_beaconTrue冻结 LLM 主干、只训练 beacon 参数是插件式训练的核心开关--use_reentrantNone梯度检查点是否使用 reentrant 实现必须设为False。源码注释强调需要保留计算图以便 beacon token 获得梯度见 src/args.py--chat_templatehf训练/评测时套用的对话模板Mistral 用mistral、Llama-3 用llama-3、Qwen-2 用qwen--attn_implflash_attention_2注意力实现配合 Flash-Attention-2 使用--eval_methodperplexity训练过程中的评测方式perplexity/generation--eval_beacon_ratio[32]评测时的压缩比--eval_beacon_ratio_mixadapt-1024评测时压缩比选取策略其中use_reentrant参数在TrainingArgs.__post_init__中被写入gradient_checkpointing_kwargs因此命令中必须同时给出--gradient_checkpointing --use_reentrant False才完整生效。预训练Mistral-7B 实例逐参数拆解以 Mistral 预训练命令为范例其余模型命令结构完全一致仅参数值不同output_namebeacon-mistral-pretrain torchrun --nproc_per_node 8 $DDP -m main.train \ --output_dir data/outputs/$output_name \ --model_name_or_path mistralai/Mistral-7B-Instruct-v0.2 \ --train_data long-llm:redpajama/train.json \ --min_length 2400 \ --max_length 20000 \ --group_by_stride strict \ --enable_beacon \ --beacon_window 2048 \ --beacon_stride 2048 \ --beacon_attn full-coverage \ --beacon_attend_prev True \ --beacon_sink_size 0 \ --beacon_ratio 2 4 8 16 32 \ --beacon_ratio_mix step-random \ --beacon_param q k v \ --beacon_pos interleave \ --attn_impl flash_attention_2 \ --gradient_checkpointing \ --use_reentrant False \ --save_only_model \ --save_strategy epoch \ --evaluation_strategy steps \ --num_train_epochs 1 \ --logging_steps 50 \ --bf16 \ --deepspeed data/deepspeed/stage2.json逐项解读关键配置--train_data long-llm:redpajama/train.json解析后为${data_root}/redpajama/train.json。预训练数据按语言建模方式处理源码 src/data.py 显示短于min_length的文本被跳过介于min_length与max_length之间的文本在末尾追加eos超过max_length的文本被截断——这与文档中每个文档加 eos、不做 packing的描述一致。--min_length 2400过滤掉过短的文档避免训练样本过短导致 beacon 压缩失去意义。--max_length 20000预训练阶段直接把序列上限拉到 20K是普通 4K 基座模型无法直接处理的长度靠 beacon 压缩使其成为可能。--group_by_stride strictstrict模式下样本按滑窗步数严格分组同一 batch 内序列的压缩结构一致减少 padding 浪费并稳定梯度。--beacon_ratio 2 4 8 16 32 --beacon_ratio_mix step-random预训练阶段在 5 档压缩比之间按步随机切换。源码 src/modeling_beacon.py 显示step-random会为每个窗口步随机抽取一个比率并循环推进而instance-random是整条样本使用同一随机比率。多档随机混合让模型同时适应 2×~32× 的压缩强度是论文中flexible compression的关键实现。--beacon_param q k v仅在 attention 的 q/k/v 投影上为 beacon token 增加可学习参数控制新增参数量。--beacon_pos interleavebeacon token 交错插入各窗口之间而非追加在序列末尾。交错模式要求beacon_window beacon_stride源码有显式断言这也是 Mistral 配置中二者均为 2048 的原因。--beacon_sink_size 0Mistral 不需要 attention sink。--attn_impl flash_attention_2使用 Flash-Attention-2 加速并节省显存训练环境需按上文安装flash-attn。--gradient_checkpointing --use_reentrant False梯度检查点必须配合非 reentrant 实现否则 beacon token 的梯度路径会被破坏。--deepspeed data/deepspeed/stage2.json使用仓库自带的 ZeRO-2 配置data/deepspeed/stage2.json其中lr、betas、weight_decay、train_batch_size等均设为auto由训练命令中的--bf16、learning_rate等参数接管目录下还提供stage3.json、stage3-offload.json、stage2-offload.json等变体data/deepspeed/显存紧张时可切换。--save_strategy epoch --evaluation_strategy steps --num_train_epochs 1每个 epoch 保存一次 checkpoint按步做困惑度评测训练 1 个 epoch。微调让 beacon 学会长指令与长文问答微调阶段在预训练产出的 checkpoint 上继续命令示例如下Mistraloutput_namebeacon-mistral-finetune torchrun --nproc_per_node 8 $DDP -m main.train \ --output_dir data/outputs/$output_name \ --model_name_or_path data/outputs/beacon-mistral-pretrain/* \ --train_data long-llm:gpt/one_detail_book.train.16K.json long-llm:gpt/one_detail_paper.train.16K.json long-llm:longalpaca/train.json long-llm:booksum/train.16K.json long-llm:needle/train.16K.json long-llm:redpajama/train.json[5000] \ --max_length 20000 \ --min_length 7200 \ --group_by_stride strict \ --enable_beacon \ --beacon_window 2048 \ --beacon_stride 2048 \ --beacon_attn full-coverage \ --beacon_attend_prev True \ --beacon_sink_size 0 \ --beacon_ratio 2 4 8 \ --beacon_ratio_mix step-random \ --beacon_param q k v \ --beacon_pos interleave \ --attn_impl flash_attention_2 \ --learning_rate 1e-5 \ --gradient_checkpointing \ --use_reentrant False \ --save_only_model \ --num_train_epochs 1 \ --save_strategy epoch \ --logging_steps 50 \ --bf16 \ --deepspeed data/deepspeed/stage2.json \ --chat_template mistral微调阶段的关键差异--model_name_or_path data/outputs/beacon-mistral-pretrain/*直接指向预训练输出目录通配符匹配具体 checkpoint 子目录。--train_data传入 6 个文件16K 合成长文问答one_detail_book/one_detail_paper、LongAlpaca、Booksum、needle长文检索任务数据以及回采的 5K 预训练数据。末尾的[5000]语法表示只从该文件采样 5000 条——源码 src/data.py 通过正则\[(\d*)\]解析该后缀并限制每个文件的样本上限。--min_length 7200微调要求每条样本至少 7200 token确保模型在真实长序列上做指令学习。--learning_rate 1e-5微调使用较低学习率且only_train_beacon默认开启主干参数被冻结只更新 beacon 相关参数训练开销小、不易灾难性遗忘。--chat_template mistral训练数据conversations字段会按指定对话模板格式化。源码 src/data.py 显示指令微调数据按apply_chat_template处理超长样本直接跳过max_length之外丢弃而非截断。微调阶段--evaluation_strategy steps被移除、beacon_ratio缩减为2 4 8聚焦于精调而非继续拓宽压缩能力。三种基座模型的差异化配置三组命令的结构完全一致差异集中在四个参数上模型基座官方命令beacon_window/beacon_stridebeacon_sink_sizechat_templateMistral-7Bmistralai/Mistral-7B-Instruct-v0.22048 / 20480mistralLlama-3-8Bmeta-llama/Meta-Llama-3-8B-Instruct1024 / 10241llama-3Qwen-2-7BQwen/Qwen2-7B-Instruct2048 / 20480qwenLlama-3 的 attention sink 是硬性要求文档中明确标注 according to our experiment, Llama-3 requires attention sink即--beacon_sink_size 1。这是因为 Llama-3 的注意力模式极度依赖序列首 token 作为 sink压缩掉头部 token 会导致困惑度显著劣化Mistral 与 Qwen-2 则不存在该问题sink 设为 0。同时 Llama-3 的窗口/步长减半为 1024滑窗颗粒度更细。其余配置full-coverage注意力、step-random比率混合、q k v参数、interleave位置、flash_attention_2、bf16 DeepSpeed ZeRO-2三模型完全一致。beacon 参数背后的实现约束源码级理解训练命令的取值有必要了解 src/modeling_beacon.py 中的硬性约束否则运行时会直接报错beacon_window beacon_stride滑窗不能小于步长否则会产生覆盖不到的序列区域beacon_attn仅接受segmentation/step-expansion/full-coverage训练命令统一使用full-coverage让每个 beacon token 能覆盖到其所在窗口的完整信息beacon_ratio_mix仅接受instance-random/step-random/sequence或以adapt-开头的策略step-random在训练中按窗口步推进比率配合--beacon_ratio 2 4 8 16 32实现同一条样本内部混合多种压缩强度beacon_posinterleave要求beacon_window beacon_stride并且该模式还要求beacon_ratio中不能出现 0即不能跳过某些窗口的压缩理论上可支持的最大序列长度与压缩比直接相关源码get_max_length按(max_position_embeddings - beacon_window) * compression_ratio beacon_window估算见 src/modeling_beacon.py这也是从 4K 到 400K的压缩式扩展原理在训练侧的具体体现。训练中的评测与检查点策略--evaluation_strategy steps预训练训练过程中按步对评测集计算困惑度便于观察长序列建模能力是否随训练提升--save_strategy epoch--save_only_model每 epoch 保存一次且只保存模型权重、不保存优化器状态节省磁盘--logging_steps 50每 50 步输出一次日志--bf16混合精度训练与 DeepSpeed 配置中的bf16: {enabled: auto}联动生效见 data/deepspeed/stage2.json训练产物为各模型后续微调--model_name_or_path指向预训练输出以及下游评测见 examples/evaluation.md其中覆盖 InfBench、LongBench、Needle-in-a-Haystack、MMLU 等任务提供基础。常见问题与注意事项$DDP环境变量命令中的$DDP需在调用前定义为分布式训练的附加参数如 master 地址与端口否则torchrun会因缺少必要参数启动失败。数据路径前缀务必保持long-llm:前缀否则路径不会被resolve_path解析会按字面相对路径查找而报文件不存在。Flash-Attention-2 未安装--attn_impl flash_attention_2要求环境中已安装flash-attn否则模型加载即失败。use_reentrant False不可省略这是 beacon 训练能正常反传梯度对 beacon token 求梯度的前提源码注释明确说明该点。Llama-3 必须设置 sink漏掉--beacon_sink_size 1会明显损害长序列困惑度。显存不足可优先切换 data/deepspeed/ 下的stage3/stage3-offload配置并调整per_device_train_batch_size默认 1该目录代码支持 DeepSpeed ZeRO-3 训练。按照本文的六条命令三个模型 × 预训练/微调配合 README.md 的环境与数据准备即可完整复现 Activation-Beacon 的长上下文训练流程并在之后用 examples/evaluation.md 中的评测脚本验证模型的长序列能力。【免费下载链接】FlagEmbeddingRetrieval and Retrieval-augmented LLMs项目地址: https://gitcode.com/GitHub_Trending/fl/FlagEmbedding创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考