diff --git a/comfy/model_management.py b/comfy/model_management.py index b15d08ba1..6f30d5a0b 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1126,7 +1126,7 @@ def text_encoder_offload_device(): def text_encoder_device(): if args.gpu_only: return get_torch_device() - elif vram_state in (VRAMState.HIGH_VRAM, VRAMState.NORMAL_VRAM) or comfy.memory_management.aimdo_enabled: + elif vram_state in (VRAMState.HIGH_VRAM, VRAMState.NORMAL_VRAM, VRAMState.SHARED) or comfy.memory_management.aimdo_enabled: if should_use_fp16(prioritize_performance=False): return get_torch_device() else: diff --git a/comfy/sd.py b/comfy/sd.py index 610c4e2b8..0cdc8e7fd 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -235,6 +235,18 @@ class CLIP: if dtype is None: dtype = model_management.text_encoder_dtype(load_device) + if not model_management.supports_cast(load_device, dtype): + load_device = offload_device + + if load_device != offload_device and not model_management.supports_cast(load_device, torch.float8_e4m3fn): + # Quantized state dicts can contain weights in dtypes the declared + # model dtypes never reflect (e.g. fp8 layers behind comfy_quant + # metadata), which devices like mps cannot cast. + for c in (state_dict if isinstance(state_dict, list) else [state_dict]): + if any(k.endswith("comfy_quant") or (isinstance(w, torch.Tensor) and w.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)) for k, w in c.items()): + load_device = offload_device + break + params['dtype'] = dtype params['device'] = model_options.get("initial_device", model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype))) params['model_options'] = model_options @@ -246,6 +258,7 @@ class CLIP: load_device = offload_device if params['device'] != offload_device: self.cond_stage_model.to(offload_device) + params['device'] = offload_device logging.warning("Had to shift TE back.") model_management.archive_model_dtypes(self.cond_stage_model) diff --git a/tests-unit/comfy_test/text_encoder_mps_test.py b/tests-unit/comfy_test/text_encoder_mps_test.py new file mode 100644 index 000000000..58f7814c1 --- /dev/null +++ b/tests-unit/comfy_test/text_encoder_mps_test.py @@ -0,0 +1,153 @@ +"""Text encoder device placement on Apple Silicon (MPS). + +MPS machines run in VRAMState.SHARED. Text encoders the device supports +(fp16/bf16/fp32) should load on the GPU, while fp8/quantized ones must +stay on the CPU because MPS cannot cast float8 dtypes. +""" + +import pytest +import torch + +import comfy.model_management as mm +import comfy.ops +import comfy.sd + +mps_only = pytest.mark.skipif( + not (torch.backends.mps.is_available() and mm.is_device_mps(mm.get_torch_device())), + reason="requires an Apple Silicon MPS device", +) + +FP8_DTYPES = [torch.float8_e4m3fn, torch.float8_e5m2] + +# Big enough that text_encoder_initial_device() picks the load device. +LARGE_PARAM_COUNT = 2 * 1024 * 1024 * 1024 + + +class DummyTokenizer: + def __init__(self, embedding_directory=None, tokenizer_data={}): + pass + + +class DummyTEModel(torch.nn.Module): + def __init__(self, device=None, dtype=None, model_options={}): + super().__init__() + operations = model_options.get("custom_operations", comfy.ops.manual_cast) + self.linear = operations.Linear(8, 8, dtype=dtype, device=device) + self.dtypes = set([dtype]) + self.construct_device = torch.device(device) if device is not None else None + + def load_sd(self, sd): + return ([], []) + + def forward(self, x): + return self.linear(x) + + +def make_clip(dtype, state_dict=[], clip_class=DummyTEModel): + class Target: + params = {} + clip = clip_class + tokenizer = DummyTokenizer + + return comfy.sd.CLIP( + target=Target(), + parameters=LARGE_PARAM_COUNT, + model_options={"dtype": dtype}, + state_dict=state_dict, + disable_dynamic=True, + ) + + +@pytest.fixture(autouse=True) +def unload_models(): + yield + mm.unload_all_models() + + +@mps_only +class TestTextEncoderDeviceMPS: + def test_vram_state_is_shared(self): + assert mm.vram_state == mm.VRAMState.SHARED + + def test_text_encoder_device_is_mps(self): + assert mm.is_device_mps(mm.text_encoder_device()) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_castable_dtype_clip_loads_on_mps(self, dtype): + clip = make_clip(dtype) + assert mm.is_device_mps(clip.patcher.load_device) + weight = clip.cond_stage_model.linear.weight + assert weight.device.type == "mps" + x = torch.ones((1, 8), dtype=dtype, device=weight.device) + assert clip.cond_stage_model(x).device.type == "mps" + + @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) + def test_fp8_clip_falls_back_to_cpu(self, fp8_dtype): + clip = make_clip(fp8_dtype) + assert clip.patcher.load_device == clip.patcher.offload_device + assert clip.cond_stage_model.construct_device.type == "cpu" + weight = clip.cond_stage_model.linear.weight + assert weight.device.type == "cpu" + x = torch.ones((1, 8), dtype=torch.float16, device=weight.device) + assert clip.cond_stage_model(x).device.type == "cpu" + + @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) + def test_fp8_state_dict_falls_back_to_cpu(self, fp8_dtype): + # fp8 weights in the state dict aren't reflected in the declared + # model dtypes, e.g. quantized checkpoints with fp16 norm layers. + sd = {"linear.weight": torch.zeros((8, 8), dtype=fp8_dtype)} + clip = make_clip(torch.float16, state_dict=[sd]) + assert not mm.is_device_mps(clip.patcher.load_device) + assert clip.cond_stage_model.construct_device.type == "cpu" + + def test_comfy_quant_state_dict_falls_back_to_cpu(self): + sd = { + "linear.weight": torch.zeros((8, 8), dtype=torch.uint8), + "linear.comfy_quant": torch.zeros(16, dtype=torch.uint8), + "spiece_model": b"not a tensor", + } + clip = make_clip(torch.float16, state_dict=[sd]) + assert not mm.is_device_mps(clip.patcher.load_device) + assert clip.cond_stage_model.construct_device.type == "cpu" + + def test_fp8_full_model_state_dict_falls_back_to_cpu(self): + # Full checkpoints pass a single dict instead of a list. + sd = {"linear.weight": torch.zeros((8, 8), dtype=torch.float8_e4m3fn)} + clip = make_clip(torch.float16, state_dict=sd) + assert not mm.is_device_mps(clip.patcher.load_device) + assert clip.cond_stage_model.construct_device.type == "cpu" + + @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) + def test_mixed_declared_dtypes_fall_back_to_cpu(self, fp8_dtype): + # A secondary declared dtype (e.g. dtype_llama) can be fp8 while the + # primary dtype is fp16. + class MixedDtypeTEModel(DummyTEModel): + def __init__(self, device=None, dtype=None, model_options={}): + super().__init__(device=device, dtype=dtype, model_options=model_options) + self.dtypes = set([dtype, fp8_dtype]) + + clip = make_clip(torch.float16, clip_class=MixedDtypeTEModel) + assert not mm.is_device_mps(clip.patcher.load_device) + assert clip.cond_stage_model.linear.weight.device.type == "cpu" + + @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) + def test_fp8_cast_still_unsupported_on_mps(self, fp8_dtype): + # If a torch release adds fp8 casts on MPS, supports_cast() can be + # updated to let fp8 text encoders onto the GPU (pytorch#132624). + assert not mm.supports_cast(mm.get_torch_device(), fp8_dtype) + t = torch.zeros(4, dtype=fp8_dtype, device="mps") + with pytest.raises((RuntimeError, TypeError)): + t.to(torch.float16) + + +@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES) +def test_fp8_capable_devices_skip_the_quant_fallback(fp8_dtype): + # The state dict scan in CLIP.__init__ is gated on this, so devices + # that can cast fp8 keep loading quantized text encoders on the GPU. + assert mm.supports_cast(torch.device("cuda"), fp8_dtype) + + +@pytest.mark.parametrize("state", [mm.VRAMState.LOW_VRAM, mm.VRAMState.NO_VRAM]) +def test_low_vram_states_keep_text_encoders_on_cpu(monkeypatch, state): + monkeypatch.setattr(mm, "vram_state", state) + assert mm.text_encoder_device() == torch.device("cpu")