mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
b20113525f
The dynamic VBAR cast path applied want_requant to both weights and biases. For an FP8 checkpoint with a low-VRAM LoRA patch, this returned an FP8 bias to the optimized linear operation even though the activation dtype was FP16. Restrict the optional requantized return value to weights while preserving resident-storage updates for both parameters. Make the legacy FP8 linear path release its VBAR/offload resources in a finally block so backend failures do not leave pages pinned before the normal fallback runs. Add regressions for the weight/bias dtype contract and exception cleanup. Fixes #14952 Signed-off-by: liminfei <liminfei@amd.com>
75 lines
2.6 KiB
Python
75 lines
2.6 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest import mock
|
|
|
|
import torch
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
|
|
from comfy.cli_args import args
|
|
|
|
if not torch.cuda.is_available():
|
|
args.cpu = True
|
|
|
|
from comfy import ops
|
|
|
|
|
|
class TestFP8Ops(unittest.TestCase):
|
|
def test_vbar_requantizes_weight_but_not_bias(self):
|
|
weight = torch.zeros((16, 16), dtype=torch.float8_e4m3fn)
|
|
bias = torch.zeros(16, dtype=torch.float8_e4m3fn)
|
|
patch = lambda tensor: tensor
|
|
layer = SimpleNamespace(
|
|
weight=weight,
|
|
bias=bias,
|
|
weight_function=[],
|
|
bias_function=[],
|
|
weight_lowvram_function=patch,
|
|
bias_lowvram_function=patch,
|
|
seed_key="layer",
|
|
_prefetch={
|
|
"resident": False,
|
|
"xfer_dest": torch.empty(1, dtype=torch.uint8),
|
|
"needs_cast": False,
|
|
"cast_geometry": None,
|
|
"signature": (1,),
|
|
},
|
|
)
|
|
|
|
with (
|
|
mock.patch.object(ops.comfy.memory_management, "interpret_gathered_like", return_value=[weight, bias]),
|
|
mock.patch.object(ops.comfy.float, "stochastic_rounding", side_effect=lambda tensor, dtype, seed: tensor.to(dtype)) as rounding,
|
|
):
|
|
cast_weight, cast_bias = ops.resolve_cast_module_with_vbar(
|
|
layer,
|
|
torch.float8_e4m3fn,
|
|
torch.device("cpu"),
|
|
torch.float16,
|
|
torch.float16,
|
|
True,
|
|
)
|
|
|
|
self.assertEqual(cast_weight.dtype, torch.float8_e4m3fn)
|
|
self.assertEqual(cast_bias.dtype, torch.float16)
|
|
self.assertEqual(rounding.call_count, 2)
|
|
|
|
def test_fp8_linear_unpins_vbar_when_linear_fails(self):
|
|
weight = torch.zeros((16, 16), dtype=torch.float8_e4m3fn)
|
|
bias = torch.zeros(16, dtype=torch.float8_e4m3fn)
|
|
vbar = object()
|
|
layer = SimpleNamespace(weight=weight, _v=vbar)
|
|
input_tensor = torch.zeros((2, 16), dtype=torch.float16)
|
|
offload_info = (None, torch.device("cpu"), None)
|
|
|
|
with (
|
|
mock.patch.object(ops, "cast_bias_weight", return_value=(weight, bias, offload_info)),
|
|
mock.patch.object(ops.comfy_aimdo.model_vbar, "vbar_unpin") as unpin,
|
|
mock.patch.object(torch.nn.functional, "linear", side_effect=RuntimeError("FP8 linear failed")),
|
|
):
|
|
with self.assertRaisesRegex(RuntimeError, "FP8 linear failed"):
|
|
ops.fp8_linear(layer, input_tensor)
|
|
|
|
unpin.assert_called_once_with(vbar)
|