Skip to content

Instantly share code, notes, and snippets.

@sshleifer
Last active June 16, 2020 11:53
Show Gist options
  • Save sshleifer/2465197070a6af6234ff537f4f21220b to your computer and use it in GitHub Desktop.
Save sshleifer/2465197070a6af6234ff537f4f21220b to your computer and use it in GitHub Desktop.
model batch_size sequence_length MB
t5-large 1 128 6558
t5-large 1 512 9568
t5-large 1 1024 23124
facebook/bart-large 1 128 3758
facebook/bart-large 1 512 4670
facebook/bart-large 1 1024 8888
t5-base 1 128 2242
t5-base 1 512 3776
t5-base 1 1024 9056
@sshleifer
Copy link
Author

sshleifer commented Jun 12, 2020

T5 Lines with top memory consumption:
=> 376: mem 1956:         weights = F.dropout(weights, p=self.dropout, training=self.training)  # (bs, n_heads, qlen, klen)
=> 359: mem 1728:         scores = torch.einsum("bnqd,bnkd->bnqk", q, k)  # (bs, n_heads, qlen, klen)
=> 375: mem 1728:         weights = F.softmax(scores.float(), dim=-1).type_as(scores)  # (bs, n_heads, qlen, klen)
=> 382: mem 600:         context = torch.matmul(weights, v)  # (bs, n_heads, qlen, dim_per_head)
=> 188: mem 396:         layer_output = hidden_states + self.dropout(y)
=> 172: mem 288:         h = F.relu(h)

@sshleifer
Copy link
Author

python run_benchmark.py --models t5-large facebook/bart-large t5-base \
   --save_to_csv --inference_time_csv_file plots/time.csv  --train_memory_csv_file plots/train_memory.csv \
   --inference_memory_csv_file plots/memory.csv --training  --batch_size 1 --sequence_lengths 128 512 1024 --trace_memory_line_by_line   --save_to_csv | tee -a trace.txt

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment