From 2f261ad4d74d672a5693487f62cb5a87635f461c Mon Sep 17 00:00:00 2001
From: mpragnay <pm3881@nyu.edu>
Date: Wed, 10 Dec 2025 19:43:00 -0500
Subject: [PATCH] print loss on only one

---
 nvshmem_train_gpt2.cu | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nvshmem_train_gpt2.cu b/nvshmem_train_gpt2.cu
index 64b7e85..a6245c0 100644
--- a/nvshmem_train_gpt2.cu
+++ b/nvshmem_train_gpt2.cu
@@ -2271,9 +2271,11 @@ int main(int argc, char *argv[]) {
         (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9;
     total_sum_iteration_time_s += time_elapsed_s;
     int tokens_per_second = (B * T) / time_elapsed_s;
+    if (my_pe == 1) {
     printf("step %4d/%d: train loss %f (%f ms, %d tok/s)\n", step + 1,
            train_num_batches, model.mean_loss, time_elapsed_s * 1000,
            tokens_per_second);
+    }
     logger_log_train(&logger, step, model.mean_loss);
   }
   // add a total average, for optimizations that are only mild improvements