| 580 | } |
| 581 | |
| 582 | void gpt2_update(GPT2 *model, float learning_rate, float beta1, float beta2, float eps, float weight_decay, int t) { |
| 583 | // reference: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html |
| 584 | |
| 585 | // lazily allocate the memory for m_memory and v_memory |
| 586 | if (model->m_memory == NULL) { |
| 587 | model->m_memory = (float*)calloc(model->num_parameters, sizeof(float)); |
| 588 | model->v_memory = (float*)calloc(model->num_parameters, sizeof(float)); |
| 589 | } |
| 590 | |
| 591 | for (size_t i = 0; i < model->num_parameters; i++) { |
| 592 | float param = model->params_memory[i]; |
| 593 | float grad = model->grads_memory[i]; |
| 594 | |
| 595 | // update the first moment (momentum) |
| 596 | float m = beta1 * model->m_memory[i] + (1.0f - beta1) * grad; |
| 597 | // update the second moment (RMSprop) |
| 598 | float v = beta2 * model->v_memory[i] + (1.0f - beta2) * grad * grad; |
| 599 | // bias-correct both moments |
| 600 | float m_hat = m / (1.0f - powf(beta1, t)); |
| 601 | float v_hat = v / (1.0f - powf(beta2, t)); |
| 602 | |
| 603 | // update |
| 604 | model->m_memory[i] = m; |
| 605 | model->v_memory[i] = v; |
| 606 | model->params_memory[i] -= learning_rate * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * param); |
| 607 | } |
| 608 | } |
| 609 | |
| 610 | void gpt2_free(GPT2 *model) { |
| 611 | free(model->params_memory); |