Update latest_finetune.py

latest_finetune.py

@@ -1,19 +1,27 @@
 """
 ================================================================================
-Tiny-Aya v3 -> v4 Sequential Fine-Tuning: NADRA/HBL Structured Transliteration
+Tiny-Aya v2 -> v4 Sequential Fine-Tuning: NADRA/HBL Structured Transliteration
 ================================================================================
 
 Purpose
 -------
-Continue fine-tuning the already-transliteration-tuned v3 checkpoint
-(CohereLabs/tiny-aya-global + merged LoRA, saved as a full HF model) on a
+Continue fine-tuning the v2 checkpoint (CohereLabs/tiny-aya-global + LoRA
+trained on ~50K synthetic Urdu-Roman name/address pairs, merged) on a
 narrower, structured task: NADRA/HBL-style pipe-separated Urdu records
 (name | father/husband name | guardian name | address | address | city,district)
 -> the equivalent Roman Urdu record, preserving field order and separators
 EXACTLY.
 
-This is NOT resumed from the old v3 LoRA adapter (that adapter is already
-merged into the v3 weights). Instead: v3's merged weights are frozen and
+Note: v2 is being used here (instead of the better-performing v3) purely
+for practical reasons -- it's already on the HBL GPU server, whereas v3's
+~6.2GB model.safetensors would take a long time to transfer over a slow
+connection. Because v2's baseline transliteration quality is weaker (only
+50K synthetic examples, output was "not that good"), this run is doing more
+substantial learning than a pure refinement pass would -- see the LORA_R/
+LORA_ALPHA and LEARNING_RATE comments below for how that changes things.
+
+This is NOT resumed from the old v2 LoRA adapter (that adapter is already
+merged into the v2 weights). Instead: v2's merged weights are frozen and
 quantized to 4-bit, and a FRESH LoRA adapter is trained on top of them for
 this specific task. That is the correct pattern for "continue fine-tuning
 a merged checkpoint."
@@ -34,9 +42,11 @@ How to use
    and run once to see real token-length stats for your NADRA records,
    then set MAX_LENGTH accordingly and set this back to False.
 3. Run: python finetune_v4_nadra.py
-4. Adapter (not merged) is saved to OUTPUT_DIR/final_adapter.
-5. Run merge_adapter() (bottom of this file, or call as a separate step)
-   to produce a merged, standalone model for deployment.
+   This trains, saves the LoRA adapter to OUTPUT_DIR/final_adapter, then
+   (by default, AUTO_MERGE_AFTER_TRAINING = True) merges it into v2 and
+   saves a standalone deployable model to OUTPUT_DIR/merged_model, all in
+   this one run. Set AUTO_MERGE_AFTER_TRAINING = False if you'd rather
+   merge later as a separate step instead.
 """
 
 import os
@@ -71,9 +81,12 @@ from peft import (
 # ==============================================================================
 
 # --- Paths -------------------------------------------------------------------
-# Local path to your v3 checkpoint (merged model.safetensors + config.json +
+# Local path to your v2 checkpoint (merged model.safetensors + config.json +
 # tokenizer files + generation_config.json + chat_template.jinja).
-V3_MODEL_PATH = r"./tiny-aya-v3-merged"
+# Local path to your v2 checkpoint on the HBL server (merged
+# model.safetensors + config.json + tokenizer files + generation_config.json
+# + chat_template.jinja). Point this at wherever v2 already lives locally.
+V2_MODEL_PATH = r"./tiny-aya-v2-merged"
 
 # Folder containing splits/train.csv, splits/val.csv, splits/test.csv
 DATA_DIR = r"./splits"
@@ -103,11 +116,15 @@ MAX_LENGTH = 384  # starting default; confirm/adjust using the analysis pass
 SEED = 42
 
 # --- LoRA config ---------------------------------------------------------------
-# Rank/alpha kept from v3 (higher capacity helps structural fidelity).
-# Dropout reduced vs. v3 (0.05 -> 0.02): 2.9M rows makes overfitting on one
-# epoch unlikely, and less dropout helps convergence on a narrower task.
-LORA_R = 32
-LORA_ALPHA = 64
+# Rank/alpha bumped up vs. the original run (32/64 -> 48/96): v2 was only
+# trained on 50K synthetic name/address pairs and its output quality was
+# admittedly weak. This adapter now has to do more of the real learning
+# (both phonetic mapping AND the NADRA structural constraints) rather than
+# just refining an already-solid base, so it gets more capacity.
+# Dropout kept low (0.02): 2.9M rows makes overfitting on one epoch unlikely
+# regardless of starting checkpoint quality.
+LORA_R = 48
+LORA_ALPHA = 96
 LORA_DROPOUT = 0.02
 LORA_TARGET_MODULES = [
     "q_proj", "k_proj", "v_proj", "o_proj",
@@ -115,16 +132,35 @@ LORA_TARGET_MODULES = [
 ]
 
 # --- Training hyperparameters ---------------------------------------------------
-# LR lowered vs. v3 (2e-4 -> 1e-4): this is a refinement pass on an already
-# fine-tuned checkpoint, not learning a domain from scratch. A gentler LR
-# reduces the risk of degrading what v3 already learned.
-LEARNING_RATE = 1e-4
+# LR raised back up vs. a "gentle refinement" rate (1e-4 -> 1.8e-4): v2 only
+# saw 50K synthetic examples and its transliteration quality was weak, so
+# there's less already-good behavior worth protecting with a low LR. This
+# run needs to do real, substantial learning (phonetic mapping AND NADRA
+# structure), not just nudge an already-capable model -- closer to the
+# original from-scratch rate (2e-4) than a fine-refinement rate.
+LEARNING_RATE = 1.8e-4
 NUM_TRAIN_EPOCHS = 1  # with ~2.9M rows, one epoch is already a large amount
-                       # of gradient updates on a fairly repetitive structure.
-                       # Bump only after checking eval loss on a pilot run.
-PER_DEVICE_TRAIN_BATCH_SIZE = 4
-PER_DEVICE_EVAL_BATCH_SIZE = 4
-GRADIENT_ACCUMULATION_STEPS = 4  # effective batch size = 16 (same as v3)
+                       # of gradient updates. Given v2's weaker starting
+                       # point, watch the eval loss curve from a LIMIT_ROWS
+                       # pilot run -- if it's still improving steadily at
+                       # the end of the epoch, a 2nd epoch is worth the
+                       # extra compute time here more than it would have
+                       # been for a v3-based refinement pass.
+PER_DEVICE_TRAIN_BATCH_SIZE = 2  # kept small deliberately -- see note below
+PER_DEVICE_EVAL_BATCH_SIZE = 2
+GRADIENT_ACCUMULATION_STEPS = 8  # effective batch size = 16 (unchanged)
+# Batch size lowered vs. the original (4 -> 2) and accumulation raised
+# (4 -> 8) to keep the SAME effective batch of 16 while cutting per-step
+# activation memory roughly in half. Two things increased memory pressure
+# vs. the original v3-refinement config: MAX_LENGTH went up (224 -> 384,
+# ~1.7x longer sequences) and LORA_R doubled (32 -> 48, more adapter
+# activations at every targeted layer). The A16 vGPU slice here is a fixed
+# 16GB budget with no headroom to discover through trial and error mid-run,
+# so this trades a bit of throughput for a solid safety margin. If a
+# LIMIT_ROWS pilot run shows plenty of free VRAM (check with `nvidia-smi`
+# while it's training), PER_DEVICE_TRAIN_BATCH_SIZE can be raised back
+# toward 4 (halve GRADIENT_ACCUMULATION_STEPS to match, keeping effective
+# batch at 16).
 WARMUP_RATIO = 0.03  # ratio, not fixed steps -- scales with LIMIT_ROWS
 
 LOGGING_STEPS = 20
@@ -140,7 +176,7 @@ USE_BNB_PAGED_OPTIMIZER = False
 
 # --- Merging ---------------------------------------------------------------
 # If True, the script merges the trained LoRA adapter into a fresh copy of
-# the v3 base model and saves a standalone deployable model immediately
+# the v2 base model and saves a standalone deployable model immediately
 # after training finishes, in the same process/run. Set False if you'd
 # rather train now and merge later as a separate step (e.g. to free up the
 # GPU sooner, or to merge from a specific earlier checkpoint by hand).
@@ -255,7 +291,7 @@ def format_example(urdu: str, roman: str, eos_token: str) -> str:
 
 
 # ==============================================================================
-# 4. Completion-only masking (offset-mapping based, same approach as v3)
+# 4. Completion-only masking (offset-mapping based, same approach as v2)
 # ==============================================================================
 def build_tokenize_fn(tokenizer, max_length):
     def tokenize_function(examples):
@@ -324,13 +360,13 @@ def analyze_lengths(tokenizer, texts, sample_size):
 
 
 # ==============================================================================
-# 6. Model loading: v3 (local, merged) -> 4-bit quantized + fresh LoRA adapter
+# 6. Model loading: v2 (local, merged) -> 4-bit quantized + fresh LoRA adapter
 # ==============================================================================
 def load_model_and_tokenizer():
-    if not os.path.isdir(V3_MODEL_PATH):
+    if not os.path.isdir(V2_MODEL_PATH):
         raise FileNotFoundError(
-            f"V3_MODEL_PATH does not exist: {V3_MODEL_PATH}. "
-            f"This must be a local directory containing the v3 merged "
+            f"V2_MODEL_PATH does not exist: {V2_MODEL_PATH}. "
+            f"This must be a local directory containing the v2 merged "
             f"checkpoint (model.safetensors, config.json, tokenizer files)."
         )
 
@@ -341,9 +377,9 @@ def load_model_and_tokenizer():
         bnb_4bit_use_double_quant=True,
     )
 
-    print(f"Loading v3 checkpoint from local path: {V3_MODEL_PATH}")
+    print(f"Loading v2 checkpoint from local path: {V2_MODEL_PATH}")
     model = AutoModelForCausalLM.from_pretrained(
-        V3_MODEL_PATH,
+        V2_MODEL_PATH,
         quantization_config=bnb_config,
         torch_dtype=torch.float16,
         device_map={"": 0},
@@ -352,10 +388,10 @@ def load_model_and_tokenizer():
     )
     model.config.use_cache = False
     model = prepare_model_for_kbit_training(model)
-    model.gradient_checkpointing_enable()
+    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
 
     tokenizer = AutoTokenizer.from_pretrained(
-        V3_MODEL_PATH, local_files_only=True
+        V2_MODEL_PATH, local_files_only=True
     )
     tokenizer.pad_token = tokenizer.eos_token
     tokenizer.padding_side = "right"
@@ -396,7 +432,7 @@ def main():
     # --- Tokenizer needed for both the length-analysis path and formatting ---
     # Loaded here (unquantized/tokenizer-only) so ANALYZE_LENGTHS_ONLY can run
     # without paying the cost of loading the full 4-bit model.
-    tokenizer = AutoTokenizer.from_pretrained(V3_MODEL_PATH, local_files_only=True)
+    tokenizer = AutoTokenizer.from_pretrained(V2_MODEL_PATH, local_files_only=True)
     tokenizer.pad_token = tokenizer.eos_token
 
     train_texts = [
@@ -440,7 +476,7 @@ def main():
             "before starting an unattended run."
         )
 
-    # --- Load model (v3 -> 4-bit -> fresh LoRA) ---
+    # --- Load model (v2 -> 4-bit -> fresh LoRA) ---
     model, model_tokenizer = load_model_and_tokenizer()
     # tokenizer used for data collation must match the one used for tokenizing
     assert model_tokenizer.get_vocab() == tokenizer.get_vocab(), (
@@ -472,7 +508,10 @@ def main():
         remove_unused_columns=False,
         optim=optim_name,
         gradient_checkpointing=True,
-        gradient_checkpointing_kwargs={"use_reentrant": True},
+        gradient_checkpointing_kwargs={"use_reentrant": False},  # more memory-
+        # efficient than reentrant checkpointing, and the current recommended
+        # setting for PEFT + quantized models -- matters more here given the
+        # fixed 16GB A16 budget.
         # No push_to_hub / hub_model_id / hub_token: server has no internet.
     )
 
@@ -518,7 +557,7 @@ def main():
     print(f"Training complete. LoRA adapter saved to: {ADAPTER_DIR}")
 
     # --- Free the 4-bit training model's GPU memory before merging ---
-    # Merging loads a fresh, unquantized copy of v3 on CPU, but the trainer
+    # Merging loads a fresh, unquantized copy of v2 on CPU, but the trainer
     # and quantized model from this run are still holding CUDA memory/cache
     # at this point. Explicitly drop references and clear the allocator
     # cache so the merge step (or anything run after it in this process)
@@ -529,7 +568,7 @@ def main():
     torch.cuda.empty_cache()
 
     if AUTO_MERGE_AFTER_TRAINING:
-        print("Proceeding to merge the trained adapter into v3 (same run)...")
+        print("Proceeding to merge the trained adapter into v2 (same run)...")
         merge_adapter()
     else:
         print(
@@ -539,7 +578,7 @@ def main():
 
 
 # ==============================================================================
-# 8. Merge LoRA adapter into v3 base -> standalone v4 model for deployment
+# 8. Merge LoRA adapter into v2 base -> standalone v4 model for deployment
 # ==============================================================================
 def merge_adapter():
     """
@@ -548,12 +587,12 @@ def merge_adapter():
     trained with AUTO_MERGE_AFTER_TRAINING = False, or want to re-merge from
     a specific checkpoint.
 
-    Loads v3 in fp16 (not quantized) on CPU, merges the trained LoRA adapter
+    Loads v2 in fp16 (not quantized) on CPU, merges the trained LoRA adapter
     into it, and saves a standalone model that can be loaded directly with
     AutoModelForCausalLM.from_pretrained() at deployment time -- no PEFT or
     bitsandbytes required for inference.
     """
-    # Priority order for adapter source, mirroring v3's approach:
+    # Priority order for adapter source, mirroring v2's approach:
     adapter_path = None
     if os.path.exists(os.path.join(ADAPTER_DIR, "adapter_config.json")):
         adapter_path = ADAPTER_DIR
@@ -585,15 +624,15 @@ def merge_adapter():
     gc.collect()
     torch.cuda.empty_cache()
 
-    print(f"Loading v3 base model (unquantized, CPU) from: {V3_MODEL_PATH}")
+    print(f"Loading v2 base model (unquantized, CPU) from: {V2_MODEL_PATH}")
     base_model = AutoModelForCausalLM.from_pretrained(
-        V3_MODEL_PATH,
+        V2_MODEL_PATH,
         torch_dtype=torch.float16,
         device_map="cpu",
         trust_remote_code=True,
         local_files_only=True,
     )
-    tokenizer = AutoTokenizer.from_pretrained(V3_MODEL_PATH, local_files_only=True)
+    tokenizer = AutoTokenizer.from_pretrained(V2_MODEL_PATH, local_files_only=True)
 
     print("Merging LoRA weights into base...")
     merged_model = PeftModel.from_pretrained(base_model, adapter_path)