#!/usr/bin/env python # -*- coding: utf-8 -*- """ Reproduction script for NAMAA-Space/alexandriax-nllb-1.3b-lora =================================================================== AlexandriaX-2026 Subtask 1 - QLoRA fine-tune of facebook/nllb-200-1.3B for English -> dialectal Arabic, with the target variety selected by NLLB's own language codes. This is the exact recipe of the original Colab notebook (`AlexandriaX_NB7_NLLB_1.3B_Full.ipynb`), extracted into a single runnable file. Every hyperparameter below is the value that produced the released adapter: 21.83 spBLEU / 38.13 chrF++ on the official 12,250-turn dev set. Two implementation details cost real debugging time and are preserved verbatim, with comments: * gradient checkpointing must be enabled EXACTLY ONCE, with use_reentrant=False, or M2M100/NLLB raises "cannot specify both decoder_input_ids and decoder_inputs_embeds"; * NLLB's native decoder_start_token_id must NOT be overridden. pip install transformers datasets peft bitsandbytes sacrebleu accelerate python train_nllb13b_qlora.py --out ./runs/nllb13b_qlora python train_nllb13b_qlora.py --out ./runs/nllb13b_qlora --epochs 1 # often as good, ~3x faster Hardware used originally: one A100-80GB. Author: NAMAA Community (Fatimah Emad Eldin, Omer Nacar, Khloud Al Jallad, Mona Abdelazim) License: CC-BY-NC-4.0 (matches the base model - non-commercial) """ import argparse import glob import json import os import re from pathlib import Path import torch from datasets import Dataset, concatenate_datasets, get_dataset_config_names, \ get_dataset_split_names, load_dataset from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig, DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments, TrainerCallback) # -------------------------------------------------------------------------------------- # Configuration - the released adapter's exact values # -------------------------------------------------------------------------------------- MODEL_NAME = "facebook/nllb-200-1.3B" DATASET = "UBC-NLP/alexandria" MAX_SRC, MAX_TGT = 200, 200 EPOCHS = 3 BATCH = 16 # per device GRAD_ACCUM = 2 # -> effective batch 32 LR = 2e-4 WARMUP_STEPS = 100 LR_SCHEDULER = "cosine" # NOTE: the repo's first auto-generated card claimed label_smoothing=0.1. It was never # applied - the string appeared only in that card template, not in TrainingArguments. OPTIM = "paged_adamw_8bit" LABEL_PAD = -100 LOGGING_STEPS, SAVE_STEPS, SAVE_LIMIT = 25, 200, 2 SEED = 42 LORA_R, LORA_ALPHA, LORA_DROPOUT = 16, 32, 0.05 LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"] # attention AND FFN GRADIENT_CKPT = True NUM_BEAMS, GEN_MAX_NEW, GEN_BATCH = 5, 128, 48 # Per-dialect NLLB target codes. Seven varieties have a native code; LY, SD, OM and MR have # none and are mapped to their nearest neighbour, which biases decoding and is then adapted # by fine-tuning. NLLB's SentencePiece model is also the one spBLEU tokenises with. SRC_LANG = "eng_Latn" FALLBACK = "arb_Arab" NLLB_CODE = {"EG": "arz_Arab", # native, Egyptian "MA": "ary_Arab", # native, Moroccan "TN": "aeb_Arab", # native, Tunisian "SA": "ars_Arab", # native, Najdi "SY": "apc_Arab", # native, North Levantine "LB": "apc_Arab", # native, North Levantine "JO": "ajp_Arab", # native, South Levantine "PS": "ajp_Arab", # native, South Levantine "YE": "acq_Arab", # native, Ta'izzi-Adeni "OM": "ars_Arab", # PROXY - no native code "MR": "ary_Arab", # PROXY - no native code "LY": "arz_Arab", # PROXY - no native code, test-only variety "SD": "arz_Arab"} # PROXY - no native code, test-only variety PROXY = {"LY", "SD", "OM", "MR"} # -------------------------------------------------------------------------------------- # Data # -------------------------------------------------------------------------------------- def _sorted_turns(turns): return sorted(turns or [], key=lambda t: int(t.get("turn_order", 0) or 0)) def _target_of(turn): for k in ("reference", "dialectal_translation", "translation", "target", "gold"): if turn.get(k): return turn[k] return None def load_pairs_by_country(split="train"): configs = set(get_dataset_config_names(DATASET)) out = {} for country in NLLB_CODE: if country not in configs or split not in set(get_dataset_split_names(DATASET, country)): continue srcs, tgts = [], [] for row in load_dataset(DATASET, country, split=split): english = {int(t.get("turn_order", i + 1)): t.get("sentence") for i, t in enumerate(_sorted_turns(row.get("turns", [])))} arabic = {int(t.get("turn_order", i + 1)): _target_of(t) for i, t in enumerate(_sorted_turns(row.get("dialectal_conversation", row.get("turns", []))))} for order, en in english.items(): if en and arabic.get(order): srcs.append(en) tgts.append(arabic[order]) out[country] = (srcs, tgts) print(f"[data] {country} ({NLLB_CODE[country]}" f"{', proxy' if country in PROXY else ''}): {len(srcs)} pairs") return out def build_dataset(tokenizer, per_country): """Tokenise each country under its own target code, then concatenate.""" parts = [] for country, (srcs, tgts) in per_country.items(): if not srcs: continue code = NLLB_CODE.get(country, FALLBACK) tokenizer.src_lang = SRC_LANG tokenizer.tgt_lang = code ds = Dataset.from_dict({"src": srcs, "tgt": tgts}) def fn(b): mi = tokenizer(b["src"], max_length=MAX_SRC, truncation=True) mi["labels"] = tokenizer(text_target=b["tgt"], max_length=MAX_TGT, truncation=True)["input_ids"] return mi parts.append(ds.map(fn, batched=True, remove_columns=ds.column_names, desc=f"tok {country}->{code}")) return concatenate_datasets(parts).shuffle(seed=SEED) # -------------------------------------------------------------------------------------- # Training # -------------------------------------------------------------------------------------- class ProgressCallback(TrainerCallback): def __init__(self, path): super().__init__() self.path = path def on_save(self, args, state, control, **kw): try: Path(self.path).write_text(json.dumps({"step": state.global_step, "epoch": state.epoch})) except Exception: pass def latest_ckpt(d): cks = [p for p in glob.glob(os.path.join(d, "checkpoint-*")) if os.path.isdir(p)] return max(cks, key=lambda p: int(re.findall(r"checkpoint-(\d+)", p)[-1])) if cks else None def main(): ap = argparse.ArgumentParser() ap.add_argument("--out", default="./runs/nllb13b_qlora") ap.add_argument("--epochs", type=int, default=EPOCHS) ap.add_argument("--batch", type=int, default=BATCH) ap.add_argument("--grad-accum", type=int, default=GRAD_ACCUM) ap.add_argument("--lr", type=float, default=LR) args = ap.parse_args() os.makedirs(args.out, exist_ok=True) bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported() cdtype = torch.bfloat16 if bf16 else torch.float16 torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True print(f"[precision] bf16={bf16} compute_dtype={cdtype}") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) for code in sorted(set(NLLB_CODE.values()) | {FALLBACK}): assert tokenizer.convert_tokens_to_ids(code) != tokenizer.unk_token_id, \ f"target code {code} is not in this tokenizer" train_ds = build_dataset(tokenizer, load_pairs_by_country("train")) print(train_ds) bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=cdtype, bnb_4bit_use_double_quant=True) model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, quantization_config=bnb, device_map="auto") # ---- the M2M100/NLLB gradient-checkpointing fix ---- # 1) prepare WITHOUT enabling GC (so it is never double-enabled), # 2) enable GC exactly once with use_reentrant=False, # 3) keep NLLB's native decoder_start_token_id, # 4) mask label padding with -100. model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=False) model.config.pad_token_id = tokenizer.pad_token_id model.config.use_cache = False lora = LoraConfig(r=LORA_R, lora_alpha=LORA_ALPHA, lora_dropout=LORA_DROPOUT, bias="none", task_type="SEQ_2_SEQ_LM", target_modules=LORA_TARGETS) model = get_peft_model(model, lora) model.print_trainable_parameters() # ~23.6M of 1.37B, about 1.7% if GRADIENT_CKPT: model.gradient_checkpointing_enable( gradient_checkpointing_kwargs={"use_reentrant": False}) model.enable_input_require_grads() # grads must reach the frozen 4-bit base under GC collator = DataCollatorForSeq2Seq(tokenizer, model=model, label_pad_token_id=LABEL_PAD) targs = Seq2SeqTrainingArguments( output_dir=args.out, num_train_epochs=args.epochs, per_device_train_batch_size=args.batch, gradient_accumulation_steps=args.grad_accum, learning_rate=args.lr, warmup_steps=WARMUP_STEPS, lr_scheduler_type=LR_SCHEDULER, logging_steps=LOGGING_STEPS, save_steps=SAVE_STEPS, save_total_limit=SAVE_LIMIT, eval_strategy="no", bf16=bf16, fp16=not bf16, gradient_checkpointing=False, # already enabled above; do NOT double-enable optim=OPTIM, report_to="none", predict_with_generate=False, remove_unused_columns=False, seed=SEED, ) trainer = Seq2SeqTrainer(model=model, args=targs, train_dataset=train_ds, data_collator=collator, processing_class=tokenizer, callbacks=[ProgressCallback(os.path.join(args.out, "progress.json"))]) ck = latest_ckpt(args.out) print(("resume " + ck) if ck else "fresh start") trainer.train(resume_from_checkpoint=ck) trainer.model.save_pretrained(args.out) tokenizer.save_pretrained(args.out) print("[done] saved LoRA adapter ->", args.out) # smoke test - beam search with the per-dialect forced BOS token model.eval() model.config.use_cache = True tokenizer.src_lang = SRC_LANG enc = tokenizer(["Good morning. How much for the whole quantity?"], return_tensors="pt").to(model.device) with torch.no_grad(): out = model.generate(**enc, forced_bos_token_id=tokenizer.convert_tokens_to_ids( NLLB_CODE["EG"]), num_beams=NUM_BEAMS, max_new_tokens=GEN_MAX_NEW) print("[smoke]", tokenizer.batch_decode(out, skip_special_tokens=True)) if __name__ == "__main__": main()