mimic-svm / code /merge_adapter.py

Upload code/merge_adapter.py with huggingface_hub

17699c4 verified 7 months ago

10.4 kB

import argparse
import json
import os
import time
from pathlib import Path
from typing import Any, Dict, Optional

import torch
from PIL import Image
from unsloth import FastVisionModel

os.environ["CUDA_VISIBLE_DEVICES"] = "7"


DEFAULT_ADAPTER_PATH = "outputs/mimic_qwen3vl_lora_8bit_5/checkpoint-17454"
DEFAULT_OUTPUT_PATH = "outputs/mimic_qwen3vl_lora_8bit_5_merged"
DEFAULT_INSTRUCTION = "Analyze this chest X-ray image and generate the corresponding radiology report."


DTYPE_MAP = {
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Merge a Qwen3-VL LoRA adapter into base weights using Unsloth. "
            "This intentionally avoids loading 8-bit modules during merge to preserve accuracy."
        )
    )

    parser.add_argument("--adapter_path", type=str, default=DEFAULT_ADAPTER_PATH)
    parser.add_argument("--output_dir", type=str, default=DEFAULT_OUTPUT_PATH)
    parser.add_argument(
        "--save_method",
        type=str,
        default="merged_16bit",
        choices=["merged_16bit", "merged_4bit", "merged_4bit_forced"],
        help=(
            "Unsloth save mode. For best fidelity use merged_16bit. "
            "Use 4bit variants only if you explicitly need compact merged weights."
        ),
    )

    parser.add_argument(
        "--dtype",
        type=str,
        default="auto",
        choices=["auto", "float16", "bfloat16", "float32"],
        help="Compute dtype used while loading model for merge.",
    )
    parser.add_argument("--device_map", type=str, default="auto")
    parser.add_argument("--cuda_visible_devices", type=str, default="")
    parser.add_argument("--max_shard_size", type=str, default="10GB")

    parser.add_argument("--verify_reload", action="store_true", help="Reload merged model after save.")
    parser.add_argument(
        "--verify_load_in_8bit",
        action="store_true",
        help="When verifying, reload merged model in 8-bit inference mode.",
    )
    parser.add_argument(
        "--verify_max_new_tokens",
        type=int,
        default=32,
        help="Max new tokens for optional verification generation.",
    )
    parser.add_argument(
        "--verify_instruction",
        type=str,
        default=DEFAULT_INSTRUCTION,
        help="Prompt text used for optional verification generation.",
    )
    parser.add_argument(
        "--skip_generate_check",
        action="store_true",
        help="If set, verification only checks loadability and skips generation.",
    )

    parser.add_argument(
        "--safe_serialization",
        action="store_true",
        default=True,
        help="Save in safetensors format when supported.",
    )
    parser.add_argument(
        "--no_safe_serialization",
        action="store_true",
        help="Disable safetensors serialization.",
    )

    return parser.parse_args()


def resolve_dtype(dtype_arg: str) -> torch.dtype:
    if dtype_arg != "auto":
        return DTYPE_MAP[dtype_arg]

    if torch.cuda.is_available():
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
        return torch.float16
    return torch.float32


def read_adapter_base_model(adapter_path: Path) -> Optional[str]:
    adapter_config_path = adapter_path / "adapter_config.json"
    if not adapter_config_path.exists():
        return None

    try:
        data = json.loads(adapter_config_path.read_text(encoding="utf-8"))
    except (OSError, ValueError, json.JSONDecodeError):
        return None

    base_model = data.get("base_model_name_or_path")
    return str(base_model) if base_model else None


def print_runtime_info(args: argparse.Namespace, merge_dtype: torch.dtype, adapter_path: Path, output_dir: Path) -> None:
    print("=" * 88)
    print("Merge configuration")
    print("=" * 88)
    print(f"adapter_path           : {adapter_path}")
    print(f"output_dir             : {output_dir}")
    print(f"save_method            : {args.save_method}")
    print(f"dtype                  : {merge_dtype}")
    print(f"device_map             : {args.device_map}")
    print(f"safe_serialization     : {args.safe_serialization and not args.no_safe_serialization}")
    print(f"verify_reload          : {args.verify_reload}")
    print(f"verify_load_in_8bit    : {args.verify_load_in_8bit}")
    print("=" * 88)


def save_merge_metadata(output_dir: Path, metadata: Dict[str, Any]) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)
    metadata_path = output_dir / "merge_metadata.json"
    metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
    print(f"Wrote merge metadata: {metadata_path}")


def merge_adapter(args: argparse.Namespace) -> Path:
    adapter_path = Path(args.adapter_path).expanduser().resolve()
    output_dir = Path(args.output_dir).expanduser().resolve()

    if not adapter_path.exists() or not adapter_path.is_dir():
        raise FileNotFoundError(f"Adapter path does not exist or is not a directory: {adapter_path}")

    if args.cuda_visible_devices:
        os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices

    merge_dtype = resolve_dtype(args.dtype)
    print_runtime_info(args, merge_dtype, adapter_path, output_dir)

    inferred_base_model = read_adapter_base_model(adapter_path)
    if inferred_base_model:
        print(f"Adapter base model from adapter_config.json: {inferred_base_model}")

    print("Loading adapter (with base) in non-quantized mode for accurate merge...")
    model, tokenizer = FastVisionModel.from_pretrained(
        model_name=str(adapter_path),
        load_in_4bit=False,
        load_in_8bit=False,
        dtype=merge_dtype,
        device_map=args.device_map,
    )

    # Ensure inference graph after load to avoid training-mode side effects.
    FastVisionModel.for_inference(model)

    output_dir.mkdir(parents=True, exist_ok=True)

    safe_serialization = args.safe_serialization and not args.no_safe_serialization
    merge_start = time.time()

    print("Saving merged model...")
    try:
        model.save_pretrained_merged(
            str(output_dir),
            tokenizer,
            save_method=args.save_method,
            safe_serialization=safe_serialization,
            max_shard_size=args.max_shard_size,
        )
    except TypeError:
        # Older Unsloth versions may not support all save kwargs.
        model.save_pretrained_merged(
            str(output_dir),
            tokenizer,
            save_method=args.save_method,
        )

    merge_seconds = time.time() - merge_start
    print(f"Merge complete in {merge_seconds:.1f}s")

    save_merge_metadata(
        output_dir,
        {
            "adapter_path": str(adapter_path),
            "inferred_base_model": inferred_base_model,
            "output_dir": str(output_dir),
            "save_method": args.save_method,
            "merge_dtype": str(merge_dtype),
            "device_map": args.device_map,
            "safe_serialization": safe_serialization,
            "merged_at_unix": int(time.time()),
        },
    )

    del model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return output_dir


def verify_merged_model(args: argparse.Namespace, merged_dir: Path) -> None:
    if not args.verify_reload:
        return

    print("Reloading merged model for verification...")

    verify_load_in_8bit = bool(args.verify_load_in_8bit)
    if verify_load_in_8bit and not torch.cuda.is_available():
        print("CUDA unavailable. Falling back to non-8bit verification load.")
        verify_load_in_8bit = False

    model, tokenizer = FastVisionModel.from_pretrained(
        model_name=str(merged_dir),
        load_in_4bit=False,
        load_in_8bit=verify_load_in_8bit,
        device_map=args.device_map,
    )
    FastVisionModel.for_inference(model)

    if args.skip_generate_check:
        print("Verification load successful (generation check skipped).")
        return

    # Minimal generation sanity check with one synthetic image.
    test_image = Image.new("RGB", (224, 224), color=(0, 0, 0))
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": test_image},
                {"type": "text", "text": args.verify_instruction},
            ],
        }
    ]

    prompt_text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = None
    tokenization_errors = []
    for image_argument in (
        {"images": [test_image]},
        {"images": [[test_image]]},
        {"image": [test_image]},
        {"image": [[test_image]]},
    ):
        try:
            inputs = tokenizer(
                text=[prompt_text],
                padding=True,
                return_tensors="pt",
                **image_argument,
            )
            break
        except Exception as error:  # pragma: no cover
            tokenization_errors.append(str(error))

    if inputs is None:
        raise RuntimeError("Tokenization failed during verify: " + " | ".join(tokenization_errors))

    model_device = next(model.parameters()).device
    inputs = {k: (v.to(model_device) if isinstance(v, torch.Tensor) else v) for k, v in inputs.items()}

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max(1, args.verify_max_new_tokens),
            do_sample=False,
        )

    input_token_count = inputs["input_ids"].shape[-1] if "input_ids" in inputs else 0
    generated_ids = outputs[:, input_token_count:]
    generated_text = tokenizer.batch_decode(
        generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False,
    )[0]

    print("Verification generation successful.")
    print("--- Verification Output (truncated to 500 chars) ---")
    print(generated_text[:500])


def main() -> None:
    args = parse_args()

    if args.verify_max_new_tokens <= 0:
        raise ValueError("--verify_max_new_tokens must be > 0")

    merged_dir = merge_adapter(args)
    verify_merged_model(args, merged_dir)

    print(f"Merged model saved to: {merged_dir}")
    print("For production inference speed/memory, load merged model with load_in_8bit=True.")


if __name__ == "__main__":
    main()

	import argparse
	import json
	import os
	import time
	from pathlib import Path
	from typing import Any, Dict, Optional

	import torch
	from PIL import Image
	from unsloth import FastVisionModel

	os.environ["CUDA_VISIBLE_DEVICES"] = "7"


	DEFAULT_ADAPTER_PATH = "outputs/mimic_qwen3vl_lora_8bit_5/checkpoint-17454"
	DEFAULT_OUTPUT_PATH = "outputs/mimic_qwen3vl_lora_8bit_5_merged"
	DEFAULT_INSTRUCTION = "Analyze this chest X-ray image and generate the corresponding radiology report."


	DTYPE_MAP = {
	"float16": torch.float16,
	"bfloat16": torch.bfloat16,
	"float32": torch.float32,
	}


	def parse_args() -> argparse.Namespace:
	parser = argparse.ArgumentParser(
	description=(
	"Merge a Qwen3-VL LoRA adapter into base weights using Unsloth. "
	"This intentionally avoids loading 8-bit modules during merge to preserve accuracy."
	)
	)

	parser.add_argument("--adapter_path", type=str, default=DEFAULT_ADAPTER_PATH)
	parser.add_argument("--output_dir", type=str, default=DEFAULT_OUTPUT_PATH)
	parser.add_argument(
	"--save_method",
	type=str,
	default="merged_16bit",
	choices=["merged_16bit", "merged_4bit", "merged_4bit_forced"],
	help=(
	"Unsloth save mode. For best fidelity use merged_16bit. "
	"Use 4bit variants only if you explicitly need compact merged weights."
	),
	)

	parser.add_argument(
	"--dtype",
	type=str,
	default="auto",
	choices=["auto", "float16", "bfloat16", "float32"],
	help="Compute dtype used while loading model for merge.",
	)
	parser.add_argument("--device_map", type=str, default="auto")
	parser.add_argument("--cuda_visible_devices", type=str, default="")
	parser.add_argument("--max_shard_size", type=str, default="10GB")

	parser.add_argument("--verify_reload", action="store_true", help="Reload merged model after save.")
	parser.add_argument(
	"--verify_load_in_8bit",
	action="store_true",
	help="When verifying, reload merged model in 8-bit inference mode.",
	)
	parser.add_argument(
	"--verify_max_new_tokens",
	type=int,
	default=32,
	help="Max new tokens for optional verification generation.",
	)
	parser.add_argument(
	"--verify_instruction",
	type=str,
	default=DEFAULT_INSTRUCTION,
	help="Prompt text used for optional verification generation.",
	)
	parser.add_argument(
	"--skip_generate_check",
	action="store_true",
	help="If set, verification only checks loadability and skips generation.",
	)

	parser.add_argument(
	"--safe_serialization",
	action="store_true",
	default=True,
	help="Save in safetensors format when supported.",
	)
	parser.add_argument(
	"--no_safe_serialization",
	action="store_true",
	help="Disable safetensors serialization.",
	)

	return parser.parse_args()


	def resolve_dtype(dtype_arg: str) -> torch.dtype:
	if dtype_arg != "auto":
	return DTYPE_MAP[dtype_arg]

	if torch.cuda.is_available():
	if torch.cuda.is_bf16_supported():
	return torch.bfloat16
	return torch.float16
	return torch.float32


	def read_adapter_base_model(adapter_path: Path) -> Optional[str]:
	adapter_config_path = adapter_path / "adapter_config.json"
	if not adapter_config_path.exists():
	return None

	try:
	data = json.loads(adapter_config_path.read_text(encoding="utf-8"))
	except (OSError, ValueError, json.JSONDecodeError):
	return None

	base_model = data.get("base_model_name_or_path")
	return str(base_model) if base_model else None


	def print_runtime_info(args: argparse.Namespace, merge_dtype: torch.dtype, adapter_path: Path, output_dir: Path) -> None:
	print("=" * 88)
	print("Merge configuration")
	print("=" * 88)
	print(f"adapter_path : {adapter_path}")
	print(f"output_dir : {output_dir}")
	print(f"save_method : {args.save_method}")
	print(f"dtype : {merge_dtype}")
	print(f"device_map : {args.device_map}")
	print(f"safe_serialization : {args.safe_serialization and not args.no_safe_serialization}")
	print(f"verify_reload : {args.verify_reload}")
	print(f"verify_load_in_8bit : {args.verify_load_in_8bit}")
	print("=" * 88)


	def save_merge_metadata(output_dir: Path, metadata: Dict[str, Any]) -> None:
	output_dir.mkdir(parents=True, exist_ok=True)
	metadata_path = output_dir / "merge_metadata.json"
	metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")
	print(f"Wrote merge metadata: {metadata_path}")


	def merge_adapter(args: argparse.Namespace) -> Path:
	adapter_path = Path(args.adapter_path).expanduser().resolve()
	output_dir = Path(args.output_dir).expanduser().resolve()

	if not adapter_path.exists() or not adapter_path.is_dir():
	raise FileNotFoundError(f"Adapter path does not exist or is not a directory: {adapter_path}")

	if args.cuda_visible_devices:
	os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda_visible_devices

	merge_dtype = resolve_dtype(args.dtype)
	print_runtime_info(args, merge_dtype, adapter_path, output_dir)

	inferred_base_model = read_adapter_base_model(adapter_path)
	if inferred_base_model:
	print(f"Adapter base model from adapter_config.json: {inferred_base_model}")

	print("Loading adapter (with base) in non-quantized mode for accurate merge...")
	model, tokenizer = FastVisionModel.from_pretrained(
	model_name=str(adapter_path),
	load_in_4bit=False,
	load_in_8bit=False,
	dtype=merge_dtype,
	device_map=args.device_map,
	)

	# Ensure inference graph after load to avoid training-mode side effects.
	FastVisionModel.for_inference(model)

	output_dir.mkdir(parents=True, exist_ok=True)

	safe_serialization = args.safe_serialization and not args.no_safe_serialization
	merge_start = time.time()

	print("Saving merged model...")
	try:
	model.save_pretrained_merged(
	str(output_dir),
	tokenizer,
	save_method=args.save_method,
	safe_serialization=safe_serialization,
	max_shard_size=args.max_shard_size,
	)
	except TypeError:
	# Older Unsloth versions may not support all save kwargs.
	model.save_pretrained_merged(
	str(output_dir),
	tokenizer,
	save_method=args.save_method,
	)

	merge_seconds = time.time() - merge_start
	print(f"Merge complete in {merge_seconds:.1f}s")

	save_merge_metadata(
	output_dir,
	{
	"adapter_path": str(adapter_path),
	"inferred_base_model": inferred_base_model,
	"output_dir": str(output_dir),
	"save_method": args.save_method,
	"merge_dtype": str(merge_dtype),
	"device_map": args.device_map,
	"safe_serialization": safe_serialization,
	"merged_at_unix": int(time.time()),
	},
	)

	del model
	if torch.cuda.is_available():
	torch.cuda.empty_cache()

	return output_dir


	def verify_merged_model(args: argparse.Namespace, merged_dir: Path) -> None:
	if not args.verify_reload:
	return

	print("Reloading merged model for verification...")

	verify_load_in_8bit = bool(args.verify_load_in_8bit)
	if verify_load_in_8bit and not torch.cuda.is_available():
	print("CUDA unavailable. Falling back to non-8bit verification load.")
	verify_load_in_8bit = False

	model, tokenizer = FastVisionModel.from_pretrained(
	model_name=str(merged_dir),
	load_in_4bit=False,
	load_in_8bit=verify_load_in_8bit,
	device_map=args.device_map,
	)
	FastVisionModel.for_inference(model)

	if args.skip_generate_check:
	print("Verification load successful (generation check skipped).")
	return

	# Minimal generation sanity check with one synthetic image.
	test_image = Image.new("RGB", (224, 224), color=(0, 0, 0))
	messages = [
	{
	"role": "user",
	"content": [
	{"type": "image", "image": test_image},
	{"type": "text", "text": args.verify_instruction},
	],
	}
	]

	prompt_text = tokenizer.apply_chat_template(
	messages,
	tokenize=False,
	add_generation_prompt=True,
	)

	inputs = None
	tokenization_errors = []
	for image_argument in (
	{"images": [test_image]},
	{"images": [[test_image]]},
	{"image": [test_image]},
	{"image": [[test_image]]},
	):
	try:
	inputs = tokenizer(
	text=[prompt_text],
	padding=True,
	return_tensors="pt",
	**image_argument,
	)
	break
	except Exception as error: # pragma: no cover
	tokenization_errors.append(str(error))

	if inputs is None:
	raise RuntimeError("Tokenization failed during verify: " + " \| ".join(tokenization_errors))

	model_device = next(model.parameters()).device
	inputs = {k: (v.to(model_device) if isinstance(v, torch.Tensor) else v) for k, v in inputs.items()}

	with torch.inference_mode():
	outputs = model.generate(
	**inputs,
	max_new_tokens=max(1, args.verify_max_new_tokens),
	do_sample=False,
	)

	input_token_count = inputs["input_ids"].shape[-1] if "input_ids" in inputs else 0
	generated_ids = outputs[:, input_token_count:]
	generated_text = tokenizer.batch_decode(
	generated_ids,
	skip_special_tokens=True,
	clean_up_tokenization_spaces=False,
	)[0]

	print("Verification generation successful.")
	print("--- Verification Output (truncated to 500 chars) ---")
	print(generated_text[:500])


	def main() -> None:
	args = parse_args()

	if args.verify_max_new_tokens <= 0:
	raise ValueError("--verify_max_new_tokens must be > 0")

	merged_dir = merge_adapter(args)
	verify_merged_model(args, merged_dir)

	print(f"Merged model saved to: {merged_dir}")
	print("For production inference speed/memory, load merged model with load_in_8bit=True.")


	if __name__ == "__main__":
	main()