Merge pull request #447 from karpathy/feature/hellaswagc

HellaSwag eval in C

.gitignore

@@ -7,7 +7,7 @@
 
 # data directories
 dev/data/__pycache__/
-dev/data/fineweb/
+dev/data/fineweb10B/
 dev/data/hellaswag/
 dev/data/mmlu/
 dev/data/tinyshakespeare/

dataloader.h

@@ -195,3 +195,269 @@ void dataloader_free(DataLoader *loader) {
     fcloseCheck(loader->tokens_file);
     globfree(&loader->glob_result);
 }
+
+// ----------------------------------------------------------------------------
+// Distributed Eval Loader
+// Many evals (like) HellaSwag and MMLU are multiple-choice
+// where there are 4 possible continuations and a label for the correct one
+// We want to load and serve these style of evals
+/*
+Copy pasting the section on the eval datafile format, from data_common.py:
+- First comes a header with 256 int32s
+- The examples follow, each example is a stream of uint16_t:
+    - <START_EXAMPLE> delimiter of 2**16-1, i.e. 65,535
+    - <EXAMPLE_BYTES>, bytes encoding this example, allowing efficient skip to next
+    - <EXAMPLE_INDEX>, the index of the example in the dataset
+    - <LABEL>, the index of the correct completion
+    - <NUM_COMPLETIONS>, indicating the number of completions (usually 4)
+    - <NUM><CONTEXT_TOKENS>, where <NUM> is the number of tokens in the context
+    - <NUM><COMPLETION_TOKENS>, repeated NUM_COMPLETIONS times
+*/
+
+// for now, could relax later
+#define ASSUMED_NUM_COMPLETIONS 4
+// helper macro for ceildiv
+#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
+
+typedef struct {
+    // variables related to distributed training
+    // each process/worker has to access different parts of the data
+    int process_rank;
+    int num_processes;
+    // hyperparameters. use size_t to prevent overflow
+    size_t B; // (micro) batch size dimension of the tensor that feeds into the model
+    size_t T; // maximum context length of the model
+    // input handling and its state
+    FILE* eval_file;
+    long file_size;
+    uint16_t* buffer; // we fread data from file into this buffer
+    // public variables that could be accessed from outside
+    int num_examples; // in total across all processes
+    int num_batches; // to process the entire dataset across all processes
+    int start_example_index; // the assignment of work for this process, start
+    int end_example_index; // and end. start is inclusive, end is exclusive
+    int current_example_index; // the next example we would read
+    int* inputs;  // input tokens into transformer
+    int* targets; // target tokens for the transformer
+    char* mask; // mask=1 at all completion token locations
+    int* label; // the correct completion labels
+    int num_completions; // number of completions for this example
+} EvalLoader;
+
+void evalloader_reset(EvalLoader *loader) {
+    // we have to be careful that each process starts at the correct offset.
+    // For example if there are N examples in the file and 4 processes,
+    // then process 0 should start at 0, process 1 at N/4, process 2 at N/2, etc.
+    // determine how much work there is for all processes
+    int examples_per_process = CEIL_DIV(loader->num_examples, loader->num_processes);
+    int can_fit_examples = loader->B / ASSUMED_NUM_COMPLETIONS;
+    loader->num_batches = CEIL_DIV(examples_per_process, can_fit_examples);
+    // determine the start and end example indices for this process
+    loader->start_example_index = examples_per_process * loader->process_rank;
+    loader->end_example_index = examples_per_process * (loader->process_rank + 1);
+    // crop the end example index to the total number of examples
+    if (loader->end_example_index > loader->num_examples) {
+        loader->end_example_index = loader->num_examples;
+    }
+    // now seek through the file to the start of that example
+    // utilize <EXAMPLE_BYTES> for efficiency
+    long header_bytes = HEADER_SIZE * sizeof(int);
+    fseekCheck(loader->eval_file, header_bytes, SEEK_SET);
+    for (int i = 0; i < loader->start_example_index; i++) {
+        uint16_t example_header[3];
+        // read 3 uint16_t values: <START_EXAMPLE>, <EXAMPLE_BYTES>, <EXAMPLE_INDEX>
+        freadCheck(&example_header[0], sizeof(uint16_t), 3, loader->eval_file);
+        // validate the <START_EXAMPLE> delimiter
+        assert(example_header[0] == 65535); // <START_EXAMPLE> delimiter
+        // validate the <EXAMPLE_INDEX>
+        assert(example_header[2] == i); // <EXAMPLE_INDEX> should match the loop index
+        // skip to the next example, keeping in mind that we already read the header
+        size_t remaining_bytes = example_header[1] - sizeof(uint16_t) * 3;
+        assert(remaining_bytes > 0); // we expect some bytes in the example
+        fseekCheck(loader->eval_file, remaining_bytes, SEEK_CUR);
+    }
+    // now we are at the start of the example we want to start at, pointing at <START_EXAMPLE>
+    loader->current_example_index = loader->start_example_index;
+}
+
+void evalloader_init(EvalLoader *loader,
+                     const char* filename,
+                     size_t B,
+                     size_t T,
+                     int process_rank,
+                     int num_processes) {
+    loader->process_rank = process_rank;
+    loader->num_processes = num_processes;
+    loader->B = B;
+    loader->T = T;
+
+    // open the file and validate the header
+    loader->eval_file = fopenCheck(filename, "rb");
+    // validate the header
+    int header[HEADER_SIZE];
+    freadCheck(header, sizeof(int), HEADER_SIZE, loader->eval_file);
+    if (header[0] != 20240522) { printf("Bad magic in eval file\n"); exit(EXIT_FAILURE); }
+    if (header[1] != 1) { printf("Bad version in data file\n"); exit(EXIT_FAILURE); }
+    loader->num_examples = header[2]; // number of tokens in the file
+    assert(loader->num_examples >= num_processes); // avoid headaches for now
+    size_t longest_example_bytes = header[3]; // longest example in the file
+    // basic sensibility check we could relax later. but roughly each example
+    // contains the prompt (or "context") and 4 completions, all of these have to be
+    // up to T tokens, and their tokens are uint16_t (so 2 bytes/token).
+    // There's a few more things in each example but they are minor.
+    // So longest example should be roughly this. Just trying to make sure it's sensible.
+    assert(longest_example_bytes > 0 && longest_example_bytes < (1+ASSUMED_NUM_COMPLETIONS)*T*2);
+
+    // allocate all the space we'll need
+    int can_fit_examples = B / ASSUMED_NUM_COMPLETIONS;
+    loader->buffer = (uint16_t*)malloc(longest_example_bytes);
+    loader->inputs = (int*)malloc(B * T * sizeof(int));
+    loader->targets = (int*)malloc(B * T * sizeof(int));
+    loader->mask = (char*)malloc(B * T * sizeof(char));
+    loader->label = (int*)malloc(can_fit_examples * sizeof(int));
+
+    // reset the loader, to initialize it
+    evalloader_reset(loader);
+}
+
+void evalloader_next_example_(EvalLoader *loader, int example_batch_index) {
+    // this function populates the inputs, targets, mask, and label fields for one example
+    // because every (B,T) tensor can fit multiple examples and we want to take advantage,
+    // we also pass in the example_batch_index to indicate which example in the batch we are loading
+    // and each example takes up ASSUMED_NUM_COMPLETIONS rows in the batch
+    size_t B = loader->B;
+    size_t T = loader->T;
+    int batch_dim_offset = example_batch_index * ASSUMED_NUM_COMPLETIONS;
+    // read the current example header
+    uint16_t example_header[3];
+    freadCheck(&example_header[0], sizeof(uint16_t), 3, loader->eval_file);
+    // validate the <START_EXAMPLE> delimiter
+    assert(example_header[0] == 65535); // <START_EXAMPLE> delimiter
+    // validate the <EXAMPLE_INDEX>
+    assert(example_header[2] == loader->current_example_index); // <EXAMPLE_INDEX> should match the loop index
+    assert(example_header[2] >= loader->start_example_index && example_header[2] < loader->end_example_index);
+    // read the rest of the example (we have space for 3 more uint16_t values in buffer, it's ok)
+    size_t example_bytes = example_header[1] - sizeof(uint16_t) * 3;
+    // read example_bytes into buffer. careful that this is actually in the units of bytes
+    freadCheck(loader->buffer, sizeof(char), example_bytes, loader->eval_file);
+    // process the example label
+    int label = (int)loader->buffer[0];
+    int can_fit_examples = loader->B / ASSUMED_NUM_COMPLETIONS;
+    assert(label >= 0 && label < ASSUMED_NUM_COMPLETIONS); // we expect the label to be in [0, 4) for right now
+    assert(example_batch_index >= 0 && example_batch_index < can_fit_examples);
+    loader->label[example_batch_index] = label; // store for output
+    // process the number of completions
+    int num_completions = (int)loader->buffer[1];
+    assert(num_completions == ASSUMED_NUM_COMPLETIONS); // we expect 4 completions for now
+    assert(batch_dim_offset + num_completions <= B); // we expect to fit in the batch
+    loader->num_completions = num_completions; // store for output
+    // process the context
+    // the context is shared for all completions, so we insert it into all data rows equally
+    int context_length = (int)loader->buffer[2];
+    uint16_t *context_tokens_start = &loader->buffer[3]; // where the tokens start
+    assert(context_length > 0 && context_length < T); // context is non-empty and up to T
+    for (int b = 0; b < num_completions; b++) {
+        for (int i = 0; i < context_length; i++) {
+            int boff = batch_dim_offset + b;
+            int tok_cur = (int)context_tokens_start[i];
+            loader->inputs[boff * T + i] = tok_cur;
+        }
+    }
+    // process the completions, insert them in their row, right after the (shared) context
+    uint16_t *completions_iter = loader->buffer + 3 + context_length;
+    for (int c = 0; c < num_completions; c++) {
+        int coff = batch_dim_offset + c;
+        int completion_length = (int)completions_iter[0];
+        uint16_t *completion_tokens_start = completions_iter + 1;
+        assert(completion_length > 0 && context_length + completion_length < T); // things fit?
+        for (int i = 0; i < completion_length; i++) {
+            int tok_cur = (int)completion_tokens_start[i];
+            // at inputs, the completions simply follow the context
+            loader->inputs[coff * T + context_length + i] = tok_cur;
+            // at targets things start to get tricky
+            // we expect the last context token to predict the first completion token
+            // and then onwards from there.
+            loader->targets[coff * T + context_length + i - 1] = tok_cur;
+            // and at these positions, we want to set mask=1, because these are the
+            // positions where we want to average the loss, in each row, to determine
+            // its overall probability of following the context.
+            loader->mask[coff * T + context_length + i - 1] = 1;
+        }
+        completions_iter += 1 + completion_length; // move to the next completion
+    }
+    // advance the current example to point to the next one we'd load
+    loader->current_example_index += 1;
+}
+
+void evalloader_next_batch(EvalLoader *loader) {
+    size_t B = loader->B;
+    size_t T = loader->T;
+    // init all inputs, targets, mask to zeros
+    // TODO: I think only mask is necessary to reset?
+    memset(loader->inputs, 0, B * T * sizeof(int));
+    memset(loader->targets, 0, B * T * sizeof(int));
+    memset(loader->mask, 0, B * T * sizeof(char));
+    // ok here is the problem we are solving
+    // we have a batch dimension of B, which we want to take full advantage of
+    // each example has some number of completions (usually 4)
+    // so we want to pack as many examples into rows of B as we can fit
+    int can_fit_examples = B / ASSUMED_NUM_COMPLETIONS; // how many examples can we fit in the batch?
+    for (int i = 0; i < can_fit_examples; i++) {
+        if (loader->current_example_index >= loader->end_example_index) {
+            break; // this process has exhausted its work, noop from here on
+        }
+        evalloader_next_example_(loader, i);
+    }
+}
+
+int evalloader_stat_losses(EvalLoader *loader, float* losses) {
+    // compute statistics of losses (B*T) resulting from a forward pass
+    // on a batch that was constructed from EvalLoader
+    // putting this functionality here because it is tightly coupled
+    // with how we construct and represent the data batches.
+    // returns the number of correct examples in this batch.
+    int correct = 0;
+    size_t B = loader->B;
+    size_t T = loader->T;
+    // iterate the examples in this batch
+    int can_fit_examples = B / ASSUMED_NUM_COMPLETIONS;
+    for (int i = 0; i < can_fit_examples; i++) {
+        float min_loss;
+        int min_loss_index;
+        char active = 0; // is this example active or fully empty?
+        // iterate the completions in this example
+        for (int b = 0; b < ASSUMED_NUM_COMPLETIONS; b++) {
+            int boff = i * ASSUMED_NUM_COMPLETIONS + b;
+            // evaluate the quality of this completion
+            // its quality is simply the average loss over the tokens
+            float average_loss = 0.0f;
+            int count = 0;
+            for (int t = 0; t < T; t++) {
+                char mask = loader->mask[boff * T + t];
+                if (mask == 1) {
+                    active = 1;
+                    average_loss += losses[boff * T + t];
+                    count++;
+                }
+            }
+            if (count > 0) { average_loss /= count; }
+            if (b == 0 || average_loss < min_loss) {
+                min_loss = average_loss;
+                min_loss_index = b;
+            }
+        }
+        if (active && (min_loss_index == loader->label[i])) {
+            correct += 1;
+        }
+    }
+    return correct;
+}
+
+void evalloader_free(EvalLoader *loader) {
+    free(loader->buffer);
+    free(loader->inputs);
+    free(loader->targets);
+    free(loader->mask);
+    free(loader->label);
+    fcloseCheck(loader->eval_file);
+}

dev/data/data_common.py

@@ -45,3 +45,64 @@ def write_datafile(filename, toks):
     with open(filename, "wb") as f:
         f.write(header.tobytes())
         f.write(toks_np.tobytes())
+
+def write_evalfile(filename, datas):
+    """
+    Saves eval data as a .bin file, for reading in C.
+    Used for multiple-choice style evals, e.g. HellaSwag and MMLU
+    - First comes a header with 256 int32s
+    - The examples follow, each example is a stream of uint16_t:
+        - <START_EXAMPLE> delimiter of 2**16-1, i.e. 65,535
+        - <EXAMPLE_BYTES>, bytes encoding this example, allowing efficient skip to next
+        - <EXAMPLE_INDEX>, the index of the example in the dataset
+        - <LABEL>, the index of the correct completion
+        - <NUM_COMPLETIONS>, indicating the number of completions (usually 4)
+        - <NUM><CONTEXT_TOKENS>, where <NUM> is the number of tokens in the context
+        - <NUM><COMPLETION_TOKENS>, repeated NUM_COMPLETIONS times
+    """
+    # construct the header
+    header = np.zeros(256, dtype=np.int32)
+    header[0] = 20240522 # magic
+    header[1] = 1 # version
+    header[2] = len(datas) # number of examples
+    header[3] = 0 # reserved for longest_example_bytes, fill in later
+    # now write the individual examples
+    longest_example_bytes = 0 # in units of uint16s
+    full_stream = [] # the stream of uint16s, we'll write a single time at the end
+    assert len(datas) < 2**16, "too many examples?"
+    for idx, data in enumerate(datas):
+        stream = []
+        # header of the example
+        stream.append(2**16-1) # <START_EXAMPLE>
+        stream.append(0) # <EXAMPLE_BYTES> (fill in later)
+        stream.append(idx) # <EXAMPLE_INDEX>
+        stream.append(data["label"]) # <LABEL>
+        ending_tokens = data["ending_tokens"]
+        assert len(ending_tokens) == 4, "expected 4 completions for now? can relax later"
+        stream.append(len(ending_tokens)) # <NUM_COMPLETIONS>
+        # the (shared) context tokens
+        ctx_tokens = data["ctx_tokens"]
+        assert all(0 <= t < 2**16-1 for t in ctx_tokens), "bad context token"
+        stream.append(len(ctx_tokens))
+        stream.extend(ctx_tokens)
+        # the completion tokens
+        for end_tokens in ending_tokens:
+            assert all(0 <= t < 2**16-1 for t in end_tokens), "bad completion token"
+            stream.append(len(end_tokens))
+            stream.extend(end_tokens)
+        # write to full stream
+        nbytes = len(stream)*2 # 2 bytes per uint16
+        assert nbytes < 2**16, "example too large?"
+        stream[1] = nbytes # fill in the <EXAMPLE_BYTES> field
+        longest_example_bytes = max(longest_example_bytes, nbytes)
+        full_stream.extend(stream)
+    # construct the numpy array
+    stream_np = np.array(full_stream, dtype=np.uint16)
+    # fill in the longest_example field
+    assert 0 < longest_example_bytes < 2**16, f"bad longest_example"
+    header[3] = longest_example_bytes
+    # write to file (for HellaSwag val this is 10,042 examples, 3.6MB file)
+    print(f"writing {len(datas):,} examples to {filename}")
+    with open(filename, "wb") as f:
+        f.write(header.tobytes())
+        f.write(stream_np.tobytes())

dev/data/hellaswag.py

@@ -1,6 +1,7 @@
 """
 Downloads and evaluates HellaSwag in Python.
 This then acts as the reference file for llm.c
+Also writes the data (tokens, labels) to .bin files for parallel evaluation in C.
 https://github.com/rowanz/hellaswag
 
 Example HellaSwag json item:
@@ -22,6 +23,8 @@ gpt2 (124M)
 gpt2-xl (1558M)
 - eleuther harness reports acc 40.04%, acc_norm 50.89% (multiple choice style)
 - this script: 10042 acc: 0.3842 acc_norm: 0.4893 (completion style)
+
+The validation set of HellaSwag has a total of 10,042 examples.
 """
 
 import os
@@ -33,7 +36,7 @@ import torch
 import torch.nn as nn
 from torch.nn import functional as F
 from transformers import GPT2LMHeadModel
-from data_common import download_file
+from data_common import download_file, write_evalfile
 
 # -----------------------------------------------------------------------------
 DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), "hellaswag")
@@ -68,14 +71,23 @@ def render_example(example):
     label = example["label"]
     endings = example["endings"]
 
+    # data needed to reproduce this eval on the C size
+    data = {
+        "label": label,
+        "ctx_tokens": None,
+        "ending_tokens": [],
+    }
+
     # gather up all the tokens
     ctx_tokens = enc.encode(ctx)
+    data["ctx_tokens"] = ctx_tokens
     tok_rows = []
     mask_rows = []
     for end in endings:
         end_tokens = enc.encode(" " + end) # note: prepending " " because GPT-2 tokenizer
         tok_rows.append(ctx_tokens + end_tokens)
         mask_rows.append([0]*len(ctx_tokens) + [1]*len(end_tokens))
+        data["ending_tokens"].append(end_tokens)
 
     # have to be careful during the collation because the number of tokens in each row can differ
     max_len = max(len(row) for row in tok_rows)
@@ -85,11 +97,10 @@ def render_example(example):
         tokens[i, :len(tok_row)] = torch.tensor(tok_row)
         mask[i, :len(mask_row)] = torch.tensor(mask_row)
 
-    return tokens, mask, label
+    return data, tokens, mask, label
 
 def iterate_examples(split):
     # there are 10,042 examples in total in val
-
     download(split)
     with open(os.path.join(DATA_CACHE_DIR, f"hellaswag_{split}.jsonl"), "r") as f:
         for line in f:
@@ -105,11 +116,13 @@ def evaluate(model_type, device):
     model.to(device)
     # model = torch.compile(model)
 
+    datas = []
     num_correct_norm = 0
     num_correct = 0
     num_total = 0
     for example in iterate_examples("val"):
-        tokens, mask, label = render_example(example)
+        data, tokens, mask, label = render_example(example)
+        datas.append(data)
         tokens = tokens.to(device)
         mask = mask.to(device)
 
@@ -137,7 +150,7 @@ def evaluate(model_type, device):
         num_total += 1
         num_correct += int(pred == label)
         num_correct_norm += int(pred_norm == label)
-        print(f"{num_total} acc: {num_correct/num_total:.4f} acc_norm: {num_correct_norm/num_total:.4f}")
+        print(f"{num_total} acc: {num_correct/num_total:.4f} acc_norm: {num_correct_norm}/{num_total}={num_correct_norm/num_total:.4f}")
 
         # debug: pretty print a few examples, and the losses in each case
         if num_total < 10:
@@ -146,7 +159,11 @@ def evaluate(model_type, device):
             print(f"Endings:")
             for i, end in enumerate(example["endings"]):
                 print(f"{i} (loss: {avg_loss[i].item():.4f}) {end}")
-            print(f"predicted: {pred}, actual: {label}")
+            print(f"predicted: {pred_norm}, actual: {label}")
+
+    # now write the data to a .bin file
+    filename = os.path.join(DATA_CACHE_DIR, f"hellaswag_val.bin")
+    write_evalfile(filename, datas)
 
 if __name__ == "__main__":
     import argparse

train_gpt2.cu

@@ -55,6 +55,7 @@ This reads & runs in fp32, B=4, T=64, LR=1e-4, val/sample never (200),
 // defines: tokenizer_init, tokenizer_decode, tokenizer_free
 #include "tokenizer.h"
 // defines: dataloader_init, dataloader_reset, dataloader_next_batch, dataloader_free
+// defines: evalloader_init, evalloader_reset, evalloader_next_batch, evalloader_free
 #include "dataloader.h"
 
 // ----------------------------------------------------------------------------
@@ -1946,6 +1947,7 @@ typedef struct {
     float mean_loss; // after a forward pass with targets, will be populated with the mean loss
     float accumulated_mean_loss; // Mean loss after aggregating it on all GPUs
     floatX* cpu_losses; // CPU buffer to copy the losses to, allocated with cudaMallocHost
+    float* cpu_losses_fp32; // same but fp32
     unsigned long long rng_state; // the RNG state for seeding stochastic rounding etc.
     int use_master_weights;
     int recompute;
@@ -2024,6 +2026,7 @@ void gpt2_build_from_checkpoint(GPT2 *model, const char* checkpoint_path) {
     model->inputs = NULL;
     model->targets = NULL;
     model->cpu_losses = NULL;
+    model->cpu_losses_fp32 = NULL;
     model->batch_size = 0;
     model->seq_len = 0;
     model->mean_loss = -1.0f; // -1.0f will designate no loss
@@ -2077,6 +2080,7 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in
         cudaCheck(cudaMalloc((void**)&model->inputs, B * T * sizeof(int)));
         cudaCheck(cudaMalloc((void**)&model->targets, B * T * sizeof(int)));
         cudaCheck(cudaMallocHost((void**)&model->cpu_losses, B * T * sizeof(floatX)));
+        cudaCheck(cudaMallocHost((void**)&model->cpu_losses_fp32, B * T * sizeof(float)));
     } else {
         // validate B,T is consistent with how we've allocated the memory before
         // in principle we could get more clever here in the future, for now this is safest
@@ -2181,7 +2185,11 @@ void gpt2_forward(GPT2 *model, int* inputs, int* targets, size_t B, size_t T, in
         // for convenience also evaluate the mean loss (TODO re-think this compute+sync point)
         cudaCheck(cudaMemcpy(model->cpu_losses, acts.losses, B * T * sizeof(floatX), cudaMemcpyDeviceToHost));
         float mean_loss = 0.0f;
-        for (int i = 0; i < B*T; i++) { mean_loss += (float)(model->cpu_losses[i]); }
+        for (int i = 0; i < B*T; i++) {
+            float loss = (float)(model->cpu_losses[i]);
+            model->cpu_losses_fp32[i] = loss;
+            mean_loss += loss;
+        }
         mean_loss /= B*T*grad_accum_steps;
         model->mean_loss = mean_loss;
     } else {
@@ -2337,13 +2345,13 @@ void gpt2_backward(GPT2 *model) {
     encoder_backward(grads.wte, grads.wpe, dresidual, model->inputs, B, T, C, random_u32(&model->rng_state));
 }
 
-// Compute a mean of a single CPU value across all GPU processes. No-op when multi-GPU is disabled.
-float multi_gpu_cpu_float_mean(float value, const MultiGpuConfig* multi_gpu_config) {
+// Compute sum of a single CPU value across all GPU processes. No-op when multi-GPU is disabled.
+float multi_gpu_cpu_float_sum(float value) {
 #ifdef MULTI_GPU
-    // MPI doesn't support all reduce with mean, so we sum up, then divide.
+    // note MPI doesn't support all reduce with mean, only sum
     float result;
     mpiCheck(MPI_Allreduce(&value, &result, 1, MPI_FLOAT, MPI_SUM, MPI_COMM_WORLD));
-    return result / multi_gpu_config->num_processes;
+    return result;
 #else
     return value;
 #endif
@@ -2356,7 +2364,7 @@ void gpt2_multi_gpu_accumulate(GPT2* model, MultiGpuConfig* multi_gpu_config) {
     NVTX_RANGE_FN();
     if (multi_gpu_config->num_processes == 1) { return; }
     // Average all losses.
-    model->accumulated_mean_loss = multi_gpu_cpu_float_mean(model->mean_loss, multi_gpu_config);
+    model->accumulated_mean_loss = multi_gpu_cpu_float_sum(model->mean_loss) / multi_gpu_config->num_processes;
     if(multi_gpu_config->zero_stage == 0) {
         //  no ZERO == standard DDP: Average all gradients.
         ncclCheck(ncclAllReduce(model->grads_memory, model->grads_memory,
@@ -2450,6 +2458,7 @@ void gpt2_free(GPT2 *model) {
     cudaCheck(cudaFree(model->inputs));
     cudaCheck(cudaFree(model->targets));
     cudaFreeHost(model->cpu_losses);
+    cudaFreeHost(model->cpu_losses_fp32);
 }
 
 // ----------------------------------------------------------------------------
@@ -2520,6 +2529,12 @@ void logger_init(Logger *logger, const char *filename) {
     if (filename != NULL) { logger->logfile = fopenCheck(filename, "w"); }
 }
 
+void logger_log_eval(Logger *logger, int step, float val_loss) {
+    if (logger->logfile != NULL) {
+        fprintf(logger->logfile, "s:%d eval:%.4f\n", step, val_loss);
+    }
+}
+
 void logger_log_val(Logger *logger, int step, float val_loss) {
     if (logger->logfile != NULL) {
         fprintf(logger->logfile, "s:%d tel:%.4f\n", step, val_loss);
@@ -2676,12 +2691,27 @@ int main(int argc, char *argv[]) {
     printf0("| val_num_batches       | %-50d |\n", val_num_batches);
     printf0("+-----------------------+----------------------------------------------------+\n");
 
+    // build an EvalLoader for HellaSwag
+    EvalLoader eval_loader;
+    const char* hellaswag_path = "dev/data/hellaswag/hellaswag_val.bin";
+    const char hellaswag_available = access(hellaswag_path, F_OK) == 0;
+    if (hellaswag_available) {
+        evalloader_init(&eval_loader, hellaswag_path, B, T, multi_gpu_config.process_rank, multi_gpu_config.num_processes);
+    }
+    printf0("| hellaswag available   | %-50s |\n", hellaswag_available ? "yes" : "no");
+    printf0("+-----------------------+----------------------------------------------------+\n");
+
     // pretty print in a table the multi-gpu configuration as well
     set_zero_configs(&multi_gpu_config, zero_stage, model.num_parameters);
     printf0("| num_processes         | %-50d |\n", multi_gpu_config.num_processes);
     printf0("| zero_stage            | %-50d |\n", multi_gpu_config.zero_stage);
     printf0("+-----------------------+----------------------------------------------------+\n");
 
+    // prints outside of pretty table to here and below
+    if (!hellaswag_available) {
+        printf0("HellaSwag eval not found at %s, skipping its evaluation\n", hellaswag_path);
+        printf0("You can run `python dev/data/hellaswag.py` to export and use it.\n");
+    }
     // more prints related to allocations from gpt2_build_from_checkpoint down here to not mess up our table above
     printf0("num_parameters: %zu => bytes: %zu\n", model.num_parameters, model.num_parameters_bytes);
     printf0("allocated %d MiB for model parameters\n", (int)round(model.num_parameters_bytes / (1024 * 1024)));
@@ -2729,11 +2759,30 @@ int main(int argc, char *argv[]) {
                 val_loss += model.mean_loss;
             }
             val_loss /= val_num_batches;
-            val_loss = multi_gpu_cpu_float_mean(val_loss, &multi_gpu_config);
+            val_loss = multi_gpu_cpu_float_sum(val_loss) / multi_gpu_config.num_processes;
             printf0("val loss %f\n", val_loss);
             logger_log_val(&logger, step, val_loss);
         }
 
+        // once in a while estimate HellaSwag accuracy
+        if (hellaswag_available &&
+           (step % val_loss_every == 0 || last_step)) {
+            NvtxRange evaluation_range("evaluation");
+            float eval_acc_norm = 0.0f;
+            evalloader_reset(&eval_loader);
+            for (int i = 0; i < eval_loader.num_batches; i++) {
+                if (i % 10 == 0) { printf("evaluating HellaSwag: %d/%d\r", i, eval_loader.num_batches); }
+                evalloader_next_batch(&eval_loader);
+                gpt2_forward(&model, eval_loader.inputs, eval_loader.targets, B, T);
+                int correct = evalloader_stat_losses(&eval_loader, model.cpu_losses_fp32);
+                eval_acc_norm += (float)correct;
+            }
+            // careful because not all ranks may have the exact same allocation of number of examples
+            eval_acc_norm = multi_gpu_cpu_float_sum(eval_acc_norm);
+            printf0("HellaSwag: %d/%d = %f\n", (int)eval_acc_norm, eval_loader.num_examples, eval_acc_norm / eval_loader.num_examples);
+            logger_log_eval(&logger, step, eval_acc_norm);
+        }
+
         // once in a while do model inference to print generated text
         if (multi_gpu_config.process_rank == 0 && (step > 0 && (step % sample_every) == 0 || last_step)) {
             NvtxRange generation_range("generation");
@@ -2844,6 +2893,7 @@ int main(int argc, char *argv[]) {
     // free and destroy everything
     cudaCheck(cudaEventDestroy(end));
     cudaCheck(cudaEventDestroy(start));
+    evalloader_free(&eval_loader);
     dataloader_free(&train_loader);
     dataloader_free(&val_loader);
     tokenizer_free(&tokenizer);
@@ -2852,7 +2902,6 @@ int main(int argc, char *argv[]) {
     free(gen_tokens);
     logger_free(&logger);
     multi_gpu_config_free(&multi_gpu_config);
-
     common_free(model);
     return 0;
 }