Add kaggle/train_grpo.py

kaggle/train_grpo.py

@@ -782,23 +782,24 @@ 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()
+    # T4 (sm75) cannot run TRL's bf16 chunked logprob kernels (need sm80+).
+    # Also: fp16 AMP GradScaler crashes if any grads stay bf16
+    # (`_amp_foreach_non_finite_check_and_unscale_cuda` not implemented for BFloat16).
+    # On T4: no AMP (fp16=False, bf16=False); keep 4-bit weights + fp16 compute dtype.
+    major, minor = (0, 0)
+    if torch.cuda.is_available():
+        major, minor = torch.cuda.get_device_capability(0)
+    use_bf16 = major >= 8 and torch.cuda.is_bf16_supported()
+    use_fp16 = False  # avoid GradScaler + QLoRA dtype mismatch on T4
     compute_dtype = torch.bfloat16 if use_bf16 else torch.float16
+    _vlog(
+        "main.dtype_choice",
+        bf16=use_bf16,
+        fp16=use_fp16,
+        compute_dtype=str(compute_dtype),
+        capability=(major, minor),
+        device=torch.cuda.get_device_name(0) if torch.cuda.is_available() else None,
+    )
     bnb = BitsAndBytesConfig(
         load_in_4bit=True,
         bnb_4bit_compute_dtype=compute_dtype,
@@ -840,6 +841,11 @@ def main(cfg: Cfg | None = None) -> None:
         ],
     )
     model = get_peft_model(model, lora)
+    # Keep trainable LoRA weights in fp32 so optimizer/grad-norm never sees bf16 under AMP-off.
+    for name, param in model.named_parameters():
+        if param.requires_grad and param.dtype != torch.float32:
+            _vlog("main.cast_trainable_fp32", name=name, from_dtype=str(param.dtype))
+            param.data = param.data.float()
 
     output_dir = "/kaggle/working/grpo-out"
     Path(output_dir).mkdir(parents=True, exist_ok=True)
@@ -936,7 +942,7 @@ def main(cfg: Cfg | None = None) -> None:
         save_total_limit=3,
         report_to=["wandb"],
         bf16=use_bf16,
-        fp16=not use_bf16,
+        fp16=use_fp16,
         remove_unused_columns=False,
         hub_model_id=cfg.hub_model_id,
         hub_token=os.environ["HF_TOKEN"],