Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -782,9 +782,26 @@ def main(cfg: Cfg | None = None) -> None:
     train, eval_problems = _load_swe_splits(cfg)
     _vlog("main.dataset_ready", train_n=len(train), eval_n=len(eval_problems))
 
+    # T4 is sm75 — TRL chunked bf16 logprob kernels need sm80+ (A100/L4/…).
+    # torch.cuda.is_bf16_supported() can still be True on T4 and then crash in backward.
+    def _use_bf16() -> bool:
+        if not torch.cuda.is_available():
+            return False
+        major, _minor = torch.cuda.get_device_capability(0)
+        ok = major >= 8 and torch.cuda.is_bf16_supported()
+        _vlog(
+            "main.dtype_choice",
+            bf16=ok,
+            capability=torch.cuda.get_device_capability(0),
+            device=torch.cuda.get_device_name(0),
+        )
+        return ok
+
+    use_bf16 = _use_bf16()
+    compute_dtype = torch.bfloat16 if use_bf16 else torch.float16
     bnb = BitsAndBytesConfig(
         load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
+        bnb_4bit_compute_dtype=compute_dtype,
         bnb_4bit_quant_type="nf4",
         bnb_4bit_use_double_quant=True,
     )
@@ -918,8 +935,8 @@ def main(cfg: Cfg | None = None) -> None:
         save_steps=cfg.save_steps,
         save_total_limit=3,
         report_to=["wandb"],
-        bf16=torch.cuda.is_bf16_supported(),
-        fp16=not torch.cuda.is_bf16_supported(),
+        bf16=use_bf16,
+        fp16=not use_bf16,
         remove_unused_columns=False,
         hub_model_id=cfg.hub_model_id,
         hub_token=os.environ["HF_TOKEN"],