Files
StarPilot/tinygrad_repo/extra/hcqfuzz/tests/bert.py
T
firestar5683 22a8575199 Vendor dependencies and resources for standalone builds
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.
2026-10-03 21:01:18 +00:00

82 lines
1.9 KiB
Python

from extra.hcqfuzz.spec import TestSpec
import random
bert_train_params = {
"DEFAULT_FLOAT": "HALF",
"SUM_DTYPE": "HALF",
"GPUS": 6,
"BS": 96,
"EVAL_BS": 96,
"BASEDIR": "/raid/datasets/wiki",
}
class TrainBert(TestSpec):
def prepare(self, dev, seed):
random.seed(seed)
self.env = {
**bert_train_params,
"IGNORE_BEAM_CACHE": 1,
"BEAM": 5,
"BEAM_UOPS_MAX": 10000,
"BEAM_UPCAST_MAX": 256,
"BEAM_LOCAL_MAX": 1024,
"BEAM_MIN_PROGRESS": 5,
"IGNORE_JIT_FIRST_BEAM": 1,
"LOGMLPERF": 0,
"SEED": seed,
}
self.cmd = "python3 examples/mlperf/model_train.py"
self.timeout = 7 * 60 * 60 # 7 hours
def get_exec_state(self): return self.env, self.cmd, self.timeout
class TrainBertShort(TestSpec):
def prepare(self, dev, seed):
random.seed(seed)
self.env = {
**bert_train_params,
"IGNORE_BEAM_CACHE": 1,
"BEAM": 5,
"BEAM_UOPS_MAX": 10000,
"BEAM_UPCAST_MAX": 256,
"BEAM_LOCAL_MAX": 1024,
"BEAM_MIN_PROGRESS": 5,
"IGNORE_JIT_FIRST_BEAM": 1,
"SEED": seed,
"BENCHMARK": 4096,
"JIT": 2
}
self.cmd = "python3 examples/mlperf/model_train.py"
self.timeout = 2 * 60 * 60 # 2 hours
def get_exec_state(self): return self.env, self.cmd, self.timeout
class BertBeam(TestSpec):
def prepare(self, dev, seed):
random.seed(seed)
self.env = {
**bert_train_params,
"IGNORE_BEAM_CACHE": 1,
"BEAM": random.choice([1, 2, 3, 4, 5]),
"BEAM_UOPS_MAX": 10000,
"BEAM_UPCAST_MAX": 256,
"BEAM_LOCAL_MAX": 1024,
"BEAM_MIN_PROGRESS": 5,
"IGNORE_JIT_FIRST_BEAM": 1,
"SEED": seed,
"RESET_STEP": 1,
"BENCHMARK": 10,
"BERT_LAYERS": 2,
"SEED": seed,
}
self.cmd = "python3 examples/mlperf/model_train.py"
self.timeout = 1 * 60 * 60 # 1 hour
def get_exec_state(self): return self.env, self.cmd, self.timeout