more general setattribute

[t-vi]

Consult the wrappers in order to build return value. This gives us a more general setattribute, e.g. for HF static KVCache.

In pseudocode

# best implemented after Polenta e Osei and Coffee

def setattr_lookaside(obj, name, attr):
    if obj is an input:
        make an anyproxy that gives obj a name
        record proxy in provenance
        # but don't replace the obj with the proxy
   do previous setattr (including recording modifications as needed)

# get build an object from a wrapped value
def build_from_wrap_value(obj):
   # this is a sibling of pretty print in the sense that it considers types
   if we have a proxy (from setattr above or otherwise):
        return existing proxy
   if it is an object of a complex datatype (currently: dataclass):
        return prims.python_dataclass_new(cls, **{k: build_from_wrapped_value(obj.attribute_wrappers[k]) for k in obj.attributes})
   if it is an object of a "printable" type:
        return unwrapped obj
   if printable collection:  # dict, list, tuple, OrderedDict, ...
        return collection(build_from_wrap_value(wrapped item) for wrapped item in wrapped collection members)
   raise NotImplementedError

def run_jit(fn, *args, **kwargs):
    ....
    interpreted_fn = interpreter.interpret(fn, .unwrap_result=False)
    wrapped_result = interpeted_fn(*args, **kwargs)
    ...
    with tracectx(epilogue_trc):
          prims.python_recturn(build_from_wrapped_value(wrapped_result)) 
    ...

So for example for the "existing" HF Static Cache, we have the following epilogue, note how the static cache (provided as an input) is inserted into the model output rather than trying to construct a new one (which was what it tried to do before).

import torch
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def epilogue(p0, logits):
  # p0: "<class 'transformers.cache_utils.StaticCache'>"
  # logits: "cuda:0 bf16[1, 6, 128256]"
  p1 = transformers_modeling_outputs_CausalLMOutputWithPast(loss=None, logits=logits, past_key_values=p0, hidden_states=None, attentions=None)
  return p1

[lantiga]

Looks great! Let's add unit tests for dataclass and setattr in a follow-up PR (to avoid conflicts on the current PR)

[lantiga]

We should have a docstring here

[lantiga]

Should it be kwargs