push

train/colab_train.py

@@ -15,6 +15,7 @@ Instructions for Colab:
 
 import collections
 import random
+import unsloth  # Must be imported before trl/transformers/peft for patching.
 from datasets import Dataset
 import torch
 import numpy as np
@@ -37,6 +38,11 @@ env = AutomathreasonerEnv(url=HF_SPACE_URL)
 max_seq_length = 1024  # Fits well within Colab T4 16GB VRAM limit
 lora_rank = 16
 
+# T4 (and many non-Ampere GPUs) do not support bf16; pick precision dynamically.
+has_cuda = torch.cuda.is_available()
+use_bf16 = has_cuda and torch.cuda.is_bf16_supported()
+use_fp16 = has_cuda and not use_bf16
+
 # 2. Load Model via Unsloth (optimized for Free Colab VRAM)
 print("Loading model via Unsloth...")
 model, tokenizer = FastLanguageModel.from_pretrained(
@@ -191,6 +197,9 @@ training_args = GRPOConfig(
     
     # 8-bit optimizer saves VRAM
     optim="adamw_8bit",
+    bf16=use_bf16,
+    fp16=use_fp16,
+    use_cpu=not has_cuda,
 )
 
 trainer = GRPOTrainer(

train/train_grpo.py

@@ -1,5 +1,6 @@
 import random
 import collections
+import unsloth  # Must be imported before trl/transformers/peft for patching.
 import torch
 import numpy as np
 from datasets import Dataset
@@ -143,6 +144,9 @@ def run_ttrl(model, tokenizer, test_problem, env, steps=5):
 def main():
     max_seq_length = 1024
     lora_rank = 16
+    has_cuda = torch.cuda.is_available()
+    use_bf16 = has_cuda and torch.cuda.is_bf16_supported()
+    use_fp16 = has_cuda and not use_bf16
     
     # Load model via Unsloth
     model, tokenizer = FastLanguageModel.from_pretrained(
@@ -354,6 +358,9 @@ def main():
         
         # Optimizer
         optim="adamw_8bit",                 # Memory-efficient
+        bf16=use_bf16,
+        fp16=use_fp16,
+        use_cpu=not has_cuda,
     )
     
     trainer = GRPOTrainer(