[Model] Avoid repeated attention and MLP work in Klein TP (#8553)
Signed-off-by: QianCyrus <101633534+QianCyrus@users.noreply.github.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
QianCyrusandHongsheng Liu authored and GitHub committed
1 parent
b21df3bcb6
commit
8ef77be66e
4 files changed
+415
-14
No files matched your search
@@ -0,0 +1,331 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""CPU regressions for Klein's packed TP layout and local-head execution."""
|
||||
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from vllm.model_executor import parameter
|
||||
from vllm.model_executor.layers import linear
|
||||
from vllm.model_executor.layers.utils import default_unquantized_gemm
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.forward_context import ForwardContext
|
||||
from vllm_omni.diffusion.models.flux2_klein import flux2_klein_transformer as klein
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _default_dtype(dtype):
|
||||
previous = torch.get_default_dtype()
|
||||
torch.set_default_dtype(dtype)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
torch.set_default_dtype(previous)
|
||||
|
||||
|
||||
class _CPUAttention(torch.nn.Module):
|
||||
"""Use SDPA while retaining the model's actual Q/K norm and RoPE."""
|
||||
|
||||
def __init__(self, num_heads, head_size, softmax_scale, causal=False, **kwargs):
|
||||
super().__init__()
|
||||
self.num_heads = num_heads
|
||||
self.scale = softmax_scale
|
||||
|
||||
def forward(self, query, key, value, metadata=None):
|
||||
assert query.shape[2] == self.num_heads
|
||||
return F.scaled_dot_product_attention(
|
||||
query.transpose(1, 2),
|
||||
key.transpose(1, 2),
|
||||
value.transpose(1, 2),
|
||||
attn_mask=None if metadata is None else metadata.attn_mask,
|
||||
scale=self.scale,
|
||||
).transpose(1, 2)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cpu_dispatch(monkeypatch):
|
||||
monkeypatch.setattr(linear, "dispatch_unquantized_gemm", lambda *a, **kw: default_unquantized_gemm)
|
||||
monkeypatch.setattr(klein, "Attention", _CPUAttention)
|
||||
monkeypatch.setattr(klein, "get_forward_context", ForwardContext)
|
||||
|
||||
|
||||
def _set_tp(monkeypatch, tp, rank):
|
||||
# Keep real vLLM parameters, constructors and checkpoint loaders. Only the
|
||||
# distributed process metadata is substituted for this CPU unit test.
|
||||
for module in (linear, parameter):
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_world_size", lambda: tp)
|
||||
monkeypatch.setattr(module, "get_tensor_model_parallel_rank", lambda: rank)
|
||||
monkeypatch.setattr(klein, "get_tensor_model_parallel_world_size", lambda: tp)
|
||||
|
||||
|
||||
def _attention(monkeypatch, tp, rank, *, local_tp=True, heads=8, dim_head=8, mlp_ratio=3, sp=1):
|
||||
_set_tp(monkeypatch, tp, rank)
|
||||
with _default_dtype(torch.bfloat16):
|
||||
attention = klein.Flux2ParallelSelfAttention(
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=tp, ulysses_degree=sp),
|
||||
query_dim=heads * dim_head,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
mlp_ratio=mlp_ratio,
|
||||
bias=True,
|
||||
local_tp=local_tp,
|
||||
)
|
||||
attention.norm_q.forward = attention.norm_q.forward_native
|
||||
attention.norm_k.forward = attention.norm_k.forward_native
|
||||
attention.rope.forward = attention.rope.forward_native
|
||||
return attention
|
||||
|
||||
|
||||
def _rank_rows(tensor, sizes, tp, rank):
|
||||
return torch.cat([part.chunk(tp, dim=0)[rank] for part in tensor.split(sizes, dim=0)])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp", [1, 2, 4, 8])
|
||||
@pytest.mark.parametrize("name", ["weight", "bias"])
|
||||
def test_packed_checkpoint_reconstruction_and_reload(monkeypatch, tp, name):
|
||||
sizes = [64, 64, 64, 192, 192]
|
||||
shape = (sum(sizes), 64) if name == "weight" else (sum(sizes),)
|
||||
canonical = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.float32).reshape(shape)
|
||||
shards = []
|
||||
for rank in range(tp):
|
||||
attention = _attention(monkeypatch, tp, rank)
|
||||
assert attention.local_tp == (tp > 1)
|
||||
assert attention.query_num_heads == 8 // tp
|
||||
projection = attention.to_qkv_mlp_proj
|
||||
param = getattr(projection, name)
|
||||
# Exact row markers verify the installed loader, without BF16 rounding
|
||||
# hiding a channel-order error. The constructor still uses BF16.
|
||||
param.data = param.data.float()
|
||||
for weights in (canonical + 1, canonical):
|
||||
param.weight_loader(param, weights)
|
||||
torch.testing.assert_close(param, _rank_rows(weights, sizes, tp, rank), rtol=0, atol=0)
|
||||
shards.append(param.detach().clone())
|
||||
local_sizes = [size // tp for size in sizes]
|
||||
reconstructed = torch.cat(
|
||||
[torch.cat([shard.split(local_sizes, dim=0)[i] for shard in shards]) for i in range(len(sizes))]
|
||||
)
|
||||
torch.testing.assert_close(reconstructed, canonical, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp,heads", [(2, 8), (4, 24), (8, 24)])
|
||||
@pytest.mark.parametrize("mask_kind", ["none", "3d", "broadcast", "per_head"])
|
||||
@torch.inference_mode()
|
||||
def test_local_heads_and_gather_order_match_dense_attention(monkeypatch, tp, heads, mask_kind):
|
||||
torch.manual_seed(17)
|
||||
dim = heads * 8
|
||||
baseline = _attention(monkeypatch, 1, 0, local_tp=False, heads=heads)
|
||||
for name, param in baseline.named_parameters():
|
||||
if name.startswith("norm_"):
|
||||
param.fill_(1)
|
||||
else:
|
||||
param.normal_(std=0.03)
|
||||
# A non-contiguous input and batch > 1 exercise the reshape boundaries.
|
||||
hidden = torch.randn(2, dim, 7, dtype=torch.bfloat16).transpose(1, 2)
|
||||
angles = torch.randn(7, 4)
|
||||
rotary = angles.cos(), angles.sin()
|
||||
mask = None
|
||||
if mask_kind != "none":
|
||||
mask_heads = heads if mask_kind == "per_head" else 1
|
||||
mask = torch.rand(2, mask_heads, 7, 7) > 0.25
|
||||
mask[..., 0] = True
|
||||
if mask_kind == "3d":
|
||||
mask = mask[:, 0]
|
||||
|
||||
captured = []
|
||||
hook = baseline.to_out.register_forward_pre_hook(lambda module, args: captured.append(args[0].clone()))
|
||||
expected = baseline(hidden, attention_mask=mask, image_rotary_emb=rotary)
|
||||
hook.remove()
|
||||
full_attention, full_mlp = captured[0].split([dim, 3 * dim], dim=-1)
|
||||
weights = {name: param.detach().clone() for name, param in baseline.named_parameters()}
|
||||
|
||||
for rank in range(tp):
|
||||
local = _attention(monkeypatch, tp, rank, heads=heads)
|
||||
for name, param in local.named_parameters():
|
||||
getattr(param, "weight_loader", default_weight_loader)(param, weights[name])
|
||||
feature_calls: list[int] = []
|
||||
output_calls: list[int] = []
|
||||
|
||||
def gather_features(value, dim=-1):
|
||||
reference = (full_attention, full_mlp)[len(feature_calls)]
|
||||
torch.testing.assert_close(value, reference.chunk(tp, dim=-1)[rank], rtol=0.02, atol=0.002)
|
||||
feature_calls.append(value.shape[-1])
|
||||
return reference
|
||||
|
||||
def gather_output(value, dim=-1):
|
||||
torch.testing.assert_close(value, expected.chunk(tp, dim=-1)[rank], rtol=0.02, atol=0.002)
|
||||
output_calls.append(value.shape[-1])
|
||||
return expected
|
||||
|
||||
# Each outgoing tensor must match the independently computed dense
|
||||
# result before the collective supplies the other ranks' slices.
|
||||
monkeypatch.setattr(klein, "tensor_model_parallel_all_gather", gather_features)
|
||||
monkeypatch.setattr(linear, "tensor_model_parallel_all_gather", gather_output)
|
||||
actual = local(hidden, attention_mask=mask, image_rotary_emb=rotary)
|
||||
torch.testing.assert_close(actual, expected)
|
||||
assert feature_calls == [dim // tp, 3 * dim // tp]
|
||||
assert output_calls == [dim // tp]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
[
|
||||
{"tp": 1},
|
||||
{"tp": 4, "local_tp": False},
|
||||
{"tp": 4, "sp": 2},
|
||||
{"tp": 4, "heads": 6},
|
||||
{"tp": 4, "mlp_ratio": 2.1},
|
||||
],
|
||||
)
|
||||
def test_ineligible_shapes_keep_original_projection(monkeypatch, kwargs):
|
||||
attention = _attention(monkeypatch, rank=0, **kwargs)
|
||||
assert not attention.local_tp
|
||||
assert type(attention.to_qkv_mlp_proj) is linear.ColumnParallelLinear
|
||||
assert attention.to_qkv_mlp_proj.gather_output
|
||||
assert attention.query_num_heads == attention.heads
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"overrides,parallel_overrides,eligible",
|
||||
[
|
||||
({}, {}, True),
|
||||
({"dtype": torch.float32}, {}, False),
|
||||
({"dtype": torch.float16}, {}, False),
|
||||
({}, {"tensor_parallel_size": 1}, False),
|
||||
({}, {"ulysses_degree": 2}, False),
|
||||
({"quantization_config": "fp8"}, {}, False),
|
||||
({"lora_path": "local-adapter"}, {}, False),
|
||||
({"enable_cpu_offload": True}, {}, False),
|
||||
({"enable_layerwise_offload": True}, {}, False),
|
||||
({"diffusion_offload_config": {"mode": "module", "components": ["dit"]}}, {}, False),
|
||||
({}, {"tensor_parallel_size": 1, "use_hsdp": True, "hsdp_shard_size": 4}, False),
|
||||
],
|
||||
)
|
||||
def test_pipeline_eligibility_uses_resolved_config(tmp_path, overrides, parallel_overrides, eligible):
|
||||
parallel = DiffusionParallelConfig(**({"tensor_parallel_size": 4} | parallel_overrides))
|
||||
config = OmniDiffusionConfig(
|
||||
**({"model": str(tmp_path), "dtype": torch.bfloat16} | overrides),
|
||||
parallel_config=parallel,
|
||||
)
|
||||
assert klein._use_local_single_stream_tp(config) is eligible
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tp", [1, 2, 4, 8])
|
||||
def test_host_restore_contract_and_model_weight_reload(monkeypatch, tp):
|
||||
_set_tp(monkeypatch, tp, 0)
|
||||
kwargs = dict(
|
||||
in_channels=8,
|
||||
num_layers=1,
|
||||
num_single_layers=1,
|
||||
attention_head_dim=8,
|
||||
num_attention_heads=8,
|
||||
joint_attention_dim=8,
|
||||
timestep_guidance_channels=8,
|
||||
axes_dims_rope=(2, 2, 2, 2),
|
||||
guidance_embeds=False,
|
||||
)
|
||||
with _default_dtype(torch.bfloat16):
|
||||
baseline = klein.Flux2Transformer2DModel(**kwargs)
|
||||
model = klein.Flux2Transformer2DModel(**kwargs, local_tp=True)
|
||||
assert baseline.host_weight_restore_contract.version == "1"
|
||||
assert model.host_weight_restore_contract.version == ("2-local-single-stream-tp" if tp > 1 else "1")
|
||||
model.validate_restored_host_weights()
|
||||
name = "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight"
|
||||
weights = torch.linspace(-1, 1, 576 * 64).reshape(576, 64).to(torch.bfloat16)
|
||||
for canonical in (weights, weights.neg()):
|
||||
assert name in model.load_weights([(name, canonical)])
|
||||
param = dict(model.named_parameters())[name]
|
||||
torch.testing.assert_close(param, _rank_rows(canonical, [64, 64, 64, 192, 192], tp, 0), rtol=0, atol=0)
|
||||
with _default_dtype(torch.bfloat16):
|
||||
restored = klein.Flux2Transformer2DModel(**kwargs, local_tp=True)
|
||||
restored.load_state_dict(model.state_dict(), strict=True)
|
||||
restored.validate_restored_host_weights()
|
||||
torch.testing.assert_close(dict(restored.named_parameters())[name], param, rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sizes", [(8, 24), (8, 8, 8, 24, 24)])
|
||||
@pytest.mark.parametrize("tp", [1, 2, 4, 8])
|
||||
@torch.inference_mode()
|
||||
def test_fused_lora_global_rows_dynamic_activation(monkeypatch, sizes, tp):
|
||||
from vllm.lora.lora_model import LoRAModel
|
||||
from vllm.lora.lora_weights import LoRALayerWeights
|
||||
from vllm.lora.peft_helper import PEFTHelper
|
||||
from vllm.lora.request import LoRARequest
|
||||
|
||||
from vllm_omni.diffusion.lora.layers.column_parallel_linear import DiffusionMergedColumnParallelLinearWithLoRA
|
||||
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
|
||||
|
||||
torch.manual_seed(29)
|
||||
weights = torch.randn(sum(sizes), 8)
|
||||
lora_a = torch.randn(2, 8)
|
||||
lora_b = torch.randn(sum(sizes), 2)
|
||||
hidden = torch.randn(2, 3, 8)
|
||||
name = "transformer.to_qkv_mlp_proj"
|
||||
peft = PEFTHelper(r=2, lora_alpha=2, target_modules=["to_qkv_mlp_proj"])
|
||||
request = LoRARequest(lora_name="test", lora_int_id=7, lora_path="unused-by-memory-loader")
|
||||
for rank in range(tp):
|
||||
_set_tp(monkeypatch, tp, rank)
|
||||
projection = linear.MergedColumnParallelLinear(8, list(sizes), bias=False, gather_output=False)
|
||||
projection.weight.weight_loader(projection.weight, weights)
|
||||
pipeline = torch.nn.Module()
|
||||
pipeline.transformer = torch.nn.Module()
|
||||
pipeline.transformer.to_qkv_mlp_proj = projection
|
||||
manager = DiffusionLoRAManager(pipeline, device=torch.device("cpu"), dtype=torch.float32)
|
||||
adapter = LoRAModel(7, 2, {name: LoRALayerWeights(name, 2, 2, lora_a, lora_b)})
|
||||
# Adapter disk I/O is outside this regression. Actual manager injection,
|
||||
# five-/two-slice wrapper, rank slicing and activation all run below.
|
||||
monkeypatch.setattr(manager, "_load_adapter", lambda request: (adapter, peft))
|
||||
base_output = F.linear(hidden, _rank_rows(weights, sizes, tp, rank))
|
||||
delta = F.linear(F.linear(hidden, lora_a), _rank_rows(lora_b, sizes, tp, rank))
|
||||
for scale in (0.5, 1.0):
|
||||
manager.set_active_adapter(request, lora_scale=scale)
|
||||
wrapped = pipeline.transformer.to_qkv_mlp_proj
|
||||
assert isinstance(wrapped, DiffusionMergedColumnParallelLinearWithLoRA)
|
||||
assert wrapped.n_slices == len(sizes)
|
||||
actual, _ = wrapped(hidden)
|
||||
torch.testing.assert_close(actual, base_output + scale * delta, rtol=1e-5, atol=1e-5)
|
||||
manager.set_active_adapter(None)
|
||||
actual, _ = wrapped(hidden)
|
||||
torch.testing.assert_close(actual, base_output, rtol=1e-5, atol=1e-5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("image_count", [1, 2])
|
||||
def test_native_editing_reaches_normal_input_validation(tmp_path, monkeypatch, image_count):
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.flux2_klein.pipeline_flux2_klein import Flux2KleinPipeline
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
class InputsReachedError(Exception):
|
||||
pass
|
||||
|
||||
pipeline = object.__new__(Flux2KleinPipeline)
|
||||
pipeline.od_config = OmniDiffusionConfig(
|
||||
model=str(tmp_path),
|
||||
dtype=torch.bfloat16,
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=4),
|
||||
)
|
||||
assert klein._use_local_single_stream_tp(pipeline.od_config)
|
||||
images = [Image.new("RGB", (16, 16)) for _ in range(image_count)]
|
||||
image_input = images[0] if image_count == 1 else images
|
||||
request = OmniDiffusionRequest(
|
||||
prompt={"prompt": "paint the boat red", "multi_modal_data": {"image": image_input}},
|
||||
request_id="native-editing",
|
||||
sampling_params=OmniDiffusionSamplingParams(height=512, width=512, num_inference_steps=4),
|
||||
)
|
||||
|
||||
def check_inputs(**kwargs):
|
||||
assert kwargs["prompt"] == "paint the boat red"
|
||||
raise InputsReachedError
|
||||
|
||||
monkeypatch.setattr(pipeline, "check_inputs", check_inputs)
|
||||
with pytest.raises(InputsReachedError):
|
||||
pipeline.forward(DiffusionRequestBatch(requests=[request]))
|
||||
assert request.prompt["multi_modal_data"]["image"] is image_input
|
||||
Reference in new issue
Block a user