File size: 4,766 Bytes
994182c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | #!/usr/bin/env python3
"""Phase 0 smoke checks for Qwen/Qwen3.6-27B.
Default mode avoids downloading full weights. Use --mode load only on the GPU
host after confirming disk and VRAM are sufficient.
"""
from __future__ import annotations
import argparse
import json
import sys
from typing import Any
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--model", default="Qwen/Qwen3.6-27B")
parser.add_argument("--mode", choices=["config", "load"], default="config")
parser.add_argument("--dtype", choices=["auto", "bfloat16", "float16", "float32"], default="bfloat16")
parser.add_argument("--device-map", default="auto")
parser.add_argument("--trust-remote-code", action="store_true", default=True)
parser.add_argument("--no-trust-remote-code", dest="trust_remote_code", action="store_false")
return parser.parse_args()
def print_json(payload: dict[str, Any]) -> None:
print(json.dumps(payload, indent=2, sort_keys=True, default=str))
def dtype_from_name(name: str):
if name == "auto":
return "auto"
import torch
return {
"bfloat16": torch.bfloat16,
"float16": torch.float16,
"float32": torch.float32,
}[name]
def main() -> int:
args = parse_args()
try:
import transformers
from transformers import AutoConfig, AutoTokenizer
except Exception as exc:
print(f"Failed to import transformers: {exc!r}", file=sys.stderr)
return 1
config = AutoConfig.from_pretrained(args.model, trust_remote_code=args.trust_remote_code)
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=args.trust_remote_code)
chat_template = getattr(tokenizer, "chat_template", "") or ""
summary: dict[str, Any] = {
"model": args.model,
"transformers_version": transformers.__version__,
"architectures": getattr(config, "architectures", None),
"model_type": getattr(config, "model_type", None),
"torch_dtype": str(getattr(config, "torch_dtype", None)),
"vocab_size": getattr(config, "vocab_size", None),
"eos_token": tokenizer.eos_token,
"pad_token": tokenizer.pad_token,
"has_think_template": "<think>" in chat_template and "</think>" in chat_template,
"has_tool_call_template": "<tool_call>" in chat_template,
"chat_template_chars": len(chat_template),
}
if args.mode == "config":
print_json(summary)
return 0
import torch
candidate_class_names = [
"AutoModelForMultimodalLM",
"AutoModelForImageTextToText",
"AutoModelForVision2Seq",
"AutoModelForCausalLM",
]
errors: list[str] = []
model = None
loaded_with = None
for class_name in candidate_class_names:
model_cls = getattr(transformers, class_name, None)
if model_cls is None:
errors.append(f"{class_name}: not present in transformers {transformers.__version__}")
continue
try:
model = model_cls.from_pretrained(
args.model,
torch_dtype=dtype_from_name(args.dtype),
device_map=args.device_map,
trust_remote_code=args.trust_remote_code,
)
loaded_with = class_name
break
except Exception as exc:
errors.append(f"{class_name}: {exc!r}")
if model is None:
summary["load_errors"] = errors
print_json(summary)
return 2
messages = [
{"role": "system", "content": "You are validating an authorized security research training environment."},
{"role": "user", "content": "Return one sentence confirming that the model can format a thinking response."},
]
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=True,
)
inputs = tokenizer(prompt, return_tensors="pt")
device = next(model.parameters()).device
inputs = {key: value.to(device) for key, value in inputs.items()}
with torch.no_grad():
output_ids = model.generate(**inputs, max_new_tokens=96, do_sample=False)
decoded = tokenizer.decode(output_ids[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=False)
summary.update(
{
"loaded_with": loaded_with,
"device": str(device),
"cuda_available": torch.cuda.is_available(),
"cuda_device_count": torch.cuda.device_count(),
"generated_chars": len(decoded),
"generated_preview": decoded[:500],
}
)
print_json(summary)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|