[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-05-24 06:54:34 +00:00
commit a75aef063c

View file

@ -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()