move things around a bit, a few comments, and also update python script to correspond to C more

train_gpt2.cu

@@ -740,10 +740,11 @@ float gpt2_validate(GPT2 *model, const int* inputs, const int* targets, size_t B
     return mean_loss;
 }
 
-
 void gpt2_zero_grad(GPT2 *model) {
     NVTX_RANGE_FN();
-    // the losses accumulate over the duration of gradient accumulation micro steps, also reset here
+    // there are currently two state vars during the gradient accumulation inner loop:
+    // 1) the losses accumulate += into acts.losses, reset here
+    // 2) the gradients accumulate += into grads_memory, reset here
     cudaCheck(cudaMemset(model->acts.losses, 0, model->batch_size * model->seq_len * sizeof(float)));
     if (model->grads_memory != NULL) {
         cudaCheck(cudaMemset(model->grads_memory, 0, model->num_parameters * sizeof(floatX)));
@@ -755,8 +756,9 @@ void gpt2_backward_and_reduce(GPT2 *model, int* inputs, const int* targets, int
     NVTX_RANGE_FN();
     bool last_step = micro_step == grad_accum_steps - 1;
 
+    // on the first micro-step zero the gradients, as we're about to += accumulate into them
     if (micro_step == 0) {
-        gpt2_zero_grad(model);  // set correct grad state on first micro-step
+        gpt2_zero_grad(model);
     }
 
     // lazily allocate the memory for gradients of the weights and activations, if needed

train_gpt2.py

@@ -698,10 +698,6 @@ if __name__ == "__main__":
         write_state(model, x, y, logits, loss, f"gpt2_{model_size_str}_debug_state.bin")
         # reset the train_loader for the optimization below
         train_loader.reset()
-        # clear the grads here explicitly because otherwise we'd have a duplicate grad accumulation
-        # since in the training loop we do a backward() and then zero_grad() at the end of the loop
-        # this would cause an incorrect first training step
-        model.zero_grad()
 
     # -------------------------------------------------------------------------
     # main training loop
@@ -794,6 +790,7 @@ if __name__ == "__main__":
 
         # --------------- TRAINING SECTION BEGIN -----------------
         model.train()
+        optimizer.zero_grad(set_to_none=True)
         # if we are trying to overfit a single batch, we reset the loader here
         if args.overfit_single_batch:
             train_loader.reset()
@@ -830,7 +827,6 @@ if __name__ == "__main__":
             param_group['lr'] = lr
         # step the optimizer
         optimizer.step()
-        optimizer.zero_grad(set_to_none=True)
         # --------------- TRAINING SECTION END -------------------
         # everything that follows now is just diagnostics, prints, logging, etc.