mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
187 lines
8.0 KiB
Python
187 lines
8.0 KiB
Python
import unittest, time, gc
|
|
import numpy as np
|
|
from tinygrad.nn import optim
|
|
from tinygrad.nn.state import get_parameters
|
|
from tinygrad.engine.jit import TinyJit
|
|
from tinygrad import Tensor, Device, GlobalCounters, dtypes, Variable
|
|
from tinygrad.helpers import Context
|
|
from test.helpers import slow, jit_cache_count, KernelCountException
|
|
from extra.lr_scheduler import OneCycleLR
|
|
from test.helpers import derandomize_model
|
|
|
|
from examples.gpt2 import Transformer as GPT2Transformer
|
|
from examples.hlb_cifar10 import SpeedyResNet, hyp
|
|
from extra.models.llama import Transformer as LLaMaTransformer
|
|
from examples.stable_diffusion import UNetModel, unet_params
|
|
from extra.models.unet import ResBlock
|
|
from extra.models.bert import BertForPretraining
|
|
|
|
global_mem_used = 0
|
|
def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jitted=False):
|
|
with Context(JIT=2):
|
|
tms = []
|
|
for _ in range(2):
|
|
early_gen = [x.realize() if isinstance(x, Tensor) else x for x in gen()]
|
|
GlobalCounters.reset()
|
|
Device[Device.DEFAULT].synchronize()
|
|
st = time.perf_counter_ns()
|
|
model(*early_gen)
|
|
Device[Device.DEFAULT].synchronize()
|
|
tms.append(time.perf_counter_ns() - st)
|
|
mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
|
|
|
|
kernels_used = jit_cache_count(model.captured.linear) if getattr(model, "captured", None) is not None else None
|
|
print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
|
|
assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
|
|
assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"
|
|
if kernels_used:
|
|
if kernels_used > max_kernels_allowed: raise KernelCountException(max_kernels_allowed, kernels_used)
|
|
if (max_kernels_allowed - kernels_used) / max_kernels_allowed >= 0.2:
|
|
raise KernelCountException(max_kernels_allowed, kernels_used)
|
|
if all_jitted:
|
|
assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted" # noqa: E501
|
|
|
|
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
|
|
|
|
class TestRealWorld(unittest.TestCase):
|
|
def setUp(self):
|
|
gc.collect()
|
|
global global_mem_used
|
|
global_mem_used = GlobalCounters.mem_used
|
|
np.random.seed(2002)
|
|
|
|
@slow
|
|
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
|
def test_stable_diffusion(self):
|
|
params = unet_params
|
|
params["model_ch"] = 8
|
|
params["ctx_dim"] = 8
|
|
params["num_res_blocks"] = 1
|
|
params["n_heads"] = 2
|
|
model = UNetModel(**params)
|
|
derandomize_model(model)
|
|
@TinyJit
|
|
def test(t, t2): return model(t, Tensor([801]), t2).realize()
|
|
helper_test("test_sd", lambda: (Tensor.randn(1, 4, 32, 32), Tensor.randn(1, 77, params["ctx_dim"])), test, 0.011, 515)
|
|
|
|
def test_unet_resblock(self):
|
|
model = [ResBlock(16, 24, 16) for _ in range(4)]
|
|
derandomize_model(model)
|
|
@TinyJit
|
|
def test(t, t2):
|
|
for l in model: t = l(t, t2)
|
|
return t.realize()
|
|
|
|
# TODO: support _offset on CL to get mem down to 0.0002
|
|
exp_mem = 0.00037 if Device.DEFAULT == "CL" else 0.0002
|
|
helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, exp_mem, 37)
|
|
|
|
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
|
def test_llama(self):
|
|
self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
|
|
|
|
args_tiny = {"dim": 1024, "hidden_dim": 2048, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-05, "vocab_size": 1000}
|
|
model = LLaMaTransformer(**args_tiny)
|
|
derandomize_model(model)
|
|
@TinyJit
|
|
def test(t): return model(t, 0).realize()
|
|
# TODO: test first token vs rest properly
|
|
helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.23, 118, all_jitted=True)
|
|
|
|
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
|
def test_gpt2(self):
|
|
self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
|
|
|
|
args_tiny = {"dim": 1024, "n_heads": 8, "n_layers": 8, "norm_eps": 1e-5, "vocab_size": 1000}
|
|
model = GPT2Transformer(**args_tiny)
|
|
derandomize_model(model)
|
|
@TinyJit
|
|
def test(t, v):
|
|
with Context(JIT=0): return model(t, v).realize()
|
|
helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23, 168, all_jitted=True)
|
|
|
|
@slow
|
|
def test_train_mnist(self):
|
|
from examples.beautiful_mnist import Model
|
|
with Context(TRAINING=1):
|
|
model = Model()
|
|
optimizer = optim.Adam(get_parameters(model))
|
|
BS = 32
|
|
|
|
@TinyJit
|
|
def train(X):
|
|
out = model(X)
|
|
loss = out.mean()
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.017, 103)
|
|
|
|
@slow
|
|
def test_forward_cifar(self):
|
|
BS = 32
|
|
# with training batchnorm still though
|
|
with Context(TRAINING=1):
|
|
model = SpeedyResNet(Tensor.ones((12,3,2,2)))
|
|
@TinyJit
|
|
def run(X): return model(X)
|
|
helper_test("forward_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), run, 0.033, 27)
|
|
|
|
@slow
|
|
def test_train_cifar(self):
|
|
with Context(TRAINING=1):
|
|
model = SpeedyResNet(Tensor.ones((12,3,2,2)))
|
|
optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=0.8, nesterov=True, weight_decay=0.15)
|
|
BS = 32
|
|
|
|
@TinyJit
|
|
def train(X):
|
|
out = model(X)
|
|
loss = out.mean()
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, 0.12, 126)
|
|
|
|
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need dtypes.float16")
|
|
def test_train_cifar_hyp(self):
|
|
self.enterContext(Context(DEFAULT_FLOAT=dtypes.float16))
|
|
with Context(TRAINING=1):
|
|
model = SpeedyResNet(Tensor.ones((12,3,2,2)))
|
|
optimizer = optim.SGD(get_parameters(model), lr=0.01, momentum=hyp['opt']['momentum'], nesterov=True, weight_decay=hyp['opt']['bias_decay'])
|
|
initial_div_factor = hyp['opt']['initial_div_factor']
|
|
final_lr_ratio = hyp['opt']['final_lr_ratio']
|
|
pct_start = hyp['opt']['percent_start']
|
|
lr_scheduler = OneCycleLR(optimizer, 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=4)
|
|
assert not np.isnan(lr_scheduler.min_lr), "lr too small or initial_div_facotr too big for half"
|
|
|
|
@slow
|
|
def test_bert(self):
|
|
with Context(TRAINING=1):
|
|
args_tiny = {"attention_probs_dropout_prob": 0.0, "hidden_dropout_prob": 0.0, "vocab_size": 30522, "type_vocab_size": 2,
|
|
"max_position_embeddings": 512, "hidden_size": 128, "intermediate_size": 512, "num_attention_heads": 2, "num_hidden_layers": 2}
|
|
model = BertForPretraining(**args_tiny)
|
|
optimizer = optim.LAMB(get_parameters(model))
|
|
|
|
@TinyJit
|
|
def train(input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
|
|
masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
|
|
lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
|
|
loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
|
|
optimizer.zero_grad()
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
from examples.mlperf.helpers import get_fake_data_bert
|
|
data = get_fake_data_bert(BS=4)
|
|
for v in data.values(): v.to_(Device.DEFAULT)
|
|
|
|
helper_test("train_bert", lambda: (data["input_ids"], data["segment_ids"], data["input_mask"], data["masked_lm_positions"], \
|
|
data["masked_lm_ids"], data["masked_lm_weights"], data["next_sentence_labels"]), train, 0.31, 400)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|