Bias tensors

[gabe-l-hart]

Dependencies

This PR is part of a sequence in support of adding Granite Code. It depends on merging the following PRs:

• Safetensors: Safetensors #1255

Issues

Closes #1250

Description

This PR adds support for models which have bias tensors for the attention and ffn modules alongside the primary weight tensors.

Changes

• Add the bias tensors to the weight_map in HF checkpoint conversion
• Handle merged wqkv tensors for bias as well as weights in HF checkpoint conversion
  • This includes changes to the permutation logic to support the shapes of the bias tensors. I leveraged the corresponding logic in llama.cpp's converter.
• Add configs to TransformerArgs to allow models to indicate the presence of attention_bias and feed_forward_bias tensors
• Populate the Attention and FeedForward modules' tensors' bias arguments based on the config args

Testing

In conjunction with my other changes for Granite Code, I've been able to validate that the results produced with this logic do produce the expected token sequence.

NOTE: If there's any preferred way to include unit tests along with the PR, please let me know and I can get them added! I don't see a familiar unit test structure in the project at this point, so I've been relying on local ad-hoc testing.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1259

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit bbea338 with merge base 766bee9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[gabe-l-hart]

Thanks for the review/merge on #1255! This PR is now ready for review

[mikekg]

Current tests are run through .github/workflows, with some scripts in .ci (including a script that can be used to ensure that code in
markdown files works).

Or were you looking for "unit tests" of subcomponents with a Python driver? If you are looking for python-level unit tests, I don't think we have any right now, but that doesn't mean we can't have any. If you want to make a proposal, you might discuss with @byjlw and @Jack-Khuu, and @lessw2020 for distributed inference.