Canlin GuoandHongsheng Liu
f53152d340
Signed-off-by: gcanlin <canlinguosdu@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Distributed utilities for vLLM-Omni diffusion models."""
|
|
|
|
from vllm_omni.diffusion.distributed.hsdp import (
|
|
HSDPInferenceConfig,
|
|
apply_hsdp_to_model,
|
|
)
|
|
from vllm_omni.diffusion.distributed.parallel_state import (
|
|
get_fs_group,
|
|
get_fully_shard_rank,
|
|
get_fully_shard_world_size,
|
|
)
|
|
from vllm_omni.diffusion.distributed.sp_plan import (
|
|
SequenceParallelConfig,
|
|
SequenceParallelInput,
|
|
SequenceParallelModelPlan,
|
|
SequenceParallelOutput,
|
|
SequenceParallelPartialInput,
|
|
get_sp_plan_from_model,
|
|
validate_sp_plan,
|
|
)
|
|
from vllm_omni.diffusion.distributed.sp_sharding import (
|
|
ShardingValidator,
|
|
get_sharding_validator,
|
|
sp_gather,
|
|
sp_shard,
|
|
sp_shard_with_padding,
|
|
)
|
|
|
|
__all__ = [
|
|
# Config
|
|
"SequenceParallelConfig",
|
|
# Plan types
|
|
"SequenceParallelInput",
|
|
"SequenceParallelOutput",
|
|
"SequenceParallelPartialInput",
|
|
"SequenceParallelModelPlan",
|
|
"validate_sp_plan",
|
|
"get_sp_plan_from_model",
|
|
# Sharding utilities
|
|
"sp_shard",
|
|
"sp_gather",
|
|
"sp_shard_with_padding",
|
|
"ShardingValidator",
|
|
"get_sharding_validator",
|
|
"HSDPInferenceConfig",
|
|
"apply_hsdp_to_model",
|
|
# FS (Fully Shard) utilities
|
|
"get_fs_group",
|
|
"get_fully_shard_rank",
|
|
"get_fully_shard_world_size",
|
|
]
|