[s2s] distill t5-large -> t5-small

[sbhaktha]

What does this PR do?

@sshleifer : This PR addresses the feature I was requesting on this thread. It add a way to specify a base model to initialize the student model with (in distillation.py), instead of necessarily starting from the same base model as the teacher and selectively copying layers.

Following are the main changes:

1. Accepts a new command line argument --student that can take a base model name, such as t5-small.
2. If this new argument is specified, simply creates a student model from scratch, if not creates the student model per the original code.
3. In the _step function, calls the teacher decoder with the teacher's encoder outputs instead of student encoder outputs.
4. When teacher and student base models are different, since they can be of different architectures, and hence have different hidden sizes, does not calculate hidden loss.

Sample command to run this:

python distillation.py --teacher t5-large --data_dir $NQOPEN_DIR \
--student t5-small --tokenizer_name t5-small \
--teacher_tokenizer_name t5-large \
--learning_rate=3e-4 --freeze_encoder --freeze_embeds \
--do_train --train_batch_size 32 \
--do_predict --n_train 3 \
--model_name_or_path t5-large --eval_beams 2 --eval_max_gen_length 142 \
--val_check_interval 0.25 --n_val 1 \
--output_dir distilled-t5 --gpus 1 --logger_name wandb

I was able to get this training to run. I am not sure what tests I need to run, and also noticed some automated test failure notifications to my email that I didn't quite follow.

This is the first PR I am making here, so hope you can point me in the right direction as to what needs to be done if anything is missing.

[x] Was this discussed/approved via a Github issue or the forum? Please add a link
to the it if that's the case.
https://discuss.huggingface.co/t/distillation-create-student-model-from-a-different-base-model-than-teacher/1501/8

Who can review?

examples/seq2seq: @sshleifer

[sshleifer]

Thanks for the contribution, I think it is a reasonable feature to add and I'm optimistic about getting this merged!

• You can fix check_code_quality by running make fixup
• You can replicate the run_examples_torch failure locally by running

pytest examples/seq2seq/test_seq2seq_examples.py -k distil

These need to pass for the PR to merge.

• It would also be good to add a test that uses the student kwarg to test_seq2_examples.py

[sshleifer]

For your test, you could check the code runs correctly with dict(teacher=BART_TINY_RANDOM, student=T5_TINY_RANDOM)

[sbhaktha]

I've pushed some fixes to get the tests to run successfully! However, I realized that an issue I was wondering about earlier is indeed important, and that is, dealing with the student and teacher models having different tokenizers, which would make the input_ids different in the two cases. However since the (language) models we are looking at all have the same output space (set of tokens) as the input space, in order for the logits from the output layers to be comparable, they have to have the same output space, and hence the same set of tokens. Does this make sense? And if so, it is a fair assumption that the teacher and student models should have the same tokenizer (specified via the existing tokenizer_name parameter), correct? In this case, unit testing distillation from BART_TINY to T5_TINY is not valid (and I do run into errors). I did try T5_TINY to T5_TINY and that works fine but of course it's not very satisfying. Can I add a test case for T5_SMALL to T5_TINY? Worried about slowing down the unit tests. Please let me know your suggestion. Sorry for the long comment!

[sshleifer]

For testing,
I just ran

python make_student.py patrickvonplaten/t5-tiny-random t5-tinier-random -e 1 -d 1
transformers-cli upload t5-tinier-random

So now you can test distilling patrickvonplaten/t5-tiny-random (2/2 layers) to sshleifer/t5-tinier-random (1/1 layer).

[sshleifer]

You also need to run make fixup

[sbhaktha]

You also need to run make fixup

I get the following message:

/bin/sh: 3: black: not found
/bin/sh: 4: isort: not found
/bin/sh: 5: flake8: not found
Makefile:7: recipe for target 'modified_only_fixup' failed
make: *** [modified_only_fixup] Error 127

Any idea?

[sbhaktha]

You also need to run make fixup

I get the following message:

/bin/sh: 3: black: not found
/bin/sh: 4: isort: not found
/bin/sh: 5: flake8: not found
Makefile:7: recipe for target 'modified_only_fixup' failed
make: *** [modified_only_fixup] Error 127

Any idea?

Ah. I wasn't familiar with these pre-commit libraries. Looks like I need to pip install those...

[sbhaktha]

I've added a unit test to distill from T5_TINIER to T5_TINY and it passes. Also ran make fixup and fixed formatting issues. Could you please review again?
All CircleCI tests above pass except the first one, check_code_quality. I checked the details on that one and they seem to be unrelated:

Run python -c "import torch; print(torch.hub.list('huggingface/transformers:$BRANCH'))"
Downloading: "https://github.com/huggingface/transformers/archive/add_student_base_model.zip" to /home/runner/.cache/torch/hub/add_student_base_model.zip
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/site-packages/torch/hub.py", line 269, in list
    repo_dir = _get_cache_or_reload(github, force_reload, True)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/site-packages/torch/hub.py", line 141, in _get_cache_or_reload
    download_url_to_file(url, cached_file, progress=False)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/site-packages/torch/hub.py", line 425, in download_url_to_file
    u = urlopen(req)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 222, in urlopen
    return opener.open(url, data, timeout)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 531, in open
    response = meth(req, response)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 641, in http_response
    'http', request, response, code, msg, hdrs)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 563, in error
    result = self._call_chain(*args)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 503, in _call_chain
    result = func(*args)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 755, in http_error_302
    return self.parent.open(new, timeout=req.timeout)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 531, in open
    response = meth(req, response)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 641, in http_response
    'http', request, response, code, msg, hdrs)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 569, in error
    return self._call_chain(*args)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 503, in _call_chain
    result = func(*args)
  File "/opt/hostedtoolcache/Python/3.7.9/x64/lib/python3.7/urllib/request.py", line 649, in http_error_default
    raise HTTPError(req.full_url, code, msg, hdrs, fp)
urllib.error.HTTPError: HTTP Error 404: Not Found
Error: Process completed with exit code 1.

[sbhaktha]

thank you for the style fix! are you going to merge?

[sshleifer]

I think in the old case of distilling to a copied student, this change uses more memory.
Trying to find a clean solution locally.

[sbhaktha]

I see. Not sure what part of the change makes it use more memory? Is there any way I can help?

[sshleifer]

We used to only extract hidden states, now we extract hidden states and encoder outputs.
Your application doesn't need hidden states, the other application doesn't need encoder outputs, so splitting the call to self.teacher.get_encoder()( to only get what's needed is the next step.

[sbhaktha]

got it.I am attempting a fix.

[sbhaktha]

It is not true that the other application doesn't need encoder outputs. When calling the teacher decoder, if student and teacher encoders are different, the right thing to do is to pass the teacher encoder outputs, and as we can see from these lines, even in the original case, we could have self.different_encoder be true if not all layers were copied from teacher to student:

transformers/examples/seq2seq/distillation.py

Lines 60 to 61 in 02bdfc0

	self.different_encoder = hparams.student_encoder_layers != teacher_encoder_layers
	self.different_decoder = hparams.student_decoder_layers != teacher_decoder_layers

.
So I think we do need to extract teacher encoder outputs inside this if different_encoder block.
However, we can extract hidden states optionally based on whether we need to calculate hidden loss or not. I have the following modified code which passes the unit tests. Please lmk what you think.

        if self.different_encoder:  # compute encoder hidden state loss
            with torch.no_grad():
                teacher_encoder = self.teacher.get_encoder()(
                    input_ids,
                    attention_mask=src_mask,
                    output_hidden_states=self.do_calc_hidden_loss,
                    return_dict=True
                )
            teacher_enc_outputs = teacher_encoder.last_hidden_state
            if self.do_calc_hidden_loss:
                teacher_enc_hid = teacher_encoder.hidden_states
                hid_loss_enc = self.calc_hidden_loss(
                    src_mask,
                    enc_hidden_state,
                    teacher_enc_hid,
                    self.e_matches,
                    normalize_hidden=self.hparams.normalize_hidden,
                )

That said, the same boolean self.do_calc_hidden_loss can be used to set output_hidden_states to True or False in the call to the teacher decoder as well . If the above makes sense to you I'll make a similar change here as well and push a commit.

[sbhaktha]

ah, thanks for the clarification! I'll take another look.

[sbhaktha]

Ok I have another fix that's passing unit tests. Please lmk what you think. I renamed the self.do_calc_hidden_loss to self.different_base_models because really that's what it was keeping track of, and the boolean value is flipped from the earlier one.

        if self.different_encoder:  # compute encoder hidden state loss
            with torch.no_grad():
                teacher_encoder = self.teacher.get_encoder()(
                    input_ids,
                    attention_mask=src_mask,
                    output_hidden_states=not self.different_base_models,
                    return_dict=True
                )
            if self.different_base_models:
                teacher_enc_outputs = teacher_encoder.last_hidden_state
            else:
                teacher_enc_hid = teacher_encoder.hidden_states
                hid_loss_enc = self.calc_hidden_loss(
                    src_mask,
                    enc_hidden_state,
                    teacher_enc_hid,
                    self.e_matches,
                    normalize_hidden=self.hparams.normalize_hidden,
                )

        teacher_mask = input_ids.ne(pad_token_id)
        with torch.no_grad():
            outputs = self.teacher(
                input_ids,
                attention_mask=teacher_mask,
                encoder_outputs=(teacher_enc_outputs, ),
                decoder_input_ids=decoder_input_ids,
                lm_labels=labels,
                output_hidden_states=not self.different_base_models,
                return_dict=True,
            )
        dec_mask = decoder_input_ids.ne(pad_token_id)
        loss_ce = self.calc_ce_loss(dec_mask, lm_logits, outputs.logits)
        if (not self.different_base_models) and self.alpha_hid > 0:  # Intermediate supervision of decoder hidden states
            tdec_hidden = outputs.decoder_hidden_states
            hid_loss_dec = self.calc_hidden_loss(
                dec_mask, dec_hidden, tdec_hidden, self.d_matches, normalize_hidden=self.hparams.normalize_hidden
            )

self.different_base_models is initialized in __init__ in the same place as the earlier self.do_calc_hidden_loss as follows:

        self.different_base_models = True

        if hparams.student is None or hparams.teacher == hparams.student:
            self.different_base_models = False
            # Intermediate supervision: Decide which layers to supervise
            if hparams.supervise_forward:
                self.e_matches = get_layers_to_supervise(n_student=len(self.e_layer_ids), n_teacher=teacher_encoder_layers)
                self.d_matches = get_layers_to_supervise(n_student=len(self.d_layer_ids), n_teacher=teacher_decoder_layers)
            else:  # student layer should emulate hidden states of the teacher layer it was copied from
                self.e_matches = self.e_layer_ids
                self.d_matches = self.d_layer_ids

[sshleifer]

yes push that, I'll make some changes and get it merged.

[sshleifer]

Thanks for iterating :)

[sbhaktha]

pushed!

[sbhaktha]

@danyaljj FYI

[sbhaktha]

CircleCI check_code_quality is failing. I see--

#!/bin/bash -eo pipefail
black --check examples tests src utils
would reformat /home/circleci/transformers/examples/seq2seq/distillation.py
Oh no! 💥 💔 💥
1 file would be reformatted, 562 files would be left unchanged.

Exited with code exit status 1
CircleCI received exit code 1

But not sure how to format it. Tried black --code examples/seq2seq/distillation.py but that doesn't work cause it expects code as a string, not a filename. How do you format a file?

[sshleifer]

pip install -e .[dev]
make fixup

[sbhaktha]

make fixup

I did run make fixup before committing. But just ran both your above steps again. I don't get any errors from distillation.py.

[sshleifer]

On a high-level, do you have a fine-tuned t5-large that you want to distill?
This makes the code a lot more complex so I am trying to understand the use-case better.

[sbhaktha]

Yes. We wish to create a small model (under 500 MB), distilled from a large or 11B T5 model. For this we want to start with a small T5 model and train using the larger T5 as the teacher. Conceptually distillation can be done from any large model to any small model as long they share the same output space, so this enables that.

[sshleifer]

For this we want to start with a small T5 model and train using the larger T5 as the teacher.

My point was that this only really makes sense if the large t5 is fine-tuned.

[sbhaktha]

ah, yes. we are using a fine-tuned large model.

[sshleifer]

Cleaned up some things and now ready to merge. Thanks for the contribution @sbhaktha !

[sbhaktha]

Thanks Sam!