OctoLong：用跨仓库代码上下文增强代码大模型长上下文能力

在真实的代码开发场景里一个功能往往跨越多个文件一段 bug 修复也经常涉及调用链上下游。代码大模型如果只能在单文件片段上做预测就很难真正理解仓库层面的依赖关系。OctoLong 给出了一条值得关注的技术路线在通用预训练之后、下游微调之前插入一个 mid-training 中间训练阶段用跨仓库代码上下文继续训练模型从而增强长上下文建模能力。文章把这条路线拆开来讲包括 mid-training 为什么能补长上下文短板、跨仓库上下文数据如何构造、训练时有哪些工程注意点、评测怎样设计才有说服力以及复现这类工作经常遇到哪些坑。1. 先从代码模型的长上下文痛点说起1.1 上下文窗口和有效上下文是两件事很多人在选代码模型时习惯先把“上下文窗口”当核心指标32K、128K、200K数字越大似乎越强。但上下文窗口只代表模型“最多能接收多长的输入”并不代表模型“能在这个长度内有效利用信息”。一个模型在 128K 窗口下可能依然只从最近的 2000 个 token 里获取关键信息窗口前段的函数定义、仓库注释、配置信息早就被注意力机制忽略了。代码场景比自然语言场景更依赖远距离信息。比如你要理解一个函数handle_payment(order)它的类型定义可能在models/order.py订单状态常量在constants/payment.py数据库查询逻辑在repositories/order_repo.py。如果模型只读到函数体片段缺少对依赖符号和返回类型的感知生成的代码大概率在类型、边界条件或异常处理上出错。所以长上下文建模的核心问题不是“窗口有多大”而是“模型在窗口内是否真的会跨片段建立引用关系”。这也是 mid-training 这类方法存在的理由。位置编码外推和注意力优化解决的是“能不能容纳长序列”而训练数据决定“模型在长序列中该学什么”。如果训练数据永远是把互不相关的短文本拼在一起模型即使有长窗口也学不会有效定位和聚合信息。1.2 现有长上下文方案在代码场景的短板当前让代码模型支持长上下文常见做法有四类。每一类都有作用但都有特定盲区。方案类型典型实现主要解决问题代码场景短板位置编码外推RoPE 扩展、NTK、YaRN、ALiBi让模型接收超出训练长度的序列只解决长度限制不解决信息利用率长上下文续训在长文档上继续预训练让模型适应长序列的数据分布网页文本堆叠多仓库结构信息弱长上下文 SFT用长指令数据微调让模型学会按指令消费长输入依赖人工标注或高质量长任务数据检索增强 RAG向量召回、BM25、GraphRAG从外部候选里找回相关片段召回质量决定上限打断整体结构细看会发现一个共性问题这些方案要么只改模型结构要么只用“长但不一定相关”的文本做训练。代码仓库本身有极强的结构信息比如 import、require、include、目录层次、符号引用、commit 依赖这些结构比普通长文本更适合训练长上下文建模。OctoLong 选择跨仓库代码上下文做 mid-training本质上是把代码仓库的结构优势变成训练信号。1.3 mid-training 在训练流程中的定位LLM 的常规流水线可以概括为三个阶段通用预训练、监督微调、偏好对齐。mid-training 被放在预训练之后、微调之前也可以理解为“领域定向继续预训练”。它和 SFT 的区别很明显SFT 会改变模型输出的格式让模型学会“用户提问 - 模型回答”的模式而 mid-training 通常不做指令格式转换而是继续用 next-token prediction 的方式把模型推向某个特定领域。它和普通继续预训练的区别则在于目标选择继续预训练的目标是继续累积领域知识mid-training 的目标往往更聚焦比如提升长上下文能力、增强多文件代码理解、适应特定领域术语。OctoLong 这个名字拆开看Octo 可以联想到“八进制”“八爪鱼”“多分支结构”的意象Long 指向 long-context。标题已经把它要做的事情讲清楚了用跨仓库代码上下文做 mid-training目的是增强长上下文建模。也就是说这套方法不改变模型的下游任务输出方式而是在模型能力矩阵里专门补“长上下文代码理解”这块短板。2. OctoLong 为什么选择跨仓库代码上下文2.1 单文件样本撑不起仓库级理解单文件数据构造很简单把文件按长度截断或者按 tokenizer 的max_length切块。这种样本的优点是干净、易处理缺点是天然丢失了仓库里的横向依赖。一个函数被截掉一半或者一个文件里反复出现的工具类信息只出现一次模型要从这些碎片里学会完整代码逻辑本质上是在做“读残卷猜全貌”。跨仓库代码上下文改变了这个状态。一个训练样本可以同时包含多个文件内容文件之间依靠真实代码结构产生关联。例如入口函数src/api.py引入了models/item.py样本里就同时出现这两个文件模型在预测src/api.py后续内容时有机会去models/item.py里寻找类型定义。这种训练迫使注意力机制做远距离检索和聚合而不是只看局部窗口。这里需要解释“跨仓库”的含义。它不一定是“把多个完全不相关的 GitHub 仓库随机拼在一起”而是指数据组织跨出了单文件边界。一个 pull request 中同时改动的多个文件、一个服务模块和它依赖的 SDK 源码、一个仓库中相近目录下的核心实现与测试文件都可以构成跨文件、甚至跨仓库依赖的训练单元。2.2 一个训练样本长什么样在没有官方数据生成工具的情况下可以按下面这种结构设计样本。核心思路是把仓库快照转成带分隔符的长文本然后用这些长文本做 next-token prediction。{ id: cross-repo-context-0001, repo: example/checkout-system, base_ref: main, context: [ { path: src/api.py, content: from models.item import Item\nfrom constants.status import OrderStatus\n\ndef create_order(item: Item):\n ..., sep: file path\src/api.py\ }, { path: models/item.py, content: class Item:\n name: str\n price: decimal.Decimal\n ..., sep: file path\models/item.py\ }, { path: constants/status.py, content: class OrderStatus:\n PENDING PENDING\n PAID PAID\n ..., sep: file path\constants/status.py\ } ], target: def mark_paid(order_id):\n ... }context是按语义关联挑选出的多文件内容target是要生成的目标片段。如果做纯续训风格可以将context里的文件依次拼接成一个大字符串去掉target直接做语言建模如果做指令风格可以在target前加一句以上代码来自多个文件请根据实现补全缺失逻辑之类的话形成监督信号。def build_training_text(sample: dict) - str: segments [] for file_obj in sample[context]: segments.append(file_obj[sep]) segments.append(file_obj[content]) segments.append(task) segments.append(sample.get(target, )) return \n.join(segments)上面的sep并不一定需要保留是否添加文件路径标记取决于模型在推理时是否会使用这类格式。如果打算把 mid-training 的产物继续做下游代码任务微调保留路径标记可以帮助模型理解文件边界。2.3 为什么跨仓库数据能增强长上下文建模第一样本长度自然变长。单文件样本通常只有几百到几千 token把几个有关联的文件拼在一起后样本很容易超过一万 token。模型被迫在更长的序列上反复计算对长序列的数值稳定性、注意力分布和梯度传播都会更适应。第二信息密度更高。代码里的符号引用是显式的约束。模型要预测OrderStatus.PAID后面的内容就要先看到constants/status.py里的定义。这种长距离依赖是天然的训练目标比人工构造的“长文本里找一句话”更丰富。第三保留了通用语言的特性。代码数据里包含注释、文档字符串、commit message、配置说明和自然语言描述这些内容仍然是通用语料的一部分。因此基于代码仓库的 mid-training 不会像纯合成文本那样导致过度的领域偏离。值得提醒的是跨仓库数据不能无脑拼接。如果只是把语料库里所有文件随机拼一大段那模型学到的依然是“长而无序”的文本甚至可能因为虚假关联而扰乱注意力。真正有效的前提是文件之间通过 import、函数调用、类型引用、commit 共现等关系连接。3. 训练环境和工程准备3.1 硬件资源要按“最长样本”规划长上下文训练对显存的消耗不是线性增长而是接近二次增长。即使使用 FlashAttention激活值、梯度、优化器状态依然会随着序列长度快速膨胀。在常见实践中复现这类方法至少需要 8 卡 A100 或同级别显卡代码模型规模在 7B 到 13B 之间时32K 上下文属于可以接受的起步配置。资源类型小规模验证正式复现生产级迭代GPU1-2 张 24G 以上8 张 A100 80G多机多卡集群系统内存128G256G512G 以上存储1-2T10T 以上对象存储 快文件系统训练框架transformers PEFTDeepSpeed / Megatron-LM自定义分布式数据管道注意力实现标准 attentionFlashAttention 2定制 kernel 序列并行学习阶段的验证不一定要完整复现。可以先用 1K 上下文跑通数据管道再用 8K 上下文训练一个很小的模型确认数据和代码都没有问题后再上完整规模。3.2 基础模型和位置编码选型标题没有指定 OctoLong 使用哪个基础模型从开源生态看这类方法通常会选择支持长窗口、代码理解能力较强的模型。选型时建议重点检查三点。位置编码是否支持扩展。RoPE 系列模型通常可以通过调整rope_scaling或 base 频率来扩展但扩展后的稳定性需要自己验证。tokenizer 对代码的压缩率。一个过度拆分代码符号的 tokenizer 会让长样本更快触达长度上限浪费有效上下文。特殊 token 是否一致。文件分隔符、任务标记、结束标记都要在训练和推理时保持一致。# 一个可参考的模型加载配置示例 model_name_or_path: Qwen/Qwen2.5-Coder-7B torch_dtype: bfloat16 attn_implementation: flash_attention_2 rope_scaling: type: yarn factor: 2.0实际使用时rope_scaling的配置需要结合基础模型的原始训练长度和你的目标长度来计算。不要盲目调大factor否则位置编码外推会降低模型的局部精度。3.3 训练框架的关键参数下面是 DeepSpeed 场景下启动训练时常见的一组参数用来解释它们各自的作用。deepspeed --num_gpus8 train.py \ --model_name_or_path Qwen/Qwen2.5-Coder-7B \ --data_path ./data/train.jsonl \ --block_size 32768 \ --bf16 true \ --per_device_train_batch_size 1 \ --gradient_accumulation_steps 8 \ --gradient_checkpointing true \ --learning_rate 2e-5 \ --lr_scheduler_type cosine \ --warmup_ratio 0.05 \ --deepspeed ds_z3_config.json参数含义常见设置调大/调小影响block_size训练序列最大长度8192~32768越大显存越高对长上下文收益越大per_device_train_batch_size单卡 batch 大小1~4长序列下通常只能设 1gradient_accumulation_steps梯度累积步数4~16越大越接近大 batch但训练更慢gradient_checkpointing重计算激活值开启明显降低显存增加训练时间learning_rate学习率1e-5~5e-5过大导致灾难性遗忘过小收敛慢warmup_ratio预热比例0.03~0.1长序列训练不稳定时可调大不要把单卡 batch size 调大作为优化目标长上下文训练中优先保证样本长度覆盖目标窗口再通过梯度累积控制有效 batch size。4. 从数据构造到最小训练闭环4.1 仓库数据收集和过滤第一步是确定数据源。可以基于公开代码语料也可以自建私有仓库合集。需要遵守各平台的许可协议并对仓库做基础过滤。收集过程中的常见过滤条件文件大小超过阈值或低于阈值比如小于 100 字节的配置碎片和大于 2MB 的生成文件明显是锁文件或构建产物的路径例如package-lock.json、dist/、node_modules/非目标语言的源文件如果只想训练 Python 相关代码就过滤掉大部分不相关内容。测试文件和样例目录是否需要保留取决于目标能力。如果想增强“修改代码后同步修改测试”的能力就保留测试文件。# 一个简单的过滤函数示例 import json from pathlib import Path def is_valid_path(path: str) - bool: block_substrings [node_modules, dist, build, .git, __pycache__] return not any(part in block_substrings for part in path.split(/)) def build_sample(repo_path: Path, candidates: list) - dict: valid [] for rel_path in candidates: full_path repo_path / rel_path if not full_path.exists(): continue if not is_valid_path(rel_path): continue content full_path.read_text(encodingutf-8, errorsignore) if len(content) 200 or len(content) 500_000: continue valid.append({path: rel_path, content: content}) return {repo: repo_path.name, context: valid[:8]}这里的candidates需要根据 import 图或 commit 关系统计得出。如果只是把一个仓库的前 8 个文件按字典序放入样本模型很难学到真正的跨文件依赖。4.2 文件关联度计算一个可以直接落地的思路是先用 AST 或正则解析 import 语句再构建“文件 - 被引用文件”的图最后从某个入口文件出发做 BFS。这个做法虽然简单但比随机挑文件更能体现跨文件依赖。# 基于 import 关系选择相关文件 from collections import deque def get_related_files(import_graph: dict, entry: str, depth: int 2): seen {entry} queue deque([(entry, 0)]) while queue: node, d queue.popleft() if d depth: continue for nxt in import_graph.get(node, []): if nxt not in seen: seen.add(nxt) queue.append((nxt, d 1)) return list(seen)如果目标是跨仓库依赖还要把子模块之间共享的基础库也纳入。例如主项目 import 了内部的common-lib那么构造样本时可以把主项目入口文件、相关业务文件和common-lib的实现文件放入同一个样本。这种跨仓库上下文对模型理解真实工程结构非常有帮助。需要注意BFS 这个策略只适合“快速验证”。如果做正式训练建议使用更丰富的关联信号函数调用关系、同名符号引用、commit 历史里同时修改的文件、README 和文档指向的模块。关联信号越丰富样本的信息结构越接近真实开发场景。4.3 tokenization 和序列切分长样本在 tokenize 后可能超过block_size。这里有两种处理方式。切分式把一个超长样本按block_size切多段每段独立作为训练样本。缺点是可能把一个跨文件上下文从中间切开损失长距离依赖。重试式如果一个样本超过max_length先通过仓库依赖信息找到更核心的入口文件重新选择相关文件直到样本长度落在目标范围。def trim_to_target_length(sample: dict, tokenizer, max_len: int) - dict: text build_training_text(sample) tokens tokenizer.encode(text, add_special_tokensFalse) if len(tokens) max_len: return sample # 重新按依赖优先级选择文件优先保留入口和核心依赖 ordered_files sorted(sample[context], keylambda f: f.get(priority, 0), reverseTrue) new_context [] total 0 for f in ordered_files: tokens_so_far tokenizer.encode(f[sep] f[content], add_special_tokensFalse) if total len(tokens_so_far) max_len - 128: continue new_context.append(f) total len(tokens_so_far) sample[context] new_context return sample这里预留 128 个 token 给后续的 task 标记和 target是一种常见做法。如果预留不足拼接后仍然会超长并产生截断噪声。4.4 训练数据配比mid-training 阶段的数据配比会直接影响效果。一个相对稳妥的起步方案是数据来源建议比例作用跨仓库/跨文件代码上下文50%~70%强化长距离代码依赖建模普通单文件代码语料15%~25%保持短代码能力防止局部建模退化通用文本和文档数据10%~20%减少通用能力遗忘不同项目的基础模型不同配比也要调。判断标准就是评测如果通用代码能力下降说明单文件代码语料占比太低如果长上下文评测没有提升说明跨仓库数据的构造质量有问题而不是比例问题。5. 评测如何判断长上下文能力真的变强5.1 评测任务要覆盖三种能力长上下文代码能力的评测至少要覆盖定位、理解、生成三层。能力层代表任务示例指标定位能力在仓库上下文中找到某个符号定义命中率理解能力基于跨文件信息回答代码逻辑问题准确率、F1生成能力根据上下文生成缺失函数或修复代码CodeBLEU、Passk只跑一个总指标不够。比如语言模型的困惑度下降不一定会带来代码生成准确率提升长文档问答提升也不代表模型在跨文件补全任务上更强。建议按能力层分别记录结果。5.2 用 needle-in-a-haystack 做快速体检针尖测试是长上下文能力的快速体检方法。做法很简单在一个长上下文里随机插入一段关键信息然后构造一个只有读过这段信息才能回答的问题。needle The flag value is OCTOLONG_NEEDLE_7. haystack build_training_text(sample) position len(haystack) // 2 haystack_with_needle haystack[:position] needle haystack[position:] prompt f{haystack_with_needle}\n\nQuestion: What is the flag value?\nAnswer:如果是代码场景针尖信息可以替换成函数签名、配置项或类型定义。例如在长代码上下文中插入一行REPO_FLAG_42 cross_repo_pass然后问模型REPO_FLAG_42是多少。评测时要注意针尖位置要有覆盖不能只放中间位置同一个问题要在不同上下文长度下测试上下文长度要覆盖训练长度和推理目标长度避免针尖信息出现在训练数据里。5.3 代码仓库级评测需要注意数据隔离跨仓库训练数据的最大风险是“训练集包含评测仓库快照”。如果一个仓库的 main 分支代码参与了训练评测时再拿同仓库其他 commit 或任务做评估很容易高估效果。建议按时间切分训练时只使用某个时间点之前的 commit 快照评测任务使用后续 commit 或专门构造的跨文件任务。这样能更接近真实场景也避免数据泄漏导致的虚高分数。消融实验至少应该包含四组完整方法跨仓库代码上下文 mid-training去掉跨仓库结构只用单文件代码语料做 mid-training去掉 mid-training直接进入下游任务微调改变训练窗口相同数据量下对比 8K、16K、32K。这样能回答三个问题跨仓库结构是否带来增益mid-training 是否比直接 SFT 更合适训练窗口大小对效果的影响有多大。6. 复现 OctoLong 类方法时的高频问题排查6.1 训练时显存溢出现象启动训练后很快报CUDA out of memory程序直接退出。排查顺序将per_device_train_batch_size降到 1确认gradient_checkpointing已开启确认attn_implementation是flash_attention_2或等价实现检查模型是否加载到了 GPU 0优化器状态是否被 ZeRO 切分如果依然 OOM把block_size减半先跑通流程再逐步加长。错误日志里如果出现torch.OutOfMemoryError不要先堆 GPU先看显存到底被谁占满。可以用nvidia-smi观察显存分布也可以在训练脚本里打印每个 batch 的输入形状。6.2 训练 loss 下降但长上下文评测没提升现象训练阶段 loss 很稳定地下降但 needle 测试命中率没有变化跨文件问答也看不到提升。可能原因数据构造里的长距离依赖不足。模型能从最近上下文猜出下一个 token根本没有被迫看远处文件评测任务格式与训练任务格式不一致训练序列长度虽然标为 32K但实际大部分样本远短于 32K模型并没有真正见到长样本。检查方式# 统计训练数据长度分布 python - PY import json from transformers import AutoTokenizer tokenizer AutoTokenizer.from_pretrained(Qwen/Qwen2.5-Coder-7B) lengths [] with open(data/train.jsonl) as f: for line in f: sample json.loads(line) text build_training_text(sample) lengths.append(len(tokenizer.encode(text))) lengths.sort() print(fmin{lengths[0]}, median{lengths[len(lengths)//2]}, max{lengths[-1]}) PY如果中位数远低于目标长度说明样本长度分布有问题。应该优先增加多文件关联样本而不是直接调大模型窗口。6.3 mid-training 后出现灾难性遗忘现象训练结束后长上下文代码理解有所提升但普通代码补全、通用对话、指令跟随能力下降。原因通常有两类一是学习率过大模型在 mid-training 数据上过度拟合二是数据配比中通用语料占比太低。建议调整顺序将学习率降到原来的 1/3 或 1/5在训练集中增加 10%~20% 的通用代码和通用文本训练结束后做一次轻量 SFT 或 adapter 回滚恢复指令能力保存 mid-training checkpoint 时保留原始模型权重便于做权重平均。6.4 评测结果异常偏高或偏低评测分数异常偏高先怀疑数据泄漏。检查训练样本是否包含评测任务相关的 commit、文件或问题答案。最容易忽略的是仓库里的test/目录如果评测任务直接来自某个测试文件而这部分测试文件参与了训练效果就会失真。评测分数异常偏低则先检查上下文长度和位置编码。如果训练时 max length 是 16K评测时突然给模型 64K 输入位置编码外推不稳定会导致分数大幅下降。评测长度应该控制在训练长度附近或者提前做好外推配置。异常现象优先检查处理方向分数异常高训练集和评测集是否有重叠文件按 commit 时间切分、按路径过滤分数异常低评测长度和训练长度是否匹配调整测试长度或外推配置结果不稳定解码参数、评测样本顺序固定 seed、多次评测取均值7. 从实验到落地工程化建议和扩展方向7.1 建立 mid-training 的迭代流程mid-training 不是一次性实验而应该成为模型迭代流水线中的一个固定环节。一个可复用的流程如下代码仓库数据采集设定许可证过滤和敏感信息过滤构建文件依赖图和 commit 关联数据生成跨仓库上下文样本做长度分布统计配置训练脚本和监控指标在 1B 或 3B 小模型上验证数据管道在目标模型上做完整训练跑长上下文评测、通用能力评测、安全评测如果通过进入 SFT 或下游任务微调阶段。每一步都要有检查点。比如第 3 步完成后必须要看样本长度分布和文件数量分布第 5 步完成后要确认 loss 能正常下降第 7 步完成后要记录基线对比表而不是只看一两个指标。7.2 推理侧还需要配套优化即使模型已经被 mid-training 增强了长上下文能力部署时仍然不建议把整个仓库无脑塞进输入。实际开发中仓库可能包含几十万行代码全部放进上下文既不经济也容易导致注意力分散。比较好的做法是结合一个仓库级检索层先用符号索引或向量索引召回与当前任务相关的文件再把这部分文件放入上下文。这样 mid-training 模型负责“在相关长上下文里精读”检索层负责“从仓库里找相关文件”两者分工明确。推理侧还需要关注 KV cache 大小。如果系统同时服务多个用户长上下文的 KV cache 会占用大量显存可以考虑 Prefix Caching、token 级别的 prompt 压缩或分段缓存。模型上下文能力强并不等于部署时一定要用满服务成本也要纳入设计。7.3 可以继续延伸的方向mid-training 在代码长上下文上的思路可以继续向几个方向延伸。第一把 mid-training 和 agent 轨迹结合。代码智能体经常需要在上下文中维护“用户需求、当前文件、测试结果、历史修改”等信息这些信息天然跨文件、跨步骤符合 mid-training 的训练结构。第二把跨仓库数据和 RAG 的检索策略结合。训练数据里的文件关联关系可以作为检索排序的监督信号改进仓库级代码检索。第三在多语言混合仓库上做实验。一个仓库里可能同时有 Python 后端、TypeScript 前端、YAML 配置和 Markdown 文档跨仓库上下文如果覆盖这些异构文件能让模型学会在多种语言之间切换和联动。如果打算在自己的代码模型上开始尝试建议从一个小型仓库集和 8K 上下文开始先建立完整的数据-训练-评测闭环再逐步扩大仓库数量和上下文窗口。跑通闭环之后再判断到底是数据关联度、上下文长度还是模型规模才是限制长上下文能力的主要瓶颈。