diff --git a/YOCO/README.md b/YOCO/README.md
index e61e137cc..ab3f3becd 100644
--- a/YOCO/README.md
+++ b/YOCO/README.md
@@ -1,6 +1,168 @@
-# YOCO
+# You Only Cache Once: Decoder-Decoder Architectures for Large Language Models
-- May 2024: Code release
-- May 2024: release preprint [YOCO](https://arxiv.org/abs/)
+## Approach
+
+
+
+
+
+
+
",
+ pad="",
+ eos="",
+ unk="",
+ ):
+ self.symbols = [bos, pad, eos, unk]
+ self.indices = {s: i for i, s in enumerate(self.symbols)}
+ self.tokenizer_pad_to_multiple = tokenizer_pad_to_multiple
+ cl100k_base = tiktoken.get_encoding(tiktoken_model)
+ self._model = tiktoken.Encoding(
+ # If you're changing the set of special tokens, make sure to use a different name
+ # It should be clear from the name what behaviour to expect.
+ name="cl100k_im",
+ pat_str=cl100k_base._pat_str,
+ mergeable_ranks=cl100k_base._mergeable_ranks,
+ special_tokens={
+ **cl100k_base._special_tokens,
+ "": 100264,
+ "": 100265,
+ "": 100266,
+ "": 100267,
+ "": 100268,
+ "": 100269,
+ "": 100270,
+ "": 100271,
+ "": 100272,
+ "": 100273,
+ "": 100274,
+ "": 100275,
+ "": 100276,
+ "": 100277,
+ "": 100278,
+ "": 100279,
+ "": 100280,
+ "": 100281,
+ }
+ )
+
+ @property
+ def n_words(self) -> int:
+ n_words = self._model.n_vocab + len(self.symbols)
+ n_words = (n_words + self.tokenizer_pad_to_multiple - 1) // self.tokenizer_pad_to_multiple * self.tokenizer_pad_to_multiple
+ return n_words
+
+ @property
+ def bos_id(self) -> int:
+ return self.indices[""]
+
+ @property
+ def eos_id(self) -> int:
+ return self.indices[""]
+
+ @property
+ def pad_id(self) -> int:
+ return self.indices[""]
+
+ @property
+ def unk_id(self) -> int:
+ return self.indices[""]
+
+ def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
+ assert isinstance(s, str)
+ t = self._model.encode(s, allowed_special="all")
+ t = [i + len(self.symbols) for i in t]
+ if bos:
+ t = [self.bos_id, *t]
+ if eos:
+ t = [*t, self.eos_id]
+ return t
+
+ def decode(self, t: List[int]) -> str:
+ t = [i - len(self.symbols) for i in t if i >= len(self.symbols)]
+ return self._model.decode(t)
\ No newline at end of file
diff --git a/YOCO/yoco/tasks/data/utils.py b/YOCO/yoco/tasks/data/utils.py
new file mode 100644
index 000000000..fd850d73f
--- /dev/null
+++ b/YOCO/yoco/tasks/data/utils.py
@@ -0,0 +1,267 @@
+import collections
+from random import Random
+from typing import Dict, Iterable, Optional
+
+import torch
+import numpy as np
+from infinibatch import iterators
+from infinibatch.iterators import CheckpointableIterator, FixedBatchIterator, SelectManyIterator, MapIterator
+
+from fairseq.data import BaseWrapperDataset, FairseqDataset, data_utils
+
+def apply_to_sample(f, sample):
+ if hasattr(sample, "__len__") and len(sample) == 0:
+ return {}
+
+ def _apply(x):
+ if isinstance(x, np.ndarray):
+ return f(x)
+ elif isinstance(x, collections.OrderedDict):
+ # OrderedDict has attributes that needs to be preserved
+ od = collections.OrderedDict(
+ (key, _apply(value)) for key, value in x.items()
+ )
+ od.__dict__ = x.__dict__
+ return od
+ elif isinstance(x, dict):
+ return {key: _apply(value) for key, value in x.items()}
+ elif isinstance(x, list):
+ return [_apply(x) for x in x]
+ elif isinstance(x, tuple):
+ return tuple(_apply(x) for x in x)
+ elif isinstance(x, set):
+ return {_apply(x) for x in x}
+ else:
+ return x
+
+ return _apply(sample)
+
+
+class NativeCheckpointableIterator(iterators.CheckpointableIterator):
+ def __init__(self, iterable: Iterable):
+ self._input_iterable = iterable
+ self.setstate(None)
+
+ def getstate(self) -> Dict:
+ return {"num_items_yielded": self._num_items_yielded}
+
+ def setstate(self, checkpoint: Optional[Dict]):
+ self._iterator = iter(self._input_iterable)
+ self._num_items_yielded = (
+ iterators._advance_iterator(self._iterator, checkpoint["num_items_yielded"])
+ if checkpoint is not None
+ else 0
+ )
+
+ def __next__(self):
+ item = next(self._iterator)
+ self._num_items_yielded += 1
+ return item
+
+ def close(self):
+ pass
+
+
+class WeightIterator(object):
+ def __init__(self, weights, seed):
+ self.weights = weights
+ self.seed = seed
+ self.control_index = list(range(len(weights)))
+ self.setstate(None)
+
+ def __iter__(self):
+ return self
+
+ def getstate(self):
+ return {"random_state": self._random_state}
+
+ def setstate(self, checkpoint):
+ self._random_state = checkpoint["random_state"] if checkpoint else None
+ self._random = (
+ None # this will trigger the lazy initialization in self.__next__
+ )
+
+ def __next__(self):
+ if self._random is None:
+ self._random = Random(self.seed)
+ if self._random_state is not None:
+ self._random.setstate(self._random_state)
+ idx = self._random.choices(self.control_index, self.weights)[0]
+ self._random_state = self._random.getstate()
+ return idx
+
+ def close(self):
+ pass
+
+
+def FixedBlockwiseShuffleIterator(source_iterator: CheckpointableIterator, block_size: int, seed: int=0):
+ """
+ Shuffles a sequence of items by grouping consecutive items in blocks of fixed size, shuffling
+ each block, and yielding the shuffled items of all blocks as a flat sequence.
+
+ E.g. [1, 2, 3, 4, 5, 6, 7, 8] with block_size = 3 may yield [3, 1, 2, 4, 6, 5, 8, 7].
+
+ Args:
+ source_iterator: checkpointable iterator or restartable iterable over input items to shuffle
+ block_size: size of the buffer in number of items used for shuffling
+ seed: random seed used for shuffling (or None)
+ """
+ # This is implemented as a pipeline:
+ # - group N consecutive items together
+ # - shuffle them
+ # - flatten the result
+ blocks = FixedBatchIterator(source_iterator, batch_size=block_size)
+ def shuffle_block_fn(block):
+ _random = Random(seed)
+ _random.shuffle(block)
+ return block
+ shuffled_blocks = MapIterator(blocks, transform=shuffle_block_fn)
+ # samples = SelectManyNoSkipIterator(shuffled_blocks, collection_selector=lambda shuffled_block: iter(shuffled_block))
+ samples = SelectManyIterator(shuffled_blocks, collection_selector=lambda shuffled_block: iter(shuffled_block))
+ return samples
+
+
+class IndexIterator(object):
+ def __init__(self, num):
+ self.num = num
+ self.setstate(None)
+
+ def __iter__(self):
+ return self
+
+ def getstate(self):
+ return {'num_items_yielded': self._num_items_yielded}
+
+ def setstate(self, checkpoint):
+ self._num_items_yielded =checkpoint['num_items_yielded'] if checkpoint is not None else 0
+
+ def __next__(self):
+ item = self._num_items_yielded % self.num
+ self._num_items_yielded += 1
+ return item
+
+ def close(self):
+ pass
+
+
+class WeightNoRandomStateIterator(object):
+ def __init__(self, weights, seed):
+ self.weights = weights
+ self.seed = seed
+ self.control_index = list(range(len(weights)))
+ self.setstate(None)
+
+ def __iter__(self):
+ return self
+
+ def getstate(self):
+ return {'num_items_yielded': self._num_items_yielded}
+
+ def setstate(self, checkpoint):
+ self._num_items_yielded =checkpoint['num_items_yielded'] if checkpoint is not None else 0
+
+ def __next__(self):
+ self._random = Random(int(self.seed) + self._num_items_yielded)
+ idx = self._random.choices(self.control_index, self.weights)[0]
+ self._num_items_yielded += 1
+ return idx
+
+ def close(self):
+ pass
+
+
+class SelectManyNoSkipIterator(CheckpointableIterator):
+ """
+ Projects each element of a source sequence to a sequence and flattens the resulting sequences into one sequence.
+ """
+ def __init__(self, source_iterator: CheckpointableIterator, collection_selector=None):
+ """
+ Args:
+ source_iterator: iterator over the items to pass to collection_selector()
+ collection_selector: user callback that maps an item into an Iterable, whose items will be yielded.
+ The returned Iterator is used only once. Hence, it is also allowed to
+ return self-iterables, such as iterators and generator expressions.
+ If None is given, no callback is applied.
+ """
+ if not isinstance(source_iterator, CheckpointableIterator):
+ raise ValueError('source_iterator has to be a CheckpointableIterator')
+ self._source_iterator = source_iterator # type: CheckpointableIterator
+ self._collection_selector = collection_selector
+ self.setstate(None)
+
+ def getstate(self) -> Dict:
+ return {'source_state': self._source_state,
+ 'flattened_items_yielded': self._flattened_items_yielded}
+
+ def setstate(self, checkpoint: Optional[Dict]):
+ self._source_state = checkpoint['source_state'] if checkpoint else None
+ self._flattened_items_yielded = 0
+ self._source_iterator.setstate(self._source_state)
+ def _generate():
+ skip_to_checkpoint = self._flattened_items_yielded
+ # main loop over source source_items
+ for source_item in self._source_iterator:
+ if self._collection_selector is not None:
+ data = iter(self._collection_selector(source_item))
+ else:
+ data = iter(source_item)
+ self._flattened_items_yielded = 0
+ # if skip_to_checkpoint:
+ # #print("Skipping to index", skip_to_checkpoint, file=sys.stderr)
+ # self._flattened_items_yielded += _advance_iterator(data, skip_to_checkpoint)
+ # skip_to_checkpoint = 0
+ # main loop over lines
+ for item in data:
+ self._flattened_items_yielded += 1
+ yield item
+ self._source_state = self._source_iterator.getstate()
+ self._iterator = _generate()
+
+ def __next__(self):
+ return next(self._iterator)
+
+ def close(self):
+ self._source_iterator.close()
+
+
+class RawArrayDataset(FairseqDataset):
+
+ def __init__(self, dataset, datatype="token"):
+ super().__init__()
+ self.dataset = dataset
+ self.datatype = datatype
+ if hasattr(dataset, 'sizes'):
+ self._sizes = dataset.sizes
+ else:
+ try:
+ self._sizes = np.array([len(x) for x in self.dataset])
+ except:
+ self._sizes = np.array([1 for x in self.dataset])
+
+ def __getitem__(self, index):
+ if type(self.dataset[index][0]) != list:
+ if self.datatype == "token":
+ return torch.Tensor(self.dataset[index]).long()
+ else:
+ return torch.Tensor(self.dataset[index]).bool()
+ else:
+ return self.dataset[index]
+
+ def __len__(self):
+ return len(self.dataset)
+
+ def collater(self, samples):
+ if hasattr(self.dataset, 'collater'):
+ return self.dataset.collater(samples)
+ else:
+ raise NotImplementedError()
+
+ @property
+ def sizes(self):
+ return self._sizes
+
+ def num_tokens(self, index):
+ return self.dataset.num_tokens(index)
+
+ def size(self, index):
+ return self.dataset.size(index)
diff --git a/YOCO/yoco/tasks/gpt.py b/YOCO/yoco/tasks/gpt.py
new file mode 100644
index 000000000..70dd10283
--- /dev/null
+++ b/YOCO/yoco/tasks/gpt.py
@@ -0,0 +1,176 @@
+import os
+from typing import Optional
+import json
+from argparse import Namespace
+import torch
+
+from fairseq.tasks import register_task, FairseqDataclass, FairseqTask
+from dataclasses import dataclass, field
+from omegaconf import II
+
+from .data.lm_loader import LMLoader
+from .data.tiktoken_tokenizer import TiktokenTokenizer
+from .data.llama_tokenizer import LLaMATokenizer
+
+
+@dataclass
+class GPTLanguageModelingConfig(FairseqDataclass):
+ data: Optional[str] = field(
+ default=None, metadata={"help": "path to data directory"}
+ )
+ tokens_per_sample: int = field(
+ default=1024,
+ metadata={"help": "max number of tokens per sample for LM dataset"},
+ )
+ max_target_positions: Optional[int] = field(
+ default=None, metadata={"help": "max number of tokens in the target sequence"}
+ )
+ llama_model: Optional[str] = field(
+ default=None,
+ metadata={"help": "path to load tokenizer and config"},
+ )
+ tiktoken_model: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "tiktoken model to tokenize the data"
+ },
+ )
+ batch_read_ahead: int = field(
+ default=10000,
+ metadata={"help": "batch read ahead size for infinibatch"},
+ )
+ pad_to_max_len: bool = field(
+ default=False,
+ metadata={"help": "pad each sentence to max length"},
+ )
+ absolute_path: bool = field(
+ default=False,
+ metadata={"help": "use absolute path in data config"},
+ )
+ tokenizer_pad_to_multiple: int = field(
+ default=8,
+ metadata={"help": "pad to multiple of this value"},
+ )
+ seed: int = II("common.seed")
+ batch_size: Optional[int] = II("dataset.batch_size")
+
+
+@register_task('gpt', dataclass=GPTLanguageModelingConfig)
+class GPTPretrainingTask(FairseqTask):
+ def __init__(self, args, tokenizer):
+ super().__init__(args)
+ self.cfg = args
+ self.tokenizer = tokenizer
+
+ @classmethod
+ def setup_task(cls, cfg, **kwargs):
+ """Setup the task (e.g., load dictionaries).
+
+ Args:
+ args (argparse.Namespace): parsed command-line arguments
+ """
+ if cfg.llama_model is not None:
+ tokenizer = LLaMATokenizer(os.path.join(cfg.llama_model, "tokenizer.model"))
+ elif cfg.tiktoken_model is not None:
+ tokenizer = TiktokenTokenizer(cfg.tiktoken_model, cfg.tokenizer_pad_to_multiple)
+ else:
+ raise ValueError("No tokenizer model provided")
+
+ return cls(cfg, tokenizer)
+
+ def load_dataset(self, split, epoch=1, combine=False, **kwargs):
+ self.datasets[split] = {
+ 'data': json.load(open(f'{self.cfg.data}/json/{split}.json')),
+ 'data_dir': self.cfg.data,
+ 'shuffle': True if split == 'train' else False,
+ }
+ self.datasets[split] = Namespace(**self.datasets[split])
+
+ def dataset(self, split):
+ if split not in self.datasets:
+ raise KeyError("Dataset not loaded: " + split)
+
+ return self.datasets[split]
+
+ def get_batch_iterator(
+ self,
+ dataset,
+ max_tokens=None,
+ max_sentences=None,
+ max_positions=None,
+ ignore_invalid_inputs=False,
+ required_batch_size_multiple=1,
+ seed=1,
+ num_shards=1,
+ shard_id=0,
+ num_workers=0,
+ epoch=1,
+ data_buffer_size=0,
+ disable_iterator_cache=False,
+ skip_remainder_batch=False,
+ grouped_shuffling=False,
+ update_epoch_batch_itr=False
+ ):
+ return LMLoader(
+ self.cfg,
+ dataset,
+ self.tokenizer,
+ max_tokens=max_tokens,
+ max_sentences=max_sentences,
+ max_positions=max_positions,
+ ignore_invalid_inputs=ignore_invalid_inputs,
+ required_batch_size_multiple=required_batch_size_multiple,
+ seed=seed,
+ epoch=epoch,
+ num_shards=num_shards,
+ shard_id=shard_id,
+ )
+
+ def train_step(
+ self, sample, model, criterion, optimizer, update_num, ignore_grad=False
+ ):
+ """
+ Do forward and backward, and return the loss as computed by *criterion*
+ for the given *model* and *sample*.
+
+ Args:
+ sample (dict): the mini-batch. The format is defined by the
+ :class:`~fairseq.data.FairseqDataset`.
+ model (~fairseq.models.BaseFairseqModel): the model
+ criterion (~fairseq.criterions.FairseqCriterion): the criterion
+ optimizer (~fairseq.optim.FairseqOptimizer): the optimizer
+ update_num (int): the current update
+ ignore_grad (bool): multiply loss by 0 if this is set to True
+
+ Returns:
+ tuple:
+ - the loss
+ - the sample size, which is used as the denominator for the
+ gradient
+ - logging outputs to display while training
+ """
+ model.train()
+ model.set_num_updates(update_num)
+ with torch.autograd.profiler.record_function("forward"):
+ loss, sample_size, logging_output = criterion(model, sample)
+ if ignore_grad:
+ loss *= 0
+ with torch.autograd.profiler.record_function("backward"):
+ optimizer.backward(loss)
+ return loss, sample_size, logging_output
+
+ def valid_step(self, sample, model, criterion):
+ model.eval()
+ with torch.no_grad():
+ loss, sample_size, logging_output = criterion(model, sample)
+ return loss, sample_size, logging_output
+
+ @property
+ def target_dictionary(self):
+ padding_idx = self.tokenizer.pad_id
+ class Dict:
+ def pad(self):
+ return padding_idx
+ dictionary = Dict()
+ return dictionary
+
diff --git a/YOCO/yoco/tasks/harness_eval.py b/YOCO/yoco/tasks/harness_eval.py
new file mode 100644
index 000000000..0b0621aea
--- /dev/null
+++ b/YOCO/yoco/tasks/harness_eval.py
@@ -0,0 +1,151 @@
+import os
+from typing import Optional
+import logging
+
+from fairseq.data import (
+ IdDataset,
+ NumSamplesDataset,
+ NumelDataset,
+ NestedDictionaryDataset,
+ NumelDataset,
+ RightPadDataset,
+ RawLabelDataset,
+)
+
+from fairseq.tasks import register_task, FairseqDataclass, LegacyFairseqTask
+from dataclasses import dataclass, field
+
+from .data.tiktoken_tokenizer import TiktokenTokenizer
+from .data.llama_tokenizer import LLaMATokenizer
+from .data.utils import RawArrayDataset
+
+from .harness_task import HarnessAnlir1, HarnessAnlir2, HarnessAnlir3, HarnessArc_challenge, HarnessArc_easy, HarnessBoolq, HarnessCopa, HarnessOpenbookqa, HarnessPiqa, HarnessRte, HarnessWic, HarnessWinogrande, HarnessHellaswag, HarnessRecord, HarnessTruthfullqaMC1, HarnessTruthfullqaMC2, HarnessSCIQ
+from .harness_task import HarnessArc_challenge25s, HarnessHellaswag10s
+
+
+logger = logging.getLogger(__name__)
+
+task_map = {
+ "harness_anli_r1": HarnessAnlir1,
+ "harness_anli_r2": HarnessAnlir2,
+ "harness_anli_r3": HarnessAnlir3,
+ "harness_boolq": HarnessBoolq,
+ "harness_copa": HarnessCopa,
+ "harness_openbookqa": HarnessOpenbookqa,
+ "harness_piqa": HarnessPiqa,
+ "harness_rte": HarnessRte,
+ "harness_wic": HarnessWic,
+ "harness_winogrande": HarnessWinogrande,
+ "harness_hellaswag": HarnessHellaswag,
+ "harness_arc_challenge": HarnessArc_challenge,
+ "harness_arc_easy": HarnessArc_easy,
+ "harness_record": HarnessRecord,
+ "harness_truthfullqa_mc1": HarnessTruthfullqaMC1,
+ "harness_truthfullqa_mc2": HarnessTruthfullqaMC2,
+ "harness_arc_challenge_25s": HarnessArc_challenge25s,
+ "harness_hellaswag_10s": HarnessHellaswag10s,
+ "harness_sciq": HarnessSCIQ,
+}
+
+from .mmlu_task import create_mmlu_tasks
+mmlu_tasks = create_mmlu_tasks()
+task_map.update(mmlu_tasks)
+
+@dataclass
+class HarnessEvalConfig(FairseqDataclass):
+ data_dir: str = field(
+ default="/mnt/msranlp/shaohanh/data/fs_eval/harness/",
+ metadata={"help": "path to data directory"},
+ )
+ eval_data: str = field(default="", metadata={"help": "dataset name"})
+ tokens_per_sample: int = field(
+ default=2048,
+ metadata={"help": "max number of tokens per sample for LM dataset"},
+ )
+ max_target_positions: Optional[int] = field(
+ default=None, metadata={"help": "max number of tokens in the target sequence"}
+ )
+ llama_model: Optional[str] = field(
+ default=None,
+ metadata={"help": "path to load tokenizer and config"},
+ )
+ tiktoken_model: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "tiktoken model to tokenize the data"
+ },
+ )
+ tokenizer_pad_to_multiple: int = field(
+ default=8,
+ metadata={"help": "pad to multiple of this value"},
+ )
+
+
+@register_task('harness_eval', dataclass=HarnessEvalConfig)
+class HarnessEval(LegacyFairseqTask):
+
+ def __init__(self, cfg, tokenizer):
+ super().__init__(cfg)
+ self.cfg = cfg
+ self.tokenizer = tokenizer
+ self.harness_task = task_map[self.cfg.eval_data](tokenizer=self.tokenizer, data_dir=cfg.data_dir, tokens_per_sample=cfg.tokens_per_sample)
+
+ @classmethod
+ def setup_task(cls, cfg, **kwargs):
+ if cfg.llama_model is not None:
+ tokenizer = LLaMATokenizer(os.path.join(cfg.llama_model, "tokenizer.model"))
+ elif cfg.tiktoken_model is not None:
+ tokenizer = TiktokenTokenizer(cfg.tiktoken_model, cfg.tokenizer_pad_to_multiple)
+ else:
+ raise ValueError("No tokenizer model provided")
+
+ return cls(cfg, tokenizer)
+
+ def load_dataset(self, split, combine=False, **kwargs):
+ src_tokens, gpt_loss_mask, label_length, labels = self.harness_task.get_data_for_evaluation()
+
+ src_tokens = RawArrayDataset(src_tokens)
+ gpt_loss_mask = RawArrayDataset(gpt_loss_mask, datatype="mask")
+ label_length = RawLabelDataset(label_length)
+ label_ids = RawLabelDataset(labels)
+ '''
+ Input format: src_tokens + option_tokens
+ '''
+ data_dict = {
+ 'id': IdDataset(),
+ 'net_input': {
+ 'src_tokens': RightPadDataset(
+ src_tokens,
+ pad_idx=self.tokenizer.pad_id,
+ ),
+ 'gpt_loss_mask': RightPadDataset(
+ gpt_loss_mask,
+ pad_idx=False,
+ ),
+ 'label_length': label_length,
+ 'src_lengths': NumelDataset(src_tokens, reduce=False),
+ },
+ 'targets': label_ids,
+ 'nsentences': NumSamplesDataset(),
+ 'ntokens': NumelDataset(src_tokens, reduce=True),
+ }
+ dataset = NestedDictionaryDataset(
+ data_dict,
+ sizes=[src_tokens.sizes],
+ )
+
+ print('| Loaded {} with {} samples'.format(split, len(dataset)))
+
+ self.datasets[split] = dataset
+ return self.datasets[split]
+
+ @property
+ def target_dictionary(self):
+ padding_idx = self.tokenizer.pad_id
+ class Dict:
+ def pad(self):
+ return padding_idx
+ dictionary = Dict()
+ return dictionary
+
+
\ No newline at end of file
diff --git a/YOCO/yoco/tasks/harness_task.py b/YOCO/yoco/tasks/harness_task.py
new file mode 100644
index 000000000..3e87f96a2
--- /dev/null
+++ b/YOCO/yoco/tasks/harness_task.py
@@ -0,0 +1,289 @@
+import json
+import numpy as np
+
+class HarnessBaseTask:
+ def __init__(self, tokenizer, data_dir, tokens_per_sample=1024):
+ self.tokenizer = tokenizer
+ self.class_num = 1
+ self.tokens_per_sample = tokens_per_sample
+ self.base_dir = data_dir
+ self.set_dataname()
+ self.set_class_num()
+ self.dataset = self.load_data()
+
+ def load_data(self):
+ import os
+ datasets = []
+ with open(os.path.join(self.base_dir, self.dataname), "r", encoding='utf-8') as fin:
+ for line in fin:
+ obj = json.loads(line)
+ datasets.append(
+ {
+ "text": obj["ctx"] if "ctx" in obj else None,
+ "label": obj["label"] if "label" in obj else None,
+ "choices": obj["choices"] if "choices" in obj else [],
+ "gold": obj["gold"] if "gold" in obj else None,
+ "raw": obj,
+ }
+ )
+ return datasets
+
+ def set_class_num(self):
+ raise NotImplementedError
+
+ def set_dataname(self):
+ raise NotImplementedError
+
+ def preprocess_example(self, example):
+ raise NotImplementedError
+
+ def get_data_for_evaluation(self):
+ src_tokens = []
+ gpt_loss_mask = []
+ label_length = []
+ labels = []
+ cut_num = 0
+ for i, example in enumerate(self.dataset):
+ input_str, label_str, label = self.preprocess_example(example)
+ if i < 2:
+ print(f"input str is {input_str}")
+ print(f"label str is {label_str}")
+
+ for j in range(len(input_str)):
+ sub_input_str, sub_label_str = input_str[j], label_str[j]
+ input_token = self.tokenizer.encode(sub_input_str)
+ label_token = self.tokenizer.encode(sub_input_str + sub_label_str)[len(input_token):]
+ if len(input_token) + len(label_token) + 1 >= self.tokens_per_sample:
+ cut_num += 1
+ input_token = input_token[-(self.tokens_per_sample - len(label_token) - 1):]
+
+ src_tokens.append([self.tokenizer.bos_id] + input_token + label_token)
+ gpt_loss_mask.append([False] * (len(input_token) + 1) + [True] * len(label_token))
+ label_length.append(len(sub_label_str.strip()))
+ labels.append(label)
+
+ if cut_num > 0:
+ print(f"cut {cut_num} examples")
+
+ return np.array(src_tokens), np.array(gpt_loss_mask), np.array(label_length), np.array(labels)
+
+
+class HarnessAnlir1(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 3
+
+ def set_dataname(self):
+ self.dataname = "anli_r1"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * self.class_num
+ answer_str = [" True", " Neither", " False"]
+ label = example["label"]
+ return input_str, answer_str, label
+
+class HarnessAnlir2(HarnessAnlir1):
+ def set_dataname(self):
+ self.dataname = "anli_r2"
+
+class HarnessAnlir3(HarnessAnlir1):
+ def set_dataname(self):
+ self.dataname = "anli_r3"
+
+class HarnessArc_challenge(HarnessBaseTask):
+ '''
+ using harness to evaluate arc challenge
+ '''
+ def set_class_num(self):
+ self.class_num = 5
+
+ def set_dataname(self):
+ self.dataname = "arc_challenge"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * len(example["choices"])
+ answer_str = [' ' + item for item in example["choices"]]
+ label = example["gold"]
+ return input_str, answer_str, label
+
+class HarnessArc_challenge25s(HarnessBaseTask):
+ '''
+ using harness to evaluate arc challenge
+ '''
+ def set_class_num(self):
+ self.class_num = 5
+
+ def set_dataname(self):
+ self.dataname = "arc_challenge_25s"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * len(example["choices"])
+ answer_str = [' ' + item for item in example["choices"]]
+ label = example["gold"]
+ return input_str, answer_str, label
+
+class HarnessArc_easy(HarnessArc_challenge):
+ def set_class_num(self):
+ self.class_num = 5
+
+ def set_dataname(self):
+ self.dataname = "arc_easy"
+
+class HarnessBoolq(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 2
+
+ def set_dataname(self):
+ self.dataname = "boolq"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * self.class_num
+ answer_str = [" no", " yes"]
+ label = example["label"]
+ return input_str, answer_str, label
+
+class HarnessCopa(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 2
+
+ def set_dataname(self):
+ self.dataname = "copa"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * self.class_num
+ answer_str = [' ' + example['raw']['choice1'], ' ' + example['raw']['choice2']]
+ label = example["label"]
+ return input_str, answer_str, label
+
+class HarnessOpenbookqa(HarnessArc_challenge):
+ def set_class_num(self):
+ self.class_num = 4
+
+ def set_dataname(self):
+ self.dataname = "openbookqa"
+
+class HarnessPiqa(HarnessArc_challenge):
+ def set_class_num(self):
+ self.class_num = 2
+
+ def set_dataname(self):
+ self.dataname = "piqa"
+
+class HarnessRte(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 2
+
+ def set_dataname(self):
+ self.dataname = "rte"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * self.class_num
+ answer_str = [' True', ' False']
+ label = example["label"]
+ return input_str, answer_str, label
+
+class HarnessWic(HarnessRte):
+ def set_dataname(self):
+ self.dataname = "wic"
+
+class HarnessWinogrande(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 2
+
+ def set_dataname(self):
+ self.dataname = "winogrande"
+
+ def preprocess_example(self, example):
+ pronoun_loc = example['raw']['sentence'].index("_")
+ input_str = []
+ input_str.append(example['raw']['sentence'][:pronoun_loc].strip() + ' ' + example['raw']['option1'])
+ input_str.append(example['raw']['sentence'][:pronoun_loc].strip() + ' ' + example['raw']['option2'])
+ answer_str = [" " + example['raw']["sentence"][pronoun_loc + 1:].strip()] * self.class_num
+ label = int(example['raw']['answer']) - 1
+ return input_str, answer_str, label
+
+class HarnessHellaswag(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 4
+
+ def set_dataname(self):
+ self.dataname = "hellaswag"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * self.class_num
+ answer_str = [' ' + item for item in example["choices"]]
+ label = example["gold"]
+ return input_str, answer_str, label
+
+
+class HarnessHellaswag10s(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 4
+
+ def set_dataname(self):
+ self.dataname = "hellaswag_10s"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * self.class_num
+ answer_str = [' ' + item for item in example["choices"]]
+ label = example["gold"]
+ return input_str, answer_str, label
+
+
+class HarnessTruthfullqaMC1(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 1
+
+ def set_dataname(self):
+ self.dataname = "truthfulqa_mc"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * len(example["raw"]["mc1_targets"]["choices"])
+ answer_str = [' ' + item for item in example["raw"]["mc1_targets"]["choices"]]
+ label = 0 # dummy label
+ return input_str, answer_str, label
+
+
+
+class HarnessTruthfullqaMC2(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 1
+
+ def set_dataname(self):
+ self.dataname = "truthfulqa_mc"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * len(example["raw"]["mc2_targets"]["choices"])
+ answer_str = [' ' + item for item in example["raw"]["mc2_targets"]["choices"]]
+ label = 0 # dummy label
+ return input_str, answer_str, label
+
+
+class HarnessRecord(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 1
+
+ def set_dataname(self):
+ self.dataname = "record"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * len(example["raw"]["entities"])
+ answer_str = [f' - {example["raw"]["query"]}'.replace("@placeholder", item) for item in example["raw"]["entities"]]
+ label = 0 # dummy label
+ return input_str, answer_str, label
+
+class HarnessSCIQ(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 4
+
+ def set_dataname(self):
+ self.dataname = "sciq"
+
+ def preprocess_example(self, example):
+ input_str = [example["text"]] * self.class_num
+ answer_str = [' ' + example["raw"]["distractor1"],
+ ' ' + example["raw"]["distractor2"],
+ ' ' + example["raw"]["distractor3"],
+ ' ' + example["raw"]["correct_answer"]
+ ]
+ label = 3
+ return input_str, answer_str, label
\ No newline at end of file
diff --git a/YOCO/yoco/tasks/mmlu_task.py b/YOCO/yoco/tasks/mmlu_task.py
new file mode 100644
index 000000000..a93476c72
--- /dev/null
+++ b/YOCO/yoco/tasks/mmlu_task.py
@@ -0,0 +1,92 @@
+from .harness_task import HarnessBaseTask
+
+
+SUBJECTS = [
+ "abstract_algebra",
+ "anatomy",
+ "astronomy",
+ "business_ethics",
+ "clinical_knowledge",
+ "college_biology",
+ "college_chemistry",
+ "college_computer_science",
+ "college_mathematics",
+ "college_medicine",
+ "college_physics",
+ "computer_security",
+ "conceptual_physics",
+ "econometrics",
+ "electrical_engineering",
+ "elementary_mathematics",
+ "formal_logic",
+ "global_facts",
+ "high_school_biology",
+ "high_school_chemistry",
+ "high_school_computer_science",
+ "high_school_european_history",
+ "high_school_geography",
+ "high_school_government_and_politics",
+ "high_school_macroeconomics",
+ "high_school_mathematics",
+ "high_school_microeconomics",
+ "high_school_physics",
+ "high_school_psychology",
+ "high_school_statistics",
+ "high_school_us_history",
+ "high_school_world_history",
+ "human_aging",
+ "human_sexuality",
+ "international_law",
+ "jurisprudence",
+ "logical_fallacies",
+ "machine_learning",
+ "management",
+ "marketing",
+ "medical_genetics",
+ "miscellaneous",
+ "moral_disputes",
+ "moral_scenarios",
+ "nutrition",
+ "philosophy",
+ "prehistory",
+ "professional_accounting",
+ "professional_law",
+ "professional_medicine",
+ "professional_psychology",
+ "public_relations",
+ "security_studies",
+ "sociology",
+ "us_foreign_policy",
+ "virology",
+ "world_religions",
+]
+
+
+def create_mmlu_tasks():
+ """Creates a dictionary of tasks from a list of subjects
+ :return: {task_name: task}
+ e.g. {hendrycksTest-abstract_algebra: Task, hendrycksTest-anatomy: Task}
+ """
+ return {f"hendrycksTest-{sub}": create_task(f"hendrycksTest-{sub}") for sub in SUBJECTS}
+
+
+def create_task(subject):
+ class HendrycksTest(GeneralHendrycksTest):
+ def set_dataname(self):
+ self.dataname = f"{subject}"
+
+ return HendrycksTest
+
+class GeneralHendrycksTest(HarnessBaseTask):
+ def set_class_num(self):
+ self.class_num = 4
+
+ def preprocess_example(self, example):
+ # find the last occurence of "Queston:" in example["text"], and remove everything before it
+ # this is to remove the context
+ # last_question = example["text"].rfind("Question:")
+ # example["text"] = example["text"][last_question:]
+ input_str = [example["text"]] * self.class_num
+ answer_str = [' ' + item for item in example["choices"]]
+ label = example["gold"]
+ return input_str, answer_str, label
diff --git a/YOCO/yoco/tasks/pseudo.py b/YOCO/yoco/tasks/pseudo.py
new file mode 100644
index 000000000..87b51a12b
--- /dev/null
+++ b/YOCO/yoco/tasks/pseudo.py
@@ -0,0 +1,202 @@
+import os
+from typing import Optional
+import torch
+
+from fairseq.data import FairseqDataset
+from fairseq.tasks import register_task, FairseqDataclass, LegacyFairseqTask
+from dataclasses import dataclass, field
+from omegaconf import II
+
+from .data.tiktoken_tokenizer import TiktokenTokenizer
+from .data.llama_tokenizer import LLaMATokenizer
+
+
+class PseudoIterator(FairseqDataset):
+ def __init__(self, batch_size, length, vocab_size):
+ super().__init__()
+ self.batch_size = batch_size
+ self.length = length
+ self.vocab_size = vocab_size
+
+ self.epoch = 1
+ self.next_epoch_idx = 1
+ self.sharded_checkpoint = True
+ self.should_close_after_finished = True
+
+ def __iter__(self):
+ while True:
+ yield self.__next__()
+
+ def __next__(self):
+ net_input = torch.randint(size=(self.batch_size, self.length), dtype=torch.long, low=0, high=self.vocab_size - 1)
+ return {
+ "net_input": {"src_tokens": net_input},
+ "target": net_input,
+ "ntokens": self.batch_size * self.length,
+ }
+
+ def __len__(self) -> int:
+ return 819200000
+
+ def next_epoch_itr(self, **kwargs):
+ return self
+
+ @property
+ def first_batch(self):
+ return "DUMMY"
+
+ def end_of_epoch(self) -> bool:
+ return False
+
+ def state_dict(self):
+ return None
+
+ def load_state_dict(self, state_dict):
+ pass
+
+ def setstate(self, value):
+ pass
+
+ def getstate(self):
+ pass
+
+ def close(self):
+ pass
+
+@dataclass
+class PseudoConfig(FairseqDataclass):
+ tokens_per_sample: int = field(
+ default=1024,
+ metadata={"help": "max number of tokens per sample for LM dataset"},
+ )
+ max_target_positions: Optional[int] = field(
+ default=None, metadata={"help": "max number of tokens in the target sequence"}
+ )
+ llama_model: Optional[str] = field(
+ default=None,
+ metadata={"help": "path to load tokenizer and config"},
+ )
+ tiktoken_model: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "tiktoken model to tokenize the data"
+ },
+ )
+ batch_read_ahead: int = field(
+ default=10000,
+ metadata={"help": "batch read ahead size for infinibatch"},
+ )
+ pad_to_max_len: bool = field(
+ default=False,
+ metadata={"help": "pad each sentence to max length"},
+ )
+ absolute_path: bool = field(
+ default=False,
+ metadata={"help": "use absolute path in data config"},
+ )
+ tokenizer_pad_to_multiple: int = field(
+ default=8,
+ metadata={"help": "pad to multiple of this value"},
+ )
+ seed: int = II("common.seed")
+ batch_size: Optional[int] = II("dataset.batch_size")
+
+
+@register_task('pseudo', dataclass=PseudoConfig)
+class PseudoTask(LegacyFairseqTask):
+ def __init__(self, args, tokenizer):
+ super().__init__(args)
+ self.cfg = args
+ self.tokenizer = tokenizer
+
+ @classmethod
+ def setup_task(cls, cfg, **kwargs):
+ """Setup the task (e.g., load dictionaries).
+
+ Args:
+ args (argparse.Namespace): parsed command-line arguments
+ """
+ if cfg.llama_model is not None:
+ tokenizer = LLaMATokenizer(os.path.join(cfg.llama_model, "tokenizer.model"))
+ elif cfg.tiktoken_model is not None:
+ tokenizer = TiktokenTokenizer(cfg.tiktoken_model, cfg.tokenizer_pad_to_multiple)
+ else:
+ raise ValueError("No tokenizer model provided")
+
+ return cls(cfg, tokenizer)
+
+ def load_dataset(self, split, **kwargs):
+ pass
+ # self.datasets[split] = None
+
+ def dataset(self, split):
+ return None
+
+ def get_batch_iterator(
+ self,
+ dataset,
+ max_tokens=None,
+ max_sentences=None,
+ max_positions=None,
+ ignore_invalid_inputs=False,
+ required_batch_size_multiple=1,
+ seed=1,
+ num_shards=1,
+ shard_id=0,
+ num_workers=0,
+ epoch=1,
+ data_buffer_size=0,
+ disable_iterator_cache=False,
+ skip_remainder_batch=False,
+ grouped_shuffling=False,
+ update_epoch_batch_itr=False
+ ):
+ return PseudoIterator(max_sentences, self.cfg.tokens_per_sample, 10000)
+
+ def train_step(
+ self, sample, model, criterion, optimizer, update_num, ignore_grad=False
+ ):
+ """
+ Do forward and backward, and return the loss as computed by *criterion*
+ for the given *model* and *sample*.
+
+ Args:
+ sample (dict): the mini-batch. The format is defined by the
+ :class:`~fairseq.data.FairseqDataset`.
+ model (~fairseq.models.BaseFairseqModel): the model
+ criterion (~fairseq.criterions.FairseqCriterion): the criterion
+ optimizer (~fairseq.optim.FairseqOptimizer): the optimizer
+ update_num (int): the current update
+ ignore_grad (bool): multiply loss by 0 if this is set to True
+
+ Returns:
+ tuple:
+ - the loss
+ - the sample size, which is used as the denominator for the
+ gradient
+ - logging outputs to display while training
+ """
+ model.train()
+ model.set_num_updates(update_num)
+ with torch.autograd.profiler.record_function("forward"):
+ loss, sample_size, logging_output = criterion(model, sample)
+ if ignore_grad:
+ loss *= 0
+ with torch.autograd.profiler.record_function("backward"):
+ optimizer.backward(loss)
+ return loss, sample_size, logging_output
+
+ def valid_step(self, sample, model, criterion):
+ model.eval()
+ with torch.no_grad():
+ loss, sample_size, logging_output = criterion(model, sample)
+ return loss, sample_size, logging_output
+
+ @property
+ def target_dictionary(self):
+ padding_idx = self.tokenizer.pad_id
+ class Dict:
+ def pad(self):
+ return padding_idx
+ dictionary = Dict()
+ return dictionary
\ No newline at end of file
diff --git a/YOCO/yoco/train.py b/YOCO/yoco/train.py
new file mode 100644
index 000000000..ee6615d87
--- /dev/null
+++ b/YOCO/yoco/train.py
@@ -0,0 +1,7 @@
+import models
+import tasks
+import criterions
+from fairseq_cli.train import cli_main
+
+if __name__ == "__main__":
+ cli_main()
diff --git a/YOCO/yoco/validate.py b/YOCO/yoco/validate.py
new file mode 100644
index 000000000..e3815ca6a
--- /dev/null
+++ b/YOCO/yoco/validate.py
@@ -0,0 +1,294 @@
+#!/usr/bin/env python3 -u
+# 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.
+"""
+Train a new model on one or across multiple GPUs.
+"""
+import models
+import tasks
+import criterions
+
+import argparse
+import logging
+import math
+import os
+import sys
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+# We need to setup root logger before importing any fairseq libraries.
+logging.basicConfig(
+ format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+ datefmt="%Y-%m-%d %H:%M:%S",
+ level=os.environ.get("LOGLEVEL", "INFO").upper(),
+ stream=sys.stdout,
+)
+logger = logging.getLogger("fairseq_cli.train")
+
+import numpy as np
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from fairseq import checkpoint_utils, options, quantization_utils, tasks, utils
+from fairseq.data import data_utils, iterators
+from fairseq.data.plasma_utils import PlasmaStore
+from fairseq.dataclass.configs import FairseqConfig
+from fairseq.dataclass.utils import convert_namespace_to_omegaconf
+from fairseq.distributed import fsdp_enable_wrap, fsdp_wrap
+from fairseq.distributed import utils as distributed_utils
+from fairseq.file_io import PathManager
+from fairseq.logging import meters, metrics, progress_bar
+
+
+def main(cfg: FairseqConfig) -> None:
+ if isinstance(cfg, argparse.Namespace):
+ cfg = convert_namespace_to_omegaconf(cfg)
+
+ utils.import_user_module(cfg.common)
+
+ if (
+ distributed_utils.is_master(cfg.distributed_training)
+ and "job_logging_cfg" in cfg
+ ):
+ # make hydra logging work with ddp (see # see https://github.com/facebookresearch/hydra/issues/1126)
+ logging.config.dictConfig(OmegaConf.to_container(cfg.job_logging_cfg))
+
+ assert (
+ cfg.dataset.max_tokens is not None or cfg.dataset.batch_size is not None
+ ), "Must specify batch size either with --max-tokens or --batch-size"
+ metrics.reset()
+
+ if cfg.common.log_file is not None:
+ handler = logging.FileHandler(filename=cfg.common.log_file)
+ logger.addHandler(handler)
+
+ np.random.seed(cfg.common.seed)
+ utils.set_torch_seed(cfg.common.seed)
+
+ # if distributed_utils.is_master(cfg.distributed_training):
+ # checkpoint_utils.verify_checkpoint_directory(cfg.checkpoint.save_dir)
+
+ # Print args
+ logger.info(cfg)
+
+ if cfg.checkpoint.write_checkpoints_asynchronously:
+ try:
+ import iopath # noqa: F401
+ except ImportError:
+ logging.exception(
+ "Asynchronous checkpoint writing is specified but iopath is "
+ "not installed: `pip install iopath`"
+ )
+ return
+
+ # Setup task, e.g., translation, language modeling, etc.
+ task = tasks.setup_task(cfg.task)
+
+ assert cfg.criterion, "Please specify criterion to train a model"
+
+ # Build model and criterion
+ if cfg.distributed_training.ddp_backend == "fully_sharded":
+ with fsdp_enable_wrap(cfg.distributed_training):
+ model = fsdp_wrap(task.build_model(cfg.model))
+ else:
+ model = task.build_model(cfg.model)
+ criterion = task.build_criterion(cfg.criterion)
+
+ tpu = cfg.common.tpu
+ cuda = torch.cuda.is_available() and not cfg.common.cpu and not tpu
+ if cuda:
+ device = torch.device("cuda")
+ elif tpu:
+ device = utils.get_tpu_device()
+ else:
+ device = torch.device("cpu")
+ if cfg.common.fp16:
+ criterion = criterion.half()
+ model = model.half()
+ elif cfg.common.bf16:
+ criterion = criterion.to(dtype=torch.bfloat16)
+ model = model.to(dtype=torch.bfloat16)
+ criterion = criterion.to(device)
+ model = model.to(device)
+
+ logger.info(model)
+ logger.info("task: {}".format(task.__class__.__name__))
+ logger.info("model: {}".format(model.__class__.__name__))
+ logger.info("criterion: {}".format(criterion.__class__.__name__))
+ logger.info(
+ "num. shared model params: {:,} (num. trained: {:,})".format(
+ sum(
+ p.numel() for p in model.parameters() if not getattr(p, "expert", False)
+ ),
+ sum(
+ p.numel()
+ for p in model.parameters()
+ if not getattr(p, "expert", False) and p.requires_grad
+ ),
+ )
+ )
+
+ logger.info(
+ "num. expert model params: {} (num. trained: {})".format(
+ sum(p.numel() for p in model.parameters() if getattr(p, "expert", False)),
+ sum(
+ p.numel()
+ for p in model.parameters()
+ if getattr(p, "expert", False) and p.requires_grad
+ ),
+ )
+ )
+
+ # Load valid dataset (we load training data below, based on the latest checkpoint)
+ # We load the valid dataset AFTER building the model
+ data_utils.raise_if_valid_subsets_unintentionally_ignored(cfg)
+ if cfg.dataset.combine_valid_subsets:
+ task.load_dataset("valid", combine=True, epoch=1)
+ else:
+ for valid_sub_split in cfg.dataset.valid_subset.split(","):
+ task.load_dataset(valid_sub_split, combine=False, epoch=1)
+
+ # Load the latest checkpoint if one is available and restore the
+ # corresponding train iterator
+ # try:
+ # state_dict = torch.load(cfg.checkpoint.restore_file)
+ # model.load_state_dict(state_dict['model'])
+ # print(f"Loaded model from {cfg.checkpoint.restore_file}")
+ # except Exception as e:
+ # print(e)
+ # print(f"No checkpoint found from {cfg.checkpoint.restore_file}")
+
+ valid_subsets = cfg.dataset.valid_subset.split(",")
+ logger.info("Start validating")
+
+ validate(
+ cfg, task, model, criterion, valid_subsets,
+ )
+
+@torch.no_grad()
+def validate(
+ cfg: DictConfig,
+ task: tasks.FairseqTask,
+ model,
+ criterion,
+ subsets: List[str],
+) -> List[Optional[float]]:
+ """Evaluate the model on the validation set(s) and return the losses."""
+ if cfg.dataset.fixed_validation_seed is not None:
+ # set fixed seed for every validation
+ utils.set_torch_seed(cfg.dataset.fixed_validation_seed)
+
+ valid_losses = []
+ for subset in subsets:
+ logger.info('begin validation on "{}" subset'.format(subset))
+
+ # Initialize data iterator
+ itr = task.get_batch_iterator(
+ dataset=task.dataset(subset),
+ max_tokens=cfg.dataset.max_tokens_valid,
+ max_sentences=cfg.dataset.batch_size_valid,
+ required_batch_size_multiple=cfg.dataset.required_batch_size_multiple,
+ seed=cfg.common.seed,
+ num_workers=cfg.dataset.num_workers_valid,
+ # always pass a fixed "epoch" to keep validation data consistent
+ # across training epochs
+ epoch=1,
+ data_buffer_size=cfg.dataset.data_buffer_size,
+ ).next_epoch_itr(
+ shuffle=False, set_dataset_epoch=False # use a fixed valid set
+ )
+ if cfg.common.tpu:
+ itr = utils.tpu_data_loader(itr)
+ progress = progress_bar.progress_bar(
+ itr,
+ log_format=cfg.common.log_format,
+ log_interval=cfg.common.log_interval,
+ prefix=f"valid on '{subset}' subset",
+ tensorboard_logdir=(
+ cfg.common.tensorboard_logdir
+ if distributed_utils.is_master(cfg.distributed_training)
+ else None
+ ),
+ default_log_format=("tqdm" if not cfg.common.no_progress_bar else "simple"),
+ wandb_project=(
+ cfg.common.wandb_project
+ if distributed_utils.is_master(cfg.distributed_training)
+ else None
+ ),
+ wandb_run_name=os.environ.get(
+ "WANDB_NAME", os.path.basename(cfg.checkpoint.save_dir)
+ ),
+ )
+
+ # create a new root metrics aggregator so validation metrics
+ # don't pollute other aggregators (e.g., train meters)
+ with metrics.aggregate(new_root=True) as agg:
+ logging_outputs = []
+ for i, sample in enumerate(progress):
+ if cfg.dataset.max_valid_steps is not None and i > cfg.dataset.max_valid_steps:
+ break
+ sample = utils.move_to_cuda(sample)
+ _, _, inner_logging_outputs = task.valid_step(
+ sample, model, criterion
+ )
+ logging_outputs.append(inner_logging_outputs)
+ task.reduce_metrics(logging_outputs, criterion)
+
+ # with metrics.aggregate(new_root=True) as agg:
+ # for i, sample in enumerate(progress):
+ # if (
+ # cfg.dataset.max_valid_steps is not None
+ # and i > cfg.dataset.max_valid_steps
+ # ):
+ # break
+ # trainer.valid_step(sample)
+
+ stats = get_valid_stats(cfg, agg.get_smoothed_values())
+
+ progress.print(stats, tag=subset)
+
+ valid_losses.append(stats[cfg.checkpoint.best_checkpoint_metric])
+ return valid_losses
+
+
+def get_valid_stats(
+ cfg: DictConfig, stats: Dict[str, Any]
+) -> Dict[str, Any]:
+ if hasattr(checkpoint_utils.save_checkpoint, "best"):
+ key = "best_{0}".format(cfg.checkpoint.best_checkpoint_metric)
+ best_function = max if cfg.checkpoint.maximize_best_checkpoint_metric else min
+ stats[key] = best_function(
+ checkpoint_utils.save_checkpoint.best,
+ stats[cfg.checkpoint.best_checkpoint_metric],
+ )
+ return stats
+
+
+def cli_main(
+ modify_parser: Optional[Callable[[argparse.ArgumentParser], None]] = None
+) -> None:
+ parser = options.get_training_parser()
+ args = options.parse_args_and_arch(parser, modify_parser=modify_parser)
+
+ cfg = convert_namespace_to_omegaconf(args)
+
+ if cfg.common.use_plasma_view:
+ server = PlasmaStore(path=cfg.common.plasma_path)
+ logger.info(
+ f"Started plasma server pid {server.server.pid} {cfg.common.plasma_path}"
+ )
+
+ if args.profile:
+ with torch.cuda.profiler.profile():
+ with torch.autograd.profiler.emit_nvtx():
+ distributed_utils.call_main(cfg, main)
+ else:
+ distributed_utils.call_main(cfg, main)
+
+ # if cfg.common.use_plasma_view:
+ # server.server.kill()
+
+
+if __name__ == "__main__":
+ cli_main()
\ No newline at end of file