Instance attributes not accessible inside VLLM trace context

[Butanium]

Instance attributes not accessible inside VLLM trace context

Summary

Instance attributes on VLLM objects are inaccessible inside the trace context. This affects both direct instance attributes and properties that access self.*.

Environment

• nnsight version: 0.5.13
• vllm version: 0.13.0
• Python: 3.10

Minimal Reproduction

from nnsight.modeling.vllm import VLLM

class MyVLLM(VLLM):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._my_attr = "hello"
    
    @property
    def my_attr(self):
        return self._my_attr  # accesses self.*
    
    @property
    def my_attr_fixed(self):
        return "fixed_value"  # no self.* access

model = MyVLLM("gpt2", gpu_memory_utilization=0.1, dispatch=True)

# Both work outside trace
print(f"outside: {model.my_attr=}")        # hello
print(f"outside: {model.my_attr_fixed=}")  # fixed_value

with model.trace("Hello world") as tracer:
    print(f"inside: {model.my_attr_fixed=}")  # WORKS: fixed_value
    print(f"inside: {model.my_attr=}")        # FAILS: AttributeError

What Works vs What Fails

Attribute Type	Outside Trace	Inside Trace
Property returning literal	✅	✅
Property accessing self.*	✅	❌
Instance attributes (self.x)	✅	❌
Dynamic attributes (model.x = ...)	✅	❌

LanguageModel handles all cases correctly.

Expected Behavior

self inside properties should refer to the wrapper class, not the underlying model.

Actual Behavior

Inside the trace, self in properties resolves to the proxied underlying model, so any self.* access fails:

AttributeError: GPT2LMHeadModel(...) has no attribute _my_attr

Use Case

This blocks patterns like:

class MyVLLM(VLLM):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.layer_names = self._compute_layer_names()  # instance attr will fail
        self._lazy_attr = None
    
    @property
    def lazy_attr(self):  # property will fail because it accesses self._lazy_attr
        if self._lazy_attr is None:
            self._lazy_attr = self._compute_lazy_attr()
        return self._lazy_attr

    def _compute_lazy_attr(self):
        return "lazy_value"

[JadenFiotto-Kaufman]

@Butanium I would add this to your class:

def __getstate__(self):
    state = super().__getstate__()
    state["my_custom_attribute"] = self.my_custom_attribute
    ...
    return state

def __setstate__(self, state):
    super().__setstate__(state)
    self.my_custom_attribute = state["my_custom_attribute"]
    ...