OctoLong：Mid-Training 如何增强跨仓库代码长上下文建模

做代码生成、仓库级问答、长链路 Bug 排查时一个非常典型的瓶颈是模型在单个文件内的局部语义理解已经不错一旦需要跨文件、跨模块甚至跨仓库组合上下文输出质量就明显下滑。这个问题的根源往往不在模型结构本身而在于训练数据的形态以及训练阶段的安排方式。OctoLong 这个研究方向正是从这个角度切入通过 Mid-Training中期训练在跨仓库代码上下文上继续训练让模型的长上下文建模能力得到针对性增强。本文会围绕这条技术路线拆解它的核心概念、数据构建思路、训练方法并给出可落地的工程实操示例帮助你理解并复现类似方案。这篇文章适合以下几类读者正在做代码大模型训练、微调和评测的算法工程师。想给已有模型增加长上下文能力的开发者和研究者。对代码数据 pipeline 感兴趣想搞清楚“跨仓库上下文”到底怎么构造的技术人。做 Code Review 助手、仓库级问答、项目级代码生成的后端开发者。读完本文你将掌握OctoLong 要解决什么问题Mid-Training 在训练流程中的定位。跨仓库代码上下文的数据构造思路。如何用常见工具链搭建一条可运行的长上下文训练与评测链路。实际训练中的效率问题、评测方法和避坑清单。1. 背景与核心概念1.1 长上下文建模的价值在真实软件开发中一个功能的实现往往散落在多个文件和仓库中。比如一个 Java 后端接口通常包含 Controller、Service、Mapper、Entity、配置文件、数据库表定义而工具类或公共 SDK 可能来自另一个内部仓库。如果模型不能同时“看到”这些跨仓库的代码片段它就只能靠记忆补全准确性会受到很大影响。长上下文建模Long-Context Modeling要解决的正是这种“输入长度超过模型默认训练长度”时的信息利用问题。它的价值体现在几个典型场景仓库级代码补全与生成根据整个仓库的代码风格、依赖关系和相似实现来生成新代码。跨文件调试给出报错信息时模型需要理解多个文件之间的调用关系。项目文档生成从分散在多个模块的代码中总结出系统设计。代码 Review需要同时看主仓变更、依赖 SDK 的接口定义、历史修改记录。这些场景都要求模型能处理几千到几十万 token 的输入并且能够在长输入中准确定位关键信息。1.2 现有代码模型的短板目前很多代码大模型包括一些知名开源模型在实际评测中仍然有几个明显短板。第一有效上下文远低于宣传窗口。很多模型宣称支持 128K 甚至 200K 上下文但实际测试中当关键信息位于长文本中段时准确率会大幅下降。这在学术上常被称为“lost in the middle”问题。第二训练数据以单文件为主。通用代码预训练语料通常来自公开仓库的单个文件模型没有充分见过“多个仓库、多个文件拼成一个长序列”的样本形态。因此即使把窗口长度撑大模型也不知道该如何利用跨文件信息。第三短上下文和长上下文能力不一致。模型可能在 8K 以内表现良好但一旦超过训练长度注意力计算、位置编码、相对位置推断都会出现退化。OctoLong 这类工作想证明的核心观点是通过 Mid-Training专门用结构化的跨仓库代码上下文去训练模型可以显著改善上述短板而且不需要重新做完整预训练。1.3 OctoLong 的核心思路从标题可以看出OctoLong 的关键词是三个Mid-Training中期训练介于预训练和指令微调之间的一个训练阶段。Cross-Repository Code Contexts跨仓库代码上下文。Long-Context Modeling长上下文建模。把三者串起来OctoLong 的做法可以理解为在基础模型已经具备一定代码能力的条件下构建专门的数据集这些数据不是简单从单个文件里截取而是按照仓库依赖、符号引用、模块调用等关系把多个仓库中相关的代码片段组合成超长训练样本。然后在这个数据上继续训练让模型学会在长跨度上关联信息。这种做法的本质是把“长上下文能力”当成一个可以定向增强的技能而不是天然从预训练中长出来的能力。2. Mid-Training一个被低估的训练阶段2.1 从预训练到微调训练流程的四阶段当前大语言模型的主流训练流程可以分成四个阶段阶段目标数据特点典型成本预训练学习通用语言与知识海量、低质量、无标注极高Mid-Training / 继续预训练强化某一领域能力领域语料、中等规模中高指令微调学会遵循指令高质量指令数据低对齐RLHF/DPO符合偏好与安全要求偏好数据低Mid-Training 在很多开源模型里也被叫做 “domain-adaptive continued pretraining” 或“二次预训练”。它的位置正好在预训练和指令微调之间。2.2 Mid-Training 的定位与作用Mid-Training 要解决的问题是基础模型在通用语料上学习了大量知识但这些知识在特定领域的组织方式、术语体系和推理模式与通用场景差异很大。直接在通用模型上做指令微调效果往往有限因为模型根本没有见过足够多的领域长文本结构。举个例子。一个模型可能理解“函数”“调用”“异常”这些词但它不一定见过“某个项目的 service 层大量依赖另一个仓库的 common 模块且异常类型要统一处理”这种真实的跨仓库代码模式。通过 Mid-Training模型先把这种模式学进来之后再做指令微调任务表现会稳定很多。OctoLong 选择在代码领域做 Mid-Training而且专门使用跨仓库上下文数据。这样做有两个好处训练数据与推理场景一致推理时需要长输入训练时也喂给模型长输入。强化信息关联能力模型学会在长文本中维护多文件引用关系而不是只做局部建模。2.3 为什么代码场景特别适合 Mid-Training代码和自然语言有一个非常大的区别代码有严格的结构和明确的引用关系。一个函数是另一个文件里定义的这个关系是可以从语法分析、依赖图中精确抽取出来的。因此代码领域可以“主动构造”出高质量的长上下文训练样本而不是简单地把随机文本拼接在一起。自然语言领域也要做长上下文增强比如拼接多篇文档但段落之间的逻辑关系往往较弱。代码不同跨文件调用是强逻辑关系。模型如果能学到这种关系长上下文能力会有质的提升。所以OctoLong 选择代码场景实践 Mid-Training不只是因为代码数据好获取更重要的是代码本身提供了“可验证的上下文关联信号”。3. 跨仓库代码上下文数据与建模思想3.1 从单文件到跨仓库数据形态的跨越大多数代码预训练语料处理流程是这样的把每个文件看成一个独立文本切分成 token 序列加入训练。这种处理方式忽略了三个层面文件内函数之间的调用关系。同一仓库内文件之间的 import 关系。不同仓库之间通过依赖管理工具建立的引用关系。跨仓库代码上下文就是把第三个层面纳入训练样本。一个典型的训练样本可能长这样[仓库A] /payment-service/src/main/java/com/example/PaymentController.java [仓库B] /common-lib/src/main/java/com/example/Result.java [仓库B] /common-lib/src/main/java/com/example/ResultCode.java [仓库A] /payment-service/src/main/java/com/example/PaymentService.java这些文件并不是随机拼接而是因为PaymentController调用了PaymentService且返回类型使用了Result所以在同一个上下文窗口中出现。3.2 跨仓库依赖的构建方式要构建这种样本核心是构建“代码依赖图”。步骤可以拆解如下克隆或拉取目标仓库集合。解析每个文件的 import / require / include 语句。解析符号定义和使用关系函数定义、类定义、函数调用、类型引用。建立文件到文件的引用边。建立仓库到仓库的依赖边通过包名、模块名、命名空间。做反向依赖查询找出“谁用了谁”。这一套逻辑在 GitHub 上有很多现成工具支持比如 tree-sitter 可以做语言级语法解析部分语言可以通过编译数据库拿到更精确的依赖。3.3 上下文打包策略得到依赖关系后还需要把相关文件组装成训练样本。这里有两个关键设计点样本长度分布和样本格式。在样本长度上建议按混合比例构建让模型同时见到不同长度的样本。比如一部分样本控制在 4K 到 8K一部分在 8K 到 32K还有一部分超过 32K。这样模型既不会丢失短文本能力又能逐步适应长序列。在样本格式上需要设计清晰的文件边界标记。一种常见做法是使用特殊 token 标记文件路径和文件内容。下面是一个简化的样本格式示例repo namepayment-service file pathPaymentController.java ...代码... /file /repo repo namecommon-lib file pathResult.java ...代码... /file /repo这样的格式可以让模型在长上下文中区分出“哪个仓库、哪个文件、哪段代码”避免文件边界混乱。4. 环境准备与工具链4.1 环境要求实际的 OctoLong 训练规模取决于基座模型大小和 GPU 资源。本文以复现实验和学习验证为目的以 7B 到 13B 级别的代码模型为例推荐环境如下操作系统LinuxUbuntu 20.04 或 22.04GPU至少 4 张 24GB 显存显卡如 RTX 3090 / 4090 / A10G建议使用 A100 或 H800CPU16 核以上内存64GB 以上磁盘至少 200GB 剩余空间用于存放模型权重和数据集需要强调的是具体资源要结合模型大小、序列长度和训练策略调整。如果只是跑通验证流程也可以使用更小的 1B 级模型和小规模语料。4.2 工具清单本文的实战环节会用到以下工具它们都是当前生态中比较通用的选择Python 3.10PyTorch 2.xTransformersDatasetsAcceleratePEFT用于 LoRA 训练tree-sitter用于代码解析networkx用于依赖图构建vllm用于推理加速可选版本需要根据你的项目实际情况调整本文示例以常见环境为例重点演示配置思路。4.3 示例项目结构建议按下面的结构组织项目octolong-lab/ ├── configs/ │ └── train.yaml ├── data/ │ └── raw_repos/ # 存放克隆的仓库 │ └── processed/ # 存放构建好的训练样本 ├── scripts/ │ ├── build_dependency.py │ ├── build_context.py │ └── train_lora.py ├── src/ │ └── octolong/ │ ├── dependency_graph.py │ └── sample_builder.py └── output/ └── checkpoints/这样一个结构可以支持从数据构建到模型训练的全流程便于后续扩展。5. 实战构建跨仓库训练样本5.1 解析仓库结构我们先用 tree-sitter 做 Python 代码的 import 解析。这个步骤的目标是从每个源码文件中提取出它导入了哪些模块。下面是一个最小示例代码路径为scripts/build_dependency.pyimport os import glob from tree_sitter import Language, Parser # 这里假设你已经编译好了 python.so 语言文件 PY_LANGUAGE Language(build/python.so, python) parser Parser(PY_LANGUAGE) def extract_imports(file_path): with open(file_path, r, encodingutf-8, errorsignore) as f: source f.read().encode(utf-8) tree parser.parse(source) imports [] def walk(node): if node.type import_statement: # 拿到 import 后面的模块名 text node.text.decode(utf-8) imports.append(text) elif node.type import_from_statement: text node.text.decode(utf-8) imports.append(text) for child in node.children: walk(child) walk(tree.root_node) return imports def scan_repository(repo_root): all_imports {} for file_path in glob.glob(os.path.join(repo_root, **, *.py), recursiveTrue): rel_path os.path.relpath(file_path, repo_root) try: imports extract_imports(file_path) all_imports[rel_path] imports except Exception as e: print(f解析失败: {file_path}, 错误: {e}) return all_imports if __name__ __main__: repo data/raw_repos/example_project result scan_repository(repo) for file, imports in result.items(): print(file, -, imports)这段代码演示了最基础的 import 提取。实际使用中还需要对相对导入、别名导入做归一化处理并且对 Java、TypeScript、Go 等语言分别配置 tree-sitter 语法。5.2 提取依赖关系得到所有文件的 import 之后下一步是把这些 import 映射到具体文件建立“文件到文件”的依赖图。import networkx as nx import os def build_dependency_graph(repo_root, import_map): graph nx.DiGraph() # 先把所有文件作为节点加入 for rel_path in import_map.keys(): graph.add_node(rel_path) # 建立模块名到文件路径的映射 module_to_file {} for rel_path in import_map.keys(): module_name rel_path.replace(os.sep, .).replace(.py, ) module_to_file[module_name] rel_path # 解析 import 关系 for rel_path, imports in import_map.items(): for imp in imports: # 简单提取 import 后的第一个模块名 # 真实场景需要处理 from x import y, import x.y.z 等情况 for candidate, target_file in module_to_file.items(): if candidate in imp and candidate ! rel_path.replace(os.sep, .).replace(.py, ): graph.add_edge(rel_path, target_file, typeimport) return graph # 使用示例 # graph build_dependency_graph(data/raw_repos/example_project, import_map) # print(nx.info(graph))这个阶段的产出是一个有向图节点是文件边是依赖关系。之后可以在图上做 BFS 或 DFS找出与某个文件强相关的文件集合。5.3 组装跨仓库上下文样本假设我们需要为PaymentService.java构造一个训练样本可以按照依赖关系找到它依赖的公共类文件再按顺序拼接到一起。这里给出一个简化版样本构建脚本的核心逻辑路径为scripts/build_context.pydef build_training_sample(anchor_file, graph, repo_roots): anchor_file: 主文件路径 graph: networkx 依赖图 repo_roots: 仓库根路径到文件路径的映射 # 使用 BFS 找与 anchor_file 相关的文件限制遍历深度为 2 related_files [] for depth, nodes in enumerate(nx.bfs_layers(graph, anchor_file)): if depth 2: break related_files.extend(nodes) # 组装上下文 sample_parts [] for file_path in related_files: repo_name get_repo_name(file_path, repo_roots) rel_path get_rel_path_in_repo(file_path, repo_roots) with open(file_path, r, encodingutf-8, errorsignore) as f: code f.read() sample_parts.append(frepo name\{repo_name}\\n) sample_parts.append(ffile path\{rel_path}\\n) sample_parts.append(code) sample_parts.append(\n/file\n/repo\n) return .join(sample_parts)这里需要注意BFS 遍历深度需要控制否则样本会过大。最终样本长度需要通过 tokenizer 做截断或拼接控制在目标上下文窗口内。需要做样本去重避免模型在训练集和评测集之间发生数据泄漏。6. 实战Long-Context 模型训练6.1 模型与长度扩展选择对于已有的开源代码模型直接做长上下文训练通常需要解决位置编码问题。常见处理方式有两种如果模型本身支持 RoPE可以通过调整 RoPE 的 base 频率来扩展窗口例如把 base 从 10000 改成 500000 或 1000000。如果模型支持 ALiBi 或相对位置编码本身对长度外推比较友好扩展成本更低。在 Transformers 中加载模型时可以通过rope_scaling参数做长度扩展。下面是一个基于 Llama 架构的示例from transformers import AutoModelForCausalLM, AutoTokenizer model_path your-code-model-path tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypeauto, device_mapauto, rope_scaling{type: linear, factor: 2.0}, )这里把 RoPE 的缩放因子设为 2.0可以支持大约 2 倍于原始窗口的长度。实际生产环境建议对缩放方式做实验对比。6.2 LoRA 训练脚本Mid-Training 如果从头训练全部参数成本非常高。更常见的是先用 LoRA 或 QLoRA 做参数高效微调验证数据效果。下面给出一个可运行的训练脚本核心片段路径为scripts/train_lora.pyfrom transformers import ( AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq, ) from datasets import load_dataset from peft import LoraConfig, get_peft_model model_path your-code-model-path data_path data/processed/train.jsonl model AutoModelForCausalLM.from_pretrained( model_path, torch_dtypebfloat16, device_mapauto, rope_scaling{type: linear, factor: 2.0}, ) tokenizer AutoTokenizer.from_pretrained(model_path) tokenizer.pad_token tokenizer.eos_token # LoRA 配置 lora_config LoraConfig( r32, lora_alpha64, target_modules[q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj], lora_dropout0.05, biasnone, task_typeCAUSAL_LM, ) model get_peft_model(model, lora_config) def preprocess_function(examples): texts examples[text] model_inputs tokenizer( texts, max_length32768, truncationTrue, paddingFalse, return_tensorsNone, ) model_inputs[labels] model_inputs[input_ids].copy() return model_inputs dataset load_dataset(json, data_filesdata_path, splittrain) tokenized_dataset dataset.map(preprocess_function, batchedTrue, remove_columnsdataset.column_names) training_args TrainingArguments( output_diroutput/checkpoints, per_device_train_batch_size1, gradient_accumulation_steps8, learning_rate2e-4, warmup_steps100, logging_steps10, save_steps500, num_train_epochs1, bf16True, gradient_checkpointingTrue, optimadamw_torch, ) trainer Trainer( modelmodel, argstraining_args, train_datasettokenized_dataset, data_collatorDataCollatorForSeq2Seq(tokenizer, paddingTrue), ) trainer.train()这段代码有几个关键点max_length32768是训练时的目标序列长度需要根据显存调整。gradient_checkpointingTrue能显著降低显存占用但会带来一点训练速度损失。per_device_train_batch_size1长序列训练下通常只能做到这种级别要靠梯度累积来模拟大的 batch size。LoRA 的target_modules需要根据模型架构调整不同模型名称可能不同。6.3 训练时的效率与稳定性长序列训练最直接的挑战是显存。即使 batch size 为 132K 长度、7B 模型的激活值也可能超过单卡 40GB。这里有几个工程手段使用 FlashAttention-2 替代标准 attention显存占用和速度都会有明显改善。使用序列打包sequence packing把多个短样本拼成一个长序列减少 padding 浪费。使用 DeepSpeed ZeRO-2 或 ZeRO-3 做参数分片。数据加载时使用 memory-mapped dataset避免把整个数据集读入内存。训练过程中要重点观察两个指标loss 是否稳定下降以及梯度范数是否出现异常放大。如果 loss 突然升高优先检查是否有样本长度分布不合理的问题。6.4 评测验证Mid-Training 的效果不能只看训练 loss还需要做针对性的评测。长上下文代码评测可以分为三类评测类型说明示例通用长文本评测检验长文本理解基本能力LongBench、L-Eval代码推理评测检验跨文件理解和修复能力CrossFileBench、RepoBench代码生成评测检验长上下文下的生成质量HumanEval、MBPP如果只做最简单的验证可以构造一组“跨文件问答”测试集。比如给定两个文件的内容问模型第二个文件中的某个函数被谁调用或者某段逻辑的返回类型是什么。这类问题对模型的跨文件关联能力非常敏感。下面给一个简单的评测脚本思路from transformers import AutoModelForCausalLM, AutoTokenizer model_path output/checkpoints/your-lora-checkpoint tokenizer AutoTokenizer.from_pretrained(model_path) model AutoModelForCausalLM.from_pretrained(model_path, device_mapauto) prompt repo namepayment-service file pathPaymentController.java ...代码... /file /repo repo namecommon-lib file pathResult.java ...代码... /file /repo 请回答PaymentController 中调用 PaymentService.createOrder 时返回的 Result 对象中code 字段在什么情况下为 500 inputs tokenizer(prompt, return_tensorspt).to(cuda) output model.generate(**inputs, max_new_tokens256) print(tokenizer.decode(output[0], skip_special_tokensTrue))这种评测方式不够系统但能很直观地看出模型在跨仓库上下文上的表现差异适合在训练过程中做快速回归。7. 常见问题与排查思路在复现和训练过程中很容易遇到各类问题。下面整理一张排查表。问题现象常见原因解决思路训练时 OOM序列过长、batch size 太大、attention 显存过高开启 gradient checkpointing使用 FlashAttention-2降低 max_length增大梯度累积步数loss 快速下降后震荡学习率过高、样本长度分布不均匀降低学习率检查数据集中超长样本占比做长度分层采样位置编码外推后效果差RoPE 缩放方式不适配当前模型用 NTK 或 YaRN 替代 linear scaling做小规模消融实验评测时模型忽略中间内容模型没有真正学到长距离关联增加长样本占比强化样本中的文件边界信息数据集样本重复率高依赖图构建过于简单同一批文件反复组合增加随机采样策略限制每个文件被选为 anchor 的次数跨仓库文件路径解析错误仓库克隆结构不一致路径映射错误统一仓库目录结构建立 repo 到根目录的映射表训练速度过慢attention 计算量过大数据加载成为瓶颈使用 FlashAttention-2开启多进程数据加载使用预分词缓存最常被忽视的问题是数据质量。很多情况下模型表现不佳不是训练参数不对而是训练样本本身没有体现真正的跨仓库依赖关系只是把多个文件机械地拼在一起。这种样本对模型来说就像一段乱序文本学不到有效信息。8. 最佳实践与工程建议8.1 数据质量优先先做小规模验证跨仓库样本的构建质量直接决定 Mid-Training 的效果。建议先用 500 到 2000 条样本做小规模实验对比训练前后模型在跨文件任务上的表现。只有数据构建逻辑验证通过后再扩大规模。验证时可以人工检查样本同一个样本中的文件之间是否有真实的调用关系文件边界标记是否清晰样本是否覆盖了不同层级的长上下文4K、8K、16K、32K是否存在训练集与真实任务分布不一致的问题8.2 训练策略窗口长度逐步拉升不要在训练一开始就使用 32K 的长序列。比较稳定的做法是分阶段第一阶段以 8K 为主让模型适应跨文件样本格式。第二阶段混合 8K 和 16K逐步加入更长样本。第三阶段加入 32K 以上的样本巩固极端长度下的表现。这个思路和人类学习类似先掌握结构再应对更长的信息跨度。同时每阶段训练结束都要做一次评测回归防止灾难性遗忘。8.3 评测与回归建立长上下文基线训练前后要固定同一套评测集。建议包含三类数据单文件代码生成任务用来监测短上下文能力是否退化。跨文件代码理解任务用来验证 Mid-Training 的核心收益。超长输入压力测试检测模型在长文本中定位关键信息的能力。每次实验只改一个变量比如只改数据构建方式或只改训练长度。这样才能定位出真正影响效果的因素。8.4 生产落地的边界跨仓库上下文训练开销较大落地上要关注几个问题模型推理时的 KV Cache 显存长输入会有很大的 KV Cache生产环境要配合 vllm 等推理框架管理显存。数据更新频率仓库代码变化快训练数据不能一成不变需要设计定期的数据重建流程。权限与合规使用真实业务仓库数据做训练时必须遵守代码保密协议确保有合法授权。涉及敏感代码的场景建议使用脱敏后的代码片段进行训练。回滚机制Mid-Training 后的模型如果在下游任务上出现退化要能快速回退到训练前版本。建议在发布流程中保留原模型和中间 checkpoints。8.5 可复现性为了让实验可复现建议在项目里固定三样东西依赖的 commit hash 或版本号。数据构建脚本的版本。训练配置文件的完整参数。每次数据更新或训练配置变化都打一个版本号。长期实验多了之后这能帮你快速定位是哪一次变更导致的效果波动。9. 总结OctoLong 给我们展示了一条清晰的路线在通用预训练和指令微调之间增加一个针对性的 Mid-Training 阶段用跨仓库代码上下文数据来增强模型的长上下文建模能力。这个思路的核心不在于“把窗口撑大”而在于让模型真正学会利用分布在不同文件、不同仓库中的信息。如果要在自己的项目中复现类似方案优先级是这样先做好数据构建真实的依赖图组装有逻辑关系的跨仓库样本。再选好训练策略用 LoRA 先验证再决定是否做全参数训练。最后完善评测设计能体现跨文件能力的评测集避免只盯着训练 loss。实际落地中最值得花时间的不是训练本身而是数据构造和评测设计。训练代码和参数都有成熟模板但“哪些文件应该出现在同一个上下文里”这个问题需要结合真实的代码依赖关系去回答。如果你正打算给代码模型做长上下文增强建议先克隆几个真实项目跑一遍依赖构建流程看看能产生多少有价值的跨仓库样本。这个实验本身成本不高但对后续训练方案的设计会很有帮助。如果你在实际操作中遇到其他问题欢迎在评论区留言交流。也可以把本文收藏起来在做跨仓库上下文训练时随时翻阅。