135 lines
4.4 KiB
Python
135 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import yaml
|
|
|
|
from common import ROOT, localize_workspace_path
|
|
|
|
DEFAULT_CONFIG = ROOT / "plugins" / "1c" / "training" / "configs" / "qwen3-coder-30b-a3b-lora.yaml"
|
|
DEFAULT_REQUIRED_BASE_FILES = [
|
|
"config.json",
|
|
"tokenizer.json",
|
|
"tokenizer_config.json",
|
|
]
|
|
|
|
|
|
def load_config(path: Path) -> dict:
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
data = yaml.safe_load(handle)
|
|
if not isinstance(data, dict):
|
|
raise ValueError(f"{path} must contain a YAML mapping")
|
|
return data
|
|
|
|
|
|
def count_jsonl(path: Path) -> int:
|
|
if not path.exists():
|
|
return 0
|
|
count = 0
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
if line.strip():
|
|
json.loads(line)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
def check_python_deps() -> list[str]:
|
|
missing = []
|
|
for module in ("torch", "transformers", "datasets", "peft", "accelerate"):
|
|
try:
|
|
__import__(module)
|
|
except Exception as exc:
|
|
missing.append(f"{module}: {type(exc).__name__}: {exc}")
|
|
return missing
|
|
|
|
|
|
def check_cuda() -> tuple[bool, str]:
|
|
try:
|
|
import torch
|
|
except Exception as exc:
|
|
return False, f"torch unavailable: {exc}"
|
|
if not torch.cuda.is_available():
|
|
return False, "torch.cuda.is_available() is false"
|
|
return True, torch.cuda.get_device_name(0)
|
|
|
|
|
|
def expected_base_files(config: dict) -> list[str]:
|
|
configured = config.get("required_base_files")
|
|
if isinstance(configured, list):
|
|
return [str(item) for item in configured if str(item).strip()]
|
|
return list(DEFAULT_REQUIRED_BASE_FILES)
|
|
|
|
|
|
def missing_weight_shards(base_model_path: Path) -> list[str]:
|
|
index_path = base_model_path / "model.safetensors.index.json"
|
|
if not index_path.exists():
|
|
return []
|
|
try:
|
|
index = json.loads(index_path.read_text(encoding="utf-8"))
|
|
except json.JSONDecodeError as exc:
|
|
return [f"invalid model.safetensors.index.json: {exc}"]
|
|
weight_map = index.get("weight_map")
|
|
if not isinstance(weight_map, dict):
|
|
return ["model.safetensors.index.json has no weight_map"]
|
|
shards = sorted({str(value) for value in weight_map.values() if str(value).strip()})
|
|
return [name for name in shards if not (base_model_path / name).exists()]
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Preflight checks for 1C LoRA training.")
|
|
parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG)
|
|
args = parser.parse_args()
|
|
|
|
errors: list[str] = []
|
|
warnings: list[str] = []
|
|
config = load_config(args.config)
|
|
|
|
dataset_path = localize_workspace_path(config["dataset_path"])
|
|
dataset_count = count_jsonl(dataset_path)
|
|
if dataset_count == 0:
|
|
errors.append(f"training dataset is missing or empty: {dataset_path}")
|
|
elif dataset_count < 50:
|
|
warnings.append(f"dataset has only {dataset_count} record(s); this is not enough for useful fine-tuning")
|
|
|
|
base_model_path = localize_workspace_path(config["base_model_path"])
|
|
missing_files = [name for name in expected_base_files(config) if not (base_model_path / name).exists()]
|
|
if missing_files:
|
|
errors.append(f"base model is incomplete at {base_model_path}: missing {', '.join(missing_files)}")
|
|
missing_shards = missing_weight_shards(base_model_path)
|
|
if missing_shards:
|
|
errors.append(f"base model shard set is incomplete at {base_model_path}: missing {', '.join(missing_shards)}")
|
|
|
|
output_dir = localize_workspace_path(config["output_dir"])
|
|
if not output_dir.parent.exists():
|
|
warnings.append(f"adapter parent directory does not exist yet: {output_dir.parent}")
|
|
|
|
missing_deps = check_python_deps()
|
|
if missing_deps:
|
|
errors.append(f"missing Python training dependencies: {'; '.join(missing_deps)}")
|
|
|
|
cuda_ok, cuda_message = check_cuda()
|
|
if cuda_ok:
|
|
print(f"CUDA: {cuda_message}")
|
|
else:
|
|
errors.append(f"GPU/CUDA unavailable: {cuda_message}")
|
|
|
|
for warning in warnings:
|
|
print(f"WARNING: {warning}")
|
|
|
|
if errors:
|
|
print("1C training preflight failed:", file=sys.stderr)
|
|
for error in errors:
|
|
print(f"- {error}", file=sys.stderr)
|
|
return 1
|
|
|
|
print("1C training preflight passed.")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|