wehub-resource-sync
314 lines
10 KiB
Python
314 lines
10 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import contextlib
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
import unittest
|
|
from io import StringIO
|
|
|
|
import torch
|
|
from fairseq import options
|
|
from fairseq_cli import train
|
|
from tests.utils import (
|
|
create_dummy_data,
|
|
generate_main,
|
|
preprocess_lm_data,
|
|
preprocess_translation_data,
|
|
train_translation_model,
|
|
)
|
|
|
|
|
|
class TestTranslationGPU(unittest.TestCase):
|
|
def setUp(self):
|
|
logging.disable(logging.CRITICAL)
|
|
|
|
def tearDown(self):
|
|
logging.disable(logging.NOTSET)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU")
|
|
def test_fp16(self):
|
|
with contextlib.redirect_stdout(StringIO()):
|
|
with tempfile.TemporaryDirectory("test_fp16") as data_dir:
|
|
create_dummy_data(data_dir)
|
|
preprocess_translation_data(data_dir)
|
|
train_translation_model(data_dir, "fconv_iwslt_de_en", ["--fp16"])
|
|
generate_main(data_dir)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU")
|
|
def test_memory_efficient_fp16(self):
|
|
with contextlib.redirect_stdout(StringIO()):
|
|
with tempfile.TemporaryDirectory("test_memory_efficient_fp16") as data_dir:
|
|
create_dummy_data(data_dir)
|
|
preprocess_translation_data(data_dir)
|
|
train_translation_model(
|
|
data_dir, "fconv_iwslt_de_en", ["--memory-efficient-fp16"]
|
|
)
|
|
generate_main(data_dir)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU")
|
|
def test_transformer_fp16(self):
|
|
with contextlib.redirect_stdout(StringIO()):
|
|
with tempfile.TemporaryDirectory("test_transformer") as data_dir:
|
|
create_dummy_data(data_dir)
|
|
preprocess_translation_data(data_dir)
|
|
train_translation_model(
|
|
data_dir,
|
|
"transformer_iwslt_de_en",
|
|
[
|
|
"--encoder-layers",
|
|
"2",
|
|
"--decoder-layers",
|
|
"2",
|
|
"--encoder-embed-dim",
|
|
"64",
|
|
"--decoder-embed-dim",
|
|
"64",
|
|
"--fp16",
|
|
],
|
|
run_validation=True,
|
|
)
|
|
generate_main(data_dir)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU")
|
|
def test_levenshtein_transformer(self):
|
|
with contextlib.redirect_stdout(StringIO()):
|
|
with tempfile.TemporaryDirectory(
|
|
"test_levenshtein_transformer"
|
|
) as data_dir:
|
|
create_dummy_data(data_dir)
|
|
preprocess_translation_data(data_dir, ["--joined-dictionary"])
|
|
train_translation_model(
|
|
data_dir,
|
|
"levenshtein_transformer",
|
|
[
|
|
"--apply-bert-init",
|
|
"--early-exit",
|
|
"6,6,6",
|
|
"--criterion",
|
|
"nat_loss",
|
|
],
|
|
task="translation_lev",
|
|
)
|
|
gen_config = [
|
|
"--task",
|
|
"translation_lev",
|
|
"--iter-decode-max-iter",
|
|
"9",
|
|
"--iter-decode-eos-penalty",
|
|
"0",
|
|
"--print-step",
|
|
]
|
|
# non-ensemble generation
|
|
generate_main(data_dir, gen_config)
|
|
# ensemble generation
|
|
generate_main(
|
|
data_dir,
|
|
gen_config,
|
|
path=os.pathsep.join([
|
|
os.path.join(data_dir, "checkpoint_last.pt"),
|
|
os.path.join(data_dir, "checkpoint_last.pt"),
|
|
]),
|
|
)
|
|
|
|
|
|
def _quantize_language_model(data_dir, arch, extra_flags=None, run_validation=False):
|
|
train_parser = options.get_training_parser()
|
|
train_args = options.parse_args_and_arch(
|
|
train_parser,
|
|
[
|
|
"--task",
|
|
"language_modeling",
|
|
data_dir,
|
|
"--arch",
|
|
arch,
|
|
"--optimizer",
|
|
"adam",
|
|
"--lr",
|
|
"0.0001",
|
|
"--criterion",
|
|
"adaptive_loss",
|
|
"--adaptive-softmax-cutoff",
|
|
"5,10,15",
|
|
"--max-tokens",
|
|
"500",
|
|
"--tokens-per-sample",
|
|
"500",
|
|
"--save-dir",
|
|
data_dir,
|
|
"--max-epoch",
|
|
"1",
|
|
"--no-progress-bar",
|
|
"--distributed-world-size",
|
|
"1",
|
|
"--ddp-backend",
|
|
"no_c10d",
|
|
"--num-workers",
|
|
"0",
|
|
]
|
|
+ (extra_flags or []),
|
|
)
|
|
train.main(train_args)
|
|
|
|
# try scalar quantization
|
|
scalar_quant_train_parser = options.get_training_parser()
|
|
scalar_quant_train_args = options.parse_args_and_arch(
|
|
scalar_quant_train_parser,
|
|
[
|
|
"--task",
|
|
"language_modeling",
|
|
data_dir,
|
|
"--arch",
|
|
arch,
|
|
"--optimizer",
|
|
"adam",
|
|
"--lr",
|
|
"0.0001",
|
|
"--criterion",
|
|
"adaptive_loss",
|
|
"--adaptive-softmax-cutoff",
|
|
"5,10,15",
|
|
"--max-tokens",
|
|
"500",
|
|
"--tokens-per-sample",
|
|
"500",
|
|
"--save-dir",
|
|
data_dir,
|
|
"--max-update",
|
|
"3",
|
|
"--no-progress-bar",
|
|
"--distributed-world-size",
|
|
"1",
|
|
"--ddp-backend",
|
|
"no_c10d",
|
|
"--num-workers",
|
|
"0",
|
|
"--quant-noise-scalar",
|
|
"0.5",
|
|
]
|
|
+ (extra_flags or []),
|
|
)
|
|
train.main(scalar_quant_train_args)
|
|
|
|
# try iterative PQ quantization
|
|
quantize_parser = options.get_training_parser()
|
|
quantize_args = options.parse_args_and_arch(
|
|
quantize_parser,
|
|
[
|
|
"--task",
|
|
"language_modeling",
|
|
data_dir,
|
|
"--arch",
|
|
arch,
|
|
"--optimizer",
|
|
"adam",
|
|
"--lr",
|
|
"0.0001",
|
|
"--criterion",
|
|
"adaptive_loss",
|
|
"--adaptive-softmax-cutoff",
|
|
"5,10,15",
|
|
"--max-tokens",
|
|
"50",
|
|
"--tokens-per-sample",
|
|
"50",
|
|
"--max-update",
|
|
"6",
|
|
"--no-progress-bar",
|
|
"--distributed-world-size",
|
|
"1",
|
|
"--ddp-backend",
|
|
"no_c10d",
|
|
"--num-workers",
|
|
"0",
|
|
"--restore-file",
|
|
os.path.join(data_dir, "checkpoint_last.pt"),
|
|
"--reset-optimizer",
|
|
"--quantization-config-path",
|
|
os.path.join(
|
|
os.path.dirname(__file__), "transformer_quantization_config.yaml"
|
|
),
|
|
]
|
|
+ (extra_flags or []),
|
|
)
|
|
train.main(quantize_args)
|
|
|
|
|
|
class TestQuantization(unittest.TestCase):
|
|
def setUp(self):
|
|
logging.disable(logging.CRITICAL)
|
|
|
|
def tearDown(self):
|
|
logging.disable(logging.NOTSET)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU")
|
|
def test_quantization(self):
|
|
with contextlib.redirect_stdout(StringIO()):
|
|
with tempfile.TemporaryDirectory("test_quantization") as data_dir:
|
|
create_dummy_data(data_dir)
|
|
preprocess_lm_data(data_dir)
|
|
# tests both scalar and iterative PQ quantization
|
|
_quantize_language_model(data_dir, "transformer_lm")
|
|
|
|
|
|
class TestOptimizersGPU(unittest.TestCase):
|
|
def setUp(self):
|
|
logging.disable(logging.CRITICAL)
|
|
|
|
def tearDown(self):
|
|
logging.disable(logging.NOTSET)
|
|
|
|
@unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU")
|
|
def test_flat_grads(self):
|
|
with contextlib.redirect_stdout(StringIO()):
|
|
with tempfile.TemporaryDirectory("test_flat_grads") as data_dir:
|
|
# Use just a bit of data and tiny model to keep this test runtime reasonable
|
|
create_dummy_data(data_dir, num_examples=10, maxlen=5)
|
|
preprocess_translation_data(data_dir)
|
|
with self.assertRaises(RuntimeError):
|
|
# adafactor isn't compatible with flat grads, which
|
|
# are used by default with --fp16
|
|
train_translation_model(
|
|
data_dir,
|
|
"lstm",
|
|
[
|
|
"--required-batch-size-multiple",
|
|
"1",
|
|
"--encoder-layers",
|
|
"1",
|
|
"--encoder-hidden-size",
|
|
"32",
|
|
"--decoder-layers",
|
|
"1",
|
|
"--optimizer",
|
|
"adafactor",
|
|
"--fp16",
|
|
],
|
|
)
|
|
# but it should pass once we set --fp16-no-flatten-grads
|
|
train_translation_model(
|
|
data_dir,
|
|
"lstm",
|
|
[
|
|
"--required-batch-size-multiple",
|
|
"1",
|
|
"--encoder-layers",
|
|
"1",
|
|
"--encoder-hidden-size",
|
|
"32",
|
|
"--decoder-layers",
|
|
"1",
|
|
"--optimizer",
|
|
"adafactor",
|
|
"--fp16",
|
|
"--fp16-no-flatten-grads",
|
|
],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|