chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,313 @@
|
||||
# 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()
|
||||
@@ -0,0 +1,28 @@
|
||||
# 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.
|
||||
|
||||
# This file defines example configuration arguments for quantizing
|
||||
# a transformer model with product quantization
|
||||
|
||||
n_centroids:
|
||||
Linear:
|
||||
key: in_features
|
||||
value: {"*": 8}
|
||||
Embedding:
|
||||
key: embedding_dim
|
||||
value: {"*": 8}
|
||||
|
||||
block_sizes:
|
||||
Linear:
|
||||
key: fuzzy_name
|
||||
value: {fc: 8, attn: 4, emb: 4}
|
||||
Embedding:
|
||||
key: fuzzy_name
|
||||
value: {emb: 8}
|
||||
|
||||
layers_to_quantize:
|
||||
- decoder\\.layers\\.\d+\\.fc[12]
|
||||
- decoder\\.embed_tokens\\.embeddings\\.[012]\\.[01]
|
||||
- decoder\\.layers\\.\d+\\.self_attn\\.(k_proj|v_proj|q_proj|out_proj)
|
||||
Reference in New Issue
Block a user