Skip to content

docs(training-hub): add QLoRA & CPT tutorials for Ascend NPU #279

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
typhoonzero wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

docs(training-hub): add QLoRA & CPT tutorials for Ascend NPU #279

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
355 changes: 355 additions & 0 deletions docs/en/training_guides/cpt-npu-tutorial.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,355 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "0-title",
"metadata": {},
"source": [
"# Continued Pre-training (CPT) on Ascend NPU\n",
"\n",
"This notebook runs **Continued Pre-training** — plain next-token prediction over a raw-text corpus — on a Huawei Ascend NPU (e.g. `Ascend910B4`) using only mainline `transformers` + `torch_npu`. No CUDA, no `bitsandbytes`, no `training_hub`.\n",
"\n",
"> Why a separate NPU notebook? `training_hub` (the library `cpt-comprehensive-tutorial.ipynb` uses) targets NVIDIA GPUs via `torch.cuda`, `bitsandbytes`, and `unsloth`. None of those run on Ascend. The recipe here reproduces the same *training* using `transformers.Trainer` on `torch_npu`, so the resulting checkpoint is a drop-in HF-format model just like the CUDA path produces.\n",
"\n",
"**Runtime prerequisites**\n",
"\n",
"- Ascend driver + CANN runtime installed on the node (matching the workbench image's CANN version).\n",
"- Workbench image: `PyTorch CANN` (Python 3.12, CANN 8.5.0, PyTorch 2.9.0, `torch_npu 2.9.0`) or equivalent.\n",
"- At least 1 NPU visible via the Kubernetes device plugin (`huawei.com/Ascend910B4` or your cluster's resource name).\n",
"- Persistent storage for the base model + checkpoints (`/opt/app-root/src/...`)."
]
},
{
"cell_type": "markdown",
"id": "1-when",
"metadata": {},
"source": [
"## When to use CPT\n",
"\n",
"CPT keeps the original **causal-LM** objective and updates *all* weights, so it's the tool for **injecting domain knowledge** before instruction tuning. Typical pipeline:\n",
"\n",
"```\n",
"raw corpus → CPT (this notebook) → SFT / LoRA → alignment\n",
"```\n",
"\n",
"CPT catastrophically forgets general-domain skills if you run it too long on a narrow corpus. Mitigations: keep learning rate small (`1e-5`–`5e-6`), mix general-domain text into the corpus, and always follow up with SFT/OSFT."
]
},
{
"cell_type": "markdown",
"id": "2-env",
"metadata": {},
"source": [
"## Step 1 — Environment check\n",
"\n",
"Confirm `torch_npu` is importable and at least one NPU is visible. `torch_npu` monkey-patches `torch` on import so `torch.npu.*` becomes available."
]
},
{
"cell_type": "code",
"id": "2-env-code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"# NPU allocator tuning — expandable segments avoid fragmentation on Ascend.\n",
"os.environ.setdefault(\"PYTORCH_NPU_ALLOC_CONF\", \"expandable_segments:True\")\n",
"\n",
"import torch\n",
"import torch_npu # noqa: F401 — side-effect import registers torch.npu\n",
"import transformers\n",
"\n",
"print(\"torch:\", torch.__version__)\n",
"print(\"torch_npu:\", torch_npu.__version__)\n",
"print(\"transformers:\", transformers.__version__)\n",
"print(\"NPU available:\", torch.npu.is_available())\n",
"print(\"NPU count:\", torch.npu.device_count())\n",
"if torch.npu.is_available():\n",
" torch.npu.set_device(0)\n",
" print(\"device 0:\", torch.npu.get_device_name(0))\n",
" # bf16 support is standard on Ascend910B — quick sanity check.\n",
" x = torch.randn(64, 64, dtype=torch.bfloat16, device=\"npu\")\n",
" y = (x @ x.T).float().mean().item()\n",
" print(f\"bf16 matmul ok, mean={y:.4f}\")"
]
},
{
"cell_type": "markdown",
"id": "3-params",
"metadata": {},
"source": [
"## Step 2 — Parameters\n",
"\n",
"Edit these to point at your model and corpus. Defaults use a small base and short training so the notebook completes as a smoke test in a few minutes on a single NPU."
]
},
{
"cell_type": "code",
"id": "3-params-code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# --- model & corpus ---\n",
"MODEL_PATH = os.environ.get(\"HF_MODEL_DIR\", \"/opt/app-root/src/models/Qwen2.5-0.5B\")\n",
"CORPUS_PATH = os.environ.get(\"CORPUS_PATH\", \"/opt/app-root/src/datasets/cpt/corpus.jsonl\")\n",
"OUTPUT_DIR = os.environ.get(\"OUTPUT_DIR\", \"/opt/app-root/src/cpt-npu-output\")\n",
"\n",
"# --- training ---\n",
"BLOCK_SIZE = 512 # chunk length used to build training samples\n",
"NUM_EPOCHS = 1\n",
"PER_DEVICE_BATCH_SIZE = 2\n",
"GRAD_ACC_STEPS = 4 # effective batch = 2 * 4 = 8 tokens/step\n",
"LEARNING_RATE = 5e-6 # keep small — CPT is easy to overfit / forget\n",
"WARMUP_STEPS = 20\n",
"SAVE_STEPS = 200\n",
"LOGGING_STEPS = 10\n",
"\n",
"# --- compute ---\n",
"USE_BF16 = True # Ascend910B supports bf16 natively\n",
"SEED = 42\n",
"\n",
"os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
"print(\"model :\", MODEL_PATH)\n",
"print(\"corpus:\", CORPUS_PATH)\n",
"print(\"output:\", OUTPUT_DIR)"
]
},
{
"cell_type": "markdown",
"id": "4-data-format",
"metadata": {},
"source": [
"## Step 3 — Data format\n",
"\n",
"CPT expects **raw text**, not chat turns. Each line of the JSONL is one document under a `text` field:\n",
"\n",
"```json\n",
"{\"text\": \"Alauda AI is an MLOps platform that runs training and inference on Kubernetes.\"}\n",
"{\"text\": \"Continued pre-training adapts a base model to a new domain using unlabeled text.\"}\n",
"```\n",
"\n",
"The next cell falls back to a tiny **built-in synthetic corpus** if `CORPUS_PATH` doesn't exist, so the notebook can be run end-to-end without any dataset upload — useful as a smoke test."
]
},
{
"cell_type": "code",
"id": "4-data-code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import json\n",
"from pathlib import Path\n",
"\n",
"if not Path(CORPUS_PATH).exists():\n",
" fallback = Path(OUTPUT_DIR) / \"synthetic_corpus.jsonl\"\n",
" docs = [\n",
" \"Alauda AI is an MLOps platform that runs fine-tuning and inference on Kubernetes.\",\n",
" \"Continued pre-training keeps the causal-LM objective and updates every weight.\",\n",
" \"Huawei Ascend NPUs are programmed through the CANN runtime and torch_npu.\",\n",
" \"A CPT run should be followed by SFT to restore instruction-following behaviour.\",\n",
" \"transformers.Trainer works on torch_npu out of the box once torch_npu is imported.\",\n",
" ] * 20\n",
" with fallback.open(\"w\") as f:\n",
" for d in docs:\n",
" json.dump({\"text\": d}, f); f.write(\"\\n\")\n",
" CORPUS_PATH = str(fallback)\n",
" print(f\"corpus not found; wrote synthetic fallback to {CORPUS_PATH}\")\n",
"print(\"corpus:\", CORPUS_PATH)"
]
},
{
"cell_type": "markdown",
"id": "5-tokenize",
"metadata": {},
"source": [
"## Step 4 — Load model + tokenizer, tokenize corpus\n",
"\n",
"Load the base model in `bf16` (Ascend-native), then tokenize the corpus and pack it into fixed-length chunks of `BLOCK_SIZE`. For CPT, `labels == input_ids` — every token contributes to the loss."
]
},
{
"cell_type": "code",
"id": "5-tokenize-code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"from datasets import load_dataset\n",
"\n",
"tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=True)\n",
"if tokenizer.pad_token is None:\n",
" tokenizer.pad_token = tokenizer.eos_token\n",
"\n",
"model = AutoModelForCausalLM.from_pretrained(\n",
" MODEL_PATH,\n",
" torch_dtype=torch.bfloat16 if USE_BF16 else torch.float32,\n",
" attn_implementation=\"eager\", # SDPA works on NPU; eager is the safe default across CANN versions\n",
")\n",
"model.to(\"npu\")\n",
"model.gradient_checkpointing_enable()\n",
"\n",
"raw = load_dataset(\"json\", data_files=CORPUS_PATH, split=\"train\")\n",
"print(\"raw documents:\", len(raw))\n",
"\n",
"def tokenize(batch):\n",
" return tokenizer(batch[\"text\"], add_special_tokens=False)\n",
"\n",
"tokenized = raw.map(tokenize, batched=True, remove_columns=raw.column_names)\n",
"\n",
"def group_texts(examples):\n",
" concatenated = {k: sum(examples[k], []) for k in examples.keys()}\n",
" total_len = (len(concatenated[\"input_ids\"]) // BLOCK_SIZE) * BLOCK_SIZE\n",
" result = {\n",
" k: [t[i : i + BLOCK_SIZE] for i in range(0, total_len, BLOCK_SIZE)]\n",
" for k, t in concatenated.items()\n",
" }\n",
" result[\"labels\"] = [ids.copy() for ids in result[\"input_ids\"]]\n",
" return result\n",
"\n",
"packed = tokenized.map(group_texts, batched=True)\n",
"print(\"packed chunks:\", len(packed), \"of\", BLOCK_SIZE, \"tokens each\")"
]
},
{
"cell_type": "markdown",
"id": "6-train",
"metadata": {},
"source": [
"## Step 5 — Train\n",
"\n",
"Use plain `Trainer` with `bf16=True`. **Do not** set `fp16=True` — Ascend prefers bf16, and mixing fp16 with older CANN builds triggers dynamic-shape recompilation.\n",
"\n",
"The optimizer is `adamw_torch` (mainline PyTorch). Do not use `adamw_bnb_8bit` — bitsandbytes is not upstream on NPU (see the QLoRA-NPU notebook for how to bring in the community `bitsandbytes-npu` fork if you need it)."
]
},
{
"cell_type": "code",
"id": "6-train-code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling\n",
"\n",
"collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n",
"\n",
"args = TrainingArguments(\n",
" output_dir=OUTPUT_DIR,\n",
" num_train_epochs=NUM_EPOCHS,\n",
" per_device_train_batch_size=PER_DEVICE_BATCH_SIZE,\n",
" gradient_accumulation_steps=GRAD_ACC_STEPS,\n",
" learning_rate=LEARNING_RATE,\n",
" warmup_steps=WARMUP_STEPS,\n",
" logging_steps=LOGGING_STEPS,\n",
" save_steps=SAVE_STEPS,\n",
" save_total_limit=2,\n",
" bf16=USE_BF16,\n",
" fp16=False,\n",
" optim=\"adamw_torch\",\n",
" lr_scheduler_type=\"cosine\",\n",
" seed=SEED,\n",
" report_to=[],\n",
" remove_unused_columns=False,\n",
" dataloader_pin_memory=False, # pinned memory is not supported on NPU\n",
")\n",
"\n",
"trainer = Trainer(\n",
" model=model,\n",
" args=args,\n",
" train_dataset=packed,\n",
" data_collator=collator,\n",
")\n",
"\n",
"train_result = trainer.train()\n",
"print(train_result.metrics)"
]
},
{
"cell_type": "markdown",
"id": "7-save",
"metadata": {},
"source": [
"## Step 6 — Save the HF-format checkpoint\n",
"\n",
"Save the tuned weights and tokenizer under `OUTPUT_DIR/hf_format` so downstream SFT / inference recipes can load it with `AutoModelForCausalLM.from_pretrained(...)`."
]
},
{
"cell_type": "code",
"id": "7-save-code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"hf_dir = os.path.join(OUTPUT_DIR, \"hf_format\")\n",
"trainer.save_model(hf_dir)\n",
"tokenizer.save_pretrained(hf_dir)\n",
"print(\"saved:\", hf_dir)\n",
"print(\"contents:\", sorted(os.listdir(hf_dir)))"
]
},
{
"cell_type": "markdown",
"id": "8-inference",
"metadata": {},
"source": [
"## Step 7 — Quick inference smoke\n",
"\n",
"Load the freshly-saved checkpoint back onto the NPU and generate a few tokens. This is a *sanity* check — the base model is small and the training corpus is tiny, so don't judge quality here."
]
},
{
"cell_type": "code",
"id": "8-infer-code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
"\n",
"tok = AutoTokenizer.from_pretrained(hf_dir)\n",
"mdl = AutoModelForCausalLM.from_pretrained(hf_dir, torch_dtype=torch.bfloat16).to(\"npu\")\n",
"mdl.eval()\n",
"\n",
"prompt = \"Alauda AI is\"\n",
"inputs = tok(prompt, return_tensors=\"pt\").to(\"npu\")\n",
"with torch.no_grad():\n",
" out = mdl.generate(**inputs, max_new_tokens=32, do_sample=False)\n",
"print(tok.decode(out[0], skip_special_tokens=True))"
]
},
{
"cell_type": "markdown",
"id": "9-notes",
"metadata": {},
"source": [
"## Notes and gotchas on NPU\n",
"\n",
"- **Import order.** `import torch_npu` **before** the first `torch.npu.*` call — it registers the backend.\n",
"- **Precision.** Use `bf16`. Avoid `fp16` unless you know your CANN build supports it end-to-end; mixing fp16 with older builds triggers dynamic-shape retracing that stalls training.\n",
"- **Optimizer.** Stick to `adamw_torch`. `adamw_bnb_8bit` requires `bitsandbytes` — see the [QLoRA-NPU tutorial](./qlora-npu-tutorial.ipynb) for how to install the Ascend fork if you need paged optimizers.\n",
"- **Attention backend.** `attn_implementation=\"eager\"` is the portable default. `sdpa` works on most CANN 8.x builds but occasionally hits an unsupported-op path on prerelease combinations. `flash_attention_2` is CUDA-only.\n",
"- **Data loader.** Set `dataloader_pin_memory=False`; NPU does not use CUDA pinned host memory and the flag can trip guard checks.\n",
"- **Multi-NPU.** For >1 NPU on a single node, launch this notebook's script form under `torch.distributed` via `torchrun --nproc_per_node=<npu_count>`. `Trainer` picks up `torch.distributed` transparently on NPU thanks to `hccl`.\n",
"- **Catastrophic forgetting.** Keep `learning_rate` in the `1e-5` – `5e-6` range for CPT and always follow up with SFT — CPT alone breaks instruction-following."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.12"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Loading