cisnlp / ofa Goto Github PK

Hi there!

Very cool work :)

One thing I wanted to point out is that your work does not correctly compare throughput of variants in Table 3.

Huggingface transformers (and your implementation) unnecessarily forwards tokens with invalid labels (i.e., tokens for which labels = -100) through the language modelling head during training training.

Consequently, your sizable vocabulary (~400K) significantly slows down variants with higher dimensional embedding matrices as a lot of unrequired dot products are computed.

As a fix, at training time your implementation should sub-select tokens with valid labels prior to the forward pass in the LMHead.

Some back of the envelope math:

• batch size: 384
• seq len: 512
• Non-valid tokens: c.80%
• Vocab size: 401K

batch_size * seq_len * 0.8 * 401000 = 62914M

unrequired inner products.. per batch I'm afraid. I don't think I'm missing something, but of course please double check :)