From 3edab5b70e393e8cd47aa205edb29596fa9ecbd7 Mon Sep 17 00:00:00 2001
From: Raju Kanumuri <vk2646@nyu.edu>
Date: Sat, 6 Dec 2025 22:31:47 -0500
Subject: [PATCH] fix 4

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

diff --git a/nvshmem_train_gpt2.cu b/nvshmem_train_gpt2.cu
index 07ef66b..9164f07 100644
--- a/nvshmem_train_gpt2.cu
+++ b/nvshmem_train_gpt2.cu
@@ -1435,6 +1435,15 @@ void gpt2_forward(GPT2 *model, int *inputs, int *targets, int B, int T) {
     nvshmem_putmem(model->nvshmem_act_buffer, layer5_output,
                    B * T * C * sizeof(float), 1); // PE 1
     nvshmem_quiet();
+
+    // GPU 0 doesn't compute loss, but needs to set mean_loss for backward pass
+    // If targets provided, set to a valid value (GPU 1 has the actual loss)
+    // If no targets (inference), set to -1.0
+    if (targets != NULL) {
+      model->mean_loss = 0.0f; // Valid placeholder for backward check
+    } else {
+      model->mean_loss = -1.0f; // No targets, inference mode
+    }
   }
 
   // GPU 1: Layers 6-11 + Final LayerNorm + Loss