Skip to content
GitLab

[Cleanup] Remove unused args in MHA, cleanup empty returns

Merged Administrator requested to merge susan/cleanup_oct15_mha1 into main

Created by: suchenzang

From MHA, removed the following args:

  • need_head_weights
  • need_weights
  • static_kv
  • before_softmax

Cleaned up some cluttered returns from MHA (attn_weights and l_aux).

Follow-up: https://github.com/facebookresearch/metaseq/issues/423

Tested with 2x model parallel 125m, and non-model-parallel 125m.