diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 140bd5dc5..63fea9199 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -77,10 +77,11 @@ steps: - "/fsx/hf_cache:/fsx/hf_cache" - label: "Diffusion Parallelism Test" - timeout_in_minutes: 15 + timeout_in_minutes: 20 depends_on: image-build commands: - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py + - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: diff --git a/docs/user_guide/acceleration/parallelism_acceleration.md b/docs/user_guide/acceleration/parallelism_acceleration.md index 0c5d1a62b..ac4c8c5fa 100644 --- a/docs/user_guide/acceleration/parallelism_acceleration.md +++ b/docs/user_guide/acceleration/parallelism_acceleration.md @@ -12,27 +12,56 @@ The following parallelism methods are currently supported in vLLM-Omni: 3. Classifier-Free-Guidance Parallel (CFG-Parallel): CFG-Parallel runs the positive/negative prompts of classifier-free guidance (CFG) on different devices, then merges on a single device to perform the scheduler step. +4. [Tensor Parallelism](#tensor-parallelism): Tensor parallelism shards model weights across devices. This can reduce per-GPU memory usage. Note that for diffusion models we currently shard the majority of layers within the DiT. + The following table shows which models are currently supported by parallelism method: ### ImageGen -| Model | Model Identifier | Ulysses-SP | Ring-SP |CFG-Parallel | -|-------|------------------|-----------|-------------|-------------| -| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ❌ | ❌ | -| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ❌ | ❌ | -| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | -| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | -| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | -| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | -| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | -| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | +| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | +|-------|------------------|------------|---------|--------------|--------------------------| +| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ❌ | ❌ | ❌ | +| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ❌ | ❌ | ❌ | +| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | +| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ | +| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ | +| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) | +| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ | + +!!! note "Why Z-Image is TP=2 only" + Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints. + For example, the model has `n_heads=30` and a final projection out dimension of `64`, so valid TP sizes must divide both 30 and 64; the only common divisors are **1 and 2**. ### VideoGen -| Model | Model Identifier | Ulysses-SP | Ring-SP | -|-------|------------------|-----------|---------| -| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | +| Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel | +|-------|------------------|------------|---------|--------------------------| +| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ | + +### Tensor Parallelism + +Tensor parallelism splits model parameters across GPUs. In vLLM-Omni, tensor parallelism is configured via `DiffusionParallelConfig.tensor_parallel_size`. + +#### Offline Inference + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + parallel_config=DiffusionParallelConfig(tensor_parallel_size=2), +) + +outputs = omni.generate( + prompt="a cat reading a book", + num_inference_steps=9, + width=512, + height=512, +) +``` ### Sequence Parallelism diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index f7398a6e0..9f06a7105 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -88,6 +88,17 @@ def parse_args() -> argparse.Namespace: choices=[1, 2], help="Number of GPUs used for classifier free guidance parallel size.", ) + parser.add_argument( + "--enforce_eager", + action="store_true", + help="Disable torch.compile and force eager execution.", + ) + parser.add_argument( + "--tensor_parallel_size", + type=int, + default=1, + help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", + ) return parser.parse_args() @@ -131,7 +142,10 @@ def main(): # assert args.ring_degree == 1, "Ring attention is not supported yet" parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size + ulysses_degree=args.ulysses_degree, + ring_degree=args.ring_degree, + cfg_parallel_size=args.cfg_parallel_size, + tensor_parallel_size=args.tensor_parallel_size, ) omni = Omni( @@ -141,6 +155,7 @@ def main(): cache_backend=args.cache_backend, cache_config=cache_config, parallel_config=parallel_config, + enforce_eager=args.enforce_eager, ) # Time profiling for generation @@ -150,7 +165,8 @@ def main(): print(f" Inference steps: {args.num_inference_steps}") print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") print( - f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}" + f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, " + f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}" ) print(f" Image size: {args.width}x{args.height}") print(f"{'=' * 60}\n") diff --git a/tests/diffusion/models/z_image/test_zimage_tp_constraints.py b/tests/diffusion/models/z_image/test_zimage_tp_constraints.py new file mode 100644 index 000000000..a276274a0 --- /dev/null +++ b/tests/diffusion/models/z_image/test_zimage_tp_constraints.py @@ -0,0 +1,44 @@ +import pytest + +from vllm_omni.diffusion.models.z_image.z_image_transformer import validate_zimage_tp_constraints + + +def test_validate_zimage_tp_constraints_tp2_ok(): + ffn_hidden_dim, final_out_dims, supported_tp = validate_zimage_tp_constraints( + dim=3840, + n_heads=30, + n_kv_heads=30, + in_channels=16, + all_patch_size=(2,), + all_f_patch_size=(1,), + tensor_parallel_size=2, + ) + assert ffn_hidden_dim == 10240 + assert final_out_dims == [64] + assert supported_tp == [1, 2] + + +def test_validate_zimage_tp_constraints_tp4_fails_on_heads(): + with pytest.raises(ValueError, match=r"n_heads % tensor_parallel_size"): + validate_zimage_tp_constraints( + dim=3840, + n_heads=30, + n_kv_heads=30, + in_channels=16, + all_patch_size=(2,), + all_f_patch_size=(1,), + tensor_parallel_size=4, + ) + + +def test_validate_zimage_tp_constraints_tp3_fails_on_ffn_hidden_dim(): + with pytest.raises(ValueError, match=r"ffn_hidden_dim % tensor_parallel_size"): + validate_zimage_tp_constraints( + dim=3840, + n_heads=30, + n_kv_heads=30, + in_channels=16, + all_patch_size=(2,), + all_f_patch_size=(1,), + tensor_parallel_size=3, + ) diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py new file mode 100644 index 000000000..d32bb2b82 --- /dev/null +++ b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import sys +import time +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from vllm.distributed.parallel_state import cleanup_dist_env_and_memory + +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.utils.platform_utils import is_npu, is_rocm + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +PROMPT = "a photo of a cat sitting on a laptop keyboard" + + +def _get_zimage_model() -> str: + # Allow overriding the model for local/offline environments. + # Can be either a HuggingFace repo id or a local path. + return os.environ.get("VLLM_TEST_ZIMAGE_MODEL", "Tongyi-MAI/Z-Image-Turbo") + + +def _pil_to_float_rgb_tensor(img: Image.Image) -> torch.Tensor: + """Convert PIL image to float32 RGB tensor in [0, 1] with shape [H, W, 3].""" + arr = np.asarray(img.convert("RGB"), dtype=np.float32) / 255.0 + return torch.from_numpy(arr) + + +def _diff_metrics(a: Image.Image, b: Image.Image) -> tuple[float, float]: + """Return (mean_abs_diff, max_abs_diff) over RGB pixels in [0, 1].""" + ta = _pil_to_float_rgb_tensor(a) + tb = _pil_to_float_rgb_tensor(b) + assert ta.shape == tb.shape, f"Image shapes differ: {ta.shape} vs {tb.shape}" + abs_diff = torch.abs(ta - tb) + return abs_diff.mean().item(), abs_diff.max().item() + + +def _extract_single_image(outputs) -> Image.Image: + first_output = outputs[0] + assert first_output.final_output_type == "image" + if not hasattr(first_output, "request_output") or not first_output.request_output: + raise ValueError("No request_output found in OmniRequestOutput") + + req_out = first_output.request_output[0] + if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): + raise ValueError("Invalid request_output structure or missing 'images' key") + + images = req_out.images + if images is None or len(images) != 1: + raise ValueError(f"Expected 1 image, got {0 if images is None else len(images)}") + return images[0] + + +def _run_zimage_generate( + *, tp_size: int, height: int, width: int, num_inference_steps: int, seed: int +) -> tuple[Image.Image, float]: + m = Omni( + model=_get_zimage_model(), + parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp_size), + ) + try: + # NOTE: Omni closes itself when a generate() call is exhausted. + # To avoid measuring teardown time (process shutdown, memory cleanup), + # we measure the latency to produce *subsequent* outputs within a single + # generator run. + # + # This also serves as a warmup: the first output may include extra + # compilation/caching overhead, while later outputs are closer to + # steady-state inference. + num_requests = 4 # 1 warmup + 3 timed + gen = m.generate( + [PROMPT] * num_requests, + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=0.0, + seed=seed, + num_outputs_per_prompt=1, + py_generator=True, + ) + + warmup_output = next(gen) + t_prev = time.perf_counter() + per_request_times_s: list[float] = [] + last_output = warmup_output + for _ in range(num_requests - 1): + last_output = next(gen) + t_now = time.perf_counter() + per_request_times_s.append(t_now - t_prev) + t_prev = t_now + + # Ensure the generator is fully consumed so it can clean up. + for _ in gen: + pass + + median_time_s = float(np.median(per_request_times_s)) + return _extract_single_image([last_output]), median_time_s + finally: + m.close() + cleanup_dist_env_and_memory() + + +@pytest.mark.integration +def test_zimage_tensor_parallel_tp2(tmp_path: Path): + if is_npu() or is_rocm(): + pytest.skip("Z-Image TP e2e test is only supported on CUDA for now.") + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.") + + height = 256 + width = 256 + num_inference_steps = 2 + seed = 42 + + tp1_img, tp1_time_s = _run_zimage_generate( + tp_size=1, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + ) + tp2_img, tp2_time_s = _run_zimage_generate( + tp_size=2, + height=height, + width=width, + num_inference_steps=num_inference_steps, + seed=seed, + ) + + tp1_path = tmp_path / "zimage_tp1.png" + tp2_path = tmp_path / "zimage_tp2.png" + tp1_img.save(tp1_path) + tp2_img.save(tp2_path) + + assert tp1_img.width == width and tp1_img.height == height + assert tp2_img.width == width and tp2_img.height == height + + mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img) + mean_threshold = 3e-2 + max_threshold = 3.5e-1 + print( + "Z-Image TP image diff stats (TP=1 vs TP=2): " + f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " + f"thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e}; " + f"tp1_img={tp1_path}, tp2_img={tp2_path}" + ) + assert mean_abs_diff <= mean_threshold and max_abs_diff <= max_threshold, ( + f"Image diff exceeded threshold: mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e} " + f"(thresholds: mean<={mean_threshold:.6e}, max<={max_threshold:.6e})" + ) + + print(f"Z-Image TP perf (lower is better): tp1_time_s={tp1_time_s:.6f}, tp2_time_s={tp2_time_s:.6f}") + assert tp2_time_s < tp1_time_s, f"Expected TP=2 to be faster than TP=1 (tp1={tp1_time_s}, tp2={tp2_time_s})" diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 585090a53..0d7e68f70 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -21,6 +21,7 @@ from collections.abc import Iterable import torch import torch.nn as nn from torch.nn.utils.rnn import pad_sequence +from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( @@ -31,11 +32,107 @@ from vllm.model_executor.layers.linear import ( from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available from vllm_omni.diffusion.layers.rope import RotaryEmbedding ADALN_EMBED_DIM = 256 SEQ_MULTI_OF = 32 +logger = init_logger(__name__) + + +def _positive_divisors(n: int) -> set[int]: + if n <= 0: + return set() + divs: set[int] = set() + for d in range(1, int(math.isqrt(n)) + 1): + if n % d == 0: + divs.add(d) + divs.add(n // d) + return divs + + +def _get_tensor_parallel_size_from_context() -> int: + if not is_forward_context_available(): + return 1 + try: + od_config = get_forward_context().omni_diffusion_config + if od_config is None: + return 1 + return int(od_config.parallel_config.tensor_parallel_size) + except Exception: + return 1 + + +def validate_zimage_tp_constraints( + *, + dim: int, + n_heads: int, + n_kv_heads: int, + in_channels: int, + all_patch_size: tuple[int, ...], + all_f_patch_size: tuple[int, ...], + tensor_parallel_size: int, +) -> tuple[int, list[int], list[int]]: + """Validate Z-Image TP constraints without requiring a distributed context. + + Returns: + (ffn_hidden_dim, final_out_dims, supported_tp_candidates) + """ + tp_size = int(tensor_parallel_size) + if tp_size <= 0: + raise ValueError(f"tensor_parallel_size must be > 0, got {tp_size}") + if dim % n_heads != 0: + raise ValueError(f"dim must be divisible by n_heads, got dim={dim}, n_heads={n_heads}") + if dim % tp_size != 0: + supported = sorted(_positive_divisors(dim)) + raise ValueError( + f"Z-Image requires dim % tensor_parallel_size == 0, but got dim={dim}, tp={tp_size}. " + f"Supported tp candidates by dim: {supported}" + ) + if n_heads % tp_size != 0: + supported = sorted(_positive_divisors(n_heads)) + raise ValueError( + f"Z-Image requires n_heads % tensor_parallel_size == 0, but got n_heads={n_heads}, tp={tp_size}. " + f"Supported tp candidates by n_heads: {supported}" + ) + if n_kv_heads % tp_size != 0: + supported = sorted(_positive_divisors(n_kv_heads)) + raise ValueError( + f"Z-Image requires n_kv_heads % tensor_parallel_size == 0, but got n_kv_heads={n_kv_heads}, " + f"tp={tp_size}. Supported tp candidates by n_kv_heads: {supported}" + ) + + ffn_hidden_dim = int(dim / 3 * 8) + if ffn_hidden_dim % tp_size != 0: + supported = sorted(_positive_divisors(ffn_hidden_dim)) + raise ValueError( + "Z-Image requires ffn_hidden_dim % tensor_parallel_size == 0 (for TP-sharded MLP), but got " + f"ffn_hidden_dim={ffn_hidden_dim}, tp={tp_size}. Supported tp candidates by ffn_hidden_dim: {supported}" + ) + + final_out_dims = [ + int(patch_size) * int(patch_size) * int(f_patch_size) * int(in_channels) + for patch_size, f_patch_size in zip(all_patch_size, all_f_patch_size) + ] + bad_final_out_dims = [d for d in final_out_dims if d % tp_size != 0] + if bad_final_out_dims: + supported = sorted(_positive_divisors(math.gcd(*final_out_dims))) + raise ValueError( + "Z-Image requires final projection out_features divisible by tensor_parallel_size, but got " + f"final_out_dims={final_out_dims}, tp={tp_size}. " + f"Supported tp candidates by final_out_dims gcd: {supported}" + ) + + supported_tp_candidates = sorted( + _positive_divisors(n_heads) + & _positive_divisors(n_kv_heads) + & _positive_divisors(dim) + & _positive_divisors(ffn_hidden_dim) + & _positive_divisors(math.gcd(*final_out_dims)) + ) + return ffn_hidden_dim, final_out_dims, supported_tp_candidates + class TimestepEmbedder(nn.Module): def __init__(self, out_size, mid_size=None, frequency_embedding_size=256): @@ -90,8 +187,8 @@ class ZImageAttention(nn.Module): ) -> None: super().__init__() self.dim = dim - self.num_heads = num_heads - self.num_kv_heads = num_kv_heads + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads self.head_dim = dim // num_heads self.qk_norm = qk_norm @@ -99,7 +196,7 @@ class ZImageAttention(nn.Module): hidden_size=dim, head_size=self.head_dim, total_num_heads=num_heads, - disable_tp=True, + total_num_kv_heads=num_kv_heads, bias=False, ) @@ -107,13 +204,27 @@ class ZImageAttention(nn.Module): self.norm_q = RMSNorm(self.head_dim, eps=eps) self.norm_k = RMSNorm(self.head_dim, eps=eps) - self.to_out = nn.ModuleList([nn.Linear(dim, dim, bias=False)]) + # NOTE: QKV is column-parallel on heads, so attention output is sharded + # on the last dim (dim / tp). Use row-parallel output projection to + # all-reduce back to full dim. + self.to_out = nn.ModuleList( + [ + RowParallelLinear( + dim, + dim, + bias=False, + input_is_parallel=True, + return_bias=False, + ) + ] + ) self.attn = Attention( - num_heads=num_heads, + num_heads=self.to_qkv.num_heads, head_size=self.head_dim, softmax_scale=1.0 / (self.head_dim**0.5), causal=False, + num_kv_heads=self.to_qkv.num_kv_heads, ) self.rope = RotaryEmbedding(is_neox_style=False) @@ -125,11 +236,13 @@ class ZImageAttention(nn.Module): sin: torch.Tensor, ): qkv, _ = self.to_qkv(hidden_states) - query, key, value = qkv.chunk(3, dim=-1) + q_size = self.to_qkv.num_heads * self.head_dim + kv_size = self.to_qkv.num_kv_heads * self.head_dim + query, key, value = qkv.split([q_size, kv_size, kv_size], dim=-1) - query = query.unflatten(-1, (self.num_heads, -1)) - key = key.unflatten(-1, (self.num_heads, -1)) - value = value.unflatten(-1, (self.num_heads, -1)) + query = query.unflatten(-1, (self.to_qkv.num_heads, -1)) + key = key.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) + value = value.unflatten(-1, (self.to_qkv.num_kv_heads, -1)) query = self.norm_q(query) key = self.norm_k(key) @@ -170,7 +283,6 @@ class FeedForward(nn.Module): dim, [hidden_dim] * 2, bias=False, - disable_tp=True, return_bias=False, ) self.act = SiluAndMul() @@ -178,7 +290,7 @@ class FeedForward(nn.Module): hidden_dim, dim, bias=False, - disable_tp=True, + input_is_parallel=True, return_bias=False, ) @@ -379,6 +491,28 @@ class ZImageTransformer2DModel(nn.Module): assert len(all_patch_size) == len(all_f_patch_size) + tp_size = _get_tensor_parallel_size_from_context() + ffn_hidden_dim, final_out_dims, supported_tp_candidates = validate_zimage_tp_constraints( + dim=dim, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + in_channels=self.out_channels, + all_patch_size=tuple(all_patch_size), + all_f_patch_size=tuple(all_f_patch_size), + tensor_parallel_size=tp_size, + ) + + logger.info_once( + "Z-Image init: dim=%d n_heads=%d n_kv_heads=%d ffn_hidden_dim=%d final_out_dims=%s tp=%d (supported_tp=%s)", + dim, + n_heads, + n_kv_heads, + ffn_hidden_dim, + tuple(final_out_dims), + tp_size, + tuple(supported_tp_candidates), + ) + all_x_embedder = {} all_final_layer = {} for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)):