# ============================================================ # AYURVEDA GEMMA-3 4B TRAINING — FINAL v5 (Local RTX A4000) # Hardware: RTX A4000 16 GB VRAM | i9-13900K | 64 GB RAM | Windows 11 # Dataset: gemma4_ayurveda_unsloth_clean_WITH_REFS.jsonl (~971 K, ~1.7 GB) # OR gemma4_ayurveda_unsloth_500mb_FINAL.jsonl (~276 K) # Model: unsloth/gemma-3-4b-it + QLoRA # Framework: Unsloth + TRL # ============================================================ # RESTART KERNEL BEFORE RUNNING CELL 1 # ============================================================ # ============================================================ # CELL 0 — ONE-TIME SETUP (Fix CPU-only PyTorch if needed) # ============================================================ # Run this cell ONCE after kernel restart. # If it installs CUDA torch, it will tell you to restart again. # ----------------------------------------------------------- import subprocess import sys # Check current torch try: import torch has_torch = True is_cpu = (not torch.cuda.is_available()) and ("+cpu" in torch.__version__) except ImportError: has_torch = False is_cpu = True if has_torch and torch.cuda.is_available(): print("✓ PyTorch CUDA is already installed:", torch.__version__) print("Proceed to Cell 1.") else: print("Detected CPU-only or missing PyTorch. Installing CUDA-enabled version...") print("This will download ~2.5 GB. Please wait...") # Uninstall CPU torch if present subprocess.check_call( [sys.executable, "-m", "pip", "uninstall", "torch", "torchvision", "torchaudio", "-y"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) # Install CUDA 12.4 enabled torch (compatible with RTX A4000 driver) subprocess.check_call( [sys.executable, "-m", "pip", "install", "torch", "torchvision", "torchaudio", "--index-url", "https://download.pytorch.org/whl/cu124"], ) print("=" * 60) print("INSTALLATION COMPLETE") print("=" * 60) print("ACTION REQUIRED: Restart the kernel now (Kernel → Restart)") print("Then skip this cell and run Cell 1 directly.") print("=" * 60) raise SystemExit("Restart kernel to load CUDA PyTorch.") # ============================================================ # CELL 1 — ENVIRONMENT & CONFIG # ============================================================ import os # --- MUST BE SET BEFORE ANY IMPORTS (torch, transformers, datasets, unsloth) --- # All caches, downloads, and model weights go to D:\Sahil\ instead of C:\ BASE_DIR = r"D:\Sahil" os.makedirs(BASE_DIR, exist_ok=True) CACHE = os.path.join(BASE_DIR, "cache") os.makedirs(CACHE, exist_ok=True) os.makedirs(os.path.join(CACHE, "unsloth"), exist_ok=True) os.makedirs(os.path.join(CACHE, "torch"), exist_ok=True) os.makedirs(os.path.join(CACHE, "triton"), exist_ok=True) os.makedirs(os.path.join(CACHE, "huggingface"), exist_ok=True) os.makedirs(os.path.join(CACHE, "datasets"), exist_ok=True) DATASET_DIR = os.path.join(BASE_DIR, "dataset") os.makedirs(DATASET_DIR, exist_ok=True) OUTPUT_DIR = os.path.join(BASE_DIR, "output") os.makedirs(OUTPUT_DIR, exist_ok=True) # Unsloth / Triton / Torch caches os.environ["UNSLOTH_COMPILED_CACHE"] = os.path.join(CACHE, "unsloth") os.environ["UNSLOTH_CACHE_DIR"] = os.path.join(CACHE, "unsloth") os.environ["TORCHINDUCTOR_CACHE_DIR"] = os.path.join(CACHE, "torch") os.environ["TRITON_CACHE_DIR"] = os.path.join(CACHE, "triton") # HuggingFace: model weights, tokenizers, config downloads os.environ["HF_HOME"] = os.path.join(CACHE, "huggingface") os.environ["TRANSFORMERS_CACHE"] = os.path.join(CACHE, "huggingface") # datasets library cache os.environ["DATASETS_CACHE"] = os.path.join(CACHE, "datasets") os.environ["XDG_CACHE_HOME"] = CACHE # PyTorch memory allocator (prevents OOM on 16 GB VRAM) os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" # Optional: suppress wandb / telemetry since we don't use them os.environ["WANDB_DISABLED"] = "true" import torch import gc # --- CONFIG --- MODEL_NAME = "unsloth/gemma-3-4b-it" # Put your dataset file inside D:\Sahil\dataset\ # Change filename here if you use the 500MB file instead DATASET_PATH = os.path.join(DATASET_DIR, "gemma4_ayurveda_unsloth_clean_WITH_REFS.jsonl") # Training outputs go here OUTPUT_DIR = os.path.join(OUTPUT_DIR, "gemma3_4b_ayurveda_run") # --- SUBSET (0 = use full dataset) --- # For the full ~971 K / 1.7 GB dataset on RTX A4000, training will take ~30–40 hours. # If you want a quick test first, set SUBSET_SIZE = 50000, run Cells 0–9, # then restart kernel and set SUBSET_SIZE = 0 for the full run. SUBSET_SIZE = 0 MAX_SEQ_LENGTH = 512 # 512 is the safe limit for RTX A4000 16 GB (OOM at 1024) BATCH_SIZE = 1 GRAD_ACCUM = 16 # Effective batch = 16 LORA_R = 16 LORA_ALPHA = 16 LR = 2e-4 EPOCHS = 1 # --- GPU CHECK with helpful error message --- if not torch.cuda.is_available(): print("=" * 60) print("ERROR: CUDA GPU not found") print("=" * 60) print("PyTorch version:", torch.__version__) print() print("If this shows '+cpu' in the version above, run Cell 0 first,") print("then RESTART the kernel (Kernel → Restart), then run Cell 1.") print("=" * 60) raise AssertionError("CUDA GPU not found. See message above.") print("=" * 60) print("GPU :", torch.cuda.get_device_name(0)) print("VRAM :", round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2), "GB") print("Model :", MODEL_NAME) print("Dataset:", DATASET_PATH) print("Output :", OUTPUT_DIR) print("Subset :", f"{SUBSET_SIZE:,}" if SUBSET_SIZE else "FULL (~971 K)") print("=" * 60) print("Cell 1 complete → Run Cell 2") # ============================================================ # CELL 2 — LOAD DATASET + SHUFFLE + OPTIONAL SUBSET # ============================================================ from datasets import load_dataset dataset = load_dataset("json", data_files=DATASET_PATH, split="train") print(f"Loaded {len(dataset):,} conversations") # Shuffle before any split so the sample is representative dataset = dataset.shuffle(seed=3407) if SUBSET_SIZE and SUBSET_SIZE < len(dataset): dataset = dataset.select(range(SUBSET_SIZE)) print(f"Using subset: {len(dataset):,} conversations") sample = dataset[0] print("Roles :", [m["role"] for m in sample["messages"]]) print("\nUser preview:") print(sample["messages"][1]["content"][:200]) gc.collect() print("Cell 2 complete → Run Cell 3") # ============================================================ # CELL 3 — LOAD MODEL & TOKENIZER # ============================================================ from unsloth import FastLanguageModel print(f"Loading {MODEL_NAME} ...") try: model, tokenizer = FastLanguageModel.from_pretrained( model_name = MODEL_NAME, max_seq_length = MAX_SEQ_LENGTH, dtype = None, # Auto-detect (bf16 on Ampere/RTX A4000) load_in_4bit = True, text_only = True, # Skip vision tower → ~2 GB VRAM saved ) except TypeError: # Fallback for older Unsloth builds that don't support text_only model, tokenizer = FastLanguageModel.from_pretrained( model_name = MODEL_NAME, max_seq_length = MAX_SEQ_LENGTH, dtype = None, load_in_4bit = True, ) print("VRAM used :", round(torch.cuda.memory_allocated() / 1024**3, 2), "GB") print("Chat template present:", tokenizer.chat_template is not None) gc.collect() torch.cuda.empty_cache() print("Cell 3 complete → Run Cell 4") # ============================================================ # CELL 4 — FORMAT DATASET + TOKEN DIAGNOSTICS + FILTER + SPLIT # ============================================================ def format_chat_template(examples): texts = [] for messages in examples["messages"]: text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=False, # False = training mode ) texts.append(text) return {"text": texts} # Apply chat template dataset = dataset.map(format_chat_template, batched=True, remove_columns=["messages"]) print(f"Formatted {len(dataset):,} examples") # --- Quick token-length peek (first 1 000 records) --- sample_tokens = [ len(tokenizer(ex["text"], add_special_tokens=False)["input_ids"]) for ex in dataset.select(range(min(1000, len(dataset)))) ] print("Max tokens (sample):", max(sample_tokens)) print("Avg tokens (sample):", round(sum(sample_tokens) / len(sample_tokens), 2)) # --- Compute exact token lengths for the full set --- # NOTE: On Windows, num_proc>1 often raises PicklingError because the tokenizer # closure isn't fork-safe. If you want faster processing, move compute_token_length # to a standalone .py file and call load_dataset().map() from there. def compute_token_length(example): return { "token_length": len( tokenizer(example["text"], add_special_tokens=False)["input_ids"] ) } dataset = dataset.map( compute_token_length, num_proc=1, # Windows-safe. On Linux with proper __main__ guard try 4–8. ) # --- Filter overlong records (critical for 16 GB VRAM) --- before = len(dataset) dataset = dataset.filter(lambda x: x["token_length"] <= MAX_SEQ_LENGTH) after = len(dataset) print(f"Removed {before - after:,} records > {MAX_SEQ_LENGTH} tokens") print(f"Remaining {after:,} records") # --- Drop the helper column (training doesn't need it) --- dataset = dataset.remove_columns(["token_length"]) # --- Train / Eval split (0.1 % eval ≈ 970 samples for full set) --- split = dataset.train_test_split(test_size=0.001, seed=3407) train_dataset = split["train"] eval_dataset = split["test"] print("Train :", len(train_dataset)) print("Eval :", len(eval_dataset)) # --- Verify chat template formatting before training --- print("\n--- Template sanity check ---") print(train_dataset[0]["text"][:1000]) gc.collect() print("Cell 4 complete → Run Cell 5") # ============================================================ # CELL 5 — APPLY QLoRA ADAPTER (GUARDED) # ============================================================ from peft import PeftModel if isinstance(model, PeftModel): print("LoRA already attached — skipping.") else: print("Attaching LoRA adapter ...") model = FastLanguageModel.get_peft_model( model, r = LORA_R, lora_alpha = LORA_ALPHA, lora_dropout = 0, bias = "none", use_gradient_checkpointing = "unsloth", random_state = 3407, target_modules = [ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], ) model.config.use_cache = False model.print_trainable_parameters() gc.collect() torch.cuda.empty_cache() print("VRAM after LoRA:", round(torch.cuda.memory_allocated() / 1024**3, 2), "GB") print("Cell 5 complete → Run Cell 6") # ============================================================ # CELL 6 — TRAINING CONFIGURATION # ============================================================ from trl import SFTConfig total_steps = (len(train_dataset) // (BATCH_SIZE * GRAD_ACCUM)) * EPOCHS print("Estimated steps:", total_steps) print("Estimated time : ~", round(total_steps * 2 / 3600, 1), "hours (assuming ~2 s/step on A4000)") # NOTE: Some newer TRL builds renamed evaluation_strategy → eval_strategy. # If you get "unexpected keyword argument", change the line below to: # eval_strategy = "steps", # # NOTE: bf16=True works on RTX A4000 (Ampere). If you get a dtype error, # swap to bf16=False, fp16=True. sft_args = SFTConfig( output_dir = OUTPUT_DIR, overwrite_output_dir = True, num_train_epochs = EPOCHS, per_device_train_batch_size = BATCH_SIZE, gradient_accumulation_steps = GRAD_ACCUM, learning_rate = LR, lr_scheduler_type = "cosine", warmup_steps = 100, optim = "adamw_8bit", weight_decay = 0.01, max_grad_norm = 0.3, bf16 = True, fp16 = False, logging_steps = 10, save_steps = 500, save_total_limit = 2, evaluation_strategy = "steps", eval_steps = 500, report_to = "none", dataset_text_field = "text", max_seq_length = MAX_SEQ_LENGTH, seed = 3407, ) print("Cell 6 complete → Run Cell 7") # ============================================================ # CELL 7 — INITIALIZE TRAINER # ============================================================ from trl import SFTTrainer trainer = SFTTrainer( model = model, train_dataset = train_dataset, eval_dataset = eval_dataset, processing_class = tokenizer, args = sft_args, packing = False, # Keep conversations separate ) print("Trainer ready") print("VRAM before training:", round(torch.cuda.memory_allocated() / 1024**3, 2), "GB") print("Cell 7 complete → Run Cell 8 (sanity) or Cell 9 (full training)") # ============================================================ # CELL 8 — OPTIONAL SANITY CHECK (100 records) # ============================================================ # Uncomment the block below to verify the pipeline on 100 samples # before launching the multi-hour full run. # # sanity_data = train_dataset.select(range(100)) # sanity_trainer = SFTTrainer( # model = model, # train_dataset = sanity_data, # processing_class = tokenizer, # args = sft_args, # packing = False, # ) # sanity_trainer.train() # print("Sanity check passed! Switch back to full dataset for real run.") # ============================================================ # CELL 9 — FULL TRAINING # ============================================================ gc.collect() torch.cuda.empty_cache() print("=" * 60) print("STARTING GEMMA-3 4B QLoRA TRAINING") print("=" * 60) print(f"Train set : {len(train_dataset):,}") print(f"Eval set : {len(eval_dataset):,}") print(f"Epochs : {EPOCHS}") print(f"Batch : {BATCH_SIZE} (accum: {GRAD_ACCUM}, effective: {BATCH_SIZE * GRAD_ACCUM})") print(f"Seq Len : {MAX_SEQ_LENGTH}") print(f"LoRA : r={LORA_R}, alpha={LORA_ALPHA}") print(f"LR : {LR}") print(f"VRAM : {torch.cuda.memory_allocated() / 1024**3:.1f} GB") print("=" * 60) # resume_from_checkpoint=True → starts from latest checkpoint if one exists in OUTPUT_DIR # safe for fresh runs (starts from scratch if none found) trainer.train(resume_from_checkpoint=True) print("=" * 60) print("TRAINING COMPLETE") print("=" * 60) print("Cell 9 complete → Run Cell 10") # ============================================================ # CELL 10 — SAVE ADAPTER + OPTIONAL MERGED 16-BIT MODEL # ============================================================ os.makedirs(OUTPUT_DIR, exist_ok=True) # --- LoRA adapter --- adapter_path = os.path.join(OUTPUT_DIR, "adapter") model.save_pretrained(adapter_path) tokenizer.save_pretrained(adapter_path) print("Adapter saved to:", adapter_path) # --- Optional: merged 16-bit weights (easier for downstream inference) --- try: merged_path = os.path.join(OUTPUT_DIR, "merged_16bit") model.save_pretrained_merged( merged_path, tokenizer, save_method="merged_16bit", ) print("Merged 16-bit model saved to:", merged_path) except Exception as e: print("Merged save skipped:", e) print("Cell 10 complete → Run Cell 11 for inference test") # ============================================================ # CELL 11 — INFERENCE TEST # ============================================================ from unsloth import FastLanguageModel FastLanguageModel.for_inference(model) messages = [ { "role": "user", "content": ( "What does Charaka say about the qualities of a good physician " "in Sutrasthana?" ), } ] inputs = tokenizer.apply_chat_template( messages, tokenize=True, return_tensors="pt", add_generation_prompt=True, # True = inference mode ).to("cuda") print("Generating ...") outputs = model.generate( inputs, max_new_tokens = 512, temperature = 0.7, top_p = 0.9, do_sample = True, ) response = tokenizer.decode(outputs[0], skip_special_tokens=True) print("\n" + "=" * 60) print("RESPONSE") print("=" * 60) print(response) print("=" * 60)