Tohirju/ameena-9B-lora

Ameena-9B LoRA Adapter

LoRA adapter for Tajik language continual pre-training of Qwen3.5-9B-Base.

Model Details

Parameter	Value
Base model	unsloth/Qwen3.5-9B-Base (9.4B params)
Method	Continual Pre-Training (CPT) with LoRA
LoRA rank	16
LoRA alpha	16
Trainable params	33,136,640 (0.35%)
Target modules	q/k/v/o_proj, gate/up/down_proj, embed_tokens, lm_head
Training data	~370M tokens of Tajik text
Dataset	Tohirju/Tajik_Pretrain_370M_Qwen35_9B
Training steps	25 (of 954)
Effective batch size	192 (48 x 4)
Learning rate	5e-5 (embeddings: 5e-6)
Precision	bf16
GPU	NVIDIA H200 (140GB)
Final loss	1.597

Usage

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained(
    "Tohirju/ameena-9B-lora",
    max_seq_length=4096,
    load_in_4bit=False,
    load_in_16bit=True,
)
FastLanguageModel.for_inference(model)

inputs = tokenizer("\u0422\u043e\u04b7\u0438\u043a\u0438\u0441\u0442\u043e\u043d \u043a\u0438\u0448\u0432\u0430\u0440\u0438 \u0437\u0435\u0431\u043e \u0430\u0441\u0442", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=100, temperature=0.7)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Merging with base model

model.save_pretrained_merged("ameena-9B-merged", tokenizer, save_method="merged_16bit")

Training Script

Click to expand full training script (run_cpt_rev2.py)

#!/usr/bin/env python3
"""Qwen3.5-9B Tajik CPT — REV2: r=16 for maximum speed
Changes from rev0:
  - LoRA r=16 (matching Unsloth official Qwen3.5 docs) — still good for CPT, 2x less compute
  - torch_compile=True — kernel fusion on H100
  - group_by_length=True — reduces padding waste
  - Cleaned up checkpoints from rev0 to start fresh
"""

import os, gc, glob, torch, time
from datetime import datetime
from huggingface_hub import login
from unsloth import FastLanguageModel
from datasets import load_from_disk, load_dataset
from transformers import TrainerCallback, Trainer, TrainingArguments
import bitsandbytes as bnb

# ============ CONFIG ============
HF_TOKEN = os.environ.get("HF_TOKEN", "")
assert HF_TOKEN, "Set HF_TOKEN env var first!"

BASE_MODEL_HF     = "unsloth/Qwen3.5-9B-Base"
MAX_SEQ_LENGTH    = 4096
LOAD_IN_4BIT      = False          # QLoRA NOT safe for Qwen3.5
LOAD_IN_16BIT     = True           # bf16 instead
DATASET_HF_REPO   = "Tohirju/Tajik_Pretrain_370M_Qwen35_9B"
PER_DEVICE_BATCH  = 48
GRAD_ACCUM        = 4              # effective batch = 64
NUM_EPOCHS        = 1
LEARNING_RATE     = 5e-5
EMBED_LR          = 5e-6           # 10x smaller for embeddings
WARMUP_RATIO      = 0.1
LORA_R            = 16             # Reduced from 128 for speed (still good for CPT)
LORA_ALPHA        = 16

WORKSPACE      = "/workspace"
HF_CACHE       = f"{WORKSPACE}/hf_cache"
MODEL_DIR      = f"{WORKSPACE}/models/Qwen3.5-9B-Base"
DATASET_DIR    = f"{WORKSPACE}/datasets/Tajik_Pretrain_Qwen35_9B"
CPT_OUTPUT_DIR = f"{WORKSPACE}/checkpoints/cpt_qwen3.5-9b-rev2"
LOGS_DIR       = f"{WORKSPACE}/training_logs/cpt_qwen3.5-9b-rev2"

for d in [HF_CACHE, MODEL_DIR, DATASET_DIR, CPT_OUTPUT_DIR, LOGS_DIR]:
    os.makedirs(d, exist_ok=True)

os.environ["HF_HOME"]            = HF_CACHE
os.environ["TRANSFORMERS_CACHE"] = HF_CACHE
os.environ["HF_DATASETS_CACHE"]  = f"{HF_CACHE}/datasets"
os.environ["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"

login(token=HF_TOKEN)
print(f"GPU    : {torch.cuda.get_device_name(0)}")
print(f"VRAM   : {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f}GB")
print("=== Cell 2 complete ===\n")

# ============ CELL 3: Load model ============
gc.collect()
torch.cuda.empty_cache()

def find_model_source():
    ckpts = sorted(glob.glob(f"{CPT_OUTPUT_DIR}/checkpoint-*"),
                   key=lambda x: int(x.split("-")[-1]))
    if ckpts:
        print(f"Found CPT checkpoint: {ckpts[-1]}")
        return ckpts[-1], "cpt_checkpoint"
    if os.path.exists(os.path.join(MODEL_DIR, "config.json")):
        print(f"Found local base model: {MODEL_DIR}")
        return MODEL_DIR, "local_base"
    print(f"Downloading {BASE_MODEL_HF} to {MODEL_DIR}...")
    from huggingface_hub import snapshot_download
    snapshot_download(repo_id=BASE_MODEL_HF, local_dir=MODEL_DIR, token=HF_TOKEN,
                      ignore_patterns=["*.pt", "original/*"])
    return MODEL_DIR, "downloaded"

MODEL_SOURCE, SOURCE_TYPE = find_model_source()
print(f"Loading model ({SOURCE_TYPE})...")

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_SOURCE, max_seq_length=MAX_SEQ_LENGTH,
    dtype=None, load_in_4bit=LOAD_IN_4BIT, load_in_16bit=LOAD_IN_16BIT, token=HF_TOKEN,
)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id

total = sum(p.numel() for p in model.parameters())
free_gb = (torch.cuda.get_device_properties(0).total_memory - torch.cuda.memory_allocated()) / 1024**3
print(f"Model loaded! Params: {total:,}, VRAM left: {free_gb:.1f}GB")
print("=== Cell 3 complete ===\n")

# ============ CELL 4: LoRA ============
has_lora = any("lora" in n.lower() for n, _ in model.named_parameters())
if has_lora:
    print("LoRA already applied (checkpoint) — skipping.")
else:
    print(f"Applying LoRA (CPT config, r={LORA_R}, rsLoRA enabled)...")
    model = FastLanguageModel.get_peft_model(
        model, r=LORA_R, lora_alpha=LORA_ALPHA,
        target_modules=["q_proj","k_proj","v_proj","o_proj",
                        "gate_proj","up_proj","down_proj",
                        "embed_tokens","lm_head"],
        lora_dropout=0.0, bias="none",
        use_gradient_checkpointing="unsloth",
        use_rslora=False,
        random_state=3407, max_seq_length=MAX_SEQ_LENGTH,
    )

trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(f"Trainable: {trainable:,} ({100*trainable/total:.2f}%)")
print("=== Cell 4 complete ===\n")

# ============ CELL 5: Dataset ============
PROCESSED_FLAG = os.path.join(DATASET_DIR, "dataset_info.json")
if os.path.exists(PROCESSED_FLAG):
    try:
        print("Loading dataset from local workspace...")
        dataset = load_from_disk(DATASET_DIR)
        print(f"Loaded: {len(dataset):,} examples")
    except Exception as e:
        print(f"Local load failed ({e}), downloading from HF...")
        dataset = load_dataset(DATASET_HF_REPO, split="train", token=HF_TOKEN,
                               cache_dir=f"{HF_CACHE}/datasets")
        dataset.save_to_disk(DATASET_DIR)
        print(f"Downloaded & saved: {len(dataset):,} examples")
else:
    print(f"Downloading dataset from HF: {DATASET_HF_REPO}")
    dataset = load_dataset(DATASET_HF_REPO, split="train", token=HF_TOKEN,
                           cache_dir=f"{HF_CACHE}/datasets")
    dataset.save_to_disk(DATASET_DIR)
    print(f"Downloaded & saved: {len(dataset):,} examples")

tok = tokenizer.tokenizer if hasattr(tokenizer, 'tokenizer') else tokenizer
print(f"Sample: {tok.decode(dataset[0]['input_ids'][:80], skip_special_tokens=True)}")
print(f"=== Cell 5 complete — {len(dataset):,} examples ===\n")

# ============ CELL 6: Verify labels ============
sample = dataset[0]
matches = sum(1 for i,l in zip(sample['input_ids'], sample['labels']) if i==l)
masked = sum(1 for l in sample['labels'] if l==-100)
print(f"labels==input_ids: {100*matches/len(sample['input_ids']):.1f}%, masked: {masked}")
if masked > 0:
    dataset = dataset.map(lambda x: {'labels': x['input_ids'].copy()}, num_proc=1)
    print("Fixed labels.")
print("=== Cell 6 complete ===\n")

# ============ CELL 7: Data collator + length column ============
train_dataset = dataset

# Add length column for group_by_length
def add_length(example):
    return {"length": len(example["input_ids"])}
train_dataset = train_dataset.map(add_length, num_proc=4)

class CPTDataCollator:
    """Pads batch to max length in batch, masks padding in labels."""
    def __init__(self, pad_token_id, max_length):
        self.pad_id = pad_token_id
        self.max_length = max_length

    def __call__(self, features):
        max_len = min(max(len(f['input_ids']) for f in features), self.max_length)
        input_ids, masks, labels = [], [], []
        for f in features:
            ids = list(f['input_ids'][:max_len])
            att = list(f['attention_mask'][:max_len])
            lab = list(f['labels'][:max_len])
            pad = max_len - len(ids)
            input_ids.append(ids + [self.pad_id] * pad)
            masks.append(att + [0] * pad)
            labels.append(lab + [-100] * pad)
        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(masks, dtype=torch.long),
            "labels": torch.tensor(labels, dtype=torch.long),
        }

pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
data_collator = CPTDataCollator(pad_id, MAX_SEQ_LENGTH)

print(f"Train: {len(train_dataset):,} examples")
print("=== Cell 7 complete ===\n")

# ============ CELL 8: Training args ============
training_args = TrainingArguments(
    output_dir=CPT_OUTPUT_DIR,
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=PER_DEVICE_BATCH,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=LEARNING_RATE,
    weight_decay=0.01,
    warmup_ratio=WARMUP_RATIO,
    lr_scheduler_type="cosine",
    max_grad_norm=1.0,
    fp16=False, bf16=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    # Speed optimizations
    torch_compile=False,  # disabled — incompatible with Unsloth patches                # Kernel fusion on H100
    #group_by_length=True,  # not supported in this transformers version
    #length_column_name="length",
    # Eval
    eval_strategy="no",
    # Logging & checkpoints
    logging_strategy="steps", logging_steps=25, logging_first_step=True,
    save_strategy="steps", save_steps=10, save_total_limit=2,
    dataloader_num_workers=0, dataloader_pin_memory=True,
    remove_unused_columns=False,
    report_to="none", seed=3407,
)

eff_batch = PER_DEVICE_BATCH * GRAD_ACCUM
total_steps = len(train_dataset) // eff_batch
print(f"Effective batch: {eff_batch}, Total steps: {total_steps:,}")
print(f"LR: {LEARNING_RATE}, Embed LR: {EMBED_LR}")
print(f"torch_compile: True, group_by_length: True")
print("=== Cell 8 complete ===\n")

# ============ CELL 9: Train ============
class VRAMMonitorCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if torch.cuda.is_available():
            alloc = torch.cuda.memory_allocated() / 1024**3
            resrv = torch.cuda.memory_reserved() / 1024**3
            total_mem = torch.cuda.get_device_properties(0).total_memory / 1024**3
            print(f"  [VRAM] Alloc: {alloc:.1f}GB | Resrv: {resrv:.1f}GB | Free: {total_mem-resrv:.1f}GB")

def find_latest_checkpoint(output_dir):
    ckpts = sorted(glob.glob(os.path.join(output_dir, "checkpoint-*")),
                   key=lambda x: int(x.split("-")[-1]))
    if ckpts:
        print(f"Resuming from: {ckpts[-1]}")
        return ckpts[-1]
    print("No checkpoint — starting fresh.")
    return None

# Build optimizer with separate embedding LR
def create_optimizer(model, lr, embed_lr, weight_decay):
    embed_params = []
    other_params = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        if "embed_tokens" in name or "lm_head" in name:
            embed_params.append(param)
        else:
            other_params.append(param)
    
    print(f"Optimizer: {len(other_params)} params @ LR={lr}, {len(embed_params)} embed params @ LR={embed_lr}")
    
    optimizer = bnb.optim.AdamW8bit([
        {"params": other_params, "lr": lr, "weight_decay": weight_decay},
        {"params": embed_params, "lr": embed_lr, "weight_decay": weight_decay},
    ])
    return optimizer

resume_from = find_latest_checkpoint(CPT_OUTPUT_DIR)
model.train()

optimizer = create_optimizer(model, LEARNING_RATE, EMBED_LR, 0.01)

trainer = Trainer(
    model=model, processing_class=tokenizer, args=training_args,
    train_dataset=train_dataset, data_collator=data_collator,
    optimizers=(optimizer, None),
    callbacks=[VRAMMonitorCallback()],
)

print(f"Training {len(train_dataset):,} examples...")
print(f"NOTE: First few steps will be slow due to torch.compile, then dramatically faster.")
try:
    result = trainer.train(resume_from_checkpoint=resume_from)
    print(f"\nTraining complete! Metrics: {result.metrics}")
    trainer.save_model(CPT_OUTPUT_DIR)
    tokenizer.save_pretrained(CPT_OUTPUT_DIR)
    trainer.save_state()
    trainer.log_metrics("train", result.metrics)
    trainer.save_metrics("train", result.metrics)
    print(f"Model saved to {CPT_OUTPUT_DIR}")
except KeyboardInterrupt:
    print("\nInterrupted — saving...")
    trainer.save_model(CPT_OUTPUT_DIR)
    tokenizer.save_pretrained(CPT_OUTPUT_DIR)
    trainer.save_state()
    print("Saved. Re-run to resume.")
except torch.cuda.OutOfMemoryError:
    print("OOM! Reduce PER_DEVICE_BATCH and restart.")
    raise

print("\n=== TRAINING COMPLETE ===")

Training Logs

Step	Loss	Grad Norm	LR	Speed
1	1.680	0.0074	0 (warmup)	149s/step
25	1.597	0.0083	1.25e-5	127s/step

License

Apache 2.0 (same as base model)