104 lines
3.9 KiB
Python
104 lines
3.9 KiB
Python
"""Phase 3 — continued pretraining of Qwen3-1.7B-Base on the packed RU corpus.
|
|
|
|
Fits a 16GB RX 7900 GRE (gfx1100). Full-weight CPT of 1.7B with plain Adam does
|
|
NOT fit; we use adafactor (no momentum states) + gradient checkpointing (path A,
|
|
plan §3.1). If it still OOMs, set PATH="dora" for high-rank DoRA on all linears.
|
|
|
|
NO bitsandbytes (ROCm flaky) → no 8-bit Adam, hence adafactor.
|
|
|
|
Critical hyperparams (plan §3.2): lr=1e-5 (10x lower than LoRA — nudge, don't
|
|
forget), 1 epoch (more = catastrophic forgetting). Checkpoints every 500 steps;
|
|
resume with --resume.
|
|
|
|
Usage:
|
|
HSA_OVERRIDE_GFX_VERSION=11.0.0 python train_cpt.py 2>&1 | tee cpt_run.log
|
|
HSA_OVERRIDE_GFX_VERSION=11.0.0 python train_cpt.py --resume
|
|
"""
|
|
import os
|
|
import sys
|
|
|
|
os.environ.setdefault("HF_HOME", "/mnt/D/.cache/huggingface")
|
|
os.environ.setdefault("HF_DATASETS_CACHE", "/mnt/D/.cache/huggingface/datasets")
|
|
os.environ.setdefault("TMPDIR", "/mnt/D/tmp")
|
|
os.environ.setdefault("HSA_OVERRIDE_GFX_VERSION", "11.0.0")
|
|
# lever 2: kill allocator fragmentation so the caching allocator doesn't hoard the card
|
|
os.environ.setdefault("PYTORCH_HIP_ALLOC_CONF", "expandable_segments:True")
|
|
|
|
import torch
|
|
from datasets import load_from_disk
|
|
from transformers import (
|
|
AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments,
|
|
DataCollatorForLanguageModeling,
|
|
)
|
|
|
|
from corpus_common import PACKED
|
|
|
|
MODEL = "Qwen/Qwen3-1.7B-Base" # BASE, not Instruct (plan hard rule 4)
|
|
OUTPUT = "./Qwen3-1.7B-ru-cpt"
|
|
PATH = os.environ.get("CPT_PATH", "full") # "full" (adafactor) | "dora"
|
|
|
|
|
|
def main():
|
|
resume = "--resume" in sys.argv
|
|
|
|
tok = AutoTokenizer.from_pretrained(MODEL)
|
|
if tok.pad_token is None:
|
|
tok.pad_token = tok.eos_token
|
|
|
|
ds = load_from_disk(str(PACKED))
|
|
print(f"[cpt] {len(ds)} blocks x 2048 = {len(ds)*2048/1e6:.0f}M tokens")
|
|
|
|
# lever 3: hard cap VRAM so an OOM raises (recoverable) instead of starving the compositor
|
|
if torch.cuda.is_available():
|
|
torch.cuda.set_per_process_memory_fraction(0.85)
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL, torch_dtype=torch.bfloat16, attn_implementation="sdpa") # lever 1: flash/mem-efficient
|
|
model.config.use_cache = False
|
|
model.gradient_checkpointing_enable()
|
|
|
|
optim = "adafactor"
|
|
if PATH == "dora": # fallback if full-weight OOMs (plan §3.1B)
|
|
from peft import LoraConfig, get_peft_model
|
|
cfg = LoraConfig(
|
|
r=64, lora_alpha=128, lora_dropout=0.05, use_dora=True,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
|
|
"gate_proj", "up_proj", "down_proj"],
|
|
task_type="CAUSAL_LM",
|
|
)
|
|
model = get_peft_model(model, cfg)
|
|
model.print_trainable_parameters()
|
|
optim = "adamw_torch_fused" # LoRA params are tiny, Adam fits
|
|
|
|
args = TrainingArguments(
|
|
output_dir=OUTPUT,
|
|
per_device_train_batch_size=1,
|
|
gradient_accumulation_steps=16, # eff batch 16 x 2048 tok
|
|
num_train_epochs=1, # ONE pass (hard rule 3)
|
|
learning_rate=1e-5, # LOW (hard rule 3)
|
|
lr_scheduler_type="cosine",
|
|
warmup_ratio=0.03,
|
|
optim=optim,
|
|
bf16=True,
|
|
gradient_checkpointing=True,
|
|
logging_steps=20,
|
|
save_steps=200, # ~1.1h between ckpts → clean daily stop/resume, low loss on interrupt
|
|
save_total_limit=3,
|
|
report_to="none",
|
|
dataloader_num_workers=0, # py3.14 mp leaks semaphores + crashes; data is pre-tokenized so workers buy nothing
|
|
)
|
|
|
|
collator = DataCollatorForLanguageModeling(tok, mlm=False)
|
|
trainer = Trainer(model=model, args=args, train_dataset=ds, data_collator=collator)
|
|
trainer.train(resume_from_checkpoint=resume)
|
|
|
|
if PATH == "dora":
|
|
model = model.merge_and_unload()
|
|
model.save_pretrained(OUTPUT)
|
|
tok.save_pretrained(OUTPUT)
|
|
print(f"[cpt] DONE -> {OUTPUT}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|