Rewrite locally

app.py

@@ -1,8 +1,9 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, DataCollatorForSeq2Seq
+from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq
 from trl import SFTTrainer, SFTConfig
 from datasets import load_dataset
 from peft import LoraConfig, get_peft_model
 import torch
+import os
 
 # -------------------------
 # Load dataset
@@ -13,37 +14,56 @@ dataset = load_dataset("mavilov/convos", split="train")
 # Load model and tokenizer
 # -------------------------
 model_id = "swiss-ai/Apertus-8B-2509"
+model_kwargs = {}
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_compute_dtype=torch.float16,  # FP16 compute
-)
+if torch.backends.mps.is_available():
+    print("⚡ Using Apple MPS backend (Metal)")
+    model_kwargs = {
+        "dtype": torch.float16,
+        "device_map": {"": "mps"},  # force load directly on MPS
+        "offload_folder": "./offload",
+        "low_cpu_mem_usage": True,   # avoid meta tensors
+    }
+elif torch.cuda.is_available():
+    print("⚡ Using CUDA with bitsandbytes quantization")
+    from transformers import BitsAndBytesConfig
+    bnb_config = BitsAndBytesConfig(
+        load_in_8bit=True,
+        llm_int8_threshold=6.0
+    )
+    model_kwargs["quantization_config"] = bnb_config
+    model_kwargs["device_map"] = "auto"
+else:
+    print("⚠️ No GPU/MPS detected, running on CPU (very slow)")
+    model_kwargs = {
+        "dtype": torch.float32,
+        "device_map": {"": "cpu"},
+        "low_cpu_mem_usage": True,
+    }
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+tokenizer.pad_token = tokenizer.eos_token
 
+# Load model safely
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    quantization_config=bnb_config,
-    device_map="auto"
+    **model_kwargs
 )
 model.config.use_cache = False
 model.config.pretraining_tp = 1
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-tokenizer.pad_token = tokenizer.eos_token
-
 # -------------------------
 # Attach LoRA adapters
 # -------------------------
 lora_config = LoraConfig(
     r=16,
     lora_alpha=32,
-    target_modules=["q_proj", "v_proj"],  # typical attention modules
+    target_modules=["q_proj", "v_proj"],
     lora_dropout=0.05,
     bias="none",
     task_type="CAUSAL_LM"
 )
-
 model = get_peft_model(model, lora_config)
 
 # -------------------------
@@ -76,7 +96,7 @@ training_args = SFTConfig(
     num_train_epochs=3,
     logging_steps=10,
     report_to="tensorboard",
-    bf16=False,  # disable bf16
+    bf16=False,
 )
 
 # -------------------------

requirements.txt

@@ -1,8 +1,9 @@
-torch>=2.1.0
-transformers>=4.35.0
-datasets>=2.15.0
+torch
+transformers
+datasets
 accelerate>=0.26.0
-trl>=0.7.0
-bitsandbytes>=0.41.0
-peft>=0.5.0
-tensorboard>=2.15.0
+trl
+bitsandbytes
+peft
+tensorboard
+huggingface_hub