#!/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()