[pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
This commit is contained in:
parent
c6d92160f6
commit
a75aef063c
1 changed files with 51 additions and 26 deletions
|
|
@ -14,6 +14,7 @@ Run inside the container:
|
|||
Skip step 5 (faster, no model download):
|
||||
docker run --rm --gpus all unsloth-blackwell:latest python /workspace/smoke_test.py --skip-train
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
|
@ -21,12 +22,13 @@ import sys
|
|||
|
||||
|
||||
def banner(title: str) -> None:
|
||||
print(f"\n=== {title} ===", flush=True)
|
||||
print(f"\n=== {title} ===", flush = True)
|
||||
|
||||
|
||||
def check_torch() -> tuple[int, int]:
|
||||
banner("torch + arch list")
|
||||
import torch
|
||||
|
||||
# Use the raw C++ accessor so this works even when CUDA isn't available
|
||||
# (lets us run a partial smoke test on a no-GPU host).
|
||||
arches = torch._C._cuda_getArchFlags().split()
|
||||
|
|
@ -47,13 +49,27 @@ def check_torch() -> tuple[int, int]:
|
|||
|
||||
def check_imports() -> None:
|
||||
banner("dep imports")
|
||||
import triton; print(f"triton {triton.__version__}")
|
||||
import xformers; print(f"xformers {xformers.__version__}")
|
||||
import bitsandbytes as bnb; print(f"bnb {bnb.__version__}")
|
||||
import transformers; print(f"transformers {transformers.__version__}")
|
||||
import trl; print(f"trl {trl.__version__}")
|
||||
import peft; print(f"peft {peft.__version__}")
|
||||
import unsloth_zoo; print(f"unsloth_zoo {unsloth_zoo.__version__}")
|
||||
import triton
|
||||
|
||||
print(f"triton {triton.__version__}")
|
||||
import xformers
|
||||
|
||||
print(f"xformers {xformers.__version__}")
|
||||
import bitsandbytes as bnb
|
||||
|
||||
print(f"bnb {bnb.__version__}")
|
||||
import transformers
|
||||
|
||||
print(f"transformers {transformers.__version__}")
|
||||
import trl
|
||||
|
||||
print(f"trl {trl.__version__}")
|
||||
import peft
|
||||
|
||||
print(f"peft {peft.__version__}")
|
||||
import unsloth_zoo
|
||||
|
||||
print(f"unsloth_zoo {unsloth_zoo.__version__}")
|
||||
|
||||
|
||||
def check_unsloth_import() -> None:
|
||||
|
|
@ -63,6 +79,7 @@ def check_unsloth_import() -> None:
|
|||
# That's fine for this smoke -- we're not training Unsloth-patched models yet.
|
||||
import unsloth
|
||||
from unsloth import FastLanguageModel
|
||||
|
||||
print(f"unsloth {unsloth.__version__}")
|
||||
print(f"FastLanguageModel {FastLanguageModel}")
|
||||
|
||||
|
|
@ -70,6 +87,7 @@ def check_unsloth_import() -> None:
|
|||
def check_tiny_train(cap: tuple[int, int]) -> None:
|
||||
banner("tiny LoRA train (5 steps)")
|
||||
import os
|
||||
|
||||
# Unsloth must be imported first.
|
||||
import unsloth # noqa: F401
|
||||
from unsloth import FastLanguageModel
|
||||
|
|
@ -79,20 +97,20 @@ def check_tiny_train(cap: tuple[int, int]) -> None:
|
|||
model_name = "unsloth/Llama-3.2-1B-Instruct-bnb-4bit"
|
||||
print(f"loading {model_name}")
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name=model_name,
|
||||
max_seq_length=512,
|
||||
dtype=None,
|
||||
load_in_4bit=True,
|
||||
model_name = model_name,
|
||||
max_seq_length = 512,
|
||||
dtype = None,
|
||||
load_in_4bit = True,
|
||||
)
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=8,
|
||||
lora_alpha=16,
|
||||
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||
lora_dropout=0.0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
random_state=0,
|
||||
r = 8,
|
||||
lora_alpha = 16,
|
||||
target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"],
|
||||
lora_dropout = 0.0,
|
||||
bias = "none",
|
||||
use_gradient_checkpointing = "unsloth",
|
||||
random_state = 0,
|
||||
)
|
||||
|
||||
prompts = [
|
||||
|
|
@ -101,26 +119,33 @@ def check_tiny_train(cap: tuple[int, int]) -> None:
|
|||
"Q: Name a primary color.\nA:",
|
||||
"Q: Hello, who are you?\nA:",
|
||||
] * 2
|
||||
enc = tokenizer(prompts, return_tensors="pt", padding=True, truncation=True, max_length=64)
|
||||
enc = tokenizer(
|
||||
prompts, return_tensors = "pt", padding = True, truncation = True, max_length = 64
|
||||
)
|
||||
enc = {k: v.cuda() for k, v in enc.items()}
|
||||
labels = enc["input_ids"].clone()
|
||||
|
||||
model.train()
|
||||
optim = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-4)
|
||||
optim = torch.optim.AdamW(
|
||||
[p for p in model.parameters() if p.requires_grad], lr = 1e-4
|
||||
)
|
||||
for step in range(5):
|
||||
out = model(**enc, labels=labels)
|
||||
out = model(**enc, labels = labels)
|
||||
out.loss.backward()
|
||||
optim.step()
|
||||
optim.zero_grad(set_to_none=True)
|
||||
print(f"step {step} loss={out.loss.item():.4f}", flush=True)
|
||||
optim.zero_grad(set_to_none = True)
|
||||
print(f"step {step} loss={out.loss.item():.4f}", flush = True)
|
||||
|
||||
print("OK: 5 LoRA steps completed")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--skip-train", action="store_true",
|
||||
help="Skip the tiny LoRA training step (no HF download).")
|
||||
ap.add_argument(
|
||||
"--skip-train",
|
||||
action = "store_true",
|
||||
help = "Skip the tiny LoRA training step (no HF download).",
|
||||
)
|
||||
args = ap.parse_args()
|
||||
|
||||
cap = check_torch()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue