Mechanistic Interpretability Nnsight API 参考

# nnsight API 参考 ## LanguageModel 用于包装具有干预能力的语言模型的主类。 ### 加载模型 python from nnsight import LanguageModel # 基础加载 model = LanguageModel("openai-community/gpt2", device_map="auto") # 更大的模型 model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto") # 自定义分词器设置 model = LanguageModel( "gpt2", device_map="auto", torch_dtype=torch.float16, ) ### 模型属性 python # 访问底层的 HuggingFace 模型 model._model # 访问分词器 model.tokenizer # 模型配置 model._model.config --- ## 追踪上下文 (Tracing Context) `trace()` 方法为延迟执行创建一个上下文。 ### 基础追踪 python with model.trace("Hello world") as tracer: # 操作被记录，不会立即执行 hidden = model.transformer.h[5].output[0].save() logits = model.output.save() # 上下文结束后，操作执行并可获取保存的值 print(hidden.shape) ### 追踪参数 python with model.trace( prompt, # 输入文本或 token remote=False, # 使用 NDIF 远程执行 validate=True, # 验证张量形状 scan=True, # 扫描形状信息 ) as tracer: ... ### 远程执行 python # 同样的代码可远程运行 with model.trace("Hello", remote=True) as tracer: hidden = model.transformer.h[5].output[0].save() --- ## 代理对象 (Proxy Objects) 在追踪上下文中，访问模块会返回代理对象。 ### 访问值 python with model.trace("Hello") as tracer: # 这些是代理对象 layer_output = model.transformer.h[5].output[0] attention = model.transformer.h[5].attn.output # 操作会创建新的代理 mean = layer_output.mean(dim=-1) normed = layer_output / layer_output.norm() ### 保存值 python with model.trace("Hello") as tracer: # 必须调用 .save() 才能在上下文之后访问 hidden = model.transformer.h[5].output[0].save() # 现在 hidden 包含实际的张量 print(hidden.shape) ### 修改值 python with model.trace("Hello") as tracer: # 原地修改 model.transformer.h[5].output[0][:] = 0 # 替换为计算出的值 model.transformer.h[5].output[0][:] = some_tensor # 算术修改 model.transformer.h[5].output[0][:] *= 0.5 model.transformer.h[5].output[0][:] += s