GPT-3 175B architecture accounting: exact params, shape trace, KV-cache (hmgpt3)

Usage

morie_geron_gpt3(
  prompt,
  n_tokens,
  n_layers = 96,
  d_model = 12288,
  n_heads = 96,
  d_ff = NULL,
  vocab_size = 50257,
  n_ctx = 2048,
  dtype_bytes = 2
)

Arguments

prompt: Token ids.
n_tokens: Tokens to generate.
n_layers, d_model, n_heads, d_ff, vocab_size, n_ctx, dtype_bytes: Config.

Value

List with total_parameters, parameters_per_layer, breakdown, shape_trace, kv_cache_bytes.

Examples

morie_geron_gpt3(prompt = c(1, 2, 3, 4, 5, 6, 7, 8), n_tokens = 5L)
#> $total_parameters
#> [1] 174604259328
#> 
#> $parameters_per_layer
#> [1] 1812099072
#> 
#> $breakdown
#> $breakdown$token_embedding
#> [1] 617558016
#> 
#> $breakdown$position_embedding
#> [1] 25165824
#> 
#> $breakdown$attention_weights
#> [1] 57982058496
#> 
#> $breakdown$attention_biases
#> [1] 4718592
#> 
#> $breakdown$feedforward_weights
#> [1] 115964116992
#> 
#> $breakdown$feedforward_biases
#> [1] 5898240
#> 
#> $breakdown$layer_norms
#> [1] 4743168
#> 
#> $breakdown$output_head
#> [1] 0
#> 
#> 
#> $shape_trace
#> $shape_trace[[1]]
#> $shape_trace[[1]][[1]]
#> [1] "token_ids"
#> 
#> $shape_trace[[1]][[2]]
#> [1] 8
#> 
#> 
#> $shape_trace[[2]]
#> $shape_trace[[2]][[1]]
#> [1] "embedded"
#> 
#> $shape_trace[[2]][[2]]
#> [1]     8 12288
#> 
#> 
#> $shape_trace[[3]]
#> $shape_trace[[3]][[1]]
#> [1] "q_per_head"
#> 
#> $shape_trace[[3]][[2]]
#> [1]  96   8 128
#> 
#> 
#> $shape_trace[[4]]
#> $shape_trace[[4]][[1]]
#> [1] "attention_scores"
#> 
#> $shape_trace[[4]][[2]]
#> [1] 96  8  8
#> 
#> 
#> $shape_trace[[5]]
#> $shape_trace[[5]][[1]]
#> [1] "attention_out"
#> 
#> $shape_trace[[5]][[2]]
#> [1]     8 12288
#> 
#> 
#> $shape_trace[[6]]
#> $shape_trace[[6]][[1]]
#> [1] "ffn_hidden"
#> 
#> $shape_trace[[6]][[2]]
#> [1]     8 49152
#> 
#> 
#> $shape_trace[[7]]
#> $shape_trace[[7]][[1]]
#> [1] "block_out"
#> 
#> $shape_trace[[7]][[2]]
#> [1]     8 12288
#> 
#> 
#> $shape_trace[[8]]
#> $shape_trace[[8]][[1]]
#> [1] "logits"
#> 
#> $shape_trace[[8]][[2]]
#> [1]     5 50257
#> 
#> 
#> 
#> $d_head
#> [1] 128
#> 
#> $n_prompt_tokens
#> [1] 8
#> 
#> $n_generated
#> [1] 5
#> 
#> $context_used
#> [1] 13
#> 
#> $kv_cache_bytes
#> [1] 61341696
#> 
#> $flops_per_token
#> [1] 349208518656
#> 
#> $estimate
#> [1] 174604259328
#> 
#> $n
#> [1] 8
#> 
#> $method
#> [1] "GPT-3 decoder-only architecture accounting"
#>