Decoder-Only Transformer

[spencerp]

Supporting decoder-only transformer training:

parlai train_model --model transformer/decoder

Summary of changes

• New TransformerDecoderLayer without encoder attention
• New TransformerDecoder that concatenates input and "encoder state"
  - Needed to pass the dictionary to TransformerDecoder to properly do concatenation
  • Right-pad the label, left-pad the context
• New PassThroughEncoder, similar to IdentityLayer but compatible with TransformerEncoder API
  • There are many assumptions about there being an "encoder" in TorchGeneratorAgent, so the path of most code reuse led to using a dummy encoder to satisfy those assumptions
• New DecoderAgent to override build_model
• Extract label from scores to calculate loss (done at end of TransformerDecoderOnly.forward)
• Only use last query_len rows from incremental attention to account for incremental step from context to first generated token
• Added DecoderIncrState and DecoderLayerIncrState type aliases
• Added unit tests

Summary of structural changes in decoder.py

Testing steps

python -m pytest tests/test_transformers.py

python -m pytest tests/test_light_whoami.py

Also ran some small training runs locally to sanity check:

TF_ARGS="--embedding-size 128 --ffn-size 512 --batchsize 16 --eval-batchsize 16 --model-parallel true --variant prelayernorm --n-heads 8 --n-positions 512 --activation gelu --text-truncate 256 --label-truncate 128 -lr 7e-05 --lr-scheduler invsqrt --optimizer adam --warmup_updates 1000 -vp 10 -vmt ppl -vmm min --load-from-checkpoint false -vstep 500 --validation-max-exs 1000 -tstep 150000 --log-every-n-secs 30 --update-freq 1 --dynamic-batching full"
parlai train_model -t convai2 --model transformer/decoder --model-file /tmp/test_decoder_only_model --n-layers 2 $TF_ARGS
parlai train_model -t convai2 --model transformer/generator --model-file /tmp/test_enc_decoder_model --n-encoder-layers 1 --n-decoder-layers 1 $TF_ARGS

Benchmark Comparison with Encoder-Decoder
Both trained to 10k steps with 120M parameters (8 layers) on 2x Quadro GP100 GPUs. The encoder/decoder is closer to 140M parameters due to the cross-attention.

TF_ARGS="--embedding-size 1024 --ffn-size 4096 --batchsize 16 --eval-batchsize 16 --model-parallel true --variant prelayernorm --n-heads 16 --n-positions 512 --activation gelu --text-truncate 256 --label-truncate 128 -lr 7e-05 --lr-scheduler invsqrt --optimizer adam --warmup_updates 1000 -vp 10 -vmt ppl -vmm min --load-from-checkpoint false -vstep 1000 --validation-max-exs 1000 -tstep 10000 --log-every-n-secs 30 --update-freq 1 --dynamic-batching full --fp16 true"

parlai train_model -t internal:new_reddit --model transformer/decoder --model-file /tmp/test_decoder_only_model --n-layers 8 $TF_ARGS
...
13:17:37 | time:6960s total_exs:433912 total_steps:9950 epochs:0.00 time_left:35s
    clen  clip  ctpb  ctps  ctrunc  ctrunclen  exps  exs  fp16_loss_scalar  gnorm  gpu_mem  llen  loss        lr  ltpb  ltps  ltrunc  ltrunclen   ppl  token_acc  token_em  total_train_updates  tpb  tps   ups
   58.72     1  2603  4957  .01423      .8327 85.63 2248            131072  1.711    .2126 26.09 4.036 2.219e-05  1173  2234       0          0 56.59      .2843         0                 9950 3776 7191 1.905

13:18:04 | Stopping from Maximum LR steps

parlai train_model -t internal:new_reddit --model transformer/generator --model-file /tmp/test_enc_decoder_model --n-encoder-layers 4 --n-decoder-layers 4 $TF_ARGS
...
10:52:31 | time:5367s total_exs:431792 total_steps:9950 epochs:0.00 time_left:27s
    clen  clip  ctpb  ctps  ctrunc  ctrunclen  exps  exs  fp16_loss_scalar  gnorm  gpu_mem  llen  loss        lr  ltpb  ltps  ltrunc  ltrunclen   ppl  token_acc  token_em  total_train_updates  tpb   tps  ups
   64.57     1  2711  6939  .02041      1.706 110.4 2156            131072  1.633    .1838 28.03 4.066 2.219e-05  1208  3094       0          0 58.33      .2803         0                 9950 3919 10033 2.56

10:52:51 | Stopping from Maximum LR steps

[stephenroller]

Yeah looks pretty good to me. Do you mind giving us a benchmark on -t internal:new_reddit with 125M params? trained to say, 10k steps?

[stephenroller]

wishlist: this would be a really great opportunity to finally implement the parser.remove_arg functionality we've desired for a long time

[stephenroller]

(y)

[stephenroller]

Alright it's a little slower but good enough for today. Let Jason know it's done!

[spencerp]

Hm, there's a distillation test failing which I don't understand yet... dec_hid_loss and enc_dec_attn_loss don't match what the test expects which is suspicious, so I'm going to wait to merge until I understand that better.

[klshuster]

cc @EricMichaelSmith: there are two distillation tests failing due to this PR, do you have any idea why that might be the case? it looks like very small (but consistent) differences in testing (that is, these tests do not seem to be failing on main but consistently fail on this branch)

[stephenroller]

Let's just let this sit a while. The regression failure is a little concerning. It's enough to make me worry the code has changed in a subtle way, but close enough that it's tempting to just --force-regen and forget about it.

[EricMichaelSmith]

Let's just let this sit a while. The regression failure is a little concerning. It's enough to make me worry the code has changed in a subtle way, but close enough that it's tempting to just --force-regen and forget about it.

Hmm seconded - yeah, @klshuster I'm not sure offhand, but perhaps this does indicate that the forward pass has indeed changed in some subtle way?

[github-actions]

This PR has not had activity in 30 days. Closing due to staleness.

[stephenroller]

[spencerp]

I will resuscitate it next week :)

[spencerp]

Finally got back to this today and figured out the issue!

Apparently the distillation code is sensitive to the order in which modules are initialized 😔. I made a branch to demonstrate the minimal change needed to repro the test break: #4526

To my knowledge, module initialization order has no implications on training dynamics/model performance. So I'm inclined to call this a distillation bug and outside the scope of this PR. Anyone know differently? @EricMichaelSmith

If I don't hear any objections by Thursday I'll merge as-is with the broken test and file an issue to fix the bug in the distillation code.

[spencerp]

Chatted with @EricMichaelSmith offline. He pointed out that this could be a result of a different order of random operations done during module initialization. I was able to confirm that's the problem, details in #4526

So I'm just going to update the test numbers and merge.