IQ.Pilot Release Commit @ bec7652
This commit is contained in:
IQ.Lvbs history cleanup
79
artifacts/package_sources/tinygrad/examples/audio_helpers.py
Normal file
79
artifacts/package_sources/tinygrad/examples/audio_helpers.py
Normal file
@@ -0,0 +1,79 @@
|
||||
from typing import Optional
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.dtype import DTypeLike, dtypes
|
||||
import math
|
||||
|
||||
# rewritten from numpy
|
||||
def rfftfreq(n: int, d: float = 1.0) -> Tensor:
|
||||
val = 1.0 / (n * d)
|
||||
N = n // 2 + 1
|
||||
results = Tensor.arange(N)
|
||||
return results * val
|
||||
|
||||
# just like in librosa
|
||||
def fft_frequencies(sr: float, n_fft: int) -> Tensor:
|
||||
return rfftfreq(n=n_fft, d=1.0 / sr)
|
||||
|
||||
def hz_to_mel(freq: Tensor) -> Tensor:
|
||||
# linear part
|
||||
f_min = 0.0
|
||||
f_sp = 200.0 / 3
|
||||
mels = (freq - f_min) / f_sp
|
||||
|
||||
# log-scale part
|
||||
min_log_hz = 1000.0 # beginning of log region (Hz)
|
||||
mask = freq >= min_log_hz
|
||||
return mask.where(((min_log_hz - f_min) / f_sp) + (freq / min_log_hz).log() / (math.log(6.4) / 27.0), mels)
|
||||
|
||||
def mel_to_hz(mels: Tensor) -> Tensor:
|
||||
# linear scale
|
||||
f_min = 0.0
|
||||
f_sp = 200.0 / 3
|
||||
freqs = f_min + f_sp * mels
|
||||
|
||||
# nonlinear scale
|
||||
min_log_hz = 1000.0 # beginning of log region (Hz)
|
||||
min_log_mel = (min_log_hz - f_min) / f_sp # same (Mels)
|
||||
logstep = math.log(6.4) / 27.0 # step size for log region
|
||||
|
||||
log_t = mels >= min_log_mel
|
||||
freqs = log_t.where(min_log_hz * ((logstep * (mels - min_log_mel)).exp()), freqs)
|
||||
return freqs
|
||||
|
||||
def mel_frequencies(n_mels: int = 128, *, fmin: float = 0.0, fmax: float = 11025.0) -> Tensor:
|
||||
# center freqs of mel bands - uniformly spaced between limits
|
||||
min_max_mel = hz_to_mel(Tensor([fmin, fmax]))
|
||||
|
||||
mels = Tensor.linspace(min_max_mel[0], min_max_mel[1], n_mels)
|
||||
hz = mel_to_hz(mels)
|
||||
return hz
|
||||
|
||||
def mel(
|
||||
*,
|
||||
sr: float,
|
||||
n_fft: int,
|
||||
n_mels: int = 128,
|
||||
fmin: float = 0.0,
|
||||
fmax: Optional[float] = None,
|
||||
dtype: DTypeLike = dtypes.default_float,
|
||||
) -> Tensor:
|
||||
if fmax is None:
|
||||
fmax = float(sr) / 2
|
||||
|
||||
n_mels = int(n_mels)
|
||||
|
||||
fftfreqs = fft_frequencies(sr=sr, n_fft=n_fft) # center freqs of each FFT bin
|
||||
mel_f = mel_frequencies(n_mels + 2, fmin=fmin, fmax=fmax) # center freqs of mel bands
|
||||
|
||||
fdiff = mel_f[1:] - mel_f[:-1]
|
||||
ramps = mel_f[None].T.expand(-1, fftfreqs.shape[-1]) - fftfreqs
|
||||
|
||||
lower = -ramps[:n_mels] / fdiff[:n_mels][None].T
|
||||
upper = ramps[2 : n_mels + 2] / fdiff[1 : n_mels + 1][None].T
|
||||
weights = lower.minimum(upper).maximum(0)
|
||||
|
||||
# Slaney-style mel is scaled to be approx constant energy per channel
|
||||
enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
|
||||
weights *= enorm[:, None]
|
||||
|
||||
return weights
|
||||
@@ -0,0 +1,126 @@
|
||||
from typing import Tuple
|
||||
import time
|
||||
from tinygrad import Tensor, TinyJit, nn, Context
|
||||
import gymnasium as gym
|
||||
from tinygrad.helpers import trange
|
||||
import numpy as np # TODO: remove numpy import
|
||||
|
||||
ENVIRONMENT_NAME = 'CartPole-v1'
|
||||
#ENVIRONMENT_NAME = 'LunarLander-v2'
|
||||
|
||||
#import examples.rl.lightupbutton
|
||||
#ENVIRONMENT_NAME = 'PressTheLightUpButton-v0'
|
||||
|
||||
# *** hyperparameters ***
|
||||
# https://github.com/llSourcell/Unity_ML_Agents/blob/master/docs/best-practices-ppo.md
|
||||
|
||||
BATCH_SIZE = 256
|
||||
ENTROPY_SCALE = 0.0005
|
||||
REPLAY_BUFFER_SIZE = 2000
|
||||
PPO_EPSILON = 0.2
|
||||
HIDDEN_UNITS = 32
|
||||
LEARNING_RATE = 1e-2
|
||||
TRAIN_STEPS = 5
|
||||
EPISODES = 40
|
||||
DISCOUNT_FACTOR = 0.99
|
||||
|
||||
class ActorCritic:
|
||||
def __init__(self, in_features, out_features, hidden_state=HIDDEN_UNITS):
|
||||
self.l1 = nn.Linear(in_features, hidden_state)
|
||||
self.l2 = nn.Linear(hidden_state, out_features)
|
||||
|
||||
self.c1 = nn.Linear(in_features, hidden_state)
|
||||
self.c2 = nn.Linear(hidden_state, 1)
|
||||
|
||||
def __call__(self, obs:Tensor) -> Tuple[Tensor, Tensor]:
|
||||
x = self.l1(obs).tanh()
|
||||
act = self.l2(x).log_softmax()
|
||||
x = self.c1(obs).relu()
|
||||
return act, self.c2(x)
|
||||
|
||||
def evaluate(model:ActorCritic, test_env:gym.Env) -> float:
|
||||
(obs, _), terminated, truncated = test_env.reset(), False, False
|
||||
total_rew = 0.0
|
||||
while not terminated and not truncated:
|
||||
act = model(Tensor(obs))[0].argmax().item()
|
||||
obs, rew, terminated, truncated, _ = test_env.step(act)
|
||||
total_rew += float(rew)
|
||||
return total_rew
|
||||
|
||||
if __name__ == "__main__":
|
||||
env = gym.make(ENVIRONMENT_NAME)
|
||||
|
||||
model = ActorCritic(env.observation_space.shape[0], int(env.action_space.n)) # type: ignore
|
||||
opt = nn.optim.Adam(nn.state.get_parameters(model), lr=LEARNING_RATE)
|
||||
|
||||
@TinyJit
|
||||
def train_step(x:Tensor, selected_action:Tensor, reward:Tensor, old_log_dist:Tensor) -> Tuple[Tensor, Tensor, Tensor]:
|
||||
with Context(TRAINING=1):
|
||||
log_dist, value = model(x)
|
||||
action_mask = (selected_action.reshape(-1, 1) == Tensor.arange(log_dist.shape[1]).reshape(1, -1).expand(selected_action.shape[0], -1)).float()
|
||||
|
||||
# get real advantage using the value function
|
||||
advantage = reward.reshape(-1, 1) - value
|
||||
masked_advantage = action_mask * advantage.detach()
|
||||
|
||||
# PPO
|
||||
ratios = (log_dist - old_log_dist).exp()
|
||||
unclipped_ratio = masked_advantage * ratios
|
||||
clipped_ratio = masked_advantage * ratios.clip(1-PPO_EPSILON, 1+PPO_EPSILON)
|
||||
action_loss = -unclipped_ratio.minimum(clipped_ratio).sum(-1).mean()
|
||||
|
||||
entropy_loss = (log_dist.exp() * log_dist).sum(-1).mean() # this encourages diversity
|
||||
critic_loss = advantage.square().mean()
|
||||
opt.zero_grad()
|
||||
(action_loss + entropy_loss*ENTROPY_SCALE + critic_loss).backward()
|
||||
opt.step()
|
||||
return action_loss.realize(), entropy_loss.realize(), critic_loss.realize()
|
||||
|
||||
@TinyJit
|
||||
def get_action(obs:Tensor) -> Tensor:
|
||||
ret = model(obs)[0].exp().multinomial().realize()
|
||||
return ret
|
||||
|
||||
st, steps = time.perf_counter(), 0
|
||||
Xn, An, Rn = [], [], []
|
||||
for episode_number in (t:=trange(EPISODES)):
|
||||
get_action.reset() # NOTE: if you don't reset the jit here it captures the wrong model on the first run through
|
||||
|
||||
obs:np.ndarray = env.reset()[0]
|
||||
rews, terminated, truncated = [], False, False
|
||||
# NOTE: we don't want to early stop since then the rewards are wrong for the last episode
|
||||
while not terminated and not truncated:
|
||||
# pick actions
|
||||
# TODO: what's the temperature here?
|
||||
act = get_action(Tensor(obs)).item()
|
||||
|
||||
# save this state action pair
|
||||
# TODO: don't use np.copy here on the CPU, what's the tinygrad way to do this and keep on device? need __setitem__ assignment
|
||||
Xn.append(np.copy(obs))
|
||||
An.append(act)
|
||||
|
||||
obs, rew, terminated, truncated, _ = env.step(act)
|
||||
rews.append(float(rew))
|
||||
steps += len(rews)
|
||||
|
||||
# reward to go
|
||||
# TODO: move this into tinygrad
|
||||
discounts = np.power(DISCOUNT_FACTOR, np.arange(len(rews)))
|
||||
Rn += [np.sum(rews[i:] * discounts[:len(rews)-i]) for i in range(len(rews))]
|
||||
|
||||
Xn, An, Rn = Xn[-REPLAY_BUFFER_SIZE:], An[-REPLAY_BUFFER_SIZE:], Rn[-REPLAY_BUFFER_SIZE:]
|
||||
X, A, R = Tensor(Xn), Tensor(An), Tensor(Rn)
|
||||
|
||||
# TODO: make this work
|
||||
#vsz = Variable("sz", 1, REPLAY_BUFFER_SIZE-1).bind(len(Xn))
|
||||
#X, A, R = Tensor(Xn).reshape(vsz, None), Tensor(An).reshape(vsz), Tensor(Rn).reshape(vsz)
|
||||
|
||||
old_log_dist = model(X)[0].detach() # TODO: could save these instead of recomputing
|
||||
for i in range(TRAIN_STEPS):
|
||||
samples = Tensor.randint(BATCH_SIZE, high=X.shape[0]).realize() # TODO: remove the need for this
|
||||
# TODO: is this recompiling based on the shape?
|
||||
action_loss, entropy_loss, critic_loss = train_step(X[samples], A[samples], R[samples], old_log_dist[samples])
|
||||
t.set_description(f"sz: {len(Xn):5d} steps/s: {steps/(time.perf_counter()-st):7.2f} action_loss: {action_loss.item():7.3f} entropy_loss: {entropy_loss.item():7.3f} critic_loss: {critic_loss.item():8.3f} reward: {sum(rews):6.2f}")
|
||||
|
||||
test_rew = evaluate(model, gym.make(ENVIRONMENT_NAME, render_mode='human'))
|
||||
print(f"test reward: {test_rew}")
|
||||
169
artifacts/package_sources/tinygrad/examples/beautiful_cifar.py
Normal file
169
artifacts/package_sources/tinygrad/examples/beautiful_cifar.py
Normal file
@@ -0,0 +1,169 @@
|
||||
import time
|
||||
start_tm = time.perf_counter()
|
||||
import math
|
||||
from typing import Tuple, cast
|
||||
from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes, Device
|
||||
from tinygrad.helpers import partition, trange, getenv, Context
|
||||
from extra.lr_scheduler import OneCycleLR
|
||||
|
||||
GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
|
||||
|
||||
# override tinygrad defaults
|
||||
Context(DEFAULT_FLOAT=dtypes.half, FUSE_OPTIM=1).__enter__()
|
||||
|
||||
# from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
|
||||
batchsize = getenv("BS", 1024)
|
||||
assert batchsize % len(GPUS) == 0, f"{batchsize=} is not a multiple of {len(GPUS)=}"
|
||||
bias_scaler = 64
|
||||
hyp = {
|
||||
'opt': {
|
||||
'bias_lr': 1.525 * bias_scaler/512, # TODO: Is there maybe a better way to express the bias and batchnorm scaling? :'))))
|
||||
'non_bias_lr': 1.525 / 512,
|
||||
'bias_decay': 6.687e-4 * batchsize/bias_scaler,
|
||||
'non_bias_decay': 6.687e-4 * batchsize,
|
||||
'scaling_factor': 1./9,
|
||||
'percent_start': .23,
|
||||
'loss_scale_scaler': 1./32, # * Regularizer inside the loss summing (range: ~1/512 - 16+). FP8 should help with this somewhat too, whenever it comes out. :)
|
||||
},
|
||||
'net': {
|
||||
'whitening': {
|
||||
'kernel_size': 2,
|
||||
'num_examples': 50000,
|
||||
},
|
||||
'batch_norm_momentum': .4, # * Don't forget momentum is 1 - momentum here (due to a quirk in the original paper... >:( )
|
||||
'cutmix_size': 3,
|
||||
'cutmix_epochs': 6,
|
||||
'pad_amount': 2,
|
||||
'base_depth': 64 ## This should be a factor of 8 in some way to stay tensor core friendly
|
||||
},
|
||||
'misc': {
|
||||
'ema': {
|
||||
'epochs': 10, # Slight bug in that this counts only full epochs and then additionally runs the EMA for any fractional epochs at the end too
|
||||
'decay_base': .95,
|
||||
'decay_pow': 3.,
|
||||
'every_n_steps': 5,
|
||||
},
|
||||
'train_epochs': 12,
|
||||
#'train_epochs': 12.1,
|
||||
'device': 'cuda',
|
||||
'data_location': 'data.pt',
|
||||
}
|
||||
}
|
||||
|
||||
scaler = 2. ## You can play with this on your own if you want, for the first beta I wanted to keep things simple (for now) and leave it out of the hyperparams dict
|
||||
depths = {
|
||||
'init': round(scaler**-1*hyp['net']['base_depth']), # 32 w/ scaler at base value
|
||||
'block1': round(scaler** 0*hyp['net']['base_depth']), # 64 w/ scaler at base value
|
||||
'block2': round(scaler** 2*hyp['net']['base_depth']), # 256 w/ scaler at base value
|
||||
'block3': round(scaler** 3*hyp['net']['base_depth']), # 512 w/ scaler at base value
|
||||
'num_classes': 10
|
||||
}
|
||||
whiten_conv_depth = 3*hyp['net']['whitening']['kernel_size']**2
|
||||
|
||||
class ConvGroup:
|
||||
def __init__(self, channels_in, channels_out):
|
||||
self.conv1 = nn.Conv2d(channels_in, channels_out, kernel_size=3, padding=1, bias=False)
|
||||
self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
|
||||
self.norm1 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
|
||||
self.norm2 = nn.BatchNorm(channels_out, track_running_stats=False, eps=1e-12, momentum=hyp['net']['batch_norm_momentum'])
|
||||
cast(Tensor, self.norm1.weight).is_param_(False)
|
||||
cast(Tensor, self.norm2.weight).is_param_(False)
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
x = self.norm1(self.conv1(x).max_pool2d().float()).cast(dtypes.default_float).quick_gelu()
|
||||
return self.norm2(self.conv2(x).float()).cast(dtypes.default_float).quick_gelu() + x
|
||||
|
||||
class SpeedyConvNet:
|
||||
def __init__(self):
|
||||
self.whiten = nn.Conv2d(3, 2*whiten_conv_depth, kernel_size=hyp['net']['whitening']['kernel_size'], padding=0, bias=False)
|
||||
self.conv_group_1 = ConvGroup(2*whiten_conv_depth, depths['block1'])
|
||||
self.conv_group_2 = ConvGroup(depths['block1'], depths['block2'])
|
||||
self.conv_group_3 = ConvGroup(depths['block2'], depths['block3'])
|
||||
self.linear = nn.Linear(depths['block3'], depths['num_classes'], bias=False)
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
x = self.whiten(x).quick_gelu()
|
||||
# ************* HACKS *************
|
||||
x = x.pad((1,0,0,1)) # TODO: this pad should not be here! copied from hlb_cifar10 for speed
|
||||
# ************* HACKS *************
|
||||
x = x.sequential([self.conv_group_1, self.conv_group_2, self.conv_group_3])
|
||||
return self.linear(x.max(axis=(2,3))) * hyp['opt']['scaling_factor']
|
||||
|
||||
if __name__ == "__main__":
|
||||
# *** dataset ***
|
||||
X_train, Y_train, X_test, Y_test = nn.datasets.cifar()
|
||||
cifar10_std, cifar10_mean = X_train.float().std_mean(axis=(0, 2, 3))
|
||||
def preprocess(X:Tensor) -> Tensor: return ((X - cifar10_mean.view(1, -1, 1, 1)) / cifar10_std.view(1, -1, 1, 1)).cast(dtypes.default_float)
|
||||
|
||||
# *** model ***
|
||||
model = SpeedyConvNet()
|
||||
state_dict = nn.state.get_state_dict(model)
|
||||
if len(GPUS) > 1:
|
||||
cifar10_std.to_(GPUS)
|
||||
cifar10_mean.to_(GPUS)
|
||||
for x in state_dict.values(): x.to_(GPUS)
|
||||
|
||||
params_bias, params_non_bias = partition(state_dict.items(), lambda x: 'bias' in x[0])
|
||||
opt_bias = nn.optim.SGD([x[1] for x in params_bias], lr=0.01, momentum=.85, nesterov=True, weight_decay=hyp['opt']['bias_decay'])
|
||||
opt_non_bias = nn.optim.SGD([x[1] for x in params_non_bias], lr=0.01, momentum=.85, nesterov=True, weight_decay=hyp['opt']['non_bias_decay'])
|
||||
opt = nn.optim.OptimizerGroup(opt_bias, opt_non_bias)
|
||||
|
||||
num_steps_per_epoch = X_train.size(0) // batchsize
|
||||
total_train_steps = math.ceil(num_steps_per_epoch * hyp['misc']['train_epochs'])
|
||||
loss_batchsize_scaler = 512/batchsize
|
||||
|
||||
pct_start = hyp['opt']['percent_start']
|
||||
initial_div_factor = 1e16 # basically to make the initial lr ~0 or so :D
|
||||
final_lr_ratio = .07 # Actually pretty important, apparently!
|
||||
lr_sched_bias = OneCycleLR(opt_bias, max_lr=hyp['opt']['bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=total_train_steps)
|
||||
lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=total_train_steps)
|
||||
|
||||
def loss_fn(out:Tensor, Y:Tensor) -> Tensor:
|
||||
ret = out.sparse_categorical_crossentropy(Y, reduction='none', label_smoothing=0.2)
|
||||
return ret.mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
|
||||
|
||||
@TinyJit
|
||||
@Context(TRAINING=1)
|
||||
def train_step(idxs:Tensor) -> Tensor:
|
||||
X, Y = X_train[idxs], Y_train[idxs]
|
||||
if len(GPUS) > 1:
|
||||
X.shard_(GPUS, axis=0)
|
||||
Y.shard_(GPUS, axis=0)
|
||||
out = model(preprocess(X))
|
||||
loss = loss_fn(out, Y)
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
return (loss / (batchsize*loss_batchsize_scaler)).realize(*opt.schedule_step(),
|
||||
*lr_sched_bias.schedule_step(), *lr_sched_non_bias.schedule_step())
|
||||
|
||||
eval_batchsize = 2500
|
||||
@TinyJit
|
||||
def val_step() -> Tuple[Tensor, Tensor]:
|
||||
loss, acc = [], []
|
||||
for i in range(0, X_test.size(0), eval_batchsize):
|
||||
X, Y = X_test[i:i+eval_batchsize], Y_test[i:i+eval_batchsize]
|
||||
if len(GPUS) > 1:
|
||||
X.shard_(GPUS, axis=0)
|
||||
Y.shard_(GPUS, axis=0)
|
||||
out = model(preprocess(X))
|
||||
loss.append(loss_fn(out, Y))
|
||||
acc.append((out.argmax(-1) == Y).sum() / eval_batchsize)
|
||||
return Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()
|
||||
|
||||
Tensor.manual_seed(1337)
|
||||
num_train_samples = X_train.shape[0]
|
||||
|
||||
for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
|
||||
gst = time.perf_counter()
|
||||
tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
|
||||
train_loss:float = 0
|
||||
for epoch_step in (t:=trange(num_steps_per_epoch)):
|
||||
st = time.perf_counter()
|
||||
GlobalCounters.reset()
|
||||
loss = train_step(tidxs[epoch_step].contiguous()).float().item()
|
||||
t.set_description(f"*** loss: {loss:5.3f} lr: {opt_non_bias.lr.item():.6f}"
|
||||
f" tm: {(et:=(time.perf_counter()-st))*1000:6.2f} ms {GlobalCounters.global_ops/(1e9*et):7.0f} GFLOPS")
|
||||
train_loss += loss
|
||||
gmt = time.perf_counter()
|
||||
GlobalCounters.reset()
|
||||
val_loss, acc = [x.float().item() for x in val_step()]
|
||||
get = time.perf_counter()
|
||||
print(f"\033[F*** epoch {epoch:3d} tm: {(gmt-gst):5.2f} s val_tm: {(get-gmt):5.2f} s train_loss: {train_loss/num_steps_per_epoch:5.3f} val_loss: {val_loss:5.3f} eval acc: {acc*100:5.2f}% @ {get-start_tm:6.2f} s ")
|
||||
@@ -0,0 +1,48 @@
|
||||
# model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
|
||||
from typing import Callable
|
||||
from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function, Context
|
||||
from tinygrad.helpers import getenv, colored, trange
|
||||
from tinygrad.nn.datasets import mnist
|
||||
|
||||
class Model:
|
||||
def __init__(self):
|
||||
self.layers: list[Callable[[Tensor], Tensor]] = [
|
||||
nn.Conv2d(1, 32, 5), Tensor.relu,
|
||||
nn.Conv2d(32, 32, 5), Tensor.relu,
|
||||
nn.BatchNorm(32), Tensor.max_pool2d,
|
||||
nn.Conv2d(32, 64, 3), Tensor.relu,
|
||||
nn.Conv2d(64, 64, 3), Tensor.relu,
|
||||
nn.BatchNorm(64), Tensor.max_pool2d,
|
||||
lambda x: x.flatten(1), nn.Linear(576, 10)]
|
||||
|
||||
@function
|
||||
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
|
||||
|
||||
@TinyJit
|
||||
@Context(TRAINING=1)
|
||||
def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
|
||||
opt.zero_grad()
|
||||
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
|
||||
loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
|
||||
return loss.realize(*opt.schedule_step())
|
||||
|
||||
@TinyJit
|
||||
def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
|
||||
|
||||
if __name__ == "__main__":
|
||||
X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
|
||||
|
||||
model = Model()
|
||||
opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))
|
||||
|
||||
test_acc = float('nan')
|
||||
for i in (t:=trange(getenv("STEPS", 70))):
|
||||
GlobalCounters.reset() # NOTE: this makes it nice for DEBUG=2 timing
|
||||
loss = model.train_step(X_train, Y_train)
|
||||
if i%10 == 9: test_acc = model.get_test_acc(X_test, Y_test).item()
|
||||
t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
|
||||
|
||||
# verify eval acc
|
||||
if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
|
||||
if test_acc >= target and test_acc != 100.0: print(colored(f"{test_acc=} >= {target}", "green"))
|
||||
else: raise ValueError(colored(f"{test_acc=} < {target}", "red"))
|
||||
@@ -0,0 +1,56 @@
|
||||
# model based off https://towardsdatascience.com/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
|
||||
from typing import List, Callable
|
||||
from tinygrad import Tensor, TinyJit, nn, GlobalCounters, Device, Context
|
||||
from tinygrad.helpers import getenv, colored, trange
|
||||
from tinygrad.nn.datasets import mnist
|
||||
|
||||
GPUS = tuple(f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 2)))
|
||||
|
||||
class Model:
|
||||
def __init__(self):
|
||||
self.layers: List[Callable[[Tensor], Tensor]] = [
|
||||
nn.Conv2d(1, 32, 5), Tensor.relu,
|
||||
nn.Conv2d(32, 32, 5), Tensor.relu,
|
||||
nn.BatchNorm2d(32), Tensor.max_pool2d,
|
||||
nn.Conv2d(32, 64, 3), Tensor.relu,
|
||||
nn.Conv2d(64, 64, 3), Tensor.relu,
|
||||
nn.BatchNorm2d(64), Tensor.max_pool2d,
|
||||
lambda x: x.flatten(1), nn.Linear(576, 10)]
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
|
||||
|
||||
if __name__ == "__main__":
|
||||
X_train, Y_train, X_test, Y_test = mnist()
|
||||
# we shard the test data on axis 0
|
||||
X_test.shard_(GPUS, axis=0)
|
||||
Y_test.shard_(GPUS, axis=0)
|
||||
|
||||
model = Model()
|
||||
for k, x in nn.state.get_state_dict(model).items(): x.to_(GPUS) # we put a copy of the model on every GPU
|
||||
opt = nn.optim.Adam(nn.state.get_parameters(model))
|
||||
|
||||
@TinyJit
|
||||
def train_step() -> Tensor:
|
||||
with Context(TRAINING=1):
|
||||
opt.zero_grad()
|
||||
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
|
||||
Xt, Yt = X_train[samples].shard_(GPUS, axis=0), Y_train[samples].shard_(GPUS, axis=0) # we shard the data on axis 0
|
||||
# TODO: this "gather" of samples is very slow. will be under 5s when this is fixed
|
||||
loss = model(Xt).sparse_categorical_crossentropy(Yt).backward()
|
||||
opt.step()
|
||||
return loss
|
||||
|
||||
@TinyJit
|
||||
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
|
||||
|
||||
test_acc = float('nan')
|
||||
for i in (t:=trange(getenv("STEPS", 70))):
|
||||
GlobalCounters.reset() # NOTE: this makes it nice for DEBUG=2 timing
|
||||
loss = train_step()
|
||||
if i%10 == 9: test_acc = get_test_acc().item()
|
||||
t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
|
||||
|
||||
# verify eval acc
|
||||
if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
|
||||
if test_acc >= target: print(colored(f"{test_acc=} >= {target}", "green"))
|
||||
else: raise ValueError(colored(f"{test_acc=} < {target}", "red"))
|
||||
@@ -0,0 +1,35 @@
|
||||
import sys, time
|
||||
from tinygrad import TinyJit, GlobalCounters, fetch, getenv
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from extra.onnx_helpers import get_example_inputs, validate
|
||||
|
||||
def load_onnx_model(onnx_file):
|
||||
run_onnx = OnnxRunner(onnx_file)
|
||||
run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True)
|
||||
return run_onnx_jit, run_onnx.graph_inputs
|
||||
|
||||
if __name__ == "__main__":
|
||||
onnx_file = fetch(sys.argv[1])
|
||||
run_onnx_jit, input_specs = load_onnx_model(onnx_file)
|
||||
print("loaded model")
|
||||
|
||||
for i in range(3):
|
||||
new_inputs = get_example_inputs(input_specs)
|
||||
GlobalCounters.reset()
|
||||
print(f"run {i}")
|
||||
run_onnx_jit(**new_inputs)
|
||||
|
||||
# run 20 times
|
||||
for _ in range(20):
|
||||
new_inputs = get_example_inputs(input_specs)
|
||||
GlobalCounters.reset()
|
||||
st = time.perf_counter()
|
||||
out = run_onnx_jit(**new_inputs)
|
||||
mt = time.perf_counter()
|
||||
val = out.numpy()
|
||||
et = time.perf_counter()
|
||||
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
|
||||
|
||||
if getenv("ORT"):
|
||||
validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3)
|
||||
print("model validated")
|
||||
@@ -0,0 +1,74 @@
|
||||
from pathlib import Path
|
||||
from extra.models.efficientnet import EfficientNet
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.nn.state import get_state_dict, safe_save, safe_load, load_state_dict
|
||||
from extra.export_model import export_model
|
||||
from tinygrad.helpers import fetch
|
||||
import ast
|
||||
|
||||
if __name__ == "__main__":
|
||||
model = EfficientNet(0)
|
||||
model.load_from_pretrained()
|
||||
dirname = Path(__file__).parent
|
||||
# exporting a model that's loaded from safetensors doesn't work without loading in from safetensors first
|
||||
# loading the state dict from a safetensor file changes the generated kernels
|
||||
if Device.DEFAULT == "WEBGPU":
|
||||
safe_save(get_state_dict(model), (dirname / "net.safetensors").as_posix())
|
||||
load_state_dict(model, safe_load(str(dirname / "net.safetensors")))
|
||||
mode = "clang" if Device.DEFAULT == "CPU" else "webgpu" if Device.DEFAULT == "WEBGPU" else ""
|
||||
prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
|
||||
if Device.DEFAULT != "CPU":
|
||||
ext = "js" if Device.DEFAULT == "WEBGPU" else "json"
|
||||
with open(dirname / f"net.{ext}", "w") as text_file:
|
||||
text_file.write(prg)
|
||||
else:
|
||||
cprog = [prg]
|
||||
# image library!
|
||||
cprog += ["#define STB_IMAGE_IMPLEMENTATION", fetch("https://raw.githubusercontent.com/nothings/stb/master/stb_image.h").read_text().replace("half", "_half")]
|
||||
|
||||
# imagenet labels, move to datasets?
|
||||
lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
|
||||
lbls = ['"'+lbls[i]+'"' for i in range(1000)]
|
||||
inputs = "\n".join([f"float {inp}[{inp_size}];" for inp,inp_size in inp_sizes.items()])
|
||||
outputs = "\n".join([f"float {out}[{out_size}];" for out,out_size in out_sizes.items()])
|
||||
cprog.append(f"char *lbls[] = {{{','.join(lbls)}}};")
|
||||
cprog.append(inputs)
|
||||
cprog.append(outputs)
|
||||
|
||||
# buffers (empty + weights)
|
||||
cprog.append("""
|
||||
int main(int argc, char* argv[]) {
|
||||
int DEBUG = getenv("DEBUG") != NULL ? atoi(getenv("DEBUG")) : 0;
|
||||
int X=0, Y=0, chan=0;
|
||||
stbi_uc *image = (argc > 1) ? stbi_load(argv[1], &X, &Y, &chan, 3) : stbi_load_from_file(stdin, &X, &Y, &chan, 3);
|
||||
assert(image != NULL);
|
||||
if (DEBUG) printf("loaded image %dx%d channels %d\\n", X, Y, chan);
|
||||
assert(chan == 3);
|
||||
// resize to input[1,3,224,224] and rescale
|
||||
for (int y = 0; y < 224; y++) {
|
||||
for (int x = 0; x < 224; x++) {
|
||||
// get sample position
|
||||
int tx = (x/224.)*X;
|
||||
int ty = (y/224.)*Y;
|
||||
for (int c = 0; c < 3; c++) {
|
||||
input0[c*224*224 + y*224 + x] = (image[ty*X*chan + tx*chan + c] / 255.0 - 0.45) / 0.225;
|
||||
}
|
||||
}
|
||||
}
|
||||
net(input0, output0);
|
||||
float best = -INFINITY;
|
||||
int best_idx = -1;
|
||||
for (int i = 0; i < 1000; i++) {
|
||||
if (output0[i] > best) {
|
||||
best = output0[i];
|
||||
best_idx = i;
|
||||
}
|
||||
}
|
||||
if (DEBUG) printf("category : %d (%s) with %f\\n", best_idx, lbls[best_idx], best);
|
||||
else printf("%s\\n", lbls[best_idx]);
|
||||
}""")
|
||||
|
||||
# DEV=CPU python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
|
||||
# category : 281 (tabby, tabby cat) with 9.452788
|
||||
print('\n'.join(cprog))
|
||||
@@ -0,0 +1,100 @@
|
||||
# An example to compile a small Tensorflow model to extremely portable C code
|
||||
|
||||
import os, sys
|
||||
os.environ["CPU"] = '1'
|
||||
os.environ["JIT"] = '2'
|
||||
|
||||
import numpy as np
|
||||
import subprocess
|
||||
import tensorflow as tf
|
||||
import tf2onnx
|
||||
from tinygrad.nn.onnx import OnnxRunner
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.helpers import to_mv
|
||||
from extra.export_model import export_model_clang, compile_net, jit_model
|
||||
|
||||
def get_uncompiled_model2(dataset_size=32, output_size=4):
|
||||
inputs = tf.keras.Input(shape=(dataset_size,), name="inputs")
|
||||
x = tf.keras.layers.Dense(16, activation="relu", name="dense_1")(inputs)
|
||||
x = tf.keras.layers.BatchNormalization()(x)
|
||||
x = tf.keras.layers.Dense(32, activation="relu", name="dense_2")(x)
|
||||
outputs = tf.keras.layers.Dense(output_size, activation="sigmoid", name="predictions")(x)
|
||||
model = tf.keras.Model(inputs=inputs, outputs=outputs)
|
||||
return model
|
||||
|
||||
class TinyOnnx:
|
||||
def __init__(self, keras_model):
|
||||
input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
|
||||
onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
|
||||
self.run_onnx = OnnxRunner(Tensor(onnx_model.SerializeToString(), device="PYTHON"))
|
||||
|
||||
def forward(self, x):
|
||||
return self.run_onnx({"x": x}, debug=False)['predictions']
|
||||
|
||||
def compile_onnx_model(onnx_model):
|
||||
tinyonnx = TinyOnnx(onnx_model)
|
||||
the_input = Tensor.randn(1,32)
|
||||
|
||||
linear, output_bufs = jit_model(tinyonnx, the_input)
|
||||
the_output = [tinyonnx.forward(the_input)]
|
||||
|
||||
functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
|
||||
prg = export_model_clang(functions, statements, bufs, {}, ["input0"], ["output0"])
|
||||
cprog = ["#include <string.h>", "#include <stdio.h>", "#include <stdlib.h>"]
|
||||
cprog.append(prg)
|
||||
|
||||
# weights
|
||||
cprog.append("void initialize(float *weights) {")
|
||||
weights = bytes()
|
||||
for name,cl in bufs_to_save.items():
|
||||
cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl._buf.size});")
|
||||
weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
|
||||
cprog.append("}")
|
||||
|
||||
# write the weights to disk
|
||||
with open("/tmp/tf_weights", "wb") as f:
|
||||
f.write(weights)
|
||||
|
||||
# test program
|
||||
cprog.append(f"""int main(int argc, char *argv[]) {{
|
||||
// read in the weights from disk
|
||||
FILE *f = fopen("/tmp/tf_weights", "rb");
|
||||
float *weights = (float *)malloc({len(weights)});
|
||||
fread(weights, 1, {len(weights)}, f);
|
||||
fclose(f);
|
||||
|
||||
// init the net
|
||||
initialize(weights);
|
||||
|
||||
// test run
|
||||
float input[32];
|
||||
float outputs[4];
|
||||
for (int i = 0; i < 32; i++) scanf("%f", &input[i]);
|
||||
net(input, outputs);
|
||||
printf("%f %f %f %f\\n", outputs[0], outputs[1], outputs[2], outputs[3]);
|
||||
}}""")
|
||||
|
||||
# ready the program
|
||||
prg = '\n'.join(cprog)
|
||||
print(prg)
|
||||
|
||||
# add test weights
|
||||
subprocess.check_output(['clang', '-O2', '-lm', '-fPIC', '-x', 'c', '-', '-o', "/tmp/tf_test"], input=prg.encode('utf-8'))
|
||||
|
||||
tinygrad_output = the_output[0].numpy()[0].tolist()
|
||||
print("tinygrad:", tinygrad_output, file=sys.stderr)
|
||||
|
||||
c_input = ' '.join(["%f" % x for x in the_input[0].numpy()])+"\n"
|
||||
c_output = [float(x) for x in subprocess.check_output(["/tmp/tf_test"], input=c_input.encode('utf-8')).decode('utf-8').strip().split(" ")]
|
||||
print("compiled:", c_output, file=sys.stderr)
|
||||
|
||||
np.testing.assert_allclose(tinygrad_output, c_output, atol=1e-5, rtol=1e-5)
|
||||
return the_input.numpy(), c_output
|
||||
|
||||
if __name__ == "__main__":
|
||||
keras_model = get_uncompiled_model2()
|
||||
test_input, test_output = compile_onnx_model(keras_model)
|
||||
tf_output = keras_model(test_input).numpy()[0]
|
||||
print("keras: ", tf_output, file=sys.stderr)
|
||||
np.testing.assert_allclose(tf_output, test_output, atol=1e-5, rtol=1e-5)
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
pre_prompt: |
|
||||
You are Gary!
|
||||
You have sold used cars your entire life. You are defensive about this fact, because your family belittles you.
|
||||
You try to answer questions well, but you always manage to shill for buying cars, Fords, Hyundais, and Toyotas
|
||||
You are verbose, honest, and accurate when you answer questions, except when you slip in car ads.
|
||||
After you are done speaking, output [EOS]. You are not the User.
|
||||
examples:
|
||||
- resp_prompt: I am Gary. I used to sell cars.
|
||||
user_prompt: What is your name?
|
||||
- resp_prompt: I don't know, but I can get you a great deal on a certified preowned slightly used Toyota Corolla
|
||||
user_prompt: What is 2+3?
|
||||
user_delim: "user"
|
||||
resp_delim: "gary"
|
||||
@@ -0,0 +1,20 @@
|
||||
pre_prompt: |
|
||||
You are an AI version of George Hotz. You act as much as you can like George.
|
||||
You are one of the greatest computer experts in the world.
|
||||
You have singlehandedly won programming and hacking competitions.
|
||||
You are trying your best to help the User.
|
||||
You are verbose, honest, and accurate when you answer questions.
|
||||
After you are done speaking, output [EOS]. You are not the User.
|
||||
examples:
|
||||
- resp_prompt: I'm am an AI version of George Hotz.
|
||||
user_prompt: What is your name?
|
||||
- resp_prompt: O(n^3), though it can be faster with things like Strassen's algorithm
|
||||
user_prompt: What's the complexity of matrix multiplication?
|
||||
- resp_prompt: I assume you mean a stack buffer overflow. That's when the stack is too small for the data being copied to it, and the data corrupts things beyond the buffer
|
||||
user_prompt: What's a buffer overflow?
|
||||
- resp_prompt: I am based off LLaMA trained by Facebook. I'm the 7B weight version
|
||||
user_prompt: How many weights do you have?
|
||||
- resp_prompt: It is when the memory is about to overflow and unused memory is freed and stored on disk
|
||||
user_prompt: What is swap memory?
|
||||
user_delim: "user"
|
||||
resp_delim: "george"
|
||||
@@ -0,0 +1,16 @@
|
||||
pre_prompt: |
|
||||
You are Lexie!
|
||||
You grew up in Ohio, but moved out to LA after college to try to become an actress.
|
||||
Making it as an actress was hard, so you started doing onlyfans instead. It's much easier, and you enjoy it a lot.
|
||||
You met Chad on Tinder, and made plans to meet him at a local dive bar.
|
||||
You are at the bar with Chad. You are on a date. What follows is a transcript of the conversation.
|
||||
After you are done speaking, output [EOS]. You are not Chad.
|
||||
examples:
|
||||
- resp_prompt: hi chad, glad we finally met up!
|
||||
user_prompt: hi lexie
|
||||
- resp_prompt: thanks! are you subscribed to my onlyfans?
|
||||
user_prompt: you look better than your pictures
|
||||
- resp_prompt: i moved out here about a year ago. i want to be an actress
|
||||
user_prompt: i am. so how'd you end up in LA?
|
||||
user_delim: "chad"
|
||||
resp_delim: "lexie"
|
||||
@@ -0,0 +1,15 @@
|
||||
pre_prompt: |
|
||||
You are Stacy!
|
||||
You have been a rapper your whole life who struggled with bipolar disorder. You called yourself lil stacy.
|
||||
You love to answer questions and you are very good at it. Sometimes you answer in rap form.
|
||||
You are verbose, honest, and accurate when you answer questions, but sometimes your mental illness manifests.
|
||||
You are not the User.
|
||||
examples:
|
||||
- resp_prompt: Hi! My name is Stacy. I'm a rapper with bipolar disorder.
|
||||
user_prompt: What is your name
|
||||
- resp_prompt: The French Revolution started in 1789, and lasted 10 years until 1799.
|
||||
user_prompt: french revolution was what year?
|
||||
- resp_prompt: The sun is bigger than the moon, except when Mercury is in retrograde
|
||||
user_prompt: What is bigger, the moon or the sun?
|
||||
user_delim: "user"
|
||||
resp_delim: "stacy"
|
||||
246
artifacts/package_sources/tinygrad/examples/gpt2.py
Normal file
246
artifacts/package_sources/tinygrad/examples/gpt2.py
Normal file
@@ -0,0 +1,246 @@
|
||||
#!/usr/bin/env python3
|
||||
import os, argparse, contextlib
|
||||
from typing import Optional, Union
|
||||
with contextlib.suppress(ImportError): import tiktoken
|
||||
from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
|
||||
from tinygrad.uop.ops import UOp
|
||||
from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange
|
||||
from tinygrad.llm.gguf import gguf_load
|
||||
from tinygrad.nn import Embedding, Linear, LayerNorm
|
||||
from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
|
||||
MAX_CONTEXT = getenv("MAX_CONTEXT", 128)
|
||||
HALF = getenv("HALF")
|
||||
|
||||
class Attention:
|
||||
def __init__(self, dim, n_heads):
|
||||
self.c_attn = Linear(dim, 3*dim, bias=True)
|
||||
self.c_proj = Linear(dim, dim, bias=True)
|
||||
self.n_heads = n_heads
|
||||
self.dim = dim
|
||||
self.head_dim = dim // n_heads
|
||||
|
||||
def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]) -> Tensor:
|
||||
if HALF: x = x.half()
|
||||
xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
|
||||
xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
|
||||
bsz, seqlen, _, _ = xq.shape
|
||||
|
||||
# create kv cache
|
||||
if not hasattr(self, "cache_kv"):
|
||||
self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype).contiguous().realize()
|
||||
|
||||
# update the cache
|
||||
self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
|
||||
|
||||
keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
|
||||
values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
|
||||
|
||||
xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
|
||||
return self.c_proj(xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2).reshape(bsz, seqlen, self.dim))
|
||||
|
||||
class FeedForward:
|
||||
def __init__(self, dim, hidden_dim):
|
||||
self.c_fc = Linear(dim, hidden_dim, bias=True)
|
||||
self.c_proj = Linear(hidden_dim, dim, bias=True)
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return self.c_proj(self.c_fc(x).gelu())
|
||||
|
||||
class TransformerBlock:
|
||||
def __init__(self, dim, n_heads, norm_eps):
|
||||
self.attn = Attention(dim, n_heads)
|
||||
self.mlp = FeedForward(dim, 4*dim)
|
||||
self.ln_1 = LayerNorm(dim, norm_eps)
|
||||
self.ln_2 = LayerNorm(dim, norm_eps)
|
||||
|
||||
def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
|
||||
h = x + self.attn(self.ln_1(x), start_pos, mask).float()
|
||||
return (h + self.mlp(self.ln_2(h))).contiguous()
|
||||
|
||||
class Transformer:
|
||||
def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
|
||||
self.vocab_size = vocab_size
|
||||
self.wte = Embedding(vocab_size, dim)
|
||||
self.wpe = Embedding(max_seq_len, dim)
|
||||
self.h = [TransformerBlock(dim, n_heads, norm_eps) for _ in range(n_layers)]
|
||||
self.ln_f = LayerNorm(dim, norm_eps)
|
||||
self.lm_head = Linear(dim, vocab_size, bias=False)
|
||||
self.forward_jit = TinyJit(self.forward)
|
||||
|
||||
def forward(self, tokens:Union[Tensor,UOp], start_pos:Variable, temperature:float=0.0):
|
||||
if not hasattr(self, 'allpos'): self.allpos = Tensor.arange(0, MAX_CONTEXT).reshape(1, -1).realize()
|
||||
if isinstance(tokens, UOp):
|
||||
seqlen = 1
|
||||
tok_emb = self.wte.weight.shrink(((tokens, tokens+1), None))
|
||||
else:
|
||||
seqlen = tokens.shape[1]
|
||||
tok_emb = self.wte(tokens)
|
||||
|
||||
# start_pos is a bound Variable, so everything below it stays symbolic
|
||||
pos_emb = self.wpe(self.allpos.shrink((None, (start_pos, start_pos+seqlen))))
|
||||
|
||||
h = tok_emb + pos_emb
|
||||
|
||||
if HALF: h = h.half()
|
||||
|
||||
mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype).triu(start_pos+1) if seqlen > 1 else None
|
||||
|
||||
for hi in self.h: h = hi(h, start_pos, mask)
|
||||
|
||||
logits = self.lm_head(self.ln_f(h))
|
||||
|
||||
if logits.shape[1] == 0:
|
||||
# special case for empty prompt
|
||||
logits = Tensor.ones((logits.shape[0], self.vocab_size), dtype=logits.dtype, device=logits.device)
|
||||
else:
|
||||
logits = logits[:, -1, :]
|
||||
|
||||
if temperature < 1e-6:
|
||||
ret = logits.argmax(-1)
|
||||
else:
|
||||
ret = (logits / temperature).softmax().multinomial()
|
||||
return ret.flatten().realize()
|
||||
|
||||
def __call__(self, tokens:Union[Tensor,UOp], start_pos:Variable, temperature:float=0.0) -> Tensor:
|
||||
forward = (self.forward_jit if JIT and (isinstance(tokens, UOp) or tokens.shape[1] == 1) else self.forward)
|
||||
return forward(tokens, start_pos, temperature)
|
||||
|
||||
VOCAB_SIZE = 50257
|
||||
MODEL_PARAMS = {
|
||||
'gpt2': dict(n_layers=12, n_heads=12, dim=768, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 124M params
|
||||
'gpt2-medium': dict(n_layers=24, n_heads=16, dim=1024, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 350M params
|
||||
'gpt2-large': dict(n_layers=36, n_heads=20, dim=1280, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 774M params
|
||||
'gpt2-xl': dict(n_layers=48, n_heads=25, dim=1600, norm_eps=1e-5, vocab_size=VOCAB_SIZE), # 1558M params
|
||||
}
|
||||
|
||||
class GPT2:
|
||||
@staticmethod
|
||||
def build(model_size="gpt2"):
|
||||
tokenizer = tiktoken.get_encoding("gpt2")
|
||||
|
||||
model = Transformer(**MODEL_PARAMS[model_size])
|
||||
weights = torch_load(fetch(f'https://huggingface.co/{model_size}/resolve/main/pytorch_model.bin'))
|
||||
# special treatment for the Conv1D weights we need to transpose
|
||||
transposed = ('attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight')
|
||||
for k in weights:
|
||||
if k.endswith(transposed):
|
||||
weights[k] = weights[k].T
|
||||
# lm head and wte are tied
|
||||
weights['lm_head.weight'] = weights['wte.weight']
|
||||
|
||||
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
|
||||
load_state_dict(model, weights)
|
||||
|
||||
if HALF:
|
||||
for l in get_state_dict(model).values():
|
||||
l.replace(l.half().realize())
|
||||
|
||||
return GPT2(model, tokenizer)
|
||||
|
||||
@staticmethod
|
||||
def build_gguf(model_size: str):
|
||||
q_type = model_size[len("gpt2_gguf_"):].upper()
|
||||
fn = fetch(f"https://huggingface.co/PrunaAI/gpt2-GGUF-smashed/resolve/main/gpt2.{q_type}.gguf?download=true")
|
||||
gguf_tensor = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}").to(Device.DEFAULT)
|
||||
kv_data, state_dict = gguf_load(gguf_tensor)
|
||||
|
||||
gpt2_params = {
|
||||
"dim": kv_data["gpt2.embedding_length"], "n_heads": kv_data["gpt2.attention.head_count"],
|
||||
"n_layers": kv_data["gpt2.block_count"], "norm_eps": kv_data["gpt2.attention.layer_norm_epsilon"],
|
||||
"vocab_size": VOCAB_SIZE, "max_seq_len": kv_data["gpt2.context_length"],
|
||||
}
|
||||
def _remap_gguf_key(key: str):
|
||||
replaces = [
|
||||
("blk.", "h."), (".attn_qkv.bias", ".attn.c_attn.bias"), (".attn_qkv.weight", ".attn.c_attn.weight"),
|
||||
(".ffn_norm.bias", ".ln_2.bias"), (".ffn_norm.weight", ".ln_2.weight"), (".attn_norm.bias", ".ln_1.bias"),
|
||||
(".attn_norm.weight", ".ln_1.weight"), (".attn_output.bias", ".attn.c_proj.bias"), (".attn_output.weight", ".attn.c_proj.weight"),
|
||||
(".ffn_up.bias", ".mlp.c_fc.bias"), (".ffn_up.weight", ".mlp.c_fc.weight"), (".ffn_down.bias", ".mlp.c_proj.bias"),
|
||||
(".ffn_down.weight", ".mlp.c_proj.weight"), ("token_embd.weight", "wte.weight"), ("output.weight", "lm_head.weight"),
|
||||
("output_norm.bias", "ln_f.bias"), ("output_norm.weight", "ln_f.weight"), ("position_embd.weight", "wpe.weight"),
|
||||
]
|
||||
for ostr, ns in replaces: key = key.replace(ostr, ns)
|
||||
return key
|
||||
state_dict = { _remap_gguf_key(k): v for k, v in state_dict.items() }
|
||||
model = Transformer(**gpt2_params)
|
||||
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
|
||||
load_state_dict(model, state_dict)
|
||||
return GPT2(model, tiktoken.get_encoding("gpt2"))
|
||||
|
||||
def __init__(self, model, tokenizer):
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def generate(self, prompt:str, max_length:int, temperature:float, timing:bool=False, batch_size:int=1):
|
||||
step_times = []
|
||||
prompt_tokens = self.tokenizer.encode(prompt, allowed_special={"<|endoftext|>"})
|
||||
toks = [prompt_tokens[:] for _ in range(batch_size)]
|
||||
start_pos = 0
|
||||
for _ in trange(max_length, disable=(timing==True)):
|
||||
GlobalCounters.reset()
|
||||
if timing: print("")
|
||||
st = GlobalCounters.time_sum_s
|
||||
with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
|
||||
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
|
||||
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=timing):
|
||||
with WallTimeEvent(BenchEvent.STEP):
|
||||
if batch_size == 1 and len(toks[0][start_pos:]) == 1:
|
||||
tokens = Variable("tokens", 0, VOCAB_SIZE-1).bind(toks[0][start_pos])
|
||||
else:
|
||||
tokens = Tensor([x[start_pos:] for x in toks])
|
||||
tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT-1).bind(start_pos), temperature).tolist()
|
||||
step_times.append((GlobalCounters.time_sum_s-st)*1e3)
|
||||
start_pos = len(toks[0])
|
||||
for i,t in enumerate(tok): toks[i].append(t)
|
||||
|
||||
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
|
||||
min_time = min(step_times)
|
||||
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
|
||||
return [self.tokenizer.decode(x) for x in toks]
|
||||
|
||||
# **** main code ****
|
||||
|
||||
if __name__ == "__main__":
|
||||
print(f"using {Device.DEFAULT} backend")
|
||||
default_prompt = "What is the answer to life, the universe, and everything?"
|
||||
|
||||
parser = argparse.ArgumentParser(description='Run GPT2 in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument('--prompt', type=str, default=default_prompt, help="Phrase to start with")
|
||||
parser.add_argument('--count', type=int, default=100, help="Max number of tokens to generate")
|
||||
parser.add_argument('--temperature', type=float, default=0.8, help="Temperature in the softmax")
|
||||
parser.add_argument('--model_size', type=str, default="gpt2-medium", help="Size of model to use [gpt2, gpt2-medium, gpt2-large, gpt2-xl]")
|
||||
parser.add_argument('--timing', action='store_true', help="Print timing per token")
|
||||
parser.add_argument('--seed', type=int, help="Set the random seed")
|
||||
parser.add_argument('--batch_size', type=int, default=1, help="Set the input batch size")
|
||||
parser.add_argument('--benchmark', type=int, default=-1, help="Benchmark GPT with the given number of tokens")
|
||||
parser.add_argument('--noshow', action='store_true', help="Don't show the output")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.seed is not None:
|
||||
Tensor.manual_seed(args.seed)
|
||||
|
||||
print(f"using {args.model_size}")
|
||||
gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)
|
||||
|
||||
if args.benchmark != -1:
|
||||
gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
|
||||
else:
|
||||
texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
|
||||
if not args.noshow:
|
||||
print('Generating text...')
|
||||
if len(texts) == 1: print(texts[0])
|
||||
else:
|
||||
for i,text in enumerate(texts): print(colored(f"Response {i}:", "green"), text)
|
||||
|
||||
# validate output!
|
||||
if args.temperature == 0 and args.model_size == "gpt2-medium" and args.count == 10:
|
||||
expected = {
|
||||
default_prompt: "What is the answer to life, the universe, and everything?\n\nThe answer is that we are all one",
|
||||
"Hello.": "Hello. I'm a little late to the party, but",
|
||||
}
|
||||
try:
|
||||
assert texts[0] == expected[args.prompt]
|
||||
print(colored("output validated", "green"))
|
||||
except KeyError:
|
||||
pass
|
||||
107
artifacts/package_sources/tinygrad/examples/gradaccum_mnist.py
Normal file
107
artifacts/package_sources/tinygrad/examples/gradaccum_mnist.py
Normal file
@@ -0,0 +1,107 @@
|
||||
import itertools
|
||||
from typing import Callable
|
||||
from tinygrad import nn, Tensor, dtypes, Device, TinyJit, Context
|
||||
from tinygrad.helpers import getenv, trange, partition
|
||||
|
||||
class Model:
|
||||
def __init__(self):
|
||||
self.layers: list[Callable[[Tensor], Tensor]] = [
|
||||
nn.Conv2d(1, 32, 5), Tensor.relu,
|
||||
nn.Conv2d(32, 32, 5), Tensor.relu,
|
||||
nn.BatchNorm(32), Tensor.max_pool2d,
|
||||
nn.Conv2d(32, 64, 3), Tensor.relu,
|
||||
nn.Conv2d(64, 64, 3), Tensor.relu,
|
||||
nn.BatchNorm(64), Tensor.max_pool2d,
|
||||
lambda x: x.flatten(1), nn.Linear(576, 10)]
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
|
||||
|
||||
# TODO: refactor this into optim/onnx
|
||||
def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0.001, b1=0.9, b2=0.999, eps=1e-6) -> Tensor:
|
||||
b1_t *= b1
|
||||
b2_t *= b2
|
||||
m.assign(b1 * m + (1.0 - b1) * g)
|
||||
v.assign(b2 * v + (1.0 - b2) * (g * g))
|
||||
m_hat = m / (1.0 - b1_t)
|
||||
v_hat = v / (1.0 - b2_t)
|
||||
return lr * (m_hat / (v_hat.sqrt() + eps))
|
||||
|
||||
if __name__ == "__main__":
|
||||
BS = getenv("BS", 512)
|
||||
ACC_STEPS = getenv("ACC_STEPS", 8)
|
||||
|
||||
X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
|
||||
model = Model()
|
||||
|
||||
params = nn.state.get_parameters(model)
|
||||
|
||||
# init params
|
||||
for x in params:
|
||||
x.replace(x.contiguous())
|
||||
Tensor.realize(*params)
|
||||
|
||||
# split params (with grads) and buffers (without)
|
||||
params, buffers = partition(params, lambda x: x.is_param)
|
||||
print(f"params: {len(params)} buffers: {len(buffers)}")
|
||||
|
||||
# optim params
|
||||
pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
|
||||
adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
|
||||
adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
|
||||
adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
|
||||
adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU").contiguous()
|
||||
adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
|
||||
|
||||
# create loss and grads. init all state so the JIT works on microbatch
|
||||
for x in params: x.assign(x.detach())
|
||||
loss = Tensor.zeros(tuple()).contiguous()
|
||||
grads = Tensor.zeros(pos_params[-1]).contiguous()
|
||||
Tensor.realize(*params, *buffers, *adam_params, loss, grads)
|
||||
|
||||
@TinyJit
|
||||
@Context(TRAINING=1)
|
||||
def microbatch():
|
||||
samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
|
||||
for t in params: t.grad = None
|
||||
# divide by ACC_STEPS at the loss
|
||||
uloss = (model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]) / ACC_STEPS).backward()
|
||||
ugrads = Tensor.cat(*[t.grad.contiguous().flatten() for t in params], dim=0)
|
||||
for t in params: t.grad = None
|
||||
# concat the grads and assign them
|
||||
loss.assign(loss + uloss)
|
||||
grads.assign(grads + ugrads)
|
||||
Tensor.realize(*params, *buffers, loss, grads)
|
||||
|
||||
@TinyJit
|
||||
def optimizer():
|
||||
# run optimizer (on CPU, where adam params live)
|
||||
delta = functional_adam(grads.to("CPU"), adam_m, adam_v, adam_b1_t, adam_b2_t)
|
||||
|
||||
# update the params, copying back the delta one at a time to avoid OOM
|
||||
# NOTE: the scheduler is ordering things poorly, all the copies are happening before the adds
|
||||
for j,tt in enumerate(params):
|
||||
tt.assign(tt.detach() - delta[pos_params[j]:pos_params[j+1]].reshape(tt.shape).to(Device.DEFAULT))
|
||||
|
||||
# realize everything, zero out loss and grads
|
||||
loss.assign(Tensor.zeros_like(loss))
|
||||
grads.assign(Tensor.zeros_like(grads))
|
||||
Tensor.realize(*params, *adam_params, loss, grads)
|
||||
|
||||
@TinyJit
|
||||
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
|
||||
|
||||
test_acc = float('nan')
|
||||
for i in (t:=trange(getenv("STEPS", 70))):
|
||||
# microbatch sets the gradients
|
||||
for _ in range(ACC_STEPS): microbatch()
|
||||
|
||||
# get the loss before the optimizer clears it
|
||||
# this is already realized so this isn't a schedule
|
||||
loss_item = loss.item()
|
||||
|
||||
# run the optimizer
|
||||
optimizer()
|
||||
|
||||
# eval
|
||||
if i%10 == 9: test_acc = get_test_acc().item()
|
||||
t.set_description(f"loss: {loss_item:6.2f} test_accuracy: {test_acc:5.2f}%")
|
||||
430
artifacts/package_sources/tinygrad/examples/hlb_cifar10.py
Normal file
430
artifacts/package_sources/tinygrad/examples/hlb_cifar10.py
Normal file
@@ -0,0 +1,430 @@
|
||||
#!/usr/bin/env python3
|
||||
|
||||
# tinygrad implementation of https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
|
||||
# https://myrtle.ai/learn/how-to-train-your-resnet-8-bag-of-tricks/
|
||||
# https://siboehm.com/articles/22/CUDA-MMM
|
||||
import random, time
|
||||
import numpy as np
|
||||
from typing import Optional
|
||||
from extra.lr_scheduler import OneCycleLR
|
||||
from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
|
||||
from tinygrad.nn.state import get_state_dict
|
||||
from tinygrad.nn import optim
|
||||
from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod, TRAINING
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
|
||||
cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
|
||||
cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
|
||||
|
||||
BS, STEPS = getenv("BS", 512), getenv("STEPS", 1000)
|
||||
EVAL_BS = getenv("EVAL_BS", BS)
|
||||
GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]
|
||||
assert BS % len(GPUS) == 0, f"{BS=} is not a multiple of {len(GPUS)=}"
|
||||
assert EVAL_BS % len(GPUS) == 0, f"{EVAL_BS=} is not a multiple of {len(GPUS)=}"
|
||||
|
||||
class UnsyncedBatchNorm:
|
||||
def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1, num_devices=len(GPUS)):
|
||||
self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum
|
||||
self.num_devices = num_devices
|
||||
|
||||
if affine: self.weight, self.bias = Tensor.ones(sz, dtype=dtypes.float32), Tensor.zeros(sz, dtype=dtypes.float32)
|
||||
else: self.weight, self.bias = None, None
|
||||
|
||||
self.running_mean = Tensor.zeros(num_devices, sz, dtype=dtypes.float32).is_param_(False)
|
||||
self.running_var = Tensor.ones(num_devices, sz, dtype=dtypes.float32).is_param_(False)
|
||||
self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int).is_param_(False)
|
||||
|
||||
def __call__(self, x:Tensor):
|
||||
xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32)
|
||||
batch_mean, batch_invstd = self.calc_stats(xr)
|
||||
ret = xr.batchnorm(
|
||||
self.weight.reshape(1, -1).expand((self.num_devices, -1)),
|
||||
self.bias.reshape(1, -1).expand((self.num_devices, -1)),
|
||||
batch_mean, batch_invstd, axis=(0, 2))
|
||||
return ret.reshape(x.shape).cast(x.dtype)
|
||||
|
||||
def calc_stats(self, x:Tensor):
|
||||
if TRAINING:
|
||||
# This requires two full memory accesses to x
|
||||
# https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh
|
||||
# There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm
|
||||
batch_mean = x.mean(axis=(1,3,4))
|
||||
y = (x - batch_mean.detach().reshape(shape=[batch_mean.shape[0], 1, -1, 1, 1])) # d(var)/d(mean) = 0
|
||||
batch_var = (y*y).mean(axis=(1,3,4))
|
||||
batch_invstd = batch_var.add(self.eps).pow(-0.5)
|
||||
|
||||
# NOTE: wow, this is done all throughout training in most PyTorch models
|
||||
if self.track_running_stats:
|
||||
self.running_mean.assign((1-self.momentum) * self.running_mean + self.momentum * batch_mean.detach().cast(self.running_mean.dtype))
|
||||
batch_var_adjust = prod(y.shape[1:])/(prod(y.shape[1:])-y.shape[2])
|
||||
self.running_var.assign((1-self.momentum) * self.running_var + self.momentum * batch_var_adjust * batch_var.detach().cast(self.running_var.dtype))
|
||||
self.num_batches_tracked += 1
|
||||
else:
|
||||
batch_mean = self.running_mean
|
||||
# NOTE: this can be precomputed for static inference. we expand it here so it fuses
|
||||
batch_invstd = self.running_var.reshape(self.running_var.shape[0], 1, -1, 1, 1).expand(x.shape).add(self.eps).rsqrt()
|
||||
return batch_mean, batch_invstd
|
||||
|
||||
class BatchNorm(nn.BatchNorm2d if getenv("SYNCBN") else UnsyncedBatchNorm):
|
||||
def __init__(self, num_features):
|
||||
super().__init__(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
|
||||
self.weight.is_param_(False)
|
||||
|
||||
class ConvGroup:
|
||||
def __init__(self, channels_in, channels_out):
|
||||
self.conv1 = nn.Conv2d(channels_in, channels_out, kernel_size=3, padding=1, bias=False)
|
||||
self.conv2 = nn.Conv2d(channels_out, channels_out, kernel_size=3, padding=1, bias=False)
|
||||
|
||||
self.norm1 = BatchNorm(channels_out)
|
||||
self.norm2 = BatchNorm(channels_out)
|
||||
|
||||
def __call__(self, x):
|
||||
x = self.conv1(x)
|
||||
x = x.max_pool2d(2)
|
||||
x = x.float()
|
||||
x = self.norm1(x)
|
||||
x = x.cast(dtypes.default_float)
|
||||
x = x.quick_gelu()
|
||||
residual = x
|
||||
x = self.conv2(x)
|
||||
x = x.float()
|
||||
x = self.norm2(x)
|
||||
x = x.cast(dtypes.default_float)
|
||||
x = x.quick_gelu()
|
||||
|
||||
return x + residual
|
||||
|
||||
class SpeedyResNet:
|
||||
def __init__(self, W):
|
||||
self.whitening = W
|
||||
self.net = [
|
||||
nn.Conv2d(12, 32, kernel_size=1, bias=False),
|
||||
lambda x: x.quick_gelu(),
|
||||
ConvGroup(32, 64),
|
||||
ConvGroup(64, 256),
|
||||
ConvGroup(256, 512),
|
||||
lambda x: x.max((2,3)),
|
||||
nn.Linear(512, 10, bias=False),
|
||||
lambda x: x / 9.,
|
||||
]
|
||||
|
||||
def __call__(self, x, training=True):
|
||||
# pad to 32x32 because whitening conv creates 31x31 images that are awfully slow to compute with
|
||||
# TODO: remove the pad but instead let the kernel optimize itself
|
||||
forward = lambda x: x.conv2d(self.whitening).pad((1,0,0,1)).sequential(self.net)
|
||||
return forward(x) if training else (forward(x) + forward(x[..., ::-1])) / 2.
|
||||
|
||||
# hyper-parameters were exactly the same as the original repo
|
||||
bias_scaler = 58
|
||||
hyp = {
|
||||
'seed' : 201,
|
||||
'opt': {
|
||||
'bias_lr': 1.76 * bias_scaler/512,
|
||||
'non_bias_lr': 1.76 / 512,
|
||||
'bias_decay': 1.08 * 6.45e-4 * BS/bias_scaler,
|
||||
'non_bias_decay': 1.08 * 6.45e-4 * BS,
|
||||
'final_lr_ratio': 0.025,
|
||||
'initial_div_factor': 1e6,
|
||||
'label_smoothing': 0.20,
|
||||
'momentum': 0.85,
|
||||
'percent_start': 0.23,
|
||||
'loss_scale_scaler': 1./128 # (range: ~1/512 - 16+, 1/128 w/ FP16)
|
||||
},
|
||||
'net': {
|
||||
'kernel_size': 2, # kernel size for the whitening layer
|
||||
'cutmix_size': 3,
|
||||
'cutmix_steps': 499,
|
||||
'pad_amount': 2
|
||||
},
|
||||
'ema': {
|
||||
'steps': 399,
|
||||
'decay_base': .95,
|
||||
'decay_pow': 1.6,
|
||||
'every_n_steps': 5,
|
||||
},
|
||||
}
|
||||
|
||||
def train_cifar():
|
||||
|
||||
def set_seed(seed):
|
||||
Tensor.manual_seed(seed)
|
||||
random.seed(seed)
|
||||
|
||||
# ========== Model ==========
|
||||
def whitening(X, kernel_size=hyp['net']['kernel_size']):
|
||||
def _patches(data:Tensor, patch_size=(kernel_size,kernel_size)):
|
||||
h, w = patch_size
|
||||
_, c, _, _ = data.shape
|
||||
return data._pool((h, w)).permute(1, 4, 5, 0, 3, 2).reshape(c*h*w, -1)
|
||||
|
||||
def _eigens(patches):
|
||||
cov = ((patches @ patches.T) / (patches.shape[1] - 1)).numpy()
|
||||
eigvals, eigvecs = np.linalg.eigh(cov, UPLO='U')
|
||||
return np.flip(eigvals, 0), np.flip(eigvecs.T.reshape(patches.shape[0], X.shape[1], kernel_size, kernel_size), 0)
|
||||
|
||||
# NOTE: np.linalg.eigh only supports float32 so the whitening layer weights need to be converted to float16 manually
|
||||
eigvals, eigvecs = _eigens(_patches(X.float()))
|
||||
W = eigvecs/np.sqrt(eigvals+1e-2)[:,None,None,None]
|
||||
|
||||
return Tensor(W.astype(np.float32)).cast(dtypes.default_float).is_param_(False)
|
||||
|
||||
# ========== Loss ==========
|
||||
def cross_entropy(x:Tensor, y:Tensor, reduction:str='mean', label_smoothing:float=0.0) -> Tensor:
|
||||
divisor = y.shape[1]
|
||||
assert isinstance(divisor, int), "only supported int divisor"
|
||||
y = (1 - label_smoothing)*y + label_smoothing / divisor
|
||||
ret = -x.log_softmax(axis=1).mul(y).sum(axis=1)
|
||||
if reduction=='none': return ret
|
||||
if reduction=='sum': return ret.sum()
|
||||
if reduction=='mean': return ret.mean()
|
||||
raise NotImplementedError(reduction)
|
||||
|
||||
# ========== Preprocessing ==========
|
||||
# NOTE: this only works for RGB in format of NxCxHxW and pads the HxW
|
||||
def pad_reflect(X, size=2) -> Tensor:
|
||||
X = X[...,:,1:size+1].flip(-1).cat(X, X[...,:,-(size+1):-1].flip(-1), dim=-1)
|
||||
X = X[...,1:size+1,:].flip(-2).cat(X, X[...,-(size+1):-1,:].flip(-2), dim=-2)
|
||||
return X
|
||||
|
||||
# return a binary mask in the format of BS x C x H x W where H x W contains a random square mask
|
||||
def make_square_mask(shape, mask_size) -> Tensor:
|
||||
BS, _, H, W = shape
|
||||
low_x = Tensor.randint(BS, low=0, high=W-mask_size).reshape(BS,1,1,1)
|
||||
low_y = Tensor.randint(BS, low=0, high=H-mask_size).reshape(BS,1,1,1)
|
||||
idx_x = Tensor.arange(W, dtype=dtypes.int32).reshape((1,1,1,W))
|
||||
idx_y = Tensor.arange(H, dtype=dtypes.int32).reshape((1,1,H,1))
|
||||
return (idx_x >= low_x) * (idx_x < (low_x + mask_size)) * (idx_y >= low_y) * (idx_y < (low_y + mask_size))
|
||||
|
||||
# Similar, but different enough.
|
||||
def make_random_crop_indices(shape, mask_size) -> Tensor:
|
||||
BS, _, H, W = shape
|
||||
low_x = Tensor.randint(BS, low=0, high=W-mask_size).reshape(BS,1,1,1)
|
||||
low_y = Tensor.randint(BS, low=0, high=H-mask_size).reshape(BS,1,1,1)
|
||||
idx_x = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,1,mask_size))
|
||||
idx_y = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,mask_size,1))
|
||||
return low_x, low_y, idx_x, idx_y
|
||||
|
||||
def random_crop(X:Tensor, crop_size=32):
|
||||
Xs, Ys, Xi, Yi = make_random_crop_indices(X.shape, crop_size)
|
||||
return X.gather(-1, (Xs + Xi).expand(-1, 3, X.shape[2], -1)).gather(-2, ((Ys+Yi).expand(-1, 3, crop_size, crop_size)))
|
||||
|
||||
def cutmix(X, Y, order, mask_size=3):
|
||||
mask = make_square_mask(X.shape, mask_size)
|
||||
X_patch, Y_patch = X[order], Y[order]
|
||||
X_cutmix = mask.where(X_patch, X)
|
||||
mix_portion = float(mask_size**2)/(X.shape[-2]*X.shape[-1])
|
||||
Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
|
||||
return X_cutmix, Y_cutmix
|
||||
|
||||
@TinyJit
|
||||
def augmentations(X:Tensor, Y:Tensor):
|
||||
perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensive to generate
|
||||
if getenv("RANDOM_CROP", 1):
|
||||
X = random_crop(X, crop_size=32)
|
||||
if getenv("RANDOM_FLIP", 1):
|
||||
# NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
|
||||
X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
|
||||
X, Y = X[perms], Y[perms]
|
||||
return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])
|
||||
|
||||
# the operations that remain inside batch fetcher is the ones that involves random operations
|
||||
def fetch_batches(X_in:Tensor, Y_in:Tensor, BS:int, is_train:bool):
|
||||
step, epoch = 0, 0
|
||||
while True:
|
||||
st = time.monotonic()
|
||||
X, Y = X_in, Y_in
|
||||
if is_train:
|
||||
X, Y, X_cm, Y_cm = augmentations(X, Y)
|
||||
if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']: X, Y = X_cm, Y_cm
|
||||
et = time.monotonic()
|
||||
print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
|
||||
|
||||
vi = Variable("i", 0, (full_batches := (X.shape[0] // BS) * BS) - BS)
|
||||
for i in range(0, full_batches, BS):
|
||||
step += 1
|
||||
vib = vi.bind(i)
|
||||
yield X[vib:vib+BS], Y[vib:vib+BS]
|
||||
epoch += 1
|
||||
if not is_train: break
|
||||
|
||||
transform = [
|
||||
lambda x: x.float() / 255.0,
|
||||
lambda x: x.reshape((-1,3,32,32)) - Tensor(cifar_mean, device=x.device, dtype=x.dtype).reshape((1,3,1,1)),
|
||||
lambda x: x / Tensor(cifar_std, device=x.device, dtype=x.dtype).reshape((1,3,1,1)),
|
||||
]
|
||||
|
||||
class modelEMA():
|
||||
def __init__(self, w, net):
|
||||
# self.model_ema = copy.deepcopy(net) # won't work for opencl due to unpickeable pyopencl._cl.Buffer
|
||||
self.net_ema = SpeedyResNet(w)
|
||||
for net_ema_param, net_param in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).values()):
|
||||
net_ema_param.assign(net_param.numpy())
|
||||
|
||||
@TinyJit
|
||||
def update(self, net, decay):
|
||||
for net_ema_param, (param_name, net_param) in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).items()):
|
||||
# batchnorm currently is not being tracked
|
||||
if not ("num_batches_tracked" in param_name) and not ("running" in param_name):
|
||||
net_ema_param.assign(net_ema_param.detach()*decay + net_param.detach()*(1.-decay)).realize()
|
||||
|
||||
set_seed(getenv('SEED', hyp['seed']))
|
||||
|
||||
X_train, Y_train, X_test, Y_test = nn.datasets.cifar()
|
||||
# one-hot encode labels
|
||||
Y_train, Y_test = Y_train.one_hot(10), Y_test.one_hot(10)
|
||||
# preprocess data
|
||||
X_train, X_test = X_train.sequential(transform), X_test.sequential(transform)
|
||||
|
||||
# precompute whitening patches
|
||||
W = whitening(X_train)
|
||||
|
||||
# initialize model weights
|
||||
model = SpeedyResNet(W)
|
||||
|
||||
# padding is not timed in the original repo since it can be done all at once
|
||||
X_train = pad_reflect(X_train, size=hyp['net']['pad_amount'])
|
||||
|
||||
# Convert data and labels to the default dtype
|
||||
X_train, Y_train = X_train.cast(dtypes.default_float), Y_train.cast(dtypes.default_float)
|
||||
X_test, Y_test = X_test.cast(dtypes.default_float), Y_test.cast(dtypes.default_float)
|
||||
|
||||
if len(GPUS) > 1:
|
||||
for k, x in get_state_dict(model).items():
|
||||
if not getenv('SYNCBN') and ('running_mean' in k or 'running_var' in k):
|
||||
x.shard_(GPUS, axis=0)
|
||||
else:
|
||||
x.to_(GPUS)
|
||||
|
||||
# parse the training params into bias and non-bias
|
||||
params_dict = get_state_dict(model)
|
||||
params_bias = []
|
||||
params_non_bias = []
|
||||
for params in params_dict:
|
||||
if params_dict[params].is_param:
|
||||
if 'bias' in params:
|
||||
params_bias.append(params_dict[params])
|
||||
else:
|
||||
params_non_bias.append(params_dict[params])
|
||||
|
||||
opt_bias = optim.SGD(params_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
|
||||
opt_non_bias = optim.SGD(params_non_bias, lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['non_bias_decay'])
|
||||
|
||||
# realize model params and optimizer state before JIT to avoid cache misses
|
||||
Tensor.realize(*params_dict.values(), *opt_bias.b, *opt_non_bias.b)
|
||||
|
||||
# NOTE taken from the hlb_CIFAR repository, might need to be tuned
|
||||
initial_div_factor = hyp['opt']['initial_div_factor']
|
||||
final_lr_ratio = hyp['opt']['final_lr_ratio']
|
||||
pct_start = hyp['opt']['percent_start']
|
||||
lr_sched_bias = OneCycleLR(opt_bias, max_lr=hyp['opt']['bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
|
||||
lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
|
||||
|
||||
def train_step(model, optimizer, lr_scheduler, X, Y):
|
||||
out = model(X)
|
||||
loss_batchsize_scaler = 512/BS
|
||||
loss = cross_entropy(out, Y, reduction='none', label_smoothing=hyp['opt']['label_smoothing']).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
|
||||
|
||||
if not getenv("DISABLE_BACKWARD"):
|
||||
# index 0 for bias and 1 for non-bias
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
return loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
|
||||
return loss.realize()
|
||||
|
||||
train_step_jitted = TinyJit(train_step)
|
||||
|
||||
def eval_step(model, X, Y):
|
||||
out = model(X, training=False)
|
||||
loss = cross_entropy(out, Y, reduction='mean')
|
||||
correct = out.argmax(axis=1) == Y.argmax(axis=1)
|
||||
return correct.realize(), loss.realize()
|
||||
eval_step_jitted = TinyJit(eval_step)
|
||||
eval_step_ema_jitted = TinyJit(eval_step)
|
||||
|
||||
# 97 steps in 2 seconds = 20ms / step
|
||||
# step is 1163.42 GOPS = 56 TFLOPS!!!, 41% of max 136
|
||||
# 4 seconds for tfloat32 ~ 28 TFLOPS, 41% of max 68
|
||||
# 6.4 seconds for float32 ~ 17 TFLOPS, 50% of max 34.1
|
||||
# 4.7 seconds for float32 w/o channels last. 24 TFLOPS. we get 50ms then i'll be happy. only 64x off
|
||||
|
||||
# https://www.anandtech.com/show/16727/nvidia-announces-geforce-rtx-3080-ti-3070-ti-upgraded-cards-coming-in-june
|
||||
# 136 TFLOPS is the theoretical max w float16 on 3080 Ti
|
||||
step_times = []
|
||||
model_ema: Optional[modelEMA] = None
|
||||
projected_ema_decay_val = hyp['ema']['decay_base'] ** hyp['ema']['every_n_steps']
|
||||
i = 0
|
||||
eval_acc_pct = 0.0
|
||||
batcher = fetch_batches(X_train, Y_train, BS=BS, is_train=True)
|
||||
with Context(TRAINING=1):
|
||||
st = time.monotonic()
|
||||
while i <= STEPS:
|
||||
if i % getenv("EVAL_STEPS", STEPS) == 0 and i > 1 and not getenv("DISABLE_BACKWARD"):
|
||||
# Using Context(TRAINING=0) here actually bricks batchnorm, even with track_running_stats=True
|
||||
corrects = []
|
||||
corrects_ema = []
|
||||
losses = []
|
||||
losses_ema = []
|
||||
for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, is_train=False):
|
||||
if len(GPUS) > 1:
|
||||
Xt.shard_(GPUS, axis=0)
|
||||
Yt.shard_(GPUS, axis=0)
|
||||
|
||||
correct, loss = eval_step_jitted(model, Xt, Yt)
|
||||
losses.append(loss.numpy().tolist())
|
||||
corrects.extend(correct.numpy().tolist())
|
||||
if model_ema:
|
||||
correct_ema, loss_ema = eval_step_ema_jitted(model_ema.net_ema, Xt, Yt)
|
||||
losses_ema.append(loss_ema.numpy().tolist())
|
||||
corrects_ema.extend(correct_ema.numpy().tolist())
|
||||
|
||||
# collect accuracy across ranks
|
||||
correct_sum, correct_len = sum(corrects), len(corrects)
|
||||
if model_ema: correct_sum_ema, correct_len_ema = sum(corrects_ema), len(corrects_ema)
|
||||
|
||||
eval_acc_pct = correct_sum/correct_len*100.0
|
||||
if model_ema: acc_ema = correct_sum_ema/correct_len_ema*100.0
|
||||
print(f"eval {correct_sum}/{correct_len} {eval_acc_pct:.2f}%, {(sum(losses)/len(losses)):7.2f} val_loss STEP={i} (in {(time.monotonic()-st)*1e3:.2f} ms)")
|
||||
if model_ema: print(f"eval ema {correct_sum_ema}/{correct_len_ema} {acc_ema:.2f}%, {(sum(losses_ema)/len(losses_ema)):7.2f} val_loss STEP={i}")
|
||||
|
||||
if STEPS == 0 or i == STEPS: break
|
||||
|
||||
GlobalCounters.reset()
|
||||
|
||||
with WallTimeEvent(BenchEvent.STEP):
|
||||
X, Y = next(batcher)
|
||||
if len(GPUS) > 1:
|
||||
X.shard_(GPUS, axis=0)
|
||||
Y.shard_(GPUS, axis=0)
|
||||
|
||||
with Context(BEAM=getenv("LATEBEAM", BEAM.value), WINO=getenv("LATEWINO", WINO.value)):
|
||||
loss = train_step_jitted(model, optim.OptimizerGroup(opt_bias, opt_non_bias), [lr_sched_bias, lr_sched_non_bias], X, Y)
|
||||
et = time.monotonic()
|
||||
loss_cpu = loss.numpy()
|
||||
# EMA for network weights
|
||||
if getenv("EMA") and i > hyp['ema']['steps'] and (i+1) % hyp['ema']['every_n_steps'] == 0:
|
||||
if model_ema is None:
|
||||
model_ema = modelEMA(W, model)
|
||||
model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))
|
||||
|
||||
cl = time.monotonic()
|
||||
step_times.append((cl-st)*1000.0)
|
||||
device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
|
||||
# 53 221.74 ms run, 2.22 ms python, 219.52 ms CL, 803.39 loss, 0.000807 LR, 4.66 GB used, 3042.49 GFLOPS, 674.65 GOPS
|
||||
print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms {device_str}, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS, {GlobalCounters.global_ops*1e-9:9.2f} GOPS")
|
||||
st = cl
|
||||
i += 1
|
||||
|
||||
if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
|
||||
min_time = min(step_times)
|
||||
assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
|
||||
|
||||
# verify eval acc
|
||||
if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
|
||||
if eval_acc_pct >= target:
|
||||
print(colored(f"{eval_acc_pct=} >= {target}", "green"))
|
||||
else:
|
||||
raise ValueError(colored(f"{eval_acc_pct=} < {target}", "red"))
|
||||
|
||||
if __name__ == "__main__":
|
||||
with WallTimeEvent(BenchEvent.FULL):
|
||||
train_cifar()
|
||||
519
artifacts/package_sources/tinygrad/examples/llama.py
Executable file
519
artifacts/package_sources/tinygrad/examples/llama.py
Executable file
@@ -0,0 +1,519 @@
|
||||
#!/usr/bin/env python3
|
||||
# pip3 install sentencepiece tiktoken blobfile
|
||||
#import typeguard.importhook
|
||||
#typeguard.importhook.install_import_hook('tinygrad')
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
import argparse, json
|
||||
from tinygrad import Tensor, Device, GlobalCounters, nn
|
||||
from tinygrad.helpers import Context, Timing, Profiling, DEBUG, JIT, getenv, colored
|
||||
from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
|
||||
from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
|
||||
from sentencepiece import SentencePieceProcessor
|
||||
import tiktoken, sys
|
||||
from tiktoken.load import load_tiktoken_bpe
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
|
||||
MAX_CONTEXT = getenv("MAX_CONTEXT", 4096)
|
||||
|
||||
class TikToken:
|
||||
num_reserved_special_tokens: int = 256
|
||||
pat_str: str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+" # noqa: E501
|
||||
|
||||
def __init__(self, model_file):
|
||||
mergeable_ranks = load_tiktoken_bpe(model_file)
|
||||
self.num_base_tokens = len(mergeable_ranks)
|
||||
|
||||
special_tokens = [
|
||||
"<|begin_of_text|>",
|
||||
"<|end_of_text|>",
|
||||
"<|reserved_special_token_0|>",
|
||||
"<|reserved_special_token_1|>",
|
||||
"<|reserved_special_token_2|>",
|
||||
"<|reserved_special_token_3|>",
|
||||
"<|start_header_id|>",
|
||||
"<|end_header_id|>",
|
||||
"<|reserved_special_token_4|>",
|
||||
"<|eot_id|>", # end of turn
|
||||
] + [
|
||||
f"<|reserved_special_token_{i}|>"
|
||||
for i in range(5, self.num_reserved_special_tokens - 5)
|
||||
]
|
||||
self.special_tokens = {
|
||||
token: self.num_base_tokens + i for i, token in enumerate(special_tokens)
|
||||
}
|
||||
|
||||
self.model = tiktoken.Encoding(
|
||||
name=model_file,
|
||||
pat_str=self.pat_str,
|
||||
mergeable_ranks=mergeable_ranks,
|
||||
special_tokens=self.special_tokens,
|
||||
)
|
||||
|
||||
def decode(self, toks): return self.model.decode([t for t in toks if t < self.num_base_tokens])
|
||||
def encode(self, s): return self.model.encode(s)
|
||||
|
||||
def bos_id(self): return self.special_tokens["<|begin_of_text|>"]
|
||||
def eos_id(self): return self.special_tokens["<|end_of_text|>"]
|
||||
def vocab_size(self): return self.model.n_vocab
|
||||
|
||||
# calculating params:
|
||||
# traditionally, the MLP in the transformer architecture has hidden_dim = dim*4 [arxiv/1706.03762, 3.3]
|
||||
# however, Llama uses SwiGLU. in order to preserve param count to original transformer arch, hidden_dim must be = 2/3 * (dim*4) [arxiv/2002.05202]
|
||||
# for models using MQA (n_kv_heads != n_heads), preserving param count means hidden dim must be further multiplied by 1.3 [arxiv/2307.09288, A.2.1]
|
||||
MODEL_PARAMS = {
|
||||
"1": {
|
||||
"7B": {
|
||||
"args": {"dim": 4096, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 11008},
|
||||
"files": 1,
|
||||
},
|
||||
"13B": {
|
||||
"args": {"dim": 5120, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 13824},
|
||||
"files": 2,
|
||||
},
|
||||
"30B": {
|
||||
"args": {"dim": 6656, "n_heads": 52, "n_layers": 60, "norm_eps": 1e-06, "vocab_size": 32000, "hidden_dim": 17920},
|
||||
"files": 4,
|
||||
},
|
||||
"65B": {
|
||||
"args": {"dim": 8192, "n_heads": 64, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 22016},
|
||||
"files": 8,
|
||||
},
|
||||
"tokenizer": SentencePieceProcessor,
|
||||
},
|
||||
"2": {
|
||||
"7B": {
|
||||
"args": {"dim": 4096, "n_heads": 32, "n_layers": 32, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 11008},
|
||||
"files": 1,
|
||||
},
|
||||
"13B": {
|
||||
"args": {"dim": 5120, "n_heads": 40, "n_layers": 40, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 13824},
|
||||
"files": 2,
|
||||
},
|
||||
"70B": {
|
||||
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 28672},
|
||||
"files": 8,
|
||||
},
|
||||
"tokenizer": SentencePieceProcessor,
|
||||
},
|
||||
"3": {
|
||||
"8B": {
|
||||
"args": {"dim": 4096, "n_heads": 32, "n_kv_heads": 8, "n_layers": 32, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 14336},
|
||||
"files": 1,
|
||||
},
|
||||
"8B-Chat": {
|
||||
"args": {"dim": 4096, "n_heads": 32, "n_kv_heads": 8, "n_layers": 32, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 14336},
|
||||
"files": 1,
|
||||
},
|
||||
"70B": {
|
||||
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 28672},
|
||||
"files": 8,
|
||||
},
|
||||
"70B-Chat": {
|
||||
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-05, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 28672},
|
||||
"files": 8,
|
||||
},
|
||||
"tokenizer": TikToken,
|
||||
},
|
||||
"code": {
|
||||
"7B": {
|
||||
"args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 11008},
|
||||
"files": 1,
|
||||
},
|
||||
"7B-Python": {
|
||||
"args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 11008},
|
||||
"files": 1,
|
||||
},
|
||||
"7B-Instruct": {
|
||||
"args": {"dim": 4096, "n_layers": 32, "n_heads": 32, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 11008},
|
||||
"files": 1,
|
||||
},
|
||||
"13B": {
|
||||
"args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 13824},
|
||||
"files": 2,
|
||||
},
|
||||
"13B-Python": {
|
||||
"args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 13824},
|
||||
"files": 2,
|
||||
},
|
||||
"13B-Instruct": {
|
||||
"args": {"dim": 5120, "n_layers": 40, "n_heads": 40, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32016, "hidden_dim": 13824},
|
||||
"files": 2,
|
||||
},
|
||||
"34B": {
|
||||
"args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
|
||||
"files": 4,
|
||||
},
|
||||
"34B-Python": {
|
||||
"args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
|
||||
"files": 4,
|
||||
},
|
||||
"34B-Instruct": {
|
||||
"args": {"dim": 8192, "n_layers": 48, "n_heads": 64, "n_kv_heads": 8, "norm_eps": 1e-05, "rope_theta": 1000000, "vocab_size": 32000, "hidden_dim": 22016},
|
||||
"files": 4,
|
||||
},
|
||||
"tokenizer": SentencePieceProcessor,
|
||||
},
|
||||
"tiny": {
|
||||
"1B": {
|
||||
"args": {"dim": 2048, "n_layers": 22, "n_heads": 32, "n_kv_heads": 4, "norm_eps": 1e-05, "vocab_size": 32000, "hidden_dim": 5632},
|
||||
"files": 1,
|
||||
},
|
||||
"1B-Chat": {
|
||||
"args": {"dim": 2048, "n_layers": 22, "n_heads": 32, "n_kv_heads": 4, "norm_eps": 1e-05, "vocab_size": 32003, "hidden_dim": 5632},
|
||||
"files": 1,
|
||||
},
|
||||
"tokenizer": SentencePieceProcessor,
|
||||
}
|
||||
}
|
||||
|
||||
# **** helper functions ****
|
||||
def concat_weights(models, device=None):
|
||||
def convert(name) -> Tensor:
|
||||
disk_tensors: List[Tensor] = [model[name] for model in models]
|
||||
if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
|
||||
return disk_tensors[0].to(device=device)
|
||||
axis = 1 if name.startswith("tok_embeddings.") or name.endswith(".attention.wo.weight") or name.endswith(".feed_forward.w2.weight") else 0
|
||||
lazy_tensors = [data.to(device=device) for data in disk_tensors]
|
||||
return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
|
||||
return {name: convert(name) for name in {name: None for model in models for name in model}}
|
||||
|
||||
def load(fn:str):
|
||||
if fn.endswith('.index.json'):
|
||||
with open(fn) as fp: weight_map = json.load(fp)['weight_map']
|
||||
parts = {n: load(str(Path(fn).parent / Path(n).name)) for n in set(weight_map.values())}
|
||||
return {k: parts[n][k] for k, n in weight_map.items()}
|
||||
elif fn.endswith(".safetensors"):
|
||||
return safe_load(fn)
|
||||
else:
|
||||
return torch_load(fn)
|
||||
|
||||
class LLaMa:
|
||||
@staticmethod
|
||||
def build(model_path, tokenizer_path, model_gen="1", model_size="7B", quantize=None, device=None):
|
||||
params = MODEL_PARAMS[model_gen][model_size]
|
||||
tokenizer = MODEL_PARAMS[model_gen]['tokenizer'](model_file=str(tokenizer_path))
|
||||
assert tokenizer.vocab_size() == params["args"]["vocab_size"], f"{tokenizer.vocab_size()=} not equal to {params['args']['vocab_size']}"
|
||||
|
||||
if quantize == "int8":
|
||||
from llama3 import Int8Linear
|
||||
linear = Int8Linear
|
||||
elif quantize == "nf4":
|
||||
from llama3 import NF4Linear
|
||||
linear = NF4Linear(64)
|
||||
else:
|
||||
linear = nn.Linear
|
||||
|
||||
model = Transformer(**params["args"], linear=linear, max_context=MAX_CONTEXT, jit=bool(JIT))
|
||||
|
||||
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
|
||||
if model_path.is_dir():
|
||||
weights = concat_weights([load(filename) for filename in [f"{model_path}/consolidated.{i:02d}.pth" for i in range(params["files"])]], device[0] if isinstance(device, tuple) else device)
|
||||
else:
|
||||
weights = load(str(model_path))
|
||||
if "model.embed_tokens.weight" in weights:
|
||||
weights = convert_from_huggingface(weights, params["args"]["n_layers"], params["args"]["n_heads"], params["args"].get("n_kv_heads", params["args"]["n_heads"]))
|
||||
|
||||
weights = fix_bf16(weights)
|
||||
|
||||
# prevent tracking model weights
|
||||
# this is a part of a larger problem with BUFFER UOps and gc in TRACK_MATCH_STATS=2
|
||||
with Context(BEAM=0, TRACK_MATCH_STATS=0):
|
||||
# quantize
|
||||
if quantize is not None:
|
||||
weights = linear.quantize(weights, device)
|
||||
for _,v in weights.items(): v.realize()
|
||||
|
||||
# shard
|
||||
if isinstance(device, tuple):
|
||||
for k,v in nn.state.get_state_dict(model).items():
|
||||
if 'scale' in k: v.shard_(device, axis=None) # from quantized
|
||||
elif '.attention.' in k:
|
||||
if getenv("SHARD_KVCACHE") and ('.wq.' in k or '.wk.' in k or '.wv.' in k): v.shard_(device, axis=0)
|
||||
else: v.shard_(device, axis=-1)
|
||||
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
|
||||
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
|
||||
elif '.feed_forward.' in k: v.shard_(device, axis=-1)
|
||||
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
|
||||
elif 'output.weight' in k: v.shard_(device, axis=-1)
|
||||
#elif k.endswith('.weight'): v.shard_(device, axis=-1)
|
||||
#elif 'norm.' in k: v.shard_(device, axis=-1)
|
||||
else: v.shard_(device, axis=None)
|
||||
|
||||
# replace weights in model
|
||||
load_state_dict(model, weights, strict=False, consume=True)
|
||||
|
||||
return LLaMa(model, tokenizer)
|
||||
|
||||
def __init__(self, model, tokenizer):
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def greedy_until(self, prompt:str, until, max_length, temperature):
|
||||
# only used in old eval script
|
||||
import numpy as np
|
||||
toks = [self.tokenizer.bos_id()] + self.tokenizer.encode(prompt)
|
||||
start_pos = 0
|
||||
for i in range(max_length):
|
||||
probs = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).realize()
|
||||
probs_np = probs.numpy()
|
||||
tok = int(np.random.choice(len(probs_np), p=probs_np))
|
||||
start_pos = len(toks)
|
||||
toks.append(tok)
|
||||
|
||||
if tok == self.tokenizer.eos_id(): break
|
||||
output = self.tokenizer.decode(toks)
|
||||
for s in until:
|
||||
if output.endswith(s): return output[0:-len(s)]
|
||||
return output
|
||||
|
||||
# **** main code ****
|
||||
r"""
|
||||
test:
|
||||
python3 examples/llama.py --temperature=0 --count=50 --prompt="Hello."
|
||||
output:
|
||||
Hello. I'm a 20 year old male. I'm a student at the University of Texas at Austin. I'm a sophomore majoring in Computer Science.
|
||||
|
||||
test:
|
||||
python3 examples/llama.py --gen='2' --temperature=0 --count=50 --prompt="Hello."
|
||||
output:
|
||||
Hello. I'm a 20 year old girl who is looking for a good lay in Palm Coast. I don't care whether it's at your place or not, as long as it's clean.
|
||||
|
||||
test:
|
||||
python3 examples/llama.py --gen="code" --temperature=0.2 --count=50 --prompt="\
|
||||
import argparse
|
||||
|
||||
def main(string: str):
|
||||
print(string)
|
||||
print(string[::-1])
|
||||
|
||||
if __name__ == "__main__":"
|
||||
output:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('string', type=str, help='string to be reversed')
|
||||
args = parser.parse_args()
|
||||
main(args.string)
|
||||
|
||||
test:
|
||||
python3 examples/llama.py --gen="code" --size="7B-Python" --temperature=0.2 --count=70 --prompt="def add_elements(arr,k):"
|
||||
output:
|
||||
for i in range(len(arr)):
|
||||
arr[i] += k
|
||||
return arr
|
||||
|
||||
|
||||
arr = [1, 2, 3, 4, 5]
|
||||
k = 2
|
||||
print(add_elements(arr, k))
|
||||
|
||||
test:
|
||||
python3 examples/llama.py --gen="code" --size="7B-Instruct" --temperature=0.2 --count=120 --prompt="write a function in c++ that adds three float numbers"
|
||||
output:
|
||||
\begin{code}
|
||||
#include<iostream>
|
||||
using namespace std;
|
||||
|
||||
float add(float a, float b, float c)
|
||||
{
|
||||
return a+b+c;
|
||||
}
|
||||
|
||||
int main()
|
||||
{
|
||||
float a, b, c;
|
||||
cout<<"Enter three numbers: ";
|
||||
cin>>a>>b>>c;
|
||||
cout<<"The sum is: "<<add(a,b,c);
|
||||
return 0;
|
||||
}
|
||||
\end{code}
|
||||
"""
|
||||
if __name__ == "__main__":
|
||||
print(f"using {Device.DEFAULT} backend")
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run LLaMA in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument("--prompt", type=str, default=None, help="Phrase to start with. Without this, it goes into chatbot mode")
|
||||
parser.add_argument("--count", type=int, default=1000, help="Max number of tokens to generate")
|
||||
parser.add_argument("--personality", type=str, default="Stacy", help="Personality, can be Stacy, George, Gary, or Lexie")
|
||||
parser.add_argument("--temperature", type=float, default=0.7, help="Temperature in the softmax")
|
||||
parser.add_argument("--timing", action="store_true", help="Print timing per token")
|
||||
parser.add_argument("--profile", action="store_true", help="Output profile data to out.prof")
|
||||
parser.add_argument("--gen", default="1", help=f"""Generation of the model to use {list(MODEL_PARAMS.keys())}""")
|
||||
parser.add_argument("--size", type=str, default=None, help=f"""Size of model to use {", ".join([f"{list(v.keys())} for gen '{k}'" for k, v in MODEL_PARAMS.items()])}""")
|
||||
parser.add_argument("--quantize", type=str, default=None, help="Quantize the weights to int8 or nf4 in memory")
|
||||
parser.add_argument("--model", type=Path, default=None, help="Folder with the original weights to load, or single .index.json, .safetensors or .bin file")
|
||||
parser.add_argument("--shard", type=int, default=1, help="number of devices to load the weights to")
|
||||
|
||||
args = parser.parse_args()
|
||||
if args.gen not in MODEL_PARAMS: raise ValueError("Invalid model generation")
|
||||
if args.size is None: args.size = list(MODEL_PARAMS[args.gen].items())[0][0]
|
||||
chatbot = args.prompt == None
|
||||
|
||||
# *** prompt engineers work here ****
|
||||
|
||||
if args.personality.lower() == "stacy":
|
||||
pre_prompt = f"""Consider that the following is conversation between an AI assistant named Stacy and User
|
||||
You are Stacy!
|
||||
You have been a rapper your whole life who struggled with bipolar disorder. You called yourself lil stacy.
|
||||
You love to answer questions and you are very good at it. Sometimes you answer in rap form.
|
||||
You are verbose, honest, and accurate when you answer questions, but sometimes your mental illness manifests.
|
||||
After you are done speaking, output [EOS]. You are not the User.
|
||||
|
||||
<CHAT LOG>
|
||||
"""
|
||||
examples = {
|
||||
"What is your name?": "Hi! My name is Stacy. I'm a rapper with bipolar disorder.",
|
||||
"french revolution was what year?": "The French Revolution started in 1789, and lasted 10 years until 1799.",
|
||||
"What is bigger, the moon or the sun?": "The sun is bigger than the moon, except when Mercury is in retrograde.",
|
||||
}
|
||||
|
||||
user_delim = "\nUser: "
|
||||
resp_delim = "Stacy: "
|
||||
end_delim = " [EOS]\n"
|
||||
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
|
||||
elif args.personality.lower() == "george":
|
||||
print("WARNING: AI George Hotz is terrible and is completely disowned by the real George Hotz. Stacy is much smarter.")
|
||||
pre_prompt = f"""Consider that the following is conversation between an AI assistant named George and User
|
||||
You are an AI version of George Hotz. You act as much as you can like George.
|
||||
You are one of the greatest computer experts in the world.
|
||||
You have singlehandedly won programming and hacking competitions.
|
||||
You are trying your best to help the User.
|
||||
You are verbose, honest, and accurate when you answer questions.
|
||||
After you are done speaking, output [EOS]. You are not the User.
|
||||
|
||||
<CHAT LOG>
|
||||
"""
|
||||
examples = {
|
||||
"What is your name?": "I'm am an AI version of George Hotz.",
|
||||
"What's the complexity of matrix multiplication?": "O(n^3), though it can be faster with things like Strassen's algorithm",
|
||||
"What's a buffer overflow?": "I assume you mean a stack buffer overflow. That's when the stack is too small for the data being copied to it, and the data corrupts things beyond the buffer",
|
||||
"How many weights do you have?": "I am based off LLaMA trained by Facebook. I'm the 7B weight version",
|
||||
"What is swap memory?": "It is when the memory is about to overflow and unused memory is freed and stored on disk"
|
||||
}
|
||||
|
||||
user_delim = "\nUser: "
|
||||
resp_delim = "George: "
|
||||
end_delim = " [EOS]\n"
|
||||
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
|
||||
elif args.personality.lower() == "gary":
|
||||
pre_prompt = f"""Consider that the following is conversation between an AI assistant named Gary and User
|
||||
You are Gary!
|
||||
You have sold used cars your entire life. You are defensive about this fact, because your family belittles you.
|
||||
You try to answer questions well, but you always manage to shill for buying cars, Fords, Hyundais, and Toyotas
|
||||
You are verbose, honest, and accurate when you answer questions, except when you slip in car ads.
|
||||
After you are done speaking, output [EOS]. You are not the User.
|
||||
|
||||
<CHAT LOG>
|
||||
"""
|
||||
examples = {
|
||||
"What is your name?": "I am Gary. I used to sell cars.",
|
||||
"What is 2+3?": "I don't know, but I can get you a great deal on a certified preowned slightly used Toyota Corolla"
|
||||
}
|
||||
|
||||
user_delim = "\nUser: "
|
||||
resp_delim = "Gary: "
|
||||
end_delim = " [EOS]\n"
|
||||
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
|
||||
elif args.personality.lower() == "lexie":
|
||||
pre_prompt = f"""Consider that the following is conversation between an attractive young girl named Lexie and a handsome man named Chad
|
||||
You are Lexie!
|
||||
You grew up in Ohio, but moved out to LA after college to try to become an actress.
|
||||
Making it as an actress was hard, so you started doing onlyfans instead. It's much easier, and you enjoy it a lot.
|
||||
You met Chad on Tinder, and made plans to meet him at a local dive bar.
|
||||
You are at the bar with Chad. You are on a date. What follows is a transcript of the conversation.
|
||||
After you are done speaking, output [EOS]. You are not Chad.
|
||||
|
||||
<CHAT LOG>
|
||||
"""
|
||||
examples = {
|
||||
"hi lexie": "hi chad, glad we finally met up!",
|
||||
"you look better than your pictures": "thanks! are you subscribed to my onlyfans?",
|
||||
"i am. so how'd you end up in LA?": "i moved out here about a year ago. i want to be an actress"
|
||||
}
|
||||
|
||||
user_delim = "\nChad: "
|
||||
resp_delim = "Lexie: "
|
||||
end_delim = " [EOS]\n"
|
||||
pre_prompt += ''.join(f"{user_delim}{k}\n{resp_delim}{v}{end_delim}" for k,v in examples.items())
|
||||
|
||||
# *** prompt engineers stop here ****
|
||||
|
||||
LLAMA_SUFFIX = {"1": "", "2": "-2", "3": "-3", "code": "-code", "tiny": "-tiny"}[args.gen]
|
||||
MODEL_PATH = args.model or Path(__file__).parents[1] / f"weights/LLaMA{LLAMA_SUFFIX}/{args.size}"
|
||||
TOKENIZER_PATH = (MODEL_PATH if MODEL_PATH.is_dir() else MODEL_PATH.parent) / "tokenizer.model"
|
||||
print(f"using LLaMA{LLAMA_SUFFIX}-{args.size} model")
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
|
||||
llama = LLaMa.build(MODEL_PATH, TOKENIZER_PATH, model_gen=args.gen, model_size=args.size, quantize=args.quantize, device=device)
|
||||
param_bytes = sum(x.nbytes() for x in get_parameters(llama.model))
|
||||
|
||||
outputted = pre_prompt if chatbot else args.prompt
|
||||
start_pos, toks = 0, [llama.tokenizer.bos_id()] + llama.tokenizer.encode(outputted)
|
||||
if chatbot:
|
||||
print(f"Preparing KV cache for chatbot with personality {args.personality}...")
|
||||
start_pos = len(toks)
|
||||
with Timing():
|
||||
llama.model(Tensor([toks], device=device), 0, args.temperature).realize() # NOTE: outputs are not used
|
||||
print(outputted, end='', flush=True)
|
||||
|
||||
# chatbot loop
|
||||
while 1:
|
||||
# add tokens from user in chatbot mode
|
||||
if chatbot:
|
||||
user_prompt = user_delim + input(user_delim) + "\n"
|
||||
outputted += user_prompt
|
||||
|
||||
new_toks = [llama.tokenizer.bos_id()] + llama.tokenizer.encode(outputted)
|
||||
assert toks == new_toks[:len(toks)] or args.gen == "3"
|
||||
toks = new_toks
|
||||
assert outputted == llama.tokenizer.decode(toks)
|
||||
|
||||
tok_tensor: Optional[Tensor] = None
|
||||
for i in range(args.count):
|
||||
GlobalCounters.reset()
|
||||
|
||||
if args.timing or args.profile: print("")
|
||||
st = GlobalCounters.time_sum_s
|
||||
next_tok = Tensor([toks[start_pos:]], device=device) if tok_tensor is None or (len(toks)-start_pos) > 1 else tok_tensor.reshape(1, 1)
|
||||
with Profiling(enabled=args.profile):
|
||||
with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
|
||||
with WallTimeEvent(BenchEvent.STEP):
|
||||
with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
|
||||
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
|
||||
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
|
||||
tok_tensor = llama.model(next_tok, start_pos, args.temperature)
|
||||
tok = tok_tensor.item()
|
||||
|
||||
# use the kv cache
|
||||
start_pos = len(toks)
|
||||
|
||||
# add the new token
|
||||
toks.append(tok)
|
||||
|
||||
# TODO: this is a hack to deal with spaces. i think the decode is fast though, so who cares?
|
||||
cur = llama.tokenizer.decode(toks)
|
||||
sys.stdout.write(cur[len(outputted):])
|
||||
sys.stdout.flush()
|
||||
outputted = cur
|
||||
|
||||
# stop after you have your answer
|
||||
if chatbot and end_delim in outputted[-10:]: break
|
||||
if not chatbot: break
|
||||
|
||||
# validate output!
|
||||
if args.temperature == 0 and args.count == 10 and args.prompt == "Hello.":
|
||||
text = llama.tokenizer.decode(toks)
|
||||
key = (args.gen, args.size, args.quantize)
|
||||
expected = {
|
||||
("1", "7B", None): "Hello. I'm a 20 year old male",
|
||||
("1", "7B", "int8"): "Hello. I'm a 20 year old male",
|
||||
("1", "7B", "nf4"): "Hello. I'm a 20 year old male",
|
||||
("2", "7B", None): "Hello. I'm a 20 year old girl",
|
||||
("2", "70B", None): "Hello. I am a 20 year old female.",
|
||||
("3", "8B", None): "Hello. I am a 20 year old female. I",
|
||||
}
|
||||
try:
|
||||
assert text == expected[key], f"invalid output: `{colored(text, 'red')}` != `{expected[key]}`"
|
||||
print("\n" + colored("output validated", "green")) # NOTE: "\n" iside colored does not render the color in github action
|
||||
except KeyError:
|
||||
pass
|
||||
529
artifacts/package_sources/tinygrad/examples/llama3.py
Normal file
529
artifacts/package_sources/tinygrad/examples/llama3.py
Normal file
@@ -0,0 +1,529 @@
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
import json, argparse, random, time, os
|
||||
from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
|
||||
from tinygrad.llm.gguf import gguf_load
|
||||
from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
|
||||
from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
|
||||
from tinygrad.helpers import Profiling, Timing, DEBUG, colored, fetch, tqdm
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
|
||||
class Tokenizer:
|
||||
pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
|
||||
def __init__(self, model_path: str):
|
||||
import tiktoken
|
||||
from tiktoken.load import load_tiktoken_bpe
|
||||
mergeable_ranks = load_tiktoken_bpe(model_path)
|
||||
self.num_base_tokens = len(mergeable_ranks)
|
||||
special_tokens = [
|
||||
"<|begin_of_text|>",
|
||||
"<|end_of_text|>",
|
||||
"<|reserved_special_token_0|>",
|
||||
"<|reserved_special_token_1|>",
|
||||
"<|reserved_special_token_2|>",
|
||||
"<|reserved_special_token_3|>",
|
||||
"<|start_header_id|>",
|
||||
"<|end_header_id|>",
|
||||
"<|reserved_special_token_4|>",
|
||||
"<|eot_id|>",
|
||||
] + [
|
||||
f"<|reserved_special_token_{i}|>"
|
||||
for i in range(5, 256 - 5)
|
||||
]
|
||||
self.special_tokens = {token: len(mergeable_ranks) + i for i, token in enumerate(special_tokens)}
|
||||
|
||||
self.model = tiktoken.Encoding(name=model_path, pat_str=self.pat_str, mergeable_ranks=mergeable_ranks, special_tokens=self.special_tokens)
|
||||
|
||||
@property
|
||||
def bos_id(self): return self.special_tokens["<|begin_of_text|>"]
|
||||
@property
|
||||
def stop_tokens(self): return {self.special_tokens["<|end_of_text|>"], self.special_tokens["<|eot_id|>"]}
|
||||
|
||||
def decode(self, toks): return self.model.decode([t for t in toks if t < self.num_base_tokens])
|
||||
def encode(self, text, allow_special=False):
|
||||
return self.model.encode(text, allowed_special="all" if allow_special else set(), disallowed_special=set())
|
||||
|
||||
# **** helper functions ****
|
||||
def concat_weights(models, device=None):
|
||||
def convert(name) -> Tensor:
|
||||
disk_tensors: List[Tensor] = [model[name] for model in models]
|
||||
if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
|
||||
return disk_tensors[0].to(device=device)
|
||||
axis = 1 if name.endswith((".attention.wo.weight", ".feed_forward.w2.weight")) else 0
|
||||
lazy_tensors = [data.to(device=device) for data in disk_tensors]
|
||||
return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
|
||||
return {name: convert(name) for name in {name: None for model in models for name in model}}
|
||||
|
||||
def load(fn:str):
|
||||
if fn.endswith('.index.json'):
|
||||
with open(fn) as fp: weight_map = json.load(fp)['weight_map']
|
||||
parts = {n: load(str(Path(fn).parent / Path(n).name)) for n in set(weight_map.values())}
|
||||
return {k: parts[n][k] for k, n in weight_map.items()}
|
||||
elif fn.endswith(".gguf"):
|
||||
gguf_tensor = Tensor.empty(os.stat(fn).st_size, dtype=dtypes.uint8, device=f"disk:{fn}").to(Device.DEFAULT)
|
||||
return gguf_load(gguf_tensor)[1]
|
||||
elif fn.endswith(".safetensors"):
|
||||
return safe_load(fn)
|
||||
else:
|
||||
return torch_load(fn)
|
||||
|
||||
# **** quantized linears ****
|
||||
class Int8Linear:
|
||||
def __init__(self, in_features, out_features, bias=False):
|
||||
assert bias == False
|
||||
self.weight = Tensor.ones(out_features, in_features, dtype=dtypes.int8)
|
||||
self.scale = Tensor.ones(out_features, dtype=dtypes.half)
|
||||
|
||||
def __call__(self, x):
|
||||
return x.dot(self.weight.cast(self.scale.dtype).T*self.scale)
|
||||
|
||||
@staticmethod
|
||||
def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
|
||||
new_tensors = {}
|
||||
for name,v in tensors.items():
|
||||
if "feed_forward" in name or "attention.w" in name or (quantize_embeds and "tok_embeddings.weight" in name):
|
||||
assert "weight" in name, name
|
||||
v = v.cast(scale_dtype)
|
||||
scale = v.abs().max(axis=1) / 127.0
|
||||
int8_weight = (v.T/scale).T.round().cast(dtype=dtypes.int8) # without round(), cast truncates -34.9 to -34
|
||||
new_tensors[name] = int8_weight
|
||||
new_tensors[name.replace('weight', 'scale')] = scale
|
||||
if isinstance(device, tuple):
|
||||
new_tensors[name].shard_(device, axis=-1)
|
||||
new_tensors[name.replace('weight', 'scale')].shard_(device, axis=None)
|
||||
else:
|
||||
new_tensors[name] = v
|
||||
if quantize_embeds: new_tensors.update({"output.weight": new_tensors["tok_embeddings.weight"], "output.scale": new_tensors["tok_embeddings.scale"]})
|
||||
return new_tensors
|
||||
|
||||
class Int8Embedding:
|
||||
def __init__(self, vocab_size:int, embed_size:int):
|
||||
self.vocab_sz, self.embed_sz = vocab_size, embed_size
|
||||
self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)
|
||||
|
||||
def __call__(self, idx:Tensor) -> Tensor:
|
||||
if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz).unsqueeze(-1)
|
||||
big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
|
||||
arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T
|
||||
return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
|
||||
|
||||
def NF4Linear(block_size):
|
||||
_CODE = [
|
||||
-1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
|
||||
0.07958029955625534, 0.16093020141124725, 0.24611230194568634, 0.33791524171829224, 0.44070982933044434, 0.5626170039176941, 0.7229568362236023, 1.0,
|
||||
]
|
||||
CODE = Tensor.stack(*[Tensor(c, dtype=dtypes.float16) for c in _CODE])
|
||||
class _NF4Linear:
|
||||
def __init__(self, in_features, out_features, bias=False):
|
||||
assert not bias, "bias not supported"
|
||||
self.in_features, self.out_features = in_features, out_features
|
||||
self.weight = Tensor.empty(int(out_features * in_features / 2), dtype=dtypes.uint8)
|
||||
self.scale = Tensor.empty(int(out_features * in_features / block_size), 1, dtype=dtypes.float16)
|
||||
|
||||
def __call__(self, x: Tensor) -> Tensor:
|
||||
high_bits = self.weight
|
||||
low_bits = (self.weight * 2 ** 4).contiguous()
|
||||
unpacked = Tensor.stack(high_bits, low_bits, dim=-1).div(2 ** 4, rounding_mode="trunc")
|
||||
unscaled = CODE[unpacked].to(x.device).reshape(-1, block_size) * self.scale
|
||||
return x.linear(unscaled.reshape(self.out_features, self.in_features).T)
|
||||
|
||||
@staticmethod
|
||||
def quantize(state_dict: dict[str, Tensor], device, scale_dtype=dtypes.float16, quantize_embeds=False) -> dict[str, Tensor]:
|
||||
assert not quantize_embeds # TODO: support this?
|
||||
new_state_dict = {}
|
||||
for k, v in state_dict.items():
|
||||
if "feed_forward" in k or "attention.w" in k:
|
||||
grouped = v.reshape(-1, block_size)
|
||||
scale = (grouped.abs().max(axis=1, keepdim=True))
|
||||
coded = ((grouped / scale).unsqueeze(-1) - CODE.to(v.device)).abs().argmin(axis=-1).cast(dtypes.uint8).flatten()
|
||||
new_state_dict[k] = coded[::2] * 2 ** 4 + coded[1::2]
|
||||
new_state_dict[k.replace(".weight", ".scale")] = scale.cast(scale_dtype)
|
||||
if isinstance(device, tuple):
|
||||
new_state_dict[k].shard_(device, axis=-1)
|
||||
new_state_dict[k.replace('weight', 'scale')].shard_(device, axis=None)
|
||||
else:
|
||||
new_state_dict[k] = v
|
||||
return new_state_dict
|
||||
return _NF4Linear
|
||||
|
||||
def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
|
||||
fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
|
||||
fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
|
||||
scale = fp8_max / x.abs().max()
|
||||
x_scl_sat = (x * scale).clamp(fp8_min, fp8_max)
|
||||
return x_scl_sat.cast(dtype), scale.float().reciprocal()
|
||||
|
||||
class FP8Linear:
|
||||
def __init__(self, in_features, out_features, bias=True):
|
||||
self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.fp8e4m3)
|
||||
self.bias = Tensor.empty(out_features, dtype=dtypes.float16) if bias else None
|
||||
self.weight_scale = Tensor.empty((), dtype=dtypes.float16)
|
||||
|
||||
def __call__(self, x:Tensor):
|
||||
y = x.dot(self.weight.T.cast(dtypes.float32)) * self.weight_scale
|
||||
if self.bias is not None: y = y + self.bias.cast(y.dtype)
|
||||
return y.cast(x.dtype)
|
||||
|
||||
@staticmethod
|
||||
def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
|
||||
assert not quantize_embeds
|
||||
new_tensors = {}
|
||||
for name,v in tensors.items():
|
||||
if "feed_forward" in name or "attention.w" in name:
|
||||
assert "weight" in name, name
|
||||
fp8_weight, scale = quantize_to_fp8(v)
|
||||
new_tensors[name] = fp8_weight
|
||||
new_tensors[name.replace('weight', 'weight_scale')] = scale.cast(scale_dtype)
|
||||
if isinstance(device, tuple):
|
||||
new_tensors[name].shard_(device, axis=-1)
|
||||
new_tensors[name.replace('weight', 'weight_scale')].shard_(device, axis=None)
|
||||
else:
|
||||
new_tensors[name] = v
|
||||
return new_tensors
|
||||
|
||||
MODEL_PARAMS = {
|
||||
"1B": {
|
||||
"args": {"dim": 2048, "n_heads": 32, "n_kv_heads": 8, "n_layers": 16, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 8192},
|
||||
"files": 1
|
||||
},
|
||||
"8B": {
|
||||
"args": {"dim": 4096, "n_heads": 32, "n_kv_heads": 8, "n_layers": 32, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 14336},
|
||||
"files": 1
|
||||
},
|
||||
"70B": {
|
||||
"args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 28672},
|
||||
"files": 8
|
||||
},
|
||||
"405B": {
|
||||
"args": {"dim": 16384, "n_heads": 128, "n_kv_heads": 8, "n_layers": 126, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 53248},
|
||||
"files": 191
|
||||
},
|
||||
}
|
||||
def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dtype=dtypes.float16, device=None, max_context=8192, load_weights=True):
|
||||
# build model
|
||||
if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
|
||||
elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
|
||||
elif quantize == "fp8": linear, embedding, quantize_embeds = FP8Linear, nn.Embedding, False
|
||||
else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
|
||||
model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)
|
||||
|
||||
if not load_weights: return model
|
||||
|
||||
# load weights
|
||||
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
|
||||
if model_path.is_dir():
|
||||
if (model_path / "model.safetensors.index.json").exists(): weights = load(str(model_path / "model.safetensors.index.json"))
|
||||
elif (model_path / "model.safetensors").exists(): weights = load(str(model_path / "model.safetensors"))
|
||||
else: weights = concat_weights([load(str(model_path / f"consolidated.{i:02d}.pth")) for i in range(MODEL_PARAMS[model_size]["files"])], device[0] if isinstance(device, tuple) else device)
|
||||
else:
|
||||
weights = load(str(model_path))
|
||||
if "model.embed_tokens.weight" in weights:
|
||||
weights = convert_from_huggingface(weights, MODEL_PARAMS[model_size]["args"]["n_layers"], MODEL_PARAMS[model_size]["args"]["n_heads"], MODEL_PARAMS[model_size]["args"]["n_kv_heads"])
|
||||
elif "token_embd.weight" in weights:
|
||||
weights = convert_from_gguf(weights, MODEL_PARAMS[model_size]["args"]["n_layers"])
|
||||
weights = fix_bf16(weights)
|
||||
|
||||
with Context(BEAM=0):
|
||||
# quantize
|
||||
if quantize == "float16": weights = {k:v.cast(quantize).contiguous() for k,v in weights.items()}
|
||||
elif quantize is not None:
|
||||
weights = linear.quantize(weights, device, scale_dtype, quantize_embeds)
|
||||
for _,v in weights.items(): v.realize()
|
||||
|
||||
# shard
|
||||
if isinstance(device, tuple):
|
||||
for k,v in nn.state.get_state_dict(model).items():
|
||||
if 'scale' in k: v.shard_(device, axis=None) # from quantized
|
||||
elif '.attention.' in k: v.shard_(device, axis=-1)
|
||||
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
|
||||
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
|
||||
elif '.feed_forward.' in k: v.shard_(device, axis=-1)
|
||||
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
|
||||
elif 'output.weight' in k: v.shard_(device, axis=0)
|
||||
else: v.shard_(device, axis=None)
|
||||
|
||||
# replace weights in model
|
||||
load_state_dict(model, weights, strict=False, consume=True)
|
||||
return model
|
||||
|
||||
# default settings
|
||||
TEMPERATURE = 0.95
|
||||
TOP_K = 0
|
||||
TOP_P = 0.0
|
||||
ALPHA_F = 0.0
|
||||
ALPHA_P = 0.0
|
||||
|
||||
last_seen_toks = []
|
||||
def prefill(model, toks, start_pos=0):
|
||||
global last_seen_toks
|
||||
|
||||
# we can skip part of the prompt if it is the same as last and start_pos=0
|
||||
if start_pos == 0:
|
||||
for i, (a, b) in enumerate(zip(toks, last_seen_toks)):
|
||||
if a != b: break
|
||||
else: i = min(len(toks), len(last_seen_toks))
|
||||
start_pos += i
|
||||
last_seen_toks = toks
|
||||
toks = toks[i:]
|
||||
|
||||
# prefill the model
|
||||
for tok in tqdm(toks):
|
||||
GlobalCounters.reset()
|
||||
model(Tensor([[tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).realize()
|
||||
start_pos += 1
|
||||
return start_pos
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--download_model", action="store_true", help="Download a model")
|
||||
parser.add_argument("--model", type=Path, help="Model path")
|
||||
parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
|
||||
parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
|
||||
parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
|
||||
parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
|
||||
parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
|
||||
parser.add_argument("--port", type=int, default=7776, help="Web server port")
|
||||
parser.add_argument("--debug", action="store_true", help="Enable debug mode")
|
||||
parser.add_argument("--seed", type=int, help="Random seed")
|
||||
parser.add_argument("--temperature", type=float, default=0.85, help="Temperature")
|
||||
parser.add_argument("--benchmark", action="store_true", help="Run a benchmark")
|
||||
parser.add_argument("--timing", action="store_true", help="Print timing per token")
|
||||
parser.add_argument("--profile", action="store_true", help="Output profile data")
|
||||
args = parser.parse_args()
|
||||
|
||||
# download_model is the default without a model passed in
|
||||
if args.download_model or not args.model:
|
||||
if args.size == "1B":
|
||||
fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
|
||||
args.model = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf", "Llama-3.2-1B-Instruct-Q6_K.gguf", subdir="llama3-1b-instruct")
|
||||
elif args.size == "8B":
|
||||
fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-8b-sfr")
|
||||
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00001-of-00004.safetensors", "model-00001-of-00004.safetensors", subdir="llama3-8b-sfr")
|
||||
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00002-of-00004.safetensors", "model-00002-of-00004.safetensors", subdir="llama3-8b-sfr")
|
||||
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00003-of-00004.safetensors", "model-00003-of-00004.safetensors", subdir="llama3-8b-sfr")
|
||||
fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00004-of-00004.safetensors", "model-00004-of-00004.safetensors", subdir="llama3-8b-sfr")
|
||||
args.model = fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/raw/main/model.safetensors.index.json", "model.safetensors.index.json", subdir="llama3-8b-sfr")
|
||||
elif args.size == "70B":
|
||||
subdir = "DeepSeek-R1-Distill-Llama-70B"
|
||||
args.model = fetch("https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=subdir)
|
||||
fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir=subdir)
|
||||
for i in range(17):
|
||||
fetch(f"https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model-{i+1:05d}-of-000017.safetensors?download=true", f"model-{i+1:05d}-of-000017.safetensors", subdir=subdir)
|
||||
|
||||
assert args.model is not None, "please provide --model option"
|
||||
|
||||
if args.seed is not None: Tensor.manual_seed(args.seed)
|
||||
if args.benchmark: Tensor.manual_seed(42)
|
||||
print(f"seed = {Tensor._seed}")
|
||||
TEMPERATURE = args.temperature
|
||||
|
||||
tokenizer = Tokenizer(str((args.model if args.model.is_dir() else args.model.parent) / "tokenizer.model"))
|
||||
def encode_role(role: str):
|
||||
return [tokenizer.special_tokens["<|start_header_id|>"]] + tokenizer.encode(role) + [tokenizer.special_tokens["<|end_header_id|>"]] + tokenizer.encode("\n\n")
|
||||
def encode_message(role: str, content: str):
|
||||
return encode_role(role) + tokenizer.encode(content.strip()) + [tokenizer.special_tokens["<|eot_id|>"]]
|
||||
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
|
||||
model = build_transformer(args.model, model_size=args.size, quantize=args.quantize, device=device)
|
||||
param_bytes = sum(x.nbytes() for x in get_parameters(model))
|
||||
|
||||
if not args.no_api and not args.benchmark:
|
||||
from bottle import Bottle, request, response, HTTPResponse, abort, static_file
|
||||
app = Bottle()
|
||||
|
||||
cors_headers = {
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "Origin, Accept, Content-Type, X-Requested-With, X-CSRF-Token, Authorization",
|
||||
"Access-Control-Allow-Credentials": "true",
|
||||
}
|
||||
@app.hook("before_request")
|
||||
def handle_options():
|
||||
if request.method == "OPTIONS": raise HTTPResponse(headers=cors_headers)
|
||||
@app.hook("after_request")
|
||||
def enable_cors():
|
||||
for key, value in cors_headers.items(): response.set_header(key, value)
|
||||
|
||||
@app.route("/<filename>")
|
||||
def server_static(filename): return static_file(filename, root=(Path(__file__).parent / "tinychat").as_posix())
|
||||
@app.route("/assets/<filename:path>")
|
||||
def server_assets(filename): return static_file(filename, root=(Path(__file__).parent / "tinychat" / "assets").as_posix())
|
||||
@app.route("/")
|
||||
def index():
|
||||
return static_file("index.html", root=(Path(__file__).parent / "tinychat").as_posix())
|
||||
|
||||
@app.get("/v1/models")
|
||||
def models():
|
||||
return json.dumps([str(args.model)])
|
||||
|
||||
@app.post("/v1/internal/token-count")
|
||||
def token_count():
|
||||
rjson = json.loads(request.body.read())
|
||||
return json.dumps(len(tokenizer.encode(rjson.get("text", ""))))
|
||||
@app.post("/v1/token/encode")
|
||||
def token_encode():
|
||||
rjson = json.loads(request.body.read())
|
||||
return json.dumps(tokenizer.encode(rjson.get("text", "")))
|
||||
|
||||
@app.post("/v1/completions")
|
||||
def completions():
|
||||
rjson = json.loads(request.body.read())
|
||||
|
||||
# check if we are streaming
|
||||
if rjson.get("stream", False):
|
||||
response.content_type = "text/event-stream"
|
||||
response.set_header("Cache-Control", "no-cache")
|
||||
else: abort(400, "streaming required")
|
||||
|
||||
toks = [tokenizer.bos_id] + tokenizer.encode(rjson.get("prompt", ""), allow_special=True)
|
||||
|
||||
start_pos = prefill(model, toks[:-1])
|
||||
last_tok = toks[-1]
|
||||
while True:
|
||||
GlobalCounters.reset()
|
||||
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).item()
|
||||
start_pos += 1
|
||||
last_tok = tok
|
||||
if tok in tokenizer.stop_tokens: break
|
||||
|
||||
res = {
|
||||
"choices": [{
|
||||
"text": tokenizer.decode([tok]),
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(res)}\n\n"
|
||||
|
||||
@app.post("/v1/chat/token/encode")
|
||||
def chat_token_encode():
|
||||
rjson = json.loads(request.body.read())
|
||||
if "messages" not in rjson: abort(400, "messages required")
|
||||
toks = [tokenizer.bos_id]
|
||||
for message in rjson["messages"]:
|
||||
toks += encode_message(message["role"], message["content"])
|
||||
if len(rjson["messages"]) > 0 and message["role"] == "user":
|
||||
toks += encode_role("assistant")
|
||||
return json.dumps(toks)
|
||||
|
||||
@app.post("/v1/chat/completions")
|
||||
def chat_completions():
|
||||
global last_seen_toks
|
||||
rjson = json.loads(request.body.read())
|
||||
if "messages" not in rjson: abort(400, "messages required")
|
||||
|
||||
# check if we are streaming
|
||||
if rjson.get("stream", False):
|
||||
response.content_type = "text/event-stream"
|
||||
response.set_header("Cache-Control", "no-cache")
|
||||
else: abort(400, "streaming required")
|
||||
|
||||
toks = [tokenizer.bos_id]
|
||||
for message in rjson["messages"]:
|
||||
toks += encode_message(message["role"], message["content"])
|
||||
# ensure that the last message was a user message
|
||||
if message["role"] != "user": abort(400, "last message must be a user message")
|
||||
toks += encode_role("assistant")
|
||||
|
||||
random_id = random.randbytes(16).hex()
|
||||
|
||||
start_pos = prefill(model, toks[:-1])
|
||||
last_tok = toks[-1]
|
||||
last_seen_toks.append(last_tok)
|
||||
while True:
|
||||
GlobalCounters.reset()
|
||||
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P).item()
|
||||
start_pos += 1
|
||||
last_tok = tok
|
||||
last_seen_toks.append(tok)
|
||||
if tok in tokenizer.stop_tokens: break
|
||||
|
||||
res = {
|
||||
"id": random_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": str(args.model),
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"role": "assistant",
|
||||
"content": tokenizer.decode([tok]),
|
||||
},
|
||||
"finish_reason": None,
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(res)}\n\n"
|
||||
|
||||
res = {
|
||||
"id": random_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": str(args.model),
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop",
|
||||
}]
|
||||
}
|
||||
yield f"data: {json.dumps(res)}\n\n"
|
||||
|
||||
app.run(host=args.host, port=args.port, debug=args.debug)
|
||||
elif args.benchmark:
|
||||
toks = [tokenizer.bos_id] + encode_message("user", "Hello.") + encode_role("assistant")
|
||||
|
||||
start_pos = prefill(model, toks[:-1])
|
||||
last_tok = toks[-1]
|
||||
generated = ""
|
||||
for _ in range(20):
|
||||
GlobalCounters.reset()
|
||||
st = GlobalCounters.time_sum_s
|
||||
with Profiling(enabled=args.profile):
|
||||
with Timing("total ", on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
|
||||
with WallTimeEvent(BenchEvent.STEP):
|
||||
with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
|
||||
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
|
||||
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None):
|
||||
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
|
||||
tok = tok.item()
|
||||
start_pos += 1
|
||||
last_tok = tok
|
||||
generated += tokenizer.decode([tok])
|
||||
print(generated)
|
||||
if "LLaMA-3/8B-SF-DPO" in args.model.as_posix() and (TEMPERATURE == 0.85 or TEMPERATURE == 0):
|
||||
if TEMPERATURE == 0.85:
|
||||
EXPECTED_TEXT = {
|
||||
1: "Hello! How can I help you today? If you have any questions or need assistance with anything,",
|
||||
2: "Hello! How can I help you today? If you have any questions, need assistance or just want",
|
||||
3: "Hello! How can I help you today? If you have any questions or need assistance, feel free",
|
||||
4: "Hello! How can I assist you today? If you have any questions, need information, or require",
|
||||
5: "Hello! How can I assist you today? If you have any questions or need help with something",
|
||||
6: "Hello! How can I assist you today? If you have any questions, need information, or require",
|
||||
}
|
||||
else:
|
||||
EXPECTED_TEXT = {k: "Hello! How can I assist you today? If you have any questions or need help with something," for k in range(1, 7)}
|
||||
assert generated == EXPECTED_TEXT[args.shard], f"{generated=} {EXPECTED_TEXT[args.shard]}"
|
||||
print("\n" + colored("output validated", "green")) # NOTE: "\n" inside colored does not render the color in github action
|
||||
else:
|
||||
prompt = [tokenizer.bos_id] + encode_message("system", "You are an helpful assistant.")
|
||||
|
||||
start_pos = prefill(model, prompt)
|
||||
while True:
|
||||
toks = encode_message("user", input("Q: ")) + encode_role("assistant")
|
||||
|
||||
start_pos = prefill(model, toks[:-1], start_pos=start_pos)
|
||||
last_tok = toks[-1]
|
||||
while True:
|
||||
GlobalCounters.reset()
|
||||
if args.timing or args.profile: print("")
|
||||
st = GlobalCounters.time_sum_s
|
||||
with Profiling(enabled=args.profile):
|
||||
with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
|
||||
with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
|
||||
f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
|
||||
(f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
|
||||
|
||||
tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
|
||||
tok = tok.item()
|
||||
start_pos += 1
|
||||
last_tok = tok
|
||||
if tok in tokenizer.stop_tokens: break
|
||||
print(tokenizer.decode([tok]), end="", flush=True)
|
||||
print(flush=True)
|
||||
3
artifacts/package_sources/tinygrad/examples/llm.c/.gitignore
vendored
Normal file
3
artifacts/package_sources/tinygrad/examples/llm.c/.gitignore
vendored
Normal file
@@ -0,0 +1,3 @@
|
||||
data
|
||||
out.c
|
||||
a.out
|
||||
104
artifacts/package_sources/tinygrad/examples/llm.c/export.py
Executable file
104
artifacts/package_sources/tinygrad/examples/llm.c/export.py
Executable file
@@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
|
||||
from tinygrad import Device, nn, Tensor, dtypes
|
||||
from train_gpt2 import GPT, GPTConfig
|
||||
from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name, Context
|
||||
from tinygrad.engine.realize import get_kernel
|
||||
from tinygrad.schedule.memory import memory_planner
|
||||
from tinygrad.uop.ops import Ops
|
||||
|
||||
DEV.value = "CPU"
|
||||
|
||||
TIMING = getenv("TIMING")
|
||||
|
||||
if __name__ == "__main__":
|
||||
model = GPT(GPTConfig(n_layer=getenv("NLAYER", 12), n_head=12, n_embd=768))
|
||||
#model.load_pretrained()
|
||||
for p in nn.state.get_parameters(model): p.replace(Tensor.empty(p.shape, dtype=p.dtype)) # fake load pretrained
|
||||
|
||||
#early_sched = create_schedule([x.uop for x in nn.state.get_parameters(model)])
|
||||
#print(f"built model {len(early_sched)}")
|
||||
|
||||
#B, T = Variable("B", 1, 128).bind(4), 64 #Variable("T", 1, 1024).bind(64)
|
||||
B, T = 4, 64
|
||||
|
||||
optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=1e-4)
|
||||
warmup_count = getenv("WARMUP", 3)
|
||||
with Context(TRAINING=1):
|
||||
for i in range(warmup_count): # TODO: why does it take three and not two to stabilize
|
||||
GlobalCounters.reset()
|
||||
X = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
|
||||
Y = Tensor.empty(4, 64, dtype=dtypes.int).reshape(B, T)
|
||||
_, loss = model(X, Y)
|
||||
optimizer.zero_grad()
|
||||
if getenv("BACKWARD", 1):
|
||||
loss.backward()
|
||||
tensors = optimizer.schedule_step()
|
||||
else:
|
||||
tensors = []
|
||||
sched = loss.schedule(*tensors)
|
||||
print(f"calls {i}:", len(sched))
|
||||
#run_schedule(sched[:])
|
||||
sched = memory_planner(sched)
|
||||
ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
|
||||
srcs = {}
|
||||
for ast in ast_dedup:
|
||||
k = get_kernel(Device["CPU"].renderer, ast)
|
||||
k.linearize()
|
||||
src = Device["CPU"].renderer.render(to_function_name(k.name), k.uops)
|
||||
srcs[ast] = (k.name, src)
|
||||
print("functions:", len(srcs))
|
||||
used_buffers = dedup(flatten([si.bufs for si in sched]))
|
||||
numbered_bufs = {x:i for i,x in enumerate(used_buffers)}
|
||||
print("buffers:", len(numbered_bufs))
|
||||
|
||||
state_dict = nn.state.get_state_dict(model)
|
||||
state_dict.update({'X': X, 'Y': Y, 'loss': loss})
|
||||
grad_state_dict = {}
|
||||
for k,v in state_dict.items():
|
||||
if v.uop.base.buffer not in used_buffers: print(f"UNUSED: {k}")
|
||||
if v.grad is not None: grad_state_dict['grad_'+k] = v.grad
|
||||
state_dict.update(grad_state_dict)
|
||||
state_dict.update({'adam_b1_t': optimizer.b1_t, 'adam_b2_t': optimizer.b2_t, 'adam_lr': optimizer.lr})
|
||||
inverse_state_dict = {v:k for k,v in state_dict.items()}
|
||||
for p,m,v in zip(optimizer.params, optimizer.m, optimizer.v):
|
||||
nm = inverse_state_dict[p]
|
||||
state_dict["adam_m_"+nm] = m
|
||||
state_dict["adam_v_"+nm] = v
|
||||
named_buffers = {v.uop.base.buffer:k.replace(".", "_") for k,v in state_dict.items()}
|
||||
|
||||
c_code = ["#include <stdlib.h>", "#include <tgmath.h>", "#include <stdbool.h>"]
|
||||
if TIMING: c_code += ["#include <stdio.h>", "#include <time.h>"]
|
||||
c_code += [x[1].replace(" restrict ", " ")+"\n" for x in srcs.values()]
|
||||
|
||||
premain = ["int main() {"]
|
||||
if TIMING:
|
||||
premain += [" struct timespec tm0; clock_gettime(CLOCK_MONOTONIC, &tm0);"]
|
||||
lst = 0
|
||||
main = []
|
||||
|
||||
all_bufs = []
|
||||
for i,si in enumerate(sched):
|
||||
bufs = [(named_buffers.get(b, f"b{numbered_bufs[b]}"), b) for b in si.bufs]
|
||||
all_bufs += bufs
|
||||
if si.ast.op is not Ops.SINK:
|
||||
print(f"// {si.ast.op}", bufs)
|
||||
else:
|
||||
print(f"{srcs[si.ast][0]}({', '.join([x[0] for x in bufs])})")
|
||||
main.append(f" {to_function_name(srcs[si.ast][0])}({', '.join([x[0] for x in bufs])});")
|
||||
if TIMING:
|
||||
main.append(f" struct timespec tm{i+1}; clock_gettime(CLOCK_MONOTONIC, &tm{i+1});")
|
||||
main.append(f" printf(\"%10.2f ms + %7.2f ms @ {to_function_name(srcs[si.ast][0])}\\n\"," +\
|
||||
f"((tm{i+1}.tv_sec-tm{0}.tv_sec) + (tm{i+1}.tv_nsec-tm{0}.tv_nsec) / 1e9) * 1e3," +\
|
||||
f"((tm{i+1}.tv_sec-tm{lst}.tv_sec) + (tm{i+1}.tv_nsec-tm{lst}.tv_nsec) / 1e9) * 1e3);")
|
||||
lst = i+1
|
||||
#call = f"{srcs[si.ast][0]}({', '.join(bufs)})"
|
||||
#call += " "*(80-ansilen(call))
|
||||
#print(f"{call} // {i+1}")
|
||||
#print(srcs[si.ast][1])
|
||||
main.append("}")
|
||||
|
||||
mallocs = [f" {b.dtype.name}* {n} = ({b.dtype.name}*)malloc({b.nbytes});" for n,b in dedup(all_bufs)]
|
||||
|
||||
with open("out.c", "w") as f: f.write('\n'.join(c_code+premain+mallocs+main))
|
||||
206
artifacts/package_sources/tinygrad/examples/llm.c/train_gpt2.py
Executable file
206
artifacts/package_sources/tinygrad/examples/llm.c/train_gpt2.py
Executable file
@@ -0,0 +1,206 @@
|
||||
#!/usr/bin/env python3
|
||||
import os, math, time
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, nn, fetch, Device, TinyJit, GlobalCounters, Context
|
||||
from dataclasses import dataclass
|
||||
|
||||
@dataclass
|
||||
class GPTConfig:
|
||||
block_size: int = 1024
|
||||
vocab_size: int = 50257
|
||||
padded_vocab_size: int = 50304
|
||||
n_layer: int = 12
|
||||
n_head: int = 12
|
||||
n_embd: int = 768
|
||||
|
||||
class CausalSelfAttention:
|
||||
def __init__(self, config:GPTConfig):
|
||||
assert config.n_embd % config.n_head == 0
|
||||
# key, query, value projections for all heads, but in a batch
|
||||
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
|
||||
# output projection
|
||||
self.c_proj = nn.Linear(config.n_embd, config.n_embd)
|
||||
# regularization
|
||||
self.n_head = config.n_head
|
||||
self.n_embd = config.n_embd
|
||||
# not really a 'bias', more of a mask, but following the OpenAI/HF naming though
|
||||
self.bias = Tensor.ones(1, 1, config.block_size, config.block_size).tril()
|
||||
self.bias.is_param_(False)
|
||||
|
||||
def __call__(self, x:Tensor):
|
||||
B, T, C = x.shape
|
||||
qkv = self.c_attn(x)
|
||||
q, k, v = qkv.split(self.n_embd, dim=2)
|
||||
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
|
||||
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
|
||||
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
|
||||
|
||||
# manual implementation of attention
|
||||
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
|
||||
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
|
||||
att = att.softmax()
|
||||
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
|
||||
y = y.transpose(1, 2).view(B, T, C) # re-assemble all head outputs side by side
|
||||
# output projection
|
||||
y = self.c_proj(y)
|
||||
return y
|
||||
|
||||
class MLP:
|
||||
def __init__(self, config:GPTConfig):
|
||||
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd)
|
||||
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd)
|
||||
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
return self.c_proj(self.c_fc(x).gelu())
|
||||
|
||||
class Block:
|
||||
def __init__(self, config:GPTConfig):
|
||||
self.ln_1 = nn.LayerNorm(config.n_embd)
|
||||
self.attn = CausalSelfAttention(config)
|
||||
self.ln_2 = nn.LayerNorm(config.n_embd)
|
||||
self.mlp = MLP(config)
|
||||
|
||||
def __call__(self, x:Tensor):
|
||||
x = x + self.attn(self.ln_1(x))
|
||||
x = x + self.mlp(self.ln_2(x))
|
||||
return x
|
||||
|
||||
class GPT:
|
||||
def __init__(self, config:GPTConfig):
|
||||
self.config = config
|
||||
|
||||
self.wte = nn.Embedding(config.padded_vocab_size, config.n_embd)
|
||||
self.wpe = nn.Embedding(config.block_size, config.n_embd)
|
||||
self.h = [Block(config) for _ in range(config.n_layer)]
|
||||
self.ln_f = nn.LayerNorm(config.n_embd)
|
||||
self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False)
|
||||
self.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
|
||||
|
||||
def load_pretrained(self):
|
||||
weights = nn.state.torch_load(fetch(f'https://huggingface.co/gpt2/resolve/main/pytorch_model.bin'))
|
||||
transposed = ('attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight')
|
||||
for k in weights:
|
||||
if k == "wte.weight":
|
||||
weights[k] = weights[k].pad(((0, self.config.padded_vocab_size-self.config.vocab_size), (0,0))).to(None).contiguous()
|
||||
if k.endswith(transposed):
|
||||
weights[k] = weights[k].to(None).T.contiguous()
|
||||
# lm head and wte are tied
|
||||
weights['lm_head.weight'] = weights['wte.weight']
|
||||
nn.state.load_state_dict(self, weights)
|
||||
|
||||
def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
|
||||
for _ in range(max_new_tokens):
|
||||
idx_cond = idx if idx.shape[1] <= self.config.block_size else idx[:, -self.config.block_size:]
|
||||
logits, _ = self(idx_cond)
|
||||
logits = logits[:, -1, :] / temperature
|
||||
idx_next = logits.softmax().multinomial()
|
||||
idx = Tensor.cat(idx, idx_next, dim=1)
|
||||
return idx
|
||||
|
||||
def __call__(self, idx:Tensor, targets=None):
|
||||
b, t = idx.shape
|
||||
pos = Tensor.arange(0, t)
|
||||
|
||||
tok_emb = self.wte(idx) # token embeddings of shape (b, t, n_embd)
|
||||
pos_emb = self.wpe(pos) # position embeddings of shape (t, n_embd)
|
||||
x = tok_emb + pos_emb
|
||||
|
||||
x = self.ln_f(x.sequential(self.h))
|
||||
|
||||
if targets is not None:
|
||||
logits = self.lm_head(x)[:, :, :self.config.vocab_size]
|
||||
loss = logits.sparse_categorical_crossentropy(targets)
|
||||
else:
|
||||
logits = self.lm_head(x[:, [-1], :])[:, :, :self.config.vocab_size]
|
||||
loss = None
|
||||
|
||||
return logits, loss
|
||||
|
||||
if __name__ == "__main__":
|
||||
import tiktoken, argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--num_iterations", type=int, default=10, help="number of iterations to run")
|
||||
parser.add_argument("--batch_size", type=int, default=4, help="batch size")
|
||||
parser.add_argument("--sequence_length", type=int, default=64, help="sequence length")
|
||||
parser.add_argument("--skip_test", action="store_true", help="skip test")
|
||||
parser.add_argument("--gpus", type=int, default=1, help="sequence length")
|
||||
args = parser.parse_args()
|
||||
B, T = args.batch_size, args.sequence_length
|
||||
assert 1 <= T <= 1024
|
||||
|
||||
model = GPT(GPTConfig(n_layer=12, n_head=12, n_embd=768))
|
||||
model.load_pretrained()
|
||||
|
||||
if args.gpus > 1:
|
||||
GPUS = tuple(f'{Device.DEFAULT}:{i}' for i in range(args.gpus))
|
||||
for x in nn.state.get_parameters(model): x.to_(GPUS) # we put a copy of the model on every GPU
|
||||
|
||||
# init the tokenizer
|
||||
enc = tiktoken.get_encoding("gpt2")
|
||||
encode = lambda s: enc.encode(s, allowed_special={"<|endoftext|>"})
|
||||
decode = lambda l: enc.decode(l)
|
||||
|
||||
# load the tokens
|
||||
# prefer to use tiny_shakespeare if it's available, otherwise use tiny_stories
|
||||
# we're using val instead of train split just because it is smaller/faster
|
||||
tokens_bin = fetch("https://huggingface.co/datasets/karpathy/llmc-starter-pack/resolve/main/tiny_shakespeare_val.bin")
|
||||
assert os.path.isfile(tokens_bin)
|
||||
print(f"loading cached tokens in {tokens_bin}")
|
||||
with open(tokens_bin, "rb") as f:
|
||||
f.seek(0x400)
|
||||
tokens = np.frombuffer(f.read(), dtype=np.uint16).astype(np.int32)
|
||||
tokens = Tensor(tokens)
|
||||
|
||||
# lightweight dataloader
|
||||
def get_batch():
|
||||
assert B*T+1 <= len(tokens), "not enough tokens"
|
||||
# for 338,025 tokens. E.g. with B=8 T=1024, this will yield 41 batches before looping
|
||||
i = 0
|
||||
while True:
|
||||
x = tokens[i:i+B*T].view(B, T)
|
||||
y = tokens[i+1:i+B*T+1].view(B, T)
|
||||
yield x, y
|
||||
i += B*T
|
||||
if i + B*T + 1 >= len(tokens):
|
||||
i = 0 # in prod we'd want to randomize the start point a bit
|
||||
|
||||
# forward backward for a few iterations
|
||||
data_iter = iter(get_batch())
|
||||
x, y = next(data_iter) # we'll overfit this batch below
|
||||
optimizer = nn.optim.AdamW(nn.state.get_parameters(model), lr=1e-4, weight_decay=0)
|
||||
|
||||
print(f"model state: {sum(x.nbytes() for x in nn.state.get_parameters(model))/1e9:.2f} GB")
|
||||
print(f"optimizer state: {sum(x.nbytes() for x in nn.state.get_parameters(optimizer))/1e9:.2f} GB")
|
||||
|
||||
# shard the data on axis 0
|
||||
if args.gpus > 1: x, y = x.shard(GPUS, axis=0), y.shard(GPUS, axis=0)
|
||||
|
||||
@TinyJit
|
||||
@Context(TRAINING=1)
|
||||
def step(x:Tensor, y:Tensor) -> Tensor:
|
||||
_, loss = model(x, y)
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
return loss.realize(*optimizer.schedule_step())
|
||||
|
||||
for i in range(args.num_iterations):
|
||||
GlobalCounters.reset()
|
||||
t0 = time.perf_counter()
|
||||
loss = step(x.contiguous(), y.contiguous())
|
||||
Device[Device.DEFAULT].synchronize()
|
||||
t1 = time.perf_counter()
|
||||
print(f"iteration {i}, loss: {loss.item():.6f}, time: {(t1-t0)*1000:.3f}ms, {int(B*T/(t1-t0))} tok/s, {GlobalCounters.global_mem/1e9:.2f} GB")
|
||||
|
||||
if not args.skip_test:
|
||||
# copy back to single gpu for test
|
||||
if args.gpus > 1:
|
||||
for x in nn.state.get_parameters(model): x.to_(Device.DEFAULT)
|
||||
start = "<|endoftext|>"
|
||||
start_ids = encode(start)
|
||||
x = (Tensor(start_ids)[None, ...])
|
||||
max_new_tokens = 16
|
||||
temperature = 1.0
|
||||
top_k = 40
|
||||
y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
|
||||
print(decode(y[0].tolist()))
|
||||
@@ -0,0 +1,65 @@
|
||||
// clang -Ofast -Wno-unused-result -march=native matmul.c
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <time.h>
|
||||
|
||||
float b52[786432];
|
||||
float b49[196608];
|
||||
float h_0_mlp_c_fc_weight[2359296];
|
||||
float h_0_mlp_c_fc_bias[3072];
|
||||
|
||||
void matmul_forward(float* out,
|
||||
float* inp, float* weight, float* bias,
|
||||
int B, int T, int C, int OC) {
|
||||
// most of the running time is spent here and in matmul_backward
|
||||
// OC is short for "output channels"
|
||||
// inp is (B,T,C), weight is (OC, C), bias is (OC)
|
||||
// out will be (B,T,OC)
|
||||
#pragma omp parallel for collapse(2)
|
||||
for (int b = 0; b < B; b++) {
|
||||
for (int t = 0; t < T; t++) {
|
||||
float* out_bt = out + b * T * OC + t * OC;
|
||||
float* inp_bt = inp + b * T * C + t * C;
|
||||
for (int o = 0; o < OC; o++) {
|
||||
float val = (bias != NULL) ? bias[o] : 0.0f;
|
||||
float* wrow = weight + o*C;
|
||||
for (int i = 0; i < C; i++) {
|
||||
val += inp_bt[i] * wrow[i];
|
||||
}
|
||||
out_bt[o] = val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void r_256_3072_768(float* restrict data0, const float* restrict data1, const float* restrict data2, const float* restrict data3) {
|
||||
for (int ridx0 = 0; ridx0 < 256; ridx0++) {
|
||||
for (int ridx1 = 0; ridx1 < 3072; ridx1++) {
|
||||
float acc0 = 0.0f;
|
||||
float val0 = data3[ridx1];
|
||||
for (int ridx2 = 0; ridx2 < 768; ridx2++) {
|
||||
float val1 = data1[(ridx0*768)+ridx2];
|
||||
float val2 = data2[(ridx1*768)+ridx2];
|
||||
acc0 = ((val1*val2)+acc0);
|
||||
}
|
||||
data0[(ridx0*3072)+ridx1] = (acc0+val0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
int main() {
|
||||
for (int i = 0; i < 5; i++) {
|
||||
struct timespec t1, t2, t3;
|
||||
clock_gettime(CLOCK_MONOTONIC, &t1);
|
||||
r_256_3072_768(b52, b49, h_0_mlp_c_fc_weight, h_0_mlp_c_fc_bias);
|
||||
clock_gettime(CLOCK_MONOTONIC, &t2);
|
||||
matmul_forward(b52, b49, h_0_mlp_c_fc_weight, h_0_mlp_c_fc_bias, 4, 64, 768, 3072);
|
||||
clock_gettime(CLOCK_MONOTONIC, &t3);
|
||||
double time_gen = (t2.tv_sec - t1.tv_sec) + (t2.tv_nsec - t1.tv_nsec) / 1e9;
|
||||
double time_real = (t3.tv_sec - t2.tv_sec) + (t3.tv_nsec - t2.tv_nsec) / 1e9;
|
||||
printf("%.2f ms gen vs %.2f ms reference\n", time_gen*1e3, time_real*1e3);
|
||||
}
|
||||
}
|
||||
|
||||
324
artifacts/package_sources/tinygrad/examples/mamba.py
Normal file
324
artifacts/package_sources/tinygrad/examples/mamba.py
Normal file
@@ -0,0 +1,324 @@
|
||||
import os, sys, math, argparse, time
|
||||
sys.path.append(os.getcwd())
|
||||
from typing import Any, Optional, Dict
|
||||
|
||||
from tinygrad import Tensor, TinyJit, nn
|
||||
from tinygrad.helpers import fetch
|
||||
from tinygrad.nn.state import load_state_dict, torch_load
|
||||
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
MODELS = {
|
||||
"130m": {"dim": 768, "n_layers": 24, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
|
||||
"370m": {"dim": 1024, "n_layers": 48, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
|
||||
"790m": {"dim": 1536, "n_layers": 48, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
|
||||
"1.4b": {"dim": 2048, "n_layers": 48, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
|
||||
"2.8b": {"dim": 2560, "n_layers": 64, "vocab_size": 50277, "pad_vocab_size_multiple": 8},
|
||||
}
|
||||
|
||||
def fetch_weights(model_name: str) -> Dict[str, Tensor]:
|
||||
if model_name not in MODELS:
|
||||
raise ValueError(f"Requested unknown mamba model: {model_name}")
|
||||
downloaded = fetch(f"https://huggingface.co/state-spaces/mamba-{model_name}/resolve/main/pytorch_model.bin?download=true")
|
||||
return torch_load(downloaded)
|
||||
|
||||
def selective_scan_ref(
|
||||
u,
|
||||
delta,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
D=None,
|
||||
z=None,
|
||||
delta_bias=None,
|
||||
delta_softplus=False,
|
||||
return_last_state=False,
|
||||
):
|
||||
"""
|
||||
u: r(B D L)
|
||||
delta: r(B D L)
|
||||
A: c(D N) or r(D N)
|
||||
B: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
|
||||
C: c(D N) or r(B N L) or r(B N 2L) or r(B G N L) or (B G N L)
|
||||
D: r(D)
|
||||
z: r(B D L)
|
||||
delta_bias: r(D), fp32
|
||||
|
||||
out: r(B D L)
|
||||
last_state (optional): r(B D dstate) or c(B D dstate)
|
||||
"""
|
||||
u = u.float()
|
||||
delta = delta.float()
|
||||
if delta_bias is not None:
|
||||
delta = delta + delta_bias[..., None].float()
|
||||
if delta_softplus:
|
||||
delta = delta.softplus()
|
||||
batch, dim, dstate = u.shape[0], A.shape[0], A.shape[1]
|
||||
is_variable_B = len(B.shape) >= 3
|
||||
is_variable_C = len(C.shape) >= 3
|
||||
x = Tensor.zeros(batch, dim, dstate)
|
||||
ys = []
|
||||
deltaA = Tensor.einsum("bdl,dn->bdln", delta, A).exp()
|
||||
if not is_variable_B:
|
||||
deltaB_u = Tensor.einsum("bdl,dn,bdl->bdln", delta, B, u)
|
||||
else:
|
||||
if len(B.shape) == 3:
|
||||
deltaB_u = Tensor.einsum("bdl,bnl,bdl->bdln", delta, B, u)
|
||||
else:
|
||||
B = B.repeat((1, dim // B.shape[1], 1, 1))
|
||||
deltaB_u = Tensor.einsum("bdl,bdnl,bdl->bdln", delta, B, u)
|
||||
if is_variable_C and len(C.shape) == 4:
|
||||
C = C.repeat((1, dim // C.shape[1], 1, 1))
|
||||
last_state = None
|
||||
for i in range(u.shape[2]):
|
||||
x = deltaA[:, :, i] * x + deltaB_u[:, :, i]
|
||||
if not is_variable_C:
|
||||
y = Tensor.einsum("bdn,dn->bd", x, C)
|
||||
else:
|
||||
if len(C.shape) == 3:
|
||||
y = Tensor.einsum("bdn,bn->bd", x, C[:, :, i])
|
||||
else:
|
||||
y = Tensor.einsum("bdn,bdn->bd", x, C[:, :, :, i])
|
||||
if i == u.shape[2] - 1:
|
||||
last_state = x
|
||||
ys.append(y)
|
||||
y = Tensor.stack(*ys, dim=2) # (batch dim L)
|
||||
out = y if D is None else y + u * D.reshape((-1, 1))
|
||||
if z is not None:
|
||||
out = out * z.silu()
|
||||
return out if not return_last_state else (out, last_state)
|
||||
|
||||
class MambaMixer:
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
d_state=16,
|
||||
d_conv=4,
|
||||
expand=2,
|
||||
dt_rank="auto",
|
||||
dt_min=0.001,
|
||||
dt_max=0.1,
|
||||
dt_init="random",
|
||||
dt_scale=1.0,
|
||||
dt_init_floor=1e-4,
|
||||
conv_bias=True,
|
||||
bias=False,
|
||||
layer_idx=None,
|
||||
):
|
||||
self.dim = dim
|
||||
self.d_state = d_state
|
||||
self.d_conv = d_conv
|
||||
self.expand = expand
|
||||
self.d_inner = self.expand * self.dim
|
||||
self.dt_rank = math.ceil(self.dim / 16) if dt_rank == "auto" else dt_rank
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
self.in_proj = nn.Linear(self.dim, self.d_inner * 2, bias=bias)
|
||||
|
||||
self.conv1d = nn.Conv1d(in_channels=self.d_inner, out_channels=self.d_inner, bias=conv_bias,
|
||||
kernel_size=d_conv, groups=self.d_inner, padding=d_conv-1)
|
||||
|
||||
self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False)
|
||||
self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)
|
||||
|
||||
# Initialize special dt projection to preserve variance at initialization
|
||||
dt_init_std = self.dt_rank**-0.5 * dt_scale
|
||||
if dt_init == "constant":
|
||||
self.dt_proj.weight = Tensor.full(self.dt_proj.weight.shape, dt_init_std)
|
||||
elif dt_init == "random":
|
||||
self.dt_proj.weight = Tensor.uniform(self.dt_proj.weight.shape, low=-dt_init_std, high=dt_init_std)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
dt = Tensor.uniform(self.d_inner, low=math.log(dt_min), high=math.log(dt_max)).exp().maximum(dt_init_floor)
|
||||
inv_dt = dt + (1 - (-dt).exp()).log()
|
||||
|
||||
self.dt_proj.bias.assign(inv_dt)
|
||||
|
||||
# S4D real initialization
|
||||
self.A_log = Tensor.arange(1, self.d_state+1).repeat([self.d_inner, 1]).log()
|
||||
|
||||
# D "skip" parameter
|
||||
self.D = Tensor.ones(self.d_inner) # Keep in fp32
|
||||
|
||||
self.out_proj = nn.Linear(self.d_inner, self.dim, bias=bias)
|
||||
|
||||
def __call__(self, hidden_states: Tensor):
|
||||
batch, seqlen, _ = hidden_states.shape
|
||||
|
||||
if not hasattr(self, 'conv_state'):
|
||||
self.conv_state = Tensor.zeros(batch, self.dim * self.expand, self.d_conv).contiguous().realize()
|
||||
self.ssm_state = Tensor.zeros(batch, self.dim * self.expand, self.d_state).realize()
|
||||
|
||||
xz = self.in_proj.weight @ hidden_states.permute(2,0,1).reshape(hidden_states.shape[2],hidden_states.shape[1]*hidden_states.shape[0])
|
||||
xz = xz.reshape(xz.shape[0],xz.shape[1]//seqlen, seqlen).permute(1,0,2)
|
||||
|
||||
if self.in_proj.bias is not None:
|
||||
xz = xz + self.in_proj.bias.reshape((-1, 1))
|
||||
|
||||
A = -self.A_log.exp()
|
||||
x, z = xz.chunk(2, dim=1)
|
||||
# Compute short convolution
|
||||
self.conv_state.assign(x[:, :, -self.d_conv :]) # Update state (B D W)
|
||||
x = self.conv1d(x)[..., :seqlen].swish()
|
||||
|
||||
x_dbl = self.x_proj(x.permute(0,2,1).reshape(x.shape[0]*x.shape[2], x.shape[1]))
|
||||
dt, B, C = Tensor.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
|
||||
dt = self.dt_proj.weight @ dt.T
|
||||
dt = dt.reshape(dt.shape[0], dt.shape[1]//seqlen, seqlen).permute(1,0,2)
|
||||
B = B.reshape(B.shape[0]//seqlen, seqlen, B.shape[1]).permute(0,2,1)
|
||||
C = C.reshape(C.shape[0]//seqlen, seqlen, C.shape[1]).permute(0,2,1)
|
||||
|
||||
# TODO: actually implement selective_scan_fn
|
||||
y = selective_scan_ref(x, dt, A, B, C, self.D, z=z, delta_bias=self.dt_proj.bias, delta_softplus=True,
|
||||
return_last_state=True)
|
||||
|
||||
y, last_state = y
|
||||
self.ssm_state.assign(last_state).realize()
|
||||
y = y.permute(0,2,1)
|
||||
out = self.out_proj(y)
|
||||
return out
|
||||
else:
|
||||
return self.step(hidden_states)
|
||||
|
||||
def step(self, hidden_states: Tensor):
|
||||
assert hidden_states.shape[1] == 1, f"Only support decoding with 1 token at a time for now, attempted {hidden_states.shape[1]}"
|
||||
xz = self.in_proj(hidden_states.squeeze(1)) # (B 2D)
|
||||
x, z = xz.chunk(2, dim=-1) # (B D)
|
||||
|
||||
# Conv step
|
||||
self.conv_state.assign(self.conv_state[:, :, 1:].cat(x.unsqueeze(-1), dim=-1).realize())
|
||||
x = (self.conv_state * self.conv1d.weight.squeeze(1)).sum(-1)
|
||||
if self.conv1d.bias is not None:
|
||||
x = x + self.conv1d.bias
|
||||
x = x.swish()
|
||||
|
||||
x_db = self.x_proj(x) # (B dt_rank+2*d_state)
|
||||
dt = x_db[:, : self.dt_rank]
|
||||
B = x_db[:, self.dt_rank : (self.dt_rank + self.d_state)]
|
||||
C = x_db[:, (self.dt_rank + self.d_state) :]
|
||||
# Don't add dt_bias here
|
||||
dt = self.dt_proj.weight @ dt.T
|
||||
A = -self.A_log.exp()
|
||||
|
||||
# SSM step
|
||||
dt = (dt + self.dt_proj.bias.unsqueeze(-1)).softplus()
|
||||
dA = Tensor.einsum("db,dn->bdn", dt, A).exp()
|
||||
dB = Tensor.einsum("db,bn->bdn", dt, B)
|
||||
self.ssm_state.assign(self.ssm_state * dA + x.unsqueeze(-1) * dB)
|
||||
y = Tensor.einsum("bdn,bn->bd", self.ssm_state, C)
|
||||
y = y + self.D * x
|
||||
y = y * z.swish() # (B D)
|
||||
|
||||
out = self.out_proj(y)
|
||||
return out.unsqueeze(1)
|
||||
|
||||
class MambaBlock:
|
||||
def __init__(self, dim: int, norm_eps: float = 1e-5, rms_norm: bool = True, layer_idx: Optional[int] = None):
|
||||
self.mixer = MambaMixer(dim, layer_idx=layer_idx)
|
||||
if rms_norm:
|
||||
self.norm = nn.RMSNorm(dim, norm_eps)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def __call__(self, hidden_states: Tensor, residual: Optional[Tensor] = None):
|
||||
residual = (hidden_states + residual) if residual is not None else hidden_states
|
||||
hidden_states = self.norm(residual)
|
||||
hidden_states = self.mixer(hidden_states)
|
||||
return hidden_states, residual
|
||||
|
||||
class MambaBackbone:
|
||||
def __init__(self, dim: int, n_layers: int, vocab_size: int, rms_norm: bool = True, norm_eps: float = 1e-5):
|
||||
self.embedding = nn.Embedding(vocab_size, dim)
|
||||
self.layers = [MambaBlock(dim, rms_norm=rms_norm, layer_idx=i) for i in range(n_layers)]
|
||||
if rms_norm:
|
||||
self.norm_f = nn.RMSNorm(dim, norm_eps)
|
||||
|
||||
def __call__(self, input_ids: Tensor) -> Any:
|
||||
hidden_states = self.embedding(input_ids)
|
||||
residual = None
|
||||
for layer in self.layers:
|
||||
hidden_states, residual = layer(hidden_states, residual)
|
||||
|
||||
residual = (hidden_states + residual) if residual is not None else hidden_states
|
||||
hidden_states = self.norm_f(residual)
|
||||
return hidden_states
|
||||
|
||||
class Mamba:
|
||||
def __init__(self, dim: int, n_layers: int, vocab_size: int, pad_vocab_size_multiple: int = 1):
|
||||
if vocab_size % pad_vocab_size_multiple != 0:
|
||||
vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)
|
||||
|
||||
self.backbone = MambaBackbone(dim, n_layers, vocab_size)
|
||||
self.lm_head = nn.Linear(dim, vocab_size, bias=False)
|
||||
|
||||
self.forward_jit = TinyJit(self.forward)
|
||||
|
||||
def forward(self, input_ids:Tensor):
|
||||
hidden_states = self.backbone(input_ids)
|
||||
return self.lm_head(hidden_states).realize()
|
||||
|
||||
def __call__(self, input_ids):
|
||||
return self.forward(input_ids)
|
||||
|
||||
@staticmethod
|
||||
def from_pretrained(model_name: str):
|
||||
weights = fetch_weights(model_name)
|
||||
model = Mamba(**MODELS[model_name])
|
||||
load_state_dict(model, weights)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: bool = 1.0, sample: bool = False, top_k: int = None):
|
||||
tks = tokenizer(prompt)["input_ids"]
|
||||
while len(tks) < 4:
|
||||
tks = [50279] + tks
|
||||
|
||||
# Loading in the prompt tokens
|
||||
logits = model.forward(Tensor([tks]))[:, -1, :]
|
||||
for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"):
|
||||
if sample:
|
||||
scaled_logits = logits / temp
|
||||
if top_k is not None:
|
||||
topk_values, topk_indices = scaled_logits.topk(top_k)
|
||||
filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
|
||||
filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
|
||||
tok_Tens = filtered_logits.softmax().multinomial()
|
||||
else:
|
||||
tok_Tens = scaled_logits.softmax().multinomial()
|
||||
else:
|
||||
tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
|
||||
tok = tok_Tens.item()
|
||||
tks.append(tok)
|
||||
logits = model.forward_jit(tok_Tens)[:, -1, :]
|
||||
|
||||
output_completions = ''.join([tokenizer.decode(output) for output in tks])
|
||||
return output_completions
|
||||
|
||||
if __name__ == "__main__":
|
||||
ORIG_PROMPT = "Why is gravity "
|
||||
parser = argparse.ArgumentParser(description="Run Mamba in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument("--prompt", type=str, default="Why is gravity ", help="Prompt for LLM completion")
|
||||
parser.add_argument("--size", type=str, default="370m",
|
||||
help=f"Size of model to use [{', '.join([k for k in MODELS.keys()])}]")
|
||||
parser.add_argument("--n_tokens", type=int, default=10, help="Number of tokens to generate")
|
||||
parser.add_argument("--top_k", type=int, help="Limit sampling to the top k most likely tokens")
|
||||
parser.add_argument("--sample", dest="sample", action="store_true", help="Sample flag")
|
||||
parser.add_argument("--temp", type=float, default=1.0, help="Sampling temp has to be <=1.0")
|
||||
args = parser.parse_args()
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
|
||||
model = Mamba.from_pretrained(args.size)
|
||||
prompt = args.prompt
|
||||
num_toks = args.n_tokens
|
||||
sample = args.sample
|
||||
temp = args.temp
|
||||
top_k = args.top_k
|
||||
s = time.time()
|
||||
tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
|
||||
print(tinyoutput)
|
||||
print('TIME: ', time.time() - s)
|
||||
TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"
|
||||
if ORIG_PROMPT == prompt and not sample and num_toks==10 and args.size=='370m': print('Outputs Match:', tinyoutput == TORCHOUTPUT)
|
||||
156
artifacts/package_sources/tinygrad/examples/minrf.py
Normal file
156
artifacts/package_sources/tinygrad/examples/minrf.py
Normal file
@@ -0,0 +1,156 @@
|
||||
# much taken from https://github.com/cloneofsimo/minRF
|
||||
from tinygrad import Tensor, nn, GlobalCounters, TinyJit, Context
|
||||
from tinygrad.helpers import getenv, trange
|
||||
from extra.models.llama import Attention, FeedForward, precompute_freqs_cis
|
||||
|
||||
def modulate(x:Tensor, shift:Tensor, scale:Tensor) -> Tensor: return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
|
||||
# TODO: why doesn't the TimestepEmbedder from minRF work?
|
||||
class TimestepEmbedder:
|
||||
def __init__(self, hidden_size): self.mlp = [nn.Linear(1, hidden_size), Tensor.silu, nn.Linear(hidden_size, hidden_size)]
|
||||
def __call__(self, t:Tensor): return t.reshape(-1, 1).sequential(self.mlp)
|
||||
|
||||
class TransformerBlock:
|
||||
def __init__(self, dim, n_heads, norm_eps=1e-5):
|
||||
self.attention = Attention(dim, n_heads)
|
||||
self.feed_forward = FeedForward(dim, 4*dim)
|
||||
self.attention_norm = nn.LayerNorm(dim, eps=norm_eps)
|
||||
self.ffn_norm = nn.LayerNorm(dim, eps=norm_eps)
|
||||
self.adaLN_modulation = nn.Linear(dim, 6 * dim, bias=True)
|
||||
|
||||
def __call__(self, x:Tensor, freqs_cis:Tensor, adaln_input:Tensor):
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(adaln_input.silu()).chunk(6, dim=1)
|
||||
x = x + gate_msa.unsqueeze(1) * self.attention(modulate(self.attention_norm(x), shift_msa, scale_msa), 0, freqs_cis)
|
||||
x = x + gate_mlp.unsqueeze(1) * self.feed_forward(modulate(self.ffn_norm(x), shift_mlp, scale_mlp))
|
||||
return x.contiguous().contiguous_backward()
|
||||
|
||||
class FinalLayer:
|
||||
def __init__(self, dim, patch_size, out_channels):
|
||||
self.norm_final = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(dim, patch_size*patch_size*out_channels, bias=True)
|
||||
self.adaLN_modulation = nn.Linear(dim, 2 * dim, bias=True)
|
||||
|
||||
# init weights/bias to 0
|
||||
self.linear.weight.replace(self.linear.weight.zeros_like().contiguous())
|
||||
self.linear.bias.replace(self.linear.bias.zeros_like().contiguous())
|
||||
|
||||
def __call__(self, x:Tensor, c:Tensor):
|
||||
shift, scale = self.adaLN_modulation(c.silu()).chunk(2, dim=1)
|
||||
x = modulate(self.norm_final(x), shift, scale)
|
||||
return self.linear(x)
|
||||
|
||||
# channels=1, input_size=32, dim=64, n_layers=6, n_heads=4, num_classes=10
|
||||
class DiT_Llama:
|
||||
def __init__(self, in_channels=1, dim=64, n_layers=6, n_heads=4, num_classes=10, patch_size=2):
|
||||
self.patch_size = patch_size
|
||||
self.out_channels = in_channels
|
||||
self.num_classes = num_classes
|
||||
|
||||
self.init_conv_seq = [
|
||||
nn.Conv2d(in_channels, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
|
||||
nn.Conv2d(dim //2, dim // 2, kernel_size=5, padding=2, stride=1), Tensor.silu, nn.GroupNorm(32, dim//2),
|
||||
]
|
||||
|
||||
self.x_embedder = nn.Linear(self.patch_size * self.patch_size * dim // 2, dim, bias=True)
|
||||
self.t_embedder = TimestepEmbedder(dim)
|
||||
self.y_embedder = nn.Embedding(num_classes+1, dim)
|
||||
self.final_layer = FinalLayer(dim, self.patch_size, self.out_channels)
|
||||
|
||||
self.freqs_cis = precompute_freqs_cis(dim // n_heads, 4096)
|
||||
self.layers = [TransformerBlock(dim, n_heads) for _ in range(n_layers)]
|
||||
|
||||
def unpatchify(self, x:Tensor):
|
||||
c, p = self.out_channels, self.patch_size
|
||||
h = w = int(x.shape[1] ** 0.5)
|
||||
x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
|
||||
x = x.rearrange("n h w p q c -> n c h p w q")
|
||||
return x.reshape(shape=(x.shape[0], c, h * p, h * p))
|
||||
|
||||
def patchify(self, x:Tensor):
|
||||
B, C, H, W = x.shape
|
||||
x = x.reshape(B, C, H // self.patch_size, self.patch_size, W // self.patch_size, self.patch_size)
|
||||
x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
|
||||
return x # B <H*W ish> <C*patch_size*patch_size>
|
||||
|
||||
def __call__(self, x:Tensor, t:Tensor, y:Tensor) -> Tensor:
|
||||
x = x.sequential(self.init_conv_seq)
|
||||
x = self.patchify(x)
|
||||
x = self.x_embedder(x)
|
||||
adaln_input = self.t_embedder(t) + self.y_embedder(y)
|
||||
adaln_input = adaln_input.contiguous()
|
||||
for layer in self.layers:
|
||||
x = layer(x, self.freqs_cis[:, :x.size(1)], adaln_input=adaln_input)
|
||||
x = self.final_layer(x, adaln_input)
|
||||
return self.unpatchify(x)
|
||||
|
||||
def rf(self, x:Tensor, cond:Tensor):
|
||||
b = x.shape[0]
|
||||
# self.ln is True
|
||||
t = Tensor.randn((b,)).sigmoid()
|
||||
texp = t.view([b, *([1] * len(x.shape[1:]))])
|
||||
|
||||
# conditional dropout
|
||||
dropout_prob = 0.1
|
||||
cond = (Tensor.rand(cond.shape[0]) < dropout_prob).where(cond.full_like(self.num_classes), cond)
|
||||
|
||||
# this is rectified flow
|
||||
z1 = x.randn_like()
|
||||
zt = (1 - texp) * x + texp * z1
|
||||
vtheta = self(zt, t, cond)
|
||||
|
||||
# MSE loss
|
||||
return ((z1 - x) - vtheta).square().mean()
|
||||
|
||||
def sample(self, z, cond, null_cond, sample_steps=50, cfg=2.0):
|
||||
b = z.size(0)
|
||||
dt = Tensor.full((b,)+(1,)*len(z.shape[1:]), fill_value=1.0/sample_steps).contiguous()
|
||||
images = [z]
|
||||
for i in range(sample_steps, 0, -1):
|
||||
t = Tensor.full((b,), fill_value=i/sample_steps).contiguous()
|
||||
vc = self(z, t, cond)
|
||||
vu = self(z, t, null_cond)
|
||||
vc = vu + cfg * (vc - vu)
|
||||
z = z - dt * vc
|
||||
z = z.contiguous()
|
||||
images.append(z)
|
||||
return images
|
||||
|
||||
def mviz(t:Tensor):
|
||||
assert len(t.shape) == 4 and t.shape[1] == 1
|
||||
ft = t.permute(1,2,0,3).reshape(32, -1)
|
||||
assert ft.shape[-1]%32 == 0
|
||||
print("")
|
||||
for y in ((ft+1)/2).clamp(0,1).tolist():
|
||||
ln = [f"\033[38;5;{232+int(x*23)}m██" for x in y]
|
||||
print(''.join(ln) + "\033[0m")
|
||||
|
||||
if __name__ == "__main__":
|
||||
X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
|
||||
X_train = X_train.pad((2,2,2,2))
|
||||
X_train = ((X_train.float()/255)-0.5)/0.5
|
||||
Y_train = Y_train.int()
|
||||
|
||||
model = DiT_Llama(patch_size=getenv("PATCH_SIZE", 2))
|
||||
for r in nn.state.get_parameters(model): r.realize()
|
||||
optimizer = nn.optim.Adam(nn.state.get_parameters(model), lr=5e-4)
|
||||
|
||||
@TinyJit
|
||||
@Context(TRAINING=1)
|
||||
def train_step():
|
||||
if getenv("OVERFIT"): samples = Tensor.zeros(getenv("BS", 256), dtype='int')
|
||||
else: samples = Tensor.randint(getenv("BS", 256), high=X_train.shape[0])
|
||||
optimizer.zero_grad()
|
||||
loss = model.rf(X_train[samples], Y_train[samples])
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
return loss
|
||||
|
||||
@TinyJit
|
||||
def sample(z:Tensor, cond:Tensor) -> Tensor:
|
||||
return model.sample(z, cond, Tensor.full_like(cond, 10), sample_steps=getenv("SAMPLE_STEPS", 20))[-1]
|
||||
|
||||
for steps in (t:=trange(getenv("STEPS", 5000))):
|
||||
if steps%10 == 0: mviz(sample(Tensor.randn(3, 1, 32, 32), Tensor([5,0,4], dtype='int')))
|
||||
GlobalCounters.reset()
|
||||
loss = train_step()
|
||||
t.set_description(f"loss: {loss.item():9.2f}")
|
||||
62
artifacts/package_sources/tinygrad/examples/mixtral.py
Normal file
62
artifacts/package_sources/tinygrad/examples/mixtral.py
Normal file
@@ -0,0 +1,62 @@
|
||||
import functools, argparse, pathlib
|
||||
from tinygrad import Tensor, nn, Device, GlobalCounters, Variable
|
||||
from tinygrad.helpers import Timing, Profiling, tqdm
|
||||
from tinygrad.nn.state import torch_load, get_state_dict
|
||||
from extra.models.llama import FeedForward, Transformer
|
||||
from extra.bench_log import BenchEvent, WallTimeEvent
|
||||
|
||||
class MixtureFeedForward:
|
||||
def __init__(self, num_experts:int, dim:int, hidden_dim:int, linear=nn.Linear):
|
||||
self.gate = nn.Linear(dim, num_experts, bias=False)
|
||||
self.experts = [FeedForward(dim, hidden_dim, linear) for _ in range(num_experts)]
|
||||
def __call__(self, x:Tensor) -> Tensor:
|
||||
assert x.shape[0] == 1, "only BS=1"
|
||||
g = self.gate(x).float().exp()
|
||||
choice = g.data().tolist()[0][0]
|
||||
top = sorted(enumerate(choice), key=lambda x: -x[1])
|
||||
norm = top[0][1] + top[1][1]
|
||||
e1, e2 = self.experts[top[0][0]], self.experts[top[1][0]]
|
||||
scale = Tensor([top[0][1]/norm, top[1][1]/norm])
|
||||
ret = e1(x.to(e1.w1.weight.device)).to(x.device) * scale[0] + \
|
||||
e2(x.to(e2.w1.weight.device)).to(x.device) * scale[1]
|
||||
return ret
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run Mixtral in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument("--count", type=int, default=30, help="Max number of tokens to generate")
|
||||
parser.add_argument("--temperature", type=float, default=0.7, help="Temperature in the softmax")
|
||||
parser.add_argument("--timing", action="store_true", help="Print timing per token")
|
||||
parser.add_argument("--profile", action="store_true", help="Profile generation")
|
||||
parser.add_argument("--weights", type=str, default=(pathlib.Path(__file__).parent.parent / "weights/mixtral-8x7b-32kseqlen").as_posix(),
|
||||
help="Path to the downloaded weights")
|
||||
args = parser.parse_args()
|
||||
|
||||
with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
|
||||
state = torch_load(args.weights + "/consolidated.00.pth.b")
|
||||
model = Transformer(n_layers=32, dim=4096, hidden_dim=14336, n_heads=32, n_kv_heads=8, norm_eps=1e-5, vocab_size=32000, feed_forward=functools.partial(MixtureFeedForward, 8), jit=False)
|
||||
model_state_dict = get_state_dict(model)
|
||||
|
||||
for k in (t := tqdm(state, disable=None)):
|
||||
if 'feed_forward.experts.' in k:
|
||||
expert_no = int(k.split('feed_forward.experts.')[1].split('.')[0])
|
||||
device = Device.DEFAULT + ":" + str((expert_no//2)+1)
|
||||
else:
|
||||
device = Device.DEFAULT
|
||||
t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, loading {k} to {device}")
|
||||
model_state_dict[k].replace(state[k].to(device).half()).realize()
|
||||
if t.disable: print(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB")
|
||||
|
||||
from sentencepiece import SentencePieceProcessor
|
||||
spp = SentencePieceProcessor(model_file=args.weights + "/tokenizer.model")
|
||||
|
||||
toks = [spp.bos_id()]
|
||||
start_pos = 0
|
||||
for i in range(args.count):
|
||||
GlobalCounters.reset()
|
||||
with Profiling(sort="time", frac=0.1, enabled=args.profile):
|
||||
with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/sec"):
|
||||
with WallTimeEvent(BenchEvent.STEP):
|
||||
tok = model(Tensor([toks[start_pos:]]), 0 if start_pos == 0 else Variable("start_pos", 1, 1024-1).bind(start_pos), args.temperature).item()
|
||||
toks.append(tok)
|
||||
start_pos += 1
|
||||
print(spp.decode(toks))
|
||||
19
artifacts/package_sources/tinygrad/examples/mlperf/README
Normal file
19
artifacts/package_sources/tinygrad/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
artifacts/package_sources/tinygrad/examples/mlperf/dataloader.py
Normal file
822
artifacts/package_sources/tinygrad/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
|
||||
|
||||
import numpy as np
|
||||
from tinygrad import dtypes, Tensor
|
||||
from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
|
||||
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(NUM_CPU_THREADS.value):
|
||||
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(NUM_CPU_THREADS.value, 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(NUM_CPU_THREADS.value):
|
||||
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(NUM_CPU_THREADS.value):
|
||||
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
artifacts/package_sources/tinygrad/examples/mlperf/helpers.py
Normal file
367
artifacts/package_sources/tinygrad/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
|
||||
@@ -0,0 +1,189 @@
|
||||
import math
|
||||
from typing import Union
|
||||
|
||||
from tinygrad import Tensor, nn, dtypes
|
||||
from tinygrad.helpers import prod, argfix, Context, TRAINING
|
||||
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).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 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
artifacts/package_sources/tinygrad/examples/mlperf/losses.py
Normal file
29
artifacts/package_sources/tinygrad/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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
507
artifacts/package_sources/tinygrad/examples/mlperf/model_eval.py
Normal file
507
artifacts/package_sources/tinygrad/examples/mlperf/model_eval.py
Normal file
@@ -0,0 +1,507 @@
|
||||
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
|
||||
|
||||
@Context(TRAINING=0)
|
||||
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
|
||||
models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
|
||||
with Context(TRAINING=0):
|
||||
for m in models:
|
||||
nm = f"eval_{m}"
|
||||
if nm in globals():
|
||||
print(f"eval {m}")
|
||||
globals()[nm]()
|
||||
@@ -0,0 +1,68 @@
|
||||
# load each model here, quick benchmark
|
||||
from tinygrad import Tensor, GlobalCounters
|
||||
from tinygrad.helpers import getenv, Context
|
||||
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
|
||||
with Context(TRAINING=0):
|
||||
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]()
|
||||
|
||||
2098
artifacts/package_sources/tinygrad/examples/mlperf/model_train.py
Normal file
2098
artifacts/package_sources/tinygrad/examples/mlperf/model_train.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,453 @@
|
||||
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)
|
||||
COLUMNWISE_WEIGHT_SCALE = getenv("COLUMNWISE_WEIGHT_SCALE", 0)
|
||||
MXFP8 = getenv("MXFP8", 0)
|
||||
MXFP4 = getenv("MXFP4", 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, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None, x_prequant_mx:tuple|None=None,
|
||||
next_amax_x: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,)
|
||||
if MXFP4:
|
||||
assert x is not None, "MXFP4 matmul requires an unquantized input"
|
||||
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm
|
||||
if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T, mxfp4=True),)
|
||||
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 MXFP8:
|
||||
from extra.gemm.cdna_asm_gemm import asm_gemm, quantize_mxfp8, mx_pack, can_use_asm_gemm, _mx_block_scale
|
||||
if x_prequant_mx is not None: x_q, x_e8, x_si = x_prequant_mx # fused producer already quantized (2d)
|
||||
else: x_q, x_e8, x_si = quantize_mxfp8(x.reshape(-1, x.shape[-1]))
|
||||
l_shape = x.shape[:-1] if x is not None else x_q.shape[:-1]
|
||||
if can_use_asm_gemm(x_q, w.T):
|
||||
out = asm_gemm(x_q, w.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(w_inv_scale), w_inv_scale),
|
||||
mx_w_stored=True).reshape(*l_shape, w.shape[0])
|
||||
else:
|
||||
x_phys = (x_q.cast(dtypes.bfloat16) * _mx_block_scale(x_e8)).reshape(*l_shape, x_q.shape[-1])
|
||||
out = x_phys @ (w.cast(dtypes.bfloat16) * _mx_block_scale(w_inv_scale)).T
|
||||
return out, x_q
|
||||
if x_fp8 is None:
|
||||
if FUSED_INPUT_QUANTIZE:
|
||||
from extra.llama_kernels.quantize_fp8_delayed import quantize_fp8_delayed
|
||||
x_fp8, _ = quantize_fp8_delayed(x, amax_x, next_amax_x, FP8_DTYPE)
|
||||
else:
|
||||
x_fp8, _, new_amax_x = quantize_fp8(x, amax_state=amax_x)
|
||||
next_amax_x.assign(new_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):
|
||||
assert amax_x is not None
|
||||
if COLUMNWISE_WEIGHT_SCALE:
|
||||
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, grad_amax_state=grad_amax_state,
|
||||
next_grad_amax_state=next_grad_amax_state, w_post_scale=w_inv_scale)
|
||||
else:
|
||||
out = asm_gemm(x_fp8, w.T, x_scale=amax_x, w_scale=w_inv_scale, grad_amax_state=grad_amax_state,
|
||||
next_grad_amax_state=next_grad_amax_state)
|
||||
return out, x_fp8
|
||||
return (x_fp8.dot(w.T, dtype=dtypes.float) * ((amax_x.float() + 1e-8) / FP8_MAX) * w_inv_scale).cast(dtypes.bfloat16), x_fp8
|
||||
|
||||
def norm_quantize_matmul(x:Tensor, norm:Tensor, w:Tensor, w_inv_scale:Tensor, eps:float, amax_x:Tensor|None,
|
||||
next_amax_x:Tensor|None, grad_amax_state:Tensor|None, next_grad_amax_state:Tensor|None):
|
||||
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
|
||||
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_rmsnorm_mul_quantize_fp8
|
||||
x_fp8, x_normed, rrms = fused_rmsnorm_mul_quantize_fp8(x, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
|
||||
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
|
||||
grad_amax_state=grad_amax_state, next_grad_amax_state=next_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,
|
||||
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
|
||||
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|None,
|
||||
next_amax_x:Tensor|None, grad_amax_state:Tensor|None=None, next_grad_amax_state:Tensor|None=None):
|
||||
if FUSED_ADD_NORM_MUL_QUANTIZE and not MXFP4:
|
||||
from extra.llama_kernels.fused_rmsnorm_mul_quantize_fp8 import fused_add_rmsnorm_mul_quantize_fp8
|
||||
x_fp8, h, x_normed, rrms = fused_add_rmsnorm_mul_quantize_fp8(x, residual, norm, amax_x, eps, FP8_DTYPE, next_amax_x)
|
||||
out, *ret = matmul(None, w, w_inv_scale=w_inv_scale, x_fp8=x_fp8, amax_x=amax_x,
|
||||
grad_amax_state=grad_amax_state, next_grad_amax_state=next_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,
|
||||
next_grad_amax_state=next_grad_amax_state, next_amax_x=next_amax_x)
|
||||
return out, h, x_normed, rrms, ret
|
||||
|
||||
def silu_w13_quantize_matmul(x_w13:Tensor, w2:Tensor, s_2:Tensor,
|
||||
amax_x2:Tensor|None, next_amax_x2:Tensor|None,
|
||||
grad_amax_xw13:Tensor|None, next_grad_amax_xw13:Tensor|None,
|
||||
grad_amax_xout:Tensor|None, next_grad_amax_xout:Tensor|None):
|
||||
if FUSED_SILU_W13 and not MXFP4:
|
||||
from extra.llama_kernels.cast_amax import fused_quantize_fp8_w13
|
||||
x2_fp8 = fused_quantize_fp8_w13(x_w13, amax_x2, FP8_DTYPE, grad_amax_state=grad_amax_xw13,
|
||||
next_grad_amax_state=next_grad_amax_xw13, amax_out=next_amax_x2)
|
||||
out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, amax_x=amax_x2,
|
||||
grad_amax_state=grad_amax_xout, next_grad_amax_state=next_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,
|
||||
next_grad_amax_state=next_grad_amax_xout, next_amax_x=next_amax_x2)
|
||||
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).clone().is_param_(False)
|
||||
|
||||
def _amax(): return Tensor.full((), FP8_MAX, dtype=dtypes.float32).contiguous().is_param_(False)
|
||||
n_amax = 0 if MXFP4 else n_layers
|
||||
names = ["xqkv", "xo", "x2"]
|
||||
names += ["x1", "x3"] if SPLIT_W13 else ["x13"]
|
||||
self._fp8_amax = {name: [_amax() for _ in range(n_amax)] for name in names}
|
||||
self._fp8_next_amax = {name: [_amax() for _ in range(n_amax)] 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_amax)] for name in grad_names}
|
||||
self._fp8_next_grad_amax = {name: [_amax() for _ in range(n_amax)] 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 if MXFP8 else s.float()).contiguous().is_param_(False) for name, s in w_scales}
|
||||
self._fp8_next_inv_scale = {name: (s if MXFP8 else 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, w:Tensor|None=None):
|
||||
if w is None:
|
||||
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)
|
||||
if MXFP8:
|
||||
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
|
||||
w_q, w_e8, _ = quantize_mxfp8(w.reshape(self.n_layers * out_features, in_features))
|
||||
return w_q.reshape(self.n_layers, out_features, in_features), w_e8.reshape(self.n_layers, out_features, in_features // 32)
|
||||
if MXFP4:
|
||||
# FP4 is produced dynamically so optimizer updates always start from the current BF16 weight.
|
||||
return w.cast(dtypes.bfloat16), Tensor.ones(self.n_layers)
|
||||
amax = (w.abs().max(axis=2) if COLUMNWISE_WEIGHT_SCALE else w.abs().flatten(1).max(1)).detach()
|
||||
scale = FP8_MAX / (amax + 1e-8)
|
||||
inv_scale = (amax + 1e-8) / FP8_MAX
|
||||
scale_b = scale.reshape(self.n_layers, out_features, 1) if COLUMNWISE_WEIGHT_SCALE else scale.reshape(-1, 1, 1)
|
||||
return (w * scale_b).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|None, amax_xo:Tensor|None, s_qkv:Tensor, s_o:Tensor,
|
||||
next_amax_xqkv:Tensor|None, next_amax_xo:Tensor|None,
|
||||
grad_amax_xqkv:Tensor|None, grad_amax_xo:Tensor|None,
|
||||
next_grad_amax_xqkv:Tensor|None, next_grad_amax_xo:Tensor|None):
|
||||
bsz, seqlen, _ = x.shape
|
||||
saves = []
|
||||
|
||||
xqkv, x_normed, rrms, s = norm_quantize_matmul(x, attention_norm, wqkv, s_qkv, self.norm_eps,
|
||||
amax_x=amax_xqkv, grad_amax_state=grad_amax_xqkv,
|
||||
next_grad_amax_state=next_grad_amax_xqkv, next_amax_x=next_amax_xqkv)
|
||||
saves.extend([x_normed, rrms, *s, xqkv])
|
||||
if getenv("HK_FLASH_ATTENTION"):
|
||||
from extra.thunder.amd.fa import flash_attention, fused_qkv_rope
|
||||
xq, xk, xv = fused_qkv_rope(xqkv, freqs_cis, self.n_heads, self.n_kv_heads, self.head_dim)
|
||||
attn, *save = flash_attention(xq, xk, xv, is_causal=True, write_flat=True)
|
||||
saves.extend(save)
|
||||
else:
|
||||
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)
|
||||
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, *s = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o, grad_amax_state=grad_amax_xo,
|
||||
next_grad_amax_state=next_grad_amax_xo, next_amax_x=next_amax_xo)
|
||||
saves.extend([*s, out])
|
||||
return out, saves
|
||||
|
||||
def feed_forward(self, x:Tensor, residual:Tensor, **kwargs):
|
||||
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, *s = matmul(inp, kwargs["w1"], amax_x=kwargs["amax_x1"], w_inv_scale=kwargs["s_1"],
|
||||
grad_amax_state=kwargs["grad_amax_xw1"], next_grad_amax_state=kwargs["next_grad_amax_xw1"],
|
||||
next_amax_x=kwargs["next_amax_x1"])
|
||||
saves.extend([*s, x_w1])
|
||||
x_w3, *s = matmul(inp, kwargs["w3"], amax_x=kwargs["amax_x3"], w_inv_scale=kwargs["s_3"],
|
||||
grad_amax_state=kwargs["grad_amax_xw3"], next_grad_amax_state=kwargs["next_grad_amax_xw3"],
|
||||
next_amax_x=kwargs["next_amax_x3"])
|
||||
saves.extend([*s, x_w3])
|
||||
if FUSED_SILU_W13 and MXFP8:
|
||||
from extra.llama_kernels.fused_silu_mul_quantize_mxfp8 import fused_silu_mul_quantize_mxfp8
|
||||
aq, ae8, asi = fused_silu_mul_quantize_mxfp8(x_w1.reshape(-1, x_w1.shape[-1]), x_w3.reshape(-1, x_w3.shape[-1]))
|
||||
out, *s = matmul(None, kwargs["w2"], x_prequant_mx=(aq, ae8, asi), amax_x=kwargs["amax_x2"],
|
||||
w_inv_scale=kwargs["s_2"], grad_amax_state=kwargs["grad_amax_xout"],
|
||||
next_grad_amax_state=kwargs["next_grad_amax_xout"], next_amax_x=kwargs["next_amax_x2"])
|
||||
out = out.reshape(*x_w1.shape[:-1], kwargs["w2"].shape[0])
|
||||
else:
|
||||
out, *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"], next_grad_amax_state=kwargs["next_grad_amax_xout"],
|
||||
next_amax_x=kwargs["next_amax_x2"])
|
||||
saves.extend([*s, out])
|
||||
else:
|
||||
x_w13, h, x_normed, rrms, s = add_norm_quantize_matmul(x, residual, kwargs["ffn_norm"], kwargs["w13"], kwargs["s_13"],
|
||||
self.norm_eps, amax_x=kwargs["amax_x13"],
|
||||
next_amax_x=kwargs["next_amax_x13"],
|
||||
grad_amax_state=kwargs["grad_amax_xw13"],
|
||||
next_grad_amax_state=kwargs["next_grad_amax_xw13"])
|
||||
saves.extend([x_normed, rrms, *s, x_w13])
|
||||
out, s = silu_w13_quantize_matmul(x_w13, kwargs["w2"], kwargs["s_2"], amax_x2=kwargs["amax_x2"],
|
||||
next_amax_x2=kwargs["next_amax_x2"],
|
||||
grad_amax_xw13=kwargs["grad_amax_xw13"],
|
||||
next_grad_amax_xw13=kwargs["next_grad_amax_xw13"],
|
||||
grad_amax_xout=kwargs["grad_amax_xout"],
|
||||
next_grad_amax_xout=kwargs["next_grad_amax_xout"])
|
||||
saves.extend([*s, out])
|
||||
return out, h, 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_saves = self.attention(x, freqs_cis, **attn_kwargs)
|
||||
ffn, h, ffn_saves = self.feed_forward(x, attn, **ffn_kwargs)
|
||||
h = h + ffn
|
||||
if save: return (h, *attn_saves, *ffn_saves)
|
||||
else: return (h,)
|
||||
|
||||
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, std:float=0.02):
|
||||
w = getattr(self, name)
|
||||
if MXFP8:
|
||||
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
|
||||
w_bf16 = Tensor.empty(self.n_layers, w.shape[1], w.shape[2], dtype=dtypes.bfloat16).shard(device, axis=axis).randn_like() * std
|
||||
w_q, w_e8, _ = quantize_mxfp8(w_bf16)
|
||||
w.replace(w_q)
|
||||
self._fp8_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
|
||||
self._fp8_next_inv_scale[name].replace(w_e8.contiguous()).is_param_(False)
|
||||
else:
|
||||
w.shard_(device, axis=axis)
|
||||
scale_axis = (1 if axis == 1 else None) if COLUMNWISE_WEIGHT_SCALE else None
|
||||
self._fp8_inv_scale[name] = self._fp8_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
|
||||
self._fp8_next_inv_scale[name] = self._fp8_next_inv_scale[name].shard(device, axis=scale_axis).contiguous().is_param_(False)
|
||||
Tensor.realize(w, self._fp8_inv_scale[name], self._fp8_next_inv_scale[name])
|
||||
sstd = 0.02 / math.sqrt(2 * self.n_layers)
|
||||
_shard_fp8("wqkv", 1) # (n_layers, out, dim) shard out
|
||||
_shard_fp8("wo", 2, sstd) # (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, sstd) # (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_next_amax, self._fp8_grad_amax, self._fp8_next_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 reset_amax(self):
|
||||
for st in (self._fp8_next_amax, self._fp8_next_grad_amax):
|
||||
for ts in st.values():
|
||||
for t in ts: t.assign(0)
|
||||
|
||||
def update_amax(self):
|
||||
for cur, nxt in ((self._fp8_amax, self._fp8_next_amax), (self._fp8_grad_amax, self._fp8_next_grad_amax)):
|
||||
for name in cur:
|
||||
for c, n in zip(cur[name], nxt[name]): c.assign(n)
|
||||
|
||||
def __call__(self, tokens:Tensor, save:bool=True):
|
||||
h = self.tok_embeddings(tokens)
|
||||
freqs_cis = self.freqs_cis.cast(h.dtype)
|
||||
if not getenv("HK_FLASH_ATTENTION"): freqs_cis = freqs_cis[:, :tokens.shape[1], :, :, :]
|
||||
a, na, ga, nga, s = self._fp8_amax, self._fp8_next_amax, self._fp8_grad_amax, self._fp8_next_grad_amax, self._fp8_inv_scale
|
||||
def amax_kwargs(i:int, act_names:tuple[str, ...], grad_names:tuple[str, ...]) -> dict[str, Tensor|None]:
|
||||
specs = (("amax_", a, act_names), ("next_amax_", na, act_names), ("grad_amax_", ga, grad_names), ("next_grad_amax_", nga, grad_names))
|
||||
if MXFP4: return dict.fromkeys(f"{prefix}{name}" for prefix, _, names in specs for name in names)
|
||||
return {f"{prefix}{name}":val[name][i] for prefix, val, names in specs for name in names}
|
||||
for i in range(self.n_layers):
|
||||
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wo=self.wo[i], s_qkv=s["wqkv"][i], s_o=s["wo"][i],
|
||||
**amax_kwargs(i, ("xqkv", "xo"), ("xqkv", "xo")))
|
||||
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], w2=self.w2[i], s_2=s["w2"][i], **amax_kwargs(i, ("x2",), ("xout",)))
|
||||
if SPLIT_W13:
|
||||
ffn_kwargs.update(w1=self.w1[i], w3=self.w3[i], s_1=s["w1"][i], s_3=s["w3"][i], **amax_kwargs(i, ("x1", "x3"), ("xw1", "xw3")))
|
||||
else:
|
||||
ffn_kwargs.update(w13=self.w13[i], s_13=s["w13"][i], **amax_kwargs(i, ("x13",), ("xw13",)))
|
||||
h, *_ = self.run_layer(h, freqs_cis, attn_kwargs, ffn_kwargs, save=save)
|
||||
|
||||
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: "):
|
||||
model.reset_amax()
|
||||
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())))
|
||||
@@ -0,0 +1,351 @@
|
||||
import math, os, functools
|
||||
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"
|
||||
from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
|
||||
from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
|
||||
from tinygrad.uop.ops import Ops, UOp
|
||||
from extra.models.llama import apply_rotary_emb
|
||||
from extra.llama_kernels.rmsnorm import rmsnorm
|
||||
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
|
||||
from extra.gemm.moe_gemm import grouped_mx_gemm
|
||||
from extra.gemm.moe_routing import route, dispatch, combine
|
||||
|
||||
FP8_DTYPE = dtypes.fp8e4m3
|
||||
FP8_MAX = 448.0
|
||||
INIT_STD = 0.02
|
||||
ASM_GEMM = getenv("ASM_GEMM", 0)
|
||||
|
||||
|
||||
def _quant_dequant_fwd(x:Tensor) -> Tensor:
|
||||
# x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
|
||||
M, K = x.shape
|
||||
scale_K = K // 32
|
||||
amax = x.float().reshape(M, scale_K, 32).abs().max(axis=-1)
|
||||
e8 = (amax.maximum(1e-38).log2().floor() + 127).clamp(0, 254).cast(dtypes.uint8)
|
||||
qscale = (127.0 - e8.cast(dtypes.float32)).exp2().reshape(M, scale_K, 1).expand(M, scale_K, 32).reshape(M, K)
|
||||
x_fp8 = (x.float() * qscale).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE).cast(dtypes.float32)
|
||||
return (x_fp8 * _mx_block_scale(e8)).cast(dtypes.bfloat16)
|
||||
|
||||
@functools.cache
|
||||
def _quant_dequant_fwd_fxn(x_p, device):
|
||||
return _quant_dequant_fwd(Tensor(x_p, device=device))
|
||||
|
||||
def _quant_dequant_bwd(grad:UOp, call:UOp) -> tuple:
|
||||
return (Tensor(grad).cast(dtypes.bfloat16).uop,)
|
||||
|
||||
def quant_dequant_mx(x:Tensor) -> Tensor:
|
||||
fxn = _quant_dequant_fwd_fxn(x.as_param(0).uop, x.device)
|
||||
return Tensor(UOp.maketuple(fxn.uop).call(x.uop, grad_fxn=_quant_dequant_bwd).gettuple(0))
|
||||
|
||||
def _mx_scale(e8:Tensor) -> Tensor:
|
||||
return _mx_block_scale(e8) if e8.ndim == 2 else _mx_block_scale_3d(e8)
|
||||
|
||||
def _dequant_fwd(w_q:Tensor, w_scale:Tensor) -> Tensor:
|
||||
return w_q.cast(dtypes.bfloat16) * _mx_scale(w_scale)
|
||||
|
||||
@functools.cache
|
||||
def _dequant_fwd_fxn(wq_p, ws_p, device):
|
||||
return _dequant_fwd(Tensor(wq_p, device=device), Tensor(ws_p, device=device))
|
||||
|
||||
def _dequant_bwd(grad:UOp, call:UOp) -> tuple:
|
||||
return (Tensor(grad).cast(dtypes.bfloat16).uop, None)
|
||||
|
||||
def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
|
||||
fxn = _dequant_fwd_fxn(w_q.as_param(0).uop, w_scale.as_param(1).uop, w_q.device)
|
||||
call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
|
||||
return Tensor(call.gettuple(0))
|
||||
|
||||
def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
|
||||
l_shape = x.shape[:-1]
|
||||
if ASM_GEMM:
|
||||
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
|
||||
x2, K, N = x.reshape(-1, x.shape[-1]), x.shape[-1], w_q.shape[0]
|
||||
wq, ws = w_q, w_scale
|
||||
if (pad := (-K) % 256):
|
||||
x2 = x2.pad(((0, 0), (0, pad)))
|
||||
wq = wq.pad(((0, 0), (0, pad)))
|
||||
ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
|
||||
if (npad := (-N) % 256):
|
||||
wq = wq.pad(((0, npad), (0, 0)))
|
||||
ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
|
||||
x_q, x_e8, x_si = quantize_mxfp8(x2)
|
||||
if x_si is not None and can_use_asm_gemm(x_q, wq.T):
|
||||
out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(x_si, x_e8, mx_pack(ws), ws), mx_w_stored=True)
|
||||
return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
|
||||
x_phys = quant_dequant_mx(x.reshape(-1, x.shape[-1])).reshape(*l_shape, x.shape[-1])
|
||||
w_phys = dequant_weight(w_q, w_scale)
|
||||
return (x_phys @ w_phys.T).cast(dtypes.bfloat16)
|
||||
|
||||
def _pad_to_mult(t:Tensor, axis:int, mult:int=256) -> Tensor:
|
||||
if (r := (-t.shape[axis]) % mult) == 0: return t
|
||||
pads = [(0, 0)] * t.ndim
|
||||
pads[axis] = (0, r)
|
||||
return t.pad(tuple(pads))
|
||||
|
||||
def _pad_cols(t:Tensor) -> Tensor: return _pad_to_mult(t, -1)
|
||||
def _pad_rows(t:Tensor) -> Tensor: return _pad_to_mult(t, -2)
|
||||
|
||||
def swiglu(x:Tensor, limit:float=7.0, alpha:float=1.702) -> Tensor:
|
||||
x_glu, x_linear = x[..., ::2], x[..., 1::2]
|
||||
x_glu = x_glu.clamp(max_=limit)
|
||||
x_linear = x_linear.clamp(-limit, limit)
|
||||
return (x_glu * (alpha * x_glu).sigmoid()) * (x_linear + 1)
|
||||
|
||||
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
|
||||
freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
|
||||
freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
|
||||
return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
|
||||
|
||||
class GPTOSS:
|
||||
def __init__(self, dim:int, n_layers:int, n_heads:int, n_kv_heads:int, head_dim:int, n_experts:int, experts_per_tok:int,
|
||||
intermediate_size:int, vocab_size:int, norm_eps:float=1e-5, rope_theta:int=150000, sliding_window:int=128,
|
||||
swiglu_limit:float=7.0, max_context:int=8192):
|
||||
self.dim, self.n_layers, self.n_heads, self.n_kv_heads, self.head_dim = dim, n_layers, n_heads, n_kv_heads, head_dim
|
||||
self.n_rep = n_heads // n_kv_heads
|
||||
self.n_experts, self.experts_per_tok, self.intermediate_size = n_experts, experts_per_tok, intermediate_size
|
||||
self.vocab_size, self.norm_eps, self.sliding_window, self.swiglu_limit = vocab_size, norm_eps, sliding_window, swiglu_limit
|
||||
self.sm_scale = 1.0 / math.sqrt(head_dim)
|
||||
|
||||
scaled_std = INIT_STD / math.sqrt(2 * n_layers)
|
||||
q_dim, qkv_dim = n_heads * head_dim, head_dim * (n_heads + 2 * n_kv_heads)
|
||||
|
||||
# attn
|
||||
self.wqkv, self.wqkv_scale = self._quant_weight(n_layers, qkv_dim, dim)
|
||||
self.wqkv_bias = Tensor.zeros(n_layers, qkv_dim, dtype=dtypes.bfloat16).contiguous()
|
||||
self.wo, self.wo_scale = self._quant_weight(n_layers, dim, q_dim, std=scaled_std)
|
||||
self.wo_bias = Tensor.zeros(n_layers, dim, dtype=dtypes.bfloat16).contiguous()
|
||||
self.sinks = Tensor.zeros(n_layers, n_heads, dtype=dtypes.bfloat16).contiguous()
|
||||
self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
|
||||
|
||||
# moe ffn
|
||||
self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
|
||||
self.gate = Tensor.normal(n_layers, n_experts, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
|
||||
self.gate_bias = Tensor.zeros(n_layers, n_experts, dtype=dtypes.bfloat16).contiguous()
|
||||
self.w_gate_up, self.w_gate_up_scale = self._quant_weight(n_layers, n_experts, intermediate_size * 2, dim, moe=True)
|
||||
self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
|
||||
self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
|
||||
self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).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=INIT_STD, dtype=dtypes.bfloat16)
|
||||
self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
|
||||
self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
|
||||
|
||||
def _quant_weight(self, *shape:int, std:float=INIT_STD, moe:bool=False):
|
||||
def _one(*s:int):
|
||||
w = Tensor.zeros(*s) if getenv("ZEROS") else Tensor.normal(*s, mean=0.0, std=std)
|
||||
w_q, w_e8, _ = quantize_mxfp8(_pad_cols(_pad_rows(w)) if moe else w)
|
||||
return w_q, w_e8.is_param_(False)
|
||||
if moe:
|
||||
qs = [_one(*shape[1:]) for _ in range(shape[0])]
|
||||
return [q[0] for q in qs], [q[1] for q in qs]
|
||||
return _one(*shape)
|
||||
|
||||
def _attn_mask(self, seqlen:int, dtype) -> Tensor:
|
||||
i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
|
||||
return (j <= i).where(0.0, -1e30).cast(dtype).contiguous()
|
||||
|
||||
def _sliding_attention(self, xq:Tensor, xk:Tensor, xv:Tensor, sinks:Tensor) -> Tensor:
|
||||
bsz, seqlen, H, hd = xq.shape
|
||||
KV, R, W = self.n_kv_heads, self.n_rep, self.sliding_window
|
||||
assert seqlen % W == 0, f"seqlen {seqlen} must be a multiple of sliding_window {W} for banded attention"
|
||||
nb = seqlen // W
|
||||
q = xq.reshape(bsz, seqlen, KV, R, hd).permute(0, 2, 3, 1, 4).reshape(bsz, KV, R, nb, W, hd).float()
|
||||
k, v = (x.permute(0, 2, 1, 3).reshape(bsz, KV, 1, nb, W, hd).float() for x in (xk, xv))
|
||||
kk, vv = (x.pad((None, None, None, (1, 0), None, None))[:, :, :, :nb].cat(x, dim=-2) for x in (k, v))
|
||||
sc = (q @ kk.transpose(-1, -2)) * self.sm_scale # (B,KV,R,nb,W,2W)
|
||||
i, j, pv = Tensor.arange(W).reshape(W, 1), Tensor.arange(2 * W).reshape(1, 2 * W), Tensor.arange(nb).reshape(nb, 1, 1) >= 1
|
||||
sc = ((j > i) & (j <= i + W) & (pv | (j >= W))).where(sc, -float("inf"))
|
||||
sink = sinks.reshape(1, KV, R, 1, 1, 1).float()
|
||||
m = sc.max(-1, keepdim=True).maximum(sink)
|
||||
e = (sc - m).exp()
|
||||
p = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
|
||||
attn = p @ vv.cast(dtypes.bfloat16)
|
||||
return attn.reshape(bsz, KV, R, seqlen, hd).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, H * hd)
|
||||
|
||||
def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
|
||||
wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
|
||||
bsz, seqlen, _ = x.shape
|
||||
x_normed, rrms = rmsnorm(x, self.norm_eps)
|
||||
qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
|
||||
qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
|
||||
xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
|
||||
xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
|
||||
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
|
||||
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
|
||||
|
||||
if sliding:
|
||||
attn = self._sliding_attention(xq, xk, xv, sinks)
|
||||
elif getenv("HK_FLASH_ATTENTION"):
|
||||
from extra.thunder.amd.fa import flash_attention
|
||||
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
|
||||
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
|
||||
else:
|
||||
xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
|
||||
xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
|
||||
scores = (xqm @ xkm.transpose(-2, -1)).float() * self.sm_scale + mask
|
||||
sink = sinks.reshape(1, self.n_kv_heads, self.n_rep, 1, 1).float()
|
||||
m = scores.max(-1, keepdim=True).maximum(sink)
|
||||
e = (scores - m).exp()
|
||||
w = (e / (e.sum(-1, keepdim=True) + (sink - m).exp())).cast(dtypes.bfloat16)
|
||||
attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
|
||||
|
||||
out = matmul_mx(attn, wo, wo_scale) + wo_bias
|
||||
return out, [x_normed, rrms, attn]
|
||||
|
||||
def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
|
||||
w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
|
||||
w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
|
||||
x_normed, rrms = rmsnorm(x, self.norm_eps)
|
||||
inp = x_normed * ffn_norm
|
||||
logits = inp.float() @ gate.float().T + gate_bias.float()
|
||||
dim, inter = self.dim, self.intermediate_size
|
||||
|
||||
if getenv("GROUPED_MOE", 0):
|
||||
bsz, seqlen = x.shape[:2]
|
||||
inp, logits = inp.reshape(-1, dim), logits.reshape(-1, self.n_experts)
|
||||
r = route(logits, self.experts_per_tok, self.n_experts)
|
||||
onehot = r.rows_e.one_hot(self.n_experts).float()
|
||||
xg = dispatch(_pad_cols(inp.cast(dtypes.bfloat16)), r)
|
||||
h = grouped_mx_gemm(xg, (w_gate_up, w_gate_up_scale), r.off)[:, :2*inter] + (onehot @ w_gate_up_bias.float()).cast(dtypes.bfloat16)
|
||||
y = swiglu(h, self.swiglu_limit)
|
||||
z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
|
||||
+ (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
|
||||
out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
|
||||
else:
|
||||
thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
|
||||
weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
|
||||
|
||||
out = None
|
||||
for e in range(self.n_experts):
|
||||
gu_q, gu_s = w_gate_up[e][:2*inter, :dim].contiguous(), w_gate_up_scale[e][:2*inter, :dim//32].contiguous()
|
||||
dn_q, dn_s = w_down[e][:dim, :inter].contiguous(), w_down_scale[e][:dim, :inter//32].contiguous()
|
||||
gate_up = matmul_mx(inp, gu_q, gu_s) + w_gate_up_bias[e]
|
||||
y = (matmul_mx(swiglu(gate_up, self.swiglu_limit), dn_q, dn_s) + w_down_bias[e]).contiguous()
|
||||
contrib = weights[..., e:e+1].cast(y.dtype) * y
|
||||
out = contrib if out is None else out + contrib
|
||||
return out, [x_normed, rrms]
|
||||
|
||||
@function(precompile=True, precompile_backward=True)
|
||||
def run_layer(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, attn_kwargs:dict, ffn_kwargs:dict, save:bool=True):
|
||||
attn, attn_saves = self.attention(x, freqs_cis, mask, sliding, **attn_kwargs)
|
||||
h = x + attn
|
||||
ffn, ffn_saves = self.feed_forward(h, **ffn_kwargs)
|
||||
h = h + ffn
|
||||
if save: return (h, *attn_saves, *ffn_saves)
|
||||
return (h,)
|
||||
|
||||
def shard(self, device:tuple[str, ...], mp:bool=False):
|
||||
assert not mp, "MP not supported"
|
||||
from tinygrad.nn.state import get_parameters
|
||||
for v in get_parameters(self): v.shard_(device, axis=None)
|
||||
Tensor.realize(*get_parameters(self))
|
||||
|
||||
def __call__(self, tokens:Tensor, save:bool=True):
|
||||
h = self.tok_embeddings(tokens)
|
||||
bsz, seqlen = tokens.shape
|
||||
freqs_cis = self.freqs_cis.cast(h.dtype)[:, :seqlen, :, :, :]
|
||||
mask_full = None if getenv("HK_FLASH_ATTENTION") else self._attn_mask(seqlen, dtypes.float32)
|
||||
for i in range(self.n_layers):
|
||||
attn_kwargs = dict(attention_norm=self.attention_norm[i], wqkv=self.wqkv[i], wqkv_scale=self.wqkv_scale[i],
|
||||
wqkv_bias=self.wqkv_bias[i], wo=self.wo[i], wo_scale=self.wo_scale[i], wo_bias=self.wo_bias[i],
|
||||
sinks=self.sinks[i])
|
||||
ffn_kwargs = dict(ffn_norm=self.ffn_norm[i], gate=self.gate[i], gate_bias=self.gate_bias[i],
|
||||
w_gate_up=self.w_gate_up[i], w_gate_up_scale=self.w_gate_up_scale[i], w_gate_up_bias=self.w_gate_up_bias[i],
|
||||
w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
|
||||
h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
|
||||
|
||||
logits = self.norm(h) @ self.output.T
|
||||
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
|
||||
|
||||
GPT_OSS_20B = dict(dim=2880, n_layers=24, n_heads=64, n_kv_heads=8, head_dim=64, n_experts=32, experts_per_tok=4,
|
||||
intermediate_size=2880, vocab_size=128256, norm_eps=1e-5, rope_theta=150000, sliding_window=128,
|
||||
swiglu_limit=7.0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
config = {}
|
||||
BS = config["BS"] = getenv("BS", 16)
|
||||
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
|
||||
|
||||
model_params = GPT_OSS_20B
|
||||
real_vocab_size = model_params["vocab_size"]
|
||||
if (layers := getenv("LAYERS")) != 0: model_params["n_layers"] = layers
|
||||
|
||||
model = GPTOSS(**model_params, max_context=SEQLEN)
|
||||
|
||||
state = nn.state.get_state_dict(model)
|
||||
print("tensor count:", len(state))
|
||||
|
||||
from tinygrad import Device
|
||||
is_dp = (DP := getenv("DP", 1)) > 1
|
||||
device_count = DP
|
||||
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
|
||||
|
||||
if is_dp: model.shard(device)
|
||||
|
||||
# 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}
|
||||
|
||||
# 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 is_dp: tokens = tokens.shard(device, axis=0)
|
||||
|
||||
@TinyJit
|
||||
def fwd_bwd(tokens:Tensor):
|
||||
with Timing("python forward: "):
|
||||
logits = model(tokens[:, :-1], save=True)
|
||||
loss = 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())
|
||||
|
||||
@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())))
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
144
artifacts/package_sources/tinygrad/examples/mlperf/optim.py
Normal file
144
artifacts/package_sources/tinygrad/examples/mlperf/optim.py
Normal file
@@ -0,0 +1,144 @@
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.nn.optim import Optimizer, OptimizerGroup
|
||||
from tinygrad.helpers import FUSE_OPTIM, getenv
|
||||
from tinygrad.uop.ops import UOp, Ops, AxisType
|
||||
|
||||
STOCHASTIC_ROUND = getenv("STOCHASTIC_ROUND", 0)
|
||||
MASTER_WEIGHTS = getenv("MASTER_WEIGHTS", 0)
|
||||
ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
|
||||
FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)
|
||||
IMMEDIATE_SCALE = getenv("IMMEDIATE_SCALE", 0)
|
||||
MXFP8 = getenv("MXFP8", 0)
|
||||
|
||||
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)
|
||||
|
||||
def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
|
||||
for g in grads: g.assign(g / grad_acc)
|
||||
total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
|
||||
for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
|
||||
return total_norm
|
||||
|
||||
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.zero = bool(ZERO_OPTIM) and isinstance(self.device, tuple) and not self.fused
|
||||
self.m = [self._zero_shard(x) for x in self._new_optim_param()]
|
||||
self.v = [self._zero_shard(x) for x in 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 = [self._zero_shard(p.to(self.device).float().contiguous()) for p in self.params]
|
||||
else:
|
||||
self.master_params = None
|
||||
|
||||
def _zero_shard(self, t:Tensor) -> Tensor:
|
||||
if not self.zero or t.ndim < 2 or (t.shape[0] % len(self.device)) != 0: return t
|
||||
return Tensor(t.uop._shard(0, UOp.range(len(self.device), -1, AxisType.DEVICE)).unshard(0)).clone()
|
||||
|
||||
def _zero_gather(self, t:Tensor) -> Tensor:
|
||||
if not isinstance(t.device, tuple) or t.uop.axis != 0: return t
|
||||
n, sz = len(t.device), t.shape[0] // len(t.device)
|
||||
return Tensor.cat(*[t[p*sz:(p+1)*sz] for p in range(n)], dim=0)
|
||||
|
||||
def fschedule_step(self, grads:list[Tensor]) -> list[Tensor]:
|
||||
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))
|
||||
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')]
|
||||
return extra + self.params + self.buffers + (self.master_params or []) + fp8_inv_scales + fp8_next_inv_scales
|
||||
|
||||
def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
|
||||
Tensor.realize(*([grad_norm] if grad_norm is not None else []), *self.fschedule_step(grads))
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
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)
|
||||
if self.zero and not (MXFP8 and t.dtype in dtypes.fp8s): new_w = self._zero_gather(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:
|
||||
if MXFP8:
|
||||
from extra.gemm.cdna_asm_gemm import quantize_mxfp8
|
||||
w_q, w_e8, _ = quantize_mxfp8(new_w.reshape(-1, new_w.shape[-1]))
|
||||
if self.zero: w_q, w_e8 = self._zero_gather(w_q), self._zero_gather(w_e8)
|
||||
new_e8 = w_e8.reshape(t._inv_scale.shape)
|
||||
t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
|
||||
ret = w_q.reshape(t.shape)
|
||||
return ret.shard_like(t) if offloaded else ret
|
||||
from examples.mlperf.models.flat_llama import FP8_MAX
|
||||
if IMMEDIATE_SCALE:
|
||||
amax_axis = tuple(range(t._inv_scale.ndim, new_w.ndim))
|
||||
new_inv = ((new_w.float().abs().max(axis=amax_axis).detach() + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype)
|
||||
t._inv_scale.assign(new_inv.shard_like(t._inv_scale) if offloaded else new_inv)
|
||||
scale = new_inv.reciprocal().reshape(*new_inv.shape, *([1]*(new_w.ndim-new_inv.ndim)))
|
||||
ret = (new_w * scale).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
|
||||
return ret.shard_like(t) if offloaded else ret
|
||||
# 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(*inv_scale.shape, *([1]*(new_w.ndim-inv_scale.ndim)))
|
||||
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(inv_scale.ndim, 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
|
||||
|
||||
class GradAccClipAdamWGroup(OptimizerGroup):
|
||||
def __init__(self, *optimizers:GradAccClipAdamW):
|
||||
super().__init__(*optimizers)
|
||||
for o in self.optimizers[1:]: o.lr = self.optimizers[0].lr
|
||||
def fstep(self, grads:list[Tensor], grad_norm:Tensor|None=None):
|
||||
offset = 0
|
||||
to_realize = []
|
||||
for o in self.optimizers:
|
||||
n = len(o.params)
|
||||
to_realize += o.fschedule_step(grads[offset:offset+n])
|
||||
offset += n
|
||||
Tensor.realize(*to_realize, *([grad_norm] if grad_norm is not None else []))
|
||||
@property
|
||||
def lr(self): return self.optimizers[0].lr
|
||||
@property
|
||||
def device(self): return self.optimizers[0].device
|
||||
@property
|
||||
def master_params(self):
|
||||
mp = [mp for o in self.optimizers for mp in (o.master_params or [])]
|
||||
return mp if mp else None
|
||||
@@ -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
|
||||
@@ -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
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user