geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

put bert data shard inside jit (#9160)

Browse Source 
python time 45ms -> 9ms, it was spending time to schedule the shard

also init bert data on CLANG since it's from numpy, so we don't create the tensor on default device then shard into GPUS
This commit is contained in:
chenyu
2025-02-18 10:36:54 -05:00
committed by GitHub
parent 679291e26a
commit ff05bff221
3 changed files with 29 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/mlperf/dataloader.py
+7 -7
View File
@@ -170,13 +170,13 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
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.float32),
"input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.default_float),
"segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.float32),
"masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.float32),
"masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.float32),
"masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32),
"next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.float32),
"input_ids": Tensor(np.concatenate([s["input_ids"] for s in data], axis=0), dtype=dtypes.float32, device="CLANG"),
"input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.default_float, device="CLANG"),
"segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.float32, device="CLANG"),
"masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.float32, device="CLANG"),
"masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.float32, device="CLANG"),
"masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32, device="CLANG"),
"next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.float32, device="CLANG"),
}
def load_file(file: str):
examples/mlperf/helpers.py
+7 -12
View File
@@ -222,18 +222,13 @@ def get_mlperf_bert_model():
config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
return BertForPretraining(**config)
def get_data_bert(GPUS:list[str], it):
data: dict[str, Tensor] = next(it)
for key in data.keys(): data[key].shard_(GPUS, axis=0)
return data
def get_fake_data_bert(BS:int):
return {
"input_ids": Tensor.empty((BS, 512), dtype=dtypes.float32),
"input_mask": Tensor.empty((BS, 512), dtype=dtypes.default_float),
"segment_ids": Tensor.empty((BS, 512), dtype=dtypes.float32),
"masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.float32),
"masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.float32),
"masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32),
"next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.float32),
"input_ids": Tensor.empty((BS, 512), dtype=dtypes.float32, device="CLANG"),
"input_mask": Tensor.empty((BS, 512), dtype=dtypes.default_float, device="CLANG"),
"segment_ids": Tensor.empty((BS, 512), dtype=dtypes.float32, device="CLANG"),
"masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CLANG"),
"masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CLANG"),
"masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CLANG"),
"next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.float32, device="CLANG"),
}
examples/mlperf/model_train.py
+15 -11
View File
@@ -573,7 +573,9 @@ def train_rnnt():
@TinyJit
def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
t.shard_(GPUS, axis=0)
optimizer.zero_grad()
lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
@@ -593,7 +595,9 @@ def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Te
@TinyJit
def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
masked_lm_weights:Tensor, next_sentence_labels:Tensor):
masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
t.shard_(GPUS, axis=0)
lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss = \
model.accuracy(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
@@ -602,7 +606,7 @@ def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:T
def train_bert():
# NOTE: pip install tensorflow, wandb required
from examples.mlperf.dataloader import batch_load_train_bert, batch_load_val_bert
from examples.mlperf.helpers import get_mlperf_bert_model, get_data_bert, get_fake_data_bert
from examples.mlperf.helpers import get_mlperf_bert_model, get_fake_data_bert
from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup
config = {}
@@ -758,7 +762,7 @@ def train_bert():
# ** train loop **
wc_start = time.perf_counter()
i, train_data = start_step, get_data_bert(GPUS, train_it)
i, train_data = start_step, next(train_it)
if RUNMLPERF:
if MLLOGGER:
@@ -771,12 +775,12 @@ def train_bert():
GlobalCounters.reset()
loss, global_norm = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler,
train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"])
train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"], GPUS)
pt = time.perf_counter()
try:
next_data = get_data_bert(GPUS, train_it)
next_data = next(train_it)
except StopIteration:
next_data = None
@@ -811,8 +815,8 @@ def train_bert():
if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK) or i == train_steps:
if MLLOGGER and RUNMLPERF:
MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*BS, "step_num": i})
if getenv("RESET_STEP", 0): train_step_bert.reset()
train_step_bert.captured.free_intermediates()
if getenv("RESET_STEP", 0) or INITMLPERF: train_step_bert.reset()
else: train_step_bert.captured.free_intermediates()
eval_lm_losses = []
eval_clsf_losses = []
eval_lm_accs = []
@@ -822,13 +826,13 @@ def train_bert():
BEAM.value = EVAL_BEAM
for j in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps, disable=BENCHMARK):
eval_data = get_data_bert(GPUS, eval_it)
eval_data = next(eval_it)
GlobalCounters.reset()
st = time.time()
lm_acc, clsf_acc, lm_loss, clsf_loss = eval_step_bert(model,
eval_data["input_ids"], eval_data["segment_ids"], eval_data["input_mask"], eval_data["masked_lm_positions"],
eval_data["masked_lm_ids"], eval_data["masked_lm_weights"], eval_data["next_sentence_labels"])
eval_data["masked_lm_ids"], eval_data["masked_lm_weights"], eval_data["next_sentence_labels"], GPUS)
lm_acc, clsf_acc, lm_loss, clsf_loss = lm_acc.item(), clsf_acc.item(), lm_loss.item(), clsf_loss.item()
eval_lm_losses.append(lm_loss)
@@ -846,7 +850,7 @@ def train_bert():
return
if getenv("RESET_STEP", 0): eval_step_bert.reset()
eval_step_bert.captured.free_intermediates()
else: eval_step_bert.captured.free_intermediates()
del eval_data
avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
avg_clsf_loss = sum(eval_clsf_losses) / len(eval_clsf_losses)