register parameters correctly in c++ MultiheadAttention #42037
Conversation
[ghstack-poisoned]
glaringlee
requested review from
anjali411,
pbelevich and
zhangguanheng66
and removed request for
ebetica and
goldsborough
|
This is to correct the parameter registration in MultiheadAttention, the root cause is described in #41951. |
| } else { | ||
| std::uniform_int_distribution<int> d(5, 20); | ||
| kv_dim = d(generator); | ||
| while (kv_dim == d_model) { |
There was a problem hiding this comment.
This branch was aimed to test kv_dim != d_model case, but the generator d could generate a same value as d_model, for eg. 9, so I add this while.
| bias_k = register_parameter("bias_k", torch::empty({1, 1, options.embed_dim()})); | ||
| bias_v = register_parameter("bias_v", torch::empty({1, 1, options.embed_dim()})); | ||
| } else { | ||
| bias_k = {}; |
There was a problem hiding this comment.
Do we need to register_parameter('bias_k', {}, /*requires_grad=*/false)?
There was a problem hiding this comment.
Based on python impl, No. It only register the non-empty case for bias_k and bias_v.
pytorch/torch/nn/modules/activation.py
Line 837 in 4121d34
zhangguanheng66
left a comment
zhangguanheng66
left a comment
There was a problem hiding this comment.
LGTM for the MHA functionality.
|
@glaringlee merged this pull request in 5246bc4. |
Stack from ghstack:
This is to fix #41951
Differential Revision: D22764717