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 在很多开源模型里也被叫做 “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 跨仓库依赖的构建方式

要构建这种样本，核心是构建“代码依赖图”。步骤可以拆解如下：

1. 克隆或拉取目标仓库集合。
2. 解析每个文件的 import / require / include 语句。
3. 解析符号定义和使用关系（函数定义、类定义、函数调用、类型引用）。
4. 建立文件到文件的引用边。
5. 建立仓库到仓库的依赖边（通过包名、模块名、命名空间）。
6. 做反向依赖查询，找出“谁用了谁”。

这一套逻辑在 GitHub 上有很多现成工具支持，比如 tree-sitter 可以做语言级语法解析，部分语言可以通过编译数据库拿到更精确的依赖。

3.3 上下文打包策略

得到依赖关系后，还需要把相关文件组装成训练样本。这里有两个关键设计点：样本长度分布和样本格式。

在样本长度上，建议按混合比例构建，让模型同时见到不同长度的样本。比如一部分样本控制在 4K 到 8K，一部分在 8K 到 32K，还有一部分超过 32K。这样模型既不会丢失短文本能力，又能逐步适应长序列。

在样本格式上，需要设计清晰的文件边界标记。一种常见做法是使用特殊 token 标记文件路径和文件内容。下面是一个简化的样本格式示例：

<repo name="payment-service"> <file path="PaymentController.java"> ...代码... </file> </repo> <repo name="common-lib"> <file path="Result.java"> ...代码... </file> </repo>

这样的格式可以让模型在长上下文中区分出“哪个仓库、哪个文件、哪段代码”，避免文件边界混乱。

4. 环境准备与工具链

4.1 环境要求

实际的 OctoLong 训练规模取决于基座模型大小和 GPU 资源。本文以复现实验和学习验证为目的，以 7B 到 13B 级别的代码模型为例，推荐环境如下：

• 操作系统：Linux（Ubuntu 20.04 或 22.04）
• GPU：至少 4 张 24GB 显存显卡（如 RTX 3090 / 4090 / A10G），建议使用 A100 或 H800
• CPU：16 核以上
• 内存：64GB 以上
• 磁盘：至少 200GB 剩余空间（用于存放模型权重和数据集）

需要强调的是，具体资源要结合模型大小、序列长度和训练策略调整。如果只是跑通验证流程，也可以使用更小的 1B 级模型和小规模语料。

4.2 工具清单

本文的实战环节会用到以下工具，它们都是当前生态中比较通用的选择：

• Python 3.10+
• PyTorch 2.x
• Transformers
• Datasets
• Accelerate
• PEFT（用于 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.py：

import 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", encoding="utf-8", errors="ignore") 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"), recursive=True): 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, type="import") 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.py：

def 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", encoding="utf-8", errors="ignore") as f: code = f.read() sample_parts.append(f"<repo name=\"{repo_name}\">\n") sample_parts.append(f"<file 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_dtype="auto", device_map="auto", rope_scaling={"type": "linear", "factor": 2.0}, )

这里把 RoPE 的缩放因子设为 2.0，可以支持大约 2 倍于原始窗口的长度。实际生产环境建议对缩放方式做实验对比。

6.2 LoRA 训练脚本

Mid-Training 如果从头训练全部参数，成本非常高。更常见的是先用 LoRA 或 QLoRA 做参数高效微调，验证数据效果。下面给出一个可运行的训练脚本核心片段，路径为scripts/train_lora.py：

from 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_dtype="bfloat16", device_map="auto", rope_scaling={"type": "linear", "factor": 2.0}, ) tokenizer = AutoTokenizer.from_pretrained(model_path) tokenizer.pad_token = tokenizer.eos_token # LoRA 配置 lora_config = LoraConfig( r=32, lora_alpha=64, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM", ) model = get_peft_model(model, lora_config) def preprocess_function(examples): texts = examples["text"] model_inputs = tokenizer( texts, max_length=32768, truncation=True, padding=False, return_tensors=None, ) model_inputs["labels"] = model_inputs["input_ids"].copy() return model_inputs dataset = load_dataset("json", data_files=data_path, split="train") tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names) training_args = TrainingArguments( output_dir="output/checkpoints", per_device_train_batch_size=1, gradient_accumulation_steps=8, learning_rate=2e-4, warmup_steps=100, logging_steps=10, save_steps=500, num_train_epochs=1, bf16=True, gradient_checkpointing=True, optim="adamw_torch", ) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_dataset, data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True), ) trainer.train()

这段代码有几个关键点：

• max_length=32768是训练时的目标序列长度，需要根据显存调整。
• gradient_checkpointing=True能显著降低显存占用，但会带来一点训练速度损失。
• per_device_train_batch_size=1长序列训练下通常只能做到这种级别，要靠梯度累积来模拟大的 batch size。
• LoRA 的target_modules需要根据模型架构调整，不同模型名称可能不同。

6.3 训练时的效率与稳定性

长序列训练最直接的挑战是显存。即使 batch size 为 1，32K 长度、7B 模型的激活值也可能超过单卡 40GB。这里有几个工程手段：

• 使用 FlashAttention-2 替代标准 attention，显存占用和速度都会有明显改善。
• 使用序列打包（sequence packing），把多个短样本拼成一个长序列，减少 padding 浪费。
• 使用 DeepSpeed ZeRO-2 或 ZeRO-3 做参数分片。
• 数据加载时使用 memory-mapped dataset，避免把整个数据集读入内存。

训练过程中要重点观察两个指标：loss 是否稳定下降，以及梯度范数是否出现异常放大。如果 loss 突然升高，优先检查是否有样本长度分布不合理的问题。

6.4 评测验证

Mid-Training 的效果不能只看训练 loss，还需要做针对性的评测。长上下文代码评测可以分为三类：

如果只做最简单的验证，可以构造一组“跨文件问答”测试集。比如给定两个文件的内容，问模型第二个文件中的某个函数被谁调用，或者某段逻辑的返回类型是什么。这类问题对模型的跨文件关联能力非常敏感。

下面给一个简单的评测脚本思路：

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_map="auto") prompt = """<repo name="payment-service"> <file path="PaymentController.java"> ...代码... </file> </repo> <repo name="common-lib"> <file path="Result.java"> ...代码... </file> </repo> 请回答：PaymentController 中调用 PaymentService.createOrder 时，返回的 Result 对象中，code 字段在什么情况下为 500？ """ inputs = tokenizer(prompt, return_tensors="pt").to("cuda") output = model.generate(**inputs, max_new_tokens=256) print(tokenizer.decode(output[0], skip_special_tokens=True))

这种评测方式不够系统，但能很直观地看出模型在跨仓库上下文上的表现差异，适合在训练过程中做快速回归。

7. 常见问题与排查思路

在复现和训练过程中，很容易遇到各类问题。下面整理一张排查表。

最常被忽视的问题是数据质量。很多情况下模型表现不佳，不是训练参数不对，而是训练样本本身没有体现真正的跨仓库依赖关系，只是把多个文件机械地拼在一起。这种样本对模型来说就像一段乱序文本，学不到有效信息。

8. 最佳实践与工程建议

8.1 数据质量优先：先做小规模验证

跨仓库样本的构建质量，直接决定 Mid-Training 的效果。建议先用 500 到 2000 条样本做小规模实验，对比训练前后模型在跨文件任务上的表现。只有数据构建逻辑验证通过后，再扩大规模。

验证时可以人工检查样本：

• 同一个样本中的文件之间是否有真实的调用关系？
• 文件边界标记是否清晰？
• 样本是否覆盖了不同层级的长上下文（4K、8K、16K、32K）？
• 是否存在训练集与真实任务分布不一致的问题？

8.2 训练策略：窗口长度逐步拉升

不要在训练一开始就使用 32K 的长序列。比较稳定的做法是分阶段：

• 第一阶段：以 8K 为主，让模型适应跨文件样本格式。
• 第二阶段：混合 8K 和 16K，逐步加入更长样本。
• 第三阶段：加入 32K 以上的样本，巩固极端长度下的表现。

这个思路和人类学习类似：先掌握结构，再应对更长的信息跨度。同时，每阶段训练结束都要做一次评测回归，防止灾难性遗忘。

8.3 评测与回归：建立长上下文基线

训练前后要固定同一套评测集。建议包含三类数据：

1. 单文件代码生成任务，用来监测短上下文能力是否退化。
2. 跨文件代码理解任务，用来验证 Mid-Training 的核心收益。
3. 超长输入压力测试，检测模型在长文本中定位关键信息的能力。

每次实验只改一个变量，比如只改数据构建方式，或只改训练长度。这样才能定位出真正影响效果的因素。

8.4 生产落地的边界

跨仓库上下文训练开销较大，落地上要关注几个问题：

• 模型推理时的 KV Cache 显存：长输入会有很大的 KV Cache，生产环境要配合 vllm 等推理框架管理显存。
• 数据更新频率：仓库代码变化快，训练数据不能一成不变，需要设计定期的数据重建流程。
• 权限与合规：使用真实业务仓库数据做训练时，必须遵守代码保密协议，确保有合法授权。涉及敏感代码的场景，建议使用脱敏后的代码片段进行训练。
• 回滚机制：Mid-Training 后的模型如果在下游任务上出现退化，要能快速回退到训练前版本。建议在发布流程中保留原模型和中间 checkpoints。

8.5 可复现性

为了让实验可复现，建议在项目里固定三样东西：

• 依赖的 commit hash 或版本号。
• 数据构建脚本的版本。
• 训练配置文件的完整参数。

每次数据更新或训练配置变化，都打一个版本号。长期实验多了之后，这能帮你快速定位是哪一次变更导致的效果波动。

9. 总结

OctoLong 给我们展示了一条清晰的路线：在通用预训练和指令微调之间，增加一个针对性的 Mid-Training 阶段，用跨仓库代码上下文数据来增强模型的长上下文建模能力。这个思路的核心不在于“把窗口撑大”，而在于让模型真正学会利用分布在不同文件、不同仓库中的信息。

如果要在自己的项目中复现类似方案，优先级是这样：

• 先做好数据：构建真实的依赖图，组装有逻辑关系的跨仓库样本。
• 再选好训练策略：用 LoRA 先验证，再决定是否做全参数训练。
• 最后完善评测：设计能体现跨文件能力的评测集，避免只盯着训练 loss。

实际落地中最值得花时间的不是训练本身，而是数据构造和评测设计。训练代码和参数都有成熟模板，但“哪些文件应该出现在同一个上下文里”，这个问题需要结合真实的代码依赖关系去回答。

如果你正打算给代码模型做长上下文增强，建议先克隆几个真实项目，跑一遍依赖构建流程，看看能产生多少有价值的跨仓库样本。这个实验本身成本不高，但对后续训练方案的设计会很有帮助。

如果你在实际操作中遇到其他问题，欢迎在评论区留言交流。也可以把本文收藏起来，在做跨仓库上下文训练时随时翻阅。