Dynamic Inputs for fx traced GPTNeoLM

[sdtblck]

🚀 Feature request

I've been looking a little into using torch.fx + HF lately, and run into a problem that the traced GPTNeoLM can't take dynamic inputs - would love to know what work would be needed to enable that, and what the current blockers are - it would be great to be able to use the traced module for inference.

I'm assuming the problem has something to do with caching of keys/values?

Your contribution

Given guidance, I could put some work into adding this feature, yes :)

[LysandreJik]

cc @michaelbenayoun

[michaelbenayoun]

Hi @sdtblck ,
The issue with dynamic inputs is not related to the past_key_values but to dynamic control-flow in the model forward pass implementation. The PR #14321, that I plan to merge at some point (hopefully before the end of the year) solves this kind of issues for the supported models.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.