File size: 3,937 Bytes
d90b3a8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
import pytest
import eval
import generate
import train
from megatron.neox_arguments import NeoXArgs
from tests.common import save_random_model, simulate_deepy_env
from tools.datasets import preprocess_data
@pytest.fixture(
params=[
"HFGPT2Tokenizer",
"HFTokenizer",
"GPT2BPETokenizer",
"CharLevelTokenizer",
"TiktokenTokenizer",
"SPMTokenizer",
]
)
def tokenizer_type(request):
return request.param
@pytest.fixture(params=[None, "tests/data/sample_prompt.txt"])
def sample_input_file(request):
return request.param
@pytest.mark.cpu
def test_preprocess_data(tokenizer_type):
if tokenizer_type == "SPMTokenizer":
pytest.xfail(
reason="Expected easy resolution: Need to provide a valid model file from somewhere"
)
vocab_file = {
"HFTokenizer": "tests/data/hf_cache/tokenizer/gpt2.json",
"TiktokenTokenizer": "cl100k_base",
"HFGPT2Tokenizer": "gpt2",
}
input_args = [
"--input",
"./tests/data/enwik8_first100.txt",
"--output-prefix",
"./tests/data/enwik8_first100",
"--vocab",
vocab_file.get(tokenizer_type, "./data/gpt2-vocab.json"),
"--tokenizer-type",
tokenizer_type,
"--merge-file",
"./data/gpt2-merges.txt",
"--append-eod",
]
preprocess_data.main(input_args)
@pytest.mark.skip(
reason="All model tests are skipped until we fix the CUDA + torch multiprocessing issue."
)
def test_generate(monkeypatch, tmpdir, tmp_path, sample_input_file):
model_dir = str(tmpdir)
sample_output_file = str(tmp_path) + ".txt"
input_args = ["generate.py", "tests/config/test_setup.yml"]
deepspeed_main_args = simulate_deepy_env(monkeypatch, input_args)
save_random_model(deepspeed_main_args, model_dir)
# Generate output
generate_args = {
"load": model_dir,
"sample_input_file": sample_input_file,
"sample_output_file": sample_output_file,
}
generate.main(input_args=deepspeed_main_args, overwrite_values=generate_args)
@pytest.mark.skip(
reason="All model tests are skipped until we fix the CUDA + torch multiprocessing issue."
)
def test_evaluate(monkeypatch, tmpdir, tmp_path):
model_dir = str(tmpdir)
sample_output_file = str(tmp_path)
input_args = ["generate.py", "tests/config/test_setup.yml"]
deepspeed_main_args = simulate_deepy_env(monkeypatch, input_args)
save_random_model(deepspeed_main_args, model_dir)
# Generate output
evaluate_args = {
"load": model_dir,
"eval_tasks": ["lambada"], # ["lambada", "hellaswag", "piqa", "sciq"],
"eval_results_prefix": sample_output_file,
}
eval.main(input_args=deepspeed_main_args, overwrite_values=evaluate_args)
@pytest.mark.skip(
reason="All model tests are skipped until we fix the CUDA + torch multiprocessing issue."
)
def test_finetuning(monkeypatch, tmpdir, tmp_path):
# Save random model, load random model, keep training
# TODO: add mocking to check that we're not ignoring the previously loaded model
model_dir = str(tmpdir)
sample_output_file = str(tmp_path)
input_args = ["generate.py", "tests/config/test_setup.yml"]
deepspeed_main_args = simulate_deepy_env(monkeypatch, input_args)
save_random_model(deepspeed_main_args, model_dir)
# Generate output
finetune_args = {"load": model_dir, "finetune": True}
train.main(input_args=deepspeed_main_args, overwrite_values=finetune_args)
@pytest.mark.skip(
reason="All model tests are skipped until we fix the CUDA + torch multiprocessing issue."
)
def test_train_launcher(monkeypatch):
input_args = ["train.py", "tests/config/test_setup.yml"]
deepspeed_main_args = simulate_deepy_env(monkeypatch, input_args)
train.main(input_args=deepspeed_main_args)
|