Add Remote LLM Support for Perturbation-Based Attribution via RemoteLLMAttribution and VLLMProvider #1544
Conversation
… hosted models that provide logprobs (like vLLM)
|
Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at [email protected]. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
|
Thank you @saichandrapandraju for the great effort! Generally, i agree this idea makes a lot of sense. @craymichael can you take a look at it? since you have studied the integration with llama-stack before |
|
Thank you for the positive feedback @aobo-y ! Happy to hear the direction makes sense. Please take your time reviewing — I’ll be around to clarify or iterate on anything as needed. Looking forward to it! |
|
Is there any update on this? Seems like a great idea! |
|
Hi, sorry for the delay on this! I actually have on my todos to dive into this PR next week. |
There was a problem hiding this comment.
I just wrapped up review of the diff - it's really in great shape, thank you for contributing this! Moreover, the remote provider abstraction is generic enough to support other frameworks, including Llama Stack (which includes support for log prob configs and prompt_logprobs depending on provider). I left some comments and questions (and some thinking out loud...). Let's discuss/iterate and we can land this contribution soon!
| attr_kws["n_samples"] = n_samples | ||
|
|
||
| # In remote mode, we don't need the actual model, this is just a placeholder | ||
| placeholder_model = torch.nn.Module() |
There was a problem hiding this comment.
For a test this is fine. In release notes/an example notebook we can instead use a simpler placeholder like
lambda *_: 0since the forward function has signature Callable[..., Union[int, float, Tensor, Future[Tensor]]]. I'm thinking as well that for convenience we can have the above lambda available as a class attribute of RemoteLLMAttribution with name placeholder_model so the syntax looks like:
AttrClass(RemoteLLMAttribution.placeholder_model)There was a problem hiding this comment.
As part of these changes, we should update (or create a counterpart) to the LLM attribution tutorial and update the docs. I'm happy to do this once this feature lands as well.
| from abc import ABC, abstractmethod | ||
| from typing import Any, List, Optional | ||
| from captum._utils.typing import TokenizerLike | ||
| from openai import OpenAI |
There was a problem hiding this comment.
Since this is being imported into attr.py, let's move this import to the VLLMProvider init wrapped around with a try-except, telling the user to install openai package if it isn't already. We'll keep it as an optional dependency as you have it in setup.py
| """ | ||
| # Parameter normalization | ||
| if 'max_tokens' not in gen_args: | ||
| gen_args['max_tokens'] = gen_args.pop('max_new_tokens', 25) |
There was a problem hiding this comment.
nit: add default to docstring
|
|
||
| try: | ||
| self.client = OpenAI(base_url=self.api_url, | ||
| api_key=os.getenv("OPENAI_API_KEY", "EMPTY") |
There was a problem hiding this comment.
nit: add env var to docstring
Looks like we don't need to handle this logic
https://github.com/openai/openai-python/blob/main/src/openai/_client.py#L114
| Initialize a vLLM provider. | ||
|
|
||
| Args: | ||
| api_url: The URL of the vLLM API |
There was a problem hiding this comment.
I would let api_url be an optional type without a default. With OpenAI API it will default to openAI's base URL which is what users may expect, but we can wrap VLLMProvider and call it OpenAIProvider for convenience.
Suggestion:
class VLLMProvider(RemoteLLMProvider):
def __init__(self, api_url: Optional[str], model_name: Optional[str] = None):Additional class (also imported into attr):
class OpenAIProvider(VLLMProvider):
def __init__(self, api_url: Optional[str] = None, model_name: Optional[str] = None):There was a problem hiding this comment.
Actually looks like OpenAI isn't letting us force logprobs with new models...feel free to ignore this for now
There was a problem hiding this comment.
this class is bound to OpenAI sdk. I am wondering how much of it is general to be extracted to another VLLMProvider. Maybe just name this OpenAIProvider?
There was a problem hiding this comment.
One thing to consider is that OpenAI doesn't support prompt logprobs at least with newer models, it's a VLLM engine parameter which is why it needs to be passed in a weird way. I actually don't think prompt_logprobs is even supported by OpenAI and there is a different way to grab logprobs, but I'm not certain.
Also VLLM server is its own library, but uses the OpenAI API and so is compatible with its SDKs as long as the base URL is set correctly. VLLM doesn't have its own SDK afaik.
| models = self.client.models.list().data | ||
| if not models: | ||
| raise ValueError("No models available from the vLLM API") | ||
| self.model_name = models[0].id |
There was a problem hiding this comment.
nit: let's log an informational message about which model is being used since none was provided
| prompt=prompt, | ||
| temperature=0.0, | ||
| max_tokens=1, | ||
| extra_body={"prompt_logprobs": 0} |
There was a problem hiding this comment.
Should prompt_logprobs be 1 here? I thought 0 would result in an exception as it needs to be > 1
| for probs in response.choices[0].prompt_logprobs[1:]: | ||
| if not probs: | ||
| raise ValueError("Empty probability data in API response") | ||
| prompt_logprobs.append(list(probs.values())[0]['logprob']) |
There was a problem hiding this comment.
If I understand correctly, prompt_logprobs is a list[dict[str, dict[str, float]]] where each element corresponds to a prompt token with contents being a map from generation token for each of the k prompt_logprobs with generation data in the token's map, including the logprob. If so, maybe we should assert that the length of probs is always 1?
Also iiuc, I can nitpick:
- We can iterate over the final
num_target_str_tokensonly - Can replace
list(probs.values())[0]withnext(iter(probs.values()))
| if not target_str: | ||
| raise ValueError("Target string cannot be empty") | ||
|
|
||
| num_target_str_tokens = len(tokenizer.encode(target_str, add_special_tokens=False)) |
There was a problem hiding this comment.
I think that each response contains token-level information. I think that it's possible that when we generate a response if none is provided that we might not need to require a tokenizer, which will minimize tokenizer-model mismatch. However, I think we can just make note of this and keep as future work.
| @@ -669,3 +670,362 @@ def test_llm_attr_with_skip_tensor_target(self) -> None: | |||
| self.assertEqual(token_attr.shape, (5, 4)) | |||
| self.assertEqual(res.input_tokens, ["<sos>", "a", "b", "c"]) | |||
| self.assertEqual(res.output_tokens, ["m", "n", "o", "p", "q"]) | |||
|
|
|||
| class DummyRemoteLLMProvider(RemoteLLMProvider): | |||
There was a problem hiding this comment.
It would be great to also have a test of VLLMProvider with a mocked openai client (method client.completions.create mainly) that solidifies the expected values of prompt_logprobs and verifies the logprob logic
|
Thank you for the comments @craymichael . I'll work on these and get back to you soon! |
This PR introduces support for applying Captum's perturbation-based attribution algorithms to remotely hosted large language models (LLMs). It enables users to perform interpretability analyses on models served via APIs, such as those using vLLM, without requiring access to model internals.
Motivation:
Captum’s current LLM attribution framework requires access to local models, limiting its usability in production and hosted environments. With the rise of scalable remote inference backends and OpenAI-compatible APIs, this PR allows Captum to be used for black-box interpretability with hosted models, as long as they return token-level log probabilities.
This integration also aligns with ongoing efforts like llama-stack, which aims to provide a unified API layer for inference (and also for RAG, Agents, Tools, Safety, Evals, and Telemetry) across multiple backends—further expanding Captum’s reach for model explainability.
Key Additions:
RemoteLLMProviderInterface:A generic interface for fetching log probabilities from remote LLMs, making it easy to plug in various inference backends.
VLLMProviderImplementation:A concrete subclass of
RemoteLLMProvidertailored for models served using vLLM, handling the specifics of communicating with vLLM endpoints to retrieve necessary data for attribution.RemoteLLMAttributionclass:A subclass of
LLMAttributionthat overrides internal methods to work with remote providers. It enables all perturbation-based algorithms (e.g., Feature Ablation, Shapley Values, KernelSHAP) using only the output logprobs from a remote LLM.Used openai client under the hood for querying remote models, as many LLM serving solutions now support the OpenAI-compatible API format (e.g., vLLM OpenAI server and projects like
llama-stack(see here for ongoing work related to this).Issue(s) related to this: