199 lines
9.1 KiB
Python
199 lines
9.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Probe InternVLA's real Cosmos image encoder on CPU and one DirectML GPU.
|
|
|
|
This loads the unchanged Omni model source as an isolated experiment because
|
|
the installed DirectML PyTorch environment is separate from the vLLM runtime.
|
|
It tests one model component, not the full action policy or robot-task quality.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import sys
|
|
import time
|
|
import traceback
|
|
import types
|
|
from pathlib import Path
|
|
|
|
|
|
def _rank(values: list[float], fraction: float) -> float:
|
|
return sorted(values)[math.ceil(len(values) * fraction) - 1]
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for block in iter(lambda: source.read(8 << 20), b""):
|
|
digest.update(block)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def _load_omni_cosmos_module(source_dir: Path):
|
|
# Keep the model implementation unchanged without importing the unrelated
|
|
# vLLM engine in a separate, older Torch/DirectML experiment environment.
|
|
package_name = "internvla_cosmos_probe"
|
|
package = types.ModuleType(package_name)
|
|
package.__path__ = [str(source_dir)]
|
|
sys.modules[package_name] = package
|
|
name = f"{package_name}.model_cosmos"
|
|
spec = importlib.util.spec_from_file_location(name, source_dir / "model_cosmos.py")
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("could not load unchanged InternVLA Cosmos source")
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--encoder", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--dml-index", type=int, default=1)
|
|
parser.add_argument("--batch-size", type=int, default=1)
|
|
parser.add_argument("--model-dtype", choices=("float32", "bfloat16"), default="float32")
|
|
parser.add_argument("--warmups", type=int, default=1)
|
|
parser.add_argument("--repeats", type=int, default=20)
|
|
args = parser.parse_args()
|
|
if args.warmups < 0 or args.repeats < 1 or args.batch_size < 1:
|
|
parser.error("invalid batch size, warmup or measured count")
|
|
|
|
import numpy as np
|
|
import psutil
|
|
import torch
|
|
import torch_directml
|
|
from safetensors import safe_open
|
|
|
|
source_dir = Path(__file__).resolve().parents[3] / "vllm_omni" / "diffusion" / "models" / "internvla_a1"
|
|
source = _load_omni_cosmos_module(source_dir)
|
|
device_name = torch_directml.device_name(args.dml_index).rstrip("\x00")
|
|
if device_name != "AMD Radeon(TM) 890M Graphics":
|
|
raise RuntimeError(f"refusing unverified DirectML device: {device_name!r}")
|
|
dml = torch_directml.device(args.dml_index)
|
|
encoder = args.encoder.resolve(strict=True)
|
|
report = {
|
|
"scope": "InternVLA real Cosmos CI8x8 image encoder component; synthetic patterned 256x256 input; no full policy or robot-task claim",
|
|
"source_files_sha256": {
|
|
name: _sha256(source_dir / name)
|
|
for name in ("model_cosmos.py", "cosmos_ci_torch.py")
|
|
},
|
|
"encoder_sha256": _sha256(encoder),
|
|
"device_name": device_name,
|
|
"torch": torch.__version__,
|
|
"torch_directml": __import__("importlib.metadata", fromlist=["version"]).version("torch-directml"),
|
|
"input_shape": [args.batch_size, 3, 256, 256],
|
|
"input_dtype": args.model_dtype,
|
|
"warmups": args.warmups,
|
|
"repeats": args.repeats,
|
|
}
|
|
try:
|
|
torch.manual_seed(42)
|
|
image = torch.zeros((1, 3, 256, 256), dtype=torch.float32)
|
|
image[:, 0, 48:208, 48:208] = 1
|
|
image[:, 1, 48:208, 48:208] = -1
|
|
image[:, 2, :, :] = torch.linspace(-1, 1, 256)[None, None, :]
|
|
image = image.repeat(args.batch_size, 1, 1, 1)
|
|
model_dtype = getattr(torch, args.model_dtype)
|
|
image = image.to(model_dtype)
|
|
|
|
started = time.perf_counter()
|
|
cpu_model = source.load_cosmos_component(encoder, component="encoder", device="cpu").to(model_dtype)
|
|
report["cpu_load_s"] = time.perf_counter() - started
|
|
with safe_open(str(encoder), framework="pt", device="cpu") as checkpoint:
|
|
checkpoint_keys = set(checkpoint.keys())
|
|
checkpoint_buffers = {
|
|
key: checkpoint.get_tensor(key)
|
|
for key in checkpoint_keys
|
|
if key in {"encoder.patcher._arange", "encoder.patcher.wavelets"}
|
|
}
|
|
model_keys = set(cpu_model.state_dict())
|
|
report["checkpoint_tensor_count"] = len(checkpoint_keys)
|
|
report["encoder_state_tensor_count"] = len(model_keys)
|
|
report["checkpoint_missing_keys"] = sorted(model_keys - checkpoint_keys)
|
|
report["checkpoint_unexpected_keys"] = sorted(checkpoint_keys - model_keys)
|
|
model_buffers = dict(cpu_model.named_buffers())
|
|
report["nonpersistent_buffer_comparison"] = {
|
|
key: {
|
|
"checkpoint_shape": list(value.shape),
|
|
"model_shape": list(model_buffers[key].shape),
|
|
"checkpoint_dtype": str(value.dtype),
|
|
"model_dtype": str(model_buffers[key].dtype),
|
|
"max_abs_error": float((value.float() - model_buffers[key].float()).abs().max()),
|
|
}
|
|
for key, value in checkpoint_buffers.items()
|
|
}
|
|
if (report["checkpoint_missing_keys"] or
|
|
set(report["checkpoint_unexpected_keys"]) != set(checkpoint_buffers)):
|
|
raise RuntimeError("Cosmos encoder checkpoint/state keys do not match")
|
|
started = time.perf_counter()
|
|
dml_model = source.load_cosmos_component(encoder, component="encoder", device="cpu").to(dtype=model_dtype, device=dml)
|
|
report["dml_load_and_transfer_s"] = time.perf_counter() - started
|
|
if {parameter.device.type for parameter in cpu_model.parameters()} != {"cpu"}:
|
|
raise RuntimeError("CPU reference encoder placement differs")
|
|
if {parameter.device.type for parameter in dml_model.parameters()} != {"privateuseone"}:
|
|
raise RuntimeError("DirectML encoder weights were not on the declared GPU")
|
|
|
|
def execute(model, input_tensor, expected_device_type):
|
|
with torch.no_grad():
|
|
output = model(input_tensor)
|
|
if not isinstance(output, torch.Tensor):
|
|
output = output[0]
|
|
if output.device.type != expected_device_type:
|
|
raise RuntimeError(f"encoder output on {output.device}, expected {expected_device_type}")
|
|
return output.cpu().float().contiguous()
|
|
|
|
cpu_times: list[float] = []
|
|
dml_times: list[float] = []
|
|
cpu_output = dml_output = None
|
|
for index in range(args.warmups + args.repeats):
|
|
started = time.perf_counter()
|
|
cpu_output = execute(cpu_model, image, "cpu")
|
|
cpu_wall = time.perf_counter() - started
|
|
started = time.perf_counter()
|
|
dml_output = execute(dml_model, image.to(dml), "privateuseone")
|
|
dml_wall = time.perf_counter() - started
|
|
if index >= args.warmups:
|
|
cpu_times.append(cpu_wall)
|
|
dml_times.append(dml_wall)
|
|
assert cpu_output is not None and dml_output is not None
|
|
if cpu_output.shape != dml_output.shape or not torch.isfinite(dml_output).all():
|
|
raise RuntimeError("DirectML encoder returned wrong shape or non-finite output")
|
|
reference = cpu_output.numpy()
|
|
actual = dml_output.numpy()
|
|
difference = reference - actual
|
|
report.update({
|
|
"status": "passed",
|
|
"output_shape": list(reference.shape),
|
|
"cpu_output_sha256_float32": hashlib.sha256(reference.tobytes()).hexdigest(),
|
|
"dml_output_sha256_float32": hashlib.sha256(actual.tobytes()).hexdigest(),
|
|
"max_abs_error": float(np.max(np.abs(difference))),
|
|
"relative_l2_error": float(np.linalg.norm(difference) / np.linalg.norm(reference)),
|
|
"cosine_similarity": float(np.dot(reference.ravel(), actual.ravel()) / (
|
|
np.linalg.norm(reference) * np.linalg.norm(actual)
|
|
)),
|
|
"cpu_nearest_rank_p50_wall_s": _rank(cpu_times, 0.5),
|
|
"cpu_nearest_rank_p95_wall_s": _rank(cpu_times, 0.95),
|
|
"dml_nearest_rank_p50_wall_s": _rank(dml_times, 0.5),
|
|
"dml_nearest_rank_p95_wall_s": _rank(dml_times, 0.95),
|
|
"cpu_times_s": cpu_times,
|
|
"dml_times_s": dml_times,
|
|
"process_rss_bytes_after": psutil.Process().memory_info().rss,
|
|
})
|
|
except BaseException as exc:
|
|
report["status"] = "failed"
|
|
report["error"] = f"{type(exc).__name__}: {exc}"
|
|
report["traceback"] = traceback.format_exc()
|
|
raise
|
|
finally:
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps({key: value for key, value in report.items() if not key.endswith("times_s")}, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|