63 lines
2.7 KiB
PowerShell
63 lines
2.7 KiB
PowerShell
$ErrorActionPreference = 'Stop'
|
|
|
|
$envPath = Join-Path $PSScriptRoot '..\..\..\.env'
|
|
Get-Content -LiteralPath $envPath | ForEach-Object {
|
|
if ($_ -match '^\s*([^#=\s]+)\s*=\s*(.*)\s*$') {
|
|
[Environment]::SetEnvironmentVariable($matches[1], $matches[2].Trim('"').Trim("'"), 'Process')
|
|
}
|
|
}
|
|
if (-not $env:RUNPOD_API_KEY) { throw 'RUNPOD_API_KEY is not set in .env.' }
|
|
|
|
$probe = @'
|
|
import json
|
|
import torch
|
|
from safetensors import safe_open
|
|
import comfy_kitchen as ck
|
|
from comfy_kitchen.tensor import TensorCoreNVFP4Layout
|
|
|
|
path = "/runpod-volume/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
|
|
prefix = "blocks.0.attn.qkv_proj"
|
|
with safe_open(path, framework="pt", device="cuda") as checkpoint:
|
|
sidecar = checkpoint.get_tensor(prefix + ".comfy_quant")
|
|
assert json.loads(bytes(sidecar.cpu().tolist())) == {"format": "nvfp4"}
|
|
weight = checkpoint.get_tensor(prefix + ".weight").contiguous()
|
|
weight_scale = checkpoint.get_tensor(prefix + ".weight_scale").view(torch.float8_e4m3fn).contiguous()
|
|
weight_scale_2 = checkpoint.get_tensor(prefix + ".weight_scale_2").to(torch.float32).contiguous()
|
|
|
|
x = torch.randn((1, weight.shape[1] * 2), device="cuda", dtype=torch.bfloat16)
|
|
packed_x, x_params = TensorCoreNVFP4Layout.quantize(x)
|
|
y = ck.scaled_mm_nvfp4(
|
|
packed_x, weight,
|
|
tensor_scale_a=x_params.scale,
|
|
tensor_scale_b=weight_scale_2,
|
|
block_scale_a=x_params.block_scale,
|
|
block_scale_b=weight_scale,
|
|
out_dtype=torch.bfloat16,
|
|
)
|
|
torch.cuda.synchronize()
|
|
print(json.dumps({
|
|
"input_shape": list(x.shape),
|
|
"packed_weight_shape": list(weight.shape),
|
|
"weight_scale_shape": list(weight_scale.shape),
|
|
"weight_scale_2_shape": list(weight_scale_2.shape),
|
|
"output_shape": list(y.shape),
|
|
"output_dtype": str(y.dtype),
|
|
"finite": bool(torch.isfinite(y).all()),
|
|
"gpu": torch.cuda.get_device_name(),
|
|
}, sort_keys=True))
|
|
'@
|
|
$encodedProbe = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($probe))
|
|
$setup = 'python -m pip install --no-cache-dir comfy-kitchen==0.2.28 "safetensors>=0.5.0"; python -c "import base64;exec(base64.b64decode(\"' + $encodedProbe + '\"))"'
|
|
$payload = @{
|
|
name = 'h3-nvfp4-linear-probe'
|
|
image = 'runpod/pytorch:1.1.0-rc.154-cu1300-torch291-ubuntu2404'
|
|
cloud = 'SECURE'
|
|
gpu = @{ id = 'NVIDIA RTX PRO 6000 Blackwell Server Edition'; count = 1 }
|
|
disk = 20
|
|
dataCenterIds = @('EUR-IS-1')
|
|
mounts = @{ network = @(@{ volumeId = 'g8r2uqufz5'; path = '/runpod-volume' }) }
|
|
args = 'bash -lc ''set -eu; ' + $setup + ''''
|
|
}
|
|
|
|
$headers = @{ Authorization = "Bearer $env:RUNPOD_API_KEY" }
|
|
Invoke-RestMethod -Method Post -Uri 'https://api.runpod.io/v2/pods' -Headers $headers -ContentType 'application/json' -Body ($payload | ConvertTo-Json -Depth 8)
|