| 142 | |
| 143 | |
| 144 | void fill_in_activation_sizes(size_t* act_sizes, GPT2Config config, int B, int T) { |
| 145 | size_t C = config.channels; |
| 146 | size_t NH = config.num_heads; |
| 147 | size_t L = config.num_layers; |
| 148 | size_t Vp = config.padded_vocab_size; |
| 149 | act_sizes[0] = B * T * C; // encoded |
| 150 | act_sizes[1] = L * B * T * C; // ln1 |
| 151 | act_sizes[2] = L * B * T; // ln1_mean |
| 152 | act_sizes[3] = L * B * T; // ln1_rstd |
| 153 | act_sizes[4] = L * B * T * 3 * C; // qkv |
| 154 | act_sizes[5] = L * B * T * C; // atty |
| 155 | act_sizes[6] = L * B * NH * T * T; // preatt |
| 156 | act_sizes[7] = L * B * NH * T * T; // att |
| 157 | act_sizes[8] = L * B * T * C; // attproj |
| 158 | act_sizes[9] = L * B * T * C; // residual2 |
| 159 | act_sizes[10] = L * B * T * C; // ln2 |
| 160 | act_sizes[11] = L * B * T; // ln2_mean |
| 161 | act_sizes[12] = L * B * T; // ln2_rstd |
| 162 | act_sizes[13] = L * B * T * 4 * C; // fch |
| 163 | act_sizes[14] = L * B * T * 4 * C; // fch_gelu |
| 164 | act_sizes[15] = L * B * T * C; // fcproj |
| 165 | act_sizes[16] = L * B * T * C; // residual3 |
| 166 | act_sizes[17] = B * T * C; // lnf |
| 167 | act_sizes[18] = B * T; // lnf_mean |
| 168 | act_sizes[19] = B * T; // lnf_rstd |
| 169 | act_sizes[20] = B * T * Vp; // logits |
| 170 | act_sizes[21] = B * T * Vp; // probs |
| 171 | act_sizes[22] = B * T; // losses |
| 172 | } |
| 173 | |
| 174 | float* malloc_and_point_activations(ActivationTensors* acts, size_t* act_sizes) { |
| 175 | size_t num_activations = 0; |