[Feature]: Return hidden states (in progress?)

[Elanmarkowitz]

🚀 The feature, motivation and pitch

I know this feature request sort of already exists: #5950
(and older, semi related requests) #3594 #1857

This is a similar pitch but I am creating a new issue as I noticed newer developments in the codebase. The pitch is to support returning hidden states when generating sequences. This enables many potential behaviors such as output classification, guardrails, etc. Whereas #5950 suggested a different step for embedding, I would suggest building it in as an option to EngineArgs or as an option that can be passed in with each generation request.

I see that in v0.5.1 there is already some new code in ModelDriverBase to support return_hidden_states. However, I don't see that supported yet in the LLM engine yet (not an input to EngineArgs). Basically, it seems like this feature is under development. I am mainly wondering what the timeline is for that? And what is the approach being taken so that I and the community can develop accordingly?

Alternatives

No response

Additional context

No response

[LiuXiaoxuanPKU]

Thanks for the question! We currently use return_hidden_states for speculative decoding. You just need to pass it a a config as here. Feel free to mimic the behavior there.

[Hambaobao]

Hi, I also have the same need. I hope to store the hidden_states during model inference so that I can conduct some interpretability research.

[PeterAdam2015]

same need, hope we can get this as an option to return embedding.

[ummagumm-a]

same need!

[freesunshine0316]

Thanks for the question! We currently use return_hidden_states for speculative decoding. You just need to pass it a a config as here. Feel free to mimic the behavior there.

Hi,
can you further specify, e.g. with demo code?

[J0hnArren]

same need

[Gxy-2001]

same need

[zkwhandan]

same need

[jvlinsta]

Same need to generate some attention heatmaps, akin to

[github-actions]

This issue has been automatically marked as stale because it has not had any activity within 90 days. It will be automatically closed if no further activity occurs within 30 days. Leave a comment if you feel this issue should remain open. Thank you!

[rogerh2]

Same need