465 lines
17 KiB
Python
465 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""
|
|
|
|
import gc
|
|
import os
|
|
import socket
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.distributed as dist
|
|
import torch.nn as nn
|
|
from torch.distributed.tensor import DeviceMesh, DTensor
|
|
|
|
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
|
from vllm_omni.diffusion.distributed.hsdp import (
|
|
HSDPInferenceConfig,
|
|
_unshardable_parameters,
|
|
shard_model,
|
|
)
|
|
|
|
pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model]
|
|
|
|
|
|
class _PackedBlock(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(2, 2))
|
|
self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False)
|
|
self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False)
|
|
|
|
|
|
class _PackedModel(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.block = _PackedBlock()
|
|
self.root_weight = nn.Parameter(torch.ones(2))
|
|
|
|
|
|
def _find_free_port() -> int:
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
|
sock.bind(("127.0.0.1", 0))
|
|
return int(sock.getsockname()[1])
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def cpu_process_group():
|
|
if dist.is_initialized():
|
|
yield
|
|
return
|
|
|
|
master_port = _find_free_port()
|
|
os.environ.update(
|
|
{
|
|
"RANK": "0",
|
|
"LOCAL_RANK": "0",
|
|
"WORLD_SIZE": "1",
|
|
"MASTER_ADDR": "127.0.0.1",
|
|
"MASTER_PORT": str(master_port),
|
|
}
|
|
)
|
|
dist.init_process_group("gloo", rank=0, world_size=1)
|
|
try:
|
|
yield
|
|
finally:
|
|
dist.destroy_process_group()
|
|
for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
|
|
os.environ.pop(key, None)
|
|
gc.collect()
|
|
|
|
|
|
def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
|
|
model = _PackedModel()
|
|
ignored_params = _unshardable_parameters(model)
|
|
expected_ignored = {model.block.packed_weight, model.block.input_global_scale}
|
|
assert ignored_params == expected_ignored
|
|
|
|
packed_weight = model.block.packed_weight
|
|
input_global_scale = model.block.input_global_scale
|
|
shard_model(
|
|
model,
|
|
mesh=DeviceMesh("cpu", [0]),
|
|
hsdp_shard_conditions=[lambda name, _module: name == "block"],
|
|
ignored_params=ignored_params,
|
|
)
|
|
|
|
assert isinstance(model.block.weight, DTensor)
|
|
assert isinstance(model.root_weight, DTensor)
|
|
assert model.block.packed_weight is packed_weight
|
|
assert model.block.input_global_scale is input_global_scale
|
|
assert not isinstance(model.block.packed_weight, DTensor)
|
|
assert not isinstance(model.block.input_global_scale, DTensor)
|
|
|
|
|
|
class TestHSDPInferenceConfig:
|
|
"""Tests for HSDPInferenceConfig dataclass."""
|
|
|
|
def test_default_values(self):
|
|
"""Test default configuration values."""
|
|
config = HSDPInferenceConfig()
|
|
assert config.enabled is False
|
|
assert config.hsdp_replicate_size == 1
|
|
assert config.hsdp_shard_size == -1
|
|
|
|
def test_custom_values(self):
|
|
"""Test custom configuration values."""
|
|
config = HSDPInferenceConfig(
|
|
enabled=True,
|
|
hsdp_replicate_size=2,
|
|
hsdp_shard_size=4,
|
|
)
|
|
assert config.enabled is True
|
|
assert config.hsdp_replicate_size == 2
|
|
assert config.hsdp_shard_size == 4
|
|
|
|
|
|
class TestDiffusionParallelConfigHSDP:
|
|
"""Tests for HSDP settings in DiffusionParallelConfig."""
|
|
|
|
def test_hsdp_disabled_by_default(self):
|
|
"""HSDP should be disabled by default."""
|
|
config = DiffusionParallelConfig()
|
|
assert config.use_hsdp is False
|
|
assert config.hsdp_shard_size == -1
|
|
assert config.hsdp_replicate_size == 1
|
|
|
|
def test_hsdp_auto_shard_size(self):
|
|
"""Test auto-calculation of hsdp_shard_size when use_hsdp=True."""
|
|
# ulysses_degree=4 -> world_size=4
|
|
# hsdp_shard_size should be auto-calculated as 4 // 1 = 4
|
|
config = DiffusionParallelConfig(
|
|
ulysses_degree=4,
|
|
use_hsdp=True,
|
|
)
|
|
assert config.world_size == 4
|
|
assert config.hsdp_shard_size == 4
|
|
assert config.hsdp_replicate_size == 1
|
|
|
|
def test_hsdp_auto_shard_size_fails_standalone(self):
|
|
"""Test that auto-calculate fails when other parallelism is all 1."""
|
|
# When all other parallelism is 1, cannot auto-calculate
|
|
# User must specify hsdp_shard_size explicitly
|
|
with pytest.raises(ValueError, match="Cannot auto-calculate hsdp_shard_size"):
|
|
DiffusionParallelConfig(
|
|
use_hsdp=True,
|
|
# All other parallelism defaults to 1
|
|
)
|
|
|
|
def test_hsdp_standalone_mode(self):
|
|
"""Test standalone HSDP (HSDP without other parallelism)."""
|
|
# Standalone HSDP: all other parallelism=1, explicit shard_size
|
|
config = DiffusionParallelConfig(
|
|
use_hsdp=True,
|
|
hsdp_shard_size=4, # Explicit shard size
|
|
hsdp_replicate_size=1,
|
|
)
|
|
# world_size should be determined by HSDP
|
|
assert config.world_size == 4
|
|
assert config.hsdp_shard_size == 4
|
|
assert config.hsdp_replicate_size == 1
|
|
|
|
def test_hsdp_standalone_with_replicate(self):
|
|
"""Test standalone HSDP with replication."""
|
|
config = DiffusionParallelConfig(
|
|
use_hsdp=True,
|
|
hsdp_shard_size=4,
|
|
hsdp_replicate_size=2,
|
|
)
|
|
# world_size = shard_size * replicate_size
|
|
assert config.world_size == 8
|
|
assert config.hsdp_shard_size == 4
|
|
assert config.hsdp_replicate_size == 2
|
|
|
|
def test_hsdp_with_replicate(self):
|
|
"""Test HSDP with replication (hybrid mode) combined with other parallelism."""
|
|
# world_size=8, replicate=2 -> shard_size should be 4
|
|
config = DiffusionParallelConfig(
|
|
ulysses_degree=8,
|
|
use_hsdp=True,
|
|
hsdp_replicate_size=2,
|
|
)
|
|
assert config.world_size == 8
|
|
assert config.hsdp_shard_size == 4
|
|
assert config.hsdp_replicate_size == 2
|
|
|
|
def test_hsdp_explicit_shard_size_valid(self):
|
|
"""Test explicit hsdp_shard_size that matches world_size."""
|
|
config = DiffusionParallelConfig(
|
|
ulysses_degree=4,
|
|
use_hsdp=True,
|
|
hsdp_shard_size=4,
|
|
hsdp_replicate_size=1,
|
|
)
|
|
assert config.hsdp_shard_size == 4
|
|
|
|
def test_hsdp_explicit_shard_size_invalid(self):
|
|
"""Test that invalid HSDP dimensions raise an error when combined with other parallelism."""
|
|
with pytest.raises(ValueError, match="HSDP dimensions"):
|
|
DiffusionParallelConfig(
|
|
ulysses_degree=4, # world_size=4
|
|
use_hsdp=True,
|
|
hsdp_shard_size=3, # 1 * 3 != 4
|
|
hsdp_replicate_size=1,
|
|
)
|
|
|
|
def test_hsdp_replicate_size_exceeds_world_size(self):
|
|
"""Test that replicate_size > world_size raises an error."""
|
|
with pytest.raises(ValueError, match="replicate_size.*must evenly divide world_size"):
|
|
DiffusionParallelConfig(
|
|
ulysses_degree=4, # world_size=4
|
|
use_hsdp=True,
|
|
hsdp_replicate_size=8, # 8 > 4, invalid
|
|
)
|
|
|
|
def test_hsdp_combined_world_size(self):
|
|
"""Test that combined HSDP matches other parallelism world_size."""
|
|
config_no_hsdp = DiffusionParallelConfig(ulysses_degree=4)
|
|
config_with_hsdp = DiffusionParallelConfig(
|
|
ulysses_degree=4,
|
|
use_hsdp=True,
|
|
hsdp_shard_size=4,
|
|
)
|
|
# When combined with other parallelism, world_size should match
|
|
assert config_no_hsdp.world_size == config_with_hsdp.world_size == 4
|
|
|
|
def test_hsdp_standalone_world_size(self):
|
|
"""Test that standalone HSDP determines world_size."""
|
|
config_hsdp = DiffusionParallelConfig(
|
|
use_hsdp=True,
|
|
hsdp_shard_size=8,
|
|
)
|
|
# Standalone HSDP: world_size is determined by HSDP
|
|
assert config_hsdp.world_size == 8
|
|
|
|
def test_hsdp_cannot_use_with_tp(self):
|
|
"""Test that HSDP and Tensor Parallelism cannot be used together."""
|
|
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
|
DiffusionParallelConfig(
|
|
tensor_parallel_size=2,
|
|
use_hsdp=True,
|
|
hsdp_shard_size=4,
|
|
)
|
|
|
|
def test_hsdp_cannot_use_with_dp(self):
|
|
"""Test that HSDP and Data Parallelism cannot be used together."""
|
|
with pytest.raises(ValueError, match="cannot be used with TP or DP"):
|
|
DiffusionParallelConfig(
|
|
data_parallel_size=2,
|
|
use_hsdp=True,
|
|
hsdp_shard_size=4,
|
|
)
|
|
|
|
def test_from_dict_with_hsdp(self):
|
|
"""Test creating config from dict with HSDP settings."""
|
|
config = DiffusionParallelConfig.from_dict(
|
|
{
|
|
"ulysses_degree": 4,
|
|
"use_hsdp": True,
|
|
"hsdp_replicate_size": 2,
|
|
}
|
|
)
|
|
assert config.use_hsdp is True
|
|
assert config.hsdp_replicate_size == 2
|
|
assert config.hsdp_shard_size == 2 # auto: 4 // 2
|
|
|
|
|
|
class TestStandaloneHSDPDetection:
|
|
"""Tests for standalone HSDP detection and dit_parallel_size calculation.
|
|
|
|
These tests verify the logic used in initialize_model_parallel() to detect
|
|
standalone HSDP mode and calculate effective parallel sizes.
|
|
|
|
Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
|
|
"""
|
|
|
|
@staticmethod
|
|
def compute_standalone_hsdp_params(
|
|
data_parallel_size: int = 1,
|
|
cfg_parallel_size: int = 1,
|
|
sequence_parallel_size: int = 1,
|
|
pipeline_parallel_size: int = 1,
|
|
tensor_parallel_size: int = 1,
|
|
fully_shard_degree: int = 1,
|
|
hsdp_replicate_size: int = 1,
|
|
) -> dict:
|
|
"""Compute standalone HSDP detection parameters.
|
|
|
|
This mirrors the logic in initialize_model_parallel().
|
|
"""
|
|
dit_parallel_size = (
|
|
data_parallel_size
|
|
* cfg_parallel_size
|
|
* sequence_parallel_size
|
|
* pipeline_parallel_size
|
|
* tensor_parallel_size
|
|
)
|
|
|
|
# Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
|
|
is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1
|
|
|
|
# For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
|
|
if is_standalone_hsdp:
|
|
effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
|
|
else:
|
|
effective_dit_parallel_size = dit_parallel_size
|
|
|
|
effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size
|
|
|
|
return {
|
|
"original_dit_parallel_size": dit_parallel_size,
|
|
"is_standalone_hsdp": is_standalone_hsdp,
|
|
"effective_dit_parallel_size": effective_dit_parallel_size,
|
|
"effective_dp_size": effective_dp_size,
|
|
}
|
|
|
|
def test_standalone_hsdp_basic(self):
|
|
"""Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
fully_shard_degree=4,
|
|
hsdp_replicate_size=1,
|
|
)
|
|
assert result["original_dit_parallel_size"] == 1
|
|
assert result["is_standalone_hsdp"] is True
|
|
assert result["effective_dit_parallel_size"] == 4
|
|
assert result["effective_dp_size"] == 4
|
|
|
|
def test_standalone_hsdp_with_replicate(self):
|
|
"""Test standalone HSDP with replication (shard_size=4, replicate=2)."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
fully_shard_degree=4,
|
|
hsdp_replicate_size=2,
|
|
)
|
|
assert result["original_dit_parallel_size"] == 1
|
|
assert result["is_standalone_hsdp"] is True
|
|
assert result["effective_dit_parallel_size"] == 8 # 4 * 2
|
|
assert result["effective_dp_size"] == 8
|
|
|
|
def test_combined_hsdp_sp_not_standalone(self):
|
|
"""Test HSDP combined with SP is NOT detected as standalone.
|
|
|
|
This is a regression test for the bug where the condition
|
|
`dit_parallel_size == fully_shard_degree` incorrectly matched
|
|
combined modes like SP=4 + HSDP=4.
|
|
"""
|
|
result = self.compute_standalone_hsdp_params(
|
|
sequence_parallel_size=4,
|
|
fully_shard_degree=4,
|
|
hsdp_replicate_size=1,
|
|
)
|
|
assert result["original_dit_parallel_size"] == 4
|
|
assert result["is_standalone_hsdp"] is False
|
|
# Should NOT override dp_size for combined mode
|
|
assert result["effective_dp_size"] == 1 # original data_parallel_size
|
|
|
|
def test_combined_hsdp_cfg_not_standalone(self):
|
|
"""Test HSDP combined with CFG is NOT detected as standalone."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
cfg_parallel_size=2,
|
|
fully_shard_degree=4,
|
|
hsdp_replicate_size=1,
|
|
)
|
|
assert result["original_dit_parallel_size"] == 2
|
|
assert result["is_standalone_hsdp"] is False
|
|
assert result["effective_dp_size"] == 1
|
|
|
|
def test_combined_hsdp_pp_not_standalone(self):
|
|
"""Test HSDP combined with PP is NOT detected as standalone."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
pipeline_parallel_size=2,
|
|
fully_shard_degree=4,
|
|
hsdp_replicate_size=1,
|
|
)
|
|
assert result["original_dit_parallel_size"] == 2
|
|
assert result["is_standalone_hsdp"] is False
|
|
assert result["effective_dp_size"] == 1
|
|
|
|
def test_no_hsdp_not_standalone(self):
|
|
"""Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
fully_shard_degree=1,
|
|
)
|
|
assert result["original_dit_parallel_size"] == 1
|
|
assert result["is_standalone_hsdp"] is False
|
|
assert result["effective_dp_size"] == 1
|
|
|
|
def test_combined_multiple_parallelism_not_standalone(self):
|
|
"""Test HSDP combined with multiple parallelism is NOT standalone."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
sequence_parallel_size=2,
|
|
cfg_parallel_size=2,
|
|
fully_shard_degree=4,
|
|
hsdp_replicate_size=1,
|
|
)
|
|
assert result["original_dit_parallel_size"] == 4 # 2 * 2
|
|
assert result["is_standalone_hsdp"] is False
|
|
assert result["effective_dp_size"] == 1
|
|
|
|
def test_standalone_hsdp_large_shard(self):
|
|
"""Test standalone HSDP with large shard size."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
fully_shard_degree=8,
|
|
hsdp_replicate_size=1,
|
|
)
|
|
assert result["is_standalone_hsdp"] is True
|
|
assert result["effective_dit_parallel_size"] == 8
|
|
assert result["effective_dp_size"] == 8
|
|
|
|
def test_standalone_hsdp_large_replicate(self):
|
|
"""Test standalone HSDP with large replicate size."""
|
|
result = self.compute_standalone_hsdp_params(
|
|
fully_shard_degree=4,
|
|
hsdp_replicate_size=4,
|
|
)
|
|
assert result["is_standalone_hsdp"] is True
|
|
assert result["effective_dit_parallel_size"] == 16 # 4 * 4
|
|
assert result["effective_dp_size"] == 16
|
|
|
|
|
|
class TestHSDPShardConditions:
|
|
"""Tests for _hsdp_shard_conditions matching logic."""
|
|
|
|
@staticmethod
|
|
def _is_transformer_block(name: str, module: nn.Module) -> bool:
|
|
"""Example shard condition matching transformer blocks."""
|
|
return "blocks" in name and name.split(".")[-1].isdigit()
|
|
|
|
def test_condition_matches_blocks(self):
|
|
"""Test that condition matches transformer block patterns."""
|
|
cond = self._is_transformer_block
|
|
# Should match
|
|
assert cond("blocks.0", nn.Linear(10, 10)) is True
|
|
assert cond("blocks.15", nn.Linear(10, 10)) is True
|
|
assert cond("transformer.blocks.0", nn.Linear(10, 10)) is True
|
|
# Should not match
|
|
assert cond("blocks", nn.Linear(10, 10)) is False
|
|
assert cond("blocks.norm", nn.Linear(10, 10)) is False
|
|
assert cond("embeddings", nn.Linear(10, 10)) is False
|
|
|
|
def test_model_with_shard_conditions(self):
|
|
"""Test model with _hsdp_shard_conditions attribute."""
|
|
|
|
class MockModel(nn.Module):
|
|
@staticmethod
|
|
def _is_block(name: str, module: nn.Module) -> bool:
|
|
return name.startswith("blocks.") and name.split(".")[-1].isdigit()
|
|
|
|
_hsdp_shard_conditions = [_is_block]
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.blocks = nn.ModuleList([nn.Linear(10, 10) for _ in range(2)])
|
|
|
|
model = MockModel()
|
|
conditions = getattr(model, "_hsdp_shard_conditions", None)
|
|
assert conditions is not None
|
|
assert len(conditions) == 1
|
|
|
|
# Verify conditions work on actual model modules
|
|
matched = []
|
|
for name, module in model.named_modules():
|
|
if any(cond(name, module) for cond in conditions):
|
|
matched.append(name)
|
|
assert "blocks.0" in matched
|
|
assert "blocks.1" in matched
|