team-qbits / training /check_env.py
ovinduG's picture
Upload folder using huggingface_hub
e55c8a4 verified
Raw History Blame Contribute Delete
6.67 kB
#!/usr/bin/env python
"""
Environment check for cpt_gemma4_sinhala.py (Unsloth + Gemma 4 E4B on an A40).
python check_env.py # quick checks (no model download)
python check_env.py --model # also loads Gemma 4 E4B, adds LoRA, runs 1 train step (~16 GB download)
Exit code 0 = everything passed.
"""
import argparse
import importlib
import importlib.metadata as md
import subprocess
import sys
OK, BAD, WARN = "[ OK ]", "[FAIL]", "[WARN]"
failures = []
def report(tag, msg):
print(f"{tag} {msg}")
if tag == BAD:
failures.append(msg)
def ver(pkg):
try:
return md.version(pkg)
except md.PackageNotFoundError:
return None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", action="store_true", help="load model + 1 training step")
ap.add_argument("--model_name", default="unsloth/gemma-4-E4B")
args = ap.parse_args()
print(f"Python {sys.version.split()[0]}")
# 1. GPU / torch (import unsloth FIRST, as in the training script)
try:
import unsloth # noqa: F401
report(OK, f"unsloth imported (v{ver('unsloth')})")
except Exception as e: # noqa: BLE001
report(BAD, f"unsloth import failed: {type(e).__name__}: {e}")
try:
import torch
report(OK, f"torch {torch.__version__}, CUDA build {torch.version.cuda}")
if torch.cuda.is_available():
p = torch.cuda.get_device_properties(0)
report(OK, f"GPU: {p.name}, {p.total_memory / 1024**3:.1f} GB, bf16={torch.cuda.is_bf16_supported()}")
else:
report(BAD, "torch.cuda.is_available() is False (Unsloth needs a GPU)")
except Exception as e: # noqa: BLE001
report(BAD, f"torch import failed: {e}")
# 2. Pinned versions
transformers_v = ver("transformers")
if transformers_v == "5.5.0":
report(OK, "transformers 5.5.0")
else:
report(BAD, f"transformers is {transformers_v}, expected 5.5.0 "
"(re-run: pip install --no-deps transformers==5.5.0)")
tok_v = ver("tokenizers")
try:
from packaging.version import Version
good = tok_v and Version("0.22.0") <= Version(tok_v) <= Version("0.23.0")
except Exception: # noqa: BLE001
good = False
report(OK if good else BAD, f"tokenizers {tok_v} (need >=0.22.0,<=0.23.0)")
hub_v = ver("huggingface_hub")
try:
from packaging.version import Version
good = hub_v and Version("1.5.0") <= Version(hub_v) < Version("2.0")
except Exception: # noqa: BLE001
good = False
report(OK if good else BAD, f"huggingface_hub {hub_v} (need >=1.5.0,<2.0)")
# 3. Other packages the script imports
for pkg in ["trl", "peft", "datasets", "pandas", "accelerate", "bitsandbytes", "timm", "hf_transfer"]:
v = ver(pkg)
report(OK if v else BAD, f"{pkg} {v or 'NOT INSTALLED'}")
# 4. Imports the training script actually does
for stmt in ["from unsloth import FastModel",
"from trl import SFTConfig, SFTTrainer",
"from datasets import Dataset",
"from huggingface_hub import hf_hub_download"]:
try:
exec(stmt, {}) # noqa: S102
report(OK, stmt)
except Exception as e: # noqa: BLE001
report(BAD, f"{stmt} -> {type(e).__name__}: {e}")
# 5. Do the SFTConfig args used by the script exist in this TRL version?
try:
from trl import SFTConfig
import inspect
fields = set(inspect.signature(SFTConfig.__init__).parameters)
for a in ["max_length", "packing", "dataset_text_field", "bf16", "optim",
"eval_strategy", "warmup_ratio"]:
report(OK if a in fields else BAD, f"SFTConfig accepts '{a}'")
except Exception as e: # noqa: BLE001
report(WARN, f"could not inspect SFTConfig: {e}")
# 6. pip's own dependency check (informational: --no-deps installs can show conflicts)
out = subprocess.run([sys.executable, "-m", "pip", "check"], capture_output=True, text=True).stdout.strip()
if out and "No broken requirements" not in out:
print(f"{WARN} pip check reported:\n{out}\n (conflicts about transformers/tokenizers are often "
"harmless with --no-deps; only worry if the model test below fails)")
else:
report(OK, "pip check: no broken requirements")
# 7. Optional end-to-end test
if args.model:
try:
import torch
from unsloth import FastModel
from datasets import Dataset
from trl import SFTConfig, SFTTrainer
print(f"\nLoading {args.model_name} ...")
model, tok = FastModel.from_pretrained(
model_name=args.model_name, dtype=None, max_seq_length=512,
load_in_4bit=False, full_finetuning=False)
report(OK, "model + tokenizer loaded")
model = FastModel.get_peft_model(
model, finetune_vision_layers=False, finetune_language_layers=True,
finetune_attention_modules=True, finetune_mlp_modules=True,
r=8, lora_alpha=8, lora_dropout=0, bias="none", use_rslora=True, random_state=3407)
report(OK, "LoRA attached (use_rslora accepted)")
sin = "ශ්‍රී ලංකාව ඉන්දියන් සාගරයේ පිහිටි ද්වීපයකි."
ids = tok(sin, add_special_tokens=False)["input_ids"]
print(f" Sinhala sample -> {len(ids)} tokens for {len(sin)} chars")
ds = Dataset.from_dict({"text": [(sin + "\n") * 20] * 16})
cfg = SFTConfig(output_dir="/tmp/envcheck", dataset_text_field="text", max_length=512,
packing=True, per_device_train_batch_size=2, gradient_accumulation_steps=1,
max_steps=2, learning_rate=5e-5, bf16=True, optim="adamw_8bit",
logging_steps=1, save_strategy="no", report_to="none")
tr = SFTTrainer(model=model, tokenizer=tok, train_dataset=ds, args=cfg)
tr.train()
report(OK, f"2 training steps ran; peak VRAM {torch.cuda.max_memory_reserved() / 1024**3:.1f} GB")
except Exception as e: # noqa: BLE001
import traceback
traceback.print_exc()
report(BAD, f"model test failed: {type(e).__name__}: {e}")
print("\n" + ("ALL CHECKS PASSED" if not failures else f"{len(failures)} PROBLEM(S):"))
for f in failures:
print(" -", f)
sys.exit(1 if failures else 0)
if __name__ == "__main__":
main()