{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3",
      "language": "python"
    },
    "language_info": {
      "name": "python",
      "version": "3.12"
    },
    "colab": {
      "name": "tbr8-colab-lora-resume.ipynb",
      "provenance": []
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "id": "cell-00",
      "cell_type": "markdown",
      "metadata": {},
      "source": "# Colab LoRA 断线续训\n\n下载后通过 Colab 的「文件 → 上传笔记本」打开，选择 GPU 运行时，从上到下运行。\n本示例使用公开 Qwen2.5-0.5B-Instruct 和 8 条自编训练样本，演示 checkpoint 持久化与恢复。\n首次挂载 Google Drive 需要你授权。请使用自己的 checkpoint。\n第二次连接新运行时后，保持 RUN_NAME、配置和数据相同，再从头运行。\n数据与模型效果只用于演示。验证记录和使用说明见 tbr8.org 配套文章。\n"
    },
    {
      "id": "cell-01",
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": "import os\nimport subprocess\nimport sys\n\nSMOKE = os.environ.get(\"TBR8_SMOKE_TEST\") == \"1\"\nPINS = [\"transformers==4.57.6\", \"peft==0.18.1\", \"accelerate==1.12.0\"]\nif not SMOKE:\n    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", *PINS])\nprint(\"Dependencies ready. Use a fresh runtime when starting or resuming.\")\n"
    },
    {
      "id": "cell-02",
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 1. 挂载存储并固定训练配置\n\nRUN_NAME 是同一次训练的目录。继续训练时保持不变，开始另一组实验时换一个新名称。\n默认演示训练 16 个 optimizer steps，每 4 步保存。实际任务自行扩大数据量与总步数。\n"
    },
    {
      "id": "cell-03",
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": "import hashlib\nimport importlib.metadata\nimport json\nfrom pathlib import Path\n\nimport torch\nfrom packaging.version import Version\nfrom huggingface_hub import HfApi\nfrom transformers import set_seed\n\nif Version(torch.__version__.split(\"+\")[0]) < Version(\"2.6\"):\n    raise RuntimeError(\"Use a runtime with PyTorch >= 2.6 for checkpoint state loading.\")\nif not SMOKE and not torch.cuda.is_available():\n    raise RuntimeError(\"Select a GPU runtime in Colab, then rerun.\")\n\nRUN_NAME = \"qwen25-05b-lora-resume-v1\"\nif SMOKE:\n    RUN_DIR = Path(os.environ[\"TBR8_SMOKE_DIR\"])\nelse:\n    from google.colab import drive\n    drive.mount(\"/content/drive\")\n    RUN_DIR = Path(\"/content/drive/MyDrive/tbr8-lora\") / RUN_NAME\nRUN_DIR.mkdir(parents=True, exist_ok=True)\n\nPAIRS = [\n    (\"训练断线后从哪里恢复？\", \"从持久化存储中最后一个完整的训练检查点恢复。\"),\n    (\"只保存 LoRA 权重够不够续训？\", \"还应保留优化器、调度器、随机状态和训练步数。\"),\n    (\"如何确认续训成功？\", \"核对恢复起始步数和训练后步数，并检查状态文件。\"),\n    (\"为什么要保存基础模型版本？\", \"同一训练需要使用相同的基础模型快照。\"),\n    (\"训练数据可以随便改吗？\", \"恢复同一训练时应保持训练数据和顺序规则一致。\"),\n    (\"检查点写到一半怎么办？\", \"跳过未完成的检查点，选择上一个通过校验的版本。\"),\n    (\"Colab 临时磁盘能长期保存吗？\", \"运行时可能回收，应将检查点放入持久化存储。\"),\n    (\"怎样开始新的训练实验？\", \"使用新的运行目录，并保存新实验的配置。\"),\n]\nconfig_path = RUN_DIR / \"run_config.json\"\nexisting_config = json.loads(config_path.read_text(\"utf-8\")) if config_path.exists() else None\nmodel_id = \"Qwen/Qwen2.5-0.5B-Instruct\"\nrevision = (\"seed-42-tiny-model\" if SMOKE else\n            existing_config[\"revision\"] if existing_config else HfApi().model_info(model_id).sha)\nMAX_STEPS = 4 if SMOKE else 16\nSAVE_STEPS = 2 if SMOKE else 4\nMAX_LENGTH = 96 if SMOKE else 256\nconfig = {\n    \"model\": \"tiny-qwen2\" if SMOKE else model_id,\n    \"revision\": revision,\n    \"seed\": 42,\n    \"max_steps\": MAX_STEPS,\n    \"save_steps\": SAVE_STEPS,\n    \"max_length\": MAX_LENGTH,\n    \"batch_size\": 1,\n    \"gradient_accumulation\": 2 if SMOKE else 4,\n    \"learning_rate\": 0.0002,\n    \"lora_r\": 4 if SMOKE else 8,\n    \"lora_alpha\": 8 if SMOKE else 16,\n    \"lora_dropout\": 0.0,\n    \"target_modules\": [\"q_proj\", \"v_proj\"],\n    \"dtype\": \"float32\" if SMOKE else \"float16\",\n    \"data_sha256\": hashlib.sha256(json.dumps(PAIRS, ensure_ascii=False).encode()).hexdigest(),\n    \"versions\": {name: importlib.metadata.version(name) for name in [\"torch\", \"transformers\", \"peft\", \"accelerate\"]},\n}\nif existing_config and existing_config != config:\n    raise RuntimeError(\"Saved run_config differs. Restore the original environment/config or use a new RUN_NAME.\")\nif not existing_config:\n    if list(RUN_DIR.glob(\"checkpoint-*\")):\n        raise RuntimeError(\"Existing checkpoints have no run_config; inspect them before continuing.\")\n    temporary = config_path.with_suffix(\".tmp\")\n    temporary.write_text(json.dumps(config, ensure_ascii=False, indent=2), encoding=\"utf-8\")\n    temporary.replace(config_path)\nRUN_SIGNATURE = hashlib.sha256(json.dumps(config, sort_keys=True).encode()).hexdigest()\nset_seed(config[\"seed\"])\nprint(json.dumps(config, ensure_ascii=False, indent=2))\n"
    },
    {
      "id": "cell-04",
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 2. 载入模型和 LoRA，构建训练样本\n\n只有 assistant 回答部分参与 loss，padding 与提示词位置使用 -100 屏蔽。\n为便于低资源试跑，使用 0.5B 模型和短序列。此处没有启用量化。\n"
    },
    {
      "id": "cell-05",
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": "from torch.utils.data import Dataset\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config, Qwen2ForCausalLM, PreTrainedTokenizerFast\nfrom peft import LoraConfig, get_peft_model\n\nif SMOKE:\n    from tokenizers import Tokenizer\n    from tokenizers.models import WordLevel\n    from tokenizers.pre_tokenizers import WhitespaceSplit\n    words = [\"[PAD]\", \"[UNK]\", \"[EOS]\", \"user\", \"assistant\", \"checkpoint\", \"resume\", \"save\", \"state\", \"step\"]\n    backend = Tokenizer(WordLevel({word: index for index, word in enumerate(words)}, unk_token=\"[UNK]\"))\n    backend.pre_tokenizer = WhitespaceSplit()\n    tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, pad_token=\"[PAD]\", unk_token=\"[UNK]\", eos_token=\"[EOS]\")\n    model = Qwen2ForCausalLM(Qwen2Config(\n        vocab_size=len(tokenizer), hidden_size=32, intermediate_size=64,\n        num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2,\n        max_position_embeddings=256, pad_token_id=tokenizer.pad_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n    ))\nelse:\n    tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)\n    if tokenizer.pad_token_id is None:\n        tokenizer.pad_token = tokenizer.eos_token\n    model = AutoModelForCausalLM.from_pretrained(\n        model_id, revision=revision, torch_dtype=torch.float16\n    )\nmodel.config.use_cache = False\nmodel = get_peft_model(model, LoraConfig(\n    task_type=\"CAUSAL_LM\", r=config[\"lora_r\"], lora_alpha=config[\"lora_alpha\"],\n    lora_dropout=config[\"lora_dropout\"], target_modules=config[\"target_modules\"],\n))\nmodel.print_trainable_parameters()\n\nrecords = []\nfor question, answer in PAIRS:\n    if SMOKE:\n        prompt = \"user checkpoint assistant \"\n        full = prompt + \"resume save state step [EOS]\"\n    else:\n        user = [{\"role\": \"user\", \"content\": question}]\n        prompt = tokenizer.apply_chat_template(user, tokenize=False, add_generation_prompt=True)\n        full = tokenizer.apply_chat_template(user + [{\"role\": \"assistant\", \"content\": answer}], tokenize=False)\n        if not full.startswith(prompt):\n            raise RuntimeError(\"Chat template changed; inspect assistant prefix before masking labels.\")\n    prompt_ids = tokenizer(prompt, add_special_tokens=False)[\"input_ids\"]\n    full_ids = tokenizer(full, add_special_tokens=False)[\"input_ids\"][:MAX_LENGTH]\n    if full_ids[:len(prompt_ids)] != prompt_ids or len(full_ids) <= len(prompt_ids):\n        raise ValueError(\"Prompt boundary mismatch or answer truncated; adjust tokenization/length.\")\n    padding = MAX_LENGTH - len(full_ids)\n    labels = [-100] * len(prompt_ids) + full_ids[len(prompt_ids):] + [-100] * padding\n    records.append({\n        \"input_ids\": full_ids + [tokenizer.pad_token_id] * padding,\n        \"attention_mask\": [1] * len(full_ids) + [0] * padding,\n        \"labels\": labels,\n    })\n\nclass ToyDataset(Dataset):\n    def __len__(self):\n        return len(records)\n    def __getitem__(self, index):\n        return records[index]\n\ndataset = ToyDataset()\nprint(\"Training examples:\", len(dataset))\n"
    },
    {
      "id": "cell-06",
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 3. 只恢复经过完整性校验的 checkpoint\n\n每次 Trainer 保存后，回调计算文件校验值，最后写入 COMPLETE.json。\n恢复时逐一验证文件并按步数选择最新的完整检查点。写入中断的目录会被跳过。\n校验值用于识别未完成或损坏的文件，不代表可以安全载入陌生来源的检查点。\n"
    },
    {
      "id": "cell-07",
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": "from transformers import TrainerCallback\n\ndef file_hash(path):\n    digest = hashlib.sha256()\n    with path.open(\"rb\") as source:\n        for block in iter(lambda: source.read(1024 * 1024), b\"\"):\n            digest.update(block)\n    return digest.hexdigest()\n\nREQUIRED = [\"adapter_model.safetensors\", \"adapter_config.json\", \"optimizer.pt\",\n            \"scheduler.pt\", \"trainer_state.json\", \"rng_state.pth\", \"training_args.bin\"]\n\nclass CompleteCheckpoint(TrainerCallback):\n    def on_save(self, args, state, control, **kwargs):\n        directory = Path(args.output_dir) / f\"checkpoint-{state.global_step}\"\n        if not all((directory / name).is_file() for name in REQUIRED):\n            raise RuntimeError(\"Checkpoint missing required state files.\")\n        files = {\n            str(path.relative_to(directory)): file_hash(path)\n            for path in directory.rglob(\"*\")\n            if path.is_file() and path.name not in [\"COMPLETE.json\", \"COMPLETE.tmp\"]\n        }\n        manifest = {\"signature\": RUN_SIGNATURE, \"step\": state.global_step, \"files\": files}\n        temporary = directory / \"COMPLETE.tmp\"\n        temporary.write_text(json.dumps(manifest, indent=2), encoding=\"utf-8\")\n        temporary.replace(directory / \"COMPLETE.json\")\n        print(\"Complete checkpoint:\", directory.name)\n\n    def on_step_end(self, args, state, control, **kwargs):\n        stop_at = int(os.environ.get(\"TBR8_INTERRUPT_AT\", \"0\")) if SMOKE else 0\n        if stop_at and state.global_step >= stop_at:\n            control.should_save = True\n            control.should_training_stop = True\n        return control\n\ndef last_complete_checkpoint():\n    valid = []\n    for directory in RUN_DIR.glob(\"checkpoint-*\"):\n        try:\n            number = int(directory.name.removeprefix(\"checkpoint-\"))\n            manifest = json.loads((directory / \"COMPLETE.json\").read_text(\"utf-8\"))\n            if manifest[\"signature\"] != RUN_SIGNATURE or manifest[\"step\"] != number:\n                continue\n            files = manifest[\"files\"]\n            if not all(name in files for name in REQUIRED):\n                continue\n            if not all((directory / name).is_file() and file_hash(directory / name) == digest\n                       for name, digest in files.items()):\n                continue\n            saved_step = json.loads((directory / \"trainer_state.json\").read_text(\"utf-8\"))[\"global_step\"]\n            if saved_step == number:\n                valid.append((number, directory))\n        except (OSError, ValueError, KeyError):\n            continue\n    return str(max(valid, key=lambda item: item[0])[1]) if valid else None\n\nlatest = last_complete_checkpoint()\nif latest is None and list(RUN_DIR.glob(\"checkpoint-*\")):\n    raise RuntimeError(\"Checkpoint directories exist but none passed verification. Inspect before starting again.\")\nstart_step = json.loads((Path(latest) / \"trainer_state.json\").read_text(\"utf-8\"))[\"global_step\"] if latest else 0\nprint(\"Resume from:\", latest, \"start_step:\", start_step)\n"
    },
    {
      "id": "cell-08",
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 4. 训练或续训\n\n断线后重连 GPU，重新运行全部单元格。不要删除 RUN_DIR。\nTrainer 根据 checkpoint 恢复 adapter、optimizer、scheduler 和步数。\n"
    },
    {
      "id": "cell-09",
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": "from transformers import Trainer, TrainingArguments, default_data_collator\n\nargs = TrainingArguments(\n    output_dir=str(RUN_DIR),\n    per_device_train_batch_size=config[\"batch_size\"],\n    gradient_accumulation_steps=config[\"gradient_accumulation\"],\n    learning_rate=config[\"learning_rate\"], max_steps=MAX_STEPS,\n    lr_scheduler_type=\"linear\", warmup_steps=0,\n    logging_steps=1, save_strategy=\"steps\", save_steps=SAVE_STEPS,\n    save_total_limit=2, save_only_model=False, save_safetensors=True,\n    report_to=\"none\", disable_tqdm=True, fp16=not SMOKE, use_cpu=SMOKE,\n    dataloader_num_workers=0, dataloader_pin_memory=not SMOKE,\n    seed=config[\"seed\"], data_seed=config[\"seed\"], optim=\"adamw_torch\",\n)\ntrainer = Trainer(\n    model=model, args=args, train_dataset=dataset,\n    data_collator=default_data_collator, processing_class=tokenizer,\n    callbacks=[CompleteCheckpoint()],\n)\nif start_step < MAX_STEPS:\n    trainer.train(resume_from_checkpoint=latest)\n    assert trainer.state.global_step > start_step\nelse:\n    print(\"This run is already complete. Use a new RUN_NAME for a new experiment.\")\nfinal_checkpoint = last_complete_checkpoint()\nassert final_checkpoint is not None\nfinal_step = json.loads((Path(final_checkpoint) / \"trainer_state.json\").read_text(\"utf-8\"))[\"global_step\"]\nsummary = {\"start_step\": start_step, \"final_step\": final_step, \"checkpoint\": Path(final_checkpoint).name}\n(RUN_DIR / \"last_run_summary.json\").write_text(json.dumps(summary, indent=2), encoding=\"utf-8\")\nprint(summary)\n"
    },
    {
      "id": "cell-10",
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 5. 从持久化 checkpoint 重新加载 adapter 做推理\n\n基础模型与 adapter 分别加载。下面的推理不会把训练状态写成新的 checkpoint。\n"
    },
    {
      "id": "cell-11",
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": "from peft import PeftModel\nimport gc\n\ndel trainer, model\ngc.collect()\nif torch.cuda.is_available():\n    torch.cuda.empty_cache()\nset_seed(config[\"seed\"])\nif SMOKE:\n    base = Qwen2ForCausalLM(Qwen2Config(\n        vocab_size=len(tokenizer), hidden_size=32, intermediate_size=64,\n        num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2,\n        max_position_embeddings=256, pad_token_id=tokenizer.pad_token_id,\n        eos_token_id=tokenizer.eos_token_id,\n    ))\n    prompt = \"user checkpoint assistant \"\nelse:\n    base = AutoModelForCausalLM.from_pretrained(model_id, revision=revision, torch_dtype=torch.float16)\n    prompt = tokenizer.apply_chat_template(\n        [{\"role\": \"user\", \"content\": \"训练断线后从哪里恢复？\"}],\n        tokenize=False, add_generation_prompt=True,\n    )\nrestored = PeftModel.from_pretrained(base, final_checkpoint)\ndevice = \"cpu\" if SMOKE else \"cuda\"\nrestored = restored.to(device).eval()\nrestored.config.use_cache = True\ninputs = tokenizer(prompt, return_tensors=\"pt\", return_token_type_ids=False).to(device)\nwith torch.inference_mode():\n    generated = restored.generate(**inputs, max_new_tokens=24, do_sample=False, pad_token_id=tokenizer.pad_token_id)\nprint(tokenizer.decode(generated[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True))\nprint(\"Adapter reload and inference completed.\")\n"
    },
    {
      "id": "cell-12",
      "cell_type": "markdown",
      "metadata": {},
      "source": "## 恢复故障排查\n\n- 新运行时找不到目录，检查是否挂载了同一 Google Drive 账号。\n- run_config 不一致，恢复原模型、依赖和数据，或者换新 RUN_NAME 开新实验。\n- 最后一个目录没有 COMPLETE.json，等待保存结束或使用之前的完整 checkpoint。\n- Drive 写入报错，检查空间与配额后重试。保存频率应结合单步耗时设置。\n- 下载整个训练目录后妥善保管。final adapter 可用于推理，续训还需要完整 checkpoint。\n\n参考 [Colab FAQ](https://research.google.com/colaboratory/faq.html)、\n[Transformers Trainer](https://huggingface.co/docs/transformers/en/main_classes/trainer)、\n[PEFT checkpoint](https://huggingface.co/docs/peft/en/developer_guides/checkpoint)。\n"
    }
  ]
}