wehub-resource-sync
60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import logging
|
|
from typing import List, Dict, Optional
|
|
from .autotp_config import TPLayerSpec, PartitionType
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SUPPORTED_STYLES = {"colwise", "rowwise"}
|
|
|
|
|
|
class TPPlanConverter:
|
|
"""Convert HuggingFace tp_plan format to DeepSpeed TPLayerSpec format."""
|
|
|
|
@staticmethod
|
|
def convert(hf_tp_plan: Dict[str, str]) -> Optional[List[TPLayerSpec]]:
|
|
"""Convert HF tp_plan to DeepSpeed layer specs.
|
|
|
|
Returns None if the plan contains any unsupported partition styles,
|
|
allowing the caller to fall back to the existing AutoTP path.
|
|
"""
|
|
unsupported = {style for style in hf_tp_plan.values() if style.lower() not in SUPPORTED_STYLES}
|
|
if unsupported:
|
|
logger.warning(
|
|
"HuggingFace tp_plan contains unsupported partition style(s): %s. "
|
|
"Falling back to AutoTP preset-based partitioning.", sorted(unsupported))
|
|
return None
|
|
|
|
layer_specs = []
|
|
|
|
for pattern, partition in hf_tp_plan.items():
|
|
regex_pattern = TPPlanConverter._wildcard_to_regex(pattern)
|
|
|
|
if partition.lower() == "colwise":
|
|
partition_type = PartitionType.COLUMN
|
|
elif partition.lower() == "rowwise":
|
|
partition_type = PartitionType.ROW
|
|
|
|
# Only add .weight suffix if not already present
|
|
if not regex_pattern.endswith(r"\.weight"):
|
|
regex_pattern += r"\.weight$"
|
|
else:
|
|
regex_pattern += r"$"
|
|
|
|
layer_specs.append(TPLayerSpec(
|
|
patterns=[regex_pattern],
|
|
partition_type=partition_type,
|
|
))
|
|
|
|
return layer_specs
|
|
|
|
@staticmethod
|
|
def _wildcard_to_regex(pattern: str) -> str:
|
|
regex = pattern.replace('.', r'\.')
|
|
regex = regex.replace('*', r'.*')
|
|
return ".*" + regex
|