mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-08 07:39:33 +01:00
159de1f1d6
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/49577 Repurposing the benchmarking from https://github.com/facebookresearch/fairscale/blob/master/benchmarks/pipe.py and pulling in a stripped down version of the benchmark into PyTorch. Sample output: ``` Running benchmark with args: Namespace(batch_size=8, checkpoint='never', chunks=4, host='localhost', max_batch=10, num_decoder_layers=10, num_devices=4) Number of parameters for model: 292833040 | batch 1 | wps 3593.07 | loss 25.98 | ppl 192556591553.37 | batch 2 | wps 4405.16 | loss 19.36 | ppl 256201548.33 | batch 3 | wps 4404.98 | loss 23.56 | ppl 17111244076.37 | batch 4 | wps 4413.25 | loss 27.11 | ppl 594561327825.83 | batch 5 | wps 4408.53 | loss 25.92 | ppl 181277705101.33 | batch 6 | wps 4385.64 | loss 24.92 | ppl 66592883598.50 | batch 7 | wps 4434.11 | loss 24.75 | ppl 56113635884.68 | batch 8 | wps 4441.25 | loss 24.88 | ppl 63666024212.82 | batch 9 | wps 4425.49 | loss 25.35 | ppl 101959669008.98 | batch 10 | wps 4421.05 | loss 25.34 | ppl 101597621863.94 Peak memory usage for GPUs: cuda:0: 2.38GiB, cuda:1: 3.04GiB, cuda:2: 3.04GiB, cuda:3: 3.67GiB, ``` ghstack-source-id: 118939686 Test Plan: sentinel Reviewed By: rohan-varma Differential Revision: D25628721 fbshipit-source-id: 41c788eed4f852aef019aec18a84cb25ad254f3a
57 lines
1.7 KiB
Python
57 lines
1.7 KiB
Python
import torch
|
|
from torch.utils.data import Dataset
|
|
|
|
|
|
def collate_sentences_lm(samples):
|
|
|
|
if len(samples) == 0:
|
|
return {}
|
|
|
|
id = torch.LongTensor([s["id"] for s in samples])
|
|
src_tokens = torch.stack([s["source"] for s in samples], 0)
|
|
tgt_tokens = torch.stack([s["target"] for s in samples], 0)
|
|
ntokens = len(samples) * len(samples[0]["target"])
|
|
src_lengths = torch.LongTensor([len(samples[0]["source"])] * len(samples))
|
|
|
|
batch = {
|
|
"id": id,
|
|
"nsentences": len(samples),
|
|
"ntokens": ntokens,
|
|
"input": src_tokens,
|
|
"target": tgt_tokens,
|
|
}
|
|
return batch
|
|
|
|
|
|
class BenchmarkLMDataset(Dataset):
|
|
"""
|
|
Dataset to benchmark a translation like seq2seq task.
|
|
Args:
|
|
vocab_size (int, optional): size of the vocabulary (default 10000).
|
|
max_source_positions (int, optional): max number of tokens in the
|
|
source sentence (default: 1024).
|
|
total_samples (int, optional): the total number of rows in the
|
|
dataset (default: 10000).
|
|
"""
|
|
|
|
def __init__(
|
|
self, vocab_size=10000, max_source_positions=1024, total_samples=10000,
|
|
):
|
|
self.vocab_size = vocab_size
|
|
self.max_source_positions = max_source_positions
|
|
self.total_samples = total_samples
|
|
self.sizes = [self.max_source_positions] * self.total_samples
|
|
|
|
def __getitem__(self, index):
|
|
length = self.sizes[index]
|
|
source = torch.randint(1, self.vocab_size, (length,))
|
|
target = source.clone()
|
|
return {
|
|
"id": index,
|
|
"source": source,
|
|
"target": target,
|
|
}
|
|
|
|
def __len__(self):
|
|
return self.total_samples
|