stanfordnlp / pyvene

Stanford NLP Python Library for Understanding and Improving PyTorch Models via Interventions
http://pyvene.ai
Apache License 2.0
545 stars 46 forks source link

[P0] First infrastructure change to support NDIF backend #167

Closed frankaging closed 3 days ago

frankaging commented 4 days ago

Description

This is only the first code change to support NDIF infra in pyvene. Multiple changes are needed to fully support NDIF backend in pyvene. Here, we will describe the overall goal of these changes.

NDIF project aims to provide a remote service for doing intervened model inference so that users do not have to host their own models. With these change, users can wrap NDIF model as an intervenable model and potentially run the intervened model inference remotely by leveraging on the NDIF infrastature.

We will provide two ways of model wrapping:

(recommended way)

from nnsight import LanguageModel
gpt2_ndif = LanguageModel('openai-community/gpt2', device_map='cpu')

# create with dict-based config
pv_config = pv.IntervenableConfig({
  "component": "transformer.h[0].attn.output",
  "intervention": pv.VanillaIntervention()})

# model builder will automatically recognize ndif model
pv_gpt2_ndif = pv.build_intervenable_model(
  pv_config, model=gpt2_ndif)

or

pv_gpt2_ndif = pv.IntervenableNdifModel(
  pv_config, model=gpt2_ndif)

All existing usages of pyvene should all be supported (pending on future changes). This change has done all the heavy-liftings and code refactory to support long-term collaborations between pyvene and NDIF.

Testing Done

Notebook testing is done for basic usages.

Checklist: