Fix stacking of pos/neg logits for RM compute_metrics
diff --git a/.vscode/launch.json b/.vscode/launch.json
index 997ef81df5..7f68ec73ef 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -17,6 +17,55 @@
       "purpose": ["debug-test"],
       "console": "integratedTerminal",
       "justMyCode": true
+    },
+    {
+      "name": "Debug Training",
+      "type": "python",
+      "request": "launch",
+      "cwd": "${workspaceFolder}/model/model_training",
+      "program": "trainer_sft.py",
+      "args": ["--configs", "defaults", "oasst_export_eu", "debug"],
+      "console": "integratedTerminal",
+      "justMyCode": true,
+      "env": {
+        "TOKENIZERS_PARALLELISM": "false",
+        "CUDA_VISIBLE_DEVICES": "1,2,3,4,5,6,7"
+      }
+    },
+    {
+      "name": "Debug RM Training",
+      "type": "python",
+      "request": "launch",
+      "cwd": "${workspaceFolder}/model/model_training",
+      "program": "trainer_rm.py",
+      "args": [
+        "--configs",
+        "defaults_rm",
+        "oasst_export_eu",
+        "pythia-1B",
+        "--cache_dir",
+        "/home/ubuntu/data_cache",
+        "--use_flash_attention",
+        "true"
+      ],
+      "console": "integratedTerminal",
+      "justMyCode": true,
+      "env": {
+        "TOKENIZERS_PARALLELISM": "false",
+        "CUDA_VISIBLE_DEVICES": "1"
+      }
+    },
+    {
+      "name": "Debug Patching",
+      "type": "python",
+      "request": "launch",
+      "cwd": "${workspaceFolder}/model/model_training",
+      "program": "models/test_patched_gpt_neox.py",
+      "console": "integratedTerminal",
+      "justMyCode": true,
+      "env": {
+        "CUDA_VISIBLE_DEVICES": "0"
+      }
     }
   ]
 }
diff --git a/model/model_training/trainer_rm.py b/model/model_training/trainer_rm.py
index 9189f50b5e..9d8770620c 100644
--- a/model/model_training/trainer_rm.py
+++ b/model/model_training/trainer_rm.py
@@ -22,7 +22,7 @@
 
 def compute_metrics(eval_pred):
     scores = eval_pred.predictions
-    pos_scores, neg_scores = scores[0::2, :], scores[1::2, :]
+    pos_scores, neg_scores = scores[:, 0], scores[:, 1]
     metrics = {
         "pos_score": np.mean(pos_scores),
         "neg_score": np.mean(neg_scores),
@@ -81,9 +81,9 @@ def prediction_step(
         pos_logits = torch.cat(pos_logits).detach()
         neg_logits = torch.cat(neg_logits).detach()
 
-        out_logits = torch.stack([pos_logits, neg_logits])
-        # we can't pass something (not nothing) for `compute_metrics` to be called`
-        # also, it has to have the same size as logits, otherwise `_pad_across_processes` hangs
+        out_logits = torch.stack([pos_logits, neg_logits], dim=1)  # shape (B, 2)
+        # need to pass something for `compute_metrics` to be called`,
+        # has to have the same size as logits, otherwise `_pad_across_processes` hangs
         labels = torch.zeros_like(out_logits[:, 0])
 
         return (loss, out_logits, labels)

Node positive & negative examples are stacked along last dimension so that compute_metrics receives predictions of size (val_examles, 2). The old code did not work with deepspeed if the number of validation examples was not divisible by the total batch size without remainder.