forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ 0798119
This commit is contained in:
IQ.Lvbs history cleanup
19
tinygrad_repo/examples/mlperf/README
Normal file
19
tinygrad_repo/examples/mlperf/README
Normal file
@@ -0,0 +1,19 @@
|
||||
Each model should be a clean single file.
|
||||
They are imported from the top level `models` directory
|
||||
|
||||
It should be capable of loading weights from the reference imp.
|
||||
|
||||
We will focus on these 5 models:
|
||||
|
||||
# Resnet50-v1.5 (classic) -- 8.2 GOPS/input
|
||||
# Retinanet
|
||||
# 3D UNET (upconvs)
|
||||
# RNNT
|
||||
# BERT-large (transformer)
|
||||
|
||||
They are used in both the training and inference benchmark:
|
||||
https://mlcommons.org/en/training-normal-21/
|
||||
https://mlcommons.org/en/inference-edge-30/
|
||||
And we will submit to both.
|
||||
|
||||
NOTE: we are Edge since we don't have ECC RAM
|
||||
822
tinygrad_repo/examples/mlperf/dataloader.py
Normal file
822
tinygrad_repo/examples/mlperf/dataloader.py
Normal file
@@ -0,0 +1,822 @@
|
||||
import os, random, pickle, queue, struct, math, functools, hashlib, time
|
||||
from typing import List
|
||||
from pathlib import Path
|
||||
from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
|
||||
|
||||
import numpy as np
|
||||
from tinygrad import dtypes, Tensor
|
||||
from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
|
||||
from tinygrad.nn.state import TensorIO
|
||||
|
||||
### ResNet
|
||||
|
||||
class MyQueue:
|
||||
def __init__(self, multiple_readers=True, multiple_writers=True):
|
||||
self._reader, self._writer = connection.Pipe(duplex=False)
|
||||
self._rlock = Lock() if multiple_readers else None
|
||||
self._wlock = Lock() if multiple_writers else None
|
||||
def get(self):
|
||||
if self._rlock: self._rlock.acquire()
|
||||
ret = pickle.loads(self._reader.recv_bytes())
|
||||
if self._rlock: self._rlock.release()
|
||||
return ret
|
||||
def put(self, obj):
|
||||
if self._wlock: self._wlock.acquire()
|
||||
self._writer.send_bytes(pickle.dumps(obj))
|
||||
if self._wlock: self._wlock.release()
|
||||
|
||||
def shuffled_indices(n, seed=None):
|
||||
rng = random.Random(seed)
|
||||
indices = {}
|
||||
for i in range(n-1, -1, -1):
|
||||
j = rng.randint(0, i)
|
||||
if i not in indices: indices[i] = i
|
||||
if j not in indices: indices[j] = j
|
||||
indices[i], indices[j] = indices[j], indices[i]
|
||||
yield indices[i]
|
||||
del indices[i]
|
||||
|
||||
def loader_process(q_in, q_out, X:Tensor, seed):
|
||||
import signal
|
||||
signal.signal(signal.SIGINT, lambda _, __: exit(0))
|
||||
|
||||
from extra.datasets.imagenet import center_crop, preprocess_train
|
||||
from PIL import Image
|
||||
|
||||
with Context(DEBUG=0):
|
||||
while (_recv := q_in.get()) is not None:
|
||||
idx, fn, val = _recv
|
||||
if fn is not None:
|
||||
img = Image.open(fn)
|
||||
img = img.convert('RGB') if img.mode != "RGB" else img
|
||||
|
||||
if val:
|
||||
# eval: 76.08%, load in 0m7.366s (0m5.301s with simd)
|
||||
# sudo apt-get install libjpeg-dev
|
||||
# CC="cc -mavx2" pip install -U --force-reinstall pillow-simd
|
||||
img = center_crop(img)
|
||||
img = np.array(img)
|
||||
else:
|
||||
# reseed rng for determinism
|
||||
if seed is not None:
|
||||
np.random.seed(seed * 2 ** 10 + idx)
|
||||
random.seed(seed * 2 ** 10 + idx)
|
||||
img = preprocess_train(img)
|
||||
else:
|
||||
# pad data with training mean
|
||||
img = np.tile(np.array([[[123.68, 116.78, 103.94]]], dtype=np.uint8), (224, 224, 1))
|
||||
X[idx].flatten().assign(img.tobytes())
|
||||
q_out.put(idx)
|
||||
q_out.put(None)
|
||||
|
||||
def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_first_batch=False):
|
||||
from extra.datasets.imagenet import get_train_files, get_val_files
|
||||
files = get_val_files() if val else get_train_files()
|
||||
from extra.datasets.imagenet import get_imagenet_categories
|
||||
cir = get_imagenet_categories()
|
||||
|
||||
if pad_first_batch:
|
||||
FIRST_BATCH_PAD = round_up(len(files), batch_size) - len(files)
|
||||
else:
|
||||
FIRST_BATCH_PAD = 0
|
||||
file_count = FIRST_BATCH_PAD + len(files)
|
||||
BATCH_COUNT = min(32, file_count // batch_size)
|
||||
|
||||
def _gen():
|
||||
for _ in range(FIRST_BATCH_PAD): yield -1
|
||||
yield from shuffled_indices(len(files), seed=seed) if shuffle else iter(range(len(files)))
|
||||
gen = iter(_gen())
|
||||
|
||||
def enqueue_batch(num):
|
||||
for idx in range(num*batch_size, (num+1)*batch_size):
|
||||
fidx = next(gen)
|
||||
if fidx != -1:
|
||||
fn = files[fidx]
|
||||
q_in.put((idx, fn, val))
|
||||
Y[idx] = cir[fn.split("/")[-2]]
|
||||
else:
|
||||
# padding
|
||||
q_in.put((idx, None, val))
|
||||
Y[idx] = -1
|
||||
|
||||
shutdown = False
|
||||
class Cookie:
|
||||
def __init__(self, num): self.num = num
|
||||
def __del__(self):
|
||||
if not shutdown:
|
||||
try: enqueue_batch(self.num)
|
||||
except StopIteration: pass
|
||||
|
||||
gotten = [0]*BATCH_COUNT
|
||||
def receive_batch():
|
||||
while 1:
|
||||
num = q_out.get()//batch_size
|
||||
gotten[num] += 1
|
||||
if gotten[num] == batch_size: break
|
||||
gotten[num] = 0
|
||||
return X[num*batch_size:(num+1)*batch_size], Y[num*batch_size:(num+1)*batch_size], Cookie(num)
|
||||
|
||||
#q_in, q_out = MyQueue(multiple_writers=False), MyQueue(multiple_readers=False)
|
||||
q_in, q_out = Queue(), Queue()
|
||||
|
||||
sz = (batch_size*BATCH_COUNT, 224, 224, 3)
|
||||
shm_name = "resnet_X_val" if val else "resnet_X_train"
|
||||
if not OSX and os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
|
||||
shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(sz))
|
||||
procs = []
|
||||
|
||||
try:
|
||||
# disk:shm is slower
|
||||
if OSX: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:shm:{shm.name}")
|
||||
else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
|
||||
Y = [None] * (batch_size*BATCH_COUNT)
|
||||
|
||||
for _ in range(cpu_count()):
|
||||
p = Process(target=loader_process, args=(q_in, q_out, X, seed))
|
||||
p.daemon = True
|
||||
p.start()
|
||||
procs.append(p)
|
||||
|
||||
for bn in range(BATCH_COUNT): enqueue_batch(bn)
|
||||
|
||||
# NOTE: this is batch aligned, last ones are ignored unless pad_first_batch is True
|
||||
for _ in range(0, file_count//batch_size): yield receive_batch()
|
||||
finally:
|
||||
shutdown = True
|
||||
# empty queues
|
||||
for _ in procs: q_in.put(None)
|
||||
q_in.close()
|
||||
for _ in procs:
|
||||
while q_out.get() is not None: pass
|
||||
q_out.close()
|
||||
# shutdown processes
|
||||
for p in procs: p.join()
|
||||
shm.close()
|
||||
try:
|
||||
shm.unlink()
|
||||
except FileNotFoundError:
|
||||
# happens with BENCHMARK set
|
||||
pass
|
||||
|
||||
### BERT
|
||||
|
||||
def process_batch_bert(data: List[dict]) -> dict[str, Tensor]:
|
||||
return {
|
||||
"input_ids": Tensor(np.concatenate([s["input_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
|
||||
"input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
|
||||
"segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
|
||||
"masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
|
||||
"masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
|
||||
"masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32, device="CPU"),
|
||||
"next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
|
||||
}
|
||||
|
||||
def load_file(file: str):
|
||||
with open(file, "rb") as f:
|
||||
return pickle.load(f)
|
||||
|
||||
class InterleavedDataset:
|
||||
def __init__(self, files:List[str], cycle_length:int):
|
||||
self.dataset = files
|
||||
self.cycle_length = cycle_length
|
||||
self.queues = [queue.Queue() for _ in range(self.cycle_length)]
|
||||
for i in range(len(self.queues)): self.queues[i].queue.extend(load_file(self.dataset.pop(0)))
|
||||
self.queue_pointer = len(self.queues) - 1
|
||||
|
||||
def get(self):
|
||||
# Round-robin across queues
|
||||
try:
|
||||
self.advance()
|
||||
return self.queues[self.queue_pointer].get_nowait()
|
||||
except queue.Empty:
|
||||
self.fill(self.queue_pointer)
|
||||
return self.get()
|
||||
|
||||
def advance(self):
|
||||
self.queue_pointer = (self.queue_pointer + 1) % self.cycle_length
|
||||
|
||||
def fill(self, queue_index: int):
|
||||
try:
|
||||
file = self.dataset.pop(0)
|
||||
except IndexError:
|
||||
return
|
||||
self.queues[queue_index].queue.extend(load_file(file))
|
||||
|
||||
# Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 394
|
||||
def batch_load_train_bert(BS:int, seed:int|None=None):
|
||||
from extra.datasets.wikipedia import get_wiki_train_files
|
||||
rng = random.Random(seed)
|
||||
fs = sorted(get_wiki_train_files())
|
||||
train_files = []
|
||||
while fs: # TF shuffle
|
||||
rng.shuffle(fs)
|
||||
train_files.append(fs.pop(0))
|
||||
|
||||
cycle_length = min(getenv("NUM_CPU_THREADS", min(os.cpu_count(), 8)), len(train_files))
|
||||
assert cycle_length > 0, "cycle_length must be greater than 0"
|
||||
|
||||
dataset = InterleavedDataset(train_files, cycle_length)
|
||||
while True:
|
||||
yield process_batch_bert([dataset.get() for _ in range(BS)])
|
||||
|
||||
# Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 416
|
||||
def batch_load_val_bert(BS:int):
|
||||
file = getenv("BASEDIR", Path(__file__).parent.parents[1] / "extra" / "datasets" / "wiki") / "eval.pkl"
|
||||
dataset = load_file(file)
|
||||
idx = 0
|
||||
while True:
|
||||
start_idx = (idx * BS) % len(dataset)
|
||||
end_idx = ((idx + 1) * BS) % len(dataset)
|
||||
if start_idx < end_idx:
|
||||
yield process_batch_bert(dataset[start_idx:end_idx])
|
||||
else: # wrap around the end to the beginning of the dataset
|
||||
yield process_batch_bert(dataset[start_idx:] + dataset[:end_idx])
|
||||
idx += 1
|
||||
|
||||
### UNET3D
|
||||
|
||||
def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tensor, Y:Tensor):
|
||||
from extra.datasets.kits19 import rand_balanced_crop, rand_flip, random_brightness_augmentation, gaussian_noise
|
||||
|
||||
while (data := queue_in.get()) is not None:
|
||||
idx, fn, val = data
|
||||
case_name = os.path.basename(fn).split("_x.npy")[0]
|
||||
x, y = np.load(preprocessed_dataset_dir / f"{case_name}_x.npy"), np.load(preprocessed_dataset_dir / f"{case_name}_y.npy")
|
||||
|
||||
if not val:
|
||||
if seed is not None:
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
|
||||
x, y = rand_balanced_crop(x, y)
|
||||
x, y = rand_flip(x, y)
|
||||
x, y = x.astype(np.float32), y.astype(np.uint8)
|
||||
x = random_brightness_augmentation(x)
|
||||
x = gaussian_noise(x)
|
||||
|
||||
X[idx].flatten().assign(x.tobytes())
|
||||
Y[idx].flatten().assign(y.tobytes())
|
||||
|
||||
queue_out.put(idx)
|
||||
queue_out.put(None)
|
||||
|
||||
def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=False, shuffle:bool=True, seed=None):
|
||||
assert preprocessed_dataset_dir is not None, "run preprocess_data on kits19"
|
||||
|
||||
files = sorted(list(preprocessed_dataset_dir.glob("*_x.npy")))
|
||||
file_indices = list(range(len(files)))
|
||||
batch_count = min(32, len(files) // batch_size)
|
||||
|
||||
queue_in, queue_out = Queue(), Queue()
|
||||
procs, data_out_count = [], [0] * batch_count
|
||||
shm_name_x, shm_name_y = "unet3d_x", "unet3d_y"
|
||||
sz = (batch_size * batch_count, 1, 128, 128, 128)
|
||||
if os.path.exists(f"/dev/shm/{shm_name_x}"): os.unlink(f"/dev/shm/{shm_name_x}")
|
||||
if os.path.exists(f"/dev/shm/{shm_name_y}"): os.unlink(f"/dev/shm/{shm_name_y}")
|
||||
shm_x = shared_memory.SharedMemory(name=shm_name_x, create=True, size=prod(sz))
|
||||
shm_y = shared_memory.SharedMemory(name=shm_name_y, create=True, size=prod(sz))
|
||||
|
||||
shutdown = False
|
||||
class Cookie:
|
||||
def __init__(self, bc):
|
||||
self.bc = bc
|
||||
def __del__(self):
|
||||
if not shutdown:
|
||||
try: enqueue_batch(self.bc)
|
||||
except StopIteration: pass
|
||||
|
||||
def enqueue_batch(bc):
|
||||
for idx in range(bc * batch_size, (bc+1) * batch_size):
|
||||
fn = files[next(ds_iter)]
|
||||
queue_in.put((idx, fn, val))
|
||||
|
||||
def shuffle_indices(file_indices, seed=None):
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(file_indices)
|
||||
|
||||
if shuffle: shuffle_indices(file_indices, seed=seed)
|
||||
ds_iter = iter(file_indices)
|
||||
|
||||
try:
|
||||
X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
|
||||
Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")
|
||||
|
||||
for _ in range(cpu_count()):
|
||||
proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
|
||||
proc.daemon = True
|
||||
proc.start()
|
||||
|
||||
procs.append(proc)
|
||||
|
||||
for bc in range(batch_count):
|
||||
enqueue_batch(bc)
|
||||
|
||||
for _ in range(len(files) // batch_size):
|
||||
while True:
|
||||
bc = queue_out.get() // batch_size
|
||||
data_out_count[bc] += 1
|
||||
if data_out_count[bc] == batch_size: break
|
||||
|
||||
data_out_count[bc] = 0
|
||||
yield X[bc * batch_size:(bc + 1) * batch_size], Y[bc * batch_size:(bc + 1) * batch_size], Cookie(bc)
|
||||
finally:
|
||||
shutdown = True
|
||||
|
||||
for _ in procs: queue_in.put(None)
|
||||
queue_in.close()
|
||||
|
||||
for _ in procs:
|
||||
while queue_out.get() is not None: pass
|
||||
queue_out.close()
|
||||
|
||||
# shutdown processes
|
||||
for proc in procs: proc.join()
|
||||
|
||||
shm_x.close()
|
||||
shm_y.close()
|
||||
try:
|
||||
shm_x.unlink()
|
||||
shm_y.unlink()
|
||||
except FileNotFoundError:
|
||||
# happens with BENCHMARK set
|
||||
pass
|
||||
|
||||
### RetinaNet
|
||||
|
||||
def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue,
|
||||
imgs:Tensor, boxes:Tensor, labels:Tensor, matches:Tensor|None=None,
|
||||
anchors:Tensor|None=None, seed:int|None=None):
|
||||
from extra.datasets.openimages import image_load, random_horizontal_flip, resize
|
||||
from examples.mlperf.helpers import box_iou, find_matches, generate_anchors
|
||||
import torch
|
||||
|
||||
while (data:=queue_in.get()) is not None:
|
||||
idx, img, tgt = data
|
||||
img = image_load(base_dir, img["subset"], img["file_name"])
|
||||
|
||||
if val:
|
||||
img = resize(img)[0]
|
||||
else:
|
||||
if seed is not None:
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
|
||||
img, tgt = random_horizontal_flip(img, tgt)
|
||||
img, tgt, _ = resize(img, tgt=tgt)
|
||||
match_quality_matrix = box_iou(tgt["boxes"], (anchor := np.concatenate(generate_anchors((800, 800)))))
|
||||
match_idxs = find_matches(match_quality_matrix, allow_low_quality_matches=True)
|
||||
clipped_match_idxs = np.clip(match_idxs, 0, None)
|
||||
clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
|
||||
|
||||
boxes[idx].flatten().assign(clipped_boxes.tobytes())
|
||||
labels[idx].flatten().assign(clipped_labels.tobytes())
|
||||
matches[idx].flatten().assign(match_idxs.tobytes())
|
||||
anchors[idx].flatten().assign(anchor.tobytes())
|
||||
|
||||
imgs[idx].flatten().assign(img.tobytes())
|
||||
|
||||
queue_out.put(idx)
|
||||
queue_out.put(None)
|
||||
|
||||
def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, shuffle:bool=True, seed:int|None=None):
|
||||
def _enqueue_batch(bc):
|
||||
from extra.datasets.openimages import prepare_target
|
||||
for idx in range(bc * batch_size, (bc+1) * batch_size):
|
||||
img = dataset.loadImgs(next(dataset_iter))[0]
|
||||
ann = dataset.loadAnns(dataset.getAnnIds(img_id:=img["id"]))
|
||||
tgt = prepare_target(ann, img_id, (img["height"], img["width"]))
|
||||
|
||||
if img_ids is not None:
|
||||
img_ids[idx] = img_id
|
||||
|
||||
if img_sizes is not None:
|
||||
img_sizes[idx] = tgt["image_size"]
|
||||
|
||||
queue_in.put((idx, img, tgt))
|
||||
|
||||
def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
|
||||
shm_name = f"{shm_name}_{os.getpid()}"
|
||||
if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
|
||||
shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
|
||||
shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
|
||||
return shm, shm_tensor
|
||||
|
||||
image_ids = sorted(dataset.imgs.keys())
|
||||
batch_count = min(32, len(image_ids) // batch_size)
|
||||
|
||||
queue_in, queue_out = Queue(), Queue()
|
||||
procs, data_out_count = [], [0] * batch_count
|
||||
|
||||
shm_imgs, imgs = _setup_shared_mem("retinanet_imgs", (batch_size * batch_count, 800, 800, 3), dtypes.uint8)
|
||||
|
||||
if val:
|
||||
boxes, labels, matches, anchors = None, None, None, None
|
||||
img_ids, img_sizes = [None] * (batch_size * batch_count), [None] * (batch_size * batch_count)
|
||||
else:
|
||||
img_ids, img_sizes = None, None
|
||||
shm_boxes, boxes = _setup_shared_mem("retinanet_boxes", (batch_size * batch_count, 120087, 4), dtypes.float32)
|
||||
shm_labels, labels = _setup_shared_mem("retinanet_labels", (batch_size * batch_count, 120087), dtypes.int64)
|
||||
shm_matches, matches = _setup_shared_mem("retinanet_matches", (batch_size * batch_count, 120087), dtypes.int64)
|
||||
shm_anchors, anchors = _setup_shared_mem("retinanet_anchors", (batch_size * batch_count, 120087, 4), dtypes.float64)
|
||||
|
||||
shutdown = False
|
||||
class Cookie:
|
||||
def __init__(self, bc):
|
||||
self.bc = bc
|
||||
def __del__(self):
|
||||
if not shutdown:
|
||||
try: _enqueue_batch(self.bc)
|
||||
except StopIteration: pass
|
||||
|
||||
def shuffle_indices(indices, seed):
|
||||
rng = random.Random(seed)
|
||||
rng.shuffle(indices)
|
||||
|
||||
if shuffle: shuffle_indices(image_ids, seed=seed)
|
||||
dataset_iter = iter(image_ids)
|
||||
|
||||
try:
|
||||
for _ in range(cpu_count()):
|
||||
proc = Process(
|
||||
target=load_retinanet_data,
|
||||
args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
|
||||
kwargs={"matches": matches, "anchors": anchors, "seed": seed}
|
||||
)
|
||||
proc.daemon = True
|
||||
proc.start()
|
||||
procs.append(proc)
|
||||
|
||||
for bc in range(batch_count):
|
||||
_enqueue_batch(bc)
|
||||
|
||||
for _ in range(len(image_ids) // batch_size):
|
||||
while True:
|
||||
bc = queue_out.get() // batch_size
|
||||
data_out_count[bc] += 1
|
||||
if data_out_count[bc] == batch_size: break
|
||||
|
||||
data_out_count[bc] = 0
|
||||
|
||||
if val:
|
||||
yield (imgs[bc * batch_size:(bc + 1) * batch_size],
|
||||
img_ids[bc * batch_size:(bc + 1) * batch_size],
|
||||
img_sizes[bc * batch_size:(bc + 1) * batch_size],
|
||||
Cookie(bc))
|
||||
else:
|
||||
yield (imgs[bc * batch_size:(bc + 1) * batch_size],
|
||||
boxes[bc * batch_size:(bc + 1) * batch_size],
|
||||
labels[bc * batch_size:(bc + 1) * batch_size],
|
||||
matches[bc * batch_size:(bc + 1) * batch_size],
|
||||
anchors[bc * batch_size:(bc + 1) * batch_size],
|
||||
Cookie(bc))
|
||||
finally:
|
||||
shutdown = True
|
||||
|
||||
for _ in procs: queue_in.put(None)
|
||||
queue_in.close()
|
||||
|
||||
for _ in procs:
|
||||
while queue_out.get() is not None: pass
|
||||
queue_out.close()
|
||||
|
||||
# shutdown processes
|
||||
for proc in procs: proc.join()
|
||||
|
||||
shm_imgs.close()
|
||||
|
||||
if not val:
|
||||
shm_boxes.close()
|
||||
shm_labels.close()
|
||||
shm_matches.close()
|
||||
shm_anchors.close()
|
||||
|
||||
try:
|
||||
shm_imgs.unlink()
|
||||
|
||||
if not val:
|
||||
shm_boxes.unlink()
|
||||
shm_labels.unlink()
|
||||
shm_matches.unlink()
|
||||
shm_anchors.unlink()
|
||||
except FileNotFoundError:
|
||||
# happens with BENCHMARK set
|
||||
pass
|
||||
|
||||
# stable diffusion callbacks to match mlperf ref; declared here because they're pickled
|
||||
def filter_dataset(sample:dict): return {k:v for k,v in sample.items() if k in {'npy', 'txt'}}
|
||||
def collate(batch:list[dict]):
|
||||
ret = {"npy": [], "txt": [], "__key__": []}
|
||||
for sample in batch:
|
||||
for k,v in sample.items():
|
||||
ret[k].append(v)
|
||||
return ret
|
||||
def collate_fn(batch): return batch
|
||||
|
||||
# Reference (code): https://github.com/mlcommons/training/blob/2f4a93fb4888180755a8ef55f4b977ef8f60a89e/stable_diffusion/ldm/data/webdatasets.py, Line 55
|
||||
# Reference (params): https://github.com/mlcommons/training/blob/ab4ae1ca718d7fe62c369710a316dff18768d04b/stable_diffusion/configs/train_01x08x08.yaml, Line 107
|
||||
def batch_load_train_stable_diffusion(urls:str, BS:int):
|
||||
import webdataset
|
||||
dataset = webdataset.WebDataset(urls=urls, resampled=True, cache_size=-1, cache_dir=None)
|
||||
dataset = dataset.shuffle(size=1000)
|
||||
dataset = dataset.decode()
|
||||
dataset = dataset.map(filter_dataset)
|
||||
dataset = dataset.batched(BS, partial=False, collation_fn=collate)
|
||||
dataset = webdataset.WebLoader(dataset, batch_size=None, shuffle=False, num_workers=1, persistent_workers=True, collate_fn=collate_fn)
|
||||
|
||||
for x in dataset:
|
||||
assert isinstance(x, dict) and all(isinstance(k, str) for k in x.keys()) and all(isinstance(v, list) for v in x.values())
|
||||
assert all(isinstance(moment_mean_logvar, np.ndarray) and moment_mean_logvar.shape==(1,8,64,64) for moment_mean_logvar in x["npy"])
|
||||
assert all(isinstance(caption, str) for caption in x["txt"])
|
||||
yield x
|
||||
|
||||
# llama3
|
||||
|
||||
class BinIdxDataset:
|
||||
def __init__(self, base_path:Path):
|
||||
self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
|
||||
self.idx = TensorIO(self.idx_t)
|
||||
|
||||
# parse idx file
|
||||
magic = self.idx.read(9)
|
||||
assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
|
||||
version, = struct.unpack("<Q", self.idx.read(8))
|
||||
assert version == 1, "unsupported index version"
|
||||
dtype_code, = struct.unpack("<B", self.idx.read(1))
|
||||
self.dtype = {1:np.dtype(np.uint8), 2:np.dtype(np.int8), 3:np.dtype(np.int16), 4:np.dtype(np.int32), 5:np.dtype(np.int64), 6:np.dtype(np.float64), 7:np.dtype(np.double), 8:np.dtype(np.uint16)}[dtype_code]
|
||||
self.count, = struct.unpack("<Q", self.idx.read(8))
|
||||
doc_count, = struct.unpack("<Q", self.idx.read(8))
|
||||
|
||||
start = self.idx.tell()
|
||||
end = start + self.count * dtypes.int32.itemsize
|
||||
self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()
|
||||
|
||||
start = end
|
||||
end = start + self.count * dtypes.int64.itemsize
|
||||
self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
|
||||
|
||||
start = end
|
||||
end = start + doc_count * dtypes.int64.itemsize
|
||||
self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
|
||||
|
||||
# bin file
|
||||
self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin")).numpy()
|
||||
|
||||
def _index(self, idx) -> tuple[int, int]:
|
||||
return int(self.pointers[idx]), int(self.sizes[idx])
|
||||
|
||||
def get(self, idx, offset:int=0, length:int|None=None):
|
||||
ptr, size = self._index(idx)
|
||||
if length is None: length = size - offset
|
||||
ptr += offset * self.dtype.itemsize
|
||||
return self.bin_t[ptr:ptr+length*self.dtype.itemsize].view(self.dtype)
|
||||
|
||||
# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
|
||||
class GPTDataset:
|
||||
def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
|
||||
self.samples, self.seqlen = samples, seqlen
|
||||
self.shuffle = shuffle
|
||||
self.rng = np.random.RandomState(seed)
|
||||
|
||||
self.indexed_dataset = BinIdxDataset(base_path)
|
||||
|
||||
# check for cache
|
||||
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
|
||||
cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
|
||||
print(f"try loading GPTDataset from {cache_path}...")
|
||||
if cache_path.exists():
|
||||
print("cache found, loading...")
|
||||
with open(cache_path, "rb") as f:
|
||||
self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
|
||||
else:
|
||||
print("cache not found, building index...")
|
||||
self.doc_idx = self._build_doc_idx()
|
||||
self.sample_idx = self._build_sample_idx()
|
||||
self.shuffle_idx = self._build_shuffle_idx()
|
||||
# save cache
|
||||
with open(cache_path, "wb") as f:
|
||||
pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if idx is None:
|
||||
text = self._get(0)
|
||||
else:
|
||||
text = self._get(idx)
|
||||
|
||||
return text
|
||||
|
||||
def _get(self, idx):
|
||||
idx = self.shuffle_idx[idx]
|
||||
|
||||
doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
|
||||
doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
|
||||
|
||||
doc_ids, sample_parts = [], []
|
||||
|
||||
if doc_idx_beg == doc_idx_end:
|
||||
doc_ids.append(self.doc_idx[doc_idx_beg])
|
||||
|
||||
sample_parts.append(
|
||||
self.indexed_dataset.get(
|
||||
int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
|
||||
else:
|
||||
for i in range(doc_idx_beg, doc_idx_end + 1):
|
||||
doc_ids.append(self.doc_idx[i])
|
||||
|
||||
offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
|
||||
length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
|
||||
sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
|
||||
|
||||
# concat all parts
|
||||
text = np.concatenate(sample_parts, axis=0)
|
||||
|
||||
return text
|
||||
|
||||
@functools.cached_property
|
||||
def tokens_per_epoch(self) -> int:
|
||||
return sum(self.indexed_dataset.sizes.tolist())
|
||||
|
||||
@functools.cached_property
|
||||
def num_epochs(self) -> int:
|
||||
# we need enough epochs to cover the requested amount of tokens
|
||||
num_epochs = 1
|
||||
num_tokens = self.tokens_per_epoch
|
||||
while num_tokens < self.samples * self.seqlen:
|
||||
num_epochs += 1
|
||||
num_tokens += self.tokens_per_epoch
|
||||
return num_epochs
|
||||
|
||||
# https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
|
||||
def _build_doc_idx(self):
|
||||
print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
|
||||
st = time.perf_counter()
|
||||
# doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
|
||||
doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
|
||||
doc_idx = doc_idx.astype(np.int32)
|
||||
at = time.perf_counter()
|
||||
if self.shuffle: self.rng.shuffle(doc_idx)
|
||||
print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
|
||||
return doc_idx
|
||||
|
||||
def _build_sample_idx(self):
|
||||
print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
|
||||
sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
|
||||
sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)
|
||||
|
||||
sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
|
||||
sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
|
||||
sample_idx_idx += 1
|
||||
|
||||
for _ in tqdm(range(1, self.samples + 1)):
|
||||
remaining_seqlen = self.seqlen + 1
|
||||
while remaining_seqlen > 0:
|
||||
doc_idx = int(self.doc_idx[doc_idx_idx])
|
||||
doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
|
||||
remaining_seqlen -= doc_len
|
||||
if remaining_seqlen <= 0:
|
||||
doc_offset += remaining_seqlen + doc_len - 1
|
||||
remaining_seqlen = 0
|
||||
else:
|
||||
if doc_idx_idx == len(self.doc_idx) - 1:
|
||||
assert sample_idx_idx == self.samples
|
||||
doc_idx = int(self.doc_idx[doc_idx_idx])
|
||||
doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
|
||||
break
|
||||
doc_idx_idx += 1
|
||||
doc_offset = 0
|
||||
|
||||
sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
|
||||
sample_idx_idx += 1
|
||||
|
||||
return sample_idx
|
||||
|
||||
def _build_shuffle_idx(self):
|
||||
print(f"building shuffle_idx for {self.samples=}")
|
||||
st = time.perf_counter()
|
||||
shuffle_idx = np.arange(self.samples, dtype=np.int32)
|
||||
at = time.perf_counter()
|
||||
if self.shuffle: self.rng.shuffle(shuffle_idx)
|
||||
print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
|
||||
return shuffle_idx
|
||||
|
||||
class BlendedGPTDataset:
|
||||
def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
|
||||
self.shuffle = shuffle
|
||||
self.rng = np.random.RandomState(seed)
|
||||
|
||||
# normalize weights
|
||||
total_weight = sum(weights)
|
||||
self.weights = [w / total_weight for w in weights]
|
||||
|
||||
self.samples = samples
|
||||
surplus = 0.005
|
||||
samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
|
||||
|
||||
self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
|
||||
|
||||
# check for cache
|
||||
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
|
||||
cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
|
||||
print(f"try loading BlendedGPTDataset from {cache_path}...")
|
||||
if cache_path.exists():
|
||||
print("cache found, loading...")
|
||||
with open(cache_path, "rb") as f:
|
||||
self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
|
||||
else:
|
||||
print("cache not found, building index...")
|
||||
self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
|
||||
# save cache
|
||||
with open(cache_path, "wb") as f:
|
||||
pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
|
||||
|
||||
def get(self, idx:int):
|
||||
tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
|
||||
return tokens
|
||||
|
||||
def _build_blend_idx(self):
|
||||
dataset_idx = np.zeros(self.samples, dtype=np.int16)
|
||||
dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
|
||||
|
||||
unspent_datasets = set(range(len(self.datasets)))
|
||||
dataset_sample_counts = [0] * len(self.datasets)
|
||||
|
||||
for i in tqdm(range(self.samples)):
|
||||
error_argmax, error_max = 0, 0.0
|
||||
for di in unspent_datasets:
|
||||
error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
|
||||
if error > error_max:
|
||||
error_max = error
|
||||
error_argmax = di
|
||||
|
||||
dataset_idx[i] = error_argmax
|
||||
dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
|
||||
|
||||
dataset_sample_counts[error_argmax] += 1
|
||||
|
||||
return dataset_idx, dataset_sample_idx
|
||||
|
||||
def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
|
||||
if small:
|
||||
if val:
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
|
||||
if val:
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
|
||||
return BlendedGPTDataset(
|
||||
[base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
|
||||
|
||||
def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
|
||||
for b in range(math.ceil(dataset.samples / bs)):
|
||||
batch = [dataset.get(b * bs + i) for i in range(bs)]
|
||||
stacked = np.stack(batch, axis=0)
|
||||
yield Tensor(stacked, device="NPY")
|
||||
|
||||
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
|
||||
return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
|
||||
|
||||
if __name__ == "__main__":
|
||||
def load_unet3d(val):
|
||||
assert not val, "validation set is not supported due to different sizes on inputs"
|
||||
|
||||
from extra.datasets.kits19 import get_train_files, get_val_files, preprocess_dataset, TRAIN_PREPROCESSED_DIR, VAL_PREPROCESSED_DIR
|
||||
preprocessed_dir = VAL_PREPROCESSED_DIR if val else TRAIN_PREPROCESSED_DIR
|
||||
files = get_val_files() if val else get_train_files()
|
||||
|
||||
if not preprocessed_dir.exists(): preprocess_dataset(files, preprocessed_dir, val)
|
||||
with tqdm(total=len(files)) as pbar:
|
||||
for x, _, _ in batch_load_unet3d(preprocessed_dir, val=val):
|
||||
pbar.update(x.shape[0])
|
||||
|
||||
def load_resnet(val):
|
||||
from extra.datasets.imagenet import get_train_files, get_val_files
|
||||
files = get_val_files() if val else get_train_files()
|
||||
with tqdm(total=len(files)) as pbar:
|
||||
for x,y,c in batch_load_resnet(val=val):
|
||||
pbar.update(x.shape[0])
|
||||
|
||||
def load_retinanet(val):
|
||||
from extra.datasets.openimages import BASEDIR, download_dataset
|
||||
from pycocotools.coco import COCO
|
||||
dataset = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), "validation" if val else "train"))
|
||||
with tqdm(total=len(dataset.imgs.keys())) as pbar:
|
||||
for x in batch_load_retinanet(dataset, val, base_dir):
|
||||
pbar.update(x[0].shape[0])
|
||||
|
||||
def load_llama3(val):
|
||||
bs = 24
|
||||
samples = 5760 if val else 1_200_000 * 1152
|
||||
seqlen = 8192
|
||||
|
||||
max_, min_ = 0, math.inf
|
||||
for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
|
||||
max_ = max(max_, tokens.shape[1])
|
||||
min_ = min(min_, tokens.shape[1])
|
||||
print(f"max seq length: {max_}")
|
||||
print(f"min seq length: {min_}")
|
||||
|
||||
load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
|
||||
if load_fn_name in globals():
|
||||
globals()[load_fn_name](getenv("VAL", 1))
|
||||
367
tinygrad_repo/examples/mlperf/helpers.py
Normal file
367
tinygrad_repo/examples/mlperf/helpers.py
Normal file
@@ -0,0 +1,367 @@
|
||||
from collections import OrderedDict
|
||||
import unicodedata
|
||||
from typing import Optional
|
||||
import math
|
||||
import numpy as np
|
||||
from tinygrad.nn import state
|
||||
from tinygrad.tensor import Tensor, dtypes
|
||||
from tinygrad.helpers import getenv
|
||||
|
||||
#
|
||||
# checkpointing utils
|
||||
#
|
||||
|
||||
def invert_dict(d): return {v: k for k, v in reversed(d.items())}
|
||||
def dedup_dict(d): return invert_dict(invert_dict(d))
|
||||
# store each tensor into the first key it appears in
|
||||
def get_training_state(model, optimizer, scheduler):
|
||||
# hack: let get_state_dict walk the tree starting with model, so that the checkpoint keys are
|
||||
# readable and can be loaded as a model for eval
|
||||
train_state = {'model': model, 'optimizer': optimizer, 'scheduler': scheduler}
|
||||
return dedup_dict(state.get_state_dict(train_state))
|
||||
def load_training_state(model, optimizer, scheduler, state_dict):
|
||||
# use fresh model to restore duplicate keys
|
||||
train_state = {'model': model, 'optimizer': optimizer, 'scheduler': scheduler}
|
||||
big_dict = state.get_state_dict(train_state)
|
||||
# hack: put back the dupes
|
||||
dupe_names = {}
|
||||
for k, v in big_dict.items():
|
||||
if v not in dupe_names:
|
||||
dupe_names[v] = k
|
||||
assert k in state_dict
|
||||
state_dict[k] = state_dict[dupe_names[v]]
|
||||
# scheduler contains optimizer and all params, load each weight only once
|
||||
scheduler_state = {'scheduler': scheduler}
|
||||
state.load_state_dict(scheduler_state, state_dict)
|
||||
|
||||
def gaussian_kernel(n, std):
|
||||
from scipy import signal
|
||||
gaussian_1d = signal.windows.gaussian(n, std)
|
||||
gaussian_2d = np.outer(gaussian_1d, gaussian_1d)
|
||||
gaussian_3d = np.outer(gaussian_2d, gaussian_1d)
|
||||
gaussian_3d = gaussian_3d.reshape(n, n, n)
|
||||
gaussian_3d = np.cbrt(gaussian_3d)
|
||||
gaussian_3d /= gaussian_3d.max()
|
||||
return gaussian_3d
|
||||
|
||||
def prepare_arrays(image, roi_shape=(128, 128, 128)):
|
||||
assert len(roi_shape) == 3 and any(roi_shape)
|
||||
image_shape = list(image.shape[2:])
|
||||
result = np.zeros((1, 3, *image_shape), dtype=image.dtype)
|
||||
norm_map = np.zeros_like(result)
|
||||
norm_patch = gaussian_kernel(roi_shape[0], 0.125 * roi_shape[0]).astype(norm_map.dtype)
|
||||
return result, norm_map, norm_patch
|
||||
|
||||
def get_slice(image, roi_shape=(128, 128, 128), overlap_factor=0.5):
|
||||
assert len(roi_shape) == 3 and any(roi_shape)
|
||||
assert 0 < overlap_factor < 1
|
||||
image_shape, dim = list(image.shape[2:]), len(image.shape[2:])
|
||||
strides = [int(roi_shape[i] * (1 - overlap_factor)) for i in range(dim)]
|
||||
size = [(image_shape[i] - roi_shape[i]) // strides[i] + 1 for i in range(dim)]
|
||||
for i in range(0, strides[0] * size[0], strides[0]):
|
||||
for j in range(0, strides[1] * size[1], strides[1]):
|
||||
for k in range(0, strides[2] * size[2], strides[2]):
|
||||
yield i, j, k
|
||||
|
||||
def _get_best_indices(logits, n_best_size):
|
||||
index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True)
|
||||
return list(map(lambda x: x[0], index_and_score))[:n_best_size]
|
||||
|
||||
def _is_punctuation(char):
|
||||
if (cp := ord(char)) in range(33, 48) or cp in range(58, 65) or cp in range(91, 97) or cp in range(123, 127):
|
||||
return True
|
||||
return unicodedata.category(char).startswith("P")
|
||||
|
||||
def _is_whitespace(char):
|
||||
if char == " " or char == "\t" or char == "\n" or char == "\r":
|
||||
return True
|
||||
return unicodedata.category(char) == "Zs"
|
||||
|
||||
def _is_control(char):
|
||||
if char == "\t" or char == "\n" or char == "\r":
|
||||
return False
|
||||
return unicodedata.category(char).startswith("C")
|
||||
|
||||
def _run_split_on_punc(text):
|
||||
if text in ("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"):
|
||||
return [text]
|
||||
start_new_word = True
|
||||
output = []
|
||||
for i in range(len(text)):
|
||||
if _is_punctuation(char := text[i]):
|
||||
output.append([char])
|
||||
start_new_word = True
|
||||
else:
|
||||
if start_new_word:
|
||||
output.append([])
|
||||
start_new_word = False
|
||||
output[-1].append(char)
|
||||
return ["".join(x) for x in output]
|
||||
|
||||
def _run_strip_accents(text):
|
||||
output = []
|
||||
for char in unicodedata.normalize("NFD", text):
|
||||
if unicodedata.category(char) != "Mn":
|
||||
output.append(char)
|
||||
return "".join(output)
|
||||
|
||||
def _clean_text(text):
|
||||
output = []
|
||||
for char in text:
|
||||
if not ((cp := ord(char)) == 0 or cp == 0xfffd or _is_control(char)):
|
||||
output.append(" " if _is_whitespace(char) else char)
|
||||
return "".join(output)
|
||||
|
||||
def _get_final_text(pred_text, orig_text):
|
||||
def _strip_spaces(text):
|
||||
ns_text = ""
|
||||
ns_to_s_map = OrderedDict()
|
||||
for i, c in enumerate(text):
|
||||
if c == " ":
|
||||
continue
|
||||
ns_to_s_map[len(ns_text)] = i
|
||||
ns_text += c
|
||||
return ns_text, ns_to_s_map
|
||||
|
||||
orig_tokens = _clean_text(orig_text).strip().split()
|
||||
split_tokens = []
|
||||
for token in orig_tokens:
|
||||
if token not in ("[UNK]", "[SEP]", "[PAD]", "[CLS]", "[MASK]"):
|
||||
token = token.lower()
|
||||
token = _run_strip_accents(token)
|
||||
split_tokens.extend(_run_split_on_punc(token))
|
||||
|
||||
tok_text = " ".join(" ".join(split_tokens).strip().split())
|
||||
start_position = tok_text.find(pred_text)
|
||||
if start_position == -1:
|
||||
return orig_text
|
||||
end_position = start_position + len(pred_text) - 1
|
||||
|
||||
orig_ns_text, orig_ns_to_s_map = _strip_spaces(orig_text)
|
||||
tok_ns_text, tok_ns_to_s_map = _strip_spaces(tok_text)
|
||||
if len(orig_ns_text) != len(tok_ns_text):
|
||||
return orig_text
|
||||
tok_s_to_ns_map = {v: k for k, v in tok_ns_to_s_map.items()}
|
||||
|
||||
orig_start_position = None
|
||||
if start_position in tok_s_to_ns_map:
|
||||
if (ns_start_position := tok_s_to_ns_map[start_position]) in orig_ns_to_s_map:
|
||||
orig_start_position = orig_ns_to_s_map[ns_start_position]
|
||||
if orig_start_position is None:
|
||||
return orig_text
|
||||
|
||||
orig_end_position = None
|
||||
if end_position in tok_s_to_ns_map:
|
||||
if (ns_end_position := tok_s_to_ns_map[end_position]) in orig_ns_to_s_map:
|
||||
orig_end_position = orig_ns_to_s_map[ns_end_position]
|
||||
if orig_end_position is None:
|
||||
return orig_text
|
||||
|
||||
output_text = orig_text[orig_start_position:(orig_end_position + 1)]
|
||||
return output_text
|
||||
|
||||
def get_bert_qa_prediction(features, example, start_end_logits):
|
||||
prelim_predictions = []
|
||||
for i, feature in enumerate(features):
|
||||
for start_index in _get_best_indices(start_end_logits[i][0], 20):
|
||||
for end_index in _get_best_indices(start_end_logits[i][1], 20):
|
||||
if start_index >= len(feature["tokens"]) or end_index >= len(feature["tokens"]):
|
||||
continue
|
||||
if start_index not in feature["token_to_orig_map"] or end_index not in feature["token_to_orig_map"]:
|
||||
continue
|
||||
if not feature["token_is_max_context"].get(start_index, False):
|
||||
continue
|
||||
if end_index < start_index or end_index - start_index + 1 > 30:
|
||||
continue
|
||||
|
||||
prelim_predictions.append({
|
||||
"feature_index": i,
|
||||
"start_index": start_index,
|
||||
"end_index": end_index,
|
||||
"start_logit": start_end_logits[i][0, start_index],
|
||||
"end_logit": start_end_logits[i][1, end_index]
|
||||
})
|
||||
predictions = sorted(prelim_predictions, key=lambda x: (x["start_logit"] + x["end_logit"]), reverse=True)
|
||||
|
||||
if len(predictions) > 0:
|
||||
feature = features[predictions[0]["feature_index"]]
|
||||
tok_tokens = feature["tokens"][predictions[0]["start_index"]:(predictions[0]["end_index"] + 1)]
|
||||
orig_doc_start = feature["token_to_orig_map"][predictions[0]["start_index"]]
|
||||
orig_doc_end = feature["token_to_orig_map"][predictions[0]["end_index"]]
|
||||
orig_tokens = example["context"][orig_doc_start:(orig_doc_end + 1)]
|
||||
tok_text = " ".join(tok_tokens).replace(" ##", "").replace("##", "")
|
||||
tok_text = " ".join(tok_text.strip().split())
|
||||
orig_text = " ".join(orig_tokens)
|
||||
return _get_final_text(tok_text, orig_text)
|
||||
return "empty"
|
||||
|
||||
def get_mlperf_bert_config():
|
||||
"""benchmark is BERT-large"""
|
||||
ret = {"attention_probs_dropout_prob": 0.1, "hidden_dropout_prob": 0.1, "vocab_size": 30522, "type_vocab_size": 2, "max_position_embeddings": 512}
|
||||
|
||||
match (bert_size:=getenv("BERT_SIZE", "large")):
|
||||
case "large": ret.update({"hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16, "num_hidden_layers": 24})
|
||||
case "tiny": ret.update({"hidden_size": 128, "intermediate_size": 512, "num_attention_heads": 2, "num_hidden_layers": 2})
|
||||
case _: raise RuntimeError(f"unhandled {bert_size=}")
|
||||
|
||||
if (bert_layers:=getenv("BERT_LAYERS")): ret["num_hidden_layers"] = bert_layers
|
||||
return ret
|
||||
|
||||
def get_mlperf_bert_model():
|
||||
from extra.models import bert
|
||||
from examples.mlperf.initializers import LinearBert, EmbeddingBert, LayerNormBert
|
||||
|
||||
bert.Linear = LinearBert
|
||||
bert.Embedding = EmbeddingBert
|
||||
bert.LayerNorm = LayerNormBert
|
||||
|
||||
from extra.models.bert import BertForPretraining
|
||||
config = get_mlperf_bert_config()
|
||||
if getenv("DISABLE_DROPOUT", 0):
|
||||
config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
|
||||
model = BertForPretraining(**config)
|
||||
if getenv("FP8_TRAIN"):
|
||||
from extra.fp8.fp8_linear import convert_to_float8_training
|
||||
def module_filter_fn(mod, fqn):
|
||||
if isinstance(mod, LinearBert):
|
||||
skip_layers = [] if (ln:=config["num_hidden_layers"]) <= 2 else ["bert.encoder.layer.0.", f"bert.encoder.layer.{ln-1}"]
|
||||
if mod.weight.shape[-1] >= 1024 and "encoder" in fqn and not any(name in fqn for name in skip_layers):
|
||||
print(f"replacing linear with fp8: {fqn} {mod.weight.shape}")
|
||||
return True
|
||||
return False
|
||||
convert_to_float8_training(model, module_filter_fn)
|
||||
return model
|
||||
|
||||
def get_fake_data_bert(BS:int):
|
||||
return {
|
||||
"input_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
|
||||
"input_mask": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
|
||||
"segment_ids": Tensor.zeros((BS, 512), dtype=dtypes.int32, device="CPU").contiguous(),
|
||||
"masked_lm_positions": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
|
||||
"masked_lm_ids": Tensor.zeros((BS, 76), dtype=dtypes.int32, device="CPU").contiguous(),
|
||||
"masked_lm_weights": Tensor.zeros((BS, 76), dtype=dtypes.float32, device="CPU").contiguous(),
|
||||
"next_sentence_labels": Tensor.zeros((BS, 1), dtype=dtypes.int32, device="CPU").contiguous(),
|
||||
}
|
||||
|
||||
def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
|
||||
BELOW_LOW_THRESHOLD, BETWEEN_THRESHOLDS = -1, -2
|
||||
|
||||
def _set_low_quality_matches_(matches:np.ndarray, all_matches:np.ndarray, match_quality_matrix:np.ndarray):
|
||||
highest_quality_foreach_gt = np.max(match_quality_matrix, axis=1)
|
||||
pred_inds_to_update = np.nonzero(match_quality_matrix == highest_quality_foreach_gt[:, None])[1]
|
||||
matches[pred_inds_to_update] = all_matches[pred_inds_to_update]
|
||||
|
||||
assert low_threshold <= high_threshold
|
||||
|
||||
matched_vals, matches = match_quality_matrix.max(axis=0), match_quality_matrix.argmax(axis=0)
|
||||
all_matches = np.copy(matches) if allow_low_quality_matches else None
|
||||
below_low_threshold = matched_vals < low_threshold
|
||||
between_thresholds = (matched_vals >= low_threshold) & (matched_vals < high_threshold)
|
||||
matches[below_low_threshold] = BELOW_LOW_THRESHOLD
|
||||
matches[between_thresholds] = BETWEEN_THRESHOLDS
|
||||
|
||||
if allow_low_quality_matches:
|
||||
assert all_matches is not None
|
||||
_set_low_quality_matches_(matches, all_matches, match_quality_matrix)
|
||||
|
||||
return matches
|
||||
|
||||
def box_iou(boxes1:np.ndarray, boxes2:np.ndarray) -> np.ndarray:
|
||||
def _box_area(boxes:np.ndarray) -> np.ndarray: return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
|
||||
|
||||
def _box_inter_union(boxes1:np.ndarray, boxes2:np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||||
area1, area2 = _box_area(boxes1), _box_area(boxes2)
|
||||
lt, rb = np.maximum(boxes1[:, None, :2], boxes2[:, :2]), np.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
|
||||
wh = np.clip(rb - lt, a_min=0, a_max=None)
|
||||
inter = wh[:, :, 0] * wh[:, :, 1]
|
||||
union = area1[:, None] + area2 - inter
|
||||
return inter, union
|
||||
|
||||
inter, union = _box_inter_union(boxes1, boxes2)
|
||||
return inter / union
|
||||
|
||||
def generate_anchors(input_size:tuple[int, int], scales:Optional[tuple[Tensor, ...]]=None, aspect_ratios:Optional[tuple[Tensor, ...]]=None) -> list[np.ndarray]:
|
||||
def _compute_grid_sizes(input_size:tuple[int, int]) -> np.ndarray:
|
||||
return np.ceil(np.array(input_size)[None, :] / 2 ** np.arange(3, 8)[:, None])
|
||||
|
||||
scales = tuple((i, int(i * 2 ** (1/3)), int(i * 2 ** (2/3))) for i in 2 ** np.arange(5, 10)) if scales is None else scales
|
||||
aspect_ratios = ((0.5, 1.0, 2.0),) * len(scales) if aspect_ratios is None else aspect_ratios
|
||||
aspect_ratios = tuple(ar for ar in aspect_ratios)
|
||||
grid_sizes = _compute_grid_sizes(input_size)
|
||||
|
||||
assert len(scales) == len(aspect_ratios) == len(grid_sizes), "scales, aspect_ratios, and grid_sizes must have the same length"
|
||||
|
||||
anchors = []
|
||||
for s, ar, gs in zip(scales, aspect_ratios, grid_sizes):
|
||||
s, ar = np.array(s), np.array(ar)
|
||||
h_ratios = np.sqrt(ar)
|
||||
w_ratios = 1 / h_ratios
|
||||
ws = (w_ratios[:, None] * s[None, :]).reshape(-1)
|
||||
hs = (h_ratios[:, None] * s[None, :]).reshape(-1)
|
||||
base_anchors = (np.stack([-ws, -hs, ws, hs], axis=1) / 2).round()
|
||||
stride_h, stride_w = input_size[0] // gs[0], input_size[1] // gs[1]
|
||||
shifts_x, shifts_y = np.meshgrid(np.arange(gs[1]) * stride_w, np.arange(gs[0]) * stride_h)
|
||||
shifts_x, shifts_y = shifts_x.reshape(-1), shifts_y.reshape(-1)
|
||||
shifts = np.stack([shifts_x, shifts_y, shifts_x, shifts_y], axis=1, dtype=np.float32)
|
||||
anchors.append((shifts[:, None] + base_anchors[None, :]).reshape(-1, 4))
|
||||
|
||||
return anchors
|
||||
|
||||
|
||||
class BoxCoder(object):
|
||||
def __init__(self, weights, bbox_xform_clip=math.log(1000. / 16), apply_to_remove=True):
|
||||
self.weights = weights
|
||||
self.bbox_xform_clip = bbox_xform_clip
|
||||
self.apply_to_remove = apply_to_remove
|
||||
|
||||
def encode(self, reference_boxes, proposals):
|
||||
TO_REMOVE = self.apply_to_remove # TODO remove
|
||||
ex_widths = proposals[..., 2] - proposals[..., 0] + TO_REMOVE
|
||||
ex_heights = proposals[..., 3] - proposals[..., 1] + TO_REMOVE
|
||||
ex_ctr_x = proposals[..., 0] + 0.5 * ex_widths
|
||||
ex_ctr_y = proposals[..., 1] + 0.5 * ex_heights
|
||||
|
||||
gt_widths = reference_boxes[..., 2] - reference_boxes[..., 0] + TO_REMOVE
|
||||
gt_heights = reference_boxes[..., 3] - reference_boxes[..., 1] + TO_REMOVE
|
||||
gt_ctr_x = reference_boxes[..., 0] + 0.5 * gt_widths
|
||||
gt_ctr_y = reference_boxes[..., 1] + 0.5 * gt_heights
|
||||
|
||||
wx, wy, ww, wh = self.weights
|
||||
targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths
|
||||
targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights
|
||||
targets_dw = ww * Tensor.log(gt_widths / ex_widths)
|
||||
targets_dh = wh * Tensor.log(gt_heights / ex_heights)
|
||||
|
||||
targets = Tensor.stack(targets_dx, targets_dy, targets_dw, targets_dh, dim=-1)
|
||||
return targets
|
||||
|
||||
def decode(self, rel_codes, boxes):
|
||||
boxes = boxes.cast(rel_codes.dtype)
|
||||
rel_codes = rel_codes
|
||||
|
||||
TO_REMOVE = self.apply_to_remove # TODO remove
|
||||
widths = boxes[:, 2] - boxes[:, 0] + TO_REMOVE
|
||||
heights = boxes[:, 3] - boxes[:, 1] + TO_REMOVE
|
||||
ctr_x = boxes[:, 0] + 0.5 * widths
|
||||
ctr_y = boxes[:, 1] + 0.5 * heights
|
||||
|
||||
wx, wy, ww, wh = self.weights
|
||||
dx = rel_codes[:, 0::4] / wx
|
||||
dy = rel_codes[:, 1::4] / wy
|
||||
dw = rel_codes[:, 2::4] / ww
|
||||
dh = rel_codes[:, 3::4] / wh
|
||||
|
||||
# Prevent sending too large values into Tensor.exp()
|
||||
dw = dw.clip(min_=dw.min(), max_=self.bbox_xform_clip)
|
||||
dh = dh.clip(min_=dh.min(), max_=self.bbox_xform_clip)
|
||||
|
||||
pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
|
||||
pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
|
||||
pred_w = dw.exp() * widths[:, None]
|
||||
pred_h = dh.exp() * heights[:, None]
|
||||
x = pred_ctr_x - 0.5 * pred_w
|
||||
y = pred_ctr_y - 0.5 * pred_h
|
||||
w = pred_ctr_x + 0.5 * pred_w - 1
|
||||
h = pred_ctr_y + 0.5 * pred_h - 1
|
||||
pred_boxes = Tensor.stack(x, y, w, h).permute(1,2,0).reshape(rel_codes.shape[0], rel_codes.shape[1])
|
||||
return pred_boxes
|
||||
189
tinygrad_repo/examples/mlperf/initializers.py
Normal file
189
tinygrad_repo/examples/mlperf/initializers.py
Normal file
@@ -0,0 +1,189 @@
|
||||
import math
|
||||
from typing import Union
|
||||
|
||||
from tinygrad import Tensor, nn, dtypes
|
||||
from tinygrad.helpers import prod, argfix, Context
|
||||
from tinygrad.nn.state import get_parameters
|
||||
from extra.models.unet import UNetModel
|
||||
|
||||
# rejection sampling truncated randn
|
||||
def rand_truncn(*shape, dtype=None, truncstds=2, **kwargs) -> Tensor:
|
||||
CNT=8
|
||||
x = Tensor.randn(*(*shape, CNT), dtype=dtype, **kwargs)
|
||||
ctr = Tensor.arange(CNT).reshape((1,) * len(x.shape[:-1]) + (CNT,)).expand(x.shape)
|
||||
take = (x.abs() <= truncstds).where(ctr, CNT).min(axis=-1, keepdim=True) # set to 0 if no good samples
|
||||
return (ctr == take).where(x, 0).sum(axis=-1)
|
||||
|
||||
# https://github.com/keras-team/keras/blob/v2.15.0/keras/initializers/initializers.py#L1026-L1065
|
||||
def he_normal(*shape, a: float = 0.00, **kwargs) -> Tensor:
|
||||
std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:])) / 0.87962566103423978
|
||||
return std * rand_truncn(*shape, **kwargs)
|
||||
|
||||
# Stable Diffusion v2 training uses default torch gelu, which doesn't use tanh approximation
|
||||
def gelu_erf(x:Tensor) -> Tensor:
|
||||
return 0.5 * x * (1.0 + (x / 1.4142135623730951).erf())
|
||||
|
||||
class Conv2dHeNormal(nn.Conv2d):
|
||||
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
|
||||
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
|
||||
self.in_channels, self.out_channels = in_channels, out_channels # for testing
|
||||
self.weight = he_normal(out_channels, in_channels//groups, *self.kernel_size, a=0.0, dtype=dtypes.float32)
|
||||
if bias: self.bias = self.bias.cast(dtypes.float32)
|
||||
def __call__(self, x: Tensor):
|
||||
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
|
||||
padding=self.padding, stride=self.stride, dilation=self.dilation, groups=self.groups)
|
||||
|
||||
class Linear(nn.Linear):
|
||||
def __init__(self, in_features, out_features, bias=True):
|
||||
super().__init__(in_features, out_features, bias=bias)
|
||||
self.weight = Tensor.normal((out_features, in_features), mean=0.0, std=0.01, dtype=dtypes.float32)
|
||||
if bias: self.bias = Tensor.zeros(out_features, dtype=dtypes.float32)
|
||||
def __call__(self, x:Tensor):
|
||||
return x.linear(self.weight.cast(dtypes.default_float).transpose(), self.bias.cast(dtypes.default_float) if self.bias is not None else None)
|
||||
|
||||
class LinearBert(nn.Linear):
|
||||
def __init__(self, in_features, out_features, bias=True, std=0.02):
|
||||
self.weight = std * rand_truncn(out_features, in_features, dtype=dtypes.float32)
|
||||
self.bias = Tensor.zeros(out_features, dtype=dtypes.float32) if bias else None
|
||||
|
||||
def __call__(self, x:Tensor):
|
||||
return x.cast(dtypes.default_float).linear(self.weight.cast(dtypes.default_float).transpose(), self.bias.cast(dtypes.default_float) if self.bias is not None else None)
|
||||
|
||||
class EmbeddingBert(nn.Embedding):
|
||||
def __init__(self, vocab_size:int, embed_size:int, std=0.02):
|
||||
self.vocab_sz, self.embed_sz = vocab_size, embed_size
|
||||
self.weight = std * rand_truncn(vocab_size, embed_size, dtype=dtypes.float32)
|
||||
|
||||
def __call__(self, idx:Tensor) -> Tensor:
|
||||
if idx.numel() == 0: return Tensor.empty(idx.shape+(self.embed_sz,), dtype=self.weight.dtype, device=self.weight.device)
|
||||
arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
|
||||
if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, device=self.weight.device).reshape(arange_shp)
|
||||
arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
|
||||
return (arange == idx).where(vals, 0).sum(2, dtype=vals.dtype)
|
||||
|
||||
class LayerNormBert:
|
||||
def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
|
||||
self.normalized_shape = (normalized_shape,) if isinstance(normalized_shape, int) else tuple(normalized_shape)
|
||||
self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
|
||||
self.weight, self.bias = (Tensor.ones(*self.normalized_shape, dtype=dtypes.float32), Tensor.zeros(*self.normalized_shape, dtype=dtypes.float32)) if elementwise_affine else (None, None)
|
||||
|
||||
def __call__(self, x:Tensor):
|
||||
assert self.normalized_shape == x.shape[-len(self.normalized_shape):], f"last dimensions of {x.shape} must match {self.normalized_shape}"
|
||||
xn = x.cast(dtypes.float32).layernorm(eps=self.eps, axis=self.axis).cast(x.dtype)
|
||||
if not self.elementwise_affine: return xn
|
||||
return (xn * self.weight.cast(dtypes.default_float) + self.bias.cast(dtypes.default_float))
|
||||
|
||||
class FrozenBatchNorm2dRetinaNet(nn.BatchNorm2d):
|
||||
def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1):
|
||||
self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
|
||||
|
||||
self.weight = Tensor.ones(sz, dtype=dtypes.float32).is_param_(False) if affine else None
|
||||
self.bias = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False) if affine else None
|
||||
|
||||
if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, dtype=dtypes.float32).is_param_(False), Tensor.ones(sz, dtype=dtypes.float32).is_param_(False)
|
||||
self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.long).is_param_(False)
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
batch_mean, batch_var = super().calc_stats(x.cast(dtypes.float32))
|
||||
if self.track_running_stats and Tensor.training:
|
||||
self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
|
||||
self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * x.numel()/(x.numel()-x.shape[1]) * batch_var.detach().cast(self.running_var.dtype))
|
||||
self.num_batches_tracked += 1
|
||||
return x.cast(dtypes.float32).batchnorm(self.weight, self.bias, batch_mean, batch_var.add(self.eps).rsqrt()).cast(x.dtype)
|
||||
|
||||
class Conv2dNormalRetinaNet(nn.Conv2d):
|
||||
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
|
||||
stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
|
||||
bias:bool=True, prior_prob:float|None=None):
|
||||
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
|
||||
self.weight = Tensor.normal(*self.weight.shape, std=0.01, dtype=dtypes.float32)
|
||||
if bias:
|
||||
if prior_prob:
|
||||
prior_prob = Tensor(prior_prob, device=self.bias.device, dtype=dtypes.float32).expand(*self.bias.shape)
|
||||
self.bias = -(((1 - prior_prob) / prior_prob).log())
|
||||
else: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
|
||||
groups=self.groups, stride=self.stride, padding=self.padding)
|
||||
|
||||
class Conv2dKaimingUniformRetinaNet(nn.Conv2d):
|
||||
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
|
||||
stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
|
||||
bias:bool=True):
|
||||
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
|
||||
self.weight = Tensor.kaiming_uniform(*self.weight.shape, a=1, dtype=dtypes.float32)
|
||||
if bias: self.bias = Tensor.zeros_like(self.bias, dtype=dtypes.float32)
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
|
||||
groups=self.groups, stride=self.stride, padding=self.padding)
|
||||
|
||||
class Conv2dRetinaNet(nn.Conv2d):
|
||||
def __init__(self, in_channels:int, out_channels:int, kernel_size:int|tuple[int, ...],
|
||||
stride:int=1, padding:int|tuple[int, ...]|str=0, dilation:int=1, groups:int=1,
|
||||
bias:bool=True):
|
||||
super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
|
||||
scale = 1 / math.sqrt(in_channels * prod(self.kernel_size))
|
||||
self.weight = Tensor.uniform(out_channels, in_channels//groups, *self.kernel_size, low=-scale, high=scale, dtype=dtypes.float32)
|
||||
self.bias: Tensor|None = Tensor.uniform(out_channels, low=-scale, high=scale, dtype=dtypes.float32) if bias else None
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
|
||||
groups=self.groups, stride=self.stride, dilation=self.dilation, padding=self.padding)
|
||||
|
||||
# copy torch AMP: isolate mixed precision to just the below autocast ops, instead of using dtypes.default_float which affects all new Tensors
|
||||
class AutocastLinear(nn.Linear):
|
||||
cast_dtype=dtypes.bfloat16 # enable monkeypatching of the mixed precision dtype
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
dtype = type(self).cast_dtype
|
||||
return x.cast(dtype).linear(self.weight.cast(dtype).transpose(), self.bias.cast(dtype) if self.bias is not None else None)
|
||||
|
||||
class AutocastConv2d(nn.Conv2d):
|
||||
cast_dtype=dtypes.bfloat16
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
dtype = type(self).cast_dtype
|
||||
return x.cast(dtype).conv2d(self.weight.cast(dtype), self.bias.cast(dtype), self.groups, self.stride, self.dilation, self.padding)
|
||||
|
||||
# copy torch AMP: upcast to float32 before GroupNorm and LayerNorm
|
||||
class AutocastGroupNorm(nn.GroupNorm):
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return super().__call__(x.cast(dtypes.float32))
|
||||
|
||||
class AutocastLayerNorm(nn.LayerNorm):
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return super().__call__(x.cast(dtypes.float32))
|
||||
|
||||
def zero_module(module):
|
||||
for p in get_parameters(module): p.assign(Tensor.zeros_like(p).contiguous())
|
||||
|
||||
# Stable Diffusion mlperf reference doesn't call scaled_dot_product_attention
|
||||
# copy torch AMP: upcast to float32 before softmax on CUDA
|
||||
def attn_f32_softmax(q:Tensor, k:Tensor, v:Tensor) -> Tensor:
|
||||
return (q.matmul(k.transpose(-2,-1), dtype=dtypes.float32) / math.sqrt(q.shape[-1])).softmax(-1).cast(q.dtype) @ v
|
||||
|
||||
def init_stable_diffusion(version:str, pretrained:str, devices:list[str]):
|
||||
from examples.stable_diffusion import StableDiffusion
|
||||
from tinygrad.nn.state import safe_load, safe_save, load_state_dict, get_state_dict
|
||||
from tempfile import TemporaryDirectory
|
||||
model = StableDiffusion(version=version, pretrained=pretrained)
|
||||
unet:UNetModel = model.model.diffusion_model
|
||||
|
||||
# this prevents extra consumption of memory, enabling much larger BS
|
||||
Tensor.realize(*get_parameters(unet))
|
||||
with TemporaryDirectory(prefix="unet_init") as tmp:
|
||||
safe_save(get_state_dict(unet), init_fn:=f"{tmp}/init_model.safetensors")
|
||||
load_state_dict(unet, safe_load(init_fn))
|
||||
|
||||
sqrt_alphas_cumprod = model.alphas_cumprod.sqrt().realize()
|
||||
sqrt_one_minus_alphas_cumprod = (1 - model.alphas_cumprod).sqrt().realize()
|
||||
|
||||
if len(devices) > 1:
|
||||
to_move = [sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod]
|
||||
if version == "v2-mlperf-train": to_move += get_parameters(unet) + get_parameters(model.cond_stage_model)
|
||||
for p in to_move:
|
||||
p.to_(devices)
|
||||
with Context(BEAM=0):
|
||||
Tensor.realize(*to_move)
|
||||
|
||||
return model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod
|
||||
29
tinygrad_repo/examples/mlperf/losses.py
Normal file
29
tinygrad_repo/examples/mlperf/losses.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from examples.mlperf.metrics import dice_score
|
||||
from tinygrad import Tensor
|
||||
|
||||
def dice_ce_loss(pred, tgt):
|
||||
ce = pred.permute(0, 2, 3, 4, 1).sparse_categorical_crossentropy(tgt.squeeze(1))
|
||||
dice = (1.0 - dice_score(pred, tgt, argmax=False, to_one_hot_x=False)).mean()
|
||||
return (dice + ce) / 2
|
||||
|
||||
def sigmoid_focal_loss(pred:Tensor, tgt:Tensor, alpha:float=0.25, gamma:float=2.0, reduction:str="none") -> Tensor:
|
||||
assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
|
||||
p, ce_loss = pred.sigmoid(), pred.binary_crossentropy_logits(tgt, reduction="none")
|
||||
p_t = p * tgt + (1 - p) * (1 - tgt)
|
||||
loss = ce_loss * ((1 - p_t) ** gamma)
|
||||
|
||||
if alpha >= 0:
|
||||
alpha_t = alpha * tgt + (1 - alpha) * (1 - tgt)
|
||||
loss = loss * alpha_t
|
||||
|
||||
if reduction == "mean": loss = loss.mean()
|
||||
elif reduction == "sum": loss = loss.sum()
|
||||
return loss
|
||||
|
||||
def l1_loss(pred:Tensor, tgt:Tensor, reduction:str="none") -> Tensor:
|
||||
assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
|
||||
loss = (pred - tgt).abs()
|
||||
|
||||
if reduction == "mean": loss = loss.mean()
|
||||
elif reduction == "sum": loss = loss.sum()
|
||||
return loss
|
||||
60
tinygrad_repo/examples/mlperf/lr_schedulers.py
Normal file
60
tinygrad_repo/examples/mlperf/lr_schedulers.py
Normal file
@@ -0,0 +1,60 @@
|
||||
import math
|
||||
from tinygrad import dtypes, Tensor
|
||||
from tinygrad.nn.optim import Optimizer
|
||||
|
||||
from extra.lr_scheduler import LR_Scheduler
|
||||
from typing import Callable
|
||||
|
||||
# https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py
|
||||
class PolynomialDecayWithWarmup(LR_Scheduler):
|
||||
def __init__(self, optimizer: Optimizer, initial_lr, end_lr, train_steps, warmup, power=2):
|
||||
super().__init__(optimizer)
|
||||
self.epoch_counter = self.epoch_counter.cast(dtypes.float32)
|
||||
assert train_steps > 0 and warmup > 0
|
||||
self.warmup = min(warmup, train_steps)
|
||||
self.initial_lr, self.end_lr, self.epochs, self.power = initial_lr, end_lr, train_steps, power
|
||||
|
||||
# set lr for first warmup step
|
||||
self.optimizer.lr.assign(self.get_lr()).realize()
|
||||
|
||||
def get_lr(self):
|
||||
# LR is 0 on the first step, matching the reference.
|
||||
warmup_lr = (self.epoch_counter * (1.0 / self.warmup)) * self.initial_lr
|
||||
x = (1 - (self.epoch_counter - self.warmup) / (self.epochs - self.warmup + 1))
|
||||
return (self.epoch_counter <= self.warmup).where(warmup_lr, (self.initial_lr - self.end_lr) * x ** self.power + self.end_lr).cast(self.optimizer.lr.dtype)
|
||||
|
||||
class CosineAnnealingLRWithWarmup(LR_Scheduler):
|
||||
def __init__(self, optimizer:Optimizer, base_lr, end_lr, warmup_steps:int, decay_steps:int):
|
||||
assert warmup_steps > 0 and decay_steps > 0
|
||||
super().__init__(optimizer)
|
||||
self.base_lr = base_lr
|
||||
self.end_lr = end_lr
|
||||
self.warmup_steps = warmup_steps
|
||||
self.decay_steps = decay_steps
|
||||
# set lr for first warmup step
|
||||
self.optimizer.lr.assign(self.get_lr()).realize()
|
||||
|
||||
def get_lr(self):
|
||||
warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr
|
||||
decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos())
|
||||
return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
|
||||
|
||||
# Reference: https://github.com/mlcommons/training/blob/64b14a9abc74e08779a175abca7d291f8c957632/stable_diffusion/ldm/lr_scheduler.py, Lines 36-97
|
||||
class LambdaLinearScheduler:
|
||||
def __init__(self, warm_up_steps:int, f_min:float, f_max:float, f_start:float, cycle_lengths:int):
|
||||
self.lr_warm_up_steps, self.f_min, self.f_max, self.f_start, self.cycle_lengths = warm_up_steps, f_min, f_max, f_start, cycle_lengths
|
||||
|
||||
def schedule(self, n:Tensor) -> Tensor:
|
||||
warm_up = (n < self.lr_warm_up_steps)
|
||||
f_warm_up = (self.f_max - self.f_start) / self.lr_warm_up_steps * n + self.f_start
|
||||
return warm_up.where(f_warm_up, self.f_min + (self.f_max - self.f_min) * (self.cycle_lengths - n) / (self.cycle_lengths))
|
||||
|
||||
# based on torch.optim.lr_scheduler.LambdaLR
|
||||
class LambdaLR(LR_Scheduler):
|
||||
def __init__(self, optimizer:Optimizer, base_lr:Tensor, lr_lambda:Callable):
|
||||
super().__init__(optimizer)
|
||||
self.base_lr, self.lr_lambda = base_lr, lr_lambda
|
||||
self.step()
|
||||
|
||||
def get_lr(self):
|
||||
return self.base_lr * self.lr_lambda(self.epoch_counter - 1)
|
||||
69
tinygrad_repo/examples/mlperf/metrics.py
Normal file
69
tinygrad_repo/examples/mlperf/metrics.py
Normal file
@@ -0,0 +1,69 @@
|
||||
import re, string
|
||||
from collections import Counter
|
||||
from tinygrad import Tensor
|
||||
|
||||
def levenshtein(a, b):
|
||||
n, m = len(a), len(b)
|
||||
if n > m:
|
||||
a, b, n, m = b, a, m, n
|
||||
|
||||
current = list(range(n + 1))
|
||||
for i in range(1, m + 1):
|
||||
previous, current = current, [i] + [0] * n
|
||||
for j in range(1, n + 1):
|
||||
add, delete = previous[j] + 1, current[j - 1] + 1
|
||||
change = previous[j - 1]
|
||||
if a[j - 1] != b[i - 1]:
|
||||
change = change + 1
|
||||
current[j] = min(add, delete, change)
|
||||
|
||||
return current[n]
|
||||
|
||||
def word_error_rate(x, y):
|
||||
scores = words = 0
|
||||
for h, r in zip(x, y):
|
||||
h_list = h.split()
|
||||
r_list = r.split()
|
||||
words += len(r_list)
|
||||
scores += levenshtein(h_list, r_list)
|
||||
return float(scores) / words, float(scores), words
|
||||
|
||||
def one_hot(x):
|
||||
return x.one_hot(3).squeeze(1).permute(0, 4, 1, 2, 3)
|
||||
|
||||
def dice_score(prediction, target, channel_axis=1, smooth_nr=1e-6, smooth_dr=1e-6, argmax=True, to_one_hot_x=True):
|
||||
channel_axis, reduce_axis = 1, tuple(range(2, len(prediction.shape)))
|
||||
if argmax: prediction = prediction.argmax(axis=channel_axis)
|
||||
else: prediction = prediction.softmax(axis=channel_axis)
|
||||
if to_one_hot_x: prediction = one_hot(prediction)
|
||||
target = one_hot(target)
|
||||
prediction, target = prediction[:, 1:], target[:, 1:]
|
||||
assert prediction.shape == target.shape, f"prediction ({prediction.shape}) and target ({target.shape}) shapes do not match"
|
||||
intersection = (prediction * target).sum(axis=reduce_axis)
|
||||
target_sum = target.sum(axis=reduce_axis)
|
||||
prediction_sum = prediction.sum(axis=reduce_axis)
|
||||
result = (2.0 * intersection + smooth_nr) / (target_sum + prediction_sum + smooth_dr)
|
||||
return result
|
||||
|
||||
def normalize_string(s):
|
||||
s = "".join(c for c in s.lower() if c not in string.punctuation)
|
||||
s = re.sub(r'\b(a|an|the)\b', ' ', s)
|
||||
return " ".join(s.split())
|
||||
|
||||
def f1_score(x, y):
|
||||
xt = normalize_string(x).split()
|
||||
yt = normalize_string(y).split()
|
||||
ct = Counter(xt) & Counter(yt)
|
||||
if (ns := sum(ct.values())) == 0:
|
||||
return 0.0
|
||||
p = ns / len(xt)
|
||||
r = ns / len(yt)
|
||||
return 2 * p * r / (p + r)
|
||||
|
||||
def log_perplexity(logit:Tensor, target:Tensor, ignore_index:int|None=None):
|
||||
# logit has shape (n_samples, seq_len, vocab_size), target has shape (n_samples, seq_len)
|
||||
assert logit.ndim == 3, logit.ndim
|
||||
assert target.ndim == 2, target.ndim
|
||||
assert logit.shape[:2] == target.shape, f"{logit.shape[:2]=}, {target.shape=}"
|
||||
log_prob = logit.log_softmax(axis=-1)
|
||||
return log_prob.transpose(1, 2).nll_loss(target, ignore_index=ignore_index)
|
||||
508
tinygrad_repo/examples/mlperf/model_eval.py
Normal file
508
tinygrad_repo/examples/mlperf/model_eval.py
Normal file
@@ -0,0 +1,508 @@
|
||||
import time, math, os
|
||||
start = time.perf_counter()
|
||||
from pathlib import Path
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, Device, dtypes, GlobalCounters, TinyJit
|
||||
from tinygrad.nn.state import get_parameters, load_state_dict, safe_load
|
||||
from tinygrad.helpers import getenv, Context, prod
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
def tlog(x): print(f"{x:25s} @ {time.perf_counter()-start:5.2f}s")
|
||||
|
||||
def eval_resnet():
|
||||
with WallTimeEvent(BenchEvent.FULL):
|
||||
# Resnet50-v1.5
|
||||
from extra.models.resnet import ResNet50
|
||||
tlog("imports")
|
||||
GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 6))]
|
||||
for x in GPUS: Device[x]
|
||||
tlog("got devices") # NOTE: this is faster with rocm-smi running
|
||||
|
||||
class ResnetRunner:
|
||||
def __init__(self, device=None):
|
||||
self.mdl = ResNet50()
|
||||
for x in get_parameters(self.mdl) if device else []: x.to_(device)
|
||||
if (fn:=getenv("RESNET_MODEL", "")): load_state_dict(self.mdl, safe_load(fn))
|
||||
else: self.mdl.load_from_pretrained()
|
||||
self.input_mean = Tensor([0.485, 0.456, 0.406], device=device).reshape(1, -1, 1, 1)
|
||||
self.input_std = Tensor([0.229, 0.224, 0.225], device=device).reshape(1, -1, 1, 1)
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
x = x.permute([0,3,1,2]).cast(dtypes.float32) / 255.0
|
||||
x -= self.input_mean
|
||||
x /= self.input_std
|
||||
return self.mdl(x).log_softmax().argmax(axis=1).realize()
|
||||
|
||||
mdl = TinyJit(ResnetRunner(GPUS))
|
||||
tlog("loaded models")
|
||||
|
||||
# evaluation on the mlperf classes of the validation set from imagenet
|
||||
from examples.mlperf.dataloader import batch_load_resnet
|
||||
iterator = batch_load_resnet(getenv("BS", 128*6), val=getenv("VAL", 1), shuffle=False, pad_first_batch=True)
|
||||
def data_get():
|
||||
x,y,cookie = next(iterator)
|
||||
return x.shard(GPUS, axis=0).realize(), y, cookie
|
||||
n,d = 0,0
|
||||
proc = data_get()
|
||||
tlog("loaded initial data")
|
||||
st = time.perf_counter()
|
||||
while proc is not None:
|
||||
GlobalCounters.reset()
|
||||
proc = (mdl(proc[0]), proc[1], proc[2]) # this frees the images
|
||||
run = time.perf_counter()
|
||||
# load the next data here
|
||||
try: next_proc = data_get()
|
||||
except StopIteration: next_proc = None
|
||||
nd = time.perf_counter()
|
||||
y = np.array(proc[1])
|
||||
proc = (proc[0].numpy() == y) & (y != -1) # this realizes the models and frees the cookies
|
||||
n += proc.sum()
|
||||
d += (y != -1).sum()
|
||||
et = time.perf_counter()
|
||||
tlog(f"****** {n:5d}/{d:5d} {n*100.0/d:.2f}% -- {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
|
||||
st = et
|
||||
proc, next_proc = next_proc, None
|
||||
tlog("done")
|
||||
|
||||
def eval_unet3d():
|
||||
# UNet3D
|
||||
from extra.models.unet3d import UNet3D
|
||||
from extra.datasets.kits19 import iterate, sliding_window_inference, get_val_files
|
||||
from examples.mlperf.metrics import dice_score
|
||||
mdl = UNet3D()
|
||||
mdl.load_from_pretrained()
|
||||
s = 0
|
||||
st = time.perf_counter()
|
||||
for i, (image, label) in enumerate(iterate(get_val_files()), start=1):
|
||||
mt = time.perf_counter()
|
||||
pred, label = sliding_window_inference(mdl, image, label)
|
||||
et = time.perf_counter()
|
||||
print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model")
|
||||
s += dice_score(Tensor(pred), Tensor(label)).mean().item()
|
||||
print(f"****** {s:.2f}/{i} {s/i:.5f} Mean DICE score")
|
||||
st = time.perf_counter()
|
||||
|
||||
def eval_retinanet():
|
||||
# RetinaNet with ResNeXt50_32X4D
|
||||
from examples.mlperf.dataloader import batch_load_retinanet
|
||||
from extra.datasets.openimages import normalize, download_dataset, BASEDIR
|
||||
from extra.models.resnet import ResNeXt50_32X4D
|
||||
from extra.models.retinanet import RetinaNet
|
||||
from pycocotools.coco import COCO
|
||||
from pycocotools.cocoeval import COCOeval
|
||||
from contextlib import redirect_stdout
|
||||
tlog("imports")
|
||||
|
||||
mdl = RetinaNet(ResNeXt50_32X4D())
|
||||
mdl.load_from_pretrained()
|
||||
tlog("loaded models")
|
||||
|
||||
coco = COCO(download_dataset(base_dir:=getenv("BASEDIR", BASEDIR), 'validation'))
|
||||
coco_eval = COCOeval(coco, iouType="bbox")
|
||||
coco_evalimgs, evaluated_imgs, ncats, narea = [], [], len(coco_eval.params.catIds), len(coco_eval.params.areaRng)
|
||||
tlog("loaded dataset")
|
||||
|
||||
iterator = batch_load_retinanet(coco, True, Path(base_dir), getenv("BS", 8), shuffle=False)
|
||||
def data_get():
|
||||
x, img_ids, img_sizes, cookie = next(iterator)
|
||||
return x.to(Device.DEFAULT).realize(), img_ids, img_sizes, cookie
|
||||
n = 0
|
||||
proc = data_get()
|
||||
tlog("loaded initial data")
|
||||
st = time.perf_counter()
|
||||
while proc is not None:
|
||||
GlobalCounters.reset()
|
||||
proc = (mdl(normalize(proc[0])), proc[1], proc[2], proc[3])
|
||||
run = time.perf_counter()
|
||||
# load the next data here
|
||||
try: next_proc = data_get()
|
||||
except StopIteration: next_proc = None
|
||||
nd = time.perf_counter()
|
||||
predictions, img_ids = mdl.postprocess_detections(proc[0].numpy(), orig_image_sizes=proc[2]), proc[1]
|
||||
pd = time.perf_counter()
|
||||
coco_results = [{"image_id": img_ids[i], "category_id": label, "bbox": box.tolist(), "score": score}
|
||||
for i, prediction in enumerate(predictions) for box, score, label in zip(*prediction.values())]
|
||||
with redirect_stdout(None):
|
||||
coco_eval.cocoDt = coco.loadRes(coco_results)
|
||||
coco_eval.params.imgIds = img_ids
|
||||
coco_eval.evaluate()
|
||||
evaluated_imgs.extend(img_ids)
|
||||
coco_evalimgs.append(np.array(coco_eval.evalImgs).reshape(ncats, narea, len(img_ids)))
|
||||
n += len(proc[0])
|
||||
et = time.perf_counter()
|
||||
tlog(f"****** {(run-st)*1000:7.2f} ms to enqueue, {(et-run)*1000:7.2f} ms to realize ({(nd-run)*1000:7.2f} ms fetching, {(pd-run)*1000:4.2f} ms postprocess_detections). {(len(proc))/(et-st):8.2f} examples/sec. {GlobalCounters.global_ops*1e-12/(et-st):5.2f} TFLOPS")
|
||||
st = et
|
||||
proc, next_proc = next_proc, None
|
||||
|
||||
coco_eval.params.imgIds = evaluated_imgs
|
||||
coco_eval._paramsEval.imgIds = evaluated_imgs
|
||||
coco_eval.evalImgs = list(np.concatenate(coco_evalimgs, -1).flatten())
|
||||
coco_eval.accumulate()
|
||||
coco_eval.summarize()
|
||||
tlog("done")
|
||||
|
||||
def eval_rnnt():
|
||||
# RNN-T
|
||||
from extra.models.rnnt import RNNT
|
||||
mdl = RNNT()
|
||||
mdl.load_from_pretrained()
|
||||
|
||||
from extra.datasets.librispeech import iterate
|
||||
from examples.mlperf.metrics import word_error_rate
|
||||
|
||||
LABELS = [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
|
||||
|
||||
c = 0
|
||||
scores = 0
|
||||
words = 0
|
||||
st = time.perf_counter()
|
||||
for X, Y in iterate():
|
||||
mt = time.perf_counter()
|
||||
tt = mdl.decode(Tensor(X[0]), Tensor([X[1]]))
|
||||
et = time.perf_counter()
|
||||
print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model")
|
||||
for n, t in enumerate(tt):
|
||||
tnp = np.array(t)
|
||||
_, scores_, words_ = word_error_rate(["".join([LABELS[int(tnp[i])] for i in range(tnp.shape[0])])], [Y[n]])
|
||||
scores += scores_
|
||||
words += words_
|
||||
c += len(tt)
|
||||
print(f"WER: {scores/words}, {words} words, raw scores: {scores}, c: {c}")
|
||||
st = time.perf_counter()
|
||||
|
||||
def eval_bert():
|
||||
# Bert-QA
|
||||
from extra.models.bert import BertForQuestionAnswering
|
||||
mdl = BertForQuestionAnswering()
|
||||
mdl.load_from_pretrained()
|
||||
|
||||
@TinyJit
|
||||
def run(input_ids, input_mask, segment_ids):
|
||||
return mdl(input_ids, input_mask, segment_ids).realize()
|
||||
|
||||
from extra.datasets.squad import iterate
|
||||
from examples.mlperf.helpers import get_bert_qa_prediction
|
||||
from examples.mlperf.metrics import f1_score
|
||||
from transformers import BertTokenizer
|
||||
|
||||
tokenizer = BertTokenizer(str(Path(__file__).parents[2] / "extra/weights/bert_vocab.txt"))
|
||||
|
||||
c = 0
|
||||
f1 = 0.0
|
||||
st = time.perf_counter()
|
||||
for X, Y in iterate(tokenizer):
|
||||
mt = time.perf_counter()
|
||||
outs = []
|
||||
for x in X:
|
||||
outs.append(run(Tensor(x["input_ids"]), Tensor(x["input_mask"]), Tensor(x["segment_ids"])).numpy())
|
||||
et = time.perf_counter()
|
||||
print(f"{(mt-st)*1000:.2f} ms loading data, {(et-mt)*1000:.2f} ms to run model over {len(X)} features")
|
||||
|
||||
pred = get_bert_qa_prediction(X, Y, outs)
|
||||
print(f"pred: {pred}\nans: {Y['answers']}")
|
||||
f1 += max([f1_score(pred, ans) for ans in Y["answers"]])
|
||||
c += 1
|
||||
print(f"f1: {f1/c}, raw: {f1}, c: {c}\n")
|
||||
|
||||
st = time.perf_counter()
|
||||
|
||||
def eval_llama3():
|
||||
from extra.models.llama import Transformer
|
||||
from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
|
||||
from tinygrad.helpers import tqdm
|
||||
|
||||
BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
|
||||
BS = getenv("BS", 4)
|
||||
SMALL = getenv("SMALL", 0)
|
||||
SEQLEN = getenv("SEQLEN", 8192)
|
||||
MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
|
||||
|
||||
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
|
||||
params = params | {"vocab_size": 32000} if not SMALL else params
|
||||
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
|
||||
model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
|
||||
|
||||
# load weights
|
||||
weights = load(str(MODEL_PATH / "model.safetensors.index.json"))
|
||||
if "model.embed_tokens.weight" in weights:
|
||||
print("converting from huggingface format")
|
||||
weights = convert_from_huggingface(weights, params["n_layers"], params["n_heads"], params["n_kv_heads"])
|
||||
|
||||
load_state_dict(model, weights, strict=False, consume=True)
|
||||
|
||||
@TinyJit
|
||||
def eval_step(model, tokens):
|
||||
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
|
||||
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
return loss.flatten().float()
|
||||
|
||||
from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
|
||||
eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
|
||||
iter = iterate_llama3_dataset(eval_dataset, BS)
|
||||
|
||||
losses = []
|
||||
for tokens in tqdm(iter, total=5760//BS):
|
||||
GlobalCounters.reset()
|
||||
losses += eval_step(model, tokens).tolist()
|
||||
tqdm.write(f"loss: {np.mean(losses)}")
|
||||
|
||||
log_perplexity = np.mean(losses)
|
||||
print(f"Log Perplexity: {log_perplexity}")
|
||||
|
||||
# NOTE: BEAM hangs on 8xmi300x with DECODE_BS=384 in final realize below; function is declared here for external testing
|
||||
@TinyJit
|
||||
def vae_decode(x:Tensor, vae, disable_beam=False) -> Tensor:
|
||||
from examples.stable_diffusion import AutoencoderKL
|
||||
assert isinstance(vae, AutoencoderKL)
|
||||
x = vae.post_quant_conv(1./0.18215 * x)
|
||||
|
||||
x = vae.decoder.conv_in(x)
|
||||
x = vae.decoder.mid(x)
|
||||
for i, l in enumerate(vae.decoder.up[::-1]):
|
||||
print("decode", x.shape)
|
||||
for b in l['block']: x = b(x)
|
||||
if 'upsample' in l:
|
||||
bs,c,py,px = x.shape
|
||||
x = x.reshape(bs, c, py, 1, px, 1).expand(bs, c, py, 2, px, 2).reshape(bs, c, py*2, px*2)
|
||||
x = l['upsample']['conv'](x)
|
||||
if i == len(vae.decoder.up) - 1 and disable_beam:
|
||||
with Context(BEAM=0): x.realize()
|
||||
else: x.realize()
|
||||
x = vae.decoder.conv_out(vae.decoder.norm_out(x).swish())
|
||||
|
||||
x = ((x + 1.0) / 2.0).clip(0.0, 1.0)
|
||||
return x
|
||||
|
||||
def eval_stable_diffusion():
|
||||
import csv, PIL, sys
|
||||
from tqdm import tqdm
|
||||
from examples.mlperf.initializers import init_stable_diffusion, gelu_erf
|
||||
from examples.stable_diffusion import AutoencoderKL
|
||||
from extra.models.unet import UNetModel
|
||||
from tinygrad.nn.state import load_state_dict, torch_load
|
||||
from tinygrad.helpers import BEAM
|
||||
from extra.models import clip
|
||||
from extra.models.clip import FrozenOpenClipEmbedder
|
||||
from extra.models.clip import OpenClipEncoder
|
||||
from extra.models.inception import FidInceptionV3
|
||||
|
||||
config = {}
|
||||
GPUS = config["GPUS"] = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
|
||||
for x in GPUS: Device[x]
|
||||
print(f"running eval on {GPUS}")
|
||||
seed = config["seed"] = getenv("SEED", 12345)
|
||||
CKPTDIR = config["CKPTDIR"] = Path(getenv("CKPTDIR", "./checkpoints"))
|
||||
DATADIR = config["DATADIR"] = Path(getenv("DATADIR", "./datasets"))
|
||||
CONTEXT_BS = config["CONTEXT_BS"] = getenv("CONTEXT_BS", 1 * len(GPUS))
|
||||
DENOISE_BS = config["DENOISE_BS"] = getenv("DENOISE_BS", 1 * len(GPUS))
|
||||
DECODE_BS = config["DECODE_BS"] = getenv("DECODE_BS", 1 * len(GPUS))
|
||||
INCEPTION_BS = config["INCEPTION_BS"] = getenv("INCEPTION_BS", 1 * len(GPUS))
|
||||
CLIP_BS = config["CLIP_BS"] = getenv("CLIP_BS", 1 * len(GPUS))
|
||||
EVAL_CKPT_DIR = config["EVAL_CKPT_DIR"] = getenv("EVAL_CKPT_DIR", "")
|
||||
STOP_IF_CONVERGED = config["STOP_IF_CONVERGED"] = getenv("STOP_IF_CONVERGED", 0)
|
||||
|
||||
if (WANDB := getenv("WANDB", "")):
|
||||
import wandb
|
||||
wandb.init(config=config, project="MLPerf-Stable-Diffusion")
|
||||
|
||||
assert EVAL_CKPT_DIR != "", "provide a directory with checkpoints to be evaluated"
|
||||
print(f"running eval on checkpoints in {EVAL_CKPT_DIR}\nSEED={seed}")
|
||||
eval_queue:list[tuple[int, Path]] = []
|
||||
for p in Path(EVAL_CKPT_DIR).iterdir():
|
||||
if p.name.endswith(".safetensors"):
|
||||
ckpt_iteration = p.name.split(".safetensors")[0]
|
||||
assert ckpt_iteration.isdigit(), f"invalid checkpoint name: {p.name}, expected <digits>.safetensors"
|
||||
eval_queue.append((int(ckpt_iteration), p))
|
||||
assert len(eval_queue), f'no files ending with ".safetensors" were found in {EVAL_CKPT_DIR}'
|
||||
print(sorted(eval_queue, reverse=True))
|
||||
|
||||
Tensor.manual_seed(seed) # seed for weight initialization
|
||||
model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-eval", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
|
||||
|
||||
# load prompts for generating images for validation; 2 MB of data total
|
||||
with open(DATADIR / "coco2014" / "val2014_30k.tsv") as f:
|
||||
reader = csv.DictReader(f, delimiter="\t")
|
||||
eval_inputs:list[dict] = [{"image_id": int(row["image_id"]), "id": int(row["id"]), "caption": row["caption"]} for row in reader]
|
||||
assert len(eval_inputs) == 30_000
|
||||
# NOTE: the clip weights are the same between model.cond_stage_model and clip_encoder
|
||||
eval_timesteps = list(reversed(range(1, 1000, 20)))
|
||||
|
||||
with Context(DEV="CPU"):
|
||||
# The choice of alphas_prev[0] = alphas_cumprod[0] seems arbitrary, but it's how the mlperf ref does it:
|
||||
# alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
|
||||
eval_alphas_prev = model.alphas_cumprod[0:1].cat(model.alphas_cumprod[list(range(1, 1000, 20))[:-1]]).to(GPUS).realize()
|
||||
inception = FidInceptionV3().load_from_pretrained(CKPTDIR / "inception" / "pt_inception-2015-12-05-6726825d.pth")
|
||||
vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
|
||||
text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
|
||||
clip.gelu = gelu_erf
|
||||
clip_encoder = OpenClipEncoder(1024, text_cfg, vision_cfg)
|
||||
loaded = torch_load(CKPTDIR / "clip" / "open_clip_pytorch_model.bin")
|
||||
loaded.update({"attn_mask": clip_encoder.attn_mask, "mean": clip_encoder.mean, "std": clip_encoder.std})
|
||||
load_state_dict(clip_encoder, loaded)
|
||||
|
||||
@TinyJit
|
||||
def denoise_step(x:Tensor, x_x:Tensor, t_t:Tensor, uc_c:Tensor, sqrt_alphas_cumprod_t:Tensor, sqrt_one_minus_alphas_cumprod_t:Tensor,
|
||||
alpha_prev:Tensor, unet:UNetModel, GPUS) -> Tensor:
|
||||
out_uncond, out = unet(x_x, t_t, uc_c).to("CPU").reshape(-1, 2, 4, 64, 64).chunk(2, dim=1)
|
||||
out_uncond = out_uncond.squeeze(1).shard(GPUS,axis=0)
|
||||
out = out.squeeze(1).shard(GPUS,axis=0)
|
||||
v_t = out_uncond + 8.0 * (out - out_uncond)
|
||||
e_t = sqrt_alphas_cumprod_t * v_t + sqrt_one_minus_alphas_cumprod_t * x
|
||||
pred_x0 = sqrt_alphas_cumprod_t * x - sqrt_one_minus_alphas_cumprod_t * v_t
|
||||
dir_xt = (1. - alpha_prev).sqrt() * e_t
|
||||
x_prev = alpha_prev.sqrt() * pred_x0 + dir_xt
|
||||
return x_prev.realize()
|
||||
|
||||
def shard_tensor(t:Tensor) -> Tensor: return t.shard(GPUS, axis=0) if len(GPUS) > 1 else t.to(GPUS[0])
|
||||
def get_batch(whole:Tensor, i:int, bs:int) -> tuple[Tensor, int]:
|
||||
batch = whole[i: i + bs].to("CPU")
|
||||
if (unpadded_bs:=batch.shape[0]) < bs:
|
||||
batch = batch.cat(batch[-1:].expand(bs - unpadded_bs, *batch[-1].shape))
|
||||
return batch, unpadded_bs
|
||||
|
||||
@Tensor.train(mode=False)
|
||||
def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
|
||||
inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
|
||||
# Eval is divided into 5 jits, one per model
|
||||
# It doesn't make sense to merge these jits, e.g. unet repeats 50 times in isolation; images fork to separate inception/clip
|
||||
# We're generating and scoring 30,000 images per eval, and all the data can flow through one jit at a time
|
||||
# To maximize throughput for each jit, we have only one model/jit on the GPU at a time, and pool outputs from each jit off-GPU
|
||||
for model in (unet, first_stage, inception, clip):
|
||||
Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
|
||||
|
||||
uc_written = False
|
||||
models = (cond_stage, unet, first_stage, inception, clip)
|
||||
jits = (jit_context:=TinyJit(cond_stage.embed_tokens), denoise_step, vae_decode, jit_inception:=TinyJit(inception),
|
||||
jit_clip:=TinyJit(clip.get_clip_score))
|
||||
all_bs = (CONTEXT_BS, DENOISE_BS, DECODE_BS, INCEPTION_BS, CLIP_BS)
|
||||
if (EVAL_SAMPLES:=getenv("EVAL_SAMPLES", 0)) and EVAL_SAMPLES > 0:
|
||||
eval_inputs = eval_inputs[0:EVAL_SAMPLES]
|
||||
output_shapes = [(ns:=len(eval_inputs),77), (ns,77,1024), (ns,4,64,64), (ns,3,512,512), (ns,2048), (ns,)]
|
||||
# Writing progress to disk lets us resume eval if we crash
|
||||
stages = ["tokens", "embeds", "latents", "imgs", "inception", "clip"]
|
||||
disk_tensor_names, disk_tensor_shapes = stages + ["end", "uc"], output_shapes + [(6,), (1,77,1024)]
|
||||
if not all(os.path.exists(f"{EVAL_CKPT_DIR}/{name}.bytes") for name in disk_tensor_names):
|
||||
for name, shape in zip(disk_tensor_names, disk_tensor_shapes):
|
||||
file = Path(f"{EVAL_CKPT_DIR}/{name}.bytes")
|
||||
file.unlink(missing_ok=True)
|
||||
with file.open("wb") as f: f.truncate(prod(shape) * 4)
|
||||
progress = {name: Tensor.empty(*shape, device=f"disk:{EVAL_CKPT_DIR}/{name}.bytes", dtype=dtypes.int if name in {"tokens", "end"} else dtypes.float)
|
||||
for name, shape in zip(disk_tensor_names, disk_tensor_shapes)}
|
||||
|
||||
def embed_tokens(tokens:Tensor) -> Tensor:
|
||||
nonlocal uc_written
|
||||
if not uc_written:
|
||||
with Context(BEAM=0): progress["uc"].assign(cond_stage.embed_tokens(cond_stage.tokenize("").to(GPUS)).to("CPU").realize()).realize()
|
||||
uc_written = True
|
||||
return jit_context(shard_tensor(tokens))
|
||||
|
||||
def generate_latents(embeds:Tensor) -> Tensor:
|
||||
uc_c = Tensor.stack(progress["uc"].to("CPU").expand(bs, 77, 1024), embeds, dim=1).reshape(-1, 77, 1024)
|
||||
uc_c = shard_tensor(uc_c)
|
||||
x = shard_tensor(Tensor.randn(bs,4,64,64))
|
||||
for step_idx, timestep in enumerate(tqdm(eval_timesteps)):
|
||||
reversed_idx = Tensor([50 - step_idx - 1], device=GPUS)
|
||||
alpha_prev = eval_alphas_prev[reversed_idx]
|
||||
ts = Tensor.full(bs, fill_value=timestep, dtype=dtypes.int, device="CPU")
|
||||
ts_ts = shard_tensor(ts.cat(ts))
|
||||
ts = shard_tensor(ts)
|
||||
sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
|
||||
sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
|
||||
x_x = shard_tensor(Tensor.stack(x.to("CPU"), x.to("CPU"), dim=1).reshape(-1, 4, 64, 64))
|
||||
x.assign(denoise_step(x, x_x, ts_ts, uc_c, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t, alpha_prev, unet, GPUS)).realize()
|
||||
return x
|
||||
|
||||
def decode_latents(latents:Tensor) -> Tensor: return vae_decode(shard_tensor(latents), first_stage, disable_beam=True)
|
||||
def generate_inception(imgs:Tensor) -> Tensor: return jit_inception(shard_tensor(imgs))[:,:,0,0]
|
||||
|
||||
def calc_clip_scores(batch:Tensor, batch_tokens:Tensor) -> Tensor:
|
||||
# Tensor.interpolate does not yet support bicubic, so we use PIL
|
||||
batch = (batch.to(GPUS[0]).permute(0,2,3,1) * 255).clip(0, 255).cast(dtypes.uint8).numpy()
|
||||
batch = [np.array(PIL.Image.fromarray(batch[i]).resize((224,224), PIL.Image.BICUBIC)) for i in range(bs)]
|
||||
batch = shard_tensor(Tensor(np.stack(batch, axis=0).transpose(0,3,1,2), device="CPU").realize())
|
||||
batch = batch.cast(dtypes.float) / 255
|
||||
batch = (batch - model.mean) / model.std
|
||||
batch = jit_clip(shard_tensor(batch_tokens), batch)
|
||||
return batch
|
||||
|
||||
callbacks = (embed_tokens, generate_latents, decode_latents, generate_inception, calc_clip_scores)
|
||||
|
||||
# save every forward pass output to disk; NOTE: this needs ~100 GB disk space because 30k images are large
|
||||
def stage_progress(stage_idx:int) -> int: return progress["end"].to("CPU")[stage_idx].item()
|
||||
if stage_progress(0) < len(eval_inputs):
|
||||
tokens = []
|
||||
for i in tqdm(range(0, len(eval_inputs), CONTEXT_BS)):
|
||||
subset = [cond_stage.tokenize(row["caption"], device="CPU") for row in eval_inputs[i: i+CONTEXT_BS]]
|
||||
tokens.append(Tensor.cat(*subset, dim=0).realize())
|
||||
progress["tokens"].assign(Tensor.cat(*tokens, dim=0).realize()).realize()
|
||||
progress["end"][0:1].assign(Tensor([len(eval_inputs)], dtype=dtypes.int)).realize()
|
||||
prev_stage = "tokens"
|
||||
tokens = progress["tokens"]
|
||||
|
||||
# wrapper code for every model
|
||||
for stage_idx, model, jit, bs, callback in zip(range(1,6), models, jits, all_bs, callbacks):
|
||||
stage = stages[stage_idx]
|
||||
if stage_progress(stage_idx) >= len(eval_inputs):
|
||||
prev_stage = stage
|
||||
continue # use cache
|
||||
t0 = time.perf_counter()
|
||||
print(f"starting eval with model: {model}")
|
||||
if stage_idx == 1: inputs = tokens
|
||||
elif stage_idx == 5: inputs = progress["imgs"]
|
||||
else: inputs = progress[prev_stage]
|
||||
|
||||
Tensor.realize(*[p.to_(GPUS) for p in get_parameters(model)])
|
||||
for batch_idx in tqdm(range(stage_progress(stage_idx), inputs.shape[0], bs)):
|
||||
t1 = time.perf_counter()
|
||||
batch, unpadded_bs = get_batch(inputs, batch_idx, bs)
|
||||
if isinstance(model, OpenClipEncoder): batch = callback(batch, get_batch(tokens, batch_idx, bs)[0].realize())
|
||||
else: batch = callback(batch)
|
||||
# to(GPUS[0]) is necessary for this to work, without that the result is still on GPUS, probably due to a bug
|
||||
batch = batch.to(GPUS[0]).to("CPU")[0:unpadded_bs].realize()
|
||||
progress[stage][batch_idx: batch_idx + bs].assign(batch).realize()
|
||||
# keep track of what our last output was, so we can resume from there if we crash in this loop
|
||||
progress["end"][stage_idx: stage_idx + 1].assign(Tensor([batch_idx + bs], dtype=dtypes.int)).realize()
|
||||
print(f"model: {model}, batch_idx: {batch_idx}, elapsed: {(time.perf_counter() - t1):.2f}")
|
||||
del batch
|
||||
|
||||
jit.reset()
|
||||
Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
|
||||
print(f"done with model: {model}, elapsed: {(time.perf_counter() - t0):.2f}")
|
||||
prev_stage = stage
|
||||
|
||||
inception_stats_fn = str(DATADIR / "coco2014" / "val2014_30k_stats.npz")
|
||||
fid_score = inception.compute_score(progress["inception"].to("CPU"), inception_stats_fn)
|
||||
clip_score = progress["clip"].to(GPUS[0]).mean().item()
|
||||
for name in disk_tensor_names:
|
||||
Path(f"{EVAL_CKPT_DIR}/{name}.bytes").unlink(missing_ok=True)
|
||||
|
||||
if EVAL_SAMPLES and BEAM:
|
||||
print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
|
||||
sys.exit() # Don't eval additional models; we don't care about clip/fid scores when running BEAM on eval sample subset
|
||||
|
||||
return clip_score, fid_score
|
||||
|
||||
# evaluate checkpoints in reverse chronological order
|
||||
for ckpt_iteration, p in sorted(eval_queue, reverse=True):
|
||||
unet_ckpt = safe_load(p)
|
||||
load_state_dict(unet, unet_ckpt)
|
||||
clip_score, fid_score = eval_unet(eval_inputs, unet, model.cond_stage_model, model.first_stage_model, inception, clip_encoder)
|
||||
converged = True if clip_score >= 0.15 and fid_score <= 90 else False
|
||||
print(f"eval results for {EVAL_CKPT_DIR}/{p.name}: clip={clip_score}, fid={fid_score}, converged={converged}")
|
||||
if WANDB:
|
||||
wandb.log({"eval/ckpt_iteration": ckpt_iteration, "eval/clip_score": clip_score, "eval/fid_score": fid_score})
|
||||
if converged and STOP_IF_CONVERGED:
|
||||
print(f"Convergence detected, exiting early before evaluating other checkpoints due to STOP_IF_CONVERGED={STOP_IF_CONVERGED}")
|
||||
sys.exit()
|
||||
|
||||
# for testing
|
||||
return clip_score, fid_score, ckpt_iteration
|
||||
|
||||
if __name__ == "__main__":
|
||||
# inference only
|
||||
Tensor.training = False
|
||||
|
||||
models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
|
||||
for m in models:
|
||||
nm = f"eval_{m}"
|
||||
if nm in globals():
|
||||
print(f"eval {m}")
|
||||
globals()[nm]()
|
||||
69
tinygrad_repo/examples/mlperf/model_spec.py
Normal file
69
tinygrad_repo/examples/mlperf/model_spec.py
Normal file
@@ -0,0 +1,69 @@
|
||||
# load each model here, quick benchmark
|
||||
from tinygrad import Tensor, GlobalCounters
|
||||
from tinygrad.helpers import getenv
|
||||
import numpy as np
|
||||
|
||||
def test_model(model, *inputs):
|
||||
GlobalCounters.reset()
|
||||
out = model(*inputs)
|
||||
if isinstance(out, Tensor): out = out.numpy()
|
||||
# TODO: return event future to still get the time_sum_s without DEBUG=2
|
||||
print(f"{GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.time_sum_s*1000:.2f} ms")
|
||||
|
||||
def spec_resnet():
|
||||
# Resnet50-v1.5
|
||||
from extra.models.resnet import ResNet50
|
||||
mdl = ResNet50()
|
||||
img = Tensor.randn(1, 3, 224, 224)
|
||||
test_model(mdl, img)
|
||||
|
||||
def spec_retinanet():
|
||||
# Retinanet with ResNet backbone
|
||||
from extra.models.resnet import ResNet50
|
||||
from extra.models.retinanet import RetinaNet
|
||||
mdl = RetinaNet(ResNet50(), num_classes=91, num_anchors=9)
|
||||
img = Tensor.randn(1, 3, 224, 224)
|
||||
test_model(mdl, img)
|
||||
|
||||
def spec_unet3d():
|
||||
# 3D UNET
|
||||
from extra.models.unet3d import UNet3D
|
||||
mdl = UNet3D()
|
||||
#mdl.load_from_pretrained()
|
||||
img = Tensor.randn(1, 1, 128, 128, 128)
|
||||
test_model(mdl, img)
|
||||
|
||||
def spec_rnnt():
|
||||
from extra.models.rnnt import RNNT
|
||||
mdl = RNNT()
|
||||
#mdl.load_from_pretrained()
|
||||
x = Tensor.randn(220, 1, 240)
|
||||
y = Tensor.randn(1, 220)
|
||||
test_model(mdl, x, y)
|
||||
|
||||
def spec_bert():
|
||||
from extra.models.bert import BertForQuestionAnswering
|
||||
mdl = BertForQuestionAnswering()
|
||||
#mdl.load_from_pretrained()
|
||||
x = Tensor.randn(1, 384)
|
||||
am = Tensor.randn(1, 384)
|
||||
tt = Tensor(np.random.randint(0, 2, (1, 384)).astype(np.float32))
|
||||
test_model(mdl, x, am, tt)
|
||||
|
||||
def spec_mrcnn():
|
||||
from extra.models.mask_rcnn import MaskRCNN, ResNet
|
||||
mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
|
||||
#mdl.load_from_pretrained()
|
||||
x = Tensor.randn(3, 224, 224)
|
||||
test_model(mdl, [x])
|
||||
|
||||
if __name__ == "__main__":
|
||||
# inference only for now
|
||||
Tensor.training = False
|
||||
|
||||
for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(","):
|
||||
nm = f"spec_{m}"
|
||||
if nm in globals():
|
||||
print(f"testing {m}")
|
||||
globals()[nm]()
|
||||
|
||||
1807
tinygrad_repo/examples/mlperf/model_train.py
Normal file
1807
tinygrad_repo/examples/mlperf/model_train.py
Normal file
File diff suppressed because it is too large
Load Diff
368
tinygrad_repo/examples/mlperf/models/flat_llama.py
Normal file
368
tinygrad_repo/examples/mlperf/models/flat_llama.py
Normal file
@@ -0,0 +1,368 @@
|
||||
import math, os
|
||||
if __name__ == "__main__":
|
||||
os.environ["DEFAULT_FLOAT"] = "bfloat16"
|
||||
os.environ["OPTIM_DTYPE"] = "bfloat16"
|
||||
if "DEV" not in os.environ: os.environ["DEV"] = "NULL::gfx950"
|
||||
# CDNA
|
||||
os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
|
||||
os.environ["ALL2ALL"] = "1"
|
||||
os.environ["USE_ATOMICS"] = "1"
|
||||
if "HK_FLASH_ATTENTION" not in os.environ:
|
||||
os.environ["HK_FLASH_ATTENTION"] = "1"
|
||||
if "ASM_GEMM" not in os.environ:
|
||||
os.environ["ASM_GEMM"] = "1"
|
||||
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
|
||||
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker, round_up
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
|
||||
from extra.llama_kernels.rmsnorm import rmsnorm
|
||||
from extra.llama_kernels import FP8_MAX, local_abs_max
|
||||
|
||||
ASM_GEMM = getenv("ASM_GEMM", 0)
|
||||
FUSED_INPUT_QUANTIZE = getenv("FUSED_INPUT_QUANTIZE", 0)
|
||||
FUSED_ADD_NORM_MUL_QUANTIZE = getenv("FUSED_ADD_NORM_MUL_QUANTIZE", 0)
|
||||
FUSED_SILU_W13 = getenv("FUSED_SILU_W13", 0)
|
||||
SPLIT_W13 = getenv("SPLIT_W13", 0)
|
||||
|
||||
FP8_DTYPE = dtypes.fp8e4m3
|
||||
FP8_GRAD_DTYPE = dtypes.fp8e5m2
|
||||
|
||||
def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
|
||||
new_amax = (local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach().cast(dtypes.float32)
|
||||
scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
|
||||
x_scaled = x * scale
|
||||
x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach()) # STE
|
||||
return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
|
||||
|
||||
def matmul(x:Tensor, w:Tensor, fp8:bool=True, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
|
||||
x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax:Tensor|None=None,
|
||||
grad_amax_state:Tensor|None=None) -> tuple[Tensor,...]:
|
||||
if not fp8:
|
||||
if ASM_GEMM:
|
||||
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
|
||||
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
|
||||
return (x @ w.T,)
|
||||
assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
|
||||
if x_fp8 is None:
|
||||
if FUSED_INPUT_QUANTIZE and amax_x is not None:
|
||||
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
|
||||
x_fp8, x_scale, x_new_amax, _ = quantize_fp8_delayed(x, amax_x, FP8_DTYPE)
|
||||
else:
|
||||
x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
|
||||
if ASM_GEMM:
|
||||
from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
|
||||
if can_use_asm_gemm(x_fp8, w.T):
|
||||
return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale, grad_amax_state=grad_amax_state), x_new_amax, x_fp8
|
||||
return (x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale).cast(dtypes.bfloat16), x_new_amax, x_fp8
|
||||
|
||||
def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor, grad_amax_state:Tensor):
|
||||
if FUSED_ADD_NORM_MUL_QUANTIZE:
|
||||
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
|
||||
x_fp8, x_inv_scale, new_amax, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE)
|
||||
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax, grad_amax_state=grad_amax_state)
|
||||
return out, x_normed, rrms, ret
|
||||
x_normed, rrms = rmsnorm(x, eps)
|
||||
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state)
|
||||
return out, x_normed, rrms, ret
|
||||
|
||||
def add_norm_quantize_matmul(x:Tensor, residual:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor,
|
||||
grad_amax_state:Tensor|None=None):
|
||||
if FUSED_ADD_NORM_MUL_QUANTIZE:
|
||||
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
|
||||
x_fp8, x_inv_scale, new_amax, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE)
|
||||
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, x_scale=x_inv_scale, x_new_amax=new_amax, grad_amax_state=grad_amax_state)
|
||||
return out, h, x_normed, rrms, ret
|
||||
h = x + residual
|
||||
x_normed, rrms = rmsnorm(h, eps)
|
||||
out, *ret = matmul(x_normed * norm, w, amax_x=amax_x, w_inv_scale=w_inv_scale, grad_amax_state=grad_amax_state)
|
||||
return out, h, x_normed, rrms, ret
|
||||
|
||||
def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
|
||||
amax_x2:Tensor,
|
||||
grad_amax_xw13:Tensor, grad_amax_xout:Tensor):
|
||||
if FUSED_SILU_W13:
|
||||
from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
|
||||
x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13)
|
||||
out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2, grad_amax_state=grad_amax_xout)
|
||||
return out, ret
|
||||
hidden = x_w13.shape[-1] // 2
|
||||
x_w1, x_w3 = x_w13[..., :hidden], x_w13[..., hidden:]
|
||||
out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2, grad_amax_state=grad_amax_xout)
|
||||
return out, ret
|
||||
|
||||
class FlatTransformer:
|
||||
def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
|
||||
rope_theta:int=10000, max_context:int=1024):
|
||||
self.vocab_size = vocab_size
|
||||
self.n_layers = n_layers
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
|
||||
self.head_dim = dim // n_heads
|
||||
self.n_rep = self.n_heads // self.n_kv_heads
|
||||
self.hidden_dim = hidden_dim
|
||||
|
||||
scaled_std = 0.02 / math.sqrt(2 * n_layers)
|
||||
|
||||
# Attention
|
||||
self.wqkv, s_qkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
|
||||
self.wo, s_o = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
|
||||
|
||||
# FeedForward
|
||||
if SPLIT_W13:
|
||||
self.w1, s_1 = self.lin_per_layer(dim, hidden_dim)
|
||||
self.w3, s_3 = self.lin_per_layer(dim, hidden_dim)
|
||||
else:
|
||||
self.w13, s_13 = self.lin_per_layer(dim, hidden_dim * 2)
|
||||
self.w2, s_2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
|
||||
|
||||
self.norm_eps = norm_eps
|
||||
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
|
||||
self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
|
||||
|
||||
# output
|
||||
self.norm = nn.RMSNorm(dim, norm_eps)
|
||||
self.tok_embeddings = nn.Embedding(vocab_size, dim)
|
||||
self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
|
||||
self.output = Tensor.normal(1, vocab_size, dim, mean=0.0, std=0.02, dtype=dtypes.bfloat16)
|
||||
self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().is_param_(False)
|
||||
|
||||
def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
|
||||
names = ["xqkv", "xo", "x2"]
|
||||
names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
|
||||
self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
|
||||
grad_names = ["xqkv", "xo", "xout"]
|
||||
grad_names += ["xw1", "xw3"] if SPLIT_W13 else ["xw13"]
|
||||
self._fp8_grad_amax = {name: [_amax() for _ in range(n_layers)] for name in grad_names}
|
||||
w_scales = [("wqkv", s_qkv), ("wo", s_o), ("w2", s_2)]
|
||||
w_scales += [("w1", s_1), ("w3", s_3)] if SPLIT_W13 else [("w13", s_13)]
|
||||
self._fp8_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales}
|
||||
self._fp8_next_inv_scale = {name: s.float().contiguous().is_param_(False) for name, s in w_scales}
|
||||
|
||||
def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
|
||||
if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
|
||||
else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
|
||||
amax = w.abs().flatten(1).max(1).detach()
|
||||
scale = FP8_MAX / (amax + 1e-8)
|
||||
inv_scale = (amax + 1e-8) / FP8_MAX
|
||||
return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE), inv_scale
|
||||
|
||||
def attention(self, x:Tensor, freqs_cis:Tensor, *, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
|
||||
amax_xqkv:Tensor, amax_xo:Tensor, s_qkv:Tensor, s_o:Tensor,
|
||||
grad_amax_xqkv:Tensor, grad_amax_xo:Tensor):
|
||||
bsz, seqlen, _ = x.shape
|
||||
amaxs, saves = [], []
|
||||
|
||||
xqkv, x_normed, rrms, (new_amax, *s) = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
|
||||
amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv)
|
||||
amaxs.append(new_amax)
|
||||
saves.extend([x_normed, rrms, *s, xqkv])
|
||||
xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
|
||||
xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
|
||||
xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
|
||||
|
||||
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
|
||||
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
|
||||
if getenv("HK_FLASH_ATTENTION"):
|
||||
from extra.thunder.amd.fa import flash_attention
|
||||
attn, *save = flash_attention(xq, xk, xv, is_causal=True)
|
||||
saves.extend(save)
|
||||
else:
|
||||
xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
|
||||
attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
|
||||
attn = attn.reshape(bsz, seqlen, -1)
|
||||
|
||||
out, new_amax, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo)
|
||||
amaxs.append(new_amax)
|
||||
saves.extend([*s, out])
|
||||
return out, amaxs, saves
|
||||
|
||||
def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
|
||||
amaxs, saves = [], []
|
||||
|
||||
if SPLIT_W13:
|
||||
h = x + residual
|
||||
x_normed, rrms = rmsnorm(h, self.norm_eps)
|
||||
saves.extend([x_normed, rrms])
|
||||
inp = x_normed * kwargs["ffn_norm"]
|
||||
x_w1, new_amax, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"], grad_amax_state=kwargs["grad_amax_xw1"])
|
||||
amaxs.append(new_amax)
|
||||
saves.extend([*s, x_w1])
|
||||
x_w3, new_amax, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"], grad_amax_state=kwargs["grad_amax_xw3"])
|
||||
amaxs.append(new_amax)
|
||||
saves.extend([*s, x_w3])
|
||||
out, new_amax, *s = matmul(x_w1.silu() * x_w3, kwargs["w2"], amax_x=kwargs["amax_x2"], w_inv_scale=kwargs["s_2"],
|
||||
grad_amax_state=kwargs["grad_amax_xout"])
|
||||
amaxs.append(new_amax)
|
||||
saves.extend([*s, out])
|
||||
else:
|
||||
x_w13, h, x_normed, rrms, (new_amax, *s) = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
|
||||
self.norm_eps, amax_x=kwargs["amax_x13"],
|
||||
grad_amax_state=kwargs["grad_amax_xw13"])
|
||||
amaxs.append(new_amax)
|
||||
saves.extend([x_normed, rrms, *s, x_w13])
|
||||
out, (new_amax, *s) = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
|
||||
grad_amax_xw13=kwargs["grad_amax_xw13"], grad_amax_xout=kwargs["grad_amax_xout"])
|
||||
amaxs.append(new_amax)
|
||||
saves.extend([*s, out])
|
||||
return out, h, amaxs, saves
|
||||
|
||||
@function(precompile=True, precompile_backward=True)
|
||||
def run_layer(self, x:Tensor, freqs_cis:Tensor, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
|
||||
attn, attn_amaxs, attn_saves = self.attention(x, freqs_cis, **attn_kwargs)
|
||||
ffn, h, ffn_amaxs, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
|
||||
h = h + ffn
|
||||
amaxs = tuple(a.detach() for a in (*attn_amaxs, *ffn_amaxs))
|
||||
if save: return (h, *amaxs, *attn_saves, *ffn_saves)
|
||||
else: return (h, *amaxs)
|
||||
|
||||
def shard(self, device:tuple[str, ...], mp:bool=False):
|
||||
from tinygrad.nn.state import get_parameters
|
||||
if not mp:
|
||||
for v in get_parameters(self): v.shard_(device, axis=None)
|
||||
else:
|
||||
# flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
|
||||
def _shard_fp8(name:str, axis:int):
|
||||
getattr(self, name).shard_(device, axis=axis)
|
||||
self._fp8_inv_scale[name] = self._fp8_inv_scale[name].to(device).contiguous().is_param_(False)
|
||||
self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].to(device).contiguous().is_param_(False)
|
||||
Tensor.realize(getattr(self, name), self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
|
||||
_shard_fp8("wqkv", 1) # (n_layers, out, dim) shard out
|
||||
_shard_fp8("wo", 2) # (n_layers, dim, in) shard in
|
||||
if SPLIT_W13:
|
||||
_shard_fp8("w1", 1)
|
||||
_shard_fp8("w3", 1)
|
||||
else:
|
||||
_shard_fp8("w13", 1) # (n_layers, hidden*2, dim) shard out
|
||||
_shard_fp8("w2", 2) # (n_layers, dim, hidden) shard in
|
||||
self.attention_norm.shard_(device, axis=None).realize()
|
||||
self.ffn_norm.shard_(device, axis=None).realize()
|
||||
self.norm.weight.shard_(device, axis=None).realize()
|
||||
self.tok_embeddings.weight.shard_(device, axis=0).realize()
|
||||
self.output.shard_(device, axis=1).realize()
|
||||
self.freqs_cis.shard_(device, axis=None).realize()
|
||||
for amax_dict in (self._fp8_amax, self._fp8_grad_amax):
|
||||
for name in amax_dict:
|
||||
for i in range(len(amax_dict[name])):
|
||||
amax_dict[name][i] = amax_dict[name][i].to(device).contiguous().is_param_(False)
|
||||
|
||||
def __call__(self, tokens:Tensor, save:bool=True):
|
||||
h = self.tok_embeddings(tokens)
|
||||
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
|
||||
a, ga, s = self._fp8_amax, self._fp8_grad_amax, self._fp8_inv_scale
|
||||
for i in range(self.n_layers):
|
||||
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i],
|
||||
amax_xqkv=a["xqkv"][i], amax_xo=a["xo"][i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
|
||||
grad_amax_xqkv=ga["xqkv"][i], grad_amax_xo=ga["xo"][i])
|
||||
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i],
|
||||
amax_x2=a["x2"][i], s_2=s["w2"][i], grad_amax_xout=ga["xout"][i])
|
||||
if SPLIT_W13:
|
||||
ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], amax_x1=a["x1"][i], amax_x3=a["x3"][i],
|
||||
s_1=s["w1"][i], s_3=s["w3"][i], grad_amax_xw1=ga["xw1"][i], grad_amax_xw3=ga["xw3"][i])
|
||||
else:
|
||||
ffn_kwargs.update(w13=self.w13[i], amax_x13=a["x13"][i], s_13=s["w13"][i], grad_amax_xw13=ga["xw13"][i])
|
||||
h, *ret = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
|
||||
amax_names = ["xqkv", "xo"] + (["x1", "x3"] if SPLIT_W13 else ["x13"]) + ["x2"]
|
||||
for name, new_val in zip(amax_names, ret[:len(amax_names)]):
|
||||
a[name][i].assign(new_val)
|
||||
|
||||
logits = matmul(self.norm(h), self.output[0], fp8=False)[0]
|
||||
return logits
|
||||
|
||||
def _get_pads(uop:UOp) -> list[UOp]:
|
||||
if uop.op == Ops.ADD: return _get_pads(uop.src[0]) + _get_pads(uop.src[1])
|
||||
return [uop]
|
||||
|
||||
def apply_grad(grad_buf:Tensor, new_grad:UOp):
|
||||
pads = _get_pads(new_grad)
|
||||
if len(pads) <= 1:
|
||||
new_grad = new_grad.cast(grad_buf.dtype)
|
||||
grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
|
||||
return
|
||||
cur = grad_buf.uop
|
||||
for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
|
||||
if pad.op == Ops.PAD:
|
||||
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
|
||||
buf_slice = cur.shrink(grad_shrink)
|
||||
cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
|
||||
else:
|
||||
cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
|
||||
grad_buf.uop = cur
|
||||
|
||||
if __name__ == "__main__":
|
||||
config = {}
|
||||
BS = config["BS"] = getenv("BS", 16)
|
||||
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
|
||||
SMALL = config["SMALL"] = getenv("SMALL", 0)
|
||||
|
||||
from examples.llama3 import MODEL_PARAMS
|
||||
model_params = MODEL_PARAMS[llama_size:=getenv("LLAMA3_SIZE", "8B")]["args"]
|
||||
# vocab_size from mixtral tokenizer
|
||||
if not SMALL: model_params |= {"vocab_size": 32000}
|
||||
real_vocab_size = model_params['vocab_size']
|
||||
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params["n_layers"] = llama_layers
|
||||
|
||||
# pad vocab
|
||||
if (MP := getenv("MP", 1)) > 1: model_params["vocab_size"] = round_up(model_params["vocab_size"], 256 * MP)
|
||||
vocab_mask:Tensor = Tensor.arange(model_params["vocab_size"]).reshape(1, 1, -1) >= real_vocab_size
|
||||
|
||||
model = FlatTransformer(**model_params, max_context=SEQLEN)
|
||||
|
||||
state = nn.state.get_state_dict(model)
|
||||
print("tensor count:", len(state))
|
||||
|
||||
# shard the model
|
||||
from tinygrad import Device
|
||||
is_dp = (DP := getenv("DP", 1)) > 1
|
||||
is_mp = (MP := getenv("MP", 1)) > 1
|
||||
is_sharding = is_dp or is_mp
|
||||
device_count = max(DP, MP)
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
|
||||
|
||||
model.shard(device, is_mp)
|
||||
|
||||
if is_dp: vocab_mask.shard_(device, axis=None).realize()
|
||||
if is_mp: vocab_mask.shard_(device, axis=2).realize()
|
||||
|
||||
# preallocate all the grad buffers and zero them out
|
||||
grad_dtype = lambda x: dtypes.bfloat16 if x.dtype in dtypes.fp8s else x.dtype
|
||||
grads = {x:x.zeros_like(dtype=grad_dtype(x)).contiguous() for x in state.values() if x.is_param}
|
||||
|
||||
fp8_amax = [t for ts in model._fp8_amax.values() for t in ts]
|
||||
fp8_grad_amax = [t for ts in model._fp8_grad_amax.values() for t in ts]
|
||||
|
||||
# print model size
|
||||
sz = 0
|
||||
for k,v in state.items():
|
||||
print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device} {v.nbytes()/1e9:.2f} GB")
|
||||
sz += v.nbytes()
|
||||
print(f"total sz: {sz/1e9:.2f} GB")
|
||||
|
||||
with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=real_vocab_size, dtype=dtypes.int)
|
||||
with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
|
||||
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
|
||||
if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
|
||||
if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
|
||||
|
||||
@TinyJit
|
||||
def fwd_bwd(tokens:Tensor):
|
||||
with Timing("python forward: "):
|
||||
logits = model(tokens[:, :-1], save=llama_size=="8B")
|
||||
loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
with Timing("python backward: "):
|
||||
for t,g in zip(grads, loss.gradient(*grads)):
|
||||
apply_grad(grads[t], g.uop)
|
||||
with Timing("run fwd_bwd: "): loss.realize(*grads.values(), *fp8_amax, *fp8_grad_amax)
|
||||
|
||||
@TinyJit
|
||||
def optim_step():
|
||||
for g in grads.values(): g.assign(g.zeros_like())
|
||||
Tensor.realize(*grads.values())
|
||||
|
||||
for i in range(6):
|
||||
GlobalCounters.reset()
|
||||
profile_marker(f"step {i}")
|
||||
with Timing(colored(f"*** step {i}: ", "red")):
|
||||
fwd_bwd(tokens)
|
||||
optim_step()
|
||||
print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
|
||||
68
tinygrad_repo/examples/mlperf/models/test_apply_grad.py
Normal file
68
tinygrad_repo/examples/mlperf/models/test_apply_grad.py
Normal file
@@ -0,0 +1,68 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, TinyJit
|
||||
from tinygrad.nn.state import get_parameters
|
||||
from examples.mlperf.models.flat_llama import apply_grad
|
||||
|
||||
class FlatModel:
|
||||
def __init__(self, n_layers:int, dim:int, hidden:int):
|
||||
self.n_layers = n_layers
|
||||
self.w1 = Tensor.uniform(n_layers, dim, hidden, low=-0.1, high=0.1)
|
||||
self.w2 = Tensor.uniform(n_layers, hidden, dim, low=-0.1, high=0.1)
|
||||
self.scale = Tensor.uniform(dim, low=0.9, high=1.1)
|
||||
self.bias = Tensor.zeros(dim).contiguous()
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
h = x
|
||||
for i in range(self.n_layers):
|
||||
h = (h @ self.w1[i]).relu() @ self.w2[i] + h
|
||||
return (h * self.scale + self.bias).sum()
|
||||
|
||||
class TestApplyGradE2E(unittest.TestCase):
|
||||
def _run_with_apply_grad(self, model, xs):
|
||||
grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
|
||||
for x in xs:
|
||||
loss = model(x)
|
||||
for p, g in zip(grads, loss.gradient(*grads)):
|
||||
apply_grad(grads[p], g.uop)
|
||||
Tensor.realize(loss, *grads.values())
|
||||
return [grads[p] for p in get_parameters(model)]
|
||||
|
||||
def _run_reference(self, model, xs):
|
||||
for x in xs: model(x).backward()
|
||||
return [p.grad for p in get_parameters(model)]
|
||||
|
||||
def _assert_close(self, got, expected, atol, rtol):
|
||||
for g, e in zip(got, expected):
|
||||
self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})")
|
||||
|
||||
def _assert_match(self, model, xs, atol, rtol):
|
||||
self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol)
|
||||
|
||||
def test_e2e_single_step(self):
|
||||
model = FlatModel(n_layers=3, dim=8, hidden=16)
|
||||
Tensor.realize(*get_parameters(model))
|
||||
self._assert_match(model, [Tensor.randn(2, 8).realize()], atol=1e-4, rtol=1e-4)
|
||||
|
||||
def test_e2e_multi_step_accumulation(self):
|
||||
model = FlatModel(n_layers=4, dim=8, hidden=16)
|
||||
Tensor.realize(*get_parameters(model))
|
||||
self._assert_match(model, [Tensor.randn(2, 8).realize() for _ in range(3)], atol=1e-4, rtol=1e-4)
|
||||
|
||||
def test_e2e_jit(self):
|
||||
model = FlatModel(n_layers=3, dim=8, hidden=16)
|
||||
Tensor.realize(*get_parameters(model))
|
||||
grads = {p: Tensor.zeros(p.shape, dtype=p.dtype).contiguous().realize() for p in get_parameters(model)}
|
||||
|
||||
@TinyJit
|
||||
def fwd_bwd(x:Tensor):
|
||||
loss = model(x)
|
||||
for p, g in zip(grads, loss.gradient(*grads)): apply_grad(grads[p], g.uop)
|
||||
Tensor.realize(loss, *grads.values())
|
||||
|
||||
xs = [Tensor.randn(2, 8).realize() for _ in range(3)]
|
||||
for x in xs: fwd_bwd(x)
|
||||
self._assert_close([grads[p] for p in get_parameters(model)], self._run_reference(model, xs), atol=1e-3, rtol=1e-3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
137
tinygrad_repo/examples/mlperf/models/test_flat_llama.py
Normal file
137
tinygrad_repo/examples/mlperf/models/test_flat_llama.py
Normal file
@@ -0,0 +1,137 @@
|
||||
import os
|
||||
os.environ["WQKV"] = "1"
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, nn, dtypes
|
||||
from tinygrad.device import Device
|
||||
from examples.mlperf.models.llama import Transformer
|
||||
from examples.mlperf.models.flat_llama import FlatTransformer
|
||||
|
||||
def copy_weights(flat:FlatTransformer, ref:Transformer):
|
||||
n_layers = flat.n_layers
|
||||
Tensor.realize(*nn.state.get_state_dict(ref).values())
|
||||
flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
|
||||
flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
|
||||
flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
|
||||
flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
|
||||
flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
|
||||
flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
|
||||
flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
|
||||
flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
|
||||
flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
|
||||
flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
|
||||
|
||||
class TestFlatLlama(unittest.TestCase):
|
||||
def test_forward_match(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2]])
|
||||
ref_logits = ref(tokens).realize()
|
||||
flat_logits = flat(tokens).realize()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
diff = (ref_logits - flat_logits).abs().max().item()
|
||||
self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
|
||||
|
||||
def test_backward_match(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2, 10]])
|
||||
|
||||
ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
ref_loss.backward()
|
||||
ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
|
||||
|
||||
flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
|
||||
flat_loss.backward()
|
||||
flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
|
||||
|
||||
# check loss matches
|
||||
self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
|
||||
|
||||
# check output weight grad matches
|
||||
diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
|
||||
self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
|
||||
|
||||
# check per-layer weight grads match
|
||||
for i in range(params["n_layers"]):
|
||||
for flat_key, ref_key in [
|
||||
("wqkv", f"layers.{i}.attention.wqkv.weight"),
|
||||
("wo", f"layers.{i}.attention.wo.weight"),
|
||||
("w1", f"layers.{i}.feed_forward.w1.weight"),
|
||||
("w2", f"layers.{i}.feed_forward.w2.weight"),
|
||||
("w3", f"layers.{i}.feed_forward.w3.weight"),
|
||||
]:
|
||||
diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
|
||||
self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
|
||||
def test_forward_match_mp(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
from tinygrad import Device
|
||||
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
flat.shard(devices, mp=True)
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
|
||||
ref_logits = ref(tokens.to(devices[0])).numpy()
|
||||
flat_logits = flat(tokens.shard(devices)).numpy()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT == "CPU", "multi-device CPU test")
|
||||
def test_forward_match_dp(self):
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
from tinygrad import Device
|
||||
devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
flat.shard(devices)
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
|
||||
ref_logits = ref(tokens.to(devices[0])).numpy()
|
||||
flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
|
||||
|
||||
@unittest.skipUnless(dtypes.fp8e4m3 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fp8 not supported on this device")
|
||||
def test_forward_fp8(self):
|
||||
import examples.mlperf.models.flat_llama as flat_llama_mod
|
||||
old_fp8 = flat_llama_mod.FP8
|
||||
try:
|
||||
flat_llama_mod.FP8 = 1
|
||||
Tensor.manual_seed(42)
|
||||
params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
|
||||
ref = Transformer(**params)
|
||||
flat = FlatTransformer(**params)
|
||||
copy_weights(flat, ref)
|
||||
Tensor.realize(*nn.state.get_state_dict(flat).values())
|
||||
|
||||
tokens = Tensor([[1, 50, 100, 999, 2]])
|
||||
ref_logits = ref(tokens).numpy()
|
||||
flat_logits = flat(tokens).numpy()
|
||||
self.assertEqual(ref_logits.shape, flat_logits.shape)
|
||||
# FP8 has lower precision, allow larger tolerance
|
||||
np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
|
||||
finally:
|
||||
flat_llama_mod.FP8 = old_fp8
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
105
tinygrad_repo/examples/mlperf/optim.py
Normal file
105
tinygrad_repo/examples/mlperf/optim.py
Normal file
@@ -0,0 +1,105 @@
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.nn.optim import Optimizer
|
||||
from tinygrad.helpers import FUSE_OPTIM, getenv
|
||||
from tinygrad.uop.ops import UOp, Ops
|
||||
|
||||
STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
|
||||
MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
|
||||
FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)
|
||||
|
||||
def stochastic_round_bf16(x:Tensor) -> Tensor:
|
||||
bits = x.bitcast(dtypes.uint32)
|
||||
if isinstance(x.device, tuple):
|
||||
shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
|
||||
noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
|
||||
else:
|
||||
noise = x.rand_like()
|
||||
noise = (noise * 0xFFFF).cast(dtypes.uint32)
|
||||
return ((bits + noise) & 0xFFFF0000).bitcast(dtypes.float32).cast(dtypes.bfloat16)
|
||||
|
||||
class GradAccClipAdamW(Optimizer):
|
||||
def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
|
||||
super().__init__(params, lr, device, fused)
|
||||
self.b1, self.b2, self.eps, self.wd = b1, b2, eps, weight_decay
|
||||
self.b1_t, self.b2_t = (Tensor.ones((1,), dtype=dtypes.float32, device=self.device) for _ in [b1, b2])
|
||||
self.m = self._new_optim_param()
|
||||
self.v = self._new_optim_param()
|
||||
self.grad_acc, self.clip_norm = grad_acc, clip_norm
|
||||
if MASTER_WEIGHTS and self.params[0].dtype != dtypes.float32:
|
||||
self.master_params:list[Tensor]|None = [p.to(self.device).float().contiguous() for p in self.params]
|
||||
else:
|
||||
self.master_params = None
|
||||
|
||||
def fstep(self, grads:list[Tensor]):
|
||||
if self.fused:
|
||||
out, extra = self._step([], grads)
|
||||
updates = [out[0][self.pos_params[i]:self.pos_params[i+1]].reshape(tt.shape) for i, tt in enumerate(self.params)]
|
||||
else:
|
||||
updates, extra = self._step([], grads)
|
||||
for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
|
||||
# collect inv_scale tensors attached to fp8 params (set by _apply_update)
|
||||
fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
|
||||
fp8_next_inv_scales = [tt._next_inv_scale for tt in self.params if hasattr(tt, '_next_inv_scale')]
|
||||
to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales+fp8_next_inv_scales
|
||||
|
||||
Tensor.realize(*to_realize)
|
||||
return extra[-1]
|
||||
|
||||
def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
|
||||
grads = list(grads)
|
||||
|
||||
for i in range(len(grads)):
|
||||
if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
|
||||
|
||||
if self.fused:
|
||||
grads[0].assign(grads[0] / self.grad_acc)
|
||||
total_norm = grads[0].float().square().sum().sqrt()
|
||||
grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
|
||||
else:
|
||||
for i in range(len(grads)):
|
||||
grads[i].assign(grads[i] / self.grad_acc)
|
||||
total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
|
||||
for i in range(len(grads)):
|
||||
grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
|
||||
|
||||
ret = []
|
||||
self.b1_t *= self.b1
|
||||
self.b2_t *= self.b2
|
||||
for i, g in enumerate(grads):
|
||||
m_new = self.b1 * self.m[i].float() + (1.0 - self.b1) * g.float()
|
||||
v_new = self.b2 * self.v[i].float() + (1.0 - self.b2) * (g.float() * g.float())
|
||||
self.m[i].assign(m_new.cast(self.m[i].dtype))
|
||||
self.v[i].assign(v_new.cast(self.v[i].dtype))
|
||||
m_hat = m_new / (1.0 - self.b1_t)
|
||||
v_hat = v_new / (1.0 - self.b2_t)
|
||||
up = m_hat / (v_hat.sqrt() + self.eps)
|
||||
ret.append(self.lr * up)
|
||||
return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]
|
||||
|
||||
def _apply_update(self, t:Tensor, up:Tensor, master:Tensor|None=None) -> Tensor:
|
||||
w = master if master is not None else t
|
||||
wd = self.wd if t.ndim >= 3 else 0.0
|
||||
up = up.float().shard_like(w) + self.lr.to(w.device) * wd * w.detach()
|
||||
new_w = w.detach() - up
|
||||
if master is not None: master.assign(new_w)
|
||||
# when master is offloaded to a different device than the param, results are resharded back onto the param's (sharded) device
|
||||
offloaded = master is not None and master.device != t.device
|
||||
if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16:
|
||||
out = stochastic_round_bf16(new_w)
|
||||
return out.shard_like(t) if offloaded else out
|
||||
if t.dtype in dtypes.fp8s:
|
||||
from examples.mlperf.models.flat_llama import FP8_MAX
|
||||
# delayed scaling: reuse previous step's inv_scale
|
||||
t._inv_scale.assign(t._next_inv_scale)
|
||||
inv_scale = t._inv_scale.to(new_w.device) if offloaded else t._inv_scale
|
||||
scale = inv_scale.reciprocal().reshape(-1, *([1]*(new_w.ndim-1)))
|
||||
scaled = (new_w * scale).clamp(-FP8_MAX, FP8_MAX)
|
||||
ret = scaled.cast(t.dtype)
|
||||
# update inv_scale for next step from quantized result
|
||||
new_amax = (ret.float().abs().max(axis=tuple(range(1, ret.ndim))) * inv_scale * FP8_AMAX_MARGIN).detach()
|
||||
new_inv = ((new_amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
|
||||
t._next_inv_scale.assign(new_inv.shard_like(t._next_inv_scale) if offloaded else new_inv)
|
||||
return ret.shard_like(t) if offloaded else ret
|
||||
out = new_w.cast(t.dtype)
|
||||
return out.shard_like(t) if offloaded else out
|
||||
19
tinygrad_repo/examples/mlperf/scripts/setup_kits19_dataset.sh
Executable file
19
tinygrad_repo/examples/mlperf/scripts/setup_kits19_dataset.sh
Executable file
@@ -0,0 +1,19 @@
|
||||
#!/bin/bash
|
||||
|
||||
if [ -z $BASEDIR ]; then
|
||||
export BASEDIR="./extra/datasets/"
|
||||
fi
|
||||
|
||||
cd $BASEDIR
|
||||
if [ -d "kits19" ]; then
|
||||
echo "kits19 dataset is already available"
|
||||
else
|
||||
echo "Downloading and preparing kits19 dataset at $BASEDIR"
|
||||
|
||||
git clone https://github.com/neheller/kits19
|
||||
cd kits19
|
||||
pip3 install -r requirements.txt
|
||||
python3 -m starter_code.get_imaging
|
||||
|
||||
echo "Done"
|
||||
fi
|
||||
57
tinygrad_repo/examples/mlperf/scripts/stable_diffusion_downloads.sh
Executable file
57
tinygrad_repo/examples/mlperf/scripts/stable_diffusion_downloads.sh
Executable file
@@ -0,0 +1,57 @@
|
||||
#!/usr/bin/env bash
|
||||
# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
|
||||
|
||||
# setup dirs
|
||||
|
||||
DATA=/raid/datasets/stable_diffusion
|
||||
|
||||
LAION=$DATA/laion-400m/webdataset-moments-filtered
|
||||
COCO=$DATA/coco2014
|
||||
mkdir -p $LAION $COCO
|
||||
|
||||
CKPT=/raid/weights/stable_diffusion
|
||||
mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
|
||||
|
||||
# download data
|
||||
|
||||
# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
|
||||
#RCLONE=""
|
||||
RCLONE="rclone"
|
||||
|
||||
## VAE-encoded image latents, from 6.1M image subset of laion-400m
|
||||
## about 1 TB for whole download
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
|
||||
cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
|
||||
echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
|
||||
|
||||
## prompts and FID statistics from 30k image subset of coco2014
|
||||
## 33 MB
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
|
||||
|
||||
$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
|
||||
$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
|
||||
|
||||
# download checkpoints
|
||||
|
||||
## clip (needed for text and vision encoders for validation)
|
||||
CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
|
||||
CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
|
||||
CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
|
||||
wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
|
||||
wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
|
||||
echo "${CLIP_WEIGHTS_SHA256} ${CKPT}/clip/open_clip_pytorch_model.bin" | sha256sum -c
|
||||
|
||||
## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
|
||||
SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
|
||||
SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
|
||||
wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
|
||||
echo "${SD_WEIGHTS_SHA256} ${CKPT}/sd/512-base-ema.ckpt" | sha256sum -c
|
||||
|
||||
## inception (needed for validation)
|
||||
FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
|
||||
FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
|
||||
wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
|
||||
echo "${FID_WEIGHTS_SHA1} ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth" | sha1sum -c
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/bin/bash
|
||||
|
||||
rocm-smi --setprofile compute
|
||||
rocm-smi --setmclk 3
|
||||
rocm-smi --setperflevel high
|
||||
|
||||
# power cap to 350W
|
||||
echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
|
||||
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"submitter": "tinycorp",
|
||||
"division": "closed",
|
||||
"status": "available",
|
||||
"system_name": "tinybox green",
|
||||
"number_of_nodes": "1",
|
||||
"host_processors_per_node": "1",
|
||||
"host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
|
||||
"host_processor_core_count": "64",
|
||||
"host_processor_frequency": "",
|
||||
"host_processor_caches": "",
|
||||
"host_processor_interconnect": "",
|
||||
"host_memory_capacity": "128GB",
|
||||
"host_storage_type": "NVMe SSD",
|
||||
"host_storage_capacity": "4 TB raid array + 1 TB boot",
|
||||
"host_networking": "",
|
||||
"host_networking_topology": "",
|
||||
"host_memory_configuration": "8x 16GB DDR4",
|
||||
"accelerators_per_node": "6",
|
||||
"accelerator_model_name": "NVIDIA GeForce RTX 4090",
|
||||
"accelerator_host_interconnect": "PCIe 4.0 x16",
|
||||
"accelerator_frequency": "",
|
||||
"accelerator_on-chip_memories": "",
|
||||
"accelerator_memory_configuration": "GDDR6X",
|
||||
"accelerator_memory_capacity": "24GB",
|
||||
"accelerator_interconnect": "",
|
||||
"accelerator_interconnect_topology": "",
|
||||
"cooling": "air",
|
||||
"hw_notes": "",
|
||||
"framework": "tinygrad, commit 0e8aa0e2886bf9a2d3ce093bce87305e182e6d4a",
|
||||
"other_software_stack": {
|
||||
"python": "3.10.12",
|
||||
"CUDA": "12.4"
|
||||
},
|
||||
"operating_system": "Ubuntu 22.04.4",
|
||||
"sw_notes": ""
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"submitter": "tinycorp",
|
||||
"division": "closed",
|
||||
"status": "available",
|
||||
"system_name": "tinybox red",
|
||||
"number_of_nodes": "1",
|
||||
"host_processors_per_node": "1",
|
||||
"host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
|
||||
"host_processor_core_count": "64",
|
||||
"host_processor_frequency": "",
|
||||
"host_processor_caches": "",
|
||||
"host_processor_interconnect": "",
|
||||
"host_memory_capacity": "128GB",
|
||||
"host_storage_type": "NVMe SSD",
|
||||
"host_storage_capacity": "4 TB raid array + 1 TB boot",
|
||||
"host_networking": "",
|
||||
"host_networking_topology": "",
|
||||
"host_memory_configuration": "8x 16GB DDR4",
|
||||
"accelerators_per_node": "6",
|
||||
"accelerator_model_name": "AMD Radeon RX 7900 XTX",
|
||||
"accelerator_host_interconnect": "PCIe 4.0 x16",
|
||||
"accelerator_frequency": "",
|
||||
"accelerator_on-chip_memories": "",
|
||||
"accelerator_memory_configuration": "GDDR6",
|
||||
"accelerator_memory_capacity": "24GB",
|
||||
"accelerator_interconnect": "",
|
||||
"accelerator_interconnect_topology": "",
|
||||
"cooling": "air",
|
||||
"hw_notes": "",
|
||||
"framework": "tinygrad, commit 0e8aa0e2886bf9a2d3ce093bce87305e182e6d4a",
|
||||
"other_software_stack": {
|
||||
"python": "3.10.12",
|
||||
"ROCm": "6.1"
|
||||
},
|
||||
"operating_system": "Ubuntu 22.04.4",
|
||||
"sw_notes": ""
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install tqdm and tensorflow.
|
||||
```
|
||||
pip install tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr + increase mes timeout.
|
||||
Install the custom amdgpu driver per [README](https://github.com/nimlgen/amdgpu_ubuntu_22_04/blob/v6.1.3/readme.md)
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### One time setup
|
||||
|
||||
```
|
||||
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
|
||||
|
||||
export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
|
||||
|
||||
export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
|
||||
|
||||
export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,73 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install tqdm and tensorflow.
|
||||
```
|
||||
pip install tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr + increase mes timeout.
|
||||
Install the custom amdgpu driver per [README](https://github.com/nimlgen/amdgpu_ubuntu_22_04/blob/v6.1.3/readme.md)
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### One time setup
|
||||
|
||||
```
|
||||
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
|
||||
|
||||
export BEAM=3
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
|
||||
|
||||
export BEAM=3
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=66 EVAL_BS=6
|
||||
|
||||
export BEAM=3
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_red_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/bin/bash
|
||||
|
||||
rocm-smi --setprofile compute
|
||||
rocm-smi --setmclk 3
|
||||
rocm-smi --setperflevel high
|
||||
|
||||
# power cap to 350W
|
||||
# echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="."
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export LAZYCACHE=0 RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/bin/bash
|
||||
|
||||
rocm-smi --setprofile compute
|
||||
rocm-smi --setmclk 3
|
||||
rocm-smi --setperflevel high
|
||||
|
||||
# power cap to 350W
|
||||
echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"submitter": "tinycorp",
|
||||
"division": "closed",
|
||||
"status": "Available on-premise",
|
||||
"system_name": "tinybox green",
|
||||
"number_of_nodes": "1",
|
||||
"host_processors_per_node": "1",
|
||||
"host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
|
||||
"host_processor_core_count": "32",
|
||||
"host_processor_vcpu_count": "64",
|
||||
"host_processor_frequency": "",
|
||||
"host_processor_caches": "",
|
||||
"host_processor_interconnect": "",
|
||||
"host_memory_capacity": "128GB",
|
||||
"host_storage_type": "NVMe SSD",
|
||||
"host_storage_capacity": "4 TB raid array + 1 TB boot",
|
||||
"host_networking": "",
|
||||
"host_networking_topology": "",
|
||||
"host_memory_configuration": "8x 16GB DDR4",
|
||||
"accelerators_per_node": "6",
|
||||
"accelerator_model_name": "NVIDIA GeForce RTX 4090",
|
||||
"accelerator_host_interconnect": "PCIe 4.0 x16",
|
||||
"accelerator_frequency": "",
|
||||
"accelerator_on-chip_memories": "",
|
||||
"accelerator_memory_configuration": "GDDR6X",
|
||||
"accelerator_memory_capacity": "24GB",
|
||||
"accelerator_interconnect": "",
|
||||
"accelerator_interconnect_topology": "",
|
||||
"cooling": "air",
|
||||
"hw_notes": "",
|
||||
"framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
|
||||
"other_software_stack": {
|
||||
"python": "3.10.12",
|
||||
"CUDA": "12.4"
|
||||
},
|
||||
"operating_system": "Ubuntu 22.04.4",
|
||||
"sw_notes": ""
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"submitter": "tinycorp",
|
||||
"division": "closed",
|
||||
"status": "Available on-premise",
|
||||
"system_name": "tinybox red",
|
||||
"number_of_nodes": "1",
|
||||
"host_processors_per_node": "1",
|
||||
"host_processor_model_name": "AMD EPYC 7532 32-Core Processor",
|
||||
"host_processor_core_count": "32",
|
||||
"host_processor_vcpu_count": "64",
|
||||
"host_processor_frequency": "",
|
||||
"host_processor_caches": "",
|
||||
"host_processor_interconnect": "",
|
||||
"host_memory_capacity": "128GB",
|
||||
"host_storage_type": "NVMe SSD",
|
||||
"host_storage_capacity": "4 TB raid array + 1 TB boot",
|
||||
"host_networking": "",
|
||||
"host_networking_topology": "",
|
||||
"host_memory_configuration": "8x 16GB DDR4",
|
||||
"accelerators_per_node": "6",
|
||||
"accelerator_model_name": "AMD Radeon RX 7900 XTX",
|
||||
"accelerator_host_interconnect": "PCIe 4.0 x16",
|
||||
"accelerator_frequency": "",
|
||||
"accelerator_on-chip_memories": "",
|
||||
"accelerator_memory_configuration": "GDDR6",
|
||||
"accelerator_memory_capacity": "24GB",
|
||||
"accelerator_interconnect": "",
|
||||
"accelerator_interconnect_topology": "",
|
||||
"cooling": "air",
|
||||
"hw_notes": "",
|
||||
"framework": "tinygrad, commit b5546912e24e0a864b35924da4efa5d71cfe368b",
|
||||
"other_software_stack": {
|
||||
"python": "3.10.12",
|
||||
"ROCm": "6.1.3"
|
||||
},
|
||||
"operating_system": "Ubuntu 22.04.4",
|
||||
"sw_notes": ""
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
# export BEAM_LOG_SURPASS_MAX=1
|
||||
# export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export RESET_STEP=1
|
||||
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,69 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install gdown (for dataset), numpy, tqdm and tensorflow.
|
||||
```
|
||||
pip install gdown numpy tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
### tinybox_8xMI300X
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,14 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,17 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
|
||||
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
|
||||
export TRAIN_STEPS=3900
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,28 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_8xMI300X"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
|
||||
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
|
||||
export TRAIN_STEPS=3900
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init # TODO: without DEBUG=2 it hangs
|
||||
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 DEBUG=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,69 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install gdown (for dataset), numpy, tqdm and tensorflow.
|
||||
```
|
||||
pip install gdown numpy tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
### tinybox_8xMI300X
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,16 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BEAM_LOG_SURPASS_MAX=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,26 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,69 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install gdown (for dataset), numpy, tqdm and tensorflow.
|
||||
```
|
||||
pip install gdown numpy tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
### tinybox_8xMI300X
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,17 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BEAM_LOG_SURPASS_MAX=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export RESET_STEP=1
|
||||
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,31 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_red_${DATETIME}_${SEED}.log"
|
||||
|
||||
export HCQDEV_WAIT_TIMEOUT_MS=100000 # prevents hang?
|
||||
|
||||
# init
|
||||
sleep 5 && sudo rmmod amdgpu || true
|
||||
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
# TODO: AM driver resulted in nan
|
||||
sudo modprobe amdgpu
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=${LOGMLPERF:-1}
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=${DEBUG:-2}
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=${LOGMLPERF:-1}
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
sleep 5 && sudo rmmod amdgpu || true
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/bin/bash
|
||||
|
||||
rocm-smi --setprofile compute
|
||||
rocm-smi --setmclk 3
|
||||
rocm-smi --setperflevel high
|
||||
|
||||
# power cap to 350W
|
||||
echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
|
||||
@@ -0,0 +1,38 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses RetinaNet for SSD.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
Also install the following dependencies:
|
||||
```
|
||||
pip install tqdm numpy pycocotools boto3 pandas torch torchvision
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download data
|
||||
|
||||
Run the following:
|
||||
```
|
||||
BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
|
||||
```
|
||||
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,14 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
# export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=5 DEBUG=2
|
||||
|
||||
python examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
# export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
export RUNMLPERF=1
|
||||
|
||||
python examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,24 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="retinanet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
|
||||
export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,14 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
# export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=5 DEBUG=2
|
||||
|
||||
python examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="retinanet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
|
||||
export BASEDIR="/raid/datasets/openimages"
|
||||
|
||||
# export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
export RUNMLPERF=1
|
||||
|
||||
python examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,72 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
DATETIME=${2:-$(date "+%m%d%H%M")}
|
||||
LOGFILE="${HOME}/logs/sd_mi300x_${DATETIME}.log"
|
||||
# UNET_CKPTDIR must be set: training saves checkpoints to this path, then a separate eval process scans this path to know which checkpoints to eval
|
||||
export UNET_CKPTDIR="${HOME}/stable_diffusion/training_checkpoints/${DATETIME}"
|
||||
mkdir -p "${HOME}/logs" "$UNET_CKPTDIR"
|
||||
|
||||
# run this script in isolation when using the --bg flag
|
||||
if [[ "${1:-}" == "--bg" ]]; then
|
||||
echo "logging output to $LOGFILE"
|
||||
echo "saving UNet checkpoints to $UNET_CKPTDIR"
|
||||
script_path="$(readlink -f "${BASH_SOURCE[0]}")"
|
||||
nohup bash "$script_path" run "$DATETIME" >"$LOGFILE" 2>&1 & disown $!
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# venv management
|
||||
if [[ -d .venv-sd-mlperf ]]; then
|
||||
. .venv-sd-mlperf/bin/activate
|
||||
else
|
||||
python3 -m venv .venv-sd-mlperf && . .venv-sd-mlperf/bin/activate
|
||||
pip install --index-url https://download.pytorch.org/whl/cpu torch && pip install tqdm numpy ftfy regex pillow scipy wandb webdataset
|
||||
fi
|
||||
pip list
|
||||
apt list --installed | grep amdgpu
|
||||
rocm-smi --version
|
||||
modinfo amdgpu | grep version
|
||||
|
||||
export BEAM=2 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 IGNORE_JIT_FIRST_BEAM=1 HCQDEV_WAIT_TIMEOUT_MS=300000
|
||||
export AMD_LLVM=0 # bf16 seems to require this
|
||||
export DATADIR="/raid/datasets/stable_diffusion"
|
||||
export CKPTDIR="/raid/weights/stable_diffusion"
|
||||
export EVAL_CKPT_DIR=$UNET_CKPTDIR
|
||||
export MODEL="stable_diffusion" PYTHONPATH="."
|
||||
export GPUS=8 BS=304
|
||||
export CONTEXT_BS=816 DENOISE_BS=600 DECODE_BS=384 INCEPTION_BS=560 CLIP_BS=240
|
||||
export WANDB=1
|
||||
export PARALLEL=4
|
||||
export PYTHONUNBUFFERED=1
|
||||
sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setperfdeterminism 1500 || exit 1
|
||||
|
||||
# Retry BEAM search if script fails before BEAM COMPLETE is printed, but don't retry after that
|
||||
run_retry(){ local try=0 max=5 code tmp py pgid kids
|
||||
while :; do
|
||||
tmp=$(mktemp)
|
||||
setsid bash -c 'exec env "$@"' _ "$@" > >(tee -a "$LOGFILE" | tee "$tmp") 2>&1 &
|
||||
py=$!; pgid=$(ps -o pgid= -p "$py" | tr -d ' ')
|
||||
wait "$py"; code=$?
|
||||
[[ -n "$pgid" ]] && { kill -TERM -"$pgid" 2>/dev/null; sleep 1; kill -KILL -"$pgid" 2>/dev/null; }
|
||||
kids=$(pgrep -P "$py" || true)
|
||||
while [[ -n "$kids" ]]; do
|
||||
kill -TERM $kids 2>/dev/null; sleep 0.5
|
||||
kids=$(for k in $kids; do pgrep -P "$k" || true; done)
|
||||
done
|
||||
grep -q 'BEAM COMPLETE' "$tmp" && { rm -f "$tmp"; return 1; }
|
||||
rm -f "$tmp"
|
||||
((code==0)) && return 0
|
||||
((try>=max)) && return 2
|
||||
((try++)); sleep 90; echo "try = ${try}"
|
||||
done
|
||||
}
|
||||
|
||||
# Power limiting to 400W is only needed if GPUs fall out of sync (causing 2.2x increased train time) at higher power, which has been observed at 450W
|
||||
sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setpoweroverdrive 750 && \
|
||||
run_retry TOTAL_CKPTS=7 python3 examples/mlperf/model_train.py; (( $? == 2 )) && { echo "training failed before BEAM completion"; exit 2; }
|
||||
sleep 90
|
||||
|
||||
run_retry EVAL_SAMPLES=600 python3 examples/mlperf/model_eval.py; (( $? == 2 )) && { echo "eval failed before BEAM completion"; exit 2; }
|
||||
# Checkpoints will be evaluated in reverse chronological order, even if above training crashed early
|
||||
# STOP_IF_CONVERGED=1: Stop the eval after the first time convergence is detected; no more checkpoints will be evaluated after that.
|
||||
STOP_IF_CONVERGED=1 python3 examples/mlperf/model_eval.py
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"submitter": "tinycorp",
|
||||
"division": "closed",
|
||||
"status": "Available on-premise",
|
||||
"system_name": "tinybox 8xMI300X",
|
||||
"number_of_nodes": "1",
|
||||
"host_processors_per_node": "2",
|
||||
"host_processor_model_name": "AMD EPYC 9354",
|
||||
"host_processor_core_count": "32",
|
||||
"host_processor_vcpu_count": "64",
|
||||
"host_processor_frequency": "",
|
||||
"host_processor_caches": "",
|
||||
"host_processor_interconnect": "",
|
||||
"host_memory_capacity": "2304GB",
|
||||
"host_storage_type": "NVMe SSD",
|
||||
"host_storage_capacity": "3x 4TB raid array",
|
||||
"host_networking": "",
|
||||
"host_networking_topology": "",
|
||||
"host_memory_configuration": "24x 96GB DDR5",
|
||||
"accelerators_per_node": "8",
|
||||
"accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
|
||||
"accelerator_host_interconnect": "PCIe 5.0 x16",
|
||||
"accelerator_frequency": "",
|
||||
"accelerator_on-chip_memories": "",
|
||||
"accelerator_memory_configuration": "HBM3",
|
||||
"accelerator_memory_capacity": "192GB",
|
||||
"accelerator_interconnect": "",
|
||||
"accelerator_interconnect_topology": "",
|
||||
"cooling": "air",
|
||||
"hw_notes": "",
|
||||
"framework": "tinygrad, branch mlperf_training_v5.0",
|
||||
"other_software_stack": {
|
||||
"python": "3.10.16",
|
||||
"ROCm": "3.0.0+94441cb"
|
||||
},
|
||||
"operating_system": "Ubuntu 24.04.1 LTS",
|
||||
"sw_notes": ""
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"submitter": "tinycorp",
|
||||
"division": "closed",
|
||||
"status": "Available on-premise",
|
||||
"system_name": "tinybox green",
|
||||
"number_of_nodes": "1",
|
||||
"host_processors_per_node": "1",
|
||||
"host_processor_model_name": "AMD EPYC 7532",
|
||||
"host_processor_core_count": "32",
|
||||
"host_processor_vcpu_count": "64",
|
||||
"host_processor_frequency": "",
|
||||
"host_processor_caches": "",
|
||||
"host_processor_interconnect": "",
|
||||
"host_memory_capacity": "128GB",
|
||||
"host_storage_type": "NVMe SSD",
|
||||
"host_storage_capacity": "4 TB raid array + 1 TB boot",
|
||||
"host_networking": "",
|
||||
"host_networking_topology": "",
|
||||
"host_memory_configuration": "8x 16GB DDR4",
|
||||
"accelerators_per_node": "6",
|
||||
"accelerator_model_name": "NVIDIA GeForce RTX 4090",
|
||||
"accelerator_host_interconnect": "PCIe 4.0 x16",
|
||||
"accelerator_frequency": "",
|
||||
"accelerator_on-chip_memories": "",
|
||||
"accelerator_memory_configuration": "GDDR6X",
|
||||
"accelerator_memory_capacity": "24GB",
|
||||
"accelerator_interconnect": "",
|
||||
"accelerator_interconnect_topology": "",
|
||||
"cooling": "air",
|
||||
"hw_notes": "",
|
||||
"framework": "tinygrad, branch mlperf_training_v5.0",
|
||||
"other_software_stack": {
|
||||
"python": "3.10.12",
|
||||
"CUDA": "12.4"
|
||||
},
|
||||
"operating_system": "Ubuntu 22.04.4",
|
||||
"sw_notes": ""
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"submitter": "tinycorp",
|
||||
"division": "closed",
|
||||
"status": "Available on-premise",
|
||||
"system_name": "tinybox red",
|
||||
"number_of_nodes": "1",
|
||||
"host_processors_per_node": "1",
|
||||
"host_processor_model_name": "AMD EPYC 7532",
|
||||
"host_processor_core_count": "32",
|
||||
"host_processor_vcpu_count": "64",
|
||||
"host_processor_frequency": "",
|
||||
"host_processor_caches": "",
|
||||
"host_processor_interconnect": "",
|
||||
"host_memory_capacity": "128GB",
|
||||
"host_storage_type": "NVMe SSD",
|
||||
"host_storage_capacity": "4 TB raid array + 1 TB boot",
|
||||
"host_networking": "",
|
||||
"host_networking_topology": "",
|
||||
"host_memory_configuration": "8x 16GB DDR4",
|
||||
"accelerators_per_node": "6",
|
||||
"accelerator_model_name": "AMD Radeon RX 7900 XTX",
|
||||
"accelerator_host_interconnect": "PCIe 4.0 x16",
|
||||
"accelerator_frequency": "",
|
||||
"accelerator_on-chip_memories": "",
|
||||
"accelerator_memory_configuration": "GDDR6",
|
||||
"accelerator_memory_capacity": "24GB",
|
||||
"accelerator_interconnect": "",
|
||||
"accelerator_interconnect_topology": "",
|
||||
"cooling": "air",
|
||||
"hw_notes": "",
|
||||
"framework": "tinygrad, branch mlperf_training_v5.0",
|
||||
"other_software_stack": {
|
||||
"python": "3.10.12"
|
||||
},
|
||||
"operating_system": "Ubuntu 22.04.4",
|
||||
"sw_notes": ""
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
# export BEAM_LOG_SURPASS_MAX=1
|
||||
# export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export RESET_STEP=1
|
||||
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,69 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install gdown (for dataset), numpy, tqdm and tensorflow.
|
||||
```
|
||||
pip install gdown numpy tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
### tinybox_8xMI300X
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,16 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export BENCHMARK=10 BERT_LAYERS=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,19 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
|
||||
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
|
||||
export TRAIN_STEPS=3900
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,30 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_8xMI300X"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
|
||||
|
||||
# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
|
||||
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
|
||||
export TRAIN_STEPS=3900
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
|
||||
|
||||
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,69 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install gdown (for dataset), numpy, tqdm and tensorflow.
|
||||
```
|
||||
pip install gdown numpy tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
### tinybox_8xMI300X
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,16 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BEAM_LOG_SURPASS_MAX=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,27 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,69 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses BERT for NLP.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
Also install gdown (for dataset), numpy, tqdm and tensorflow.
|
||||
```
|
||||
pip install gdown numpy tqdm tensorflow
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
### 1. Download raw data
|
||||
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
|
||||
```
|
||||
|
||||
### 2. Preprocess train and validation data
|
||||
|
||||
Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended.
|
||||
|
||||
#### Training:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
|
||||
```
|
||||
|
||||
Generating a specific topic (Between 0 and 499)
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
|
||||
```
|
||||
|
||||
#### Validation:
|
||||
```
|
||||
BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
|
||||
```
|
||||
## Running
|
||||
|
||||
### tinybox_green
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
|
||||
```
|
||||
|
||||
### tinybox_red
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
### tinybox_8xMI300X
|
||||
|
||||
#### Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,17 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BEAM_LOG_SURPASS_MAX=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export RESET_STEP=1
|
||||
export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
RUNMLPERF=1 python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,30 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="bert"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
|
||||
|
||||
export IGNORE_OOB=1
|
||||
|
||||
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
|
||||
export IGNORE_JIT_FIRST_BEAM=1
|
||||
export BASEDIR="/raid/datasets/wiki"
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=1
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="bert_red_${DATETIME}_${SEED}.log"
|
||||
|
||||
export HCQDEV_WAIT_TIMEOUT_MS=100000 # prevents hang?
|
||||
|
||||
# init
|
||||
sleep 5 && sudo rmmod amdgpu || true
|
||||
BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=2
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=NV
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_green"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=${LOGMLPERF:-1}
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
@@ -0,0 +1,50 @@
|
||||
# 1. Problem
|
||||
|
||||
This problem uses the ResNet-50 CNN to do image classification.
|
||||
|
||||
## Requirements
|
||||
|
||||
Install tinygrad and mlperf-logging from master.
|
||||
```
|
||||
git clone https://github.com/tinygrad/tinygrad.git
|
||||
python3 -m pip install -e ".[mlperf]"
|
||||
```
|
||||
|
||||
### tinybox_green
|
||||
Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
|
||||
This is the default on production tinybox green.
|
||||
|
||||
### tinybox_red
|
||||
Disable cwsr
|
||||
This is the default on production tinybox red.
|
||||
```
|
||||
sudo vi /etc/modprobe.d/amdgpu.conf
|
||||
cat <<EOF > /etc/modprobe.d/amdgpu.conf
|
||||
options amdgpu cwsr_enable=0
|
||||
EOF
|
||||
sudo update-initramfs -u
|
||||
sudo reboot
|
||||
|
||||
# validate
|
||||
sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
|
||||
```
|
||||
|
||||
# 2. Directions
|
||||
|
||||
## Steps to download and verify data
|
||||
|
||||
```
|
||||
IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
|
||||
```
|
||||
|
||||
## Steps for one time setup
|
||||
|
||||
### tinybox_red
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
|
||||
```
|
||||
|
||||
## Steps to run benchmark
|
||||
```
|
||||
examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
|
||||
```
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export BENCHMARK=10 DEBUG=${DEBUG:-2}
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/bin/bash
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="resnet"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
export EVAL_START_EPOCH=3 EVAL_FREQ=4
|
||||
|
||||
export WANDB=1 PARALLEL=0
|
||||
|
||||
python3 examples/mlperf/model_train.py
|
||||
@@ -0,0 +1,26 @@
|
||||
#!/bin/bash
|
||||
set -e # Exit on any error
|
||||
set -o pipefail # Make pipeline fail if any command fails
|
||||
|
||||
export PYTHONPATH="." DEV=AMD
|
||||
export MODEL="resnet"
|
||||
export SUBMISSION_PLATFORM="tinybox_red"
|
||||
export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
|
||||
|
||||
export RESET_STEP=0
|
||||
|
||||
export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
|
||||
|
||||
# pip install -e ".[mlperf]"
|
||||
export LOGMLPERF=${LOGMLPERF:-1}
|
||||
|
||||
export SEED=$RANDOM
|
||||
DATETIME=$(date "+%m%d%H%M")
|
||||
LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
|
||||
|
||||
# init
|
||||
sleep 5 && sudo rmmod amdgpu || true
|
||||
BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
|
||||
|
||||
# run
|
||||
PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user