Upload folder using huggingface_hub

src/mathcompose/eval/runners.py

@@ -35,8 +35,15 @@ class HFRunner:
         self.max_context = max_context
         self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
 
-        torch_dtype = {"auto": "auto", "bf16": torch.bfloat16, "fp16": torch.float16,
-                       "fp32": torch.float32}.get(dtype, "auto")
+        if dtype == "auto":
+            if self.device == "cuda":
+                bf16_ok = torch.cuda.get_device_capability()[0] >= 8  # Ampere+; T4 is fp16
+                torch_dtype = torch.bfloat16 if bf16_ok else torch.float16
+            else:
+                torch_dtype = torch.float32
+        else:
+            torch_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16,
+                           "fp32": torch.float32}.get(dtype, torch.float32)
         # transformers 5.x: `dtype=` (torch_dtype is deprecated)
         model = AutoModelForCausalLM.from_pretrained(base_id, dtype=torch_dtype)
         if adapter_path:

src/mathcompose/train/lora.py

@@ -17,13 +17,14 @@ def build_lora_config(cfg: dict):
     )
 
 
-def build_bnb_config(cfg: dict):
-    """4-bit NF4 config. Requires bitsandbytes + CUDA."""
+def build_bnb_config(cfg: dict, compute_dtype=None):
+    """4-bit NF4 config. Requires bitsandbytes + CUDA. ``compute_dtype`` overrides
+    the config value (used to force fp16 on GPUs without bf16, e.g. T4)."""
     import torch
     from transformers import BitsAndBytesConfig
 
     q = cfg["quant"]
-    compute_dtype = getattr(torch, q.get("bnb_4bit_compute_dtype", "bfloat16"))
+    compute_dtype = compute_dtype or getattr(torch, q.get("bnb_4bit_compute_dtype", "bfloat16"))
     return BitsAndBytesConfig(
         load_in_4bit=q.get("load_in_4bit", True),
         bnb_4bit_quant_type=q.get("bnb_4bit_quant_type", "nf4"),

src/mathcompose/train/train.py

@@ -43,6 +43,13 @@ def build_and_train(args) -> str:
     val_file = args.val_file or cfg["data"].get("val_path")
 
     has_cuda = torch.cuda.is_available()
+    # bf16 needs Ampere+ (sm_80): A100/L4/H100 yes, T4 (sm_75) no -> use fp16 there.
+    bf16_ok = has_cuda and torch.cuda.get_device_capability()[0] >= 8
+    compute_dtype = torch.bfloat16 if bf16_ok else (torch.float16 if has_cuda else torch.float32)
+    if has_cuda:
+        cap = torch.cuda.get_device_capability()
+        print(f"[gpu] {torch.cuda.get_device_name(0)} (sm_{cap[0]}{cap[1]}) -> "
+              f"{'bf16' if bf16_ok else 'fp16'} compute")
     use_4bit = (not args.no_4bit) and cfg["quant"]["load_in_4bit"] and has_cuda
     if cfg["quant"]["load_in_4bit"] and not has_cuda and not args.no_4bit:
         print("[warn] 4-bit requested but no CUDA — falling back to full-precision "
@@ -50,12 +57,9 @@ def build_and_train(args) -> str:
 
     tok = load_tokenizer(base_id, pad_token=cfg["model"]["pad_token"])
 
-    model_kwargs = {}
+    model_kwargs = {"dtype": compute_dtype}
     if use_4bit:
-        model_kwargs["quantization_config"] = build_bnb_config(cfg)
-        model_kwargs["dtype"] = torch.bfloat16
-    else:
-        model_kwargs["dtype"] = torch.bfloat16 if has_cuda else torch.float32
+        model_kwargs["quantization_config"] = build_bnb_config(cfg, compute_dtype=compute_dtype)
     model = AutoModelForCausalLM.from_pretrained(base_id, **model_kwargs)
 
     data_files = {"train": train_file}
@@ -84,7 +88,8 @@ def build_and_train(args) -> str:
         packing=t.get("packing", False),
         gradient_checkpointing=grad_ckpt,
         optim=optim,
-        bf16=has_cuda,
+        bf16=bf16_ok,
+        fp16=(has_cuda and not bf16_ok),
         seed=t.get("seed", 0),
         report_to=(["wandb"] if args.wandb else "none"),
         max_steps=args.max_steps if args.max_steps else -1,