Merge branch 'karpathy:master' into train_gp2cu-porting-fixes

profile_gpt2cu.py

@@ -39,7 +39,8 @@ metrics = [
     "dram__bytes_write.sum",                    # DRAM writes
     "lts__t_sectors_srcunit_tex_op_read.sum",   # L2 reads (sectors -- 32B)
     "lts__t_sectors_srcunit_tex_op_write.sum",  # L2 reads (sectors -- 32B)
-    "smsp__inst_executed.sum",                   # instructions
+    "sm__pipe_tensor_op_hmma_cycles_active.avg.pct_of_peak_sustained_active", # % of peak tensor core utilization
+    "smsp__inst_executed.sum",                  # instructions
 ]
 cmd = [NCU, "-i", "profile.ncu-rep", "--csv", "--page", "raw", "--metrics", ",".join(metrics)]
 result = subprocess.check_output(cmd, text=True).strip()
@@ -72,16 +73,35 @@ for rid, row in kernel_profile_data:
 
 assert CLS_START != -1
 
+# Check every kernel to find the maximum DRAM bandwidth and Tensor Core utilisation values
+max_dram_bw = 0.0
+max_tensor = 0.0
+for rid, row in kernel_profile_data:
+    if rid <= 2:
+        continue
+    time = float(row[13])
+    read = float(row[11])
+    write = float(row[12])
+    tensor = float(row[16])
+    dram_bw = (read + write) / (time / 1000.0)
+    max_dram_bw = max(max_dram_bw, dram_bw)
+    max_tensor = max(max_tensor, tensor)
+
+# round the maximum tensor core utilisation to 50% or 100%
+# consumer GPUs can only achieve 50% of peak tensor throughput on this counter
+# and for GPUs without tensor cores, we set the value to 50% to avoid division by zero
+max_tensor = (max_tensor > 50.0) and 100.0 or 50.0
+
 print()
 print("Kernel calls:")
 for rid, row in kernel_profile_data:
     if rid == 0:
         #  headings
-        print(f"id pass    {'name':<40} {'time':>8} {'RAM rd':>8} {'RAM wt':>8} {'L2 rd':>8} {'L2 wt':>8} {'inst':>8}")
+        print(  f"id pass    {'name':<40} {'time':>8} {'RAM BW':>8} {'tensor':>8} {'RAM rd':>8} {'RAM wt':>8} {'L2 rd':>8} {'L2 wt':>8} {'inst':>8}")
         continue
     if rid == 1:
         # units
-        units = f"           {'':<40} {'ms':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'MInst':>8}"
+        units = f"           {'':<40} {'ms':>8} {'GB/s':>8} {'core %':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'GiB':>8} {'MInst':>8}"
         print(units)
         print("." * len(units))
         continue
@@ -95,7 +115,9 @@ for rid, row in kernel_profile_data:
     write = float(row[12])
     l2_read = float(row[14])
     l2_write = float(row[15])
-    inst = float(row[16]) / 1e6
+    tensor = float(row[16])
+    inst = float(row[17]) / 1e6
+    dram_bw = (read + write) / (time / 1000.0)
 
     kid = rid - 2
 
@@ -149,6 +171,7 @@ for rid, row in kernel_profile_data:
     l2_read = l2_read * 32 / 1024 / 1024 / 1024
     l2_write = l2_write * 32 / 1024 / 1024 / 1024
 
+    efficiency = max(dram_bw / max_dram_bw, tensor / max_tensor)
     summaries[fn_name] += time
     counts[fn_name] += multiplier
     passes[pass_name] += time
@@ -159,13 +182,18 @@ for rid, row in kernel_profile_data:
         total['l2_read'] += l2_read
         total['l2_write'] += l2_write
         total['inst'] += inst
+        total['tensor'] += tensor * time # % so multiplied by time
+        total['efficiency'] += efficiency * time
 
     pass_info = f"{pass_name}×{multiplier}"
-    print(f"{kid:02} {pass_info:7} {fn_name:<40} {time:8.2f} {read:8.2f} {write:8.2f} {l2_read:8.2f} {l2_write:8.2f} {inst:8.2f}")
+    print(f"{kid:02} {pass_info:7} {fn_name:<40} {time:8.2f} {dram_bw:8.1f} {tensor:8.1f} {read:8.2f} {write:8.2f} {l2_read:8.2f} {l2_write:8.2f} {inst:8.2f}")
+
 
 total_time = total['time']
+avg_dram_bw = (total['read'] + total['write']) / (total_time / 1000.0)
+avg_tensor_util = total['tensor'] / total_time
 print("." * len(units))
-print(f"           {'Total':<40} {total['time']:8.2f} {total['read']:8.2f} {total['write']:8.2f} {total['l2_read']:8.2f} {total['l2_write']:8.2f} {total['inst']:8.2f}")
+print(f"           {'Total':<40} {total['time']:8.2f} {avg_dram_bw:8.1f} {avg_tensor_util:8.1f} {total['read']:8.2f} {total['write']:8.2f} {total['l2_read']:8.2f} {total['l2_write']:8.2f} {total['inst']:8.2f}")
 
 print()
 print("Kernel type summaries:")
@@ -192,5 +220,9 @@ In total, a training step takes {total_time:.1f}ms, distributed as:
 We read {total['read']:.1f}GiB ({total['read']/ts:.1f}GB/s) and write {total['write']:.1f}GiB ({total['write']/ts:.1f}GB/s) to DRAM,
 read {total['l2_read']:.1f}GiB ({total['l2_read']/ts:.1f}GB/s) and write {total['l2_write']:.1f}GiB ({total['l2_write']/ts:.1f}GB/s) to L2,
 and execute {total['inst'] / 1000:.1f} billion instructions ({total['inst'] / 1000 / ts:.1f} GInst/s).
+
+Assuming that every kernel should be either fully DRAM bandwidth or tensor core limited,
+with a peak DRAM bandwidth of {max_dram_bw:.1f}GB/s and a peak tensor throughput of {max_tensor:.1f}%,
+our overall efficiency is {(total['efficiency'] * 100.0 / total_time):.1f}%.
 """
 print(summary)

train_gpt2.py

@@ -147,12 +147,11 @@ class GPT(nn.Module):
             # if we are given some desired targets also calculate the loss
             logits = self.lm_head(x)
             loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+            return None, loss
         else:
             # inference-time mini-optimization: only forward the lm_head on the very last position
             logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
-            loss = None
-
-        return logits, loss
+            return logits, None
 
     @classmethod
     def from_pretrained(cls, model_type):