add download of owt2

configs/base_deepspeed.json

@@ -1,6 +1,7 @@
 {
  "train_batch_size": 8,
  "gradient_accumulation_steps": 1,
+ "gradient_clipping": 1.0,
  "tensorboard": {
  "enabled": true,
  "output_path": "./logs",
@@ -12,12 +13,20 @@
  "lr": 0.00015
  }
  },
+ "scheduler": {
+ "type": "WarmupLR",
+ "params": {
+ "warmup_min_lr": 0,
+ "warmup_max_lr": 0.00015,
+ "warmup_num_steps": 5000
+ }
+ },
  "fp16": {
  "enabled": true
  },
  "zero_optimization": {
  "stage": 2,
  "contiguous_gradients" : true,
  "cpu_offload": true
-  }
+ }
 }

configs/gpt3_small.json

@@ -5,8 +5,8 @@
  "add_padding_token": false
  },
  "dataset": {
- "name": "OWT2",
- "dir": "/root/data",
+ "name": "owt2",
+ "dir": "./data",
  "seed": 1,
  "shuffle_input_filenames": true,
  "pretokenized": true,

gpt_neox/autoregressive_wrapper.py

@@ -86,5 +86,8 @@ def forward(self, x, **kwargs):
  kwargs.update(mask = mask)
 
  out = self.net(xi, **kwargs)
- loss = F.cross_entropy(out.transpose(1, 2), xo, ignore_index = self.ignore_index)
+
+ losses = F.cross_entropy(out.transpose(1, 2), xo, reduction='none', ignore_index = self.ignore_index)
+ loss = losses.mean()
+
  return loss

gpt_neox/downloader.py

@@ -1,23 +1,24 @@
 import os
-import tarfile
 
+DATASETS = {
+ "owt2": "http:https://eaidata.bmk.sh/data/owt2_new.tar.gz",
+ "enwiki8": "http:https://eaidata.bmk.sh/data/enwik8.gz"
+}
 
-def download_dataset(dataset, dataset_dir="/root/data"):
- if dataset == "OWT2":
- _download_owt2(dataset_dir)
+
+def download_dataset(dataset, dataset_dir="./data"):
+ if DATASETS.get(dataset, False):
+ return _download_dataset(DATASETS[dataset], os.path.join(dataset_dir, dataset))
  else:
- raise NotImplementedError # TODO: tokenize text data on the fly
+ raise NotImplementedError
 
 
-def _download_owt2(dataset_dir):
- download_url = "http:https://eaidata.bmk.sh/data/owt2_new.tar.gz"
- file_name = os.path.basename(download_url)
+def _download_dataset(dataset_url, dataset_dir):
+ file_name = os.path.basename(dataset_url)
  output_path = os.path.join(dataset_dir, file_name)
 
  if not os.path.isfile(output_path):
- os.system('mkdir -p {}'.format(dir))
- os.system('wget -O {}'.format(output_path))
+ os.system('mkdir -p {}'.format(dataset_dir))
+ os.system('wget {} -O {}'.format(dataset_url, output_path))
 
- dataset_tar = tarfile.open(output_path)
- dataset_tar.extractall(dataset_dir)
- dataset_tar.close()
+ return output_path

gpt_neox/gpt_neox.py

@@ -99,7 +99,7 @@ def forward(self, x, **kwargs):
  i, j = q.shape[-2], k.shape[-2]
  bool_mask = torch.ones(i, j, device=device).triu_(j - i + 1).bool()
  mask = torch.zeros(i, j, device=device).to(q)
- mask_value = -torch.finfo(q.dtype).max
+ mask_value = -(torch.finfo(q.dtype).max / 2)
  mask.masked_fill_(bool_mask, mask_value)
 
  out = self.attn_fn(q, k, v, attn_mask=mask)

gpt_neox/utils.py

@@ -1,17 +1,14 @@
 import gzip
 import os
+import tarfile
 
 import numpy as np
 import torch
 
 
 # helpers
-def prepare_enwik8_data():
- if not os.path.isfile('./data/enwik8.gz'):
- os.system('mkdir -p ./data')
- os.system('wget http:https://eaidata.bmk.sh/data/enwik8.gz -O ./data/enwik8.gz')
-
- with gzip.open('./data/enwik8.gz') as file:
+def prepare_enwik8_data(data_path):
+ with gzip.open(data_path) as file:
  X = np.fromstring(file.read(int(95e6)), dtype=np.uint8)
  trX, vaX = np.split(X, [int(90e6)])
  data_train, data_val = torch.from_numpy(trX), torch.from_numpy(vaX)
@@ -28,6 +25,13 @@ def get_all_files(filetype, files_dir):
  return files
 
 
+def extract_tarfile(tarfile_path, extract_dir=None):
+ dataset_tar = tarfile.open(tarfile_path)
+ os.makedirs(extract_dir, exist_ok=False)
+ dataset_tar.extractall(extract_dir)
+ dataset_tar.close()
+
+
 def cycle(loader):
  while True:
  for data in loader:

train.py

@@ -1,5 +1,6 @@
 import argparse
 import json
+import os
 import random
 from collections import defaultdict
 
@@ -8,7 +9,7 @@
 from torch.utils.data import DataLoader
 from tqdm.auto import trange
 
-from gpt_neox import (GPTNeoX, AutoregressiveWrapper, GPT2Dataset,
+from gpt_neox import (GPTNeoX, AutoregressiveWrapper, GPT2Dataset, extract_tarfile,
  prepare_optimizer_parameters, get_tokenizer, download_dataset, get_all_files)
 
 
@@ -55,9 +56,10 @@ def get_params(model):
 dset_params = params["dataset"]
 assert dset_params is not None
 
-download_dataset(dataset=params["name"], dataset_dir=params["dir"])
-files = get_all_files(filetype=params["filetype"], files_dir=params["dir"])
-# TODO: SPLIT?
+data_path = download_dataset(dataset=dset_params["name"], dataset_dir=dset_params["dir"])
+data_dir = os.path.dirname(data_path)
+extract_tarfile(tarfile_path=data_path, extract_dir=data_dir)
+files = get_all_files(filetype=dset_params["filetype"], files_dir=data_dir)
 
 train_dataset = GPT2Dataset(files=files,
  seq_len=params["seq_len"],

train_enwik8.py

@@ -1,13 +1,15 @@
-from gpt_neox import (GPTNeoX, AutoregressiveWrapper, TextSamplerDataset,
- cycle, prepare_optimizer_parameters, decode_tokens, prepare_enwik8_data)
+import argparse
+import json
 import random
+from collections import defaultdict
+
+import deepspeed
 import torch
 from torch.utils.data import DataLoader
-import deepspeed
 from tqdm.auto import trange
-import argparse
-import json
-from collections import defaultdict
+
+from gpt_neox import (GPTNeoX, AutoregressiveWrapper, TextSamplerDataset, download_dataset,
+ cycle, prepare_optimizer_parameters, decode_tokens, prepare_enwik8_data)
 
 
 def get_args():
@@ -44,7 +46,8 @@ def get_params(model):
 model = AutoregressiveWrapper(model)
 
 # prepare enwik8 data
-data_train, data_val = prepare_enwik8_data()
+data_path = download_dataset(dataset="enwiki8")
+data_train, data_val = prepare_enwik8_data(data_path=data_path)
 train_dataset = TextSamplerDataset(data_train, params["seq_len"])
 val_dataset = TextSamplerDataset(data_val, params["seq_len"])
 val_loader = cycle(DataLoader(val_dataset, batch_size=params["batch_size"]))
@@ -60,7 +63,7 @@ def get_params(model):
  model=model,
  optimizer=optim,
  model_parameters=ds_model_params,
- training_data=train_dataset)
+  training_data=train_dataset)
 
 pbar = trange(params["num_epochs"], mininterval=10., desc='Training Model', dynamic_ncols=True)
 for _ in pbar: