Keep training stack parity: CPU fallback path and mixed-precision safeguards.

scripts/train_sota.py

@@ -432,8 +432,12 @@ def build_model_and_tokenizer(model_cfg: Dict[str, Any], training_defaults: Dict
     if not base_model:
         raise ValueError("model.base_model is required.")
 
-    use_bf16 = bool(model_cfg.get("use_bf16", True))
-    dtype = torch.bfloat16 if use_bf16 else torch.float16
+    use_cuda = torch.cuda.is_available()
+    requested_bf16 = bool(model_cfg.get("use_bf16", True))
+    if use_cuda:
+        dtype = torch.bfloat16 if requested_bf16 else torch.float16
+    else:
+        dtype = torch.float32
 
     tokenizer = build_tokenizer(model_cfg)
 
@@ -445,10 +449,11 @@ def build_model_and_tokenizer(model_cfg: Dict[str, Any], training_defaults: Dict
     if attn_impl:
         model_kwargs["attn_implementation"] = attn_impl
 
-    load_in_4bit = bool(model_cfg.get("load_in_4bit", True))
+    requested_load_in_4bit = bool(model_cfg.get("load_in_4bit", True))
+    load_in_4bit = requested_load_in_4bit and use_cuda
+    if requested_load_in_4bit and not load_in_4bit:
+        print("CUDA unavailable. Disabling 4-bit loading and using full-precision CPU fallback.")
     if load_in_4bit:
-        if not torch.cuda.is_available():
-            raise RuntimeError("4-bit loading requested but CUDA is not available.")
         model_kwargs["quantization_config"] = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type=as_text(model_cfg.get("bnb_4bit_quant_type")) or "nf4",
@@ -565,6 +570,9 @@ def build_training_args(
     has_eval_split: bool,
 ) -> TrainingArguments:
     output_dir.mkdir(parents=True, exist_ok=True)
+    use_cuda = torch.cuda.is_available()
+    bf16_runtime = bool(use_cuda and use_bf16)
+    fp16_runtime = bool(use_cuda and not bf16_runtime)
     return TrainingArguments(
         output_dir=str(output_dir),
         num_train_epochs=as_float(training_cfg.get("num_train_epochs"), 1.0),
@@ -582,8 +590,8 @@ def build_training_args(
         save_total_limit=as_int(training_cfg.get("save_total_limit"), 3),
         dataloader_num_workers=as_int(training_cfg.get("dataloader_num_workers"), 0),
         seed=as_int(training_cfg.get("seed"), 17),
-        bf16=use_bf16,
-        fp16=not use_bf16,
+        bf16=bf16_runtime,
+        fp16=fp16_runtime,
         remove_unused_columns=False,
         report_to="none",
         evaluation_strategy="steps" if has_eval_split else "no",
@@ -860,6 +868,10 @@ def main() -> None:
         model = None
     else:
         model, tokenizer = build_model_and_tokenizer(cfg["model"], cfg.get("training_defaults", {}))
+        if torch.cuda.is_available():
+            print("Compute mode: GPU")
+        else:
+            print("Compute mode: CPU fallback (no CUDA detected)")
 
     data_cfg = cfg["data"]
     stage_reports: List[Dict[str, Any]] = []