Tests

[thomasw21]

Context

I started experimenting with tests after coming up with an issue regarding prefix lm where the deepspeed version of gpt has some hacks that changes some behaviours concerning attention mechanisms. https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/main/pretrain_gpt.py#L56-L71

basically the issue, is that deepspeed would override some arguments passed down forcing the model to ignore whatever mask I would feed to the rest, causin everything to run smoothly but using a causal mask instead of a prefix one.

• Fixes: Add tests & setup CI bigscience-workshop/Megatron-DeepSpeed#53

What I want to test

Ideally I would want to test out deepspeed --num_gpus 1 pretrain_gpt.py ${ARGS}. However that prevents us from easily handling different types of inputs because data is fed via a file. Instead I did the following:

model, _, _ = setup_model_and_optimizer(gpt_model_provider)
model = model[0]
input_batch = get_gpt_batch_pipe({"text": token_ids})[0]
output = model(*input_batch)

It's not ideal, as it feels like a hacky way to intialise the model correctly. But it allows to use the model "interactively", ie I can process an input and then test outputs.

We use patch to simulate different arguments.

How I ran test

# Prevent hf transformers and datasets library to access internet
export TRANSFORMERS_OFFLINE=1
export HF_DATASETS_OFFLINE=1

# Emulate a distributed env
export MASTER_ADDR=localhost
export MASTER_PORT=9994
export RANK=0
export LOCAL_RANK=0
export WORLD_SIZE=1

# TESTING SOME STUFF

python -m unittest Megatron-DeepSpeed.megatron.model.test.test_gpt_model

TODO

• Write test for causal gpt
• Write test for prefix lm
• Write test for rotary embeddings (I'm not sure how to test that exactly, besides it's running)

[thomasw21]

This is okay as all values beyong the 3rd index are ignored. Which is nice to be used to expose variables needed for testing.

[stas00]

Glad to see you started working on it, @thomasw21

FYI, there will be more tests coming from: bigscience-workshop#47

• So, first we will build the test suite around pytest since we have a rich library of helper test utils already at HF transformers and use unittest on top of it where it makes sense.

• In particular the distributed target application support that is implemented here:
  https://github.com/huggingface/transformers/blob/7fcee113c163a95d1b125ef35dc49a0a1aa13a50/src/transformers/testing_utils.py#L1200
  If you want to see a good example of how it's being used, please see:
  https://github.com/huggingface/transformers/blob/7fcee113c163a95d1b125ef35dc49a0a1aa13a50/tests/deepspeed/test_deepspeed.py#L780

• Then the important one is TestCasePlus that provides lots of great features.
  https://github.com/huggingface/transformers/blob/7fcee113c163a95d1b125ef35dc49a0a1aa13a50/src/transformers/testing_utils.py#L753
  the key is managing temp dirs and and much more.

• Then log and std stream catchers...

well a lot more

I guess we can just drop the helper library into tests for now.

• The base for the tests will be under ./tests and run as pytest tests. i.e. all tests will be there.

Let me know if you prefer I set things up first and then you can drop in all your awesome additions.

[thomasw21]

Cool! So I don't have much experience with pytest (none actually) or how it's done in transformers. Maybe you could setup it up and then I can migrate to the new framework?

Either way I think I can still merge this (on my prefix lm branch but also on master) as those are runnable tests (see commands in the description). Which might be helpful for engs to at least check they haven't broken anything on their setup. Once we setup CI it'll be much smoother (and costly). The tests here don't guarantee you haven't broken anything, but IMO they are a good attempt at catching some basic things (like my attention issue).

[stas00]

Sounds good, @thomasw21. You can merge and I will setup the test suite tomorrow and then you will move/adapt it if you prefer that order.