tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

casts the prepared data to int16 (doesn't help with training memory)

Browse Source
This commit is contained in:
Wing Lian
2023-04-17 21:36:02 -04:00
parent 120e7df7df
commit 2db9436410
2 changed files with 13 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
src/axolotl/datasets.py
View File

@@ -30,7 +30,6 @@ class TokenizedPromptDataset(IterableDataset):
except InvalidDataException:
pass
# TODO this isn't the best since it can't interleave datasets
class ConstantLengthDataset(IterableDataset):
"""
@@ -40,7 +39,6 @@ class ConstantLengthDataset(IterableDataset):
dataset (dataset.Dataset): Dataset with text files.
seq_length (int): Length of token sequences to return.
"""
def __init__(
self,
tokenizer,
@@ -52,6 +50,15 @@ class ConstantLengthDataset(IterableDataset):
self.datasets: List[IterableDataset] = datasets
self.seq_length = seq_length
vocab_size = len(tokenizer.get_vocab())
if vocab_size <= torch.iinfo(torch.int16).max:
self.tokens_dtype = torch.int16
elif vocab_size <= torch.iinfo(torch.int32).max:
self.tokens_dtype = torch.int32
else:
self.tokens_dtype = torch.int64
def __iter__(self):
buffer = {"input_ids": [], "attention_mask": [], "labels": []}
buffer_len = 0
@@ -105,11 +112,11 @@ class ConstantLengthDataset(IterableDataset):
attention_mask.append(1)
labels.append(self.concat_token_id)
input_ids_with_concat = torch.tensor(input_ids, dtype=torch.long)
input_ids_with_concat = torch.tensor(input_ids, dtype=self.tokens_dtype)
attention_mask_with_concat = torch.tensor(
attention_mask, dtype=torch.long
attention_mask, dtype=self.tokens_dtype
)
labels_with_concat = torch.tensor(labels, dtype=torch.long)
labels_with_concat = torch.tensor(labels, dtype=self.tokens_dtype)
buffer["input_ids"].append(input_ids_with_concat)
buffer["attention_mask"].append(attention_mask_with_concat)