fix: check for wte.weight along with embed_tokens.weight

foundation-model-stack / fms-hf-tuning

🚀 Collection of tuning recipes with HuggingFace SFTTrainer and PyTorch FSDP.

Apache License 2.0

28 stars 48 forks source link

fix: check for wte.weight along with embed_tokens.weight #356

Closed willmj closed 2 months ago

willmj commented 2 months ago

Description of the change

Some granite models use naming convention wte.weight instead of embed_tokens.weight. Accounting for the difference for v2.0.0 release until we can get something more generic.

Related issue number

How to verify the PR

Was the PR tested

[ ] I have added >=1 unit test(s) for every new method I have added.
[ ] I have ensured all unit tests pass

github-actions[bot] commented 2 months ago

Thanks for making a pull request! 😃 One of the maintainers will review and advise on the next steps.