# ============================================================ # AYURVEDA GEMMA-3 4B TRAINING — HF STANDARD QLoRA FINAL # 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 (auto-downloaded from HuggingFace Hub) # Method: Standard HF QLoRA (transformers + bitsandbytes + peft + trl) # NOTE: This notebook uses STANDARD HuggingFace instead of Unsloth. # Slightly slower (~20%) but rock-solid on Windows. # ============================================================ # RESTART KERNEL BEFORE RUNNING CELL 1 # ============================================================ # ============================================================ # CELL 1 — ENVIRONMENT & CONFIG # ============================================================ import os # --- MUST BE SET BEFORE ANY IMPORTS (torch, transformers, datasets) --- # 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, "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) # PyTorch / HuggingFace caches os.environ["TORCHINDUCTOR_CACHE_DIR"] = os.path.join(CACHE, "torch") os.environ["TRITON_CACHE_DIR"] = os.path.join(CACHE, "triton") os.environ["HF_HOME"] = os.path.join(CACHE, "huggingface") os.environ["TRANSFORMERS_CACHE"] = os.path.join(CACHE, "huggingface") os.environ["HUGGINGFACE_HUB_CACHE"] = os.path.join(CACHE, "huggingface") 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 telemetry # WANDB disabled via report_to="none" in SFTConfig (Cell 6) import torch import gc # --- CONFIG --- MODEL_NAME = "unsloth/gemma-3-4b-it" # Put your dataset file inside D:\Sahil\dataset\ 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) --- SUBSET_SIZE = 0 # 50000 for quick test; 0 for full ~971K MAX_SEQ_LENGTH = 512 # Safe for RTX A4000 16 GB BATCH_SIZE = 1 GRAD_ACCUM = 16 # Effective batch = 16 LORA_R = 16 LORA_ALPHA = 16 LR = 2e-4 EPOCHS = 1 # --- GPU CHECK --- if not torch.cuda.is_available(): print("=" * 60) print("ERROR: CUDA GPU not found") print("=" * 60) print("PyTorch version:", torch.__version__) print() print("You need CUDA-enabled PyTorch. Run this in Anaconda Prompt:") print(' pip install torch==2.6.0+cu124 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124') print("=" * 60) raise AssertionError("CUDA GPU not found.") 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 HUGGINGFACE HUB # ============================================================ # IMPORTANT: This cell downloads ~2.5 GB of model weights on FIRST RUN. # Download goes to: D:\Sahil\cache\huggingface\hub\models--unsloth--gemma-3-4b-it\... # On subsequent runs it loads from cache instantly. # ------------------------------------------------------------- from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, ) print(f"Loading {MODEL_NAME} from HuggingFace Hub...") print("First download: ~2.5 GB. Will be cached to D:\\Sahil\\cache\\huggingface\\") print("Please wait...") # 4-bit quantization config (QLoRA) bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16, ) tokenizer = AutoTokenizer.from_pretrained( MODEL_NAME, trust_remote_code=True, cache_dir=os.environ["HF_HOME"], ) # Gemma-3 tokenizer: set pad token if missing if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.pad_token_id = tokenizer.eos_token_id model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, quantization_config=bnb_config, device_map="auto", max_length=MAX_SEQ_LENGTH, trust_remote_code=True, cache_dir=os.environ["HF_HOME"], dtype=torch.float16, ) print("VRAM used :", round(torch.cuda.memory_allocated() / 1024**3, 2), "GB") print("Chat template present:", tokenizer.chat_template is not None) print("Model device map:", model.hf_device_map if hasattr(model, "hf_device_map") else "auto") gc.collect() torch.cuda.empty_cache() print("Cell 3 complete → Run Cell 4") # ============================================================ # CELL 4 — FORMAT DATASET + TOKEN DIAGNOSTICS + FILTER + TRAIN/TEST SPLIT # ============================================================ # This cell does 4 things: # 1. Applies the Gemma chat template to convert [system, user, assistant] → text # 2. Samples token lengths to estimate memory usage # 3. Filters out conversations longer than MAX_SEQ_LENGTH (prevents OOM) # 4. Splits into train (99.9%) and eval (0.1%) sets # ------------------------------------------------------------- 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 filtering --- # NOTE: We pass tokenizer via fn_kwargs because Windows cannot pickle # objects defined in the outer scope when using multiprocessing. def compute_token_length(example, tok): return { "token_length": len( tok(example["text"], add_special_tokens=False)["input_ids"] ) } dataset = dataset.map( compute_token_length, fn_kwargs={"tok": tokenizer}, ) # --- 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 --- # test_size=0.001 means 0.1% of data is held out for evaluation. # For 971K records: ~970 eval samples, ~970K train samples. # seed=3407 ensures reproducibility. 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 (PEFT) # ============================================================ from peft import LoraConfig, get_peft_model, TaskType print("Attaching LoRA adapter ...") peft_config = LoraConfig( task_type=TaskType.CAUSAL_LM, r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=0, bias="none", target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], ) model = get_peft_model(model, peft_config) model.print_trainable_parameters() # Disable cache during training (required for gradient checkpointing) model.config.use_cache = False model.gradient_checkpointing_enable() 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 * 0.9 / 3600, 1), "hours (fp16, expected ~0.9 s/step on A4000)") # NOTE: Some newer TRL builds renamed evaluation_strategy → eval_strategy. # If you get "unexpected keyword argument", change to: # eval_strategy = "steps", 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 = False, fp16 = True, logging_steps = 10, save_steps = 500, save_total_limit = 2, eval_strategy = "steps", eval_steps = 500, report_to = "none", dataset_text_field = "text", 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, ) 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, # eval_dataset = eval_dataset, # processing_class = tokenizer, # args = sft_args, # ) # 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 (Standard HF)") 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 MODEL # ============================================================ import os os.makedirs(OUTPUT_DIR, exist_ok=True) # --- LoRA adapter (small, ~10-50 MB) --- 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: merge LoRA into base model for easier inference --- # This creates a full 16-bit model (~8 GB) that doesn't need PEFT at runtime. try: from peft import AutoPeftModelForCausalLM merged_path = os.path.join(OUTPUT_DIR, "merged_16bit") model.save_pretrained(merged_path) # PEFT merges automatically on save tokenizer.save_pretrained(merged_path) 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 # ============================================================ model.eval() # Set to eval mode 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 ...") with torch.no_grad(): 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)