Adapter code is undocumented and over-specialized to Intrinsics.

[nrfulton]

This issue might be capturing my own misconcptions about the current state of the code-base.

I'm opening the issue to get feedback from @jakelorocco and @frreiss. It's possible that I am misunderstanding things and all that's needed is better documentation. If that's the case, then we can probably iterate on this issue description to explain how things work (or should work).

If I state anything false or misleading in this issue, please quote and explain what I have wrong.

The Immediate Problem

While fixing issues related to our m tune command (see PR #422 ), I found that none of our custom alora stuff has worked since the intrinsics change.

(These examples aren't in the code-base's test suite because they require training and loading a custom model; we should find a way to mock that someone in the future.)

The cause for the break is as follows:

• all requirement checking adapters must be intrinsics, but
• all intrinsics must be listed in _INTRINSICS_CATALOG_ENTRIES, and
• there is no way to add custom stuff to _INTRINSICS_CATALOG_ENTRIES (and even if there were that would be ugly).

The failure point is in the Intrinsic.__init__ call within the AloraRequirement initializer:

        Intrinsic.__init__(
            self,
            intrinsic_name=intrinsic_name,
            intrinsic_kwargs={"requirement": f"{self.description}"},
        )

which calls adapters.catalog.fetch_intrinsic_metadata(intrinsic_name: str) function, which does a lookup against the _INTRINSICS_CATALOG_ENTRIES table. This means that you cannot use any adapter for requirement checking that is not part of the _INTRINSICS_CATALOG_ENTRIES table. That's a problem because we might want custom lora/alora requirement checkers that are implemented as adapters (as in the canonical stembolt example).

I tried to hack around this by adding to the catalog, but that also failed because the adapter must be a GraniteCommonAdapter (huggingface.py:354), which opened another whole can of worms.

Larger Issue

After digging through the current Intrinsic code-base, I think this is a symptom of a larger problem. We have not done a proper job at documenting the architectural choices associated with intrinsics and adapters. We may have also made some bad decisions, but I'm not sure because I am still not quite sure how everything is supposed to work.

I'm capturing some other alora questions here, which aren't directly related to getting m tune working again.

1. It's unclear to me how you're supposed to write the logic for activating a generic adapter without turning it into an intrinsic. Our current training code doesn't auto-generate the JSONs for eg activation sequences (trivial to fix), and I think you might also need to modify granite.common itself in order to use a custom adapter, at least for our current code base (less trivial to fix / does require some redesign). In any case, it should be possible to train an adapter and use it to implement a component's generation without bundling it as an intrinsic. Open to being
2. We tout that the primary advantage of aloras is that you re-use kv cache, but then implement aloras in terms of chat rewrites. If you rewite prior history, then you have to recompute kv cache anyways. So why not use a LoRA?
3. Our code currently allows the use of multiple aLoRAs against a single model. The original paper states that this is not possible. Have changed since then / do we have changes in the pipeline to enable this? If not, we should enforce the invariant that only one alora may be loaded / used with any particular backend. You must unload the old alora before loading the new one, and that triggers a cache invalidation.

[nrfulton]

Note: there are other hacky solutions, or perhaps even principled incremental (small bug fix) solutions. The larger issue I'm trying to point out is that the way that this whole system works is opaque and not very well documented.

[nrfulton]

Below I capture a proposed re-design. I want to first address the top-level concerns; it may not make sense to proceed down this path, so I'm burying it in a comment for now.

Proposed Solution

we've conflated intrinsics and adapters, or perhaps prioritized intrinsics over adapters.

I think that the actual substantive difference between an intrinsic and any other (alora) adapter is that intrinsics ship with three things:

1. A name
2. JSON file telling you how to load and activate the intrinsic
3. some chat rewrite logic implement in granite-common (looked up by name)

So "Intrinsic", in the current code-base, is just a really name for concept of a JSON file format and unique name which together entail a checkpoint, an activation sequence, and some chat rewrite logic. This means that it's a specific implementation of an Adapter.

AFAICT, there is no way to use a LoRA or aLoRA adapter that isn't shipped exactly the way that granite intrinsics are shipped. This is a problem because approximately all current users of LoRAs do not do things in the granite intrinsics way. For this reason, our code base should treat intrinsics as a special case of adapters -- not the other way around.

In mu proposed alternative design, we have Adapters and AdapterBackedComponents. Intrinsics are just special cases of both that implement the Adapter interface / behavior using the name, JSON, granite.common impl discussed above.

My belief is that

Intrinsic <: Adapter

and any action : Component backed by an intrinsic is ultimately backed by a particular type of Adapter, so there does not even need to exist an IntrinsicBackedComponent. Those are just AdapterBackedComponents whose adapters happen to be of type Intrinsic. I.e., what we currently call an Intrinsic is just a AdapterBackedComponent whose adapter happens to be an intrinsic.

I'll briefly sketch the adapter lifecycle and protocol, both because that's the basis of my belief and because we should flesh this out and document it somewhere (probably in docs/rewrite/adapters_and_intrinsics.md with corresponding examples).

Adapter Lifecycle

Adapters have a lifecycle that's mostly identical to how things work now. Parts of this protocol might need to be implemented in the backend rather than in the intrinsic. I'm not sure exactly what the split is; I think @jakelorocco has the most context to answer that question.

Loading (done by backends)

load
    * as-is

unload
    * as-is

Activation (done by backends)

activate
    * as-is

detect_activation
    * detects when an adapter has been activated by a token stream. Optional for now since we don't have this functionality atm. But we should design for it.

Deactivation (done by backends)

deactivate
    * as-is

detect_deactivation
    * detects when an adapter has been deactivated by a token stream. Optional for now since we don't have this functionality atm. 

Preparing (done by Adapaters or backends? IDK)

Adapter.rewrite(action, context, backend) -> action, context

Adapter.parse_result <- not necessary, used parsed_repr for this???

Specifics of Proposed Change

AdapterBackedComponent

The only way to use an adapter (therefore, intrinsic) is by passing an AdapaterBackedComponent as the action.

The AdapterBackedComponent is just a Component with an adapter : str | Adapter property. The union type works like our model_ids -- if it's a str then we look it up in the _INTRINSICS_CATALOG_ENTRIES but otherwise we use the adapter.

For example, from the current code base, we would need ALoRARequirement <: AdapaterBackedComponent with adapter="requirement_checking" by default.

This solves the lookup problem that led to my deep-dive of the curretn adapter/intrinsics code base.

Intrinsics

There is an IntrinsicAdapter <: Adapter which uses the Intrinsic JSON format to implement each of the methods. the backend might also dispatch on this subtype to do special handling if necessary (although I think this should not be necessary because Adapter.rewrite should do everything that's intrinsic-specific).

For example, there is no

Backend._generate_from_intrinsic

There is only

Backend._generate_from_adapter_backed_component

which operates over things that implement the AdapterBackedComponent interface using things that implement the Adapter interface.

@jakelorocco please let me know what you think.

Other design considerations:

1. Adapters can only be added to a single backend. Add lifecycle checks that enforce this invariant.
2. Adapters have only a single type. Is this correct? Or should Adapters have types, like Intrinsic does now? (This is an additional code smell that we're conflating a few different things, imo.)

[nrfulton]

I went ahead and hacked up some dicts and upload a model to huggingface, leaving a trace of issues as I worked. I now need a io.yaml file but cannot figure out what's supposed to go in those files or what the semantics are even after reviewing all of the ones in the rag lib. Do we have infra for bundling and uploading these intrinsic huggingface repo things?

Edit: figured out the io.yaml format. It looks like all intrinsics must return json; is that correct?