← back to Exo
Generalizing some of the dataset biz while also creating uniform batches
a6fd7a34309d32caa09621951aa80c3910f1dc69 · 2024-11-26 00:18:58 -0800 · Nel Nibcord
Files touched
M exo/main.pyM exo/train/dataset.py
Diff
commit a6fd7a34309d32caa09621951aa80c3910f1dc69
Author: Nel Nibcord <blindcrone@tuta.io>
Date: Tue Nov 26 00:18:58 2024 -0800
Generalizing some of the dataset biz while also creating uniform batches
---
exo/main.py | 38 ++++++++++++++++++-------------------
exo/train/dataset.py | 53 +++++++++++++++++++++++++++++++++-------------------
2 files changed, 52 insertions(+), 39 deletions(-)
diff --git a/exo/main.py b/exo/main.py
index da792906..1f26cd4c 100644
--- a/exo/main.py
+++ b/exo/main.py
@@ -11,6 +11,7 @@ import time
import traceback
import uuid
import numpy as np
+from functools import partial
from tqdm import tqdm
from tqdm.asyncio import tqdm_asyncio
from exo.train.dataset import load_dataset, iterate_batches
@@ -215,43 +216,41 @@ def clean_path(path):
path = path.strip('Optional("').rstrip('")')
return os.path.expanduser(path)
-async def eval_model_cli(node: Node, inference_engine: InferenceEngine, model_name, dataset, batch_size, num_batches=-1):
+async def eval_model_cli(node: Node, inference_engine: InferenceEngine, model_name, dataloader, batch_size, num_batches=-1):
inference_class = inference_engine.__class__.__name__
shard = build_base_shard(model_name, inference_class)
if not shard:
print(f"Error: Unsupported model '{model_name}' for inference engine {inference_engine.__class__.__name__}")
return
tokenizer = await resolve_tokenizer(get_repo(shard.model_id, inference_class))
- all_losses = []
- ntokens = 0
-
+ train, val, test = dataloader(tokenizer)
+ dataset = test
print(f"Evaluating {len(dataset)} examples with batch_size {batch_size}")
losses = []
tokens = []
- for batch in tqdm(iterate_batches(dataset, tokenizer, batch_size), total=len(dataset) // batch_size):
+ for batch in tqdm(iterate_batches(test, batch_size), total=len(dataset) // batch_size):
_, _, lengths = batch
losses.append(np.sum(lengths * await node.enqueue_example(shard, *batch)))
tokens.append(np.sum(lengths))
total_loss = np.sum(losses) / np.sum(tokens)
print(f"total | loss: {total_loss}, tokens: {np.sum(tokens)}")
-async def train_model_cli(node: Node, inference_engine: InferenceEngine, model_name, dataset, batch_size, iters):
+async def train_model_cli(node: Node, inference_engine: InferenceEngine, model_name, dataloader, batch_size, iters):
inference_class = inference_engine.__class__.__name__
shard = build_base_shard(model_name, inference_class)
if not shard:
print(f"Error: Unsupported model '{model_name}' for inference engine {inference_engine.__class__.__name__}")
return
tokenizer = await resolve_tokenizer(get_repo(shard.model_id, inference_class))
- all_losses = []
- ntokens = 0
-
- print(f"Training on {len(dataset)} examples with batch_size {batch_size}")
- losses = []
- tokens = []
- for batch in tqdm(iterate_batches(dataset, tokenizer, batch_size), total=len(dataset) // batch_size):
- _, _, lengths = batch
- losses.append(np.sum(lengths * await node.enqueue_example(shard, *batch)))
- tokens.append(np.sum(lengths))
+ train, val, test = dataloader(tokenizer)
+ print(f"Training on {len(train)} examples with batch_size {batch_size}")
+ for epoch in range(iters):
+ losses = []
+ tokens = []
+ for batch in tqdm(iterate_batches(train, batch_size), total=len(dataset) // batch_size):
+ _, _, lengths = batch
+ losses.append(np.sum(lengths * await node.enqueue_example(shard, *batch)))
+ tokens.append(np.sum(lengths))
total_loss = np.sum(losses) / np.sum(tokens)
print(f"total | loss: {total_loss}, tokens: {np.sum(tokens)}")
@@ -301,19 +300,18 @@ async def main():
return
await run_model_cli(node, inference_engine, model_name, args.prompt)
elif args.command == "eval" or args.command == 'train':
- data_path = args.data
- train, val, test = load_dataset(data_path)
model_name = args.model_name
+ dataloader = lambda tok: load_dataset(args.data, preprocess=lambda i: tok.encode(i["text"]))
if args.command == 'eval':
if not model_name:
print("Error: Much like a human, I can't evaluate anything without a model")
return
- await eval_model_cli(node, inference_engine, model_name, test, args.batch_size)
+ await eval_model_cli(node, inference_engine, model_name, dataloader, args.batch_size)
else:
if not model_name:
print("Error: This train ain't leaving the station without a model")
return
- await train_model_cli(node, inference_engine, model_name, test, args.batch_size, args.iters)
+ await train_model_cli(node, inference_engine, model_name, dataloader, args.batch_size, args.iters)
else:
asyncio.create_task(api.run(port=args.chatgpt_api_port)) # Start the API server as a non-blocking task
diff --git a/exo/train/dataset.py b/exo/train/dataset.py
index 2f31b4a0..8a5f1d7b 100644
--- a/exo/train/dataset.py
+++ b/exo/train/dataset.py
@@ -2,63 +2,78 @@
from pathlib import Path
import numpy as np
import json
+from functools import partial, reduce
+def compose(*funcs):
+ return reduce(lambda f, g: lambda x: f(g(x)), funcs, lambda x : x)
-def make_batch(tokens):
+def batch_with_lengths(tokens, maxlen = None):
lengths = [len(x) for x in tokens]
batch_size = len(lengths)
-
- # Check if any sequence is longer than 2048 tokens
- if max(lengths) > 2048:
- print("You've got sequences with over 2048 tokens in here! Split your data fool!")
+ if maxlen is None:
+ maxlen = max(lengths)
+ else:
+ lengths = [min(maxlen, l) for l in lengths]
# Pad to the max length
- batch_arr = np.zeros((batch_size, max(lengths)), np.int32)
+ batch_arr = np.zeros((batch_size, maxlen), np.int32)
for j in range(batch_size):
batch_arr[j, : lengths[j]] = tokens[j]
batch = np.array(batch_arr)
return batch[:, :-1], batch[:, 1:], np.array(lengths)
-def iterate_batches(dset, tokenizer, batch_size, train=False):
+def batch_chunk(batch_size):
+ return lambda d, i: d[i:i + batch_size]
+
+
+def iterate_batches(dset, batch_size, train=False, uniform_length=True):
# Shuffle indices
+ make_batch = lambda b: batch_with_lengths(b, maxlen=dset._maxlen if uniform_length else None)
+ chunk = batch_chunk(batch_size)
while True:
indices = np.arange(len(dset))
if train:
indices = np.random.permutation(indices)
+ batch = compose(make_batch, lambda i: [dset[k] for k in i], partial(chunk, indices))
# Collect batches from dataset
for i in range(0, len(indices) - batch_size + 1, batch_size):
- # Encode batch
- yield make_batch([tokenizer.encode(dset[indices[i + j]]) for j in range(batch_size)])
+ yield batch(i)
if not train:
break
class Dataset:
- """
- Light-weight wrapper to hold lines from a jsonl file
- """
-
- def __init__(self, path: Path, key: str = "text"):
+ preprocess = lambda item: item
+ load = lambda line: line
+ def __init__(self, path: Path, preprocess=None, load=None, metrics={}):
if not path.exists():
self._data = None
else:
+ if preprocess is not None:
+ self.preprocess = preprocess
+ if load is not None:
+ self.load = load
with open(path, "r") as fid:
- self._data = [json.loads(l) for l in fid]
- self._key = key
+ self._data = [load(l) for l in fid]
+ self._maxlen = max([len(self.preprocess(x)) for x in self._data])
+ # Check if any sequence is longer than 2048 tokens
+ if self._maxlen > 2048:
+ print("You've got sequences with over 2048 tokens in here! Split your data fool!")
+
def __getitem__(self, idx: int):
- return self._data[idx][self._key]
+ return self.preprocess(self._data[idx])
def __len__(self):
return len(self._data)
-def load_dataset(data_path: str):
+def load_dataset(data_path: str, preprocess=None):
def load_and_check(name):
dataset_path = Path(data_path) / f"{name}.jsonl"
try:
- return Dataset(dataset_path)
+ return Dataset(dataset_path, preprocess=preprocess, load=json.loads)
except Exception as e:
print(f"Unable to build dataset {dataset_path} ({e})")
raise
← f5efbe1b Initial distributed evaluation implementation
·
back to Exo
·
WIP: Training works on mlx 83685682 →