Transformer Debugger开发者指南：如何扩展和自定义功能

Transformer Debugger（TDB）是由OpenAI超级对齐团队开发的强大工具，旨在支持对小语言模型特定行为的深入调查。本指南将为您详细介绍如何扩展和自定义Transformer Debugger的功能，让您能够更好地适应特定的研究需求。

🔧 项目架构概览

Transformer Debugger采用前后端分离架构，主要由三个核心组件构成：

• 激活服务器（Activation Server）：位于neuron_explainer/activation_server/的后端服务，负责从模型推理中获取激活数据
• 神经元查看器（Neuron Viewer）：位于neuron_viewer/的React前端应用
• 模型库（Models）：位于neuron_explainer/models/的推理库

🚀 快速入门：启动开发环境

后端服务器配置

要启动GPT-2小模型的激活服务器，运行以下命令：

python neuron_explainer/activation_server/main.py --model_name gpt2-small --port 8000

前端应用启动

进入前端目录并启动开发服务器：

cd neuron_viewer
npm install
npm start

访问http://localhost:1234即可开始使用Transformer Debugger。

🎯 核心扩展机制详解

1. 模拟器系统架构

Transformer Debugger提供了灵活的模拟器系统，位于neuron_explainer/explanations/simulator.py，包含三种主要模拟器类型：

• ExplanationNeuronSimulator：一次性模拟所有令牌激活
• ExplanationTokenByTokenSimulator：逐个令牌模拟激活
• LogprobFreeExplanationTokenSimulator：不依赖对数概率的模拟器

2. 解释器自定义

在neuron_explainer/explanations/explainer.py中，您可以自定义神经元行为的解释生成：

class TokenActivationPairExplainer(NeuronExplainer):
    def __init__(
        self,
        model_name: str,
        prompt_format: PromptFormat = PromptFormat.CHAT_MESSAGES,
        context_size: ContextSize = ContextSize.FOUR_K,
        few_shot_example_set: FewShotExampleSet = FewShotExampleSet.ORIGINAL,
        repeat_non_zero_activations: bool = False,
    ):
        # 自定义初始化逻辑

3. 激活数据处理

激活数据处理的核心理念在neuron_explainer/activations/derived_scalars/目录中实现，支持多种派生标量类型。

📊 高级自定义技巧

1. 自定义派生标量

在neuron_explainer/activations/derived_scalars/目录中，您可以创建新的标量派生器：

from neuron_explainer.activations.derived_scalars.scalar_deriver import ScalarDeriver

class CustomScalarDeriver(ScalarDeriver):
    def __init__(self, ...):
        # 实现自定义的激活数据处理逻辑

2. 前端组件扩展

前端组件位于neuron_viewer/src/TransformerDebugger/目录，您可以根据需要添加新的可视化组件。

3. 自动编码器集成

要启用自动编码器功能，启动服务器时指定相应的自动编码器名称：

python neuron_explainer/activation_server/main.py --model_name gpt2-small --port 8000 --mlp_autoencoder_name ae-resid-delta-mlp-v4

🔍 调试与优化技巧

1. CUDA内存调试

启用CUDA内存调试功能：

python neuron_explainer/activation_server/main.py --cuda_memory_debugging True

2. 客户端库生成

修改后端代码后，需要重新生成TypeScript客户端库：

# 启动无模型激活服务器
python neuron_explainer/activation_server/main.py --run_model False --port 8000

# 在另一个终端中生成客户端
cd neuron_viewer
npm run generate-client

💡 最佳实践建议

1. 代码格式检查：运行npm run check-code-format确保代码规范
2. 类型验证：使用npm run check-type-warnings检查类型问题
3. 构建验证：修改后运行npm run build确认功能正常

🎉 开始您的自定义之旅

通过本指南，您已经掌握了Transformer Debugger的核心扩展机制。现在可以：

• 创建自定义的模拟器来适应特定的研究需求
• 扩展前端组件以提供更好的可视化体验
• 集成新的模型架构和自动编码器

记住，Transformer Debugger的强大之处在于其灵活的可扩展性。根据您的具体需求，大胆地进行定制和优化吧！