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.
103 lines
4.3 KiB
Python
103 lines
4.3 KiB
Python
import unittest, time
|
|
|
|
from tinygrad import Tensor, TinyJit, GlobalCounters, Device
|
|
from tinygrad.helpers import getenv, Context
|
|
from tinygrad.nn.optim import LAMB
|
|
from tinygrad.nn.state import get_parameters
|
|
from tinygrad.engine.realize import run_linear
|
|
|
|
from extra.models import bert
|
|
|
|
bs = getenv("BS", 16)
|
|
seq_len = getenv("SEQ_LEN", 512)
|
|
|
|
class BenchmarkBertTrain(unittest.TestCase):
|
|
def _get_layer(self, layer_id):
|
|
if not hasattr(self, "model"):
|
|
dropout_prob = 0.0 if getenv("DISABLE_DROPOUT") else 0.1
|
|
self.model = bert.BertForPretraining(attention_probs_dropout_prob=dropout_prob, hidden_dropout_prob=dropout_prob)
|
|
hidden_size = self.model.bert.embeddings.word_embeddings.embed_sz
|
|
intermediate_size = self.model.bert.encoder.layer[0].intermediate.dense.weight.shape[0]
|
|
|
|
layer_map = {
|
|
"embedding": self.model.bert.embeddings,
|
|
"attention_self": self.model.bert.encoder.layer[0].attention.self,
|
|
"attention_output": self.model.bert.encoder.layer[0].attention.output,
|
|
"intermediate": self.model.bert.encoder.layer[0].intermediate,
|
|
"output": self.model.bert.encoder.layer[0].output
|
|
}
|
|
input_shapes = {
|
|
"embedding": [(bs, seq_len), (bs, seq_len)],
|
|
"attention_self": [(bs, seq_len, hidden_size), (bs, 1, 1, seq_len)],
|
|
"attention_output": [(bs, seq_len, hidden_size), (bs, seq_len, 1)],
|
|
"intermediate": [(bs, seq_len, hidden_size)],
|
|
"output": [(bs, seq_len, intermediate_size), (bs, seq_len, 1)]
|
|
}.get(layer_id)
|
|
|
|
return f"{layer_id}-layer, Input: {input_shapes}", layer_map.get(layer_id), input_shapes
|
|
|
|
def _test_layer(self, name, layer, input_shapes):
|
|
optim = LAMB(get_parameters(layer))
|
|
with Context(TRACK_MATCH_STATS=0): Tensor.realize(*[t.assign(t.detach().contiguous()) for t in get_parameters(optim)])
|
|
|
|
JITCNT = getenv("JITCNT", 1)
|
|
@TinyJit
|
|
@Context(TRAINING=1)
|
|
def step(inputs):
|
|
optim.zero_grad()
|
|
for i in inputs: i.grad = None
|
|
|
|
y = layer(*inputs).contiguous().contiguous_backward()
|
|
y.sum().backward()
|
|
if getenv("ASSIGN", 1): linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *optim.schedule_step())
|
|
else: linear, var_vals = Tensor.linear_with_vars(y, *list(inputs), *[t.grad for t in optim.params])
|
|
|
|
for _ in range(JITCNT):
|
|
run_linear(linear, var_vals)
|
|
|
|
CNT = getenv("CNT", 5)
|
|
best_tm = None
|
|
flops, mem_used, mem, kernels = None, None, None, None
|
|
for _ in range(CNT):
|
|
with Context(TRACK_MATCH_STATS=0): inputs = [Tensor.randn(*shape).realize() for shape in input_shapes]
|
|
GlobalCounters.reset()
|
|
|
|
st = time.perf_counter()
|
|
step(inputs)
|
|
Device[Device.DEFAULT].synchronize()
|
|
et = time.perf_counter()
|
|
|
|
flops = GlobalCounters.global_ops / JITCNT
|
|
mem_used = GlobalCounters.mem_used
|
|
mem = GlobalCounters.global_mem / JITCNT
|
|
if kernels is None: kernels = GlobalCounters.kernel_count // JITCNT
|
|
tm = (et-st) / JITCNT
|
|
if best_tm is None or tm < best_tm: best_tm = tm
|
|
print(f"\r{name:70s}: {best_tm * 1000:>9.2f} ms, {flops / 10**12 / best_tm:>6.2f} TFLOPS, {mem / 10**9 / best_tm:>5.0f} GB/s, "
|
|
f"{mem_used / 10**9: 6.2f} GB used, {kernels:>5d} kernels")
|
|
return best_tm, flops, mem, kernels
|
|
|
|
def test_embedding_layer(self): self._est(*self._test_layer(*self._get_layer("embedding")), 1)
|
|
def test_attention_self_layer(self): self._est(*self._test_layer(*self._get_layer("attention_self")), 24) # Assumes BERT-large
|
|
def test_attention_output_layer(self): self._est(*self._test_layer(*self._get_layer("attention_output")), 24)
|
|
def test_intermediate_layer(self): self._est(*self._test_layer(*self._get_layer("intermediate")), 24)
|
|
def test_output_layer(self): self._est(*self._test_layer(*self._get_layer("output")), 24)
|
|
|
|
est_tm, est_flops, est_mem, est_kernels = 0, 0, 0, 0
|
|
|
|
@classmethod
|
|
def _est(cls, tm, flops, mem, kernels, mult):
|
|
cls.est_tm += tm * mult
|
|
cls.est_flops += flops * mult
|
|
cls.est_mem += mem * mult
|
|
cls.est_kernels += kernels * mult
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
print(f"\restimated step tm: {cls.est_tm * 1000.0:.2f} ms, {cls.est_flops / 10 ** 12 / cls.est_tm:.3f} tflops, "
|
|
f"{cls.est_mem / 10 ** 9 / cls.est_tm:.2f} GB/s, {cls.est_kernels} kernels")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|