diff --git a/HF5-ENVIRONMENT.md b/HF5-ENVIRONMENT.md new file mode 100644 index 000000000000..a1e46521ad11 --- /dev/null +++ b/HF5-ENVIRONMENT.md @@ -0,0 +1,111 @@ +# HF5 环境留档与服务器安装 + +本次固定的是 2026-09-22 Mac 上 `colossalai-hf5` 的实际已安装版本：Python 3.11.16、PyTorch 2.11.0、TorchVision 0.26.0、TorchAudio 2.11.0、Transformers 5.17.0、pytest 8.4.2。原有 requirements 文件未修改。 + +这是当前开发环境的版本快照，**不是已经验证通过的完整 Linux GPU 环境**。Mac 快照没有 Linux PyTorch 所需的 NVIDIA/Triton 等平台依赖，也没有锁定系统驱动、编译器、CUDA Toolkit 和包下载哈希；这些需要在服务器安装并验证后再补充留档。 + +## 文件分别有什么用 + +| 文件 | 用途 | +| --- | --- | +| `requirements-hf5.lock.txt` | 当前已安装 Python 包的精确版本，已移除本机路径及 ColossalAI Git 安装地址。服务器安装 CUDA PyTorch 后再安装此文件。 | +| `requirements-hf5-cu128.txt` | 仅安装 PyTorch 三件套，单独使用官方 CUDA 12.8 wheel 索引。 | +| `environment-hf5-server.yml` | 创建 Python 3.11.16 Conda 环境的起点；它本身不包含下面的 pip 安装步骤。 | +| `logs/2026-09-22-hf5-environment-lock/` | 原始 pip/Conda 导出、实际依赖元数据、检查结果和校验和。原始 Conda 导出用于 Mac 留档，不用于 Linux 安装。 | + +## 1. 上传当前源码 + +把当前修改后的仓库目录上传到服务器，包含新增的根目录环境文件、`tests/test_shardformer/test_hf5_compat/` 和 `logs/2026-09-22-qwen-gpu-validation/`。不要上传 `anaconda3/`、`__pycache__/`、`.pytest_cache/`；Mac 环境目录不能直接搬到 Linux 使用。 + +只克隆远端仓库、只上传 Git 已跟踪文件，或者直接使用原始 `pip freeze` 中的 `-e git+...`，都会遗漏尚未提交的修改或新增文件。开发安装必须指向这份实际上传的源码。 + +## 2. 在 Linux NVIDIA 服务器安装 + +以下命令在**服务器上传后的仓库根目录**依次执行，假设已安装 Conda。每一步成功后再执行下一步。 + +```bash +conda env create -f environment-hf5-server.yml +conda activate colossalai-hf5-gpu + +python -m pip install pip==26.2.1 setuptools==81.0.0 wheel==0.48.0 +python -m pip install -r requirements-hf5-cu128.txt +python -m pip install -r requirements-hf5.lock.txt + +BUILD_EXT=0 python -m pip install -e . --no-deps --no-build-isolation +``` + +CUDA wheel 版本组合来自 [PyTorch 官方版本页](https://pytorch.org/get-started/previous-versions/)。必须在新环境中先安装 CUDA 文件，再安装普通依赖文件，避免使用其他来源的 Torch 构建。普通清单中的 `torch==2.11.0` 也允许 `2.11.0+cu128`，后续用 `torch.version.cuda` 确认 CUDA 构建。 + +Python 的精确补丁版本在 Linux Conda 源上的可用性尚未实测；如果解析失败，保留错误日志，不要把失败当作环境配置完成。其他包在 Linux 上也可能需要不同平台的 wheel 或编译工具。安装器会补齐平台专属的传递依赖，其版本尚未被本次 Mac 快照锁定。 + +`--no-deps` 只用于最后的 ColossalAI 源码安装，避免仓库旧声明把 Torch/Transformers 降级。它不会补齐缺失依赖。`BUILD_EXT=0` 表示安装时不预编译扩展；运行某些 GPU kernel 时仍可能需要 CUDA Toolkit 和编译工具。此轮不要直接安装原有 `requirements/requirements.txt` 或 `requirements/requirements-test.txt` 来补包，它们仍包含旧版约束。 + +## 3. 验证安装结果 + +```bash +nvidia-smi +python - <<'PY' +from pathlib import Path +import torch +import torch.distributed as dist +import transformers +import colossalai +from packaging.version import Version + +print('torch:', torch.__version__, 'CUDA:', torch.version.cuda) +print('transformers:', transformers.__version__) +print('colossalai:', colossalai.__file__) +print('GPUs:', torch.cuda.device_count()) +assert Version(torch.__version__).base_version == '2.11.0' +assert torch.version.cuda == '12.8' +assert transformers.__version__ == '5.17.0' +assert Path(colossalai.__file__).resolve().parent == Path.cwd() / 'colossalai' +assert torch.cuda.is_available() +assert dist.is_nccl_available() +x = torch.arange(6, dtype=torch.float32, device='cuda').reshape(2, 3) +torch.testing.assert_close((x @ x.T).cpu(), torch.tensor([[5., 14.], [14., 50.]])) +print('Source import and basic CUDA computation: PASS') +PY +``` + +这一步只验证环境和一个 CUDA 计算，不代表模型数值对齐或多卡通信通过。接着按照 [Qwen GPU 验证指南](logs/2026-09-22-qwen-gpu-validation/GPU-validation-guide.md) 执行单卡数值检查，再收集和运行 TP/PP/SP 测试。 + +## 4. 已知缺口：当前 pip check 不通过 + +留档时实际发现： + +- 仓库声明 `torch>=2.2.0,<=2.5.1`、`transformers==4.51.3`、`peft>=0.7.1,<=0.13.2`，与当前安装的 2.11.0、5.17.0、0.19.1 不一致。 +- 仓库声明但本地未安装：`diffusers`、`fabric`、`fastapi`、`google`、`pre-commit`、`protobuf`、`ray`、`rpyc`、`sentencepiece`、`uvicorn`。 + +这些缺口被如实保留，没有为了让快照看起来完整而新增未经验证的包。已完成的 Qwen 本地测试不覆盖所有仓库功能；原有 GPU 测试收集会经过 model_zoo 的聚合导入，之前已经遇到缺少 `diffusers`。完整 GPU 测试仍需按实际导入错误补齐依赖、确认 HF5 兼容版本并重新留档。不要把所有 `pip check` 错误都视为可以忽略：上面是已知情况，新增冲突需要单独排查。 + +## 5. 服务器验证后补充精确快照 + +在同一服务器环境、仓库根目录执行： + +```bash +hf5_run_dir="logs/hf5-server-env-$(date +%Y%m%d-%H%M%S)" +mkdir -p "$hf5_run_dir" +python -m pip freeze --all > "$hf5_run_dir/pip-freeze-raw.txt" +python -m pip check > "$hf5_run_dir/pip-check.txt" 2>&1 +python -m pip inspect > "$hf5_run_dir/pip-inspect.json" +conda env export > "$hf5_run_dir/conda-server.yml" +nvidia-smi > "$hf5_run_dir/nvidia-smi.txt" +git rev-parse HEAD > "$hf5_run_dir/git-commit.txt" +git diff --binary > "$hf5_run_dir/working-tree.patch" +git ls-files --others --exclude-standard > "$hf5_run_dir/untracked-files.txt" +python - <<'PY' > "$hf5_run_dir/requirements-server-resolved.txt" +import importlib.metadata as md +import re +pins = [] +for d in md.distributions(): + name = re.sub(r'[-_.]+', '-', d.metadata['Name']).lower() + if name != 'colossalai': + pins.append(f'{name}=={d.version}') +print('# Installed server package versions; install transferred ColossalAI source separately.') +print('# CUDA wheels require the matching PyTorch index; this file has no artifact hashes.') +print('\n'.join(sorted(pins))) +PY +``` + +`pip check` 当前可能返回非零，需阅读其输出；如果 shell 开启了 `set -e`，在保存后续快照前先处理这一退出。记录中包含 Linux 新增的 NVIDIA/Triton 等版本。Git diff 不包含未跟踪文件的内容，仍需保存实际新增源码及测试文件。只有完成相应 GPU 验证后，才能把服务器快照标记为已验证基线。 diff --git a/colossalai/_compat.py b/colossalai/_compat.py new file mode 100644 index 000000000000..07b184a4bbee --- /dev/null +++ b/colossalai/_compat.py @@ -0,0 +1,225 @@ +r"""Internal transformers compatibility shim for ColossalAI (not exported). + +Background +---------- +ColossalAI pins transformers to ``==4.51.3`` (``requirements/requirements.txt``). +During the migration to v5, several referenced transformers internals changed in two ways: + +1. **Moved** — for example, ``no_init_weights`` moved from ``transformers.modeling_utils`` + to ``transformers.initialization``. +2. **Removed** — ``is_remote_url`` / ``download_url`` / ``is_safetensors_available`` + have no remaining definitions in the v5 transformers installation tree + (see ``scripts/N11_symbol_relocation_probe.py``). + +This module centralizes both differences so version-specific branches do not spread across +``lazy`` / ``shardformer`` / ``checkpoint_io``. + +Conventions (following existing repository style, not introducing a new pattern) +---------------------------------------------------------------------------------- +* **Prefer feature detection over version strings**: moved symbols use two + ``try/except ImportError`` import paths, covering v4 and v5 without parsing + ``transformers.__version__``. An existing example is + ``shardformer/policies/qwen2.py``, which uses + ``hasattr(self.model.config, "num_key_value_heads")`` for feature detection. +* **Do not silently change semantics when a capability is removed**: return ``None`` and + let the call site handle it explicitly and raise an explanatory error. +* **Importing this module has no side effects**: do not import heavyweight transformers + modules or log at module scope. As with ``shardformer/_utils.py``, all ``get_*`` helpers + resolve dependencies lazily inside functions. + +Usage +----- + from colossalai._compat import get_no_init_weights, is_safetensors_available + + with get_no_init_weights()(): + ... + + use_safetensors = kwargs.pop("use_safetensors", None if is_safetensors_available() else False) +""" + +from importlib import import_module +from importlib.util import find_spec +from inspect import signature +from typing import Any, Callable, Dict, Optional +from urllib.parse import urlparse + +# Cache the result of `get_auth_kwarg_name` lazily to avoid inspecting the signature per load. +_AUTH_KWARG_NAME: Optional[str] = None +_IS_TRANSFORMERS_V5: Optional[bool] = None + + +def get_no_init_weights() -> Callable[[], Any]: + r"""Return the ``no_init_weights`` context-manager factory. + + Transformers v5 defines ``no_init_weights`` in ``transformers.initialization`` + (v5.17.0: ``initialization.py:254``); v4 keeps it in ``transformers.modeling_utils``. + + Returns: + Callable[[], Any]: A no-argument callable that returns a context manager. + """ + try: + from transformers.initialization import no_init_weights as _fn + except ImportError: + from transformers.modeling_utils import no_init_weights as _fn + return _fn + + +def is_safetensors_available() -> bool: + r"""Return whether ``safetensors`` is available. + + v4 provides ``transformers.utils.is_safetensors_available``; v5 removed that helper + because safetensors became a hard dependency, so this falls back to ``importlib``. + + The repository also has a same-named implementation at + ``colossalai/checkpoint_io/utils.py:62``. Its body is + ``try: return True / except ImportError: return False`` and therefore always returns + ``True`` (a remnant of an early module-level import). Do not reuse it here, or the + availability check would become a constant. + + Returns: + bool: ``True`` when the package is available. + """ + try: + from transformers.utils import is_safetensors_available as _fn + except ImportError: + return find_spec("safetensors") is not None + return bool(_fn()) + + +def is_remote_url(url_or_filename: str) -> bool: + r"""Return whether a string is a raw URL (``http`` / ``https``). + + Transformers v5 removed this symbol, so it is implemented here with the v4 semantics + (the v4 implementation also checks ``urlparse(...).scheme in ("http", "https")``). + It is retained to support an explanatory error; see :func:`get_download_url`. + + Args: + url_or_filename (str): Path or URL to inspect. + + Returns: + bool: ``True`` when the input is a raw URL. + """ + return urlparse(str(url_or_filename)).scheme in ("http", "https") + + +def get_download_url() -> Optional[Callable[..., str]]: + r"""Return the implementation for downloading weights from a raw URL. + + v4 returns ``transformers.modeling_utils.download_url``. **v5 removed this + capability**: ``download_url`` and ``is_remote_url`` have no remaining definitions, + and ``transformers.utils.hub.cached_file`` now accepts + ``(path_or_repo_id, filename, **kwargs)`` rather than URLs. + + Therefore this returns ``None`` under v5. Callers must handle ``None`` explicitly and + raise an explanatory error instead of silently treating the URL as a repository ID, + which would expose an opaque ``OSError``. + + Returns: + Optional[Callable[..., str]]: The v4 callable, or ``None`` under v5. + """ + try: + from transformers.modeling_utils import download_url as _fn + except ImportError: + return None + return _fn + + +def get_auth_kwarg_name() -> str: + r"""Return the authentication parameter name used by the installed transformers. + + v5 renamed ``use_auth_token`` to ``token``. On v5.17.0, all four relevant entry points + expose ``token`` and no longer expose ``use_auth_token``: + + * ``PretrainedConfig.from_pretrained`` + * ``GenerationConfig.from_pretrained`` + * ``transformers.utils.hub.has_file`` + * ``transformers.utils.hub.cached_file`` + + ``cached_file`` itself accepts ``(path_or_repo_id, filename, **kwargs)`` and forwards + kwargs to ``cached_files``, which accepts ``token``. + + This is required because ``lazy/pretrained.py`` passes the argument explicitly to + ``PretrainedConfig.from_pretrained(..., return_unused_kwargs=True)``. v4 recognizes it + and does not return it as unused; v5 would return it in ``model_kwargs`` as an unknown + argument and eventually raise ``TypeError: ... unexpected keyword argument + 'use_auth_token'``. **Even omitting the argument triggers the failure**, because the + returned value is ``None``. + + Probe ``token`` first. In transitional versions that expose both names, this selects the + new name without triggering a deprecation warning for ``use_auth_token``. + + Returns: + str: ``"token"`` or ``"use_auth_token"``. Falls back to ``"token"`` if signature + inspection fails. + """ + global _AUTH_KWARG_NAME + if _AUTH_KWARG_NAME is None: + try: + from transformers.configuration_utils import PretrainedConfig as _Config + + _params = signature(_Config.from_pretrained).parameters + _AUTH_KWARG_NAME = "token" if "token" in _params else "use_auth_token" + except Exception: + _AUTH_KWARG_NAME = "token" + return _AUTH_KWARG_NAME + + +def auth_kwargs(token: Any = None) -> Dict[str, Any]: + r"""Build authentication kwargs with the name used by the installed transformers. + + Call sites replace ``use_auth_token=use_auth_token`` with:: + + config, model_kwargs = cls.config_class.from_pretrained( + config_path, ..., **auth_kwargs(use_auth_token) + ) + + Args: + token (Any): Authentication token. ``None`` means unauthenticated, but is still + passed explicitly; see :func:`get_auth_kwarg_name`. + + Returns: + Dict[str, Any]: Either ``{"token": token}`` or ``{"use_auth_token": token}``. + """ + return {get_auth_kwarg_name(): token} + + +def is_transformers_v5() -> bool: + r"""Return whether the installed transformers is v5 or later. + + This is a deliberate exception to the feature-detection convention above. It picks a whole + implementation (see :func:`reexport_transformers_impl`), and the boundary between the two is + the v5 architecture change rather than any single symbol: probing e.g. + ``transformers.masking_utils`` would route late v4 releases, which already ship it, to the v5 + implementation. + + Returns: + bool: ``True`` when ``transformers.__version__`` has major version 5 or later. + """ + global _IS_TRANSFORMERS_V5 + if _IS_TRANSFORMERS_V5 is None: + import transformers + + _IS_TRANSFORMERS_V5 = int(transformers.__version__.split(".", 1)[0]) >= 5 + return _IS_TRANSFORMERS_V5 + + +def reexport_transformers_impl(namespace: Dict[str, Any], package: str, stem: str) -> None: + r"""Fill a module namespace with the implementation matching the transformers major version. + + Model families whose v4 and v5 implementations diverge keep them in two sibling modules, + ``_tf4`` (the pinned v4, ``transformers==4.51.3``) and ``_tf5``. The public + module calls this helper so that existing imports of it keep working under both versions. + Every name except dunders is copied, private helpers included, because tests and policies + import some of them. + + The names are references to the implementation module's objects, and its functions resolve + globals in that module. Patching (``mock.patch``, ``monkeypatch.setattr``) must therefore + target ``_tf4`` / ``_tf5``; patching the public module does not change behaviour. + + Args: + namespace (Dict[str, Any]): ``globals()`` of the public module. + package (str): ``__package__`` of the public module. + stem (str): Module name prefix, e.g. ``"_llama"``. + """ + impl = import_module(f"{package}.{stem}_tf5" if is_transformers_v5() else f"{package}.{stem}_tf4") + namespace.update({k: v for k, v in vars(impl).items() if not (k.startswith("__") and k.endswith("__"))}) diff --git a/colossalai/accelerator/cpu_accelerator.py b/colossalai/accelerator/cpu_accelerator.py index 080aa61e8e3a..e87f8ac1a179 100644 --- a/colossalai/accelerator/cpu_accelerator.py +++ b/colossalai/accelerator/cpu_accelerator.py @@ -100,10 +100,8 @@ def utilization(self, device=None) -> int: # random number generator APIs # ======================= def get_rng_state(self, device=None) -> torch.Tensor: - """ - Returns the random number generator state of the specified GPU as a ByteTensor. - """ - return torch.get_rng_state(device) + """Return the CPU generator state; device is retained for API compatibility.""" + return torch.get_rng_state() def get_rng_state_all(self) -> List[torch.Tensor]: """ diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py index 20c3c7e8c469..a4806bf7e55e 100644 --- a/colossalai/booster/plugin/hybrid_parallel_plugin.py +++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py @@ -40,7 +40,7 @@ from colossalai.shardformer.policies.base_policy import Policy from colossalai.shardformer.shard.shard_config import SUPPORT_SP_MODE from colossalai.tensor.colo_parameter import ColoParameter -from colossalai.tensor.d_tensor.api import is_distributed_tensor +from colossalai.tensor.d_tensor.api import is_customized_distributed_tensor, is_distributed_tensor from colossalai.tensor.param_op_hook import ColoParamOpHookManager from colossalai.zero.low_level import LowLevelZeroOptimizer from colossalai.zero.low_level.zero_hook import ZeroOpHook, wait_all_gather_handle @@ -56,6 +56,12 @@ def _convert_floating_point(x, dtype: torch.dtype = torch.float16): return x +def _is_tp_split(param: Tensor) -> bool: + """Whether ``param`` is split across the tp group: a distributed tensor with a layout, or a customized one that + keeps its own shard / gather functions (e.g. fused QKV or fused MoE experts).""" + return is_distributed_tensor(param) or is_customized_distributed_tensor(param) + + class HybridParallelModule(ModelWrapper, AMPModelMixin): def __init__( self, @@ -424,7 +430,7 @@ def _compute_grad_norm(self, param_gradient_pairs: List[Tuple[Tensor]], norm_typ # To ensure mathematical equivalence, we divide the 'grad_norm' by 'tp_size.' if self.tp_size > 1: param_for_grad = grad_to_param_mapping[id(grad)] - if not is_distributed_tensor(param_for_grad): + if not _is_tp_split(param_for_grad): grad_norm_exponentiated /= self.tp_size # If 'pp_size' is greater than 1 and the gradient belongs to shared parameters, @@ -630,7 +636,7 @@ def _compute_grad_norm(self, param_gradient_pairs: List[Tuple[Tensor]], norm_typ # To ensure mathematical equivalence, we divide the 'grad_norm' by 'tp_size.' if self.tp_size > 1: param_for_grad = grad_to_param_mapping[id(grad)] - if not is_distributed_tensor(param_for_grad): + if not _is_tp_split(param_for_grad): grad_norm_exponentiated /= self.tp_size # If 'pp_size' is greater than 1 and the gradient belongs to shared parameters, @@ -854,7 +860,7 @@ def _grad_norm_weight(self, grad: Tensor) -> float: # a parameter not split by tp is replicated on every tp rank: each copy counts 1 / tp_size if self.tp_pg is not None and get_world_size(self.tp_pg) > 1: param_for_grad = ctypes.cast(self.get_param_id_for_grad(grad), ctypes.py_object).value - if not is_distributed_tensor(param_for_grad): + if not _is_tp_split(param_for_grad): weight /= get_world_size(self.tp_pg) # a parameter shared by several pipeline stages (e.g. tied embeddings) counts 1 / num_stages per stage if self.pp_pg is not None and get_world_size(self.pp_pg) > 1: diff --git a/colossalai/checkpoint_io/general_checkpoint_io.py b/colossalai/checkpoint_io/general_checkpoint_io.py index fccc813ccfe0..bda06a338a88 100644 --- a/colossalai/checkpoint_io/general_checkpoint_io.py +++ b/colossalai/checkpoint_io/general_checkpoint_io.py @@ -292,14 +292,24 @@ def load_sharded_model( ckpt_index_file = CheckpointIndexFile.from_file(checkpoint_index_file) checkpoint_files, _ = ckpt_index_file.get_checkpoint_filenames() missing_keys = [] + loaded_keys = set() for state_dict in load_state_dict_shards(checkpoint_files, False, use_safetensors, low_cpu_mem_mode): + loaded_keys.update(state_dict) if not low_cpu_mem_mode: state_dict = create_pinned_state_dict(state_dict, empty=False, num_threads=num_threads) load_state_dict_into_model(model, state_dict, missing_keys, strict, load_sub_module) if strict: remain_keys = reduce(lambda a, b: a & b, map(set, missing_keys)) + # Safetensors may store only one name of a tied Parameter. A + # missing alias is satisfied only by a loaded name of that object. + aliases = {} + for name, parameter in model.named_parameters(remove_duplicate=False): + aliases.setdefault(id(parameter), set()).add(name) + for names in aliases.values(): + if names & loaded_keys: + remain_keys.difference_update(names) if len(remain_keys) > 0: error_msgs = [ "Missing key(s) in state_dict: {}. ".format(", ".join('"{}"'.format(k) for k in remain_keys)) diff --git a/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py b/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py index 9d972635214d..4749a9ab479c 100644 --- a/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py +++ b/colossalai/checkpoint_io/hybrid_parallel_checkpoint_io.py @@ -56,6 +56,53 @@ _EXTRA_STATE_KEY_SUFFIX = "_extra_state" +def _route_padded_parameter_aliases(model: nn.Module, state_dict: Dict[str, torch.Tensor]) -> None: + """Route shared padded weights only through padding-aware module loaders.""" + from colossalai.shardformer.layer.parallel_module import PaddingParallelModule + + modules = dict(model.named_modules(remove_duplicate=False)) + aliases = defaultdict(lambda: {"managed": [], "regular": []}) + for name, param in model.named_parameters(remove_duplicate=False): + if not is_padded_tensor(param): + continue + module_name, _, _ = name.rpartition(".") + module = modules[module_name] + managed = isinstance(module, PaddingParallelModule) + aliases[id(param)]["managed" if managed else "regular"].append(name) + aliases[id(param)]["parameter"] = param + routes = [] + for names in aliases.values(): + if not names["managed"] or not names["regular"]: + continue + present = [name for name in names["managed"] + names["regular"] if name in state_dict] + if not present: + continue + parameter = names["parameter"] + + def logical_value(value): + dim = parameter._padding_dim + if value.shape[dim] == parameter._current_length: + padding = value.narrow( + dim, parameter._origin_length, parameter._current_length - parameter._origin_length + ) + if torch.count_nonzero(padding): + raise ValueError(f"Nonzero padding in checkpoint aliases: {present}") + return value.narrow(dim, 0, parameter._origin_length) + return value + + value = logical_value(state_dict[present[0]]) + for name in present[1:]: + other = logical_value(state_dict[name]) + if value.shape != other.shape or not torch.equal(value, other): + raise ValueError(f"Conflicting checkpoint values for shared padded aliases: {present}") + routes.append((names, value)) + for names, value in routes: + for name in names["managed"]: + state_dict[name] = value + for name in names["regular"]: + state_dict.pop(name, None) + + class HybridParallelCheckpointIO(GeneralCheckpointIO): """ CheckpointIO for Hybrid Parallel Training. @@ -851,6 +898,7 @@ def load_unsharded_model( # has been implemented by _load_from_state_dict method of ParallelModule in Shardformer, # model.load_state_dict can be directly called. state_dict = load_state_dict(checkpoint) + _route_padded_parameter_aliases(model, state_dict) if not low_cpu_mem_mode: state_dict = create_pinned_state_dict(state_dict, empty=False, num_threads=num_threads) model.load_state_dict(state_dict, strict=strict) diff --git a/colossalai/checkpoint_io/moe_checkpoint.py b/colossalai/checkpoint_io/moe_checkpoint.py index 2989146bcedb..a24cd4d638d8 100644 --- a/colossalai/checkpoint_io/moe_checkpoint.py +++ b/colossalai/checkpoint_io/moe_checkpoint.py @@ -17,6 +17,7 @@ from colossalai.checkpoint_io.index_file import CheckpointIndexFile from colossalai.checkpoint_io.utils import ( StateDictSharder, + async_save_state_dict_shards, gather_distributed_param, gather_state_dict_fast, get_lora_state_dict, @@ -33,7 +34,8 @@ sharded_optimizer_loading_epilogue, ) from colossalai.interface import ModelWrapper, OptimizerWrapper -from colossalai.tensor.moe_tensor.api import is_moe_tensor +from colossalai.tensor.d_tensor import is_customized_distributed_tensor +from colossalai.tensor.moe_tensor.api import is_moe_tensor, is_moe_tensor_ep_fused try: from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX @@ -74,11 +76,21 @@ def _model_sharder( keep_vars: bool = False, size_per_shard: int = 1024, param_name_pattern: Optional[str] = None, + pinned_state_dicts: Optional[Dict[str, torch.Tensor]] = None, ) -> Iterator[Tuple[OrderedDict, int]]: # An internel method that breaks state_dict of model into shards within limited size. state_dict_sharder = StateDictSharder(size_per_shard) + def maybe_pin_tensor(name: str, tensor: torch.Tensor) -> torch.Tensor: + if pinned_state_dicts is None or not torch.is_tensor(tensor) or tensor.device.type != "cuda": + return tensor + full_name = prefix + name + if full_name not in pinned_state_dicts: + pinned_state_dicts[full_name] = torch.empty_like(tensor, device="cpu", pin_memory=True) + pinned_state_dicts[full_name].copy_(tensor) + return pinned_state_dicts[full_name] + # Save parameters. for name, param in model.named_parameters(): if param is None: @@ -87,6 +99,21 @@ def _model_sharder( continue # Gather tensor pieces when using tensor parallel. param_ = gather_distributed_param(param, keep_vars=False) + # Fused experts have the same name on every EP rank, while a sharded checkpoint has one + # weight-map entry per name. Gather their EP-local slices to the EP root so every rank can + # reload the same full tensor and let the model's load hook select its local slice. + # Per-expert parameters already have distinct names on each EP rank and are not gathered. + if is_moe_tensor_ep_fused(param): + ep_group = param.ep_group + ep_size = dist.get_world_size(ep_group) + if ep_size > 1: + ep_rank = dist.get_rank(ep_group) + ep_root = get_global_rank(ep_group, 0) + gathered = [torch.empty_like(param_) for _ in range(ep_size)] if ep_rank == 0 else None + dist.gather(param_.contiguous(), gather_list=gathered, dst=ep_root, group=ep_group) + if ep_rank == 0: + param_ = torch.cat(gathered, dim=0) + param_ = maybe_pin_tensor(name, param_) block, block_size = state_dict_sharder.append_param(prefix + name, param_) if block is not None: yield block, block_size @@ -95,6 +122,7 @@ def _model_sharder( for name, buf in model.named_buffers(): if buf is not None and name not in model._non_persistent_buffers_set: buffer = buf if keep_vars else buf.detach() + buffer = maybe_pin_tensor(name, buffer) block, block_size = state_dict_sharder.append_param(prefix + name, buffer) if block is not None: yield block, block_size @@ -106,6 +134,7 @@ def _model_sharder( is not torch.nn.Module.get_extra_state ): extra_state = model.get_extra_state() + extra_state = maybe_pin_tensor(_EXTRA_STATE_KEY_SUFFIX, extra_state) block, block_size = state_dict_sharder.append_param(extra_state_key, extra_state) if block is not None: yield block, block_size @@ -154,20 +183,33 @@ def save_sharded_model( dist.barrier() return - # ep_rank 0 saves all the parameters and buffers. - # other ep_ranks save only experts + # Fused experts are gathered across EP by `_model_sharder`, so only ep_rank 0 writes them. + # Otherwise ep_rank 0 saves all the parameters and buffers, and other ep_ranks save their own experts. + gathers_fused_experts = self.ep_size > 1 and any(is_moe_tensor_ep_fused(p) for p in model.parameters()) # Then collect the sharded parameters & buffers along tp_group. # Only devices with tp_rank == 0 are responsible for model saving. - state_dict_shard = MoECheckpointIO._model_sharder(model, size_per_shard=size_per_shard) weights_name, save_index_file = get_model_base_filenames(prefix, use_safetensors) index_file = CheckpointIndexFile(checkpoint) - control_saving = self.tp_rank == 0 and self.sp_rank == 0 + control_saving = self.tp_rank == 0 and self.sp_rank == 0 and (self.ep_rank == 0 or not gathers_fused_experts) + pinned_state_dicts = None + if self.pp_size == 1 and self.ep_size == 1 and use_async and control_saving: + pinned_state_dicts = self.pinned_state_dicts.setdefault(hash(model), {}) + state_dict_shard = MoECheckpointIO._model_sharder( + model, size_per_shard=size_per_shard, pinned_state_dicts=pinned_state_dicts + ) if self.pp_size == 1 and self.ep_size == 1: # When pipeline is not used, save the model shards as in general checkpointIO if use_async: - super().save_unsharded_model(model, checkpoint, gather_dtensor, use_safetensors, use_async=use_async) + total_size, writers = async_save_state_dict_shards( + sharded_state_dict=state_dict_shard, + checkpoint=checkpoint, + index_file=index_file, + base_filename=weights_name, + is_master=control_saving, + ) + self.async_writers.extend(writers) else: total_size = save_state_dict_shards( sharded_state_dict=state_dict_shard, @@ -177,21 +219,21 @@ def save_sharded_model( is_master=control_saving, use_safetensors=use_safetensors, ) - if control_saving: - index_file.append_meta_data("total_size", total_size) - index_file.write_index_file(save_index_file) - save_config_file(model, checkpoint) - if self.verbose and self.coordinator.is_master(): - logging.info( - f"The model is split into checkpoint shards. " - f"You can find where each parameters has been saved in the " - f"index located at {save_index_file}." - ) + if control_saving: + index_file.append_meta_data("total_size", total_size) + index_file.write_index_file(save_index_file) + save_config_file(model, checkpoint) + if self.verbose and self.coordinator.is_master(): + logging.info( + f"The model is split into checkpoint shards. " + f"You can find where each parameters has been saved in the " + f"index located at {save_index_file}." + ) dist.barrier() else: - # When pipeline is used, each stage produces its own shard files and index files. - # Index files belonging to each stage are saved under a temporary folder ./tmp_index_files/ + # When pipeline or expert parallelism is used, each stage (and each EP rank that saves) produces + # its own shard files and index files under a temporary folder ./tmp_index_files/. # After all the state_dicts have been saved, the master rank integrates all the index files into one final index file and deletes the tmp folder. final_index_file_path = copy.deepcopy(save_index_file) @@ -282,6 +324,9 @@ def gather_from_sharded_optimizer_state( for k, v in state_.items(): if isinstance(v, torch.Tensor) and k != "step": v = v.cuda() + # Only the root of the Zero group gets the complete state; the other ranks keep a flat shard. + # Ranks in one TP / EP group share their Zero group position, so they skip the gathers below together. + holds_full_state = True # First gather Zero shards. if use_zero and is_moe_param and moe_dp_size > 1: @@ -293,6 +338,7 @@ def gather_from_sharded_optimizer_state( v = torch.stack(gather_tensor).view(-1)[: param.numel()].reshape_as(param) else: dist.gather(v, group=moe_dp_group, dst=dst) + holds_full_state = False elif use_zero and not is_moe_param and global_dp_size > 1: dp_rank = dist.get_rank(global_dp_group) @@ -303,6 +349,19 @@ def gather_from_sharded_optimizer_state( v = torch.stack(gather_tensor).view(-1)[: param.numel()].reshape_as(param) else: dist.gather(v, group=global_dp_group, dst=dst) + holds_full_state = False + + if not holds_full_state: + # This rank does not save the state; a flat shard cannot be concatenated along a TP dimension. + state_[k] = v.detach().clone().to(device) + continue + + if is_moe_tensor_ep_fused(param): + # Fused experts are sharded along dim 0 by EP (and possibly along other dims by TP), + # which `search_tp_partition_dim` cannot tell apart. Gather them like the model weights. + v = MoECheckpointIO._gather_fused_expert_state(v, param) + state_[k] = v.detach().clone().to(device) + continue # Then gather TP shards. partition_dim = search_tp_partition_dim(current_shape, original_shape, tp_size) @@ -358,6 +417,9 @@ def _optimizer_sharder( if only_moe_param and not is_moe_tensor(working_param): continue + # Fused expert states were gathered to the EP root, which saves them. + if is_moe_tensor_ep_fused(working_param) and dist.get_rank(working_param.ep_group) != 0: + continue block, block_size = state_dict_sharder.append_optim_state(param_id, state_) if block is not None: @@ -614,6 +676,8 @@ def _get_param_id_from_optimizer_param( device=device, inplace=True, is_moe_param=is_moe_tensor(working_param), + working_param=working_param, + master_param=param, ) optimizer.optim.state[param] = sharded_state @@ -621,6 +685,52 @@ def _get_param_id_from_optimizer_param( if self.verbose and self.coordinator.is_master(): logging.info(f"The optimizer has been successfully loaded from sharded checkpoint: {ckpt_root_path}.") + @staticmethod + def _gather_fused_expert_state(state: torch.Tensor, param: torch.Tensor) -> torch.Tensor: + """ + Gather the optimizer state of a fused expert parameter into the full ``[num_experts, ...]`` layout + on the EP root. ``state`` holds this rank's complete (Zero-gathered) state, possibly flattened. + """ + state = state.reshape(-1)[: param.numel()].reshape(param.shape) + # Expert TP slices use the same gather function as the model checkpoint. + if is_customized_distributed_tensor(param): + state = param.gather_fn(state) + ep_group = param.ep_group + ep_size = dist.get_world_size(ep_group) + if ep_size > 1: + ep_rank = dist.get_rank(ep_group) + gathered = [torch.empty_like(state) for _ in range(ep_size)] if ep_rank == 0 else None + dist.gather(state.contiguous(), gather_list=gathered, dst=get_global_rank(ep_group, 0), group=ep_group) + if ep_rank == 0: + state = torch.cat(gathered, dim=0) + return state + + def _shard_fused_expert_state( + self, state: torch.Tensor, param: torch.Tensor, master_param: Optional[torch.Tensor] + ) -> torch.Tensor: + """ + Slice this rank's part out of the full ``[num_experts, ...]`` optimizer state of a fused expert parameter. + """ + if is_customized_distributed_tensor(param): + # Selects this rank's EP experts and TP columns, as when loading the model weights. + state = param.shard_fn(state) + else: + num_local_experts = param.shape[0] + start = dist.get_rank(param.ep_group) * num_local_experts + state = state[start : start + num_local_experts] + + if self.use_zero and self.moe_dp_size > 1: + # LowLevelZeRO pads each param by the size of the group it is sharded over + padding_size = (self.moe_dp_size - state.numel() % self.moe_dp_size) % self.moe_dp_size + state = state.flatten() + if padding_size > 0: + state = torch.nn.functional.pad(state, [0, padding_size]) + state = state.split(state.numel() // self.moe_dp_size, dim=0)[self.moe_dp_rank] + elif master_param is not None and state.numel() == master_param.numel(): + # Zero keeps flattened master parameters. + state = state.reshape(master_param.shape) + return state + def shard_from_complete_optimizer_state( self, state: OrderedDict, @@ -629,6 +739,8 @@ def shard_from_complete_optimizer_state( device: torch.device, inplace: bool, is_moe_param: bool, + working_param: Optional[torch.Tensor] = None, + master_param: Optional[torch.Tensor] = None, ) -> OrderedDict: """ With complete optimizer states of a specific parameter loaded from checkpoint, @@ -640,6 +752,9 @@ def shard_from_complete_optimizer_state( original_shape (torch.Size): The size of parameter before sharding. device (torch.device): The destination device of loaded optimizer states. inplace (bool): If set to True, will update the values of argument 'state' in place. Else will make a copy of state. + working_param (torch.Tensor, optional): The model parameter; needed to shard fused expert states. + master_param (torch.Tensor, optional): The parameter the optimizer keeps states for (the fp32 master + parameter when using Zero); fused expert states are reshaped to it. Returns: OrderedDict: The sharded optimizer state of the given parameter. @@ -647,6 +762,11 @@ def shard_from_complete_optimizer_state( state_ = state if inplace else copy.deepcopy(state) for k, v in state_.items(): if isinstance(v, torch.Tensor) and k != "step": + if working_param is not None and is_moe_tensor_ep_fused(working_param): + v = self._shard_fused_expert_state(v, working_param, master_param) + state_[k] = v.detach().clone().to(device) + continue + # Shard state along tensor parallel group. partition_dim = search_tp_partition_dim(current_shape, original_shape, self.tp_size) if partition_dim is not None: diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py index 6b838eb1ce4a..58aba87537dc 100644 --- a/colossalai/checkpoint_io/utils.py +++ b/colossalai/checkpoint_io/utils.py @@ -568,8 +568,10 @@ def save_config_file(model: nn.Module, checkpoint_path: str, is_master: bool = T checkpoint_path (str): Path to the checkpoint directory. is_master (bool): Whether current rank is main process. """ + # `get_parameter_dtype` was removed in Transformers v5; importing it here made this function return + # silently, so sharded checkpoints lost config.json. `PreTrainedModel.dtype` exists in v4 and v5. try: - from transformers.modeling_utils import PreTrainedModel, get_parameter_dtype + from transformers.modeling_utils import PreTrainedModel from transformers.modeling_utils import unwrap_model as unwrap_huggingface_model except ImportError: return @@ -581,8 +583,12 @@ def save_config_file(model: nn.Module, checkpoint_path: str, is_master: bool = T model = unwrap_huggingface_model(model) # save the string version of dtype to the config, e.g. convert torch.float32 => "float32" - dtype = get_parameter_dtype(model) - model.config.torch_dtype = str(dtype).split(".")[1] + dtype_name = str(model.dtype).split(".")[1] + # v5 renamed `config.torch_dtype` to `config.dtype` and keeps the old name as a deprecated property. + if isinstance(getattr(type(model.config), "torch_dtype", None), property): + model.config.dtype = dtype_name + else: + model.config.torch_dtype = dtype_name # Attach architecture to the config model.config.architectures = [model.__class__.__name__] diff --git a/colossalai/lazy/lazy_init.py b/colossalai/lazy/lazy_init.py index cc3137f2df3a..466fd6cac88f 100644 --- a/colossalai/lazy/lazy_init.py +++ b/colossalai/lazy/lazy_init.py @@ -15,6 +15,7 @@ import colossalai._analyzer._subclasses._meta_registration # noqa + # reference: https://pytorch.org/cppdocs/notes/tensor_creation.html _NORMAL_FACTORY = [ "arange", @@ -313,6 +314,9 @@ def unwrap(x): return target else: meta_to_lazy = {} + recorded_inplace = set() + # torch.nn.init forwards its tensor as a keyword argument. + inplace_target = args[0] if args else kwargs.get("tensor", kwargs.get("input", kwargs.get("self"))) def unwrap(x): if isinstance(x, LazyTensor): @@ -320,8 +324,14 @@ def unwrap(x): # for early materialized tensor, use its materialized data directly return x._materialized_data if is_change_meta_op else x._materialized_data.data t = x if is_inplace else x.clone() - if func.__name__ not in _NO_RERUN_OPS: + # An in-place operation mutates its first argument only. + # Recording copy_(destination, source) on source as well + # creates a cycle when either tensor is materialized. + record_op = not is_inplace or (x is inplace_target and id(x) not in recorded_inplace) + if func.__name__ not in _NO_RERUN_OPS and record_op: t._op_buffer.append((func, args, kwargs)) + if is_inplace: + recorded_inplace.add(id(x)) meta = x._meta_data if is_change_meta_op else x._meta_data.data meta_to_lazy[meta] = t return meta @@ -520,25 +530,48 @@ def wrapper(*args, **kwargs): def wrap_legacy_constructor(target, dtype): # legacy constructor (e.g. torch.LongTensor()) def wrapper(*args, **kwargs): - if len(args) == 1 and isinstance(args[0], torch.Tensor): + if _is_int_tuple(args): + # (tuple of ints size, *, torch.device device) + kwargs = {**kwargs, "dtype": dtype} + replaced, orig = overrides["empty"] + return replaced(*(args or (0,)), **kwargs) + elif len(args) == 1 and isinstance(args[0], torch.Tensor): # (Tensor other) return args[0] elif len(args) == 1: # (object data, *, torch.device device) kwargs = {**kwargs, "dtype": dtype} - replaced, orig = self.overrides["tensor"] - return replaced(*args, **kwargs) - elif _is_int_tuple(args): - # (tuple of ints size, *, torch.device device) - kwargs = {**kwargs, "dtype": dtype} - replaced, orig = self.overrides["empty"] + replaced, orig = overrides["tensor"] return replaced(*args, **kwargs) else: raise TypeError( f"new() received an invalid combination of arguments - got {tuple(type(x) for x in args)}, but expected one of:\n * (Tensor other)\n * (tuple of ints size, *, torch.device device)\n * (object data, *, torch.device device)" ) - return wrapper, target + # HF evaluates annotations such as torch.LongTensor | None while + # importing modeling modules inside this context. Keep a type-like + # constructor instead of replacing the class with a function. + class LegacyTensorMeta(type): + def __call__(cls, *args, **kwargs): + return wrapper(*args, **kwargs) + + def __instancecheck__(cls, instance): + return isinstance(instance, target) + + def __subclasscheck__(cls, subclass): + return issubclass(subclass, target) + + def __or__(cls, other): + return target | other + + def __ror__(cls, other): + return other | target + + def __getattr__(cls, name): + return getattr(target, name) + + proxy = LegacyTensorMeta(target.__name__, (), {"__module__": target.__module__}) + return proxy, target def wrap_no_meta_factory(target): # factory functions which don't support meta tensor backend diff --git a/colossalai/lazy/pretrained.py b/colossalai/lazy/pretrained.py index 684be922332b..a087aa988e99 100644 --- a/colossalai/lazy/pretrained.py +++ b/colossalai/lazy/pretrained.py @@ -38,25 +38,23 @@ def new_from_pretrained( ) -> Module: from transformers import GenerationConfig from transformers.configuration_utils import PretrainedConfig - from transformers.modeling_utils import ( - ContextManagers, - _add_variant, - cached_file, - download_url, - has_file, - is_offline_mode, + from transformers.modeling_utils import ContextManagers, _add_variant, cached_file, has_file, is_offline_mode + from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME, logging + + from colossalai._compat import ( + auth_kwargs, + get_download_url, + get_no_init_weights, is_remote_url, - no_init_weights, - ) - from transformers.utils import ( - SAFE_WEIGHTS_INDEX_NAME, - SAFE_WEIGHTS_NAME, - WEIGHTS_INDEX_NAME, - WEIGHTS_NAME, is_safetensors_available, - logging, ) + # Transformers v5 compatibility helpers (see `colossalai/_compat.py`): + # `no_init_weights` moved to `transformers.initialization`; `download_url` was removed + # entirely and therefore returns `None`, which the raw-URL branch handles explicitly. + no_init_weights = get_no_init_weights() + download_url = get_download_url() + logger = logging.get_logger(__name__) config = kwargs.pop("config", None) @@ -64,13 +62,26 @@ def new_from_pretrained( force_download = kwargs.pop("force_download", False) proxies = kwargs.pop("proxies", None) local_files_only = kwargs.pop("local_files_only", False) + # v4 called this argument `use_auth_token`; v5 renamed it to `token` (see + # `colossalai/_compat.py`). Accept both names so an explicit v5 `token` does not + # collide with the keyword passed to the config loader below. use_auth_token = kwargs.pop("use_auth_token", None) + token = kwargs.pop("token", None) + if use_auth_token is not None and token is not None and use_auth_token != token: + raise ValueError("Conflicting token and use_auth_token arguments") + if token is not None: + use_auth_token = token revision = kwargs.pop("revision", None) _ = kwargs.pop("mirror", None) from_pipeline = kwargs.pop("_from_pipeline", None) from_auto_class = kwargs.pop("_from_auto", False) kwargs.pop("_fast_init", True) torch_dtype = kwargs.pop("torch_dtype", None) + dtype = kwargs.pop("dtype", None) + if dtype is not None and torch_dtype is not None and dtype != torch_dtype: + raise ValueError("Conflicting dtype and torch_dtype arguments") + if dtype is not None: + torch_dtype = dtype subfolder = kwargs.pop("subfolder", "") commit_hash = kwargs.pop("_commit_hash", None) variant = kwargs.pop("variant", None) @@ -117,7 +128,7 @@ def new_from_pretrained( force_download=force_download, proxies=proxies, local_files_only=local_files_only, - use_auth_token=use_auth_token, + **auth_kwargs(use_auth_token), revision=revision, subfolder=subfolder, _from_auto=from_auto_class, @@ -178,6 +189,16 @@ def new_from_pretrained( archive_file = pretrained_model_name_or_path is_local = True elif is_remote_url(pretrained_model_name_or_path): + if download_url is None: + # Transformers v5 removed raw-URL loading: `download_url` and + # `is_remote_url` no longer exist in the installed tree, and `cached_file` + # no longer accepts URLs. Fail explicitly instead of treating the URL as a + # repository id and exposing an opaque cache lookup error. + raise NotImplementedError( + "Transformers v5 no longer supports loading weights from a raw URL; received " + f"{pretrained_model_name_or_path!r}. Download the weights to a local directory " + "first, or use a Hugging Face repository ID." + ) filename = pretrained_model_name_or_path resolved_archive_file = download_url(pretrained_model_name_or_path) else: @@ -194,7 +215,7 @@ def new_from_pretrained( "force_download": force_download, "proxies": proxies, "local_files_only": local_files_only, - "use_auth_token": use_auth_token, + **auth_kwargs(use_auth_token), "user_agent": user_agent, "revision": revision, "subfolder": subfolder, @@ -239,7 +260,7 @@ def new_from_pretrained( has_file_kwargs = { "revision": revision, "proxies": proxies, - "use_auth_token": use_auth_token, + **auth_kwargs(use_auth_token), } if variant is not None and has_file(pretrained_model_name_or_path, WEIGHTS_NAME, **has_file_kwargs): raise EnvironmentError( @@ -281,19 +302,16 @@ def new_from_pretrained( if torch_dtype is not None: if not isinstance(torch_dtype, torch.dtype): raise ValueError(f"`torch_dtype` can be either `torch.dtype` or `None`, but received {torch_dtype}") - dtype_orig = cls._set_default_torch_dtype(torch_dtype) - - config.name_or_path = pretrained_model_name_or_path - - # Instantiate model. - init_contexts = [no_init_weights()] - - with ContextManagers(init_contexts): - model = cls(config, *model_args, **model_kwargs) - - if from_pt: - # restore default dtype - if dtype_orig is not None: + dtype_orig = torch.get_default_dtype() + torch.set_default_dtype(torch_dtype) + + try: + config.name_or_path = pretrained_model_name_or_path + init_contexts = [no_init_weights()] + with ContextManagers(init_contexts): + model = cls(config, *model_args, **model_kwargs) + finally: + if from_pt and dtype_orig is not None: torch.set_default_dtype(dtype_orig) # make sure token embedding weights are still tied if needed @@ -311,7 +329,7 @@ def new_from_pretrained( force_download=force_download, proxies=proxies, local_files_only=local_files_only, - use_auth_token=use_auth_token, + **auth_kwargs(use_auth_token), revision=revision, subfolder=subfolder, _from_auto=from_auto_class, diff --git a/colossalai/pipeline/schedule/_utils.py b/colossalai/pipeline/schedule/_utils.py index ded75d9684a6..232669bf0091 100644 --- a/colossalai/pipeline/schedule/_utils.py +++ b/colossalai/pipeline/schedule/_utils.py @@ -53,6 +53,35 @@ def tree_flatten_hf(pytree: Any) -> Tuple[List[Any], TreeSpec]: return result, tree_spec +def collect_gradients(pytree: Any) -> Any: + """Collect gradients from every tensor leaf in a pipeline pytree. + + Pipeline stage outputs may contain nested tensor structures, such as the tuple of + Mixtral router logits. Returning ``None`` when a pytree has no tensor gradient lets + callers preserve the existing convention of omitting non-differentiable outputs. + """ + + gradients = tree_map( + lambda value: value.grad if isinstance(value, torch.Tensor) and value.grad is not None else None, + pytree, + ) + return gradients if any(isinstance(value, torch.Tensor) for value in tree_flatten(gradients)[0]) else None + + +def flatten_tensors_and_grads(tensors: Any, gradients: Any) -> List[Tuple[torch.Tensor, torch.Tensor]]: + """Flatten matching tensor/gradient pytrees for ``autograd.backward``.""" + + tensor_leaves, tensor_spec = tree_flatten(tensors) + gradient_leaves, gradient_spec = tree_flatten(gradients) + if tensor_spec != gradient_spec: + raise ValueError("Pipeline tensor and gradient structures must match") + return [ + (tensor, gradient) + for tensor, gradient in zip(tensor_leaves, gradient_leaves) + if isinstance(tensor, torch.Tensor) and isinstance(gradient, torch.Tensor) + ] + + def to_device(x: Any, device: Optional[torch.device] = None) -> Any: """Move object to device if it is a tensor. diff --git a/colossalai/pipeline/schedule/interleaved_pp.py b/colossalai/pipeline/schedule/interleaved_pp.py index eeab080af328..9c3ce5204515 100644 --- a/colossalai/pipeline/schedule/interleaved_pp.py +++ b/colossalai/pipeline/schedule/interleaved_pp.py @@ -13,7 +13,17 @@ from colossalai.quantization.fp8 import cast_from_fp8_pipeline, cast_to_fp8_pipeline from colossalai.utils import get_current_device -from ._utils import detach, get_batch_size, get_micro_batch, merge_batch, model_forward, retain_grad, to_device +from ._utils import ( + collect_gradients, + detach, + flatten_tensors_and_grads, + get_batch_size, + get_micro_batch, + merge_batch, + model_forward, + retain_grad, + to_device, +) from .base import PipelineSchedule @@ -50,12 +60,12 @@ def __init__( self.last_batch_size: Optional[int] = None self.microbatch_offset: List[int] - # P2PMeta cache + # P2P metadata is chunk-local because virtual chunks may send different pytrees. self.enable_metadata_cache = enable_metadata_cache - self.send_tensor_metadata = True - self.send_grad_metadata = True - self.tensor_metadata_recv = None - self.grad_metadata_recv = None + self.send_tensor_metadata = [True] * num_model_chunks + self.send_grad_metadata = [True] * num_model_chunks + self.tensor_metadata_recv = [None] * num_model_chunks + self.grad_metadata_recv = [None] * num_model_chunks self.fp8_communication = fp8_communication @@ -94,10 +104,10 @@ def load_batch(self, data_iter: Iterable, device: Optional[torch.device] = None) # NOTE: disable metadata cache when batch size changes (not valid anymore) if self.batch_size != self.last_batch_size: self.enable_metadata_cache = False - self.send_tensor_metadata = True - self.send_grad_metadata = True - self.tensor_metadata_recv = None - self.grad_metadata_recv = None + self.send_tensor_metadata = [True] * self.num_model_chunks + self.send_grad_metadata = [True] * self.num_model_chunks + self.tensor_metadata_recv = [None] * self.num_model_chunks + self.grad_metadata_recv = [None] * self.num_model_chunks self.last_batch_size = self.batch_size @@ -149,10 +159,12 @@ def recv_forward(self, model_chunk_id: int, prev_rank: int = None) -> Tuple[Any, """ with self.stage_manager.switch_model_chunk_id(model_chunk_id): if not self.stage_manager.is_first_stage(): - input_tensor, wait_handles = self.comm.recv_forward(prev_rank, metadata_recv=self.tensor_metadata_recv) + input_tensor, wait_handles = self.comm.recv_forward( + prev_rank, metadata_recv=self.tensor_metadata_recv[model_chunk_id] + ) - if self.enable_metadata_cache and self.tensor_metadata_recv is None: - self.tensor_metadata_recv = create_send_metadata(input_tensor) + if self.enable_metadata_cache and self.tensor_metadata_recv[model_chunk_id] is None: + self.tensor_metadata_recv[model_chunk_id] = create_send_metadata(input_tensor) return input_tensor, wait_handles return None, [] @@ -172,10 +184,10 @@ def recv_backward(self, model_chunk_id: int, next_rank: int = None) -> Tuple[Any with self.stage_manager.switch_model_chunk_id(model_chunk_id): if not self.stage_manager.is_last_stage(): output_tensor_grad, wait_handles = self.comm.recv_backward( - next_rank, metadata_recv=self.grad_metadata_recv + next_rank, metadata_recv=self.grad_metadata_recv[model_chunk_id] ) - if self.enable_metadata_cache and self.grad_metadata_recv is None: - self.grad_metadata_recv = create_send_metadata(output_tensor_grad) + if self.enable_metadata_cache and self.grad_metadata_recv[model_chunk_id] is None: + self.grad_metadata_recv[model_chunk_id] = create_send_metadata(output_tensor_grad) return output_tensor_grad, wait_handles return None, [] @@ -196,8 +208,10 @@ def send_forward(self, model_chunk_id: int, output_tensor: Any, next_rank: int = if not self.stage_manager.is_last_stage(): if self.fp8_communication: cast_to_fp8_pipeline(output_tensor) - send_handles = self.comm.send_forward(output_tensor, next_rank, send_metadata=self.send_tensor_metadata) - self.send_tensor_metadata = not self.enable_metadata_cache + send_handles = self.comm.send_forward( + output_tensor, next_rank, send_metadata=self.send_tensor_metadata[model_chunk_id] + ) + self.send_tensor_metadata[model_chunk_id] = not self.enable_metadata_cache if self.fp8_communication: cast_from_fp8_pipeline(output_tensor) return send_handles @@ -220,9 +234,9 @@ def send_backward(self, model_chunk_id: int, input_tensor_grad: Any, prev_rank: if self.fp8_communication: cast_to_fp8_pipeline(input_tensor_grad) send_handles = self.comm.send_backward( - input_tensor_grad, prev_rank, send_metadata=self.send_grad_metadata + input_tensor_grad, prev_rank, send_metadata=self.send_grad_metadata[model_chunk_id] ) - self.send_grad_metadata = not self.enable_metadata_cache + self.send_grad_metadata[model_chunk_id] = not self.enable_metadata_cache if self.fp8_communication: cast_from_fp8_pipeline(input_tensor_grad) return send_handles @@ -241,14 +255,15 @@ def send_forward_recv_forward( output_tensor, is_send, is_recv, - send_metadata=self.send_tensor_metadata, - metadata_recv=self.tensor_metadata_recv, + send_metadata=self.send_tensor_metadata[model_chunk_id_send], + metadata_recv=self.tensor_metadata_recv[model_chunk_id_recv], send_first=send_first, ) # Cache metadata - self.send_tensor_metadata = not self.enable_metadata_cache and is_send - if is_recv and self.enable_metadata_cache and self.tensor_metadata_recv is None: - self.tensor_metadata_recv = create_send_metadata(input_tensor) + if is_send: + self.send_tensor_metadata[model_chunk_id_send] = not self.enable_metadata_cache + if is_recv and self.enable_metadata_cache and self.tensor_metadata_recv[model_chunk_id_recv] is None: + self.tensor_metadata_recv[model_chunk_id_recv] = create_send_metadata(input_tensor) if self.fp8_communication: cast_from_fp8_pipeline(output_tensor) @@ -267,14 +282,15 @@ def send_backward_recv_backward( input_tensor_grad, is_send, is_recv, - send_metadata=self.send_grad_metadata, - metadata_recv=self.grad_metadata_recv, + send_metadata=self.send_grad_metadata[model_chunk_id_send], + metadata_recv=self.grad_metadata_recv[model_chunk_id_recv], send_first=send_first, ) # Cache metadata - self.send_grad_metadata = not self.enable_metadata_cache and is_send - if is_recv and self.enable_metadata_cache and self.grad_metadata_recv is None: - self.grad_metadata_recv = create_send_metadata(output_tensor_grad) + if is_send: + self.send_grad_metadata[model_chunk_id_send] = not self.enable_metadata_cache + if is_recv and self.enable_metadata_cache and self.grad_metadata_recv[model_chunk_id_recv] is None: + self.grad_metadata_recv[model_chunk_id_recv] = create_send_metadata(output_tensor_grad) if self.fp8_communication: cast_from_fp8_pipeline(input_tensor_grad) return output_tensor_grad, wait_handles @@ -354,8 +370,13 @@ def backward_step( tensors_to_backward = [] grads_to_backward = [] for k in keys: - tensors_to_backward.append(output_obj[k]) - grads_to_backward.append(output_obj_grad[k]) + if k not in output_obj_grad: + continue + for tensor, grad in flatten_tensors_and_grads(output_obj[k], output_obj_grad[k]): + tensors_to_backward.append(tensor) + grads_to_backward.append(grad) + if not tensors_to_backward: + return None if len(tensors_to_backward) == 1: optimizer.backward_by_grad(tensors_to_backward[0], grads_to_backward[0]) else: @@ -366,8 +387,9 @@ def backward_step( if input_obj is not None: input_obj_grad = {} for k, v in input_obj.items(): - if isinstance(v, torch.Tensor) and v.grad is not None: - input_obj_grad[k] = v.grad + gradients = collect_gradients(v) + if gradients is not None: + input_obj_grad[k] = gradients return input_obj_grad def run_forward_only( diff --git a/colossalai/pipeline/schedule/one_f_one_b.py b/colossalai/pipeline/schedule/one_f_one_b.py index dcffa858c5c4..2c7aeb334e24 100644 --- a/colossalai/pipeline/schedule/one_f_one_b.py +++ b/colossalai/pipeline/schedule/one_f_one_b.py @@ -13,7 +13,9 @@ from colossalai.utils import get_current_device from ._utils import ( + collect_gradients, detach, + flatten_tensors_and_grads, get_batch_size, get_micro_batch, merge_batch, @@ -308,8 +310,13 @@ def backward_step( tensors_to_backward = [] grads_to_backward = [] for k in keys: - tensors_to_backward.append(output_obj[k]) - grads_to_backward.append(output_obj_grad[k]) + if k not in output_obj_grad: + continue + for tensor, grad in flatten_tensors_and_grads(output_obj[k], output_obj_grad[k]): + tensors_to_backward.append(tensor) + grads_to_backward.append(grad) + if not tensors_to_backward: + return None if len(tensors_to_backward) == 1: optimizer.backward_by_grad(tensors_to_backward[0], grads_to_backward[0]) else: @@ -320,8 +327,9 @@ def backward_step( if input_obj is not None: input_obj_grad = {} for k, v in input_obj.items(): - if isinstance(v, torch.Tensor) and v.grad is not None: - input_obj_grad[k] = v.grad + gradients = collect_gradients(v) + if gradients is not None: + input_obj_grad[k] = gradients return input_obj_grad def run_forward_only( diff --git a/colossalai/pipeline/schedule/zero_bubble_pp.py b/colossalai/pipeline/schedule/zero_bubble_pp.py index edbb7118aa1a..045c86c9f57d 100644 --- a/colossalai/pipeline/schedule/zero_bubble_pp.py +++ b/colossalai/pipeline/schedule/zero_bubble_pp.py @@ -16,6 +16,7 @@ from ._utils import ( clone, + collect_gradients, detach, get_batch_size, get_micro_batch, @@ -584,8 +585,9 @@ def backward_b_step( pass else: for k, v in input_obj.items(): - if isinstance(v, torch.Tensor) and v.grad is not None: - input_obj_grad[k] = v.grad + gradients = collect_gradients(v) + if gradients is not None: + input_obj_grad[k] = gradients return input_obj_grad def backward_w_step( diff --git a/colossalai/shardformer/_utils.py b/colossalai/shardformer/_utils.py index e956279ba11b..3303bb910ea2 100644 --- a/colossalai/shardformer/_utils.py +++ b/colossalai/shardformer/_utils.py @@ -1,6 +1,189 @@ import re +def is_hf5(): + from packaging.version import Version + from transformers import __version__ + + return Version(__version__).major >= 5 + + +def hf5_policy_error(family, shard_config=None): + """Return a declared HF5 support-boundary error, not a compatibility guess.""" + if not is_hf5(): + return None + reasons = { + "chatglm2": "remote-code revision and HF5 compatibility have not been validated", + "deepseek": "legacy remote-code policy is not validated for HF5 native Deepseek models", + } + if family in reasons: + reason = reasons[family] + elif family in {"opt", "command", "blip2"} and shard_config is not None: + from packaging.version import Version + from transformers import __version__ + + if Version(__version__) < Version("5.17.0"): + return ( + f"ShardFormer {family} HF5 compatibility requires transformers>=5.17.0; installed {__version__}. " + "See docs/HF5-COMPATIBILITY.md." + ) + unsupported = [ + name + for name in ( + "enable_tensor_parallelism", + "enable_sequence_parallelism", + "enable_flash_attention", + "enable_jit_fused", + "enable_fused_normalization", + "enable_all_optimization", + "fp8_communication", + ) + if getattr(shard_config, name, False) + ] + if getattr(shard_config, "pipeline_stage_manager", None) is not None: + unsupported.append("pipeline_parallelism") + if not unsupported: + return None + reason = ( + "HF5 migration is currently validated only without parallel/optional optimization; requested " + + ", ".join(unsupported) + ) + else: + return None + from transformers import __version__ + + return ( + f"ShardFormer {family} policy does not support this configuration with transformers {__version__}: " + f"{reason}. See docs/HF5-COMPATIBILITY.md. Use the documented validated configuration, " + "or a separate legacy environment following the repository's transformers==4.51.3 declaration " + "(not a guarantee that every model/configuration is validated)." + ) + + +def check_hf5_policy(family, shard_config=None): + error = hf5_policy_error(family, shard_config) + if error is not None: + raise NotImplementedError(error) + + +def get_hf_mask_config(config, *, eager=False): + """Select an explicit mask backend without mutating the model's shared config.""" + from copy import copy + + if not eager: + return config + config = copy(config) + config._attn_implementation = "eager" + return config + + +def prepare_hf_bidirectional_mask(config, inputs_embeds, attention_mask=None, encoder_hidden_states=None): + """Build an explicit additive mask for vision or encoder cross-attention.""" + from transformers.masking_utils import create_bidirectional_mask + + return create_bidirectional_mask( + config=get_hf_mask_config(config, eager=True), + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + +def prepare_hf_encoder_decoder_cache(config, past_key_values=None, *, encoder_hidden_states=None): + """Allocate decoder-sized self/cross caches, preserving an existing Cache object. + + Whisper maps num_hidden_layers to encoder_layers. Override it on a private + copy so distilled models allocate only the actual number of decoder slots. + """ + from copy import deepcopy + + from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache + + if past_key_values is not None and not isinstance(past_key_values, Cache): + raise TypeError("HF 5.17 requires a Cache object, not legacy tuple past_key_values") + needs_cross = encoder_hidden_states is not None or config.is_encoder_decoder + if needs_cross and isinstance(past_key_values, DynamicCache): + if past_key_values.get_seq_length() != 0: + raise ValueError("Use EncoderDecoderCache for nonempty encoder-decoder past_key_values") + past_key_values = None + if past_key_values is None: + cache_config = deepcopy(config) + cache_config.num_hidden_layers = config.decoder_layers + past_key_values = DynamicCache(config=cache_config) + if needs_cross: + past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=cache_config)) + return past_key_values + + +def get_hf_cache_seq_length(past_key_values, layer_idx: int = 0): + """Read the number of seen tokens using the HF Cache contract, including an empty cache.""" + if past_key_values is None: + return 0 + return past_key_values.get_seq_length(layer_idx) + + +def prepare_hf_causal_masks(config, inputs_embeds, attention_mask, past_key_values, position_ids, *, eager=False): + """Build HF 5 masks by layer type without changing the model's backend config. + + ShardFormer's eager attention and ColoAttention need explicit additive masks; + an SDPA mask may otherwise be None (implicit causal) or boolean. Cache-aware + HF mask creation also handles sliding-window offsets and precomputed masks. + Imports stay local so unrelated models do not require this HF 5 API. + """ + from copy import copy + + import torch + from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask + + if eager: + config = copy(config) + config._attn_implementation = "eager" + kwargs = dict( + config=config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + factories = {"full_attention": create_causal_mask} + if "sliding_attention" in config.layer_types: + factories["sliding_attention"] = create_sliding_window_causal_mask + masks = {} + for layer_type, factory in factories.items(): + if isinstance(attention_mask, dict): + mask = attention_mask[layer_type] + # SDPA's None means implicit causality, not unrestricted attention. + if mask is None and eager: + mask = factory(**{**kwargs, "attention_mask": None}) + else: + mask = factory(**kwargs) + if eager and mask is not None and mask.dtype == torch.bool: + additive_mask = torch.zeros_like(mask, dtype=inputs_embeds.dtype) + mask = additive_mask.masked_fill(~mask, torch.finfo(inputs_embeds.dtype).min) + masks[layer_type] = mask + return masks + + +def get_hf_flash_attention_kwargs(attention_mask): + """Pass an explicit HF additive mask to ColoAttention, including cache offsets. + + Fully masked query rows (e.g. left padding) are unused by the caller's loss. + Unmask those rows to satisfy the custom-mask kernel's nonempty-row contract; + valid queries retain their original causal, padding and sliding constraints. + """ + import torch + + from .layer import AttnMaskType + + if attention_mask is None: + return {} + if attention_mask.dtype == torch.bool: + raise TypeError("ColoAttention requires an additive floating-point HF mask") + empty_rows = (attention_mask == torch.finfo(attention_mask.dtype).min) | torch.isneginf(attention_mask) + attention_mask = attention_mask.masked_fill(empty_rows.all(dim=-1, keepdim=True), 0) + return {"attention_mask": attention_mask, "attention_mask_type": AttnMaskType.CUSTOM} + + def get_obj_list_element(obj, attr: str): r""" Get the element of the list in the object diff --git a/colossalai/shardformer/layer/_operation.py b/colossalai/shardformer/layer/_operation.py index 0252f90e1c27..0bc2332487fd 100644 --- a/colossalai/shardformer/layer/_operation.py +++ b/colossalai/shardformer/layer/_operation.py @@ -461,13 +461,15 @@ def switch_step(): handles = communicate_step() # actual computation - input_tensors.append(send_tensors) + # The communication dictionaries and their buffers are reused by the + # next round. Backward needs an immutable snapshot of each input shard. + input_tensors.append({k: v.clone() for k, v in send_tensors.items()}) output_tensors.append(func(**send_tensors, **input_local)) # final round: special case, no need to send/recv again for handle in handles: handle.wait() - input_tensors.append(send_tensors) + input_tensors.append({k: v.clone() for k, v in recv_tensors.items()}) output_tensors.append(func(**recv_tensors, **input_local)) gathered_input = {} @@ -552,6 +554,7 @@ def forward(ctx, input_, weight, bias, process_group, async_grad_reduce_scatter, input_to_gather=input_to_gather, input_local=input_local, process_group=process_group, + gather_dim=dim, ) ctx.gathered_input = input_dict["input"] diff --git a/colossalai/shardformer/modeling/_bert_tf4.py b/colossalai/shardformer/modeling/_bert_tf4.py new file mode 100644 index 000000000000..b788cbf58644 --- /dev/null +++ b/colossalai/shardformer/modeling/_bert_tf4.py @@ -0,0 +1,1293 @@ +import warnings +from typing import List, Optional, Tuple, Union + +import torch +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers.modeling_outputs import ( + BaseModelOutputWithPoolingAndCrossAttentions, + CausalLMOutputWithCrossAttentions, + MaskedLMOutput, + MultipleChoiceModelOutput, + NextSentencePredictorOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from transformers.models.bert.modeling_bert import ( + BertForMaskedLM, + BertForMultipleChoice, + BertForNextSentencePrediction, + BertForPreTraining, + BertForPreTrainingOutput, + BertForQuestionAnswering, + BertForSequenceClassification, + BertForTokenClassification, + BertLMHeadModel, + BertModel, +) +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer import ShardConfig +from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward + + +class BertPipelineForwards: + """ + This class serves as a micro library for forward function substitution of Bert models + under pipeline setting. + """ + + @staticmethod + def bert_model_forward( + self: BertModel, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, # this is from the previous stage + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: + # TODO(jianghai): add explaination of the output here. + r""" + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + """ + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + batch_size, seq_length = input_shape + device = input_ids.device if input_ids is not None else inputs_embeds.device + if token_type_ids is None: + if hasattr(self.embeddings, "token_type_ids"): + buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] + buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) + token_type_ids = buffered_token_type_ids_expanded + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + else: + input_shape = hidden_states.size()[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + # past_key_values_length + past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 + + if attention_mask is None: + attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) + attention_mask = extended_attention_mask + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_extended_attention_mask = None + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + hidden_states = hidden_states if hidden_states is not None else None + + if stage_manager.is_first_stage(): + hidden_states = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + + # inherit from bert_layer,this should be changed when we add the feature to record hidden_states + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None + + if self.encoder.gradient_checkpointing and self.encoder.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + next_decoder_cache = () if use_cache else None + + start_idx, end_idx = stage_index[0], stage_index[1] + # layer_outputs + layer_outputs = hidden_states if hidden_states is not None else None + + # split the input tensor along sequence dimension + # [batch_size, seq_len, hidden_size] -> [batch_size, seq_len/TP_size, hidden_size] + if shard_config is not None and shard_config.enable_sequence_parallelism: + if shard_config.sequence_parallelism_mode == "split_gather": + hidden_states = split_forward_gather_backward( + hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + if encoder_hidden_states is not None: + encoder_hidden_states = split_forward_gather_backward( + encoder_hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + for idx, encoder_layer in enumerate(self.encoder.layer[start_idx:end_idx], start=start_idx): + if stage_manager.is_first_stage() and idx == 0: + encoder_attention_mask = encoder_extended_attention_mask + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_head_mask = head_mask[idx] if head_mask is not None else None + past_key_value = past_key_values[idx] if past_key_values is not None else None + + if self.encoder.gradient_checkpointing and self.encoder.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs, past_key_value, output_attentions) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(encoder_layer), + hidden_states, + attention_mask, + layer_head_mask, + encoder_hidden_states, + encoder_attention_mask, + ) + else: + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + layer_head_mask, + encoder_hidden_states, + encoder_attention_mask, + past_key_value, + output_attentions, + ) + hidden_states = layer_outputs[0] + if use_cache: + next_decoder_cache += (layer_outputs[-1],) + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + if self.config.add_cross_attention: + all_cross_attentions = all_cross_attentions + (layer_outputs[2],) + + # When sequence parallelism done, gather the output tensor in forward and split it in backward + if shard_config is not None and shard_config.enable_sequence_parallelism: + if shard_config.sequence_parallelism_mode == "split_gather": + hidden_states = gather_forward_split_backward( + hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # end of a stage loop + sequence_output = hidden_states if hidden_states is not None else None + + if stage_manager.is_last_stage(): + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + if not return_dict: + return (sequence_output, pooled_output) + layer_outputs[1:] + # return dict is not supported at this moment + else: + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=next_decoder_cache, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + cross_attentions=all_cross_attentions, + ) + + # output of non-first and non-last stages: must be a dict + else: + # intermediate stage always return dict + return { + "hidden_states": hidden_states, + } + + @staticmethod + def bert_for_pretraining_forward( + self: BertForPreTraining, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + next_sentence_label: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + # TODO(jianghai) left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states if hidden_states is not None else None, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + sequence_output, pooled_output = outputs[:2] + prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) + # the last stage for pretraining model + total_loss = None + if labels is not None and next_sentence_label is not None: + loss_fct = CrossEntropyLoss() + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1)) + total_loss = masked_lm_loss + next_sentence_loss + + if not return_dict: + output = (prediction_scores, seq_relationship_score) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return BertForPreTrainingOutput( + loss=total_loss, + prediction_logits=prediction_scores, + seq_relationship_logits=seq_relationship_score, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + + # intermediate stage always return dict + return { + "hidden_states": hidden_states, + } + + @staticmethod + def bert_lm_head_model_forward( + self: BertLMHeadModel, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.Tensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in + `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are + ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]` + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + """ + logger = logging.get_logger(__name__) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if labels is not None: + use_cache = False + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states if hidden_states is not None else None, + stage_index=stage_index, + shard_config=shard_config, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + lm_loss = None + if labels is not None: + # we are doing next-token prediction; shift prediction scores and input ids by one + shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous() + labels = labels[:, 1:].contiguous() + loss_fct = CrossEntropyLoss() + lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ((lm_loss,) + output) if lm_loss is not None else output + + return CausalLMOutputWithCrossAttentions( + loss=lm_loss, + logits=prediction_scores, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + cross_attentions=outputs.cross_attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + # intermediate stage always return dict + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_masked_lm_forward( + self: BertForMaskedLM, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_next_sentence_prediction_forward( + self: BertForNextSentencePrediction, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + # -> Union[Tuple[torch.Tensor], NextSentencePredictorOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair + (see `input_ids` docstring). Indices should be in `[0, 1]`: + + - 0 indicates sequence B is a continuation of sequence A, + - 1 indicates sequence B is a random sequence. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, BertForNextSentencePrediction + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") + >>> model = BertForNextSentencePrediction.from_pretrained("bert-base-uncased") + + >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light." + >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt") + + >>> outputs = model(**encoding, labels=torch.LongTensor([1])) + >>> logits = outputs.logits + >>> assert logits[0, 0] < logits[0, 1] # next sentence was random + ``` + """ + logger = logging.get_logger(__name__) + + if "next_sentence_label" in kwargs: + warnings.warn( + "The `next_sentence_label` argument is deprecated and will be removed in a future version, use" + " `labels` instead.", + FutureWarning, + ) + labels = kwargs.pop("next_sentence_label") + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + pooled_output = outputs[1] + seq_relationship_scores = self.cls(pooled_output) + + next_sentence_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1)) + + if not return_dict: + output = (seq_relationship_scores,) + outputs[2:] + return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output + + return NextSentencePredictorOutput( + loss=next_sentence_loss, + logits=seq_relationship_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + # intermediate stage always return dict + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_sequence_classification_forward( + self: BertForSequenceClassification, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + pooled_output = outputs[1] + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_token_classification_forward( + self: BertForTokenClassification, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_multiple_choice_forward( + self: BertForMultipleChoice, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., + num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See + `input_ids` above) + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # in our pipeline design,input ids are copied for every stage and shouldn't be none + # the input_ids for multiple choice model is [batch_size, num_choices, sequence_length] + if stage_manager.is_last_stage(): + num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None + attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None + token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None + position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None + inputs_embeds = ( + inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None + else None + ) + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + if stage_manager.is_last_stage(): + pooled_output = outputs[1] + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + reshaped_logits = logits.view(-1, num_choices) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(reshaped_logits, labels) + + if not return_dict: + output = (reshaped_logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return MultipleChoiceModelOutput( + loss=loss, + logits=reshaped_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_question_answering_forward( + self: BertForQuestionAnswering, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + start_positions: Optional[torch.Tensor] = None, + end_positions: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + # NOTE: the arg start_position and end_position are used only for the last stage + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_jit_fused_bert_self_output_forward(): + from transformers.models.bert.modeling_bert import BertSelfOutput + + def forward(self: BertSelfOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + return forward + + +def get_jit_fused_bert_output_forward(): + from transformers.models.bert.modeling_bert import BertOutput + + def forward(self: BertOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + return forward + + +# Fix the tgt_len size in sequence parallel attention: +# same with the one in BertSdpaSelfAttention forward in v4.51.3 transformers except the +def get_bert_sequence_parallel_attention_forward(shard_config: ShardConfig): + from transformers.models.bert.modeling_bert import BertSdpaSelfAttention + + def forward( + self: BertSdpaSelfAttention, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.FloatTensor] = None, + head_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + output_attentions: Optional[bool] = False, + ) -> Tuple[torch.Tensor]: + + bsz, tgt_len, _ = hidden_states.size() + + query_layer = self.transpose_for_scores(self.query(hidden_states)) + + # If this is instantiated as a cross-attention module, the keys and values come from an encoder; the attention + # mask needs to be such that the encoder's padding tokens are not attended to. + is_cross_attention = encoder_hidden_states is not None + + current_states = encoder_hidden_states if is_cross_attention else hidden_states + attention_mask = encoder_attention_mask if is_cross_attention else attention_mask + + # Check `seq_length` of `past_key_value` == `len(current_states)` to support prefix tuning + if is_cross_attention and past_key_value and past_key_value[0].shape[2] == current_states.shape[1]: + key_layer, value_layer = past_key_value + else: + key_layer = self.transpose_for_scores(self.key(current_states)) + value_layer = self.transpose_for_scores(self.value(current_states)) + if past_key_value is not None and not is_cross_attention: + key_layer = torch.cat([past_key_value[0], key_layer], dim=2) + value_layer = torch.cat([past_key_value[1], value_layer], dim=2) + + if self.is_decoder: + # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # Further calls to cross_attention layer can then reuse all cross-attention + # key/value_states (first "if" case) + # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # all previous decoder key/value_states. Further calls to uni-directional self-attention + # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) + # if encoder bi-directional self-attention `past_key_value` is always `None` + past_key_value = (key_layer, value_layer) + + # SDPA with memory-efficient backend is broken in torch==2.1.2 when using non-contiguous inputs and a custom + # attn_mask, so we need to call `.contiguous()` here. This was fixed in torch==2.2.0. + # Reference: https://github.com/pytorch/pytorch/issues/112577 + if self.require_contiguous_qkv and query_layer.device.type == "cuda" and attention_mask is not None: + query_layer = query_layer.contiguous() + key_layer = key_layer.contiguous() + value_layer = value_layer.contiguous() + + # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment + # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. + # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create + # a causal mask in case tgt_len == 1. + is_causal = ( + True if self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 else False + ) + attn_output = torch.nn.functional.scaled_dot_product_attention( + query_layer, + key_layer, + value_layer, + attn_mask=attention_mask, + dropout_p=self.dropout_prob if self.training else 0.0, + is_causal=is_causal, + ) + + attn_output = attn_output.transpose(1, 2) + _, _, tgt_len, _ = query_layer.shape + attn_output = attn_output.reshape(bsz, tgt_len, self.all_head_size) + + outputs = (attn_output,) + if self.is_decoder: + outputs = outputs + (past_key_value,) + return outputs + + return forward + + +def bert_sequence_parallel_forward_fn(shard_config: ShardConfig): + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: + r""" + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + batch_size, seq_length = input_shape + device = input_ids.device if input_ids is not None else inputs_embeds.device + + # past_key_values_length + past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 + + if attention_mask is None: + attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) + + if token_type_ids is None: + if hasattr(self.embeddings, "token_type_ids"): + buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] + buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) + token_type_ids = buffered_token_type_ids_expanded + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] + # ourselves in which case we just need to make it broadcastable to all heads. + extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) + else: + encoder_extended_attention_mask = None + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + + # split the input tensor along sequence dimension + # [batch_size, seq_len, hidden_size] -> [batch_size, seq_len/TP_size, hidden_size] + embedding_output = split_forward_gather_backward( + embedding_output, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + if encoder_hidden_states is not None: + encoder_hidden_states = split_forward_gather_backward( + encoder_hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=extended_attention_mask, + head_mask=head_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_extended_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = encoder_outputs[0] + + # When sequence parallelism done, gather the output tensor in forward and split it in backward + sequence_output = gather_forward_split_backward( + sequence_output, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + return (sequence_output, pooled_output) + encoder_outputs[1:] + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + return forward + + +def get_jit_fused_bert_intermediate_forward(): + from transformers.models.bert.modeling_bert import BertIntermediate + + from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction + + def forward(self: BertIntermediate, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, bias = self.dense(hidden_states) + hidden_states = JitGeLUFunction.apply(hidden_states, bias) + return hidden_states + + return forward diff --git a/colossalai/shardformer/modeling/_bert_tf5.py b/colossalai/shardformer/modeling/_bert_tf5.py new file mode 100644 index 000000000000..0dfa1d4ed0dc --- /dev/null +++ b/colossalai/shardformer/modeling/_bert_tf5.py @@ -0,0 +1,1282 @@ +import warnings +from typing import List, Optional, Tuple, Union + +import torch +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers.cache_utils import Cache, DynamicCache, EncoderDecoderCache +from transformers.masking_utils import create_bidirectional_mask, create_causal_mask +from transformers.modeling_outputs import ( + BaseModelOutputWithPoolingAndCrossAttentions, + CausalLMOutputWithCrossAttentions, + MaskedLMOutput, + MultipleChoiceModelOutput, + NextSentencePredictorOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from transformers.models.bert.modeling_bert import ( + BertForMaskedLM, + BertForMultipleChoice, + BertForNextSentencePrediction, + BertForPreTraining, + BertForPreTrainingOutput, + BertForQuestionAnswering, + BertForSequenceClassification, + BertForTokenClassification, + BertLMHeadModel, + BertModel, +) +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer import ShardConfig +from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward + + +class BertPipelineForwards: + """ + This class serves as a micro library for forward function substitution of Bert models + under pipeline setting. + """ + + @staticmethod + def bert_model_forward( + self: BertModel, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, # this is from the previous stage + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: + # TODO(jianghai): add explaination of the output here. + r""" + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + """ + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + batch_size, seq_length = input_shape + device = input_ids.device if input_ids is not None else inputs_embeds.device + if token_type_ids is None: + if hasattr(self.embeddings, "token_type_ids"): + buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] + buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) + token_type_ids = buffered_token_type_ids_expanded + else: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + else: + input_shape = hidden_states.size()[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + # `head_mask` is gone in v5 (`get_head_mask` was removed too). The other PP forwards still + # pass it by name, so the parameter stays; a non-None value warns instead of being + # swallowed by v5's `**kwargs` like upstream does. + if head_mask is not None: + logger.warning_once( + "`head_mask` is not supported for pipeline models at the moment and was removed in " + "transformers v5; the value is ignored." + ) + + # past_key_values_length + # v5 dropped `DynamicCache.from_legacy_cache` / `to_legacy_cache` / + # `Cache.get_usable_length`; lengths come from `get_seq_length()`. A non-`Cache` argument + # raises rather than falling back to the old `past_key_values[0][0].shape[2]`. + if past_key_values is not None and not isinstance(past_key_values, Cache): + raise TypeError( + "`past_key_values` must be a `transformers.cache_utils.Cache` under transformers v5 " + f"(got {type(past_key_values).__name__}); the legacy tuple format was removed upstream." + ) + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + + if attention_mask is None: + attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) + + if stage_manager.is_first_stage(): + hidden_states = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + + # The mask is built after the embeddings because v5's helpers infer shape / dtype / device + # from the embedding output (their second parameter is `inputs_embeds`, where v4's + # `get_extended_attention_mask` only wanted `input_shape`). Off the first stage, + # `hidden_states` holds the activations from the previous stage, which are equivalent for + # that purpose. + if self.config.is_decoder: + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + ) + else: + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + ) + + # If a 2D or 3D attention mask is provided for the cross-attention + # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] + if self.config.is_decoder and encoder_hidden_states is not None: + encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() + encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) + if encoder_attention_mask is None: + encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) + encoder_extended_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + else: + encoder_extended_attention_mask = None + + # inherit from bert_layer,this should be changed when we add the feature to record hidden_states + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None + + # The old `self.encoder.gradient_checkpointing` guard would now raise: v5's `BertEncoder` is + # a plain `nn.Module` with no such attribute. Checking pointing is handled by each + # `BertLayer` (a `GradientCheckpointingLayer`) instead -- see the loop below. + next_decoder_cache = () if use_cache else None + + start_idx, end_idx = stage_index[0], stage_index[1] + # layer_outputs + layer_outputs = hidden_states if hidden_states is not None else None + + # split the input tensor along sequence dimension + # [batch_size, seq_len, hidden_size] -> [batch_size, seq_len/TP_size, hidden_size] + if shard_config is not None and shard_config.enable_sequence_parallelism: + if shard_config.sequence_parallelism_mode == "split_gather": + hidden_states = split_forward_gather_backward( + hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + if encoder_hidden_states is not None: + encoder_hidden_states = split_forward_gather_backward( + encoder_hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + for idx, encoder_layer in enumerate(self.encoder.layer[start_idx:end_idx], start=start_idx): + if stage_manager.is_first_stage() and idx == 0: + encoder_attention_mask = encoder_extended_attention_mask + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # v5's `BertLayer.forward` takes `(hidden_states, attention_mask, encoder_hidden_states, + # encoder_attention_mask, past_key_values, **kwargs)`: no `head_mask` and no + # `output_attentions`, and it returns a bare tensor rather than `(hidden, [attn...])`. + # Passing positionally would land arguments on the wrong parameters and mis-compute + # silently, so `past_key_values` goes by name here (one whole `Cache`, indexed inside + # `Cache.update`, not `past_key_values[idx]` per layer) and the result is used directly. + # `hidden_states` stays positional, as `BertLayer.__call__` requires. + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + ) + hidden_states = layer_outputs + if use_cache: + next_decoder_cache = past_key_values + # `all_self_attentions` / `all_cross_attentions` stay None: this path forces + # `output_attentions` to False above, and v5's `BertLayer` no longer returns them. + + # When sequence parallelism done, gather the output tensor in forward and split it in backward + if shard_config is not None and shard_config.enable_sequence_parallelism: + if shard_config.sequence_parallelism_mode == "split_gather": + hidden_states = gather_forward_split_backward( + hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + # end of a stage loop + sequence_output = hidden_states if hidden_states is not None else None + + if stage_manager.is_last_stage(): + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + outputs = BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=next_decoder_cache, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + cross_attentions=all_cross_attentions, + ) + if not return_dict: + # The old `(sequence_output, pooled_output) + layer_outputs[1:]` would slice a + # tensor now that `layer_outputs` is a bare tensor. `to_tuple()` follows the + # dataclass field order instead of a hand-written literal. + return outputs.to_tuple() + return outputs + + # output of non-first and non-last stages: must be a dict + else: + # intermediate stage always return dict + return { + "hidden_states": hidden_states, + } + + @staticmethod + def bert_for_pretraining_forward( + self: BertForPreTraining, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + next_sentence_label: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + # TODO(jianghai) left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states if hidden_states is not None else None, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + sequence_output, pooled_output = outputs[:2] + prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) + # the last stage for pretraining model + total_loss = None + if labels is not None and next_sentence_label is not None: + loss_fct = CrossEntropyLoss() + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1)) + total_loss = masked_lm_loss + next_sentence_loss + + if not return_dict: + output = (prediction_scores, seq_relationship_score) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return BertForPreTrainingOutput( + loss=total_loss, + prediction_logits=prediction_scores, + seq_relationship_logits=seq_relationship_score, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + + # intermediate stage always return dict + return { + "hidden_states": hidden_states, + } + + @staticmethod + def bert_lm_head_model_forward( + self: BertLMHeadModel, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + past_key_values: Optional[List[torch.Tensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if + the model is configured as a decoder. + encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in + the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in + `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are + ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]` + past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): + Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that + don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all + `decoder_input_ids` of shape `(batch_size, sequence_length)`. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + """ + logger = logging.get_logger(__name__) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if labels is not None: + use_cache = False + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states if hidden_states is not None else None, + stage_index=stage_index, + shard_config=shard_config, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + lm_loss = None + if labels is not None: + # we are doing next-token prediction; shift prediction scores and input ids by one + shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous() + labels = labels[:, 1:].contiguous() + loss_fct = CrossEntropyLoss() + lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ((lm_loss,) + output) if lm_loss is not None else output + + return CausalLMOutputWithCrossAttentions( + loss=lm_loss, + logits=prediction_scores, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + cross_attentions=outputs.cross_attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + # intermediate stage always return dict + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_masked_lm_forward( + self: BertForMaskedLM, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., + config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the + loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() # -100 index = padding token + masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (prediction_scores,) + outputs[2:] + return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + + return MaskedLMOutput( + loss=masked_lm_loss, + logits=prediction_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_next_sentence_prediction_forward( + self: BertForNextSentencePrediction, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + # -> Union[Tuple[torch.Tensor], NextSentencePredictorOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair + (see `input_ids` docstring). Indices should be in `[0, 1]`: + + - 0 indicates sequence B is a continuation of sequence A, + - 1 indicates sequence B is a random sequence. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, BertForNextSentencePrediction + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") + >>> model = BertForNextSentencePrediction.from_pretrained("bert-base-uncased") + + >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." + >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light." + >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt") + + >>> outputs = model(**encoding, labels=torch.LongTensor([1])) + >>> logits = outputs.logits + >>> assert logits[0, 0] < logits[0, 1] # next sentence was random + ``` + """ + logger = logging.get_logger(__name__) + + if "next_sentence_label" in kwargs: + warnings.warn( + "The `next_sentence_label` argument is deprecated and will be removed in a future version, use" + " `labels` instead.", + FutureWarning, + ) + labels = kwargs.pop("next_sentence_label") + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + pooled_output = outputs[1] + seq_relationship_scores = self.cls(pooled_output) + + next_sentence_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1)) + + if not return_dict: + output = (seq_relationship_scores,) + outputs[2:] + return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output + + return NextSentencePredictorOutput( + loss=next_sentence_loss, + logits=seq_relationship_scores, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + # intermediate stage always return dict + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_sequence_classification_forward( + self: BertForSequenceClassification, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + pooled_output = outputs[1] + + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + + loss = None + if labels is not None: + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_token_classification_forward( + self: BertForTokenClassification, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_multiple_choice_forward( + self: BertForMultipleChoice, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., + num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See + `input_ids` above) + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # in our pipeline design,input ids are copied for every stage and shouldn't be none + # the input_ids for multiple choice model is [batch_size, num_choices, sequence_length] + if stage_manager.is_last_stage(): + num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None + attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None + token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None + position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None + inputs_embeds = ( + inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None + else None + ) + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + if stage_manager.is_last_stage(): + pooled_output = outputs[1] + pooled_output = self.dropout(pooled_output) + logits = self.classifier(pooled_output) + reshaped_logits = logits.view(-1, num_choices) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + loss = loss_fct(reshaped_logits, labels) + + if not return_dict: + output = (reshaped_logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return MultipleChoiceModelOutput( + loss=loss, + logits=reshaped_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def bert_for_question_answering_forward( + self: BertForQuestionAnswering, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + start_positions: Optional[torch.Tensor] = None, + end_positions: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.Tensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + # NOTE: the arg start_position and end_position are used only for the last stage + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + outputs = BertPipelineForwards.bert_model_forward( + self.bert, + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + hidden_states=hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=shard_config, + ) + if stage_manager.is_last_stage(): + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + total_loss = None + if start_positions is not None and end_positions is not None: + # If we are on multi-GPU, split add a dimension + if len(start_positions.size()) > 1: + start_positions = start_positions.squeeze(-1) + if len(end_positions.size()) > 1: + end_positions = end_positions.squeeze(-1) + # sometimes the start/end positions are outside our model inputs, we ignore these terms + ignored_index = start_logits.size(1) + start_positions = start_positions.clamp(0, ignored_index) + end_positions = end_positions.clamp(0, ignored_index) + + loss_fct = CrossEntropyLoss(ignore_index=ignored_index) + start_loss = loss_fct(start_logits, start_positions) + end_loss = loss_fct(end_logits, end_positions) + total_loss = (start_loss + end_loss) / 2 + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((total_loss,) + output) if total_loss is not None else output + + return QuestionAnsweringModelOutput( + loss=total_loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_jit_fused_bert_self_output_forward(): + from transformers.models.bert.modeling_bert import BertSelfOutput + + def forward(self: BertSelfOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + return forward + + +def get_jit_fused_bert_output_forward(): + from transformers.models.bert.modeling_bert import BertOutput + + def forward(self: BertOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) + hidden_states = self.LayerNorm(hidden_states) + return hidden_states + + return forward + + +# Fix the tgt_len size in sequence parallel attention. Under ``split_gather``, ``self.query`` / +# ``self.key`` / ``self.value`` are ``Linear1D_Col(seq_parallel_mode=...)``, which **gather the full +# sequence** in forward (``linear_gather_forward_reducescatter_backward``), while ``hidden_states`` +# only holds this rank's shard. v5's ``BertSelfAttention.forward`` derives the output shape from +# ``hidden_states`` (``input_shape = hidden_states.shape[:-1]``), so it would reshape a +# full-sequence ``attn_output`` into a shard-shaped tensor -- the element counts do not even match. +# Taking the shape from the projection output instead is the only reason this override exists; +# everything else follows v5's own ``BertSelfAttention.forward``. +def get_bert_sequence_parallel_attention_forward(shard_config: ShardConfig): + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + from transformers.models.bert.modeling_bert import BertSelfAttention, eager_attention_forward + + def forward( + self: BertSelfAttention, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.FloatTensor] = None, + past_key_values: Optional[Cache] = None, + **kwargs, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # `full_shape` is derived from the projection output, not from `hidden_states` -- see above. + projected = self.query(hidden_states) + full_shape = (*projected.shape[:-1], -1, self.attention_head_size) + query_layer = projected.view(*full_shape).transpose(1, 2) + key_layer = self.key(hidden_states).view(*full_shape).transpose(1, 2) + value_layer = self.value(hidden_states).view(*full_shape).transpose(1, 2) + + if past_key_values is not None: + # One whole `Cache`, indexed inside `update(k, v, layer_idx)` -- no per-layer tuple + # subscripts, and no v4 prefix-tuning branch on `past_key_value[0]`. + current_past_key_values = past_key_values + if isinstance(past_key_values, EncoderDecoderCache): + current_past_key_values = past_key_values.self_attention_cache + key_layer, value_layer = current_past_key_values.update(key_layer, value_layer, self.layer_idx) + + # `attention_mask` comes from v5's `create_bidirectional_mask` / `create_causal_mask`, + # already in the form the attention interface consumes, so the old hand-rolled `is_causal` + # dispatch to SDPA is gone. + attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface( + self.config._attn_implementation, eager_attention_forward + ) + attn_output, attn_weights = attention_interface( + self, + query_layer, + key_layer, + value_layer, + attention_mask, + dropout=0.0 if not self.training else self.dropout.p, + scaling=self.scaling, + **kwargs, + ) + # `full_shape[1]` is the gathered, full sequence length, not this rank's shard. + attn_output = attn_output.reshape(full_shape[0], full_shape[1], -1).contiguous() + return attn_output, attn_weights + + return forward + + +def bert_sequence_parallel_forward_fn(shard_config: ShardConfig): + r"""Replacement of ``BertModel.forward`` for ``split_gather`` sequence parallelism. + + Only the sequence splitting/gathering at the two ends of the encoder is ours; the rest is kept + line-for-line equal to v5's ``BertModel.forward`` so that the v5 contracts (mask construction, + ``Cache``, declarative output capture) hold: + + * ``embedding_output`` is split along the sequence dim before the encoder, and + ``sequence_output`` is gathered back before the pooler. + * The attention mask is built from the **full** (pre-split) embeddings -- v5's mask helpers infer + the shapes from ``inputs_embeds``, and the q/k/v projections gather the sequence back anyway, + so attention runs over the full sequence on every rank. Building the mask after the split + would produce a shard-sized mask that cannot broadcast against full-sequence attention. + * ``@capture_outputs`` is re-attached: it is v5's only producer of ``attentions`` / + ``hidden_states`` (via hooks on ``BertLayer`` / ``BertSelfAttention``). Without it a replaced + ``BertModel.forward`` silently loses both, and ``BertForPreTraining.forward`` reads + ``outputs.attentions`` off our return value. + """ + from transformers.utils.generic import merge_with_config_defaults + from transformers.utils.output_capturing import capture_outputs + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + token_type_ids: Optional[torch.Tensor] = None, + position_ids: Optional[torch.Tensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs, + ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: + r""" + `output_attentions` / `output_hidden_states` / `return_dict` are not read by this body under + v5: the first two are read off the call-site kwargs by the `@capture_outputs` wrapper, which + then injects what it recorded, and it pops the third one and calls `to_tuple()` for `False`. + All three are kept for API shape only. + """ + logger = logging.get_logger(__name__) + + # `head_mask` was removed from `BertModel.forward` in v5 (along with `get_head_mask`). + # Upstream lets it ride `**kwargs` into attention; popping it here makes "passed but ignored" + # visible instead. + head_mask = kwargs.pop("head_mask", None) + if head_mask is not None: + logger.warning_once("`head_mask` was removed from `BertModel.forward` in transformers v5 and is ignored.") + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.config.is_decoder: + use_cache = use_cache if use_cache is not None else self.config.use_cache + else: + use_cache = False + + if past_key_values is not None and not isinstance(past_key_values, Cache): + raise TypeError( + "`past_key_values` must be a `transformers.cache_utils.Cache` under transformers v5 " + f"(got {type(past_key_values).__name__}); the legacy tuple format was removed upstream." + ) + if use_cache and past_key_values is None: + # v5 dropped `DynamicCache.from_legacy_cache` / `to_legacy_cache` / + # `get_usable_length`: the cache is constructed directly and lengths come from + # `get_seq_length()`. + past_key_values = ( + EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config)) + if encoder_hidden_states is not None or self.config.is_encoder_decoder + else DynamicCache(config=self.config) + ) + + past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0 + + embedding_output = self.embeddings( + input_ids=input_ids, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + past_key_values_length=past_key_values_length, + ) + + # The mask is built after the embeddings (v5's helpers need the embedding output to infer + # shape/dtype/device) but before the split -- see the function docstring. + if self.config.is_decoder: + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + past_key_values=past_key_values, + ) + else: + attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=attention_mask, + ) + + if encoder_attention_mask is not None: + encoder_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=embedding_output, + attention_mask=encoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + ) + + # split the input tensor along sequence dimension + # [batch_size, seq_len, hidden_size] -> [batch_size, seq_len/TP_size, hidden_size] + embedding_output = split_forward_gather_backward( + embedding_output, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + if encoder_hidden_states is not None: + encoder_hidden_states = split_forward_gather_backward( + encoder_hidden_states, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + encoder_outputs = self.encoder( + embedding_output, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + position_ids=position_ids, + **kwargs, + ) + + # v5's `BertEncoder.forward` returns only `last_hidden_state` / `past_key_values` and its + # layers return bare tensors: `attentions` / `hidden_states` are no longer accumulated, they + # are injected by `@capture_outputs`. + sequence_output = encoder_outputs.last_hidden_state + + # When sequence parallelism done, gather the output tensor in forward and split it in backward + sequence_output = gather_forward_split_backward( + sequence_output, + dim=1, + process_group=shard_config.tensor_parallel_process_group, + fp8_communication=shard_config.fp8_communication, + ) + + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + return BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + past_key_values=encoder_outputs.past_key_values, + ) + + return forward + + +def get_jit_fused_bert_intermediate_forward(): + from transformers.models.bert.modeling_bert import BertIntermediate + + from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction + + def forward(self: BertIntermediate, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, bias = self.dense(hidden_states) + hidden_states = JitGeLUFunction.apply(hidden_states, bias) + return hidden_states + + return forward diff --git a/colossalai/shardformer/modeling/_deepseek_v3_tf4.py b/colossalai/shardformer/modeling/_deepseek_v3_tf4.py new file mode 100644 index 000000000000..c582ac9a3265 --- /dev/null +++ b/colossalai/shardformer/modeling/_deepseek_v3_tf4.py @@ -0,0 +1,370 @@ +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup +from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast + +from colossalai.lazy import LazyInitContext +from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer.linear import ParallelModule +from colossalai.shardformer.shard.utils import set_tensors_to_none +from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_group + + +class EpDeepseekV3MoE(ParallelModule): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}") + + def setup_process_groups( + self, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + ): + assert moe_dp_group is not None + assert ep_group is not None + + self.ep_size = dist.get_world_size(ep_group) + self.ep_rank = dist.get_rank(ep_group) + self.num_experts = self.config.n_routed_experts + assert self.num_experts % self.ep_size == 0 + + self.ep_group = ep_group + self.num_experts_per_ep = self.num_experts // self.ep_size + self.experts_per_rank = self.num_experts_per_ep + self.expert_start_idx = self.ep_rank * self.num_experts_per_ep + held_experts = self.experts[self.expert_start_idx : self.expert_start_idx + self.num_experts_per_ep] + + set_tensors_to_none(self.experts, exclude=set(held_experts)) + + # setup moe_dp group + self.moe_dp_group = moe_dp_group + self.moe_dp_size = dist.get_world_size(moe_dp_group) + + for p in self.experts.parameters(): + set_moe_tensor_ep_group(p, ep_group) + + @staticmethod + def from_native_module( + module, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + *args, + **kwargs, + ) -> "EpDeepseekV3MoE": + if module.__class__.__name__ != "DeepseekV3MLP": + module.__class__ = EpDeepseekV3MoE + module.setup_process_groups(moe_dp_group, ep_group) + LazyInitContext.materialize(module) + return module + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + orig_shape = hidden_states.shape + topk_idx, topk_weight = self.gate(hidden_states) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape) + if self.config.n_shared_experts is not None: + y = y + self.shared_experts(identity) + return y + + def moe_forward(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: + cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) + cnts.scatter_(1, topk_ids, 1) + tokens_per_expert = cnts.sum(dim=0) + idxs = topk_ids.view(-1).argsort() + sorted_tokens = x[idxs // topk_ids.shape[1]] + if self.ep_size > 1: + tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1) + tokens_per_expert_group = tokens_per_expert.new_empty(tokens_per_expert.shape[0]) + dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert, group=self.ep_group) + + output_splits = tokens_per_expert_group.view(self.ep_size, -1).sum(1).tolist() + input_split_sizes = tokens_per_ep_rank.tolist() + + gathered_tokens, _ = all_to_all_uneven(sorted_tokens, input_split_sizes, output_splits, self.ep_group) + tokens_per_expert_post_gather = tokens_per_expert_group.view(self.ep_size, self.experts_per_rank).sum(dim=0) + gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) + s = 0 + for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): + gatherd_idxs[s : s + k] = i % self.experts_per_rank + s += k + gatherd_idxs = gatherd_idxs.argsort() + sorted_tokens = gathered_tokens[gatherd_idxs] + tokens_per_expert = tokens_per_expert_post_gather + + # ep related code + sorted_tokens = EPGradScalerIn.apply(sorted_tokens, self.ep_size) + + tokens_per_expert = tokens_per_expert.cpu().numpy() + + outputs = [] + start_idx = 0 + for i, num_tokens in enumerate(tokens_per_expert): + end_idx = start_idx + num_tokens + if num_tokens == 0: + continue + expert = self.experts[i + self.ep_rank * self.experts_per_rank] + tokens_for_this_expert = sorted_tokens[start_idx:end_idx] + expert_out = expert(tokens_for_this_expert) + outputs.append(expert_out) + start_idx = end_idx + + if len(outputs) > 0: + outs = torch.cat(outputs, dim=0) + else: + assert sorted_tokens.numel() == 0, f"sorted_tokens: should be empty, but got {sorted_tokens.shape}" + outs = sorted_tokens + + if self.ep_size > 1: + outs = EPGradScalerOut.apply(outs, self.ep_size) + new_x = torch.empty_like(outs) + new_x[gatherd_idxs] = outs + gathered_tokens, _ = all_to_all_uneven(new_x, output_splits, input_split_sizes, self.ep_group) + outs = gathered_tokens + + new_x = torch.empty_like(outs) + new_x[idxs] = outs + final_out = ( + (new_x.view(*topk_ids.shape, -1).type(topk_weight.dtype) * topk_weight.unsqueeze(dim=-1)) + .sum(dim=1) + .type(new_x.dtype) + ) + + return final_out + + +def deepseek_v3_model_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + hidden_states_internal: Optional[torch.Tensor] = None, +) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + past_key_values_length = 0 + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = past_key_values.get_usable_length(seq_length) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0) + + if stage_manager is None or stage_manager.is_first_stage(): + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + else: + inputs_embeds = hidden_states_internal + + if self._use_flash_attention_2: + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + ) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + if stage_index is not None: + start_idx, end_idx = stage_index + else: + start_idx, end_idx = 0, len(self.layers) + for i, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and i > 0: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if stage_manager is None or stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache + if stage_manager is not None and not stage_manager.is_last_stage(): + return { + "hidden_states_internal": hidden_states, + } + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +def deepseek_v3_for_causal_lm_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + hidden_states_internal: Optional[torch.Tensor] = None, +) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`. + Returns: + Example: + ```python + >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM + >>> model = DeepseekV3ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = deepseek_v3_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + stage_index=stage_index, + hidden_states_internal=hidden_states_internal, + ) + if stage_manager is not None and not stage_manager.is_last_stage(): + return outputs + + hidden_states = outputs[0] + + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/colossalai/shardformer/modeling/_deepseek_v3_tf5.py b/colossalai/shardformer/modeling/_deepseek_v3_tf5.py new file mode 100644 index 000000000000..b8935d766708 --- /dev/null +++ b/colossalai/shardformer/modeling/_deepseek_v3_tf5.py @@ -0,0 +1,408 @@ +from typing import List, Optional, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed import ProcessGroup +from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache, DynamicCache +from transformers.masking_utils import create_causal_mask +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.utils.output_capturing import capture_outputs + +from colossalai.lazy import LazyInitContext +from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer.linear import ParallelModule +from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_fused, set_moe_tensor_ep_group + + +class EpDeepseekV3MoE(ParallelModule): + """ + A mixed expert module containing shared experts. + """ + + def __init__(self, config): + raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}") + + def setup_process_groups( + self, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + ): + assert moe_dp_group is not None + assert ep_group is not None + + self.ep_size = dist.get_world_size(ep_group) + self.ep_rank = dist.get_rank(ep_group) + self.ep_group = ep_group + + # v5 stores experts as fused 3D parameters (`gate_up_proj [E,2I,H]` / `down_proj [E,H,I]`), so + # `num_experts` lives on `self.experts`, not on the block (docs/30 §二). + num_experts = self.experts.num_experts + if num_experts % self.ep_size != 0: + raise ValueError("The number of experts must be divisible by the number of expert parallel groups.") + + self.num_experts = num_experts + self.num_experts_per_ep = num_experts // self.ep_size + self.expert_start_idx = self.ep_rank * self.num_experts_per_ep + + # primitive ① (docs/30 §4.1): slice the fused params to the local experts and release the rest + # (P5). `.clone()` is required, not `.contiguous()`: the dim-0 slice of a contiguous tensor is + # already contiguous, so `.contiguous()` would return the view and keep the full `[E, ...]` alive. + experts = self.experts + s = self.expert_start_idx + n = self.num_experts_per_ep + # Swap the storage but keep the Parameter objects: optimizers are built from the unsharded model and + # the MoE plugin only prunes (never re-adds) optimizer params, so a new Parameter would drop out of training. + experts.gate_up_proj.data = experts.gate_up_proj.data[s : s + n].clone() + experts.down_proj.data = experts.down_proj.data[s : s + n].clone() + experts.num_experts = n + + # setup moe_dp group + self.moe_dp_group = moe_dp_group + self.moe_dp_size = dist.get_world_size(moe_dp_group) + + # primitive ② (docs/30 §4.2): mark the sliced fused params so the sharded loader slices dim 0 + # (the expert dimension) by `ep_group`. + for p in self.experts.parameters(): + set_moe_tensor_ep_group(p, ep_group) + set_moe_tensor_ep_fused(p) + + @staticmethod + def from_native_module( + module, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + *args, + **kwargs, + ) -> "EpDeepseekV3MoE": + # Materialize first, then slice: primitive ① clones the fused params, which must already be + # real tensors under lazy-init (mirrors EPMixtralSparseMoeBlock, docs/30 §四). + LazyInitContext.materialize(module) + # The first `first_k_dense_replace` layers keep a dense `DeepseekV3MLP` (no experts); only + # the MoE layers (`DeepseekV3MoE`) carry the fused 3D params to slice. + if module.__class__.__name__ != "DeepseekV3MLP": + module.__class__ = EpDeepseekV3MoE + module.setup_process_groups(moe_dp_group, ep_group) + return module + + def _expert_forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor: + r"""Run one local expert on its dispatched tokens (docs/30 §4.3). + + Mirrors v5's ``DeepseekV3Experts.forward``: ``gate, up = linear(x, gate_up_proj[e]).chunk(2, + dim=-1)``, ``act_fn(gate) * up``, then ``linear(down_proj[e])``. The fused params are already + sliced to ``[E/ep, ...]``, so ``expert_idx`` is a local index. + """ + gate, up = F.linear(x, self.experts.gate_up_proj[expert_idx]).chunk(2, dim=-1) + return F.linear(self.experts.act_fn(gate) * up, self.experts.down_proj[expert_idx]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + # v5's `DeepseekV3TopkRouter` returns the (logits, top-k weights, top-k indices) triple; group + # scoring / top-k / normalisation / `routed_scaling_factor` are already applied inside it, so + # there is no local routing to redo here. + _, topk_weights, topk_indices = self.gate(hidden_states) + routing_weights = topk_weights.to(hidden_states.dtype) + + # Unlike Mixtral, `DeepseekV3MoE` carries no `self.top_k`; it lives on the router. + top_k = self.gate.top_k + + selected_experts = topk_indices.t().reshape(-1) + selected_experts_idx = selected_experts.argsort() + dispatch_states = hidden_states.repeat(top_k, 1)[selected_experts_idx] + input_split_sizes = selected_experts.bincount(minlength=self.num_experts) + + output_split_sizes = torch.zeros_like(input_split_sizes) + + dist.all_to_all_single(output_split_sizes, input_split_sizes, group=self.ep_group) + + input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() + output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() + + output_states, _ = all_to_all_uneven( + dispatch_states, + input_split_list, + output_split_list, + self.ep_group, + ) + # compute expert output + output_states = EPGradScalerIn.apply(output_states, self.ep_size) + if output_states.size(0) > 0: + if self.num_experts_per_ep == 1: + # no need to split + output_states = self._expert_forward(output_states, 0) + else: + output_states_splits = output_states.split(output_split_sizes.tolist()) + output_states_list = [] + for i, split_states in enumerate(output_states_splits): + if split_states.size(0) == 0: + continue + expert_idx = i % self.num_experts_per_ep + split_states = self._expert_forward(split_states, expert_idx) + output_states_list.append(split_states) + output_states = torch.cat(output_states_list) + + output_states = EPGradScalerOut.apply(output_states, self.ep_size) + dispatch_states, _ = all_to_all_uneven(output_states, output_split_list, input_split_list, self.ep_group) + + recover_experts_idx = torch.empty_like(selected_experts_idx) + recover_experts_idx[selected_experts_idx] = torch.arange( + selected_experts_idx.size(0), device=selected_experts_idx.device + ) + dispatch_states = dispatch_states[recover_experts_idx] + k_hidden_states = dispatch_states.chunk(top_k) + output_states = k_hidden_states[0] * routing_weights[:, 0, None] + for i in range(1, top_k): + output_states += k_hidden_states[i] * routing_weights[:, i, None] + output_states = output_states.view(*orig_shape) + + if self.config.n_shared_experts is not None: + output_states = output_states + self.shared_experts(identity) + return output_states + + +# v5 records `attentions` / `hidden_states` declaratively and injects them via `@capture_outputs`. +# v5's own forward carries that decorator and ours replaces it, so it has to be put back. Unlike +# Mixtral, the recorded classes survive here (the policy keeps `self_attn`). +@capture_outputs +def deepseek_v3_model_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + hidden_states_internal: Optional[torch.Tensor] = None, + **kwargs, +) -> Union[Tuple, BaseModelOutputWithPast]: + # Not consumed here under v5: the `@capture_outputs` wrapper reads these from `kwargs` to + # decide what to record, then injects the result. Kept for API shape, and deliberately not + # resolved -- a resolved value would look like this body accumulates them itself. + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + past_key_values_length = 0 + if use_cache: + # v5 dropped `DynamicCache.from_legacy_cache` / `to_legacy_cache` and + # `Cache.get_usable_length`; `ddp_cache_data` still takes the per-layer + # `(key, value[, sliding_window])` tuples, and the length comes from `get_seq_length()`. + if past_key_values is None: + past_key_values = DynamicCache() + elif not isinstance(past_key_values, Cache): + past_key_values = DynamicCache(past_key_values) + past_key_values_length = past_key_values.get_seq_length() + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0) + + if stage_manager is None or stage_manager.is_first_stage(): + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + else: + inputs_embeds = hidden_states_internal + + # v5 dropped `DeepseekV3Model._update_causal_mask` in favour of `create_causal_mask`, with + # `allow_is_causal_skip` left at its default: DeepSeek-v3 runs v5's native attention (the + # policy keeps `self_attn`), which lets the backend decide, and `_use_flash_attention_2` -- + # the attribute the old manual dispatch read -- no longer exists. + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + # embed positions + hidden_states = inputs_embeds + + # decoder layers + # Neither is accumulated by hand under v5: `DecoderLayer.forward` returns a bare tensor and + # `@capture_outputs` overwrites both fields afterwards. Accumulating would also give wrong + # results for v5's new "layer-index list" form of `output_hidden_states`. + all_hidden_states = None + all_self_attns = None + next_decoder_cache = None + + if stage_index is not None: + start_idx, end_idx = stage_index + else: + start_idx, end_idx = 0, len(self.layers) + position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids) + for i, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if self.gradient_checkpointing: + # v5 moves the recompute decision into `GradientCheckpointingLayer.__call__`, driven by + # a per-layer flag (which is what `gradient_checkpointing_enable()` sets -- the flag + # relies on it). Calling `_gradient_checkpointing_func` directly, as before, now + # misaligns the trailing positional args and fails inside attention. Layer 0 keeps + # its old exemption. + decoder_layer.gradient_checkpointing = i > 0 + + # Keyword names must match v5's `DeepseekV3DecoderLayer.forward`: the old `past_key_value=` + # (singular) and `output_attentions=` no longer exist and are swallowed by `**kwargs` -- + # that is how the KV cache was being lost. + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + ) + # v5's `DeepseekV3DecoderLayer.forward` returns a bare tensor (v4 returned a tuple), so + # the old `layer_outputs[0]` would index into the batch dim instead. + hidden_states = layer_outputs + + if use_cache: + # v5 updates the passed-in `Cache` in place instead of returning it + next_decoder_cache = past_key_values + + if stage_manager is None or stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache + if stage_manager is not None and not stage_manager.is_last_stage(): + return { + "hidden_states_internal": hidden_states, + } + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + +def deepseek_v3_for_causal_lm_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + stage_manager: Optional[PipelineStageManager] = None, + stage_index: Optional[List[int]] = None, + hidden_states_internal: Optional[torch.Tensor] = None, + **kwargs, +) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`. + Returns: + Example: + ```python + >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM + >>> model = DeepseekV3ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + # v5 moves `output_attentions` / `output_hidden_states` into `**kwargs`, so pull them back out + # here -- otherwise the same key travels both ways and collides on the way down. + output_attentions = kwargs.pop("output_attentions", output_attentions) + output_hidden_states = kwargs.pop("output_hidden_states", output_hidden_states) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # These must be resolved booleans, never `None`: the downstream `@capture_outputs` wrapper + # tests `kwargs.get("output_attentions", config.output_attentions)`, and an explicit `None` + # does not fall back to the config -- it silently turns recording off. + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + # `return_dict` is deliberately not forwarded: the wrapper pops it and unconditionally calls + # `outputs.to_tuple()` when it is False, but the non-last PP stage returns a plain dict, which + # has no `to_tuple`. The parameter below is therefore unreachable from outside, kept for API. + outputs = deepseek_v3_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + stage_manager=stage_manager, + stage_index=stage_index, + hidden_states_internal=hidden_states_internal, + ) + if stage_manager is not None and not stage_manager.is_last_stage(): + return outputs + + hidden_states = outputs[0] + + # v5 addition: logits for the last `logits_to_keep` positions only (`0` = all, the old behaviour). + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/colossalai/shardformer/modeling/_llama_tf4.py b/colossalai/shardformer/modeling/_llama_tf4.py new file mode 100644 index 000000000000..3f9b7b12880c --- /dev/null +++ b/colossalai/shardformer/modeling/_llama_tf4.py @@ -0,0 +1,616 @@ +import math +import warnings +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torch.distributed +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.models.llama.modeling_llama import ( + LlamaForCausalLM, + LlamaForSequenceClassification, + LlamaModel, + StaticCache, + apply_rotary_pos_emb, + repeat_kv, +) +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import all_to_all_comm, gather_sp_output, split_forward_gather_backward +from colossalai.shardformer.layer.utils import is_share_sp_tp, split_batch_zigzag +from colossalai.shardformer.shard import ShardConfig + +from ..layer import ColoAttention, RingAttention, dist_cross_entropy + +_SUPPORTED_SP_MODE = ["all_to_all", "split_gather", "ring", "ring_attn"] + + +class LlamaPipelineForwards: + """ + This class serves as a micro library for forward function substitution of Llama models + under pipeline setting. + """ + + @staticmethod + def llama_model_forward( + self: LlamaModel, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + force_sp_gather: bool = True, # Set to false only when computing cross entropy + ): + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with pipeline parallelism. Setting `use_cache=False`..." + ) + use_cache = False + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + disable_pp = stage_manager is None + # retrieve input_ids and inputs_embeds + if disable_pp or stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + device = hidden_states.device + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + # Support SP + PP + sp_mode = shard_config.sequence_parallelism_mode + sp_group = shard_config.sequence_parallel_process_group + sp_size = shard_config.sequence_parallel_size + # Generating full positions ids for modes that gather sequence before attn + if stage_manager and (sp_mode != "ring_attn" and not stage_manager.is_first_stage()): + seq_length *= sp_size + + past_seen_tokens = 0 + if use_cache: # kept for BC (cache positions) + if not isinstance(past_key_values, StaticCache): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_seen_tokens = past_key_values.get_seq_length() + if cache_position is None: + if isinstance(past_key_values, StaticCache): + raise ValueError("cache_position is a required argument when using StaticCache.") + cache_position = torch.arange(past_seen_tokens, past_seen_tokens + seq_length, device=device) + + seq_length_with_past = seq_length + past_seen_tokens + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + no_split_input = disable_pp or not stage_manager.is_first_stage() + if no_split_input and sp_mode == "ring_attn": + _, attn_kwargs, _ = RingAttention.prepare_varlen_batch(attention_mask, sp_group) + elif shard_config.enable_flash_attention: + mask_shape = (batch_size, 1, seq_length_with_past, seq_length_with_past) + attn_kwargs: dict = ColoAttention.prepare_attn_kwargs( + mask_shape, + hidden_states.dtype, + hidden_states.device, + q_padding_mask=attention_mask, + is_causal=True, + invert=(sp_mode != "ring_attn"), + ) + else: + # Later pipeline stages hold an SP slice, whereas Q/K are gathered + # before attention. Build the mask from the full sequence shape, + # not the local hidden_states length used by _update_causal_mask. + attn_kwargs: torch.Tensor = _prepare_4d_causal_attention_mask( + attention_mask, (batch_size, seq_length), hidden_states, past_seen_tokens + ) + + # Support SP + PP. Later stages have already received the split input. + split_input = disable_pp or stage_manager.is_first_stage() + if split_input: + # Ring Attention zigzag batch processing + if sp_mode == "ring_attn": + assert shard_config.enable_flash_attention, "Ring Attention inherently requires Flash Attention." + if not attention_mask.bool().all(): + hidden_states, attn_kwargs, position_ids = RingAttention.prepare_varlen_batch( + attention_mask, sp_group, hidden_states, position_ids + ) + else: + hidden_states, position_ids = split_batch_zigzag([hidden_states, position_ids], sp_group) + + elif is_share_sp_tp(sp_mode): + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + + if self.gradient_checkpointing and self.training and use_cache: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + start_idx, end_idx = (0, len(self.layers)) if disable_pp else (stage_index[0], stage_index[1]) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + num_ckpt_layers = 0 + if self.gradient_checkpointing and self.training: + num_ckpt_layers = end_idx - start_idx + # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer + if shard_config.gradient_checkpoint_config is not None: + num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( + stage=stage_manager.stage, + num_stages=stage_manager.num_stages, + num_layers=end_idx - start_idx, + model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), + num_model_chunks=stage_manager.num_model_chunks, + ) + assert num_ckpt_layers <= end_idx - start_idx + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if output_hidden_states: + all_hidden_states += (hidden_states,) + if idx - start_idx < num_ckpt_layers: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attn_kwargs, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attn_kwargs, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if disable_pp or stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + if (not shard_config.parallel_output) or force_sp_gather or is_share_sp_tp(sp_mode): # noqa + hidden_states = gather_sp_output(hidden_states, shard_config) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + next_cache = next_decoder_cache if use_cache else None + if disable_pp or stage_manager.is_last_stage(): + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + next_cache, + all_hidden_states, + all_self_attns, + ] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + # always return dict for intermediate stage + return {"hidden_states": hidden_states} + + @staticmethod + def llama_for_causal_lm_forward( + self: LlamaForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, LlamaForCausalLM + + >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + if shard_config.sequence_parallelism_mode == "ring_attn" and shard_config.parallel_output: + # Split labels in a zigzag fashion too + sp_group = shard_config.sequence_parallel_process_group + if attention_mask.bool().all(): + labels = split_batch_zigzag(labels, sp_group, seq_dim=1, is_label=True) + else: + # [B, max_seqlen // sp_size] + labels, _, _ = RingAttention.prepare_varlen_batch(attention_mask, sp_group, labels, is_label=True) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = LlamaPipelineForwards.llama_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + force_sp_gather=False, + ) + past_key_values = None + + disable_pp = stage_manager is None + if disable_pp or stage_manager.is_last_stage(): + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, self.model.dtype) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def llama_for_sequence_classification_forward( + self: LlamaForSequenceClassification, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + transformer_outputs = LlamaPipelineForwards.llama_model_forward( + self.model, + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + ) + + if input_ids is not None: + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + batch_size = inputs_embeds.shape[0] + else: + batch_size = hidden_states.shape[0] + + if stage_manager.is_last_stage(): + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + else: + hidden_states = transformer_outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[Union[torch.Tensor, Dict]] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if sp_mode is not None: + assert sp_mode in _SUPPORTED_SP_MODE, f"SP mode {sp_mode} is not supported by {type(self)} yet" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + # sp: modify sp_len when sequence parallel mode is ring + if is_share_sp_tp(sp_mode): + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + if sp_mode == "ring_attn": + attn_output = RingAttention.attention( + query_states, + key_states, + value_states, + sp_axis=shard_config.sp_axis, + **attention_mask, + inner_ring_size=shard_config.inner_ring_size, + pg_mesh=shard_config.pg_mesh, + ) + + elif shard_config.enable_flash_attention: + assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict." + attn_output = ColoAttention.attention(query_states, key_states, value_states, **attention_mask) + else: + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + # Modern Transformers may reserve one extra cache position in + # its causal mask; native attention slices it to actual keys. + attention_mask = attention_mask[:, :, :, :kv_seq_len] + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) + else: + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + return attn_output, attn_weights + + return forward diff --git a/colossalai/shardformer/modeling/_llama_tf5.py b/colossalai/shardformer/modeling/_llama_tf5.py new file mode 100644 index 000000000000..01013a310592 --- /dev/null +++ b/colossalai/shardformer/modeling/_llama_tf5.py @@ -0,0 +1,658 @@ +import math +import warnings +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torch.distributed +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers.cache_utils import Cache, DynamicCache, StaticCache +from transformers.masking_utils import create_causal_mask +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.models.llama.modeling_llama import ( + LlamaForCausalLM, + LlamaForSequenceClassification, + LlamaModel, + apply_rotary_pos_emb, + repeat_kv, +) +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import all_to_all_comm, gather_sp_output, split_forward_gather_backward +from colossalai.shardformer.layer.utils import is_share_sp_tp, split_batch_zigzag +from colossalai.shardformer.shard import ShardConfig + +from ..layer import ColoAttention, RingAttention, dist_cross_entropy + +_SUPPORTED_SP_MODE = ["all_to_all", "split_gather", "ring", "ring_attn"] + + +def _get_causal_mask_input_embeddings(hidden_states, query_length: int, sequence_is_sharded: bool): + """Return a shape-only full-sequence view when PP receives an SP shard. + + Transformers' ``create_causal_mask`` infers query length, dtype, and device + from ``inputs_embeds``; it does not inspect embedding values. Under SP + (all_to_all, split_gather, ring), later PP stages receive a sequence shard, while + attention first gathers q/k/v back to the full sequence. Expanding one token as a + view gives mask construction the full query length without copying activation data. + """ + if not sequence_is_sharded or hidden_states.shape[1] == query_length: + return hidden_states + return hidden_states[:, :1, :].expand(-1, query_length, -1) + + +class _EagerMaskConfig: + """Read-through view of a model config under which ``create_causal_mask`` builds an eager mask. + + Without flash attention, the policy replaces attention under SP with one that adds the mask to + its scores before an explicit softmax, so the mask must follow eager semantics (additive float) + whatever ``config._attn_implementation`` says: ``sdpa`` yields a boolean mask, ``flash_attention_2`` + ``None`` or a 2D padding mask, ``flex_attention`` a ``BlockMask``. ``create_causal_mask`` picks the + mask format from ``config._attn_implementation`` only, so this view overrides that field and reads + every other attribute from the live config, which is neither copied nor modified. + """ + + _attn_implementation = "eager" + + def __init__(self, config): + self._config = config + + def __getattr__(self, name): + if name == "_config": # not yet set (e.g. during copy / unpickling): avoid infinite recursion + raise AttributeError(name) + return getattr(self._config, name) + + +class LlamaPipelineForwards: + """ + This class serves as a micro library for forward function substitution of Llama models + under pipeline setting. + """ + + @staticmethod + def llama_model_forward( + self: LlamaModel, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + force_sp_gather: bool = True, # Set to false only when computing cross entropy + ): + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + # This forward also replaces `LlamaModel.forward` when PP is off (plain TP / SP / FA / + # inference), so the cache must only be disabled when PP is actually on. + if use_cache and stage_manager is not None: + logger.warning_once( + "`use_cache=True` is incompatible with pipeline parallelism. Setting `use_cache=False`..." + ) + use_cache = False + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + disable_pp = stage_manager is None + # retrieve input_ids and inputs_embeds + if disable_pp or stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape[:2] + elif inputs_embeds is not None: + batch_size, seq_length = inputs_embeds.shape[:2] + else: + raise ValueError("You have to specify either input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + device = hidden_states.device + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + # Support SP + PP + sp_mode = shard_config.sequence_parallelism_mode + sp_group = shard_config.sequence_parallel_process_group + sp_size = shard_config.sequence_parallel_size + # Generating full positions ids for modes that gather sequence before attn + if stage_manager and (sp_mode != "ring_attn" and not stage_manager.is_first_stage()): + seq_length *= sp_size + + past_seen_tokens = 0 + if use_cache: # kept for BC (cache positions) + if not isinstance(past_key_values, StaticCache): + # v5 dropped `from_legacy_cache` / `to_legacy_cache`; `DynamicCache(...)` + # accepts the same per-layer `(key, value)` tuples. + if past_key_values is None: + past_key_values = DynamicCache() + elif not isinstance(past_key_values, Cache): + past_key_values = DynamicCache(past_key_values) + past_seen_tokens = past_key_values.get_seq_length() + if cache_position is None: + if isinstance(past_key_values, StaticCache): + raise ValueError("cache_position is a required argument when using StaticCache.") + cache_position = torch.arange(past_seen_tokens, past_seen_tokens + seq_length, device=device) + + seq_length_with_past = seq_length + past_seen_tokens + + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + # `output_attentions` / `output_hidden_states` stay disabled: the values accumulated + # below are per-rank shards taken before `gather_sp_output`, i.e. not the full sequence. + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + no_split_input = disable_pp or not stage_manager.is_first_stage() + if no_split_input and sp_mode == "ring_attn": + _, attn_kwargs, _ = RingAttention.prepare_varlen_batch(attention_mask, sp_group) + elif shard_config.enable_flash_attention: + mask_shape = (batch_size, 1, seq_length_with_past, seq_length_with_past) + attn_kwargs: dict = ColoAttention.prepare_attn_kwargs( + mask_shape, + hidden_states.dtype, + hidden_states.device, + q_padding_mask=attention_mask, + is_causal=True, + invert=(sp_mode != "ring_attn"), + ) + else: + # `allow_is_causal_skip=False` is required: the attention below uses an explicit + # softmax, so a skipped (`None`) mask would silently attend to every position. + causal_mask_inputs_embeds = _get_causal_mask_input_embeddings( + hidden_states, + seq_length, + sequence_is_sharded=( + stage_manager is not None + and not stage_manager.is_first_stage() + and sp_mode in ("all_to_all", "split_gather", "ring") + ), + ) + # Without flash attention, the policy replaces attention only under SP, with one that adds the + # mask to its scores; HF attention layers interpret the mask per `_attn_implementation`. + attn_kwargs: torch.Tensor = create_causal_mask( + config=_EagerMaskConfig(self.config) if shard_config.enable_sequence_parallelism else self.config, + inputs_embeds=causal_mask_inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + allow_is_causal_skip=False, + ) + + # Support SP + PP. Later stages have already received the split input. + split_input = disable_pp or stage_manager.is_first_stage() + if split_input: + # Ring Attention zigzag batch processing + if sp_mode == "ring_attn": + assert shard_config.enable_flash_attention, "Ring Attention inherently requires Flash Attention." + if not attention_mask.bool().all(): + hidden_states, attn_kwargs, position_ids = RingAttention.prepare_varlen_batch( + attention_mask, sp_group, hidden_states, position_ids + ) + else: + hidden_states, position_ids = split_batch_zigzag([hidden_states, position_ids], sp_group) + + elif is_share_sp_tp(sp_mode): + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + + # v5 resolves the checkpointing / cache conflict inside `GradientCheckpointingLayer`. + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + start_idx, end_idx = (0, len(self.layers)) if disable_pp else (stage_index[0], stage_index[1]) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + num_ckpt_layers = 0 + if self.gradient_checkpointing and self.training: + num_ckpt_layers = end_idx - start_idx + # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer + if shard_config.gradient_checkpoint_config is not None: + num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( + stage=stage_manager.stage, + num_stages=stage_manager.num_stages, + num_layers=end_idx - start_idx, + model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), + num_model_chunks=stage_manager.num_model_chunks, + ) + assert num_ckpt_layers <= end_idx - start_idx + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if output_hidden_states: + all_hidden_states += (hidden_states,) + if self.gradient_checkpointing and self.training: + # v5 checkpoints inside `GradientCheckpointingLayer.__call__`, based on a per-layer + # flag; set it directly to keep ColossalAI's "first N layers only" extension. + decoder_layer.gradient_checkpointing = idx - start_idx < num_ckpt_layers + # Keyword names must match v5's `DecoderLayer.forward`; unknown ones land in `**kwargs`. + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attn_kwargs, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + ) + # v5's `DecoderLayer.forward` returns a bare tensor. + hidden_states = layer_outputs + + if use_cache: + # v5 updates the cache in place instead of returning it. + next_decoder_cache = past_key_values + + if disable_pp or stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + if shard_config.enable_sequence_parallelism and ( + (not shard_config.parallel_output) or force_sp_gather or is_share_sp_tp(sp_mode) + ): + hidden_states = gather_sp_output(hidden_states, shard_config) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + next_cache = next_decoder_cache if use_cache else None + if disable_pp or stage_manager.is_last_stage(): + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + next_cache, + all_hidden_states, + all_self_attns, + ] + if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + # always return dict for intermediate stage + return {"hidden_states": hidden_states} + + @staticmethod + def llama_for_causal_lm_forward( + self: LlamaForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, LlamaForCausalLM + + >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + if shard_config.sequence_parallelism_mode == "ring_attn" and shard_config.parallel_output: + # Split labels in a zigzag fashion too + sp_group = shard_config.sequence_parallel_process_group + if attention_mask.bool().all(): + labels = split_batch_zigzag(labels, sp_group, seq_dim=1, is_label=True) + else: + # [B, max_seqlen // sp_size] + labels, _, _ = RingAttention.prepare_varlen_batch(attention_mask, sp_group, labels, is_label=True) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = LlamaPipelineForwards.llama_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + force_sp_gather=False, + ) + past_key_values = None + + disable_pp = stage_manager is None + if disable_pp or stage_manager.is_last_stage(): + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, self.model.dtype) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def llama_for_sequence_classification_forward( + self: LlamaForSequenceClassification, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + transformer_outputs = LlamaPipelineForwards.llama_model_forward( + self.model, + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + ) + + if input_ids is not None: + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + batch_size = inputs_embeds.shape[0] + else: + batch_size = hidden_states.shape[0] + + if stage_manager.is_last_stage(): + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + else: + hidden_states = transformer_outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[Union[torch.Tensor, Dict]] = None, + # v5 renamed this argument to the plural `past_key_values`; the old name would be + # swallowed by `**kwargs` and the cache silently dropped. + past_key_values: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if sp_mode is not None: + assert sp_mode in _SUPPORTED_SP_MODE, f"SP mode {sp_mode} is not supported by {type(self)} yet" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + bsz, q_len, _ = hidden_states.size() + input_shape = hidden_states.shape[:-1] + # v5's attention has no `num_heads` attribute; it is only written by the TP and + # SP-all_to_all policies, so read it defensively. + num_heads = getattr(self, "num_heads", None) or self.config.num_attention_heads + # sp: modify sp_len when sequence parallel mode is ring + if is_share_sp_tp(sp_mode): + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_values is not None and self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + # v5 removed `Cache.get_usable_length`; `update` returns k/v including history. + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + kv_seq_len = key_states.shape[-2] + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + if sp_mode == "ring_attn": + attn_output = RingAttention.attention( + query_states, + key_states, + value_states, + sp_axis=shard_config.sp_axis, + **attention_mask, + inner_ring_size=shard_config.inner_ring_size, + pg_mesh=shard_config.pg_mesh, + ) + + elif shard_config.enable_flash_attention: + assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict." + attn_output = ColoAttention.attention(query_states, key_states, value_states, **attention_mask) + else: + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + if attn_weights.size() != (bsz, num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, num_heads * self.head_dim) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) + else: + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + return attn_output, attn_weights + + return forward diff --git a/colossalai/shardformer/modeling/_mixtral_tf4.py b/colossalai/shardformer/modeling/_mixtral_tf4.py new file mode 100644 index 000000000000..3d1d01ead731 --- /dev/null +++ b/colossalai/shardformer/modeling/_mixtral_tf4.py @@ -0,0 +1,866 @@ +import warnings +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed import ProcessGroup +from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_causal_attention_mask, + _prepare_4d_causal_attention_mask_for_sdpa, +) +from transformers.models.mixtral.modeling_mixtral import ( + MixtralModel, + MixtralSparseMoeBlock, + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + apply_rotary_pos_emb, + load_balancing_loss_func, +) +from transformers.utils import is_flash_attn_2_available, logging + +from colossalai.lazy import LazyInitContext +from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import ( + all_to_all_comm, + gather_forward_split_backward, + split_forward_gather_backward, +) +from colossalai.shardformer.layer.linear import Linear1D_Col, Linear1D_Row, ParallelModule +from colossalai.shardformer.shard import ShardConfig +from colossalai.shardformer.shard.utils import set_tensors_to_none +from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_group + + +class EPMixtralSparseMoeBlock(ParallelModule): + def __init__(self, *args, **kwargs): + raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}") + + def setup_process_groups( + self, + tp_group: ProcessGroup, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + fp8_communication: bool = False, + use_zbv: bool = False, + ): + assert tp_group is not None + assert moe_dp_group is not None + assert ep_group is not None + + # setup ep group + self.ep_size = dist.get_world_size(ep_group) + self.ep_rank = dist.get_rank(ep_group) + self.ep_group = ep_group + self.fp8_communication = fp8_communication + self.use_zbv = use_zbv + + if self.num_experts % self.ep_size != 0: + raise ValueError("The number of experts must be divisible by the number of expert parallel groups.") + + self.num_experts_per_ep = self.num_experts // self.ep_size + self.expert_start_idx = self.ep_rank * self.num_experts_per_ep + held_experts = self.experts[self.expert_start_idx : self.expert_start_idx + self.num_experts_per_ep] + + set_tensors_to_none(self.experts, exclude=set(held_experts)) + + # setup moe_dp group + self.moe_dp_group = moe_dp_group + self.moe_dp_size = moe_dp_group.size() + + # setup global tp group + self.tp_group = tp_group + if self.tp_group.size() > 1: + for expert in held_experts: + expert.w1 = Linear1D_Col.from_native_module( + expert.w1, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv + ) + expert.w3 = Linear1D_Col.from_native_module( + expert.w3, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv + ) + expert.w2 = Linear1D_Row.from_native_module( + expert.w2, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv + ) + + for p in self.experts.parameters(): + set_moe_tensor_ep_group(p, ep_group) + + @staticmethod + def from_native_module( + module: MixtralSparseMoeBlock, + tp_group: ProcessGroup, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + *args, + **kwargs, + ) -> "EPMixtralSparseMoeBlock": + # TODO: better init + LazyInitContext.materialize(module) + module.__class__ = EPMixtralSparseMoeBlock + fp8_communication = kwargs.get("fp8_communication", False) + module.setup_process_groups(tp_group, moe_dp_group, ep_group, fp8_communication) + return module + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # router_logits: (batch * sequence_length, n_experts) + router_logits = self.gate(hidden_states) + + routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + # we cast back to the input dtype + routing_weights = routing_weights.to(hidden_states.dtype) + + selected_experts = selected_experts.t().reshape(-1) + selected_experts_idx = selected_experts.argsort() + dispatch_states = hidden_states.repeat(self.top_k, 1)[selected_experts_idx] + input_split_sizes = selected_experts.bincount(minlength=self.num_experts) + + output_split_sizes = torch.zeros_like(input_split_sizes) + + dist.all_to_all_single(output_split_sizes, input_split_sizes, group=self.ep_group) + + input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() + output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() + + output_states, _ = all_to_all_uneven( + dispatch_states, + input_split_list, + output_split_list, + self.ep_group, + fp8_communication=self.fp8_communication, + ) + # compute expert output + output_states = EPGradScalerIn.apply(output_states, self.ep_size) + if output_states.size(0) > 0: + if self.num_experts_per_ep == 1: + # no need to split + expert = self.experts[self.expert_start_idx] + output_states = expert.act_fn(expert.w1(output_states)) * expert.w3(output_states) + output_states = expert.w2(output_states) + else: + output_states_splits = output_states.split(output_split_sizes.tolist()) + output_states_list = [] + for i, split_states in enumerate(output_states_splits): + if split_states.size(0) == 0: + continue + expert = self.experts[self.expert_start_idx + i % self.num_experts_per_ep] + split_states = expert.act_fn(expert.w1(split_states)) * expert.w3(split_states) + split_states = expert.w2(split_states) + output_states_list.append(split_states) + output_states = torch.cat(output_states_list) + + output_states = EPGradScalerOut.apply(output_states, self.ep_size) + dispatch_states, _ = all_to_all_uneven( + output_states, output_split_list, input_split_list, self.ep_group, fp8_communication=self.fp8_communication + ) + + recover_experts_idx = torch.empty_like(selected_experts_idx) + recover_experts_idx[selected_experts_idx] = torch.arange( + selected_experts_idx.size(0), device=selected_experts_idx.device + ) + dispatch_states = dispatch_states[recover_experts_idx] + k_hidden_states = dispatch_states.chunk(self.top_k) + output_states = k_hidden_states[0] * routing_weights[:, 0, None] + for i in range(1, self.top_k): + output_states += k_hidden_states[i] * routing_weights[:, i, None] + output_states = output_states.reshape(batch_size, sequence_length, hidden_dim) + return output_states, router_logits + + +class MixtralPipelineForwards: + """ + This class serves as a micro library for forward function substitution of Mixtral models + under pipeline setting. + """ + + @staticmethod + def mixtral_model_forward( + self: MixtralModel, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + past_router_logits: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if stage_manager.is_first_stage(): + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + device = input_ids.device if input_ids is not None else inputs_embeds.device + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + seq_length_with_past = seq_length + past_key_values_length = 0 + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + + if position_ids is None: + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + # embed positions, for the first stage, hidden_states is the input embeddings, + # for the other stages, hidden_states is the output of the previous stage + if is_flash_attn_2_available(): + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + hidden_states, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + hidden_states.shape[1], device=hidden_states.device + ) + + start_idx, end_idx = stage_index[0], stage_index[1] + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + past_key_value = past_key_values[idx] if past_key_values is not None else None + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + # None for past_key_value + return module(*inputs) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(decoder_layer), + hidden_states, + attention_mask, + position_ids, + None, + output_attentions, + output_router_logits, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask, + position_ids, + past_key_value, + output_attentions, + output_router_logits, + use_cache, + cache_position, + position_embeddings, + ) + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = (layer_outputs[2 if output_attentions else 1],) + if output_attentions: + all_self_attns += (layer_outputs[1],) + if output_router_logits: + all_router_logits += (layer_outputs[-1],) + + if stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + next_cache = next_decoder_cache if use_cache else None + + if output_router_logits and past_router_logits is not None: + all_router_logits = past_router_logits + all_router_logits + + if stage_manager.is_last_stage(): + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] + if v is not None + ) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + else: + if output_router_logits: + return { + "hidden_states": hidden_states, + "past_router_logits": all_router_logits, + } + else: + return { + "hidden_states": hidden_states, + } + + @staticmethod + def mixtral_for_causal_lm_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + past_router_logits: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = MixtralPipelineForwards.mixtral_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + past_router_logits=past_router_logits, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func(outputs[-1], self.num_experts, self.num_experts_per_tok) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=None, + hidden_states=outputs[0], + attentions=None, + router_logits=outputs[-1], + ) + else: + out = {} + hidden_states = outputs.get("hidden_states") + out["hidden_states"] = hidden_states + if output_router_logits: + out["past_router_logits"] = outputs["past_router_logits"] + return out + + +def get_mixtral_flash_attention_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None): + logger = logging.get_logger(__name__) + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + from transformers.models.mixtral.modeling_mixtral import eager_attention_forward + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if sp_mode is not None: + assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + # overwrite attention_mask with padding_mask + attention_mask = kwargs.pop("padding_mask") + bsz, q_len, _ = hidden_states.size() + + # sp: modify sp_len when sequence parallel mode is ring + if sp_mode in ["split_gather", "ring"]: + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # Transformers attention interfaces handle grouped-query heads and + # receive tensors in [batch, heads, sequence, head_dim] order. + + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama + **kwargs, + ) + + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim).contiguous() # (1, 8, 128) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) # (1, 4, 256) + else: + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + return attn_output, attn_weights + + return forward + + +def get_mixtral_flash_attention_model_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None): + logger = logging.get_logger(__name__) + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, MoeModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + past_key_values_length = 0 + + if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + if use_cache: + use_legacy_cache = not isinstance(past_key_values, Cache) + if use_legacy_cache: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + past_key_values_length = past_key_values.get_usable_length(seq_length) + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if attention_mask is not None and self.config._attn_implementation == "flash_attention_2" and use_cache: + is_padding_right = attention_mask[:, -1].sum().item() != batch_size + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Mixtral. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if self.config._attn_implementation == "flash_attention_2": + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self.config._attn_implementation == "sdpa" and not output_attentions: + # output_attentions=True can not be supported when using SDPA, and we fall back on + # the manual implementation that requires a 4D causal mask in all cases. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) + + if sp_mode in ["ring", "split_gather"]: + inputs_embeds = split_forward_gather_backward( + inputs_embeds, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + inputs_embeds = split_forward_gather_backward( + inputs_embeds, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + hidden_states = inputs_embeds + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_router_logits = () if output_router_logits else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + output_router_logits, + use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if output_router_logits: + all_router_logits += (layer_outputs[-1],) + + hidden_states = self.norm(hidden_states) + + if sp_mode == "ring" or sp_mode == "split_gather": + hidden_states = gather_forward_split_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = gather_forward_split_backward( + hidden_states, 1, sp_group, grad_scale=sp_size, fp8_communication=shard_config.fp8_communication + ) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache + + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] + if v is not None + ) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + + return forward diff --git a/colossalai/shardformer/modeling/_mixtral_tf5.py b/colossalai/shardformer/modeling/_mixtral_tf5.py new file mode 100644 index 000000000000..3e2f22c482e7 --- /dev/null +++ b/colossalai/shardformer/modeling/_mixtral_tf5.py @@ -0,0 +1,1246 @@ +import inspect +import warnings +from typing import Callable, List, Optional, Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed import ProcessGroup +from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache, DynamicCache +from transformers.masking_utils import create_causal_mask +from transformers.models.mixtral.modeling_mixtral import ( + MixtralModel, + MixtralSparseMoeBlock, + MixtralTopKRouter, + MoeCausalLMOutputWithPast, + MoeModelOutputWithPast, + apply_rotary_pos_emb, + load_balancing_loss_func, +) +from transformers.utils import is_flash_attn_2_available, logging +from transformers.utils.output_capturing import capture_outputs + +from colossalai.lazy import LazyInitContext +from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import ( + all_to_all_comm, + gather_forward_split_backward, + linear_with_async_comm, + reduce_forward, + split_forward_gather_backward, +) +from colossalai.shardformer.layer.linear import Linear1D_Col, LinearWithGradAccum, ParallelModule +from colossalai.shardformer.shard import ShardConfig +from colossalai.tensor.d_tensor import init_tensor_as_customization_distributed +from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_fused, set_moe_tensor_ep_group + +# v4-era flag, read by `get_mixtral_flash_attention_forward` for a one-off warning. It is only +# populated below when flash-attn is actually installed; default to "supported" so the reference +# stays in scope (and no spurious sliding-window warning fires) on flash-attn-free environments. +_flash_supports_window_size = True + +if is_flash_attn_2_available(): + from flash_attn import flash_attn_func + + from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa + + _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) + + +def _slice_fused_expert_param( + tensor: torch.Tensor, + name: str, + expert_start_idx: int, + num_experts_per_ep: int, + num_experts: int, + tp_rank: int, + tp_size: int, + local_shape: torch.Size, +) -> torch.Tensor: + """Select this rank's EP experts and TP columns from a full or EP-local checkpoint tensor.""" + if tensor.shape == local_shape: + return tensor + if tensor.ndim != len(local_shape): + raise ValueError(f"Unexpected {name} rank in the Mixtral expert checkpoint.") + + if tensor.shape[0] == num_experts: + tensor = tensor[expert_start_idx : expert_start_idx + num_experts_per_ep] + elif tensor.shape[0] != num_experts_per_ep: + raise ValueError(f"Unexpected expert dimension in the Mixtral {name} checkpoint.") + + if tp_size > 1: + if name == "gate_up_proj": + if tensor.shape[1] % 2 != 0: + raise ValueError("The fused gate_up_proj dimension must contain an even gate/up pair.") + intermediate_size = tensor.shape[1] // 2 + if intermediate_size % tp_size != 0: + raise ValueError("The fused gate_up_proj dimension must be divisible by tp_size.") + tp_partition_size = intermediate_size // tp_size + tp_start = tp_rank * tp_partition_size + tensor = torch.cat( + ( + tensor.narrow(1, tp_start, tp_partition_size), + tensor.narrow(1, intermediate_size + tp_start, tp_partition_size), + ), + dim=1, + ) + else: + if tensor.shape[2] % tp_size != 0: + raise ValueError("The fused down_proj intermediate dimension must be divisible by tp_size.") + tp_partition_size = tensor.shape[2] // tp_size + tensor = tensor.narrow(2, tp_rank * tp_partition_size, tp_partition_size) + + if tensor.shape != local_shape: + raise ValueError(f"The sliced Mixtral {name} shape does not match the local expert parameter.") + return tensor.contiguous() + + +def _shard_param_in_place(param: torch.nn.Parameter, local: torch.Tensor) -> None: + """Swap a parameter's storage for its local shard while keeping the Parameter object. + + Optimizers are built from the unsharded model and HybridParallel/MoE plugins only prune (never + re-add) optimizer params, so replacing the Parameter object would silently drop it from training. + """ + param.data = local + + +def _register_fused_expert_checkpoint_hook( + experts: torch.nn.Module, + expert_start_idx: int, + num_experts_per_ep: int, + num_experts: int, + tp_rank: int = 0, + tp_size: int = 1, +) -> None: + """Shard full fused-expert parameters before they reach the local EP module.""" + + def _slice_full_expert_params( + module: torch.nn.Module, + state_dict, + prefix: str, + local_metadata, + strict: bool, + missing_keys, + unexpected_keys, + error_msgs, + ) -> None: + del local_metadata, strict, missing_keys, unexpected_keys, error_msgs + for name in ("gate_up_proj", "down_proj"): + key = prefix + name + checkpoint_param = state_dict.get(key) + local_param = getattr(module, name) + if checkpoint_param is None: + continue + if checkpoint_param.shape == local_param.shape: + continue + try: + state_dict[key] = _slice_fused_expert_param( + checkpoint_param, + name, + expert_start_idx, + num_experts_per_ep, + num_experts, + tp_rank, + tp_size, + local_param.shape, + ) + except ValueError: + continue + + experts.register_load_state_dict_pre_hook(_slice_full_expert_params) + + +def _mark_fused_expert_tp_shard( + param: torch.nn.Parameter, + name: str, + tp_group: ProcessGroup, + expert_start_idx: int, + num_experts_per_ep: int, + num_experts: int, +) -> None: + """Attach checkpoint gather/scatter functions to a manually TP-sharded fused expert parameter.""" + tp_rank = dist.get_rank(tp_group) + tp_size = dist.get_world_size(tp_group) + local_shape = param.shape + + def shard_fn(tensor: torch.Tensor) -> torch.Tensor: + return _slice_fused_expert_param( + tensor, + name, + expert_start_idx, + num_experts_per_ep, + num_experts, + tp_rank, + tp_size, + local_shape, + ) + + def gather_fn(tensor: torch.Tensor) -> torch.Tensor: + gathered = [torch.empty_like(tensor) for _ in range(tp_size)] + dist.all_gather(gathered, tensor.contiguous(), group=tp_group) + if name == "gate_up_proj": + partition_size = tensor.shape[1] // 2 + gate = torch.cat([shard[:, :partition_size] for shard in gathered], dim=1) + up = torch.cat([shard[:, partition_size:] for shard in gathered], dim=1) + return torch.cat((gate, up), dim=1) + return torch.cat(gathered, dim=2) + + init_tensor_as_customization_distributed(param, shard_fn, gather_fn) + + +class _MixtralTopKRouterMixin: + r"""Sharded / grad-accumulation variants of the v5 ``MixtralTopKRouter``. + + Only the matmul is replaced; the post-processing matches v5's + ``MixtralTopKRouter.forward`` line for line. + + Args: + hidden_dim (int): second dimension of the router weight, used to reshape the input + into ``(N, hidden_dim)``. + top_k (int): number of experts selected per token. + + v5 collects `router_logits` declaratively: `MixtralModel._can_record_outputs` declares + `OutputRecorder(MixtralTopKRouter, index=0)` and the capturing hook is matched with + `isinstance` (`output_capturing.py:165`). Both subclasses below therefore also inherit from + `MixtralTopKRouter`, so the hook attaches to the sharded router as well and records the first + element of its (logits, scores, indices) triple. No extra communication is needed for that: + `Linear1D_Col` runs with `gather_output=True`, so the expert dimension of the logits is already + gathered back on every rank. + """ + + def forward(self, hidden_states: torch.Tensor): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = super().forward(hidden_states) + router_probs = F.softmax(router_logits.float(), dim=-1) + router_top_value, router_indices = torch.topk(router_probs, self.top_k, dim=-1) + router_top_value /= router_top_value.sum(dim=-1, keepdim=True) + return router_logits, router_top_value, router_indices + + +# `MixtralTopKRouter` is the last base on purpose: it is only there for the `isinstance` match of +# v5's capturing hook. The mixin must come first so that `super().forward()` inside it reaches +# `Linear1D_Col.forward` (sharded matmul + gather); putting the native router before the parallel +# linear would silently take v5's `F.linear` path instead, i.e. topk over a single rank's experts. +class MixtralTopKRouter1D(_MixtralTopKRouterMixin, Linear1D_Col, MixtralTopKRouter): + r"""Tensor-parallel Mixtral router. + + v5 replaced ``MixtralSparseMoeBlock.gate`` with ``MixtralTopKRouter`` and unpacks a triple + from it:: + + _, top_k_weights, top_k_index = self.gate(hidden_states) + + ``Linear1D_Col`` returns a single tensor, so substituting it directly would silently compute + the wrong thing: unpacking an ``[N, E]`` output iterates over dim 0 (the tokens), which raises + only when ``N == 3``. This class keeps the weight sharded over the expert dimension while + matching v5's forward contract. + + The weight stays under the name ``weight`` (the sharding writes back in place), so the state + dict key ``...mlp.gate.weight`` matches v5's HF key. + + Inherits from `MixtralTopKRouter` so that v5's output-capturing hook recognises it; see the + mixin docstring for why that is the only requirement (the logits are already gathered). + """ + + @staticmethod + def from_native_module( + module: "MixtralTopKRouter", process_group: ProcessGroup = None, **kwargs + ) -> "MixtralTopKRouter1D": + r"""Convert a native ``MixtralTopKRouter`` to a tensor-parallel one. + + Args: + module (MixtralTopKRouter): the native router to be converted. + process_group (ProcessGroup): the process group of tensor parallelism. + **kwargs: passed to ``Linear1D_Col`` (e.g. ``fp8_communication`` / ``use_zbv``). + """ + LazyInitContext.materialize(module) + + top_k, hidden_dim = module.top_k, module.hidden_dim + num_experts, in_features = module.weight.shape + + # Under sequence parallelism (split_gather / ring / all_to_all) the sequence — not the batch + # — is split across the TP ranks, so the standard expert-dim gather below would pair each + # local token with a *different* rank's local token. The router weight is tiny + # ([num_experts, hidden]), so keep it unsharded: the native forward already computes the full + # logits for the local tokens, and no gather is needed. + if kwargs.get("seq_parallel_mode", None) is not None: + return module + + tp_size = dist.get_world_size(process_group) if process_group is not None else 1 + if num_experts < tp_size: + # Same as `Linear1D_Col.from_native_module` (linear.py:301-303): skip sharding when it + # would not divide evenly. + return module + + # `Linear1D_Col` shards by writing back in place (`sharded_tensor_to_existing_param`), so + # the router's own `Parameter` ends up holding the shard. + shim = torch.nn.Linear(in_features, num_experts, bias=False) + shim.weight = module.weight + + kwargs = dict(kwargs) + kwargs.setdefault("gather_output", True) + router = Linear1D_Col.from_native_module(shim, process_group, **kwargs) + router.__class__ = MixtralTopKRouter1D + router.top_k = top_k + router.hidden_dim = hidden_dim + return router + + +class MixtralTopKRouterWithGradAccum(_MixtralTopKRouterMixin, LinearWithGradAccum, MixtralTopKRouter): + r"""Mixtral router for Zero-Bubble-V (``use_zbv``, no TP). + + Same reasoning as ``MixtralTopKRouter1D``: the ZBV branch used to substitute + ``LinearWithGradAccum`` for the gate, which breaks v5's triple contract in the same way. Only + the matmul is swapped, and nothing is sharded, so the logits keep v5's shapes. + + Inherits from `MixtralTopKRouter` so that v5's output-capturing hook recognises it; see the + mixin docstring for why that is the only requirement (nothing is sharded here, so the logits + are complete by construction). + """ + + @staticmethod + def from_native_module(module: "MixtralTopKRouter", **kwargs) -> "MixtralTopKRouterWithGradAccum": + r"""Convert a native ``MixtralTopKRouter`` to a grad-accumulation one.""" + LazyInitContext.materialize(module) + + top_k, hidden_dim = module.top_k, module.hidden_dim + num_experts, in_features = module.weight.shape + + shim = torch.nn.Linear(in_features, num_experts, bias=False) + shim.weight = module.weight + + router = LinearWithGradAccum.from_native_module(shim, **kwargs) + router.__class__ = MixtralTopKRouterWithGradAccum + router.top_k = top_k + router.hidden_dim = hidden_dim + return router + + +class EPMixtralSparseMoeBlock(ParallelModule): + def __init__(self, *args, **kwargs): + raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}") + + def setup_process_groups( + self, + tp_group: ProcessGroup, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + fp8_communication: bool = False, + use_zbv: bool = False, + ): + assert tp_group is not None + assert moe_dp_group is not None + assert ep_group is not None + + # setup ep group + self.ep_size = dist.get_world_size(ep_group) + self.ep_rank = dist.get_rank(ep_group) + self.ep_group = ep_group + self.fp8_communication = fp8_communication + self.use_zbv = use_zbv + + # v5 stores experts as fused 3D parameters (`gate_up_proj [E,2I,H]` / `down_proj [E,H,I]`), so + # `num_experts` lives on `self.experts`, not on the block (docs/30 §二). + num_experts = self.experts.num_experts + if num_experts % self.ep_size != 0: + raise ValueError("The number of experts must be divisible by the number of expert parallel groups.") + + self.num_experts = num_experts + self.num_experts_per_ep = num_experts // self.ep_size + self.expert_start_idx = self.ep_rank * self.num_experts_per_ep + + # primitive ① (docs/30 §4.1): slice the fused params to the local experts and release the rest + # (P5). `.clone()` is required, not `.contiguous()`: the dim-0 slice of a contiguous tensor is + # already contiguous, so `.contiguous()` would return the view and keep the full `[E, ...]` alive. + experts = self.experts + s = self.expert_start_idx + n = self.num_experts_per_ep + _shard_param_in_place(experts.gate_up_proj, experts.gate_up_proj.data[s : s + n].clone()) + _shard_param_in_place(experts.down_proj, experts.down_proj.data[s : s + n].clone()) + experts.num_experts = n + + # setup moe_dp group + self.moe_dp_group = moe_dp_group + self.moe_dp_size = moe_dp_group.size() + + # setup global tp group + self.tp_group = tp_group + self.tp_size = tp_group.size() + self.tp_rank = dist.get_rank(tp_group) + if self.tp_size > 1: + gate_up_proj = experts.gate_up_proj.data + down_proj = experts.down_proj.data + if gate_up_proj.shape[1] % 2 != 0: + raise ValueError("The fused gate_up_proj dimension must contain an even gate/up pair.") + intermediate_size = gate_up_proj.shape[1] // 2 + if intermediate_size != down_proj.shape[2]: + raise ValueError("The fused gate_up_proj and down_proj intermediate dimensions must match.") + if intermediate_size % self.tp_size != 0: + raise ValueError("The fused intermediate dimension must be divisible by tp_size.") + if down_proj.shape[2] % self.tp_size != 0: + raise ValueError("The fused down_proj intermediate dimension must be divisible by tp_size.") + + gate_up_partition_size = intermediate_size // self.tp_size + down_partition_size = down_proj.shape[2] // self.tp_size + tp_start = self.tp_rank * gate_up_partition_size + _shard_param_in_place( + experts.gate_up_proj, + torch.cat( + ( + gate_up_proj[:, tp_start : tp_start + gate_up_partition_size], + gate_up_proj[ + :, intermediate_size + tp_start : intermediate_size + tp_start + gate_up_partition_size + ], + ), + dim=1, + ).clone(), + ) + tp_start = self.tp_rank * down_partition_size + _shard_param_in_place(experts.down_proj, down_proj[:, :, tp_start : tp_start + down_partition_size].clone()) + + # These manual slices are not DTensors; checkpoint sharding needs explicit TP gather metadata. + _mark_fused_expert_tp_shard(experts.gate_up_proj, "gate_up_proj", tp_group, s, n, num_experts) + _mark_fused_expert_tp_shard(experts.down_proj, "down_proj", tp_group, s, n, num_experts) + + _register_fused_expert_checkpoint_hook(experts, s, n, num_experts, tp_rank=self.tp_rank, tp_size=self.tp_size) + + # primitive ② (docs/30 §4.2): mark the sliced fused params so the sharded loader slices dim 0 + # (the expert dimension) by `ep_group`. + for p in self.experts.parameters(): + set_moe_tensor_ep_group(p, ep_group) + set_moe_tensor_ep_fused(p) + + @staticmethod + def from_native_module( + module: MixtralSparseMoeBlock, + tp_group: ProcessGroup, + moe_dp_group: ProcessGroup, + ep_group: ProcessGroup, + *args, + **kwargs, + ) -> "EPMixtralSparseMoeBlock": + # TODO: better init + LazyInitContext.materialize(module) + module.__class__ = EPMixtralSparseMoeBlock + fp8_communication = kwargs.get("fp8_communication", False) + use_zbv = kwargs.get("use_zbv", False) + module.setup_process_groups(tp_group, moe_dp_group, ep_group, fp8_communication, use_zbv) + return module + + def _expert_forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor: + r"""Run one local expert on its dispatched tokens (docs/30 §4.3). + + Mirrors v5's ``MixtralExperts.forward``: ``gate, up = linear(x, gate_up_proj[e]).chunk(2, + dim=-1)``, ``act_fn(gate) * up``, then ``linear(down_proj[e])``. The fused params are already + sliced to ``[E/ep, ...]``, so ``expert_idx`` is a local index. With TP over experts, each local + ``gate_up_proj`` keeps matching gate/up slices and ``down_proj`` keeps the corresponding input slice; + the partial output is then reduced. + """ + gate_up_weight = self.experts.gate_up_proj[expert_idx] + if self.tp_size == 1: + gate_up = F.linear(x, gate_up_weight) + else: + gate_up = linear_with_async_comm( + x, + gate_up_weight, + None, + self.tp_group, + async_grad_allreduce=True, + fp8_communication=self.fp8_communication, + use_zbv=self.use_zbv, + ) + gate, up = gate_up.chunk(2, dim=-1) + output = F.linear(self.experts.act_fn(gate) * up, self.experts.down_proj[expert_idx]) + if self.tp_size > 1: + output = reduce_forward(output, self.tp_group, fp8_communication=self.fp8_communication) + return output + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + # v5's `MixtralTopKRouter` returns the (logits, normalised top-k weights, top-k indices) triple; + # softmax / top-k / normalisation are already applied, so there is no local routing to redo here. + _, top_k_weights, top_k_index = self.gate(hidden_states) + routing_weights = top_k_weights.to(hidden_states.dtype) + + selected_experts = top_k_index.t().reshape(-1) + selected_experts_idx = selected_experts.argsort() + dispatch_states = hidden_states.repeat(self.top_k, 1)[selected_experts_idx] + input_split_sizes = selected_experts.bincount(minlength=self.num_experts) + + output_split_sizes = torch.zeros_like(input_split_sizes) + + dist.all_to_all_single(output_split_sizes, input_split_sizes, group=self.ep_group) + + input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() + output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() + + output_states, _ = all_to_all_uneven( + dispatch_states, + input_split_list, + output_split_list, + self.ep_group, + fp8_communication=self.fp8_communication, + ) + # compute expert output + output_states = EPGradScalerIn.apply(output_states, self.ep_size) + if output_states.size(0) > 0: + if self.num_experts_per_ep == 1: + # no need to split + output_states = self._expert_forward(output_states, 0) + else: + output_states_splits = output_states.split(output_split_sizes.tolist()) + output_states_list = [] + for i, split_states in enumerate(output_states_splits): + if split_states.size(0) == 0: + continue + expert_idx = i % self.num_experts_per_ep + split_states = self._expert_forward(split_states, expert_idx) + output_states_list.append(split_states) + output_states = torch.cat(output_states_list) + + output_states = EPGradScalerOut.apply(output_states, self.ep_size) + dispatch_states, _ = all_to_all_uneven( + output_states, output_split_list, input_split_list, self.ep_group, fp8_communication=self.fp8_communication + ) + + recover_experts_idx = torch.empty_like(selected_experts_idx) + recover_experts_idx[selected_experts_idx] = torch.arange( + selected_experts_idx.size(0), device=selected_experts_idx.device + ) + dispatch_states = dispatch_states[recover_experts_idx] + k_hidden_states = dispatch_states.chunk(self.top_k) + output_states = k_hidden_states[0] * routing_weights[:, 0, None] + for i in range(1, self.top_k): + output_states += k_hidden_states[i] * routing_weights[:, i, None] + output_states = output_states.reshape(batch_size, sequence_length, hidden_dim) + return output_states + + +class MixtralPipelineForwards: + """ + This class serves as a micro library for forward function substitution of Mixtral models + under pipeline setting. + """ + + @staticmethod + def mixtral_model_forward( + self: MixtralModel, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + past_router_logits: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if stage_manager.is_first_stage(): + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + device = input_ids.device if input_ids is not None else inputs_embeds.device + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + seq_length_with_past = seq_length + past_key_values_length = 0 + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + if past_key_values is not None: + # v5's `past_key_values` is a `Cache` object and not subscriptable; the old + # `past_key_values[0][0].shape[2]` only worked on v4's legacy tuple. + past_key_values_length = past_key_values.get_seq_length() + seq_length_with_past = seq_length_with_past + past_key_values_length + + if position_ids is None: + position_ids = torch.arange( + past_key_values_length, + seq_length + past_key_values_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + # embed positions, for the first stage, hidden_states is the input embeddings, + # for the other stages, hidden_states is the output of the previous stage + if self.config._attn_implementation == "flash_attention_2": + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + else: + # v5 deprecates `_prepare_4d_causal_attention_mask*`; `create_causal_mask` covers both + # the sdpa and the eager branch. This must follow the actual configured + # backend, not merely whether FlashAttention is installed. + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + allow_is_causal_skip=False, + ) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + hidden_states.shape[1], device=hidden_states.device + ) + + start_idx, end_idx = stage_index[0], stage_index[1] + + # v5 collects router logits declaratively and activates the collector per call: wrapping only + # this stage's layer loop in `capture_outputs` yields exactly this stage's logits, which + # `past_router_logits` then accumulates across stages (as in v4). + @capture_outputs + def _local_layers_forward( + _model, + hidden_states, + *, + output_hidden_states=False, + output_attentions=False, + output_router_logits=False, + ): + r"""Run this stage's layers. `_model` only exists so that `capture_outputs` can reach + `_can_record_outputs` and install the hooks; the closure's `self` does the work.""" + local_hidden_states = () if output_hidden_states else None + local_self_attns = () if output_attentions else None + next_decoder_cache = None + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if output_hidden_states: + local_hidden_states += (hidden_states,) + + # The cache is updated in place per layer; gradient checkpointing is handled by + # `GradientCheckpointingLayer.__call__`. + if self.gradient_checkpointing and self.training: + decoder_layer.gradient_checkpointing = True + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + ) + # v5's `DecoderLayer.forward` returns a bare tensor (v4 returned a tuple) + hidden_states = layer_outputs + + if use_cache: + next_decoder_cache = past_key_values + + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_decoder_cache, + hidden_states=local_hidden_states, + attentions=local_self_attns, + ) + + local_outputs = _local_layers_forward( + self, + hidden_states, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_router_logits=output_router_logits, + ) + hidden_states = local_outputs.last_hidden_state + all_hidden_states = local_outputs.hidden_states + all_self_attns = local_outputs.attentions + next_decoder_cache = local_outputs.past_key_values + # `capture_outputs` writes `outputs[key] = tuple(collected)` unconditionally, so "nothing + # collected" is an empty tuple, not None -- and `load_balancing_loss_func((), ...)` raises + # `IndexError` on `gate_logits[0]`. Normalise empty to None here; None returns 0 there. + local_router_logits = tuple(local_outputs.router_logits) if local_outputs.router_logits else None + all_router_logits = local_router_logits if output_router_logits else None + + if stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + next_cache = next_decoder_cache if use_cache else None + + if output_router_logits and past_router_logits is not None: + # `all_router_logits` is None when this stage collected none (see above). + all_router_logits = past_router_logits + (all_router_logits or ()) + + if stage_manager.is_last_stage(): + if not return_dict: + return tuple( + v + for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] + if v is not None + ) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + router_logits=all_router_logits, + ) + else: + if output_router_logits: + return { + "hidden_states": hidden_states, + "past_router_logits": all_router_logits, + } + else: + return { + "hidden_states": hidden_states, + } + + @staticmethod + def mixtral_for_causal_lm_forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + past_router_logits: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = MixtralPipelineForwards.mixtral_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_router_logits=output_router_logits, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + past_router_logits=past_router_logits, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + loss_fct = CrossEntropyLoss() + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Enable model parallelism + shift_labels = shift_labels.to(shift_logits.device) + loss = loss_fct(shift_logits, shift_labels) + + # Read by field name, not by position: v5's `ModelOutput` only stores non-None fields, + # so `outputs[-1]` silently degenerates to `hidden_states` whenever `router_logits` is + # absent (the field-level read yields the class default `None` instead, which is what + # `load_balancing_loss_func` wants). + if output_router_logits: + # `mixtral_model_forward` returns a tuple when `return_dict=False`; in that + # form `router_logits` is the final non-None entry (see its tuple construction + # above). Field access is only valid for the ModelOutput form. + router_logits = outputs.router_logits if return_dict else outputs[-1] + else: + router_logits = None + + aux_loss = None + if output_router_logits: + aux_loss = load_balancing_loss_func(router_logits, self.num_experts, self.num_experts_per_tok) + if labels is not None: + loss += self.router_aux_loss_coef * aux_loss + + if not return_dict: + output = (logits,) + outputs[1:] + if output_router_logits: + output = (aux_loss,) + output + return (loss,) + output if loss is not None else output + + return MoeCausalLMOutputWithPast( + loss=loss, + aux_loss=aux_loss, + logits=logits, + past_key_values=None, + hidden_states=outputs[0], + attentions=None, + router_logits=router_logits, + ) + else: + out = {} + hidden_states = outputs.get("hidden_states") + out["hidden_states"] = hidden_states + if output_router_logits: + out["past_router_logits"] = outputs["past_router_logits"] + return out + + +def get_mixtral_flash_attention_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None): + logger = logging.get_logger(__name__) + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + from transformers.models.mixtral.modeling_mixtral import eager_attention_forward + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + # Keep the plural name: v5's call sites pass `past_key_values=`, so a singular parameter + # would swallow it into `**kwargs` and silently disable the KV cache. + past_key_values: Optional[Cache] = None, + output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, + use_cache: bool = False, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: + if sp_mode is not None: + assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + + if "padding_mask" in kwargs: + warnings.warn( + "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" + ) + + # overwrite attention_mask with padding_mask + attention_mask = kwargs.pop("padding_mask") + bsz, q_len, _ = hidden_states.size() + # v5's `MixtralAttention` has no `num_heads` / `num_key_value_heads` attributes of its own + # and reads the config directly. Only the TP and SP-all_to_all policies write those names, + # so a plain SP run would hit an `AttributeError` on `self.num_heads` -- hence the fallback. + num_heads = getattr(self, "num_heads", None) or self.config.num_attention_heads + num_key_value_heads = getattr(self, "num_key_value_heads", None) or self.config.num_key_value_heads + + # sp: modify sp_len when sequence parallel mode is ring + if sp_mode in ["split_gather", "ring"]: + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = query_states.view(bsz, q_len, num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, num_key_value_heads, self.head_dim).transpose(1, 2) + if past_key_values is not None and self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if not _flash_supports_window_size: + logger.warning_once( + "The current flash attention version does not support sliding window attention, for a more memory efficient implementation" + " make sure to upgrade flash-attn library." + ) + if past_key_values is not None: + # v5 dropped `Cache.get_usable_length`; `update` now returns the full (history + # included) k/v, matching v5's own `MixtralAttention`. + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + + # v5's `eager_attention_forward` / `sdpa_attention_forward` both do their own `repeat_kv` + # (keyed on `num_key_value_groups`) and expect q/k/v in `[bsz, heads, q_len, head_dim]`. + # The v4 leftovers here -- an explicit `repeat_kv` plus a `transpose(1, 2)` before the + # interface -- double-count the head dim and shift the sequence into the head axis. + # In PEFT, usually we cast the layer norms in float32 for training stability reasons + # therefore the input hidden states gets silently casted in float32. Hence, we need + # cast them back in float16 just to be sure everything works as expected. + input_dtype = query_states.dtype + if input_dtype == torch.float32: + if torch.is_autocast_enabled(): + target_dtype = torch.get_autocast_gpu_dtype() + # Handle the case where the model is quantized + elif hasattr(self.config, "_pre_quantization_dtype"): + target_dtype = self.config._pre_quantization_dtype + else: + target_dtype = self.q_proj.weight.dtype + + logger.warning_once( + f"The input hidden states seems to be silently casted in float32, this might be related to" + f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" + f" {target_dtype}." + ) + + query_states = query_states.to(target_dtype) + key_states = key_states.to(target_dtype) + value_states = value_states.to(target_dtype) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama + **kwargs, + ) + + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, num_heads * self.head_dim).contiguous() # (1, 8, 128) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) # (1, 4, 256) + else: + attn_output = attn_output.reshape(bsz, q_len, num_heads * self.head_dim) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + return attn_output, attn_weights + + return forward + + +def get_mixtral_flash_attention_model_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None): + logger = logging.get_logger(__name__) + + # v5 collects `router_logits` declaratively (`OutputRecorder(MixtralTopKRouter, index=0)` in + # `MixtralModel._can_record_outputs`) and injects it via `@capture_outputs`. v5's own forward + # carries that decorator and ours replaces it, so it has to be put back here. + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + output_router_logits: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, MoeModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_router_logits = ( + output_router_logits if output_router_logits is not None else self.config.output_router_logits + ) + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + past_key_values_length = 0 + + if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + if use_cache: + # v5 dropped `DynamicCache.from_legacy_cache` / `to_legacy_cache`, but the + # `ddp_cache_data` argument of `DynamicCache.__init__` is exactly the per-layer + # `(key, value[, sliding_window])` tuples, so constructing it is the old semantics. + if past_key_values is None: + past_key_values = DynamicCache() + elif not isinstance(past_key_values, Cache): + past_key_values = DynamicCache(past_key_values) + # v5 dropped `Cache.get_usable_length` in favour of `get_seq_length`. + past_key_values_length = past_key_values.get_seq_length() + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if attention_mask is not None and self.config._attn_implementation == "flash_attention_2" and use_cache: + is_padding_right = attention_mask[:, -1].sum().item() != batch_size + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Mixtral. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if self.config._attn_implementation == "flash_attention_2": + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + else: + # v5 deprecates `modeling_attn_mask_utils._prepare_4d_causal_attention_mask*` in + # favour of `create_causal_mask`, which also folds in the sdpa / eager split and + # sliding window. `allow_is_causal_skip=False` because this branch ends in explicit + # softmax rather than delegating causality to the backend. + attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + allow_is_causal_skip=False, + ) + + if sp_mode in ["ring", "split_gather"]: + inputs_embeds = split_forward_gather_backward( + inputs_embeds, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + inputs_embeds = split_forward_gather_backward( + inputs_embeds, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + hidden_states = inputs_embeds + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + # v5 hands gradient checkpointing to `GradientCheckpointingLayer.__call__`, but + # `_gradient_checkpointing_func` is still required (it is called there) and is only + # installed by going through `gradient_checkpointing_enable()` -- hence the flag. + if self.gradient_checkpointing and self.training: + decoder_layer.gradient_checkpointing = True + + # Keyword names follow v5's `MixtralDecoderLayer.forward`. The old positional call + # fails outright (7 args against 6 parameters), and would misalign `position_embeddings` + # even if the count matched, since v5 puts it second. + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + ) + + # v5's `DecoderLayer.forward` returns a bare tensor (v4 returned a tuple) + hidden_states = layer_outputs + + if use_cache: + # v5 updates the passed-in `Cache` in place instead of returning it + next_decoder_cache = past_key_values + + hidden_states = self.norm(hidden_states) + + if sp_mode == "ring" or sp_mode == "split_gather": + hidden_states = gather_forward_split_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = gather_forward_split_backward( + hidden_states, 1, sp_group, grad_scale=sp_size, fp8_communication=shard_config.fp8_communication + ) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = None + if use_cache: + # v5 dropped `to_legacy_cache`: the cache comes back as a `Cache` object, as v5 + # itself does, no longer downgraded to a legacy tuple to match the input. + next_cache = next_decoder_cache + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + # Left to the `@capture_outputs` decorator (see the top of this function). + router_logits=None, + ) + + return forward diff --git a/colossalai/shardformer/modeling/_qwen2_tf4.py b/colossalai/shardformer/modeling/_qwen2_tf4.py new file mode 100644 index 000000000000..80dfb003dccf --- /dev/null +++ b/colossalai/shardformer/modeling/_qwen2_tf4.py @@ -0,0 +1,841 @@ +import math +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_causal_attention_mask, + _prepare_4d_causal_attention_mask_for_sdpa, +) +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2Attention, + Qwen2ForCausalLM, + Qwen2ForSequenceClassification, + Qwen2Model, + apply_rotary_pos_emb, + repeat_kv, +) +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import all_to_all_comm, split_forward_gather_backward +from colossalai.shardformer.shard import ShardConfig + +from ..layer import ColoAttention, dist_cross_entropy +from ..layer._operation import gather_sp_output +from ..layer.utils import is_share_sp_tp + + +class Qwen2PipelineForwards: + """ + This class serves as a micro library for forward function substitution of Qwen2 models + under pipeline setting. + """ + + @staticmethod + def qwen2_model_forward( + self: Qwen2Model, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + force_sp_output_gather: bool = True, + ) -> Union[Tuple, BaseModelOutputWithPast]: + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + seq_length_with_past = seq_length + past_key_values_length = 0 + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + # assert past_key_values is None, "past_key_values is not supported for Qwen2 models at the moment." + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + + # Support SP + PP + sp_size = shard_config.sequence_parallel_size + sp_group = shard_config.sequence_parallel_process_group + sp_mode = shard_config.sequence_parallelism_mode + # For generating full positions ids (the states will be gathered along the seq dim before attention fwd). + if shard_config.enable_sequence_parallelism and sp_mode != "ring_attn" and not stage_manager.is_first_stage(): + seq_length *= sp_size + # Incoming pipeline states contain only the local SP slice. Both mask + # dimensions must describe the full sequence used by attention. + seq_length_with_past = seq_length + past_key_values_length + + if position_ids is None: + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + # embed positions, for the first stage, hidden_states is the input embeddings, + # for the other stages, hidden_states is the output of the previous stage + if shard_config.enable_flash_attention: + # in this case, attention_mask is a dict rather than a tensor + mask_shape = (batch_size, 1, seq_length, seq_length_with_past) + attention_mask = ColoAttention.prepare_attn_kwargs( + mask_shape, + hidden_states.dtype, + hidden_states.device, + q_padding_mask=attention_mask, + is_causal=True, + ) + else: + if self.config._attn_implementation == "flash_attention_2": + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self.config._attn_implementation == "sdpa" and not output_attentions: + # output_attentions=True can not be supported when using SDPA, and we fall back on + # the manual implementation that requires a 4D causal mask in all cases. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (batch_size, seq_length), + hidden_states, + past_key_values_length, + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + hidden_states, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) + + if stage_manager.is_first_stage(): + if shard_config.enable_sequence_parallelism: + if is_share_sp_tp(sp_mode): + hidden_states = split_forward_gather_backward( + hidden_states, + dim=1, + process_group=sp_group, + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, + dim=1, + process_group=sp_group, + grad_scale=1 / sp_size, + ) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + start_idx, end_idx = stage_index[0], stage_index[1] + num_ckpt_layers = 0 + if self.gradient_checkpointing and self.training: + num_ckpt_layers = end_idx - start_idx + # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer + if shard_config.gradient_checkpoint_config is not None: + num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( + stage=stage_manager.stage, + num_stages=stage_manager.num_stages, + num_layers=end_idx - start_idx, + model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), + num_model_chunks=stage_manager.num_model_chunks, + ) + assert num_ckpt_layers <= end_idx - start_idx + + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + past_key_values[idx] if past_key_values is not None else None + + if idx - start_idx < num_ckpt_layers: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if stage_manager.is_last_stage(): + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + # always return dict for imediate stage + return {"hidden_states": hidden_states} + + @staticmethod + def qwen2_for_causal_lm_forward( + self: Qwen2ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you consciours? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = Qwen2PipelineForwards.qwen2_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + force_sp_output_gather=False, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + hidden_states = outputs[0] + if hidden_states.shape[1] == 2: + pass + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def qwen2_for_sequence_classification_forward( + self: Qwen2ForSequenceClassification, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + transformer_outputs = Qwen2PipelineForwards.qwen2_model_forward( + self.model, + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + ) + + if input_ids is not None: + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + batch_size = inputs_embeds.shape[0] + else: + batch_size = hidden_states.shape[0] + + if stage_manager.is_last_stage(): + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(pooled_logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(pooled_logits, labels) + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + else: + hidden_states = transformer_outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_qwen2_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + def forward( + self: Qwen2Attention, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if sp_mode is not None: + assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + + bsz, q_len, _ = hidden_states.size() + # sp: modify sp_len when sequence parallel mode is ring + if sp_mode in ["split_gather", "ring"]: + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # Activate slicing cache only if the config has a value `sliding_windows` attribute + cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 + if ( + getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and cache_has_contents + ): + slicing_tokens = 1 - self.config.sliding_window + + past_key = past_key_value[self.layer_idx][0] + past_value = past_key_value[self.layer_idx][1] + + past_key = past_key[:, :, slicing_tokens:, :].contiguous() + past_value = past_value[:, :, slicing_tokens:, :].contiguous() + + if past_key.shape[-2] != self.config.sliding_window - 1: + raise ValueError( + f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" + f" {past_key.shape}" + ) + + if attention_mask is not None: + attention_mask = attention_mask[:, slicing_tokens:] + attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) + + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + if shard_config.enable_flash_attention: + assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict." + attn_output = ColoAttention.attention(query_states, key_states, value_states, **attention_mask) + else: + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + attn_output = attn_output.transpose(1, 2).contiguous() + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) + else: + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output, None + + return forward + + +def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + logger = logging.get_logger(__name__) + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + force_sp_output_gather: bool = True, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + seq_length_with_past = seq_length + past_key_values_length = 0 + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + # embed positions + hidden_states = inputs_embeds + + if shard_config.enable_flash_attention: + # in this case, attention_mask is a dict rather than a tensor + mask_shape = (batch_size, 1, seq_length, seq_length_with_past) + attention_mask = ColoAttention.prepare_attn_kwargs( + mask_shape, + hidden_states.dtype, + hidden_states.device, + q_padding_mask=attention_mask, + is_causal=True, + ) + else: + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) + + if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + if sp_mode in ["ring", "split_gather"]: + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + return forward + + +def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig): + def forward( + self: Qwen2ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + force_sp_output_gather=False, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return forward diff --git a/colossalai/shardformer/modeling/_qwen2_tf5.py b/colossalai/shardformer/modeling/_qwen2_tf5.py new file mode 100644 index 000000000000..d9edd942f401 --- /dev/null +++ b/colossalai/shardformer/modeling/_qwen2_tf5.py @@ -0,0 +1,618 @@ +from functools import partial +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2Attention, + Qwen2ForCausalLM, + Qwen2ForSequenceClassification, + Qwen2Model, + apply_rotary_pos_emb, + repeat_kv, +) +from transformers.utils import logging +from transformers.utils.generic import merge_with_config_defaults +from transformers.utils.output_capturing import capture_outputs + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import all_to_all_comm, split_forward_gather_backward +from colossalai.shardformer.shard import ShardConfig + +from .._utils import get_hf_cache_seq_length, get_hf_flash_attention_kwargs, prepare_hf_causal_masks +from ..layer import ColoAttention, dist_cross_entropy +from ..layer._operation import gather_sp_output +from ..layer.utils import is_share_sp_tp + + +class Qwen2PipelineForwards: + """ + This class serves as a micro library for forward function substitution of Qwen2 models + under pipeline setting. + """ + + @staticmethod + def qwen2_model_forward( + self: Qwen2Model, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + force_sp_output_gather: bool = True, + ) -> Union[Tuple, BaseModelOutputWithPast]: + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # retrieve input_ids and inputs_embeds + if stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + # Pipeline generation was not supported by this forward. Do not silently + # consume/mutate a cache while returning no cache to the caller. + if past_key_values is not None: + raise ValueError("past_key_values is not supported in the Qwen pipeline forward") + if output_attentions or output_hidden_states: + logger.warning_once("Recording attentions/hidden states is not supported in pipeline models.") + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models; disabling it.") + use_cache = False + + sp_size = shard_config.sequence_parallel_size + sp_group = shard_config.sequence_parallel_process_group + sp_mode = shard_config.sequence_parallelism_mode + if shard_config.enable_sequence_parallelism and sp_mode != "ring_attn" and not stage_manager.is_first_stage(): + seq_length *= sp_size + if position_ids is None: + positions = cache_position if cache_position is not None else torch.arange(seq_length, device=device) + position_ids = positions.unsqueeze(0) + # Later stages may only own a sequence shard. HF uses embeddings here + # solely for mask shape/dtype; positions and attention operate on the full sequence. + mask_inputs = hidden_states + if hidden_states.shape[1] != seq_length: + mask_inputs = hidden_states.new_empty((batch_size, seq_length, 1)) + causal_masks = prepare_hf_causal_masks( + self.config, + mask_inputs, + attention_mask, + None, + position_ids, + eager=shard_config.enable_flash_attention or shard_config.enable_sequence_parallelism, + ) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + if stage_manager.is_first_stage() and shard_config.enable_sequence_parallelism: + if is_share_sp_tp(sp_mode): + hidden_states = split_forward_gather_backward(hidden_states, dim=1, process_group=sp_group) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, dim=1, process_group=sp_group, grad_scale=1 / sp_size + ) + + start_idx, end_idx = stage_index + num_ckpt_layers = 0 + if self.gradient_checkpointing and self.training: + num_ckpt_layers = end_idx - start_idx + # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer + if shard_config.gradient_checkpoint_config is not None: + num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( + stage=stage_manager.stage, + num_stages=stage_manager.num_stages, + num_layers=end_idx - start_idx, + model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), + num_model_chunks=stage_manager.num_model_chunks, + ) + assert num_ckpt_layers <= end_idx - start_idx + + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + layer_kwargs = dict( + attention_mask=causal_masks[self.config.layer_types[idx]], + position_ids=position_ids, + past_key_values=None, + use_cache=False, + position_embeddings=position_embeddings, + ) + # PP selects checkpointed layers itself. Bypass HF's automatic + # checkpoint wrapper, but retain nn.Module hooks and keyword binding. + layer_call = partial(nn.Module.__call__, decoder_layer, **layer_kwargs) + if idx - start_idx < num_ckpt_layers: + hidden_states = decoder_layer._gradient_checkpointing_func(layer_call, hidden_states) + else: + hidden_states = layer_call(hidden_states) + + if stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + if not return_dict: + return (hidden_states,) + return BaseModelOutputWithPast(last_hidden_state=hidden_states) + # always return dict for imediate stage + return {"hidden_states": hidden_states} + + @staticmethod + def qwen2_for_causal_lm_forward( + self: Qwen2ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you consciours? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = Qwen2PipelineForwards.qwen2_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + force_sp_output_gather=False, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + # Match HF's FP32 cross entropy while preserving the output logits dtype. + loss = dist_cross_entropy( + labels, logits.float(), shard_config, self.lm_head.out_features, torch.float32 + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def qwen2_for_sequence_classification_forward( + self: Qwen2ForSequenceClassification, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + transformer_outputs = Qwen2PipelineForwards.qwen2_model_forward( + self.model, + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + ) + + if input_ids is not None: + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + batch_size = inputs_embeds.shape[0] + else: + batch_size = hidden_states.shape[0] + + if stage_manager.is_last_stage(): + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + else: + hidden_states = transformer_outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_qwen2_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + def forward( + self: Qwen2Attention, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if sp_mode is not None: + assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + + bsz, q_len, _ = hidden_states.size() + # sp: modify sp_len when sequence parallel mode is ring + if sp_mode in ["split_gather", "ring"]: + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + num_heads = query_states.shape[1] + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = None + dropout = self.attention_dropout if self.training else 0.0 + if shard_config.enable_flash_attention: + attn_output = ColoAttention.attention( + query_states, + key_states, + value_states, + dropout_p=dropout, + scale=self.scaling, + **get_hf_flash_attention_kwargs(attention_mask), + ) + else: + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling + if attention_mask is not None: + if attention_mask.dtype == torch.bool: + attn_weights = attn_weights.masked_fill(~attention_mask, torch.finfo(attn_weights.dtype).min) + else: + attn_weights = attn_weights + attention_mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + if attn_output.shape != (bsz, num_heads, q_len, self.head_dim): + raise ValueError(f"Unexpected attention output shape: {attn_output.shape}") + attn_output = attn_output.transpose(1, 2).contiguous() + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, num_heads * self.head_dim) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) + else: + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + return forward + + +def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + force_sp_output_gather: bool = True, + **kwargs, + ) -> BaseModelOutputWithPast: + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + if self.gradient_checkpointing and self.training: + # Do not let backward recomputation mutate an existing cache. + use_cache = False + past_key_values = None + if sp_mode in ["ring", "all_to_all"] and self.training: + use_cache = False + past_key_values = None + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + if position_ids is None: + positions = cache_position + if positions is None: + positions = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + positions = positions + get_hf_cache_seq_length(past_key_values) + position_ids = positions.unsqueeze(0) + + causal_masks = prepare_hf_causal_masks( + self.config, inputs_embeds, attention_mask, past_key_values, position_ids, eager=True + ) + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + if sp_mode in ["ring", "split_gather"]: + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + + for idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_masks[self.config.layer_types[idx]], + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + return forward + + +def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig): + def forward( + self: Qwen2ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + force_sp_output_gather=False, + **kwargs, + ) + + hidden_states = outputs[0] + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + logits = logits.float() + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return forward diff --git a/colossalai/shardformer/modeling/_qwen3_tf4.py b/colossalai/shardformer/modeling/_qwen3_tf4.py new file mode 100644 index 000000000000..9b6670ba051b --- /dev/null +++ b/colossalai/shardformer/modeling/_qwen3_tf4.py @@ -0,0 +1,834 @@ +# Modifed from qwen2 modeling +import math +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_causal_attention_mask, + _prepare_4d_causal_attention_mask_for_sdpa, +) +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3Attention, + Qwen3ForCausalLM, + Qwen3ForSequenceClassification, + Qwen3Model, + apply_rotary_pos_emb, + repeat_kv, +) +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import all_to_all_comm, split_forward_gather_backward +from colossalai.shardformer.shard import ShardConfig + +from ..layer import ColoAttention, dist_cross_entropy +from ..layer._operation import gather_sp_output +from ..layer.utils import is_share_sp_tp + + +class Qwen3PipelineForwards: + """ + This class serves as a micro library for forward function substitution of Qwen3 models + under pipeline setting. + """ + + @staticmethod + def qwen3_model_forward( + self: Qwen3Model, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + force_sp_output_gather: bool = True, + ) -> Union[Tuple, BaseModelOutputWithPast]: + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + seq_length_with_past = seq_length + past_key_values_length = 0 + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + + # Support SP + PP + sp_size = shard_config.sequence_parallel_size + sp_group = shard_config.sequence_parallel_process_group + sp_mode = shard_config.sequence_parallelism_mode + # For generating full positions ids (the states will be gathered along the seq dim before attention fwd). + if shard_config.enable_sequence_parallelism and sp_mode != "ring_attn" and not stage_manager.is_first_stage(): + seq_length *= sp_size + # Incoming pipeline states contain only the local SP slice. Both mask + # dimensions must describe the full sequence used by attention. + seq_length_with_past = seq_length + past_key_values_length + + if position_ids is None: + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + # embed positions, for the first stage, hidden_states is the input embeddings, + # for the other stages, hidden_states is the output of the previous stage + if shard_config.enable_flash_attention: + # in this case, attention_mask is a dict rather than a tensor + mask_shape = (batch_size, 1, seq_length, seq_length_with_past) + attention_mask = ColoAttention.prepare_attn_kwargs( + mask_shape, + hidden_states.dtype, + hidden_states.device, + q_padding_mask=attention_mask, + is_causal=True, + ) + else: + if self.config._attn_implementation == "flash_attention_2": + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self.config._attn_implementation == "sdpa" and not output_attentions: + # output_attentions=True can not be supported when using SDPA, and we fall back on + # the manual implementation that requires a 4D causal mask in all cases. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (batch_size, seq_length), + hidden_states, + past_key_values_length, + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + hidden_states, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) + + if stage_manager.is_first_stage(): + if shard_config.enable_sequence_parallelism: + if is_share_sp_tp(sp_mode): + hidden_states = split_forward_gather_backward( + hidden_states, + dim=1, + process_group=sp_group, + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, + dim=1, + process_group=sp_group, + grad_scale=1 / sp_size, + ) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + start_idx, end_idx = stage_index[0], stage_index[1] + num_ckpt_layers = 0 + if self.gradient_checkpointing and self.training: + num_ckpt_layers = end_idx - start_idx + # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer + if shard_config.gradient_checkpoint_config is not None: + num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( + stage=stage_manager.stage, + num_stages=stage_manager.num_stages, + num_layers=end_idx - start_idx, + model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), + num_model_chunks=stage_manager.num_model_chunks, + ) + assert num_ckpt_layers <= end_idx - start_idx + + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + past_key_values[idx] if past_key_values is not None else None + + if idx - start_idx < num_ckpt_layers: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if stage_manager.is_last_stage(): + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + # always return dict for imediate stage + return {"hidden_states": hidden_states} + + @staticmethod + def qwen3_for_causal_lm_forward( + self: Qwen3ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you consciours? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = Qwen3PipelineForwards.qwen3_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + force_sp_output_gather=False, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + hidden_states = outputs[0] + if hidden_states.shape[1] == 2: + pass + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def qwen3_for_sequence_classification_forward( + self: Qwen3ForSequenceClassification, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + transformer_outputs = Qwen3PipelineForwards.qwen3_model_forward( + self.model, + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + ) + + if input_ids is not None: + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + batch_size = inputs_embeds.shape[0] + else: + batch_size = hidden_states.shape[0] + + if stage_manager.is_last_stage(): + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + else: + hidden_states = transformer_outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_qwen3_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + def forward( + self: Qwen3Attention, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if sp_mode is not None: + assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + + bsz, q_len, _ = hidden_states.size() + # sp: modify sp_len when sequence parallel mode is ring + if sp_mode in ["split_gather", "ring"]: + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = self.q_norm(query_states.view(bsz, q_len, -1, self.head_dim)).transpose(1, 2) + key_states = self.k_norm(key_states.view(bsz, q_len, -1, self.head_dim)).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + kv_seq_len = key_states.shape[-2] + if past_key_value is not None: + if self.layer_idx is None: + raise ValueError( + f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " + "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " + "with a layer index." + ) + kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) + # Because the input can be padded, the absolute sequence length depends on the max position id. + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # Activate slicing cache only if the config has a value `sliding_windows` attribute + cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 + if ( + getattr(self.config, "sliding_window", None) is not None + and kv_seq_len > self.config.sliding_window + and cache_has_contents + ): + slicing_tokens = 1 - self.config.sliding_window + + past_key = past_key_value[self.layer_idx][0] + past_value = past_key_value[self.layer_idx][1] + + past_key = past_key[:, :, slicing_tokens:, :].contiguous() + past_value = past_value[:, :, slicing_tokens:, :].contiguous() + + if past_key.shape[-2] != self.config.sliding_window - 1: + raise ValueError( + f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" + f" {past_key.shape}" + ) + + if attention_mask is not None: + attention_mask = attention_mask[:, slicing_tokens:] + attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) + + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + if shard_config.enable_flash_attention: + assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict." + attn_output = ColoAttention.attention( + query_states, + key_states, + value_states, + dropout_p=0.0 if not self.training else self.attention_dropout, + **attention_mask, + ) + else: + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): + raise ValueError( + f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" + f" {attn_weights.size()}" + ) + + if attention_mask is not None: + if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): + raise ValueError( + f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" + ) + attn_weights = attn_weights + attention_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + attn_output = attn_output.transpose(1, 2).contiguous() + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) + else: + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output, None + + return forward + + +def get_qwen3_model_forward_for_flash_attn(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + logger = logging.get_logger(__name__) + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + force_sp_output_gather: bool = True, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + seq_length_with_past = seq_length + past_key_values_length = 0 + + if past_key_values is not None: + past_key_values_length = past_key_values[0][0].shape[2] + seq_length_with_past = seq_length_with_past + past_key_values_length + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device + ) + position_ids = position_ids.unsqueeze(0).view(-1, seq_length) + else: + position_ids = position_ids.view(-1, seq_length).long() + + # embed positions + hidden_states = inputs_embeds + + if shard_config.enable_flash_attention: + # in this case, attention_mask is a dict rather than a tensor + mask_shape = (batch_size, 1, seq_length, seq_length_with_past) + attention_mask = ColoAttention.prepare_attn_kwargs( + mask_shape, + hidden_states.dtype, + hidden_states.device, + q_padding_mask=attention_mask, + is_causal=True, + ) + else: + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length, + sliding_window=self.config.sliding_window, + ) + + if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training: + if use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + if sp_mode in ["ring", "split_gather"]: + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + return forward + + +def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig): + def forward( + self: Qwen3ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + force_sp_output_gather=False, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return forward diff --git a/colossalai/shardformer/modeling/_qwen3_tf5.py b/colossalai/shardformer/modeling/_qwen3_tf5.py new file mode 100644 index 000000000000..326d87a8a6e6 --- /dev/null +++ b/colossalai/shardformer/modeling/_qwen3_tf5.py @@ -0,0 +1,620 @@ +# Modifed from qwen2 modeling +from functools import partial +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3Attention, + Qwen3ForCausalLM, + Qwen3ForSequenceClassification, + Qwen3Model, + apply_rotary_pos_emb, + repeat_kv, +) +from transformers.utils import logging +from transformers.utils.generic import merge_with_config_defaults +from transformers.utils.output_capturing import capture_outputs + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer._operation import all_to_all_comm, split_forward_gather_backward +from colossalai.shardformer.shard import ShardConfig + +from .._utils import get_hf_cache_seq_length, get_hf_flash_attention_kwargs, prepare_hf_causal_masks +from ..layer import ColoAttention, dist_cross_entropy +from ..layer._operation import gather_sp_output +from ..layer.utils import is_share_sp_tp + + +class Qwen3PipelineForwards: + """ + This class serves as a micro library for forward function substitution of Qwen3 models + under pipeline setting. + """ + + @staticmethod + def qwen3_model_forward( + self: Qwen3Model, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + force_sp_output_gather: bool = True, + ) -> Union[Tuple, BaseModelOutputWithPast]: + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # retrieve input_ids and inputs_embeds + if stage_manager.is_first_stage(): + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + device = input_ids.device if input_ids is not None else inputs_embeds.device + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + hidden_states = inputs_embeds + else: + input_shape = hidden_states.shape[:-1] + batch_size, seq_length = input_shape + device = hidden_states.device + + # Pipeline generation was not supported by this forward. Do not silently + # consume/mutate a cache while returning no cache to the caller. + if past_key_values is not None: + raise ValueError("past_key_values is not supported in the Qwen pipeline forward") + if output_attentions or output_hidden_states: + logger.warning_once("Recording attentions/hidden states is not supported in pipeline models.") + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models; disabling it.") + use_cache = False + + sp_size = shard_config.sequence_parallel_size + sp_group = shard_config.sequence_parallel_process_group + sp_mode = shard_config.sequence_parallelism_mode + if shard_config.enable_sequence_parallelism and sp_mode != "ring_attn" and not stage_manager.is_first_stage(): + seq_length *= sp_size + if position_ids is None: + positions = cache_position if cache_position is not None else torch.arange(seq_length, device=device) + position_ids = positions.unsqueeze(0) + # Later stages may only own a sequence shard. HF uses embeddings here + # solely for mask shape/dtype; positions and attention operate on the full sequence. + mask_inputs = hidden_states + if hidden_states.shape[1] != seq_length: + mask_inputs = hidden_states.new_empty((batch_size, seq_length, 1)) + causal_masks = prepare_hf_causal_masks( + self.config, + mask_inputs, + attention_mask, + None, + position_ids, + eager=shard_config.enable_flash_attention or shard_config.enable_sequence_parallelism, + ) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + if stage_manager.is_first_stage() and shard_config.enable_sequence_parallelism: + if is_share_sp_tp(sp_mode): + hidden_states = split_forward_gather_backward(hidden_states, dim=1, process_group=sp_group) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, dim=1, process_group=sp_group, grad_scale=1 / sp_size + ) + + start_idx, end_idx = stage_index + num_ckpt_layers = 0 + if self.gradient_checkpointing and self.training: + num_ckpt_layers = end_idx - start_idx + # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer + if shard_config.gradient_checkpoint_config is not None: + num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( + stage=stage_manager.stage, + num_stages=stage_manager.num_stages, + num_layers=end_idx - start_idx, + model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), + num_model_chunks=stage_manager.num_model_chunks, + ) + assert num_ckpt_layers <= end_idx - start_idx + + for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): + layer_kwargs = dict( + attention_mask=causal_masks[self.config.layer_types[idx]], + position_ids=position_ids, + past_key_values=None, + use_cache=False, + position_embeddings=position_embeddings, + ) + # PP selects checkpointed layers itself. Bypass HF's automatic + # checkpoint wrapper, but retain nn.Module hooks and keyword binding. + layer_call = partial(nn.Module.__call__, decoder_layer, **layer_kwargs) + if idx - start_idx < num_ckpt_layers: + hidden_states = decoder_layer._gradient_checkpointing_func(layer_call, hidden_states) + else: + hidden_states = layer_call(hidden_states) + + if stage_manager.is_last_stage(): + hidden_states = self.norm(hidden_states) + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + if not return_dict: + return (hidden_states,) + return BaseModelOutputWithPast(last_hidden_state=hidden_states) + # always return dict for imediate stage + return {"hidden_states": hidden_states} + + @staticmethod + def qwen3_for_causal_lm_forward( + self: Qwen3ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + **kwargs, + ): + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you consciours? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you." + ```""" + logger = logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = Qwen3PipelineForwards.qwen3_model_forward( + self.model, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + force_sp_output_gather=False, + ) + past_key_values = None + + if stage_manager.is_last_stage(): + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + loss = None + if labels is not None: + # Match HF's FP32 cross entropy while preserving the output logits dtype. + loss = dist_cross_entropy( + labels, logits.float(), shard_config, self.lm_head.out_features, torch.float32 + ) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + else: + hidden_states = outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + @staticmethod + def qwen3_for_sequence_classification_forward( + self: Qwen3ForSequenceClassification, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + shard_config: ShardConfig = None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + logger = logging.get_logger(__name__) + + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + transformer_outputs = Qwen3PipelineForwards.qwen3_model_forward( + self.model, + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + shard_config=shard_config, + ) + + if input_ids is not None: + batch_size = input_ids.shape[0] + elif inputs_embeds is not None: + batch_size = inputs_embeds.shape[0] + else: + batch_size = hidden_states.shape[0] + + if stage_manager.is_last_stage(): + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + else: + hidden_states = transformer_outputs.get("hidden_states") + return {"hidden_states": hidden_states} + + +def get_qwen3_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + def forward( + self: Qwen3Attention, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor] = None, + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + if sp_mode is not None: + assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" + assert (sp_size is not None) and ( + sp_group is not None + ), "Must specify sp_size and sp_group for sequence parallel" + + bsz, q_len, _ = hidden_states.size() + # sp: modify sp_len when sequence parallel mode is ring + if sp_mode in ["split_gather", "ring"]: + q_len *= sp_size + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + # sp: all-to-all comminucation when introducing sequence parallel + if sp_mode == "all_to_all": + query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) + key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) + value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) + bsz, q_len, _ = query_states.size() + + query_states = self.q_norm(query_states.view(bsz, q_len, -1, self.head_dim)).transpose(1, 2) + key_states = self.k_norm(key_states.view(bsz, q_len, -1, self.head_dim)).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + if past_key_values is not None: + key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx) + num_heads = query_states.shape[1] + + # repeat k/v heads if n_kv_heads < n_heads + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = None + dropout = self.attention_dropout if self.training else 0.0 + if shard_config.enable_flash_attention: + attn_output = ColoAttention.attention( + query_states, + key_states, + value_states, + dropout_p=dropout, + scale=self.scaling, + **get_hf_flash_attention_kwargs(attention_mask), + ) + else: + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scaling + if attention_mask is not None: + if attention_mask.dtype == torch.bool: + attn_weights = attn_weights.masked_fill(~attention_mask, torch.finfo(attn_weights.dtype).min) + else: + attn_weights = attn_weights + attention_mask + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + if attn_output.shape != (bsz, num_heads, q_len, self.head_dim): + raise ValueError(f"Unexpected attention output shape: {attn_output.shape}") + attn_output = attn_output.transpose(1, 2).contiguous() + if sp_mode == "all_to_all": + attn_output = attn_output.reshape(bsz, q_len, num_heads * self.head_dim) + attn_output = all_to_all_comm( + attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication + ) + else: + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + return attn_output, attn_weights + + return forward + + +def get_qwen3_model_forward_for_flash_attn(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + force_sp_output_gather: bool = True, + **kwargs, + ) -> BaseModelOutputWithPast: + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + if self.gradient_checkpointing and self.training: + # Do not let backward recomputation mutate an existing cache. + use_cache = False + past_key_values = None + if sp_mode in ["ring", "all_to_all"] and self.training: + use_cache = False + past_key_values = None + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + if position_ids is None: + positions = cache_position + if positions is None: + positions = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + positions = positions + get_hf_cache_seq_length(past_key_values) + position_ids = positions.unsqueeze(0) + + causal_masks = prepare_hf_causal_masks( + self.config, inputs_embeds, attention_mask, past_key_values, position_ids, eager=True + ) + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + if sp_mode in ["ring", "split_gather"]: + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication + ) + elif sp_mode == "all_to_all": + hidden_states = split_forward_gather_backward( + hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication + ) + + for idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_masks[self.config.layer_types[idx]], + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = self.norm(hidden_states) + if shard_config.enable_sequence_parallelism: + if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): + hidden_states = gather_sp_output(hidden_states, shard_config) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + return forward + + +def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig): + def forward( + self: Qwen3ForCausalLM, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen2ForCausalLM + + >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) + >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + force_sp_output_gather=False, + **kwargs, + ) + + hidden_states = outputs[0] + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + logits = logits.float() + loss = None + if labels is not None: + loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return forward diff --git a/colossalai/shardformer/modeling/_vit_tf4.py b/colossalai/shardformer/modeling/_vit_tf4.py new file mode 100644 index 000000000000..5106d97cf4bc --- /dev/null +++ b/colossalai/shardformer/modeling/_vit_tf4.py @@ -0,0 +1,386 @@ +from typing import List, Optional, Tuple, Union + +import torch +from transformers.models.vit.modeling_vit import BaseModelOutput, ViTEncoder +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer import ColoAttention + + +def _encoder_forward( + encoder: ViTEncoder, + start_idx: int, + end_idx: int, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + stage_manager: PipelineStageManager = None, +) -> Union[tuple, BaseModelOutput]: + for i in range(start_idx, end_idx): + layer_module = encoder.layer[i] + + layer_head_mask = head_mask[i] if head_mask is not None else None + + if encoder.gradient_checkpointing and encoder.training: + layer_outputs = encoder._gradient_checkpointing_func( + layer_module.__call__, + hidden_states, + layer_head_mask, + output_attentions, + ) + else: + layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions) + + hidden_states = layer_outputs[0] + if not stage_manager.is_last_stage(): + return hidden_states + else: + if not return_dict: + return tuple(hidden_states) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=None, + attentions=None, + ) + + +def ViTModel_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): + from transformers.models.vit.modeling_vit import BaseModelOutputWithPooling + + def pp_forward( + self, + pixel_values: Optional[torch.Tensor] = None, + bool_masked_pos: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + r""" + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*): + Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + logger = logging.get_logger(__name__) + + # Preprocess passed in arguments + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + if stage_manager.is_first_stage(): + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + # TODO(FoolPlayer): maybe have a cleaner way to cast the input (from `ImageProcessor` side?) + expected_dtype = self.embeddings.patch_embeddings.projection.weight.dtype + if pixel_values.dtype != expected_dtype: + pixel_values = pixel_values.to(expected_dtype) + + embedding_output = self.embeddings( + pixel_values, + bool_masked_pos=bool_masked_pos, + interpolate_pos_encoding=interpolate_pos_encoding, + ) + hidden_states = embedding_output + else: + assert ( + hidden_states is not None + ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" + + encoder_outputs = _encoder_forward( + encoder=self.encoder, + start_idx=stage_index[0], + end_idx=stage_index[1], + hidden_states=hidden_states, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + ) + if not stage_manager.is_last_stage(): + return {"hidden_states": encoder_outputs} + + sequence_output = encoder_outputs[0] + sequence_output = self.layernorm(sequence_output) + pooled_output = self.pooler(sequence_output) if self.pooler is not None else None + + if not return_dict: + head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,) + return head_outputs + encoder_outputs[1:] + + return BaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) + + return pp_forward + + +def ViTForImageClassification_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): + from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + from transformers.models.vit.modeling_vit import ImageClassifierOutput + + def pp_forward( + self, + pixel_values: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + ) -> Union[tuple, ImageClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if not stage_manager.is_first_stage(): + assert ( + hidden_states is not None + ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" + + outputs = self.vit( + pixel_values, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=return_dict, + hidden_states=hidden_states, + ) + + # not last stage, return hidden_states + if not stage_manager.is_last_stage(): + return outputs + else: + sequence_output = outputs[0] + + # last stage + logits = self.classifier(sequence_output[:, 0, :]) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return pp_forward + + +def ViTForMaskedImageModeling_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): + import math + + import torch.nn as nn + from transformers.models.vit.modeling_vit import ImageClassifierOutput, MaskedImageModelingOutput + + def pp_forward( + self, + pixel_values: Optional[torch.Tensor] = None, + bool_masked_pos: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + ) -> Union[tuple, ImageClassifierOutput]: + r""" + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`): + Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). + + Returns: + + Examples: + ```python + >>> from transformers import AutoImageProcessor, ViTForMaskedImageModeling + >>> import torch + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k") + >>> model = ViTForMaskedImageModeling.from_pretrained("google/vit-base-patch16-224-in21k") + + >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2 + >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values + >>> # create random boolean mask of shape (batch_size, num_patches) + >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool() + + >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos) + >>> loss, reconstructed_pixel_values = outputs.loss, outputs.reconstruction + >>> list(reconstructed_pixel_values.shape) + [1, 3, 224, 224] + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if bool_masked_pos is not None and (self.config.patch_size != self.config.encoder_stride): + raise ValueError( + "When `bool_masked_pos` is provided, `patch_size` must be equal to `encoder_stride` to ensure that " + "the reconstructed image has the same dimensions as the input." + f"Got `patch_size` = {self.config.patch_size} and `encoder_stride` = {self.config.encoder_stride}." + ) + + if not stage_manager.is_first_stage(): + assert ( + hidden_states is not None + ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" + + outputs = self.vit( + pixel_values, + bool_masked_pos=bool_masked_pos, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=return_dict, + hidden_states=hidden_states, + ) + if not stage_manager.is_last_stage(): + return outputs + else: + sequence_output = outputs[0] + + # Reshape to (batch_size, num_channels, height, width) + sequence_output = sequence_output[:, 1:] + batch_size, sequence_length, num_channels = sequence_output.shape + height = width = math.floor(sequence_length**0.5) + sequence_output = sequence_output.permute(0, 2, 1).reshape(batch_size, num_channels, height, width) + + # Reconstruct pixel values + reconstructed_pixel_values = self.decoder(sequence_output) + + masked_im_loss = None + if bool_masked_pos is not None: + size = self.config.image_size // self.config.patch_size + bool_masked_pos = bool_masked_pos.reshape(-1, size, size) + mask = ( + bool_masked_pos.repeat_interleave(self.config.patch_size, 1) + .repeat_interleave(self.config.patch_size, 2) + .unsqueeze(1) + .contiguous() + ) + reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none") + masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels + + if not return_dict: + output = (reconstructed_pixel_values,) + outputs[1:] + return ((masked_im_loss,) + output) if masked_im_loss is not None else output + + return MaskedImageModelingOutput( + loss=masked_im_loss, + reconstruction=reconstructed_pixel_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return pp_forward + + +def get_vit_flash_self_attention_forward(): + from transformers.models.vit.modeling_vit import ViTSelfAttention + + def forward( + self: ViTSelfAttention, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + assert head_mask is None, "head_mask is not supported for FlashAttention" + mixed_query_layer = self.query(hidden_states) + + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + query_layer = self.transpose_for_scores(mixed_query_layer) + + dropout_p = self.dropout_prob if self.training else 0.0 + context_layer = ColoAttention.attention(query_layer, key_layer, value_layer, dropout_p=dropout_p) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + + outputs = (context_layer, None) if output_attentions else (context_layer,) + + return outputs + + return forward + + +def get_jit_fused_vit_output_forward(): + from transformers.models.vit.modeling_vit import ViTOutput + + def forward(self: ViTOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) + return hidden_states + + return forward + + +def get_jit_fused_vit_intermediate_forward(): + from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, bias = self.dense(hidden_states) + hidden_states = JitGeLUFunction.apply(hidden_states, bias) + + return hidden_states + + return forward diff --git a/colossalai/shardformer/modeling/_vit_tf5.py b/colossalai/shardformer/modeling/_vit_tf5.py new file mode 100644 index 000000000000..8edc0598bd81 --- /dev/null +++ b/colossalai/shardformer/modeling/_vit_tf5.py @@ -0,0 +1,299 @@ +"""ViT forwards for the Transformers 5.17 flat layer/attention layout.""" + +from typing import List, Optional, Union + +import torch +from transformers.modeling_outputs import BaseModelOutputWithPooling +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer import ColoAttention + +from .._utils import get_hf_flash_attention_kwargs, prepare_hf_bidirectional_mask + + +def ViTModel_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): + def pp_forward( + self, + pixel_values=None, + bool_masked_pos=None, + head_mask=None, + output_attentions=None, + output_hidden_states=None, + interpolate_pos_encoding=None, + return_dict=None, + hidden_states=None, + attention_mask=None, + **kwargs, + ): + if head_mask is not None: + raise ValueError("HF 5.17 ViT does not support head_mask") + if output_attentions or output_hidden_states: + logging.get_logger(__name__).warning_once("Pipeline ViT does not collect attentions/hidden states.") + return_dict = self.config.use_return_dict if return_dict is None else return_dict + if stage_manager.is_first_stage(): + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + pixel_values = pixel_values.to(self.embeddings.patch_embeddings.projection.weight.dtype) + hidden_states = self.embeddings( + pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding + ) + elif hidden_states is None: + raise ValueError("hidden_states must be provided to a non-first ViT stage") + + # Explicit masks work for both HF attention and the ColoAttention replacement. + attention_mask = prepare_hf_bidirectional_mask(self.config, hidden_states, attention_mask) + indices = stage_index[stage_manager.model_chunk_id] if stage_manager.is_interleave else stage_index + for i in range(*indices): + # HF GradientCheckpointingLayer owns checkpointing; do not checkpoint twice. + hidden_states = self.layers[i](hidden_states, attention_mask=attention_mask, **kwargs) + if not stage_manager.is_last_stage(): + return {"hidden_states": hidden_states} + hidden_states = self.layernorm(hidden_states) + pooled_output = self.pooler(hidden_states) if self.pooler is not None else None + outputs = BaseModelOutputWithPooling(last_hidden_state=hidden_states, pooler_output=pooled_output) + return outputs if return_dict else outputs.to_tuple() + + return pp_forward + + +def ViTForImageClassification_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): + from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + from transformers.modeling_outputs import ImageClassifierOutput + + def pp_forward( + self, + pixel_values: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> Union[tuple, ImageClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if not stage_manager.is_first_stage(): + assert ( + hidden_states is not None + ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" + + outputs = self.vit( + pixel_values, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=return_dict, + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + + # not last stage, return hidden_states + if not stage_manager.is_last_stage(): + return outputs + else: + sequence_output = outputs[0] + + # last stage + logits = self.classifier(sequence_output[:, 0, :]) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return pp_forward + + +def ViTForMaskedImageModeling_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): + import math + + import torch.nn as nn + from transformers.modeling_outputs import ImageClassifierOutput, MaskedImageModelingOutput + + def pp_forward( + self, + pixel_values: Optional[torch.Tensor] = None, + bool_masked_pos: Optional[torch.BoolTensor] = None, + head_mask: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + interpolate_pos_encoding: Optional[bool] = None, + return_dict: Optional[bool] = None, + hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> Union[tuple, ImageClassifierOutput]: + r""" + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`): + Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). + + Returns: + + Examples: + ```python + >>> from transformers import AutoImageProcessor, ViTForMaskedImageModeling + >>> import torch + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k") + >>> model = ViTForMaskedImageModeling.from_pretrained("google/vit-base-patch16-224-in21k") + + >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2 + >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values + >>> # create random boolean mask of shape (batch_size, num_patches) + >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool() + + >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos) + >>> loss, reconstructed_pixel_values = outputs.loss, outputs.reconstruction + >>> list(reconstructed_pixel_values.shape) + [1, 3, 224, 224] + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if bool_masked_pos is not None and (self.config.patch_size != self.config.encoder_stride): + raise ValueError( + "When `bool_masked_pos` is provided, `patch_size` must be equal to `encoder_stride` to ensure that " + "the reconstructed image has the same dimensions as the input." + f"Got `patch_size` = {self.config.patch_size} and `encoder_stride` = {self.config.encoder_stride}." + ) + + if not stage_manager.is_first_stage(): + assert ( + hidden_states is not None + ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" + + outputs = self.vit( + pixel_values, + bool_masked_pos=bool_masked_pos, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + return_dict=return_dict, + hidden_states=hidden_states, + attention_mask=attention_mask, + **kwargs, + ) + if not stage_manager.is_last_stage(): + return outputs + else: + sequence_output = outputs[0] + + # Reshape to (batch_size, num_channels, height, width) + sequence_output = sequence_output[:, 1:] + batch_size, sequence_length, num_channels = sequence_output.shape + height = width = math.floor(sequence_length**0.5) + sequence_output = sequence_output.permute(0, 2, 1).reshape(batch_size, num_channels, height, width) + + # Reconstruct pixel values + reconstructed_pixel_values = self.decoder(sequence_output) + + masked_im_loss = None + if bool_masked_pos is not None: + size = self.config.image_size // self.config.patch_size + bool_masked_pos = bool_masked_pos.reshape(-1, size, size) + mask = ( + bool_masked_pos.repeat_interleave(self.config.patch_size, 1) + .repeat_interleave(self.config.patch_size, 2) + .unsqueeze(1) + .contiguous() + ) + reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none") + masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels + + if not return_dict: + output = (reconstructed_pixel_values,) + outputs[1:] + return ((masked_im_loss,) + output) if masked_im_loss is not None else output + + return MaskedImageModelingOutput( + loss=masked_im_loss, + reconstruction=reconstructed_pixel_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + return pp_forward + + +def get_vit_flash_self_attention_forward(): + def forward(self, hidden_states, attention_mask=None, **kwargs): + input_shape = hidden_states.shape[:-1] + shape = (*input_shape, -1, self.head_dim) + query = self.q_proj(hidden_states).view(shape).transpose(1, 2) + key = self.k_proj(hidden_states).view(shape).transpose(1, 2) + value = self.v_proj(hidden_states).view(shape).transpose(1, 2) + if attention_mask is not None and attention_mask.dtype == torch.bool: + attention_mask = torch.zeros_like(attention_mask, dtype=query.dtype).masked_fill( + ~attention_mask, torch.finfo(query.dtype).min + ) + output = ColoAttention.attention( + query, + key, + value, + **get_hf_flash_attention_kwargs(attention_mask), + dropout_p=self.attention_dropout if self.training else 0.0, + scale=self.scaling, + ) + output = output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + return self.o_proj(output), None + + return forward + + +def get_jit_fused_vit_layer_forward(): + def forward(self, hidden_states, attention_mask=None, **kwargs): + residual = hidden_states + hidden_states, _ = self.attention(self.layernorm_before(hidden_states), attention_mask, **kwargs) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout.p, self.training) + residual = hidden_states + hidden_states = self.mlp(self.layernorm_after(hidden_states)) + return self.dropout_add(hidden_states, residual, self.dropout.p, self.training) + + return forward diff --git a/colossalai/shardformer/modeling/_whisper_tf4.py b/colossalai/shardformer/modeling/_whisper_tf4.py new file mode 100644 index 000000000000..619bbc98e3a0 --- /dev/null +++ b/colossalai/shardformer/modeling/_whisper_tf4.py @@ -0,0 +1,1193 @@ +import logging +import random +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_causal_attention_mask, + _prepare_4d_causal_attention_mask_for_sdpa, +) +from transformers.modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, + SequenceClassifierOutput, +) +from transformers.models.whisper.modeling_whisper import ( + _HIDDEN_STATES_START_POSITION, + WhisperDecoder, + WhisperEncoder, + WhisperForAudioClassification, + WhisperForConditionalGeneration, + WhisperModel, + shift_tokens_right, +) +from transformers.utils import logging + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer import ColoAttention +from colossalai.shardformer.shard import ShardConfig + +logger = logging.get_logger(__name__) + + +def _get_attention_mask( + self: WhisperDecoder, + shard_config: ShardConfig, + hidden_states: torch.Tensor, + past_key_values_length: int, + attention_mask: Optional[torch.FloatTensor], + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, +): + batch_size, seq_length = hidden_states.shape[:2] + mask_seq_length = past_key_values_length + seq_length + if shard_config.enable_flash_attention: + attention_mask = ColoAttention.prepare_attn_kwargs( + (batch_size, 1, seq_length, mask_seq_length), + hidden_states.dtype, + hidden_states.device, + attention_mask, + is_causal=True, + ) + else: + input_shape = (batch_size, seq_length) + if self._use_flash_attention_2: + # 2d mask is passed through the layers + attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None + elif self._use_sdpa and head_mask is None and not output_attentions: + # output_attentions=True & head_mask can not be supported when using SDPA. + attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, input_shape, hidden_states, past_key_values_length + ) + else: + # 4d mask is passed through the layers + attention_mask = _prepare_4d_causal_attention_mask( + attention_mask, input_shape, hidden_states, past_key_values_length + ) + return attention_mask + + +def get_whisper_flash_attention_forward(): + from transformers.models.whisper.modeling_whisper import WhisperAttention + + def forward( + self: WhisperAttention, + hidden_states: torch.Tensor, + key_value_states: Optional[torch.Tensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + attention_mask: Optional[dict] = None, + layer_head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + cache_position: Optional[torch.LongTensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + """Input shape: Batch x Time x Channel""" + assert layer_head_mask is None, "layer_head_mask is not supported for FlashAttention" + # for encoder, attention_mask is None + if attention_mask is None: + attention_mask = {} + # if key_value_states are provided this layer is used as a cross-attention layer + # for the decoder + is_cross_attention = key_value_states is not None + + bsz, tgt_len, _ = hidden_states.size() + + # get query proj + query_states = self.q_proj(hidden_states) + # get key, value proj + # `past_key_value[0].shape[2] == key_value_states.shape[1]` + # is checking that the `sequence_length` of the `past_key_value` is the same as + # the provided `key_value_states` to support prefix tuning + if ( + is_cross_attention + and past_key_value is not None + and past_key_value[0].shape[2] == key_value_states.shape[1] + ): + # reuse k,v, cross_attentions + key_states = past_key_value[0] + value_states = past_key_value[1] + elif is_cross_attention: + # cross_attentions + key_states = self._shape(self.k_proj(key_value_states), -1, bsz) + value_states = self._shape(self.v_proj(key_value_states), -1, bsz) + elif past_key_value is not None: + # reuse k, v, self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + key_states = torch.cat([past_key_value[0], key_states], dim=2) + value_states = torch.cat([past_key_value[1], value_states], dim=2) + else: + # self_attention + key_states = self._shape(self.k_proj(hidden_states), -1, bsz) + value_states = self._shape(self.v_proj(hidden_states), -1, bsz) + + if self.is_decoder: + # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. + # Further calls to cross_attention layer can then reuse all cross-attention + # key/value_states (first "if" case) + # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of + # all previous decoder key/value_states. Further calls to uni-directional self-attention + # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) + # if encoder bi-directional self-attention `past_key_value` is always `None` + past_key_value = (key_states, value_states) + + query_states = self._shape(query_states, tgt_len, bsz) + + dropout_p = self.dropout if self.training else 0.0 + attn_output = ColoAttention.attention( + query_states, + key_states, + value_states, + **attention_mask, + dropout_p=dropout_p, + scale=self.scaling, + ) + attn_output = attn_output.transpose(1, 2) + + # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be + # partitioned across GPUs when using tensor-parallelism. + attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) + + attn_output = self.out_proj(attn_output) + + return attn_output, None, past_key_value + + return forward + + +def get_whisper_decoder_forward_for_flash_attention(shard_config: ShardConfig): + def forward( + self: WhisperDecoder, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + head_mask=None, + cross_attn_head_mask=None, + past_key_values=None, + inputs_embeds=None, + position_ids=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + cache_position=None, + ): + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + # past_key_values_length + past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + attention_mask = _get_attention_mask(self, shard_config, inputs_embeds, past_key_values_length, attention_mask) + + # embed positions + if input_ids is not None: + positions = self.embed_positions( + input_ids, past_key_values_length=past_key_values_length, position_ids=position_ids + ) + else: + positions = self.embed_positions( + inputs_embeds, past_key_values_length=past_key_values_length, position_ids=position_ids + ) + + hidden_states = inputs_embeds + positions + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`..." + ) + use_cache = False + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None + next_decoder_cache = () if use_cache else None + + # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired + for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]): + if attn_mask is not None: + assert attn_mask.size()[0] == (len(self.layers)), ( + f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for" + f" {head_mask.size()[0]}." + ) + for idx, decoder_layer in enumerate(self.layers): + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + if output_hidden_states: + all_hidden_states += (hidden_states,) + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + + past_key_value = past_key_values[idx] if past_key_values is not None else None + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + # None for past_key_value + return module(*inputs, output_attentions, use_cache) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(decoder_layer), + hidden_states, + attention_mask, + encoder_hidden_states, + None, # encoder attention mask + head_mask[idx] if head_mask is not None else None, + (cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), + None, # past_key_value + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + cross_attn_layer_head_mask=( + cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None + ), + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache += (layer_outputs[3 if output_attentions else 1],) + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if encoder_hidden_states is not None: + all_cross_attentions += (layer_outputs[2],) + + hidden_states = self.layer_norm(hidden_states) + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + next_cache, + all_hidden_states, + all_self_attns, + all_cross_attentions, + ] + if v is not None + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + cross_attentions=all_cross_attentions, + ) + + return forward + + +def get_jit_fused_whisper_encoder_layer_forward(): + from transformers.models.whisper.modeling_whisper import WhisperEncoderLayer + + def forward( + self: WhisperEncoderLayer, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + layer_head_mask: torch.Tensor, + output_attentions: bool = False, + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size + `(encoder_attention_heads,)`. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + hidden_states, attn_weights, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + output_attentions=output_attentions, + ) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + + if hidden_states.dtype == torch.float16 and ( + torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any() + ): + clamp_value = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (attn_weights,) + + return outputs + + return forward + + +def get_jit_fused_whisper_decoder_layer_forward(): + from transformers.models.whisper.modeling_whisper import WhisperDecoderLayer + + def forward( + self: WhisperDecoderLayer, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + encoder_hidden_states: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + layer_head_mask: Optional[torch.Tensor] = None, + cross_attn_layer_head_mask: Optional[torch.Tensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = True, + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + encoder_hidden_states (`torch.FloatTensor`): + cross attention input to the layer of shape `(batch, seq_len, embed_dim)` + encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size + `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. + layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size + `(encoder_attention_heads,)`. + cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of + size `(decoder_attention_heads,)`. + past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + hidden_states = self.self_attn_layer_norm(hidden_states) + + # Self Attention + # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 + self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None + # add present self-attn cache to positions 1,2 of present_key_value tuple + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + past_key_value=self_attn_past_key_value, + attention_mask=attention_mask, + layer_head_mask=layer_head_mask, + output_attentions=output_attentions, + ) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + + # Cross-Attention Block + cross_attn_present_key_value = None + cross_attn_weights = None + if encoder_hidden_states is not None: + residual = hidden_states + hidden_states = self.encoder_attn_layer_norm(hidden_states) + + # cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple + cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None + hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + layer_head_mask=cross_attn_layer_head_mask, + past_key_value=cross_attn_past_key_value, + output_attentions=output_attentions, + ) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + + # add cross-attn to positions 3,4 of present_key_value tuple + present_key_value = present_key_value + cross_attn_present_key_value + + # Fully Connected + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states = self.activation_fn(self.fc1(hidden_states)) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.fc2(hidden_states) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights, cross_attn_weights) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + return forward + + +class WhisperPipelineForwards: + """ + This class serves as a micro library for forward function substitution of Llama models + under pipeline setting. + """ + + @staticmethod + def whisper_encoder_forward( + self: WhisperEncoder, + input_features, + attention_mask=None, + head_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + encoder_states=None, + all_attentions=None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ): + r""" + Args: + input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`): + Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be + obtained by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a + `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into + `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding + and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] + attention_mask (`torch.Tensor`)`, *optional*): + Whisper does not support masking of the `input_features`, this argument is preserved for compatibility, + but it is not used. By default the silence in the input log mel spectrogram are ignored. + head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + logging.get_logger(__name__) + + stage = stage_manager.stage + at_first_stage = stage == 0 + at_last_stage = stage == decoder_starting_stage - 1 + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Process inputs if at the first stage of encoder. + if at_first_stage: + inputs_embeds = nn.functional.gelu(self.conv1(input_features)) + inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) + + inputs_embeds = inputs_embeds.permute(0, 2, 1) + embed_pos = self.embed_positions.weight + + hidden_states = inputs_embeds + embed_pos + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + + # check if head_mask has a correct number of layers specified if desired + if head_mask is not None: + assert head_mask.size()[0] == ( + len(self.layers) + ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." + + else: + if hidden_states is None: + raise ValueError( + "hidden_states shouldn't be None for stages other than the first stage of encoder/decoder." + ) + + start_idx, end_idx = stage_index[0], stage_index[1] + + for idx in range(start_idx, end_idx): + encoder_layer = self.layers[idx] + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + dropout_probability = random.uniform(0, 1) + if self.training and (dropout_probability < self.layerdrop): # skip the layer + layer_outputs = (None, None) + else: + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + encoder_layer.__call__, + hidden_states, + None, + (head_mask[idx] if head_mask is not None else None), + output_attentions, + ) + else: + layer_outputs = encoder_layer( + hidden_states, + None, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + output_attentions=output_attentions, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_attentions = all_attentions + (layer_outputs[1],) + + if at_last_stage: + hidden_states = self.layer_norm(hidden_states) + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + if not return_dict: + return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, + ) + + else: + return {"hidden_states": hidden_states, "head_mask": head_mask} + + @staticmethod + def whisper_decoder_forward( + self: WhisperDecoder, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + head_mask=None, + cross_attn_head_mask=None, + past_key_values=None, + inputs_embeds=None, + position_ids=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ): + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you + provide it. + + Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention + of the decoder. + head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): + Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention + on hidden heads. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): + Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of + shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. + + Contains pre-computed hidden-states (key and values in the self-attention blocks and in the + cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. + + If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those + that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of + all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of + shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing + `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more + control over how to convert `input_ids` indices into associated vectors than the model's internal + embedding lookup matrix. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + logger = logging.get_logger(__name__) + stage = stage_manager.stage + at_first_stage = stage == decoder_starting_stage + at_last_stage = stage == stage_manager.num_stages - 1 + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None + next_decoder_cache = () if use_cache else None + + # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired + for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]): + if attn_mask is not None: + assert attn_mask.size()[0] == (len(self.layers)), ( + f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for" + f" {head_mask.size()[0]}." + ) + + # past_key_values_length + past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 + + if at_first_stage: + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + attention_mask = _get_attention_mask( + self, shard_config, inputs_embeds, past_key_values_length, attention_mask + ) + + # embed positions + if input_ids is not None: + positions = self.embed_positions( + input_ids, past_key_values_length=past_key_values_length, position_ids=position_ids + ) + else: + positions = self.embed_positions( + inputs_embeds, past_key_values_length=past_key_values_length, position_ids=position_ids + ) + + hidden_states = inputs_embeds + positions + hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) + + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once( + "`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`..." + ) + use_cache = False + + else: + if hidden_states is None: + raise ValueError( + "hidden_states shouldn't be None for stages other than the first stage of encoder/decoder." + ) + input_shape = hidden_states.size()[:-1] + attention_mask = _get_attention_mask( + self, + shard_config, + hidden_states, + past_key_values_length, + attention_mask, + ) + + start_idx, end_idx = stage_index[0], stage_index[1] + + for idx in range(start_idx, end_idx): + # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) + decoder_layer = self.layers[idx] + + if output_hidden_states: + all_hidden_states += (hidden_states,) + dropout_probability = random.uniform(0, 1) + if self.training and (dropout_probability < self.layerdrop): + continue + + past_key_value = past_key_values[idx] if past_key_values is not None else None + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + attention_mask, + encoder_hidden_states, + None, # encoder attention mask + head_mask[idx] if head_mask is not None else None, + (cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), + None, # past_key_value + output_attentions, + use_cache, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + layer_head_mask=(head_mask[idx] if head_mask is not None else None), + cross_attn_layer_head_mask=( + cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None + ), + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + ) + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache += (layer_outputs[3 if output_attentions else 1],) + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + if encoder_hidden_states is not None: + all_cross_attentions += (layer_outputs[2],) + + if at_last_stage: + hidden_states = self.layer_norm(hidden_states) + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + next_cache = next_decoder_cache if use_cache else None + if not return_dict: + return tuple( + v + for v in [ + hidden_states, + next_cache, + all_hidden_states, + all_self_attns, + all_cross_attentions, + ] + if v is not None + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + cross_attentions=all_cross_attentions, + ) + + else: + return { + "head_mask": head_mask, + "cross_attn_head_mask": cross_attn_head_mask, + "hidden_states": hidden_states, + } + + @staticmethod + def whisper_model_forward( + self: WhisperModel, + input_features: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.LongTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + decoder_head_mask: Optional[torch.Tensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, + decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ): + r""" + Returns: + + Example: + ```python + >>> import torch + >>> from transformers import AutoFeatureExtractor, WhisperModel + >>> from datasets import load_dataset + + >>> model = WhisperModel.from_pretrained("openai/whisper-base") + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id + >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state + >>> list(last_hidden_state.shape) + [1, 2, 512] + ```""" + # TODO: left the recording kv-value tensors as () or None type, this feature may be added in the future. + if past_key_values: + logger.warning_once("Non-empty past_key_values is not supported for pipeline models at the moment.") + past_key_values = None + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + logging.get_logger(__name__) + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + in_decoder = stage_manager.stage >= decoder_starting_stage + if not in_decoder: + if encoder_outputs is None: + input_features = self._mask_input_features(input_features, attention_mask=attention_mask) + + encoder_outputs = WhisperPipelineForwards.whisper_encoder_forward( + self.encoder, + input_features, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + ) + + if stage_manager.stage == decoder_starting_stage - 1: + # last stage of encoder + return {"encoder_hidden_states": encoder_outputs[0]} + else: + return encoder_outputs + # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True + elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + encoder_outputs = BaseModelOutput( + last_hidden_state=encoder_outputs[0], + hidden_states=(encoder_outputs[1] if len(encoder_outputs) > 1 else None), + attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + ) + + at_last_decoder_stage = stage_manager.is_last_stage() + at_first_decoder_stage = stage_manager.stage == decoder_starting_stage + if encoder_outputs is not None: + encoder_hidden_states = encoder_outputs[0] + elif encoder_hidden_states is None: + raise ValueError("Non-empty encoder_hidden_states should be passed in at decoder stages.") + + if not at_first_decoder_stage and hidden_states is None: + raise ValueError("If not at the first layer of decoder, non-empty hidden_states must be provided.") + + # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn) + decoder_outputs = WhisperPipelineForwards.whisper_decoder_forward( + self.decoder, + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + position_ids=decoder_position_ids, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + shard_config=shard_config, + ) + + # Directly return outputs of overloaded Whisper forward if not at last stage. + if not at_last_decoder_stage: + # encoder_hidden_states should be passed to the next stage + decoder_outputs["encoder_hidden_states"] = encoder_hidden_states + return decoder_outputs + + if not return_dict: + return decoder_outputs + encoder_outputs + + return Seq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_hidden_states, + ) + + @staticmethod + def whisper_for_conditional_generation_forward( + self: WhisperForConditionalGeneration, + input_features: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.LongTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + decoder_head_mask: Optional[torch.Tensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, + decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` + or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is + only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, WhisperForConditionalGeneration + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") + >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + + >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + + >>> generated_ids = model.generate(inputs=input_features) + + >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + >>> transcription + ' Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.' + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if labels is not None: + if decoder_input_ids is None and decoder_inputs_embeds is None: + decoder_input_ids = shift_tokens_right( + labels, self.config.pad_token_id, self.config.decoder_start_token_id + ) + in_decoder = stage_manager.stage >= decoder_starting_stage + at_last_decoder_stage = stage_manager.is_last_stage() + outputs = WhisperPipelineForwards.whisper_model_forward( + self.model, + input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + head_mask=head_mask, + decoder_head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + decoder_position_ids=decoder_position_ids, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + shard_config=shard_config, + ) + if not in_decoder: + return outputs + + if not at_last_decoder_stage: + # encoder_hidden_states should be passed to the next stage + outputs["encoder_hidden_states"] = encoder_hidden_states + return outputs + + lm_logits = self.proj_out(outputs[0]) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + # move labels to correct device to enable PP + labels = labels.to(lm_logits.device) + loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.reshape(-1)) + + if not return_dict: + output = (lm_logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return Seq2SeqLMOutput( + loss=loss, + logits=lm_logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + @staticmethod + def whisper_for_audio_classification_forward( + self: WhisperForAudioClassification, + input_features: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + encoder_states=None, + all_attentions=None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ): + r""" + This function is modified on the basis of transformers.models.whisper.modeling_whisper.WhisperForAudioClassification.forward. + Please refer to original code of transformers for more details. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + if self.config.use_weighted_layer_sum: + output_hidden_states = True + elif output_hidden_states is None: + output_hidden_states = self.config.output_hidden_states + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # audio_classification only holds encoder + encoder_outputs = WhisperPipelineForwards.whisper_encoder_forward( + self.encoder, + input_features, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + ) + + if not stage_manager.is_last_stage(): + return encoder_outputs + + if self.config.use_weighted_layer_sum: + hidden_states = encoder_outputs[_HIDDEN_STATES_START_POSITION] + hidden_states = torch.stack(hidden_states, dim=1) + norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) + hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) + else: + hidden_states = encoder_outputs[0] + + hidden_states = self.projector(hidden_states) + pooled_output = hidden_states.mean(dim=1) + + logits = self.classifier(pooled_output) + + loss = None + + if labels is not None: + loss_fct = CrossEntropyLoss() + # move labels to correct device to enable PP + labels = labels.to(logits.device) + loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + encoder_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) diff --git a/colossalai/shardformer/modeling/_whisper_tf5.py b/colossalai/shardformer/modeling/_whisper_tf5.py new file mode 100644 index 000000000000..380784b22dda --- /dev/null +++ b/colossalai/shardformer/modeling/_whisper_tf5.py @@ -0,0 +1,709 @@ +"""Whisper adapters for Transformers 5.17 Cache and Tensor layer contracts.""" + +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn +from torch.nn import CrossEntropyLoss +from transformers.cache_utils import Cache, EncoderDecoderCache +from transformers.masking_utils import create_causal_mask +from transformers.modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPastAndCrossAttentions, + Seq2SeqLMOutput, + Seq2SeqModelOutput, + SequenceClassifierOutput, +) +from transformers.models.whisper.modeling_whisper import ( + _HIDDEN_STATES_START_POSITION, + WhisperForAudioClassification, + WhisperForConditionalGeneration, + WhisperModel, + shift_tokens_right, +) +from transformers.utils import logging +from transformers.utils.generic import merge_with_config_defaults +from transformers.utils.output_capturing import capture_outputs + +from colossalai.pipeline.stage_manager import PipelineStageManager +from colossalai.shardformer.layer import ColoAttention +from colossalai.shardformer.shard import ShardConfig + +from .._utils import ( + get_hf_cache_seq_length, + get_hf_flash_attention_kwargs, + get_hf_mask_config, + prepare_hf_encoder_decoder_cache, +) + +logger = logging.get_logger(__name__) + + +def _get_attention_mask( + self, + shard_config, + hidden_states, + past_key_values_length, + attention_mask, + head_mask=None, + output_attentions=False, + *, + past_key_values=None, + position_ids=None, +): + if head_mask is not None: + raise ValueError("HF 5.17 Whisper does not support head_mask") + flash = shard_config is not None and shard_config.enable_flash_attention + config = get_hf_mask_config(self.config, eager=flash) + if past_key_values is None and past_key_values_length: + # Preserve the length-only helper interface for callers without a Cache. + # Normal decode below always supplies the real Cache to masking_utils. + from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask + + return _prepare_4d_causal_attention_mask( + attention_mask, hidden_states.shape[:2], hidden_states, past_key_values_length + ) + return create_causal_mask( + config=config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + + +def get_whisper_flash_attention_forward(): + def forward( + self, + hidden_states, + key_value_states=None, + past_key_values=None, + attention_mask=None, + output_attentions=False, + **kwargs, + ): + is_cross = key_value_states is not None + input_shape = hidden_states.shape[:-1] + # Match HF's order of scaling (Whisper is sensitive to moving this + # multiplication after the QK product). + query = ( + (self.q_proj(hidden_states) * self.scaling) + .view(*input_shape, -1, self.head_dim) + .transpose(1, 2) + .contiguous() + ) + cache = past_key_values + updated = False + if isinstance(cache, EncoderDecoderCache): + updated = cache.is_updated.get(self.layer_idx, False) + cache = cache.cross_attention_cache if is_cross else cache.self_attention_cache + if is_cross and cache is not None and updated: + key = cache.layers[self.layer_idx].keys + value = cache.layers[self.layer_idx].values + else: + states = key_value_states if is_cross else hidden_states + shape = (input_shape[0], -1, self.num_heads, self.head_dim) + key = self.k_proj(states).view(shape).transpose(1, 2).contiguous() + value = self.v_proj(states).view(shape).transpose(1, 2).contiguous() + if cache is not None: + key, value = cache.update(key, value, self.layer_idx) + if is_cross and isinstance(past_key_values, EncoderDecoderCache): + past_key_values.is_updated[self.layer_idx] = True + output = ColoAttention.attention( + query, + key, + value, + **get_hf_flash_attention_kwargs(attention_mask), + dropout_p=self.dropout if self.training else 0.0, + scale=1.0, + ) + output = output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + return self.out_proj(output), None + + return forward + + +def get_whisper_decoder_forward_for_flash_attention(shard_config): + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + past_key_values=None, + inputs_embeds=None, + position_ids=None, + use_cache=None, + **kwargs, + ): + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids and inputs_embeds") + if self.gradient_checkpointing and self.training: + if use_cache: + logger.warning_once("Whisper checkpointing disables use_cache.") + use_cache = False + past_key_values = None + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + if use_cache: + past_key_values = prepare_hf_encoder_decoder_cache( + self.config, past_key_values, encoder_hidden_states=encoder_hidden_states + ) + past_length = get_hf_cache_seq_length(past_key_values) + if position_ids is None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_length + position_ids = position_ids.unsqueeze(0).expand(inputs_embeds.shape[0], -1) + positions = self.embed_positions(inputs_embeds, past_key_values_length=past_length, position_ids=position_ids) + hidden_states = nn.functional.dropout( + inputs_embeds + positions.to(inputs_embeds.device), p=self.dropout, training=self.training + ) + mask = _get_attention_mask( + self, + shard_config, + inputs_embeds, + past_length, + attention_mask, + past_key_values=past_key_values, + position_ids=position_ids, + ) + for layer in self.layers: + if self.training and torch.rand([]) < self.layerdrop: + continue + # Encoder states need positional binding for reentrant checkpoint + # to track the cross-attention gradient back into the encoder. + hidden_states = layer( + hidden_states, + mask, + encoder_hidden_states, + encoder_attention_mask=None, + past_key_values=past_key_values if use_cache else None, + use_cache=use_cache, + **kwargs, + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=self.layer_norm(hidden_states), + past_key_values=past_key_values, + ) + + return forward + + +def get_jit_fused_whisper_encoder_layer_forward(): + def forward(self, hidden_states, attention_mask=None, **kwargs): + residual = hidden_states + hidden_states, _ = self.self_attn( + self.self_attn_layer_norm(hidden_states), attention_mask=attention_mask, **kwargs + ) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + residual = hidden_states + hidden_states = self.activation_fn(self.fc1(self.final_layer_norm(hidden_states))) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + hidden_states = self.dropout_add(self.fc2(hidden_states), residual, self.dropout, self.training) + if hidden_states.dtype == torch.float16: + bound = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = hidden_states.clamp(-bound, bound) + return hidden_states + + return forward + + +def get_jit_fused_whisper_decoder_layer_forward(): + def forward( + self, + hidden_states, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + use_cache=True, + **kwargs, + ): + residual = hidden_states + hidden_states, _ = self.self_attn( + self.self_attn_layer_norm(hidden_states), + attention_mask=attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + if encoder_hidden_states is not None: + residual = hidden_states + hidden_states, _ = self.encoder_attn( + self.encoder_attn_layer_norm(hidden_states), + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + **kwargs, + ) + hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) + residual = hidden_states + hidden_states = self.activation_fn(self.fc1(self.final_layer_norm(hidden_states))) + hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) + return self.dropout_add(self.fc2(hidden_states), residual, self.dropout, self.training) + + return forward + + +class WhisperPipelineForwards: + """Pipeline stage functions; Cache generation remains a non-pipeline path.""" + + @staticmethod + def whisper_encoder_forward( + self, + input_features, + attention_mask=None, + head_mask=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + stage_manager=None, + hidden_states=None, + encoder_states=None, + all_attentions=None, + stage_index=None, + decoder_starting_stage=None, + shard_config=None, + ): + if head_mask is not None: + raise ValueError("HF 5.17 Whisper does not support head_mask") + if output_attentions: + logger.warning_once("Pipeline Whisper does not collect attention weights.") + return_dict = self.config.use_return_dict if return_dict is None else return_dict + output_hidden_states = ( + self.config.output_hidden_states if output_hidden_states is None else output_hidden_states + ) + if stage_manager.stage == 0: + expected = self.config.max_source_positions * self.conv1.stride[0] * self.conv2.stride[0] + if input_features.shape[-1] != expected: + raise ValueError(f"Whisper expects mel input length {expected}, got {input_features.shape[-1]}") + hidden_states = nn.functional.gelu(self.conv1(input_features)) + hidden_states = nn.functional.gelu(self.conv2(hidden_states)).permute(0, 2, 1) + positions = torch.arange(self.embed_positions.num_embeddings, device=hidden_states.device) + hidden_states = nn.functional.dropout( + hidden_states + self.embed_positions(positions), p=self.dropout, training=self.training + ) + encoder_states = () if output_hidden_states else None + elif hidden_states is None: + raise ValueError("hidden_states is required at non-first encoder stages") + if output_hidden_states and encoder_states is None: + raise ValueError("encoder_states must be forwarded when collecting pipeline hidden states") + for idx in range(*stage_index): + if output_hidden_states: + encoder_states += (hidden_states,) + if self.training and torch.rand([]) < self.layerdrop: + continue + hidden_states = self.layers[idx](hidden_states, attention_mask=None) + if stage_manager.stage == decoder_starting_stage - 1: + hidden_states = self.layer_norm(hidden_states) + if output_hidden_states: + encoder_states += (hidden_states,) + outputs = BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states) + return outputs if return_dict else outputs.to_tuple() + outputs = {"hidden_states": hidden_states} + if output_hidden_states: + outputs["encoder_states"] = encoder_states + return outputs + + @staticmethod + def whisper_decoder_forward( + self, + input_ids=None, + attention_mask=None, + encoder_hidden_states=None, + head_mask=None, + cross_attn_head_mask=None, + past_key_values=None, + inputs_embeds=None, + position_ids=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + stage_manager=None, + hidden_states=None, + stage_index=None, + decoder_starting_stage=None, + shard_config=None, + ): + if head_mask is not None or cross_attn_head_mask is not None: + raise ValueError("HF 5.17 Whisper does not support head masks") + if past_key_values is not None or use_cache: + logger.warning_once("Pipeline Whisper does not support Cache generation; cache is disabled.") + return_dict = self.config.use_return_dict if return_dict is None else return_dict + if stage_manager.stage == decoder_starting_stage: + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids and inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + positions = self.embed_positions(inputs_embeds, position_ids=position_ids) + hidden_states = nn.functional.dropout( + inputs_embeds + positions.to(inputs_embeds.device), p=self.dropout, training=self.training + ) + elif hidden_states is None: + raise ValueError("hidden_states is required at non-first decoder stages") + mask = _get_attention_mask(self, shard_config, hidden_states, 0, attention_mask, position_ids=position_ids) + for idx in range(*stage_index): + if self.training and torch.rand([]) < self.layerdrop: + continue + hidden_states = self.layers[idx]( + hidden_states, + mask, + encoder_hidden_states, + encoder_attention_mask=None, + past_key_values=None, + use_cache=False, + ) + if not stage_manager.is_last_stage(): + return {"hidden_states": hidden_states} + outputs = BaseModelOutputWithPastAndCrossAttentions(last_hidden_state=self.layer_norm(hidden_states)) + return outputs if return_dict else outputs.to_tuple() + + @staticmethod + def whisper_model_forward( + self: WhisperModel, + input_features: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.LongTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + decoder_head_mask: Optional[torch.Tensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Cache] = None, + decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, + decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ): + r""" + Returns: + + Example: + ```python + >>> import torch + >>> from transformers import AutoFeatureExtractor, WhisperModel + >>> from datasets import load_dataset + + >>> model = WhisperModel.from_pretrained("openai/whisper-base") + >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base") + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id + >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state + >>> list(last_hidden_state.shape) + [1, 2, 512] + ```""" + # Pipeline generation is not supported; never re-enable cache from config defaults. + if past_key_values is not None: + logger.warning_once("Non-empty past_key_values is not supported for pipeline models at the moment.") + past_key_values = None + if output_attentions: + logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") + output_attentions = False + if output_hidden_states: + logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") + output_hidden_states = False + if use_cache: + logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") + use_cache = False + + use_cache = False + output_attentions = False + output_hidden_states = False + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + in_decoder = stage_manager.stage >= decoder_starting_stage + if not in_decoder: + if encoder_outputs is None: + if stage_manager.stage == 0: + input_features = self._mask_input_features(input_features, attention_mask=attention_mask) + + encoder_outputs = WhisperPipelineForwards.whisper_encoder_forward( + self.encoder, + input_features, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + ) + + if stage_manager.stage == decoder_starting_stage - 1: + # last stage of encoder + return {"encoder_hidden_states": encoder_outputs[0]} + else: + return encoder_outputs + # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True + elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): + encoder_outputs = BaseModelOutput( + last_hidden_state=encoder_outputs[0], + hidden_states=(encoder_outputs[1] if len(encoder_outputs) > 1 else None), + attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, + ) + + at_last_decoder_stage = stage_manager.is_last_stage() + at_first_decoder_stage = stage_manager.stage == decoder_starting_stage + if encoder_outputs is not None: + encoder_hidden_states = encoder_outputs[0] + elif encoder_hidden_states is None: + raise ValueError("Non-empty encoder_hidden_states should be passed in at decoder stages.") + + if not at_first_decoder_stage and hidden_states is None: + raise ValueError("If not at the first layer of decoder, non-empty hidden_states must be provided.") + + # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn) + decoder_outputs = WhisperPipelineForwards.whisper_decoder_forward( + self.decoder, + input_ids=decoder_input_ids, + attention_mask=decoder_attention_mask, + encoder_hidden_states=encoder_hidden_states, + head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + past_key_values=past_key_values, + inputs_embeds=decoder_inputs_embeds, + position_ids=decoder_position_ids, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + shard_config=shard_config, + ) + + # Directly return outputs of overloaded Whisper forward if not at last stage. + if not at_last_decoder_stage: + # encoder_hidden_states should be passed to the next stage + decoder_outputs["encoder_hidden_states"] = encoder_hidden_states + return decoder_outputs + + if not return_dict: + return decoder_outputs + (encoder_hidden_states,) + + return Seq2SeqModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + past_key_values=decoder_outputs.past_key_values, + decoder_hidden_states=decoder_outputs.hidden_states, + decoder_attentions=decoder_outputs.attentions, + cross_attentions=decoder_outputs.cross_attentions, + encoder_last_hidden_state=encoder_hidden_states, + ) + + @staticmethod + def whisper_for_conditional_generation_forward( + self: WhisperForConditionalGeneration, + input_features: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.LongTensor] = None, + decoder_input_ids: Optional[torch.LongTensor] = None, + decoder_attention_mask: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + decoder_head_mask: Optional[torch.Tensor] = None, + cross_attn_head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + past_key_values: Optional[Cache] = None, + decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, + decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` + or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is + only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> import torch + >>> from transformers import AutoProcessor, WhisperForConditionalGeneration + >>> from datasets import load_dataset + + >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") + >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") + + >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") + + >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") + >>> input_features = inputs.input_features + + >>> generated_ids = model.generate(inputs=input_features) + + >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] + >>> transcription + ' Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.' + ```""" + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if labels is not None: + if decoder_input_ids is None and decoder_inputs_embeds is None: + decoder_input_ids = shift_tokens_right( + labels, self.config.pad_token_id, self.config.decoder_start_token_id + ) + in_decoder = stage_manager.stage >= decoder_starting_stage + at_last_decoder_stage = stage_manager.is_last_stage() + outputs = WhisperPipelineForwards.whisper_model_forward( + self.model, + input_features, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + encoder_outputs=encoder_outputs, + decoder_attention_mask=decoder_attention_mask, + head_mask=head_mask, + decoder_head_mask=decoder_head_mask, + cross_attn_head_mask=cross_attn_head_mask, + past_key_values=past_key_values, + decoder_inputs_embeds=decoder_inputs_embeds, + decoder_position_ids=decoder_position_ids, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + shard_config=shard_config, + ) + if not in_decoder: + return outputs + + if not at_last_decoder_stage: + # encoder_hidden_states should be passed to the next stage + outputs["encoder_hidden_states"] = encoder_hidden_states + return outputs + + lm_logits = self.proj_out(outputs[0]) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + # move labels to correct device to enable PP + labels = labels.to(lm_logits.device) + loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.reshape(-1)) + + if not return_dict: + output = (lm_logits,) + outputs[1:] + return ((loss,) + output) if loss is not None else output + + return Seq2SeqLMOutput( + loss=loss, + logits=lm_logits, + past_key_values=outputs.past_key_values, + decoder_hidden_states=outputs.decoder_hidden_states, + decoder_attentions=outputs.decoder_attentions, + cross_attentions=outputs.cross_attentions, + encoder_last_hidden_state=outputs.encoder_last_hidden_state, + encoder_hidden_states=outputs.encoder_hidden_states, + encoder_attentions=outputs.encoder_attentions, + ) + + @staticmethod + def whisper_for_audio_classification_forward( + self: WhisperForAudioClassification, + input_features: Optional[torch.LongTensor] = None, + head_mask: Optional[torch.Tensor] = None, + encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, + labels: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + stage_manager: Optional[PipelineStageManager] = None, + hidden_states: Optional[torch.FloatTensor] = None, + encoder_states=None, + all_attentions=None, + stage_index: Optional[List[int]] = None, + decoder_starting_stage: Optional[int] = None, + shard_config: Optional[ShardConfig] = None, + ): + r""" + This function is modified on the basis of transformers.models.whisper.modeling_whisper.WhisperForAudioClassification.forward. + Please refer to original code of transformers for more details. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + if self.config.use_weighted_layer_sum: + output_hidden_states = True + elif output_hidden_states is None: + output_hidden_states = self.config.output_hidden_states + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # audio_classification only holds encoder + encoder_outputs = WhisperPipelineForwards.whisper_encoder_forward( + self.encoder, + input_features, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + stage_manager=stage_manager, + hidden_states=hidden_states, + encoder_states=encoder_states, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + ) + + if not stage_manager.is_last_stage(): + return encoder_outputs + + if self.config.use_weighted_layer_sum: + hidden_states = encoder_outputs[_HIDDEN_STATES_START_POSITION] + hidden_states = torch.stack(hidden_states, dim=1) + norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) + hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) + else: + hidden_states = encoder_outputs[0] + + hidden_states = self.projector(hidden_states) + pooled_output = hidden_states.mean(dim=1) + + logits = self.classifier(pooled_output) + + loss = None + + if labels is not None: + loss_fct = CrossEntropyLoss() + # move labels to correct device to enable PP + labels = labels.to(logits.device) + loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) + + if not return_dict: + output = (logits,) + encoder_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) diff --git a/colossalai/shardformer/modeling/bert.py b/colossalai/shardformer/modeling/bert.py index b788cbf58644..6de3149f9d09 100644 --- a/colossalai/shardformer/modeling/bert.py +++ b/colossalai/shardformer/modeling/bert.py @@ -1,1293 +1,9 @@ -import warnings -from typing import List, Optional, Tuple, Union +r"""BERT modeling, selected by the installed transformers major version. -import torch -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from transformers.modeling_outputs import ( - BaseModelOutputWithPoolingAndCrossAttentions, - CausalLMOutputWithCrossAttentions, - MaskedLMOutput, - MultipleChoiceModelOutput, - NextSentencePredictorOutput, - QuestionAnsweringModelOutput, - SequenceClassifierOutput, - TokenClassifierOutput, -) -from transformers.models.bert.modeling_bert import ( - BertForMaskedLM, - BertForMultipleChoice, - BertForNextSentencePrediction, - BertForPreTraining, - BertForPreTrainingOutput, - BertForQuestionAnswering, - BertForSequenceClassification, - BertForTokenClassification, - BertLMHeadModel, - BertModel, -) -from transformers.utils import logging +``_bert_tf5`` targets transformers v5; ``_bert_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer import ShardConfig -from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward +from colossalai._compat import reexport_transformers_impl - -class BertPipelineForwards: - """ - This class serves as a micro library for forward function substitution of Bert models - under pipeline setting. - """ - - @staticmethod - def bert_model_forward( - self: BertModel, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, # this is from the previous stage - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: - # TODO(jianghai): add explaination of the output here. - r""" - encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): - Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if - the model is configured as a decoder. - encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in - the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): - Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that - don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all - `decoder_input_ids` of shape `(batch_size, sequence_length)`. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see - `past_key_values`). - """ - logger = logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if self.config.is_decoder: - use_cache = use_cache if use_cache is not None else self.config.use_cache - else: - use_cache = False - - if stage_manager.is_first_stage(): - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - batch_size, seq_length = input_shape - device = input_ids.device if input_ids is not None else inputs_embeds.device - if token_type_ids is None: - if hasattr(self.embeddings, "token_type_ids"): - buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] - buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) - token_type_ids = buffered_token_type_ids_expanded - else: - token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) - else: - input_shape = hidden_states.size()[:-1] - batch_size, seq_length = input_shape - device = hidden_states.device - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - if use_cache: - logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") - use_cache = False - - # past_key_values_length - past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 - - if attention_mask is None: - attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) - - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) - attention_mask = extended_attention_mask - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if self.config.is_decoder and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_extended_attention_mask = None - - # Prepare head mask if needed - # 1.0 in head_mask indicate we keep the head - # attention_probs has shape bsz x n_heads x N x N - # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] - # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] - head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) - hidden_states = hidden_states if hidden_states is not None else None - - if stage_manager.is_first_stage(): - hidden_states = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - token_type_ids=token_type_ids, - inputs_embeds=inputs_embeds, - past_key_values_length=past_key_values_length, - ) - - # inherit from bert_layer,this should be changed when we add the feature to record hidden_states - all_hidden_states = () if output_hidden_states else None - all_self_attentions = () if output_attentions else None - all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None - - if self.encoder.gradient_checkpointing and self.encoder.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - next_decoder_cache = () if use_cache else None - - start_idx, end_idx = stage_index[0], stage_index[1] - # layer_outputs - layer_outputs = hidden_states if hidden_states is not None else None - - # split the input tensor along sequence dimension - # [batch_size, seq_len, hidden_size] -> [batch_size, seq_len/TP_size, hidden_size] - if shard_config is not None and shard_config.enable_sequence_parallelism: - if shard_config.sequence_parallelism_mode == "split_gather": - hidden_states = split_forward_gather_backward( - hidden_states, - dim=1, - process_group=shard_config.tensor_parallel_process_group, - fp8_communication=shard_config.fp8_communication, - ) - if encoder_hidden_states is not None: - encoder_hidden_states = split_forward_gather_backward( - encoder_hidden_states, - dim=1, - process_group=shard_config.tensor_parallel_process_group, - fp8_communication=shard_config.fp8_communication, - ) - - for idx, encoder_layer in enumerate(self.encoder.layer[start_idx:end_idx], start=start_idx): - if stage_manager.is_first_stage() and idx == 0: - encoder_attention_mask = encoder_extended_attention_mask - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - layer_head_mask = head_mask[idx] if head_mask is not None else None - past_key_value = past_key_values[idx] if past_key_values is not None else None - - if self.encoder.gradient_checkpointing and self.encoder.training: - - def create_custom_forward(module): - def custom_forward(*inputs): - return module(*inputs, past_key_value, output_attentions) - - return custom_forward - - layer_outputs = torch.utils.checkpoint.checkpoint( - create_custom_forward(encoder_layer), - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - ) - else: - layer_outputs = encoder_layer( - hidden_states, - attention_mask, - layer_head_mask, - encoder_hidden_states, - encoder_attention_mask, - past_key_value, - output_attentions, - ) - hidden_states = layer_outputs[0] - if use_cache: - next_decoder_cache += (layer_outputs[-1],) - if output_attentions: - all_self_attentions = all_self_attentions + (layer_outputs[1],) - if self.config.add_cross_attention: - all_cross_attentions = all_cross_attentions + (layer_outputs[2],) - - # When sequence parallelism done, gather the output tensor in forward and split it in backward - if shard_config is not None and shard_config.enable_sequence_parallelism: - if shard_config.sequence_parallelism_mode == "split_gather": - hidden_states = gather_forward_split_backward( - hidden_states, - dim=1, - process_group=shard_config.tensor_parallel_process_group, - fp8_communication=shard_config.fp8_communication, - ) - - if output_hidden_states: - all_hidden_states = all_hidden_states + (hidden_states,) - - # end of a stage loop - sequence_output = hidden_states if hidden_states is not None else None - - if stage_manager.is_last_stage(): - pooled_output = self.pooler(sequence_output) if self.pooler is not None else None - if not return_dict: - return (sequence_output, pooled_output) + layer_outputs[1:] - # return dict is not supported at this moment - else: - return BaseModelOutputWithPoolingAndCrossAttentions( - last_hidden_state=sequence_output, - pooler_output=pooled_output, - past_key_values=next_decoder_cache, - hidden_states=all_hidden_states, - attentions=all_self_attentions, - cross_attentions=all_cross_attentions, - ) - - # output of non-first and non-last stages: must be a dict - else: - # intermediate stage always return dict - return { - "hidden_states": hidden_states, - } - - @staticmethod - def bert_for_pretraining_forward( - self: BertForPreTraining, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - next_sentence_label: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # TODO(jianghai) left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states if hidden_states is not None else None, - stage_index=stage_index, - shard_config=shard_config, - ) - - if stage_manager.is_last_stage(): - sequence_output, pooled_output = outputs[:2] - prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) - # the last stage for pretraining model - total_loss = None - if labels is not None and next_sentence_label is not None: - loss_fct = CrossEntropyLoss() - masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) - next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1)) - total_loss = masked_lm_loss + next_sentence_loss - - if not return_dict: - output = (prediction_scores, seq_relationship_score) + outputs[2:] - return ((total_loss,) + output) if total_loss is not None else output - - return BertForPreTrainingOutput( - loss=total_loss, - prediction_logits=prediction_scores, - seq_relationship_logits=seq_relationship_score, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - - # intermediate stage always return dict - return { - "hidden_states": hidden_states, - } - - @staticmethod - def bert_lm_head_model_forward( - self: BertLMHeadModel, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.Tensor]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): - Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if - the model is configured as a decoder. - encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in - the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in - `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are - ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]` - past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): - Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that - don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all - `decoder_input_ids` of shape `(batch_size, sequence_length)`. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see - `past_key_values`). - """ - logger = logging.get_logger(__name__) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if labels is not None: - use_cache = False - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states if hidden_states is not None else None, - stage_index=stage_index, - shard_config=shard_config, - ) - past_key_values = None - - if stage_manager.is_last_stage(): - sequence_output = outputs[0] - prediction_scores = self.cls(sequence_output) - - lm_loss = None - if labels is not None: - # we are doing next-token prediction; shift prediction scores and input ids by one - shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous() - labels = labels[:, 1:].contiguous() - loss_fct = CrossEntropyLoss() - lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) - - if not return_dict: - output = (prediction_scores,) + outputs[2:] - return ((lm_loss,) + output) if lm_loss is not None else output - - return CausalLMOutputWithCrossAttentions( - loss=lm_loss, - logits=prediction_scores, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - cross_attentions=outputs.cross_attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - # intermediate stage always return dict - return {"hidden_states": hidden_states} - - @staticmethod - def bert_for_masked_lm_forward( - self: BertForMaskedLM, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.Tensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., - config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the - loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - hidden_states=hidden_states, - stage_manager=stage_manager, - stage_index=stage_index, - shard_config=shard_config, - ) - - if stage_manager.is_last_stage(): - sequence_output = outputs[0] - prediction_scores = self.cls(sequence_output) - - masked_lm_loss = None - if labels is not None: - loss_fct = CrossEntropyLoss() # -100 index = padding token - masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)) - - if not return_dict: - output = (prediction_scores,) + outputs[2:] - return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output - - return MaskedLMOutput( - loss=masked_lm_loss, - logits=prediction_scores, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - @staticmethod - def bert_for_next_sentence_prediction_forward( - self: BertForNextSentencePrediction, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.Tensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - **kwargs, - ): - # -> Union[Tuple[torch.Tensor], NextSentencePredictorOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair - (see `input_ids` docstring). Indices should be in `[0, 1]`: - - - 0 indicates sequence B is a continuation of sequence A, - - 1 indicates sequence B is a random sequence. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, BertForNextSentencePrediction - >>> import torch - - >>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") - >>> model = BertForNextSentencePrediction.from_pretrained("bert-base-uncased") - - >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced." - >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light." - >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt") - - >>> outputs = model(**encoding, labels=torch.LongTensor([1])) - >>> logits = outputs.logits - >>> assert logits[0, 0] < logits[0, 1] # next sentence was random - ``` - """ - logger = logging.get_logger(__name__) - - if "next_sentence_label" in kwargs: - warnings.warn( - "The `next_sentence_label` argument is deprecated and will be removed in a future version, use" - " `labels` instead.", - FutureWarning, - ) - labels = kwargs.pop("next_sentence_label") - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - hidden_states=hidden_states, - stage_manager=stage_manager, - stage_index=stage_index, - shard_config=shard_config, - ) - - if stage_manager.is_last_stage(): - pooled_output = outputs[1] - seq_relationship_scores = self.cls(pooled_output) - - next_sentence_loss = None - if labels is not None: - loss_fct = CrossEntropyLoss() - next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1)) - - if not return_dict: - output = (seq_relationship_scores,) + outputs[2:] - return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output - - return NextSentencePredictorOutput( - loss=next_sentence_loss, - logits=seq_relationship_scores, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - # intermediate stage always return dict - return {"hidden_states": hidden_states} - - @staticmethod - def bert_for_sequence_classification_forward( - self: BertForSequenceClassification, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.Tensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If - `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - hidden_states=hidden_states, - stage_manager=stage_manager, - stage_index=stage_index, - shard_config=shard_config, - ) - - if stage_manager.is_last_stage(): - pooled_output = outputs[1] - - pooled_output = self.dropout(pooled_output) - logits = self.classifier(pooled_output) - - loss = None - if labels is not None: - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.num_labels == 1: - loss = loss_fct(logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - if not return_dict: - output = (logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutput( - loss=loss, - logits=logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - @staticmethod - def bert_for_token_classification_forward( - self: BertForTokenClassification, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.Tensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - hidden_states=hidden_states, - stage_manager=stage_manager, - stage_index=stage_index, - shard_config=shard_config, - ) - - if stage_manager.is_last_stage(): - sequence_output = outputs[0] - - sequence_output = self.dropout(sequence_output) - logits = self.classifier(sequence_output) - - loss = None - if labels is not None: - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - - if not return_dict: - output = (logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output - - return TokenClassifierOutput( - loss=loss, - logits=logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - @staticmethod - def bert_for_multiple_choice_forward( - self: BertForMultipleChoice, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.Tensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., - num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See - `input_ids` above) - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - # in our pipeline design,input ids are copied for every stage and shouldn't be none - # the input_ids for multiple choice model is [batch_size, num_choices, sequence_length] - if stage_manager.is_last_stage(): - num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] - - input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None - attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None - token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None - position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None - inputs_embeds = ( - inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) - if inputs_embeds is not None - else None - ) - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - hidden_states=hidden_states, - stage_manager=stage_manager, - stage_index=stage_index, - shard_config=shard_config, - ) - if stage_manager.is_last_stage(): - pooled_output = outputs[1] - pooled_output = self.dropout(pooled_output) - logits = self.classifier(pooled_output) - reshaped_logits = logits.view(-1, num_choices) - - loss = None - if labels is not None: - loss_fct = CrossEntropyLoss() - loss = loss_fct(reshaped_logits, labels) - - if not return_dict: - output = (reshaped_logits,) + outputs[2:] - return ((loss,) + output) if loss is not None else output - - return MultipleChoiceModelOutput( - loss=loss, - logits=reshaped_logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - @staticmethod - def bert_for_question_answering_forward( - self: BertForQuestionAnswering, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - start_positions: Optional[torch.Tensor] = None, - end_positions: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.Tensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - # NOTE: the arg start_position and end_position are used only for the last stage - r""" - start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the start of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence - are not taken into account for computing the loss. - end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for position (index) of the end of the labelled span for computing the token classification loss. - Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence - are not taken into account for computing the loss. - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - outputs = BertPipelineForwards.bert_model_forward( - self.bert, - input_ids, - attention_mask=attention_mask, - token_type_ids=token_type_ids, - position_ids=position_ids, - head_mask=head_mask, - inputs_embeds=inputs_embeds, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - hidden_states=hidden_states, - stage_manager=stage_manager, - stage_index=stage_index, - shard_config=shard_config, - ) - if stage_manager.is_last_stage(): - sequence_output = outputs[0] - - logits = self.qa_outputs(sequence_output) - start_logits, end_logits = logits.split(1, dim=-1) - start_logits = start_logits.squeeze(-1).contiguous() - end_logits = end_logits.squeeze(-1).contiguous() - - total_loss = None - if start_positions is not None and end_positions is not None: - # If we are on multi-GPU, split add a dimension - if len(start_positions.size()) > 1: - start_positions = start_positions.squeeze(-1) - if len(end_positions.size()) > 1: - end_positions = end_positions.squeeze(-1) - # sometimes the start/end positions are outside our model inputs, we ignore these terms - ignored_index = start_logits.size(1) - start_positions = start_positions.clamp(0, ignored_index) - end_positions = end_positions.clamp(0, ignored_index) - - loss_fct = CrossEntropyLoss(ignore_index=ignored_index) - start_loss = loss_fct(start_logits, start_positions) - end_loss = loss_fct(end_logits, end_positions) - total_loss = (start_loss + end_loss) / 2 - - if not return_dict: - output = (start_logits, end_logits) + outputs[2:] - return ((total_loss,) + output) if total_loss is not None else output - - return QuestionAnsweringModelOutput( - loss=total_loss, - start_logits=start_logits, - end_logits=end_logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - -def get_jit_fused_bert_self_output_forward(): - from transformers.models.bert.modeling_bert import BertSelfOutput - - def forward(self: BertSelfOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: - hidden_states = self.dense(hidden_states) - hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) - hidden_states = self.LayerNorm(hidden_states) - return hidden_states - - return forward - - -def get_jit_fused_bert_output_forward(): - from transformers.models.bert.modeling_bert import BertOutput - - def forward(self: BertOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: - hidden_states = self.dense(hidden_states) - hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) - hidden_states = self.LayerNorm(hidden_states) - return hidden_states - - return forward - - -# Fix the tgt_len size in sequence parallel attention: -# same with the one in BertSdpaSelfAttention forward in v4.51.3 transformers except the -def get_bert_sequence_parallel_attention_forward(shard_config: ShardConfig): - from transformers.models.bert.modeling_bert import BertSdpaSelfAttention - - def forward( - self: BertSdpaSelfAttention, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.FloatTensor] = None, - head_mask: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - encoder_attention_mask: Optional[torch.FloatTensor] = None, - past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - output_attentions: Optional[bool] = False, - ) -> Tuple[torch.Tensor]: - - bsz, tgt_len, _ = hidden_states.size() - - query_layer = self.transpose_for_scores(self.query(hidden_states)) - - # If this is instantiated as a cross-attention module, the keys and values come from an encoder; the attention - # mask needs to be such that the encoder's padding tokens are not attended to. - is_cross_attention = encoder_hidden_states is not None - - current_states = encoder_hidden_states if is_cross_attention else hidden_states - attention_mask = encoder_attention_mask if is_cross_attention else attention_mask - - # Check `seq_length` of `past_key_value` == `len(current_states)` to support prefix tuning - if is_cross_attention and past_key_value and past_key_value[0].shape[2] == current_states.shape[1]: - key_layer, value_layer = past_key_value - else: - key_layer = self.transpose_for_scores(self.key(current_states)) - value_layer = self.transpose_for_scores(self.value(current_states)) - if past_key_value is not None and not is_cross_attention: - key_layer = torch.cat([past_key_value[0], key_layer], dim=2) - value_layer = torch.cat([past_key_value[1], value_layer], dim=2) - - if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. - # Further calls to cross_attention layer can then reuse all cross-attention - # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of - # all previous decoder key/value_states. Further calls to uni-directional self-attention - # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) - # if encoder bi-directional self-attention `past_key_value` is always `None` - past_key_value = (key_layer, value_layer) - - # SDPA with memory-efficient backend is broken in torch==2.1.2 when using non-contiguous inputs and a custom - # attn_mask, so we need to call `.contiguous()` here. This was fixed in torch==2.2.0. - # Reference: https://github.com/pytorch/pytorch/issues/112577 - if self.require_contiguous_qkv and query_layer.device.type == "cuda" and attention_mask is not None: - query_layer = query_layer.contiguous() - key_layer = key_layer.contiguous() - value_layer = value_layer.contiguous() - - # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment - # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling. - # The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create - # a causal mask in case tgt_len == 1. - is_causal = ( - True if self.is_decoder and not is_cross_attention and attention_mask is None and tgt_len > 1 else False - ) - attn_output = torch.nn.functional.scaled_dot_product_attention( - query_layer, - key_layer, - value_layer, - attn_mask=attention_mask, - dropout_p=self.dropout_prob if self.training else 0.0, - is_causal=is_causal, - ) - - attn_output = attn_output.transpose(1, 2) - _, _, tgt_len, _ = query_layer.shape - attn_output = attn_output.reshape(bsz, tgt_len, self.all_head_size) - - outputs = (attn_output,) - if self.is_decoder: - outputs = outputs + (past_key_value,) - return outputs - - return forward - - -def bert_sequence_parallel_forward_fn(shard_config: ShardConfig): - def forward( - self, - input_ids: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - token_type_ids: Optional[torch.Tensor] = None, - position_ids: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - inputs_embeds: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]: - r""" - encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): - Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if - the model is configured as a decoder. - encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in - the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`): - Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that - don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all - `decoder_input_ids` of shape `(batch_size, sequence_length)`. - use_cache (`bool`, *optional*): - If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see - `past_key_values`). - """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if self.config.is_decoder: - use_cache = use_cache if use_cache is not None else self.config.use_cache - else: - use_cache = False - - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - batch_size, seq_length = input_shape - device = input_ids.device if input_ids is not None else inputs_embeds.device - - # past_key_values_length - past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 - - if attention_mask is None: - attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device) - - if token_type_ids is None: - if hasattr(self.embeddings, "token_type_ids"): - buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length] - buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length) - token_type_ids = buffered_token_type_ids_expanded - else: - token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) - - # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length] - # ourselves in which case we just need to make it broadcastable to all heads. - extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape) - - # If a 2D or 3D attention mask is provided for the cross-attention - # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length] - if self.config.is_decoder and encoder_hidden_states is not None: - encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size() - encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length) - if encoder_attention_mask is None: - encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device) - encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask) - else: - encoder_extended_attention_mask = None - - # Prepare head mask if needed - # 1.0 in head_mask indicate we keep the head - # attention_probs has shape bsz x n_heads x N x N - # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] - # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] - head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) - - embedding_output = self.embeddings( - input_ids=input_ids, - position_ids=position_ids, - token_type_ids=token_type_ids, - inputs_embeds=inputs_embeds, - past_key_values_length=past_key_values_length, - ) - - # split the input tensor along sequence dimension - # [batch_size, seq_len, hidden_size] -> [batch_size, seq_len/TP_size, hidden_size] - embedding_output = split_forward_gather_backward( - embedding_output, - dim=1, - process_group=shard_config.tensor_parallel_process_group, - fp8_communication=shard_config.fp8_communication, - ) - if encoder_hidden_states is not None: - encoder_hidden_states = split_forward_gather_backward( - encoder_hidden_states, - dim=1, - process_group=shard_config.tensor_parallel_process_group, - fp8_communication=shard_config.fp8_communication, - ) - - encoder_outputs = self.encoder( - embedding_output, - attention_mask=extended_attention_mask, - head_mask=head_mask, - encoder_hidden_states=encoder_hidden_states, - encoder_attention_mask=encoder_extended_attention_mask, - past_key_values=past_key_values, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - - sequence_output = encoder_outputs[0] - - # When sequence parallelism done, gather the output tensor in forward and split it in backward - sequence_output = gather_forward_split_backward( - sequence_output, - dim=1, - process_group=shard_config.tensor_parallel_process_group, - fp8_communication=shard_config.fp8_communication, - ) - - pooled_output = self.pooler(sequence_output) if self.pooler is not None else None - - if not return_dict: - return (sequence_output, pooled_output) + encoder_outputs[1:] - - return BaseModelOutputWithPoolingAndCrossAttentions( - last_hidden_state=sequence_output, - pooler_output=pooled_output, - past_key_values=encoder_outputs.past_key_values, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - cross_attentions=encoder_outputs.cross_attentions, - ) - - return forward - - -def get_jit_fused_bert_intermediate_forward(): - from transformers.models.bert.modeling_bert import BertIntermediate - - from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction - - def forward(self: BertIntermediate, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states, bias = self.dense(hidden_states) - hidden_states = JitGeLUFunction.apply(hidden_states, bias) - return hidden_states - - return forward +reexport_transformers_impl(globals(), __package__, "_bert") diff --git a/colossalai/shardformer/modeling/command.py b/colossalai/shardformer/modeling/command.py index aa11b7043c37..ade0a269ee67 100644 --- a/colossalai/shardformer/modeling/command.py +++ b/colossalai/shardformer/modeling/command.py @@ -2,14 +2,13 @@ import torch from torch import nn -from transformers.cache_utils import Cache, DynamicCache +from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.cohere.modeling_cohere import ( CohereAttention, CohereForCausalLM, CohereModel, - StaticCache, apply_rotary_pos_emb, repeat_kv, ) diff --git a/colossalai/shardformer/modeling/deepseek_v3.py b/colossalai/shardformer/modeling/deepseek_v3.py index c582ac9a3265..7ad71ec9d690 100644 --- a/colossalai/shardformer/modeling/deepseek_v3.py +++ b/colossalai/shardformer/modeling/deepseek_v3.py @@ -1,370 +1,9 @@ -from typing import List, Optional, Tuple, Union +r"""DeepSeek-v3 modeling, selected by the installed transformers major version. -import numpy as np -import torch -import torch.distributed as dist -from torch.distributed import ProcessGroup -from torch.nn import CrossEntropyLoss -from transformers.cache_utils import Cache, DynamicCache -from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask -from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +``_deepseek_v3_tf5`` targets transformers v5; ``_deepseek_v3_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -from colossalai.lazy import LazyInitContext -from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer.linear import ParallelModule -from colossalai.shardformer.shard.utils import set_tensors_to_none -from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_group +from colossalai._compat import reexport_transformers_impl - -class EpDeepseekV3MoE(ParallelModule): - """ - A mixed expert module containing shared experts. - """ - - def __init__(self, config): - raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}") - - def setup_process_groups( - self, - moe_dp_group: ProcessGroup, - ep_group: ProcessGroup, - ): - assert moe_dp_group is not None - assert ep_group is not None - - self.ep_size = dist.get_world_size(ep_group) - self.ep_rank = dist.get_rank(ep_group) - self.num_experts = self.config.n_routed_experts - assert self.num_experts % self.ep_size == 0 - - self.ep_group = ep_group - self.num_experts_per_ep = self.num_experts // self.ep_size - self.experts_per_rank = self.num_experts_per_ep - self.expert_start_idx = self.ep_rank * self.num_experts_per_ep - held_experts = self.experts[self.expert_start_idx : self.expert_start_idx + self.num_experts_per_ep] - - set_tensors_to_none(self.experts, exclude=set(held_experts)) - - # setup moe_dp group - self.moe_dp_group = moe_dp_group - self.moe_dp_size = dist.get_world_size(moe_dp_group) - - for p in self.experts.parameters(): - set_moe_tensor_ep_group(p, ep_group) - - @staticmethod - def from_native_module( - module, - moe_dp_group: ProcessGroup, - ep_group: ProcessGroup, - *args, - **kwargs, - ) -> "EpDeepseekV3MoE": - if module.__class__.__name__ != "DeepseekV3MLP": - module.__class__ = EpDeepseekV3MoE - module.setup_process_groups(moe_dp_group, ep_group) - LazyInitContext.materialize(module) - return module - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - identity = hidden_states - orig_shape = hidden_states.shape - topk_idx, topk_weight = self.gate(hidden_states) - hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) - y = self.moe_forward(hidden_states, topk_idx, topk_weight).view(*orig_shape) - if self.config.n_shared_experts is not None: - y = y + self.shared_experts(identity) - return y - - def moe_forward(self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor) -> torch.Tensor: - cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts))) - cnts.scatter_(1, topk_ids, 1) - tokens_per_expert = cnts.sum(dim=0) - idxs = topk_ids.view(-1).argsort() - sorted_tokens = x[idxs // topk_ids.shape[1]] - if self.ep_size > 1: - tokens_per_ep_rank = tokens_per_expert.view(self.ep_size, -1).sum(dim=1) - tokens_per_expert_group = tokens_per_expert.new_empty(tokens_per_expert.shape[0]) - dist.all_to_all_single(tokens_per_expert_group, tokens_per_expert, group=self.ep_group) - - output_splits = tokens_per_expert_group.view(self.ep_size, -1).sum(1).tolist() - input_split_sizes = tokens_per_ep_rank.tolist() - - gathered_tokens, _ = all_to_all_uneven(sorted_tokens, input_split_sizes, output_splits, self.ep_group) - tokens_per_expert_post_gather = tokens_per_expert_group.view(self.ep_size, self.experts_per_rank).sum(dim=0) - gatherd_idxs = np.zeros(shape=(gathered_tokens.shape[0],), dtype=np.int32) - s = 0 - for i, k in enumerate(tokens_per_expert_group.cpu().numpy()): - gatherd_idxs[s : s + k] = i % self.experts_per_rank - s += k - gatherd_idxs = gatherd_idxs.argsort() - sorted_tokens = gathered_tokens[gatherd_idxs] - tokens_per_expert = tokens_per_expert_post_gather - - # ep related code - sorted_tokens = EPGradScalerIn.apply(sorted_tokens, self.ep_size) - - tokens_per_expert = tokens_per_expert.cpu().numpy() - - outputs = [] - start_idx = 0 - for i, num_tokens in enumerate(tokens_per_expert): - end_idx = start_idx + num_tokens - if num_tokens == 0: - continue - expert = self.experts[i + self.ep_rank * self.experts_per_rank] - tokens_for_this_expert = sorted_tokens[start_idx:end_idx] - expert_out = expert(tokens_for_this_expert) - outputs.append(expert_out) - start_idx = end_idx - - if len(outputs) > 0: - outs = torch.cat(outputs, dim=0) - else: - assert sorted_tokens.numel() == 0, f"sorted_tokens: should be empty, but got {sorted_tokens.shape}" - outs = sorted_tokens - - if self.ep_size > 1: - outs = EPGradScalerOut.apply(outs, self.ep_size) - new_x = torch.empty_like(outs) - new_x[gatherd_idxs] = outs - gathered_tokens, _ = all_to_all_uneven(new_x, output_splits, input_split_sizes, self.ep_group) - outs = gathered_tokens - - new_x = torch.empty_like(outs) - new_x[idxs] = outs - final_out = ( - (new_x.view(*topk_ids.shape, -1).type(topk_weight.dtype) * topk_weight.unsqueeze(dim=-1)) - .sum(dim=1) - .type(new_x.dtype) - ) - - return final_out - - -def deepseek_v3_model_forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - hidden_states_internal: Optional[torch.Tensor] = None, -) -> Union[Tuple, BaseModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape[:2] - elif inputs_embeds is not None: - batch_size, seq_length = inputs_embeds.shape[:2] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - - past_key_values_length = 0 - if use_cache: - use_legacy_cache = not isinstance(past_key_values, Cache) - if use_legacy_cache: - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - past_key_values_length = past_key_values.get_usable_length(seq_length) - - if position_ids is None: - device = input_ids.device if input_ids is not None else inputs_embeds.device - position_ids = torch.arange( - past_key_values_length, - seq_length + past_key_values_length, - dtype=torch.long, - device=device, - ) - position_ids = position_ids.unsqueeze(0) - - if stage_manager is None or stage_manager.is_first_stage(): - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - else: - inputs_embeds = hidden_states_internal - - if self._use_flash_attention_2: - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, - (batch_size, seq_length), - inputs_embeds, - past_key_values_length, - ) - - # embed positions - hidden_states = inputs_embeds - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - - if stage_index is not None: - start_idx, end_idx = stage_index - else: - start_idx, end_idx = 0, len(self.layers) - for i, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and i > 0: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - ) - - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache = layer_outputs[2 if output_attentions else 1] - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if stage_manager is None or stage_manager.is_last_stage(): - hidden_states = self.norm(hidden_states) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = None - if use_cache: - next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache - if stage_manager is not None and not stage_manager.is_last_stage(): - return { - "hidden_states_internal": hidden_states, - } - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - - -def deepseek_v3_for_causal_lm_forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - stage_index: Optional[List[int]] = None, - hidden_states_internal: Optional[torch.Tensor] = None, -) -> Union[Tuple, CausalLMOutputWithPast]: - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`. - Returns: - Example: - ```python - >>> from transformers import AutoTokenizer, DeepseekV3ForCausalLM - >>> model = DeepseekV3ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = deepseek_v3_model_forward( - self.model, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - stage_index=stage_index, - hidden_states_internal=hidden_states_internal, - ) - if stage_manager is not None and not stage_manager.is_last_stage(): - return outputs - - hidden_states = outputs[0] - - logits = self.lm_head(hidden_states) - logits = logits.float() - - loss = None - if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) +reexport_transformers_impl(globals(), __package__, "_deepseek_v3") diff --git a/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py index 3f9b7b12880c..82b2427175e0 100644 --- a/colossalai/shardformer/modeling/llama.py +++ b/colossalai/shardformer/modeling/llama.py @@ -1,616 +1,9 @@ -import math -import warnings -from typing import Dict, List, Optional, Tuple, Union +r"""Llama modeling, selected by the installed transformers major version. -import torch -import torch.distributed -import torch.utils.checkpoint -from torch import nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from transformers.cache_utils import Cache, DynamicCache -from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask -from transformers.modeling_outputs import ( - BaseModelOutputWithPast, - CausalLMOutputWithPast, - SequenceClassifierOutputWithPast, -) -from transformers.models.llama.modeling_llama import ( - LlamaForCausalLM, - LlamaForSequenceClassification, - LlamaModel, - StaticCache, - apply_rotary_pos_emb, - repeat_kv, -) -from transformers.utils import logging +``_llama_tf5`` targets transformers v5; ``_llama_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer._operation import all_to_all_comm, gather_sp_output, split_forward_gather_backward -from colossalai.shardformer.layer.utils import is_share_sp_tp, split_batch_zigzag -from colossalai.shardformer.shard import ShardConfig +from colossalai._compat import reexport_transformers_impl -from ..layer import ColoAttention, RingAttention, dist_cross_entropy - -_SUPPORTED_SP_MODE = ["all_to_all", "split_gather", "ring", "ring_attn"] - - -class LlamaPipelineForwards: - """ - This class serves as a micro library for forward function substitution of Llama models - under pipeline setting. - """ - - @staticmethod - def llama_model_forward( - self: LlamaModel, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - force_sp_gather: bool = True, # Set to false only when computing cross entropy - ): - logger = logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with pipeline parallelism. Setting `use_cache=False`..." - ) - use_cache = False - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - disable_pp = stage_manager is None - # retrieve input_ids and inputs_embeds - if disable_pp or stage_manager.is_first_stage(): - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape[:2] - elif inputs_embeds is not None: - batch_size, seq_length = inputs_embeds.shape[:2] - else: - raise ValueError("You have to specify either input_ids or inputs_embeds") - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - hidden_states = inputs_embeds - device = hidden_states.device - else: - input_shape = hidden_states.shape[:-1] - batch_size, seq_length = input_shape - device = hidden_states.device - - # Support SP + PP - sp_mode = shard_config.sequence_parallelism_mode - sp_group = shard_config.sequence_parallel_process_group - sp_size = shard_config.sequence_parallel_size - # Generating full positions ids for modes that gather sequence before attn - if stage_manager and (sp_mode != "ring_attn" and not stage_manager.is_first_stage()): - seq_length *= sp_size - - past_seen_tokens = 0 - if use_cache: # kept for BC (cache positions) - if not isinstance(past_key_values, StaticCache): - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - past_seen_tokens = past_key_values.get_seq_length() - if cache_position is None: - if isinstance(past_key_values, StaticCache): - raise ValueError("cache_position is a required argument when using StaticCache.") - cache_position = torch.arange(past_seen_tokens, past_seen_tokens + seq_length, device=device) - - seq_length_with_past = seq_length + past_seen_tokens - - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - if use_cache: - logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") - use_cache = False - - if position_ids is None: - position_ids = cache_position.unsqueeze(0) - - no_split_input = disable_pp or not stage_manager.is_first_stage() - if no_split_input and sp_mode == "ring_attn": - _, attn_kwargs, _ = RingAttention.prepare_varlen_batch(attention_mask, sp_group) - elif shard_config.enable_flash_attention: - mask_shape = (batch_size, 1, seq_length_with_past, seq_length_with_past) - attn_kwargs: dict = ColoAttention.prepare_attn_kwargs( - mask_shape, - hidden_states.dtype, - hidden_states.device, - q_padding_mask=attention_mask, - is_causal=True, - invert=(sp_mode != "ring_attn"), - ) - else: - # Later pipeline stages hold an SP slice, whereas Q/K are gathered - # before attention. Build the mask from the full sequence shape, - # not the local hidden_states length used by _update_causal_mask. - attn_kwargs: torch.Tensor = _prepare_4d_causal_attention_mask( - attention_mask, (batch_size, seq_length), hidden_states, past_seen_tokens - ) - - # Support SP + PP. Later stages have already received the split input. - split_input = disable_pp or stage_manager.is_first_stage() - if split_input: - # Ring Attention zigzag batch processing - if sp_mode == "ring_attn": - assert shard_config.enable_flash_attention, "Ring Attention inherently requires Flash Attention." - if not attention_mask.bool().all(): - hidden_states, attn_kwargs, position_ids = RingAttention.prepare_varlen_batch( - attention_mask, sp_group, hidden_states, position_ids - ) - else: - hidden_states, position_ids = split_batch_zigzag([hidden_states, position_ids], sp_group) - - elif is_share_sp_tp(sp_mode): - hidden_states = split_forward_gather_backward( - hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication - ) - elif sp_mode == "all_to_all": - hidden_states = split_forward_gather_backward( - hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication - ) - - if self.gradient_checkpointing and self.training and use_cache: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - start_idx, end_idx = (0, len(self.layers)) if disable_pp else (stage_index[0], stage_index[1]) - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - num_ckpt_layers = 0 - if self.gradient_checkpointing and self.training: - num_ckpt_layers = end_idx - start_idx - # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer - if shard_config.gradient_checkpoint_config is not None: - num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( - stage=stage_manager.stage, - num_stages=stage_manager.num_stages, - num_layers=end_idx - start_idx, - model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), - num_model_chunks=stage_manager.num_model_chunks, - ) - assert num_ckpt_layers <= end_idx - start_idx - for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): - if output_hidden_states: - all_hidden_states += (hidden_states,) - if idx - start_idx < num_ckpt_layers: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attn_kwargs, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attn_kwargs, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - ) - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache = layer_outputs[2 if output_attentions else 1] - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if disable_pp or stage_manager.is_last_stage(): - hidden_states = self.norm(hidden_states) - if (not shard_config.parallel_output) or force_sp_gather or is_share_sp_tp(sp_mode): # noqa - hidden_states = gather_sp_output(hidden_states, shard_config) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = next_decoder_cache if use_cache else None - if disable_pp or stage_manager.is_last_stage(): - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - next_cache, - all_hidden_states, - all_self_attns, - ] - if v is not None - ) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - # always return dict for intermediate stage - return {"hidden_states": hidden_states} - - @staticmethod - def llama_for_causal_lm_forward( - self: LlamaForCausalLM, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - **kwargs, - ): - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, LlamaForCausalLM - - >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - logger = logging.get_logger(__name__) - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - if shard_config.sequence_parallelism_mode == "ring_attn" and shard_config.parallel_output: - # Split labels in a zigzag fashion too - sp_group = shard_config.sequence_parallel_process_group - if attention_mask.bool().all(): - labels = split_batch_zigzag(labels, sp_group, seq_dim=1, is_label=True) - else: - # [B, max_seqlen // sp_size] - labels, _, _ = RingAttention.prepare_varlen_batch(attention_mask, sp_group, labels, is_label=True) - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = LlamaPipelineForwards.llama_model_forward( - self.model, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - cache_position=cache_position, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - shard_config=shard_config, - force_sp_gather=False, - ) - past_key_values = None - - disable_pp = stage_manager is None - if disable_pp or stage_manager.is_last_stage(): - hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - loss = None - if labels is not None: - loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, self.model.dtype) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - @staticmethod - def llama_for_sequence_classification_forward( - self: LlamaForSequenceClassification, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If - `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - transformer_outputs = LlamaPipelineForwards.llama_model_forward( - self.model, - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - cache_position=cache_position, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - shard_config=shard_config, - ) - - if input_ids is not None: - batch_size = input_ids.shape[0] - elif inputs_embeds is not None: - batch_size = inputs_embeds.shape[0] - else: - batch_size = hidden_states.shape[0] - - if stage_manager.is_last_stage(): - hidden_states = transformer_outputs[0] - logits = self.score(hidden_states) - - if self.config.pad_token_id is None and batch_size != 1: - raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") - if self.config.pad_token_id is None: - sequence_lengths = -1 - else: - if input_ids is not None: - sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) - else: - sequence_lengths = -1 - - pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.num_labels == 1: - loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(pooled_logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(pooled_logits, labels) - if not return_dict: - output = (pooled_logits,) + transformer_outputs[1:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutputWithPast( - loss=loss, - logits=pooled_logits, - past_key_values=transformer_outputs.past_key_values, - hidden_states=transformer_outputs.hidden_states, - attentions=transformer_outputs.attentions, - ) - - else: - hidden_states = transformer_outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - -def get_llama_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], - attention_mask: Optional[Union[torch.Tensor, Dict]] = None, - past_key_value: Optional[Cache] = None, - output_attentions: bool = False, - use_cache: bool = False, - cache_position: Optional[torch.LongTensor] = None, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: - if sp_mode is not None: - assert sp_mode in _SUPPORTED_SP_MODE, f"SP mode {sp_mode} is not supported by {type(self)} yet" - assert (sp_size is not None) and ( - sp_group is not None - ), "Must specify sp_size and sp_group for sequence parallel" - if "padding_mask" in kwargs: - warnings.warn( - "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" - ) - - bsz, q_len, _ = hidden_states.size() - input_shape = hidden_states.shape[:-1] - # sp: modify sp_len when sequence parallel mode is ring - if is_share_sp_tp(sp_mode): - q_len *= sp_size - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - # sp: all-to-all comminucation when introducing sequence parallel - if sp_mode == "all_to_all": - query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) - key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) - value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) - bsz, q_len, _ = query_states.size() - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - - kv_seq_len = key_states.shape[-2] - if past_key_value is not None: - if self.layer_idx is None: - raise ValueError( - f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " - "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " - "with a layer index." - ) - - kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) - - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - - if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) - - # repeat k/v heads if n_kv_heads < n_heads - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) - - if sp_mode == "ring_attn": - attn_output = RingAttention.attention( - query_states, - key_states, - value_states, - sp_axis=shard_config.sp_axis, - **attention_mask, - inner_ring_size=shard_config.inner_ring_size, - pg_mesh=shard_config.pg_mesh, - ) - - elif shard_config.enable_flash_attention: - assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict." - attn_output = ColoAttention.attention(query_states, key_states, value_states, **attention_mask) - else: - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): - raise ValueError( - f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) - - if attention_mask is not None: - # Modern Transformers may reserve one extra cache position in - # its causal mask; native attention slices it to actual keys. - attention_mask = attention_mask[:, :, :, :kv_seq_len] - if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights + attention_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_output = torch.matmul(attn_weights, value_states) - - if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" - f" {attn_output.size()}" - ) - - attn_output = attn_output.transpose(1, 2).contiguous() - # sp: all-to-all comminucation when introducing sequence parallel - if sp_mode == "all_to_all": - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) - attn_output = all_to_all_comm( - attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication - ) - else: - attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() - - attn_output = self.o_proj(attn_output) - - if not output_attentions: - attn_weights = None - return attn_output, attn_weights - - return forward +reexport_transformers_impl(globals(), __package__, "_llama") diff --git a/colossalai/shardformer/modeling/mixtral.py b/colossalai/shardformer/modeling/mixtral.py index 3d1d01ead731..cbd4896d7821 100644 --- a/colossalai/shardformer/modeling/mixtral.py +++ b/colossalai/shardformer/modeling/mixtral.py @@ -1,866 +1,9 @@ -import warnings -from typing import Callable, List, Optional, Tuple, Union +r"""Mixtral modeling, selected by the installed transformers major version. -import torch -import torch.distributed as dist -import torch.nn.functional as F -from torch.distributed import ProcessGroup -from torch.nn import CrossEntropyLoss -from transformers.cache_utils import Cache, DynamicCache -from transformers.modeling_attn_mask_utils import ( - _prepare_4d_causal_attention_mask, - _prepare_4d_causal_attention_mask_for_sdpa, -) -from transformers.models.mixtral.modeling_mixtral import ( - MixtralModel, - MixtralSparseMoeBlock, - MoeCausalLMOutputWithPast, - MoeModelOutputWithPast, - apply_rotary_pos_emb, - load_balancing_loss_func, -) -from transformers.utils import is_flash_attn_2_available, logging +``_mixtral_tf5`` targets transformers v5; ``_mixtral_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -from colossalai.lazy import LazyInitContext -from colossalai.moe._operation import EPGradScalerIn, EPGradScalerOut, all_to_all_uneven -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer._operation import ( - all_to_all_comm, - gather_forward_split_backward, - split_forward_gather_backward, -) -from colossalai.shardformer.layer.linear import Linear1D_Col, Linear1D_Row, ParallelModule -from colossalai.shardformer.shard import ShardConfig -from colossalai.shardformer.shard.utils import set_tensors_to_none -from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_group +from colossalai._compat import reexport_transformers_impl - -class EPMixtralSparseMoeBlock(ParallelModule): - def __init__(self, *args, **kwargs): - raise RuntimeError(f"Please use `from_native_module` to create an instance of {self.__class__.__name__}") - - def setup_process_groups( - self, - tp_group: ProcessGroup, - moe_dp_group: ProcessGroup, - ep_group: ProcessGroup, - fp8_communication: bool = False, - use_zbv: bool = False, - ): - assert tp_group is not None - assert moe_dp_group is not None - assert ep_group is not None - - # setup ep group - self.ep_size = dist.get_world_size(ep_group) - self.ep_rank = dist.get_rank(ep_group) - self.ep_group = ep_group - self.fp8_communication = fp8_communication - self.use_zbv = use_zbv - - if self.num_experts % self.ep_size != 0: - raise ValueError("The number of experts must be divisible by the number of expert parallel groups.") - - self.num_experts_per_ep = self.num_experts // self.ep_size - self.expert_start_idx = self.ep_rank * self.num_experts_per_ep - held_experts = self.experts[self.expert_start_idx : self.expert_start_idx + self.num_experts_per_ep] - - set_tensors_to_none(self.experts, exclude=set(held_experts)) - - # setup moe_dp group - self.moe_dp_group = moe_dp_group - self.moe_dp_size = moe_dp_group.size() - - # setup global tp group - self.tp_group = tp_group - if self.tp_group.size() > 1: - for expert in held_experts: - expert.w1 = Linear1D_Col.from_native_module( - expert.w1, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv - ) - expert.w3 = Linear1D_Col.from_native_module( - expert.w3, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv - ) - expert.w2 = Linear1D_Row.from_native_module( - expert.w2, self.tp_group, fp8_communication=self.fp8_communication, use_zbv=self.use_zbv - ) - - for p in self.experts.parameters(): - set_moe_tensor_ep_group(p, ep_group) - - @staticmethod - def from_native_module( - module: MixtralSparseMoeBlock, - tp_group: ProcessGroup, - moe_dp_group: ProcessGroup, - ep_group: ProcessGroup, - *args, - **kwargs, - ) -> "EPMixtralSparseMoeBlock": - # TODO: better init - LazyInitContext.materialize(module) - module.__class__ = EPMixtralSparseMoeBlock - fp8_communication = kwargs.get("fp8_communication", False) - module.setup_process_groups(tp_group, moe_dp_group, ep_group, fp8_communication) - return module - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - batch_size, sequence_length, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) - # router_logits: (batch * sequence_length, n_experts) - router_logits = self.gate(hidden_states) - - routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float) - routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) - routing_weights /= routing_weights.sum(dim=-1, keepdim=True) - # we cast back to the input dtype - routing_weights = routing_weights.to(hidden_states.dtype) - - selected_experts = selected_experts.t().reshape(-1) - selected_experts_idx = selected_experts.argsort() - dispatch_states = hidden_states.repeat(self.top_k, 1)[selected_experts_idx] - input_split_sizes = selected_experts.bincount(minlength=self.num_experts) - - output_split_sizes = torch.zeros_like(input_split_sizes) - - dist.all_to_all_single(output_split_sizes, input_split_sizes, group=self.ep_group) - - input_split_list = input_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() - output_split_list = output_split_sizes.view(self.ep_size, self.num_experts_per_ep).sum(dim=-1).tolist() - - output_states, _ = all_to_all_uneven( - dispatch_states, - input_split_list, - output_split_list, - self.ep_group, - fp8_communication=self.fp8_communication, - ) - # compute expert output - output_states = EPGradScalerIn.apply(output_states, self.ep_size) - if output_states.size(0) > 0: - if self.num_experts_per_ep == 1: - # no need to split - expert = self.experts[self.expert_start_idx] - output_states = expert.act_fn(expert.w1(output_states)) * expert.w3(output_states) - output_states = expert.w2(output_states) - else: - output_states_splits = output_states.split(output_split_sizes.tolist()) - output_states_list = [] - for i, split_states in enumerate(output_states_splits): - if split_states.size(0) == 0: - continue - expert = self.experts[self.expert_start_idx + i % self.num_experts_per_ep] - split_states = expert.act_fn(expert.w1(split_states)) * expert.w3(split_states) - split_states = expert.w2(split_states) - output_states_list.append(split_states) - output_states = torch.cat(output_states_list) - - output_states = EPGradScalerOut.apply(output_states, self.ep_size) - dispatch_states, _ = all_to_all_uneven( - output_states, output_split_list, input_split_list, self.ep_group, fp8_communication=self.fp8_communication - ) - - recover_experts_idx = torch.empty_like(selected_experts_idx) - recover_experts_idx[selected_experts_idx] = torch.arange( - selected_experts_idx.size(0), device=selected_experts_idx.device - ) - dispatch_states = dispatch_states[recover_experts_idx] - k_hidden_states = dispatch_states.chunk(self.top_k) - output_states = k_hidden_states[0] * routing_weights[:, 0, None] - for i in range(1, self.top_k): - output_states += k_hidden_states[i] * routing_weights[:, i, None] - output_states = output_states.reshape(batch_size, sequence_length, hidden_dim) - return output_states, router_logits - - -class MixtralPipelineForwards: - """ - This class serves as a micro library for forward function substitution of Mixtral models - under pipeline setting. - """ - - @staticmethod - def mixtral_model_forward( - self: MixtralModel, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_router_logits: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - past_router_logits: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MixtralForCausalLM - - >>> model = MixtralForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - logger = logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if stage_manager.is_first_stage(): - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape - elif inputs_embeds is not None: - batch_size, seq_length, _ = inputs_embeds.shape - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - device = input_ids.device if input_ids is not None else inputs_embeds.device - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - hidden_states = inputs_embeds - else: - input_shape = hidden_states.shape[:-1] - batch_size, seq_length = input_shape - device = hidden_states.device - - seq_length_with_past = seq_length - past_key_values_length = 0 - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - if use_cache: - logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") - use_cache = False - - if past_key_values is not None: - past_key_values_length = past_key_values[0][0].shape[2] - seq_length_with_past = seq_length_with_past + past_key_values_length - - if position_ids is None: - position_ids = torch.arange( - past_key_values_length, - seq_length + past_key_values_length, - dtype=torch.long, - device=device, - ) - position_ids = position_ids.unsqueeze(0).view(-1, seq_length) - else: - position_ids = position_ids.view(-1, seq_length).long() - - # embed positions, for the first stage, hidden_states is the input embeddings, - # for the other stages, hidden_states is the output of the previous stage - if is_flash_attn_2_available(): - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, - (batch_size, seq_length), - hidden_states, - past_key_values_length, - sliding_window=self.config.sliding_window, - ) - - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_router_logits = () if output_router_logits else None - if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - cache_position = torch.arange( - past_seen_tokens, past_seen_tokens + hidden_states.shape[1], device=hidden_states.device - ) - - start_idx, end_idx = stage_index[0], stage_index[1] - for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - past_key_value = past_key_values[idx] if past_key_values is not None else None - - if self.gradient_checkpointing and self.training: - - def create_custom_forward(module): - def custom_forward(*inputs): - # None for past_key_value - return module(*inputs) - - return custom_forward - - layer_outputs = torch.utils.checkpoint.checkpoint( - create_custom_forward(decoder_layer), - hidden_states, - attention_mask, - position_ids, - None, - output_attentions, - output_router_logits, - use_cache, - cache_position, - position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask, - position_ids, - past_key_value, - output_attentions, - output_router_logits, - use_cache, - cache_position, - position_embeddings, - ) - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache = (layer_outputs[2 if output_attentions else 1],) - if output_attentions: - all_self_attns += (layer_outputs[1],) - if output_router_logits: - all_router_logits += (layer_outputs[-1],) - - if stage_manager.is_last_stage(): - hidden_states = self.norm(hidden_states) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = next_decoder_cache if use_cache else None - - if output_router_logits and past_router_logits is not None: - all_router_logits = past_router_logits + all_router_logits - - if stage_manager.is_last_stage(): - if not return_dict: - return tuple( - v - for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] - if v is not None - ) - return MoeModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - router_logits=all_router_logits, - ) - else: - if output_router_logits: - return { - "hidden_states": hidden_states, - "past_router_logits": all_router_logits, - } - else: - return { - "hidden_states": hidden_states, - } - - @staticmethod - def mixtral_for_causal_lm_forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_router_logits: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - past_router_logits: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, MixtralForCausalLM - - >>> model = MixtralForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - logger = logging.get_logger(__name__) - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = MixtralPipelineForwards.mixtral_model_forward( - self.model, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - output_router_logits=output_router_logits, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - past_router_logits=past_router_logits, - ) - past_key_values = None - - if stage_manager.is_last_stage(): - hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - logits = logits.float() - loss = None - if labels is not None: - # Shift so that tokens < n predict n - shift_logits = logits[..., :-1, :].contiguous() - shift_labels = labels[..., 1:].contiguous() - # Flatten the tokens - loss_fct = CrossEntropyLoss() - shift_logits = shift_logits.view(-1, self.config.vocab_size) - shift_labels = shift_labels.view(-1) - # Enable model parallelism - shift_labels = shift_labels.to(shift_logits.device) - loss = loss_fct(shift_logits, shift_labels) - - aux_loss = None - if output_router_logits: - aux_loss = load_balancing_loss_func(outputs[-1], self.num_experts, self.num_experts_per_tok) - if labels is not None: - loss += self.router_aux_loss_coef * aux_loss - - if not return_dict: - output = (logits,) + outputs[1:] - if output_router_logits: - output = (aux_loss,) + output - return (loss,) + output if loss is not None else output - - return MoeCausalLMOutputWithPast( - loss=loss, - aux_loss=aux_loss, - logits=logits, - past_key_values=None, - hidden_states=outputs[0], - attentions=None, - router_logits=outputs[-1], - ) - else: - out = {} - hidden_states = outputs.get("hidden_states") - out["hidden_states"] = hidden_states - if output_router_logits: - out["past_router_logits"] = outputs["past_router_logits"] - return out - - -def get_mixtral_flash_attention_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None): - logger = logging.get_logger(__name__) - from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS - from transformers.models.mixtral.modeling_mixtral import eager_attention_forward - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_value: Optional[Cache] = None, - output_attentions: bool = False, - cache_position: Optional[torch.LongTensor] = None, - use_cache: bool = False, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Cache]]: - if sp_mode is not None: - assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" - assert (sp_size is not None) and ( - sp_group is not None - ), "Must specify sp_size and sp_group for sequence parallel" - - if "padding_mask" in kwargs: - warnings.warn( - "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`" - ) - - # overwrite attention_mask with padding_mask - attention_mask = kwargs.pop("padding_mask") - bsz, q_len, _ = hidden_states.size() - - # sp: modify sp_len when sequence parallel mode is ring - if sp_mode in ["split_gather", "ring"]: - q_len *= sp_size - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - - # sp: all-to-all comminucation when introducing sequence parallel - if sp_mode == "all_to_all": - query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) - key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) - value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) - bsz, q_len, _ = query_states.size() - - query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) - kv_seq_len = key_states.shape[-2] - if past_key_value is not None: - if self.layer_idx is None: - raise ValueError( - f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " - "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " - "with a layer index." - ) - kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) - - # Because the input can be padded, the absolute sequence length depends on the max position id. - cos, sin = position_embeddings - - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - - if past_key_value is not None: - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) - - # Transformers attention interfaces handle grouped-query heads and - # receive tensors in [batch, heads, sequence, head_dim] order. - - # In PEFT, usually we cast the layer norms in float32 for training stability reasons - # therefore the input hidden states gets silently casted in float32. Hence, we need - # cast them back in float16 just to be sure everything works as expected. - input_dtype = query_states.dtype - if input_dtype == torch.float32: - if torch.is_autocast_enabled(): - target_dtype = torch.get_autocast_gpu_dtype() - # Handle the case where the model is quantized - elif hasattr(self.config, "_pre_quantization_dtype"): - target_dtype = self.config._pre_quantization_dtype - else: - target_dtype = self.q_proj.weight.dtype - - logger.warning_once( - f"The input hidden states seems to be silently casted in float32, this might be related to" - f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in" - f" {target_dtype}." - ) - - query_states = query_states.to(target_dtype) - key_states = key_states.to(target_dtype) - value_states = value_states.to(target_dtype) - attention_interface: Callable = eager_attention_forward - if self.config._attn_implementation != "eager": - if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): - logger.warning_once( - "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " - 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' - ) - else: - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - - attn_output, attn_weights = attention_interface( - self, - query_states, - key_states, - value_states, - attention_mask, - dropout=0.0 if not self.training else self.attention_dropout, - scaling=self.scaling, - sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama - **kwargs, - ) - - # sp: all-to-all comminucation when introducing sequence parallel - if sp_mode == "all_to_all": - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim).contiguous() # (1, 8, 128) - attn_output = all_to_all_comm( - attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication - ) # (1, 4, 256) - else: - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) - - attn_output = self.o_proj(attn_output) - - if not output_attentions: - attn_weights = None - return attn_output, attn_weights - - return forward - - -def get_mixtral_flash_attention_model_forward(shard_config, sp_mode=None, sp_size=None, sp_group=None): - logger = logging.get_logger(__name__) - - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - output_router_logits: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, MoeModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_router_logits = ( - output_router_logits if output_router_logits is not None else self.config.output_router_logits - ) - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape - elif inputs_embeds is not None: - batch_size, seq_length, _ = inputs_embeds.shape - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - past_key_values_length = 0 - - if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - if use_cache: - use_legacy_cache = not isinstance(past_key_values, Cache) - if use_legacy_cache: - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - past_key_values_length = past_key_values.get_usable_length(seq_length) - - if position_ids is None: - device = input_ids.device if input_ids is not None else inputs_embeds.device - position_ids = torch.arange( - past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device - ) - position_ids = position_ids.unsqueeze(0).view(-1, seq_length) - else: - position_ids = position_ids.view(-1, seq_length).long() - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - if attention_mask is not None and self.config._attn_implementation == "flash_attention_2" and use_cache: - is_padding_right = attention_mask[:, -1].sum().item() != batch_size - if is_padding_right: - raise ValueError( - "You are attempting to perform batched generation with padding_side='right'" - " this may lead to unexpected behaviour for Flash Attention version of Mixtral. Make sure to " - " call `tokenizer.padding_side = 'left'` before tokenizing the input. " - ) - if self.config._attn_implementation == "flash_attention_2": - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - elif self.config._attn_implementation == "sdpa" and not output_attentions: - # output_attentions=True can not be supported when using SDPA, and we fall back on - # the manual implementation that requires a 4D causal mask in all cases. - attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - (batch_size, seq_length), - inputs_embeds, - past_key_values_length, - ) - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, - (batch_size, seq_length), - inputs_embeds, - past_key_values_length, - sliding_window=self.config.sliding_window, - ) - - if sp_mode in ["ring", "split_gather"]: - inputs_embeds = split_forward_gather_backward( - inputs_embeds, 1, sp_group, fp8_communication=shard_config.fp8_communication - ) - elif sp_mode == "all_to_all": - inputs_embeds = split_forward_gather_backward( - inputs_embeds, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication - ) - hidden_states = inputs_embeds - - if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - cache_position = torch.arange( - past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device - ) - if position_ids is None: - position_ids = cache_position.unsqueeze(0) - - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_router_logits = () if output_router_logits else None - next_decoder_cache = None - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - output_router_logits, - use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - output_router_logits=output_router_logits, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - ) - - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache = layer_outputs[2 if output_attentions else 1] - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if output_router_logits: - all_router_logits += (layer_outputs[-1],) - - hidden_states = self.norm(hidden_states) - - if sp_mode == "ring" or sp_mode == "split_gather": - hidden_states = gather_forward_split_backward( - hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication - ) - elif sp_mode == "all_to_all": - hidden_states = gather_forward_split_backward( - hidden_states, 1, sp_group, grad_scale=sp_size, fp8_communication=shard_config.fp8_communication - ) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = None - if use_cache: - next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache - - if not return_dict: - return tuple( - v - for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, all_router_logits] - if v is not None - ) - return MoeModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - router_logits=all_router_logits, - ) - - return forward +reexport_transformers_impl(globals(), __package__, "_mixtral") diff --git a/colossalai/shardformer/modeling/qwen2.py b/colossalai/shardformer/modeling/qwen2.py index 80dfb003dccf..258e7c565613 100644 --- a/colossalai/shardformer/modeling/qwen2.py +++ b/colossalai/shardformer/modeling/qwen2.py @@ -1,841 +1,3 @@ -import math -from typing import List, Optional, Tuple, Union +from colossalai._compat import reexport_transformers_impl -import torch -from torch import nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from transformers.modeling_attn_mask_utils import ( - _prepare_4d_causal_attention_mask, - _prepare_4d_causal_attention_mask_for_sdpa, -) -from transformers.modeling_outputs import ( - BaseModelOutputWithPast, - CausalLMOutputWithPast, - SequenceClassifierOutputWithPast, -) -from transformers.models.qwen2.modeling_qwen2 import ( - Qwen2Attention, - Qwen2ForCausalLM, - Qwen2ForSequenceClassification, - Qwen2Model, - apply_rotary_pos_emb, - repeat_kv, -) -from transformers.utils import logging - -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer._operation import all_to_all_comm, split_forward_gather_backward -from colossalai.shardformer.shard import ShardConfig - -from ..layer import ColoAttention, dist_cross_entropy -from ..layer._operation import gather_sp_output -from ..layer.utils import is_share_sp_tp - - -class Qwen2PipelineForwards: - """ - This class serves as a micro library for forward function substitution of Qwen2 models - under pipeline setting. - """ - - @staticmethod - def qwen2_model_forward( - self: Qwen2Model, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - force_sp_output_gather: bool = True, - ) -> Union[Tuple, BaseModelOutputWithPast]: - logger = logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if stage_manager.is_first_stage(): - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape - elif inputs_embeds is not None: - batch_size, seq_length, _ = inputs_embeds.shape - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - device = input_ids.device if input_ids is not None else inputs_embeds.device - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - hidden_states = inputs_embeds - else: - input_shape = hidden_states.shape[:-1] - batch_size, seq_length = input_shape - device = hidden_states.device - - seq_length_with_past = seq_length - past_key_values_length = 0 - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - if use_cache: - logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") - use_cache = False - - # assert past_key_values is None, "past_key_values is not supported for Qwen2 models at the moment." - - if past_key_values is not None: - past_key_values_length = past_key_values[0][0].shape[2] - seq_length_with_past = seq_length_with_past + past_key_values_length - - # Support SP + PP - sp_size = shard_config.sequence_parallel_size - sp_group = shard_config.sequence_parallel_process_group - sp_mode = shard_config.sequence_parallelism_mode - # For generating full positions ids (the states will be gathered along the seq dim before attention fwd). - if shard_config.enable_sequence_parallelism and sp_mode != "ring_attn" and not stage_manager.is_first_stage(): - seq_length *= sp_size - # Incoming pipeline states contain only the local SP slice. Both mask - # dimensions must describe the full sequence used by attention. - seq_length_with_past = seq_length + past_key_values_length - - if position_ids is None: - position_ids = torch.arange( - past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device - ) - position_ids = position_ids.unsqueeze(0).view(-1, seq_length) - else: - position_ids = position_ids.view(-1, seq_length).long() - - # embed positions, for the first stage, hidden_states is the input embeddings, - # for the other stages, hidden_states is the output of the previous stage - if shard_config.enable_flash_attention: - # in this case, attention_mask is a dict rather than a tensor - mask_shape = (batch_size, 1, seq_length, seq_length_with_past) - attention_mask = ColoAttention.prepare_attn_kwargs( - mask_shape, - hidden_states.dtype, - hidden_states.device, - q_padding_mask=attention_mask, - is_causal=True, - ) - else: - if self.config._attn_implementation == "flash_attention_2": - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - elif self.config._attn_implementation == "sdpa" and not output_attentions: - # output_attentions=True can not be supported when using SDPA, and we fall back on - # the manual implementation that requires a 4D causal mask in all cases. - attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - (batch_size, seq_length), - hidden_states, - past_key_values_length, - ) - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, - (batch_size, seq_length), - hidden_states, - past_key_values_length, - sliding_window=self.config.sliding_window, - ) - - if stage_manager.is_first_stage(): - if shard_config.enable_sequence_parallelism: - if is_share_sp_tp(sp_mode): - hidden_states = split_forward_gather_backward( - hidden_states, - dim=1, - process_group=sp_group, - ) - elif sp_mode == "all_to_all": - hidden_states = split_forward_gather_backward( - hidden_states, - dim=1, - process_group=sp_group, - grad_scale=1 / sp_size, - ) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - start_idx, end_idx = stage_index[0], stage_index[1] - num_ckpt_layers = 0 - if self.gradient_checkpointing and self.training: - num_ckpt_layers = end_idx - start_idx - # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer - if shard_config.gradient_checkpoint_config is not None: - num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( - stage=stage_manager.stage, - num_stages=stage_manager.num_stages, - num_layers=end_idx - start_idx, - model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), - num_model_chunks=stage_manager.num_model_chunks, - ) - assert num_ckpt_layers <= end_idx - start_idx - - for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - past_key_values[idx] if past_key_values is not None else None - - if idx - start_idx < num_ckpt_layers: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if stage_manager.is_last_stage(): - hidden_states = self.norm(hidden_states) - if shard_config.enable_sequence_parallelism: - if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): - hidden_states = gather_sp_output(hidden_states, shard_config) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - - if stage_manager.is_last_stage(): - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - # always return dict for imediate stage - return {"hidden_states": hidden_states} - - @staticmethod - def qwen2_for_causal_lm_forward( - self: Qwen2ForCausalLM, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - **kwargs, - ): - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, Qwen2ForCausalLM - - >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - - >>> prompt = "Hey, are you consciours? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you." - ```""" - logger = logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = Qwen2PipelineForwards.qwen2_model_forward( - self.model, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - shard_config=shard_config, - force_sp_output_gather=False, - ) - past_key_values = None - - if stage_manager.is_last_stage(): - hidden_states = outputs[0] - if hidden_states.shape[1] == 2: - pass - logits = self.lm_head(hidden_states) - loss = None - if labels is not None: - loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - @staticmethod - def qwen2_for_sequence_classification_forward( - self: Qwen2ForSequenceClassification, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If - `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - transformer_outputs = Qwen2PipelineForwards.qwen2_model_forward( - self.model, - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - shard_config=shard_config, - ) - - if input_ids is not None: - batch_size = input_ids.shape[0] - elif inputs_embeds is not None: - batch_size = inputs_embeds.shape[0] - else: - batch_size = hidden_states.shape[0] - - if stage_manager.is_last_stage(): - hidden_states = transformer_outputs[0] - logits = self.score(hidden_states) - - if self.config.pad_token_id is None and batch_size != 1: - raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") - if self.config.pad_token_id is None: - sequence_lengths = -1 - else: - if input_ids is not None: - sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device) - else: - sequence_lengths = -1 - - pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] - - loss = None - if labels is not None: - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.num_labels == 1: - loss = loss_fct(pooled_logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(pooled_logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(pooled_logits, labels) - if not return_dict: - output = (pooled_logits,) + transformer_outputs[1:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutputWithPast( - loss=loss, - logits=pooled_logits, - past_key_values=transformer_outputs.past_key_values, - hidden_states=transformer_outputs.hidden_states, - attentions=transformer_outputs.attentions, - ) - - else: - hidden_states = transformer_outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - -def get_qwen2_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): - def forward( - self: Qwen2Attention, - hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], - attention_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - cache_position: Optional[torch.LongTensor] = None, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - if sp_mode is not None: - assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" - assert (sp_size is not None) and ( - sp_group is not None - ), "Must specify sp_size and sp_group for sequence parallel" - - bsz, q_len, _ = hidden_states.size() - # sp: modify sp_len when sequence parallel mode is ring - if sp_mode in ["split_gather", "ring"]: - q_len *= sp_size - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - # sp: all-to-all comminucation when introducing sequence parallel - if sp_mode == "all_to_all": - query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) - key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) - value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) - bsz, q_len, _ = query_states.size() - - query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - - kv_seq_len = key_states.shape[-2] - if past_key_value is not None: - if self.layer_idx is None: - raise ValueError( - f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " - "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " - "with a layer index." - ) - kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) - # Because the input can be padded, the absolute sequence length depends on the max position id. - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - - if past_key_value is not None: - # Activate slicing cache only if the config has a value `sliding_windows` attribute - cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 - if ( - getattr(self.config, "sliding_window", None) is not None - and kv_seq_len > self.config.sliding_window - and cache_has_contents - ): - slicing_tokens = 1 - self.config.sliding_window - - past_key = past_key_value[self.layer_idx][0] - past_value = past_key_value[self.layer_idx][1] - - past_key = past_key[:, :, slicing_tokens:, :].contiguous() - past_value = past_value[:, :, slicing_tokens:, :].contiguous() - - if past_key.shape[-2] != self.config.sliding_window - 1: - raise ValueError( - f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" - f" {past_key.shape}" - ) - - if attention_mask is not None: - attention_mask = attention_mask[:, slicing_tokens:] - attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) - - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) - - # repeat k/v heads if n_kv_heads < n_heads - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) - - if shard_config.enable_flash_attention: - assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict." - attn_output = ColoAttention.attention(query_states, key_states, value_states, **attention_mask) - else: - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - - if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): - raise ValueError( - f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) - - if attention_mask is not None: - if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights + attention_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_output = torch.matmul(attn_weights, value_states) - - if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" - f" {attn_output.size()}" - ) - attn_output = attn_output.transpose(1, 2).contiguous() - if sp_mode == "all_to_all": - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) - attn_output = all_to_all_comm( - attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication - ) - else: - attn_output = attn_output.reshape(bsz, q_len, -1) - - attn_output = self.o_proj(attn_output) - - return attn_output, None - - return forward - - -def get_qwen2_model_forward_for_flash_attn(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): - logger = logging.get_logger(__name__) - - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - return_dict: Optional[bool] = None, - force_sp_output_gather: bool = True, - ) -> Union[Tuple, BaseModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape - elif inputs_embeds is not None: - batch_size, seq_length, _ = inputs_embeds.shape - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - seq_length_with_past = seq_length - past_key_values_length = 0 - - if past_key_values is not None: - past_key_values_length = past_key_values[0][0].shape[2] - seq_length_with_past = seq_length_with_past + past_key_values_length - - if position_ids is None: - device = input_ids.device if input_ids is not None else inputs_embeds.device - position_ids = torch.arange( - past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device - ) - position_ids = position_ids.unsqueeze(0).view(-1, seq_length) - else: - position_ids = position_ids.view(-1, seq_length).long() - - # embed positions - hidden_states = inputs_embeds - - if shard_config.enable_flash_attention: - # in this case, attention_mask is a dict rather than a tensor - mask_shape = (batch_size, 1, seq_length, seq_length_with_past) - attention_mask = ColoAttention.prepare_attn_kwargs( - mask_shape, - hidden_states.dtype, - hidden_states.device, - q_padding_mask=attention_mask, - is_causal=True, - ) - else: - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, - (batch_size, seq_length), - inputs_embeds, - past_key_values_length, - sliding_window=self.config.sliding_window, - ) - - if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - if sp_mode in ["ring", "split_gather"]: - hidden_states = split_forward_gather_backward( - hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication - ) - elif sp_mode == "all_to_all": - hidden_states = split_forward_gather_backward( - hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication - ) - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - hidden_states = self.norm(hidden_states) - - if shard_config.enable_sequence_parallelism: - if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): - hidden_states = gather_sp_output(hidden_states, shard_config) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - - return forward - - -def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig): - def forward( - self: Qwen2ForCausalLM, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, CausalLMOutputWithPast]: - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, Qwen2ForCausalLM - - >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - force_sp_output_gather=False, - ) - - hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - logits = logits.float() - loss = None - if labels is not None: - loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - return forward +reexport_transformers_impl(globals(), __package__, "_qwen2") diff --git a/colossalai/shardformer/modeling/qwen3.py b/colossalai/shardformer/modeling/qwen3.py index 9b6670ba051b..19d46e525ca6 100644 --- a/colossalai/shardformer/modeling/qwen3.py +++ b/colossalai/shardformer/modeling/qwen3.py @@ -1,834 +1,3 @@ -# Modifed from qwen2 modeling -import math -from typing import List, Optional, Tuple, Union +from colossalai._compat import reexport_transformers_impl -import torch -from torch import nn -from transformers.modeling_attn_mask_utils import ( - _prepare_4d_causal_attention_mask, - _prepare_4d_causal_attention_mask_for_sdpa, -) -from transformers.modeling_outputs import ( - BaseModelOutputWithPast, - CausalLMOutputWithPast, - SequenceClassifierOutputWithPast, -) -from transformers.models.qwen3.modeling_qwen3 import ( - Qwen3Attention, - Qwen3ForCausalLM, - Qwen3ForSequenceClassification, - Qwen3Model, - apply_rotary_pos_emb, - repeat_kv, -) -from transformers.utils import logging - -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer._operation import all_to_all_comm, split_forward_gather_backward -from colossalai.shardformer.shard import ShardConfig - -from ..layer import ColoAttention, dist_cross_entropy -from ..layer._operation import gather_sp_output -from ..layer.utils import is_share_sp_tp - - -class Qwen3PipelineForwards: - """ - This class serves as a micro library for forward function substitution of Qwen3 models - under pipeline setting. - """ - - @staticmethod - def qwen3_model_forward( - self: Qwen3Model, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - force_sp_output_gather: bool = True, - ) -> Union[Tuple, BaseModelOutputWithPast]: - logger = logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if stage_manager.is_first_stage(): - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape - elif inputs_embeds is not None: - batch_size, seq_length, _ = inputs_embeds.shape - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - device = input_ids.device if input_ids is not None else inputs_embeds.device - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - hidden_states = inputs_embeds - else: - input_shape = hidden_states.shape[:-1] - batch_size, seq_length = input_shape - device = hidden_states.device - - seq_length_with_past = seq_length - past_key_values_length = 0 - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - if use_cache: - logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") - use_cache = False - - if past_key_values is not None: - past_key_values_length = past_key_values[0][0].shape[2] - seq_length_with_past = seq_length_with_past + past_key_values_length - - # Support SP + PP - sp_size = shard_config.sequence_parallel_size - sp_group = shard_config.sequence_parallel_process_group - sp_mode = shard_config.sequence_parallelism_mode - # For generating full positions ids (the states will be gathered along the seq dim before attention fwd). - if shard_config.enable_sequence_parallelism and sp_mode != "ring_attn" and not stage_manager.is_first_stage(): - seq_length *= sp_size - # Incoming pipeline states contain only the local SP slice. Both mask - # dimensions must describe the full sequence used by attention. - seq_length_with_past = seq_length + past_key_values_length - - if position_ids is None: - position_ids = torch.arange( - past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device - ) - position_ids = position_ids.unsqueeze(0).view(-1, seq_length) - else: - position_ids = position_ids.view(-1, seq_length).long() - - # embed positions, for the first stage, hidden_states is the input embeddings, - # for the other stages, hidden_states is the output of the previous stage - if shard_config.enable_flash_attention: - # in this case, attention_mask is a dict rather than a tensor - mask_shape = (batch_size, 1, seq_length, seq_length_with_past) - attention_mask = ColoAttention.prepare_attn_kwargs( - mask_shape, - hidden_states.dtype, - hidden_states.device, - q_padding_mask=attention_mask, - is_causal=True, - ) - else: - if self.config._attn_implementation == "flash_attention_2": - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - elif self.config._attn_implementation == "sdpa" and not output_attentions: - # output_attentions=True can not be supported when using SDPA, and we fall back on - # the manual implementation that requires a 4D causal mask in all cases. - attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - (batch_size, seq_length), - hidden_states, - past_key_values_length, - ) - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, - (batch_size, seq_length), - hidden_states, - past_key_values_length, - sliding_window=self.config.sliding_window, - ) - - if stage_manager.is_first_stage(): - if shard_config.enable_sequence_parallelism: - if is_share_sp_tp(sp_mode): - hidden_states = split_forward_gather_backward( - hidden_states, - dim=1, - process_group=sp_group, - ) - elif sp_mode == "all_to_all": - hidden_states = split_forward_gather_backward( - hidden_states, - dim=1, - process_group=sp_group, - grad_scale=1 / sp_size, - ) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - start_idx, end_idx = stage_index[0], stage_index[1] - num_ckpt_layers = 0 - if self.gradient_checkpointing and self.training: - num_ckpt_layers = end_idx - start_idx - # TODO: We can replace `gradient_checkpointing_enable` fn and initialize a gradient_checkpointing (List[bool]) for each layer - if shard_config.gradient_checkpoint_config is not None: - num_ckpt_layers = shard_config.gradient_checkpoint_config.get_num_ckpt_layers( - stage=stage_manager.stage, - num_stages=stage_manager.num_stages, - num_layers=end_idx - start_idx, - model_chunk_id=(stage_manager.model_chunk_id if stage_manager.is_interleave else 0), - num_model_chunks=stage_manager.num_model_chunks, - ) - assert num_ckpt_layers <= end_idx - start_idx - - for idx, decoder_layer in enumerate(self.layers[start_idx:end_idx], start=start_idx): - if output_hidden_states: - all_hidden_states += (hidden_states,) - - past_key_values[idx] if past_key_values is not None else None - - if idx - start_idx < num_ckpt_layers: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if stage_manager.is_last_stage(): - hidden_states = self.norm(hidden_states) - if shard_config.enable_sequence_parallelism: - if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): - hidden_states = gather_sp_output(hidden_states, shard_config) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - - if stage_manager.is_last_stage(): - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - # always return dict for imediate stage - return {"hidden_states": hidden_states} - - @staticmethod - def qwen3_for_causal_lm_forward( - self: Qwen3ForCausalLM, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - **kwargs, - ): - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, Qwen2ForCausalLM - - >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - - >>> prompt = "Hey, are you consciours? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you consciours? Can you talk to me?\nI'm not consciours, but I can talk to you." - ```""" - logger = logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = Qwen3PipelineForwards.qwen3_model_forward( - self.model, - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - shard_config=shard_config, - force_sp_output_gather=False, - ) - past_key_values = None - - if stage_manager.is_last_stage(): - hidden_states = outputs[0] - if hidden_states.shape[1] == 2: - pass - logits = self.lm_head(hidden_states) - loss = None - if labels is not None: - loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - else: - hidden_states = outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - @staticmethod - def qwen3_for_sequence_classification_forward( - self: Qwen3ForSequenceClassification, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - shard_config: ShardConfig = None, - ): - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If - `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - """ - logger = logging.get_logger(__name__) - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # TODO(jianghai): left the recording kv-value tensors as () or None type, this feature may be added in the future. - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - transformer_outputs = Qwen3PipelineForwards.qwen3_model_forward( - self.model, - input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - shard_config=shard_config, - ) - - if input_ids is not None: - batch_size = input_ids.shape[0] - elif inputs_embeds is not None: - batch_size = inputs_embeds.shape[0] - else: - batch_size = hidden_states.shape[0] - - if stage_manager.is_last_stage(): - hidden_states = transformer_outputs[0] - logits = self.score(hidden_states) - - if self.config.pad_token_id is None and batch_size != 1: - raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") - - if self.config.pad_token_id is None: - last_non_pad_token = -1 - elif input_ids is not None: - # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id - non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) - token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) - last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) - else: - last_non_pad_token = -1 - logger.warning_once( - f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " - "unexpected if using padding tokens in conjunction with `inputs_embeds.`" - ) - - pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] - - loss = None - if labels is not None: - loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) - - if not return_dict: - output = (pooled_logits,) + transformer_outputs[1:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutputWithPast( - loss=loss, - logits=pooled_logits, - past_key_values=transformer_outputs.past_key_values, - hidden_states=transformer_outputs.hidden_states, - attentions=transformer_outputs.attentions, - ) - - else: - hidden_states = transformer_outputs.get("hidden_states") - return {"hidden_states": hidden_states} - - -def get_qwen3_flash_attention_forward(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): - def forward( - self: Qwen3Attention, - hidden_states: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], - attention_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - cache_position: Optional[torch.LongTensor] = None, - **kwargs, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - if sp_mode is not None: - assert sp_mode in ["all_to_all", "split_gather", "ring"], "Invalid sp_mode" - assert (sp_size is not None) and ( - sp_group is not None - ), "Must specify sp_size and sp_group for sequence parallel" - - bsz, q_len, _ = hidden_states.size() - # sp: modify sp_len when sequence parallel mode is ring - if sp_mode in ["split_gather", "ring"]: - q_len *= sp_size - - query_states = self.q_proj(hidden_states) - key_states = self.k_proj(hidden_states) - value_states = self.v_proj(hidden_states) - # sp: all-to-all comminucation when introducing sequence parallel - if sp_mode == "all_to_all": - query_states = all_to_all_comm(query_states, sp_group, fp8_communication=shard_config.fp8_communication) - key_states = all_to_all_comm(key_states, sp_group, fp8_communication=shard_config.fp8_communication) - value_states = all_to_all_comm(value_states, sp_group, fp8_communication=shard_config.fp8_communication) - bsz, q_len, _ = query_states.size() - - query_states = self.q_norm(query_states.view(bsz, q_len, -1, self.head_dim)).transpose(1, 2) - key_states = self.k_norm(key_states.view(bsz, q_len, -1, self.head_dim)).transpose(1, 2) - value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - - kv_seq_len = key_states.shape[-2] - if past_key_value is not None: - if self.layer_idx is None: - raise ValueError( - f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} " - "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class " - "with a layer index." - ) - kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx) - # Because the input can be padded, the absolute sequence length depends on the max position id. - cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - - if past_key_value is not None: - # Activate slicing cache only if the config has a value `sliding_windows` attribute - cache_has_contents = past_key_value.get_seq_length(self.layer_idx) > 0 - if ( - getattr(self.config, "sliding_window", None) is not None - and kv_seq_len > self.config.sliding_window - and cache_has_contents - ): - slicing_tokens = 1 - self.config.sliding_window - - past_key = past_key_value[self.layer_idx][0] - past_value = past_key_value[self.layer_idx][1] - - past_key = past_key[:, :, slicing_tokens:, :].contiguous() - past_value = past_value[:, :, slicing_tokens:, :].contiguous() - - if past_key.shape[-2] != self.config.sliding_window - 1: - raise ValueError( - f"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1, head_dim`), got" - f" {past_key.shape}" - ) - - if attention_mask is not None: - attention_mask = attention_mask[:, slicing_tokens:] - attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) - - cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models - key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) - - # repeat k/v heads if n_kv_heads < n_heads - key_states = repeat_kv(key_states, self.num_key_value_groups) - value_states = repeat_kv(value_states, self.num_key_value_groups) - - if shard_config.enable_flash_attention: - assert isinstance(attention_mask, dict), "Flash Attention Error: attention_mask should be a dict." - attn_output = ColoAttention.attention( - query_states, - key_states, - value_states, - dropout_p=0.0 if not self.training else self.attention_dropout, - **attention_mask, - ) - else: - attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) - - if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): - raise ValueError( - f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" - f" {attn_weights.size()}" - ) - - if attention_mask is not None: - if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): - raise ValueError( - f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" - ) - attn_weights = attn_weights + attention_mask - - # upcast attention to fp32 - attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) - attn_output = torch.matmul(attn_weights, value_states) - - if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): - raise ValueError( - f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" - f" {attn_output.size()}" - ) - attn_output = attn_output.transpose(1, 2).contiguous() - if sp_mode == "all_to_all": - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim) - attn_output = all_to_all_comm( - attn_output, sp_group, scatter_dim=1, gather_dim=2, fp8_communication=shard_config.fp8_communication - ) - else: - attn_output = attn_output.reshape(bsz, q_len, -1) - - attn_output = self.o_proj(attn_output) - - return attn_output, None - - return forward - - -def get_qwen3_model_forward_for_flash_attn(shard_config: ShardConfig, sp_mode=None, sp_size=None, sp_group=None): - logger = logging.get_logger(__name__) - - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - return_dict: Optional[bool] = None, - force_sp_output_gather: bool = True, - ) -> Union[Tuple, BaseModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - batch_size, seq_length = input_ids.shape - elif inputs_embeds is not None: - batch_size, seq_length, _ = inputs_embeds.shape - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - seq_length_with_past = seq_length - past_key_values_length = 0 - - if past_key_values is not None: - past_key_values_length = past_key_values[0][0].shape[2] - seq_length_with_past = seq_length_with_past + past_key_values_length - - if position_ids is None: - device = input_ids.device if input_ids is not None else inputs_embeds.device - position_ids = torch.arange( - past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device - ) - position_ids = position_ids.unsqueeze(0).view(-1, seq_length) - else: - position_ids = position_ids.view(-1, seq_length).long() - - # embed positions - hidden_states = inputs_embeds - - if shard_config.enable_flash_attention: - # in this case, attention_mask is a dict rather than a tensor - mask_shape = (batch_size, 1, seq_length, seq_length_with_past) - attention_mask = ColoAttention.prepare_attn_kwargs( - mask_shape, - hidden_states.dtype, - hidden_states.device, - q_padding_mask=attention_mask, - is_causal=True, - ) - else: - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, - (batch_size, seq_length), - inputs_embeds, - past_key_values_length, - sliding_window=self.config.sliding_window, - ) - - if (self.gradient_checkpointing or sp_mode in ["ring", "all_to_all"]) and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - if sp_mode in ["ring", "split_gather"]: - hidden_states = split_forward_gather_backward( - hidden_states, 1, sp_group, fp8_communication=shard_config.fp8_communication - ) - elif sp_mode == "all_to_all": - hidden_states = split_forward_gather_backward( - hidden_states, 1, sp_group, 1 / sp_size, fp8_communication=shard_config.fp8_communication - ) - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - hidden_states = self.norm(hidden_states) - - if shard_config.enable_sequence_parallelism: - if (not shard_config.parallel_output) or force_sp_output_gather or is_share_sp_tp(sp_mode): - hidden_states = gather_sp_output(hidden_states, shard_config) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - - return forward - - -def get_lm_forward_with_dist_cross_entropy(shard_config: ShardConfig): - def forward( - self: Qwen3ForCausalLM, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - ) -> Union[Tuple, CausalLMOutputWithPast]: - r""" - Args: - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> from transformers import AutoTokenizer, Qwen2ForCausalLM - - >>> model = Qwen2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS) - >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER) - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) - outputs = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - force_sp_output_gather=False, - ) - - hidden_states = outputs[0] - logits = self.lm_head(hidden_states) - logits = logits.float() - loss = None - if labels is not None: - loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) - - if not return_dict: - output = (logits,) + outputs[1:] - return (loss,) + output if loss is not None else output - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - return forward +reexport_transformers_impl(globals(), __package__, "_qwen3") diff --git a/colossalai/shardformer/modeling/vit.py b/colossalai/shardformer/modeling/vit.py index 5106d97cf4bc..cf60db7af5e5 100644 --- a/colossalai/shardformer/modeling/vit.py +++ b/colossalai/shardformer/modeling/vit.py @@ -1,386 +1,3 @@ -from typing import List, Optional, Tuple, Union +from colossalai._compat import reexport_transformers_impl -import torch -from transformers.models.vit.modeling_vit import BaseModelOutput, ViTEncoder -from transformers.utils import logging - -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer import ColoAttention - - -def _encoder_forward( - encoder: ViTEncoder, - start_idx: int, - end_idx: int, - hidden_states: torch.Tensor, - head_mask: Optional[torch.Tensor] = None, - output_attentions: bool = False, - output_hidden_states: bool = False, - return_dict: bool = True, - stage_manager: PipelineStageManager = None, -) -> Union[tuple, BaseModelOutput]: - for i in range(start_idx, end_idx): - layer_module = encoder.layer[i] - - layer_head_mask = head_mask[i] if head_mask is not None else None - - if encoder.gradient_checkpointing and encoder.training: - layer_outputs = encoder._gradient_checkpointing_func( - layer_module.__call__, - hidden_states, - layer_head_mask, - output_attentions, - ) - else: - layer_outputs = layer_module(hidden_states, layer_head_mask, output_attentions) - - hidden_states = layer_outputs[0] - if not stage_manager.is_last_stage(): - return hidden_states - else: - if not return_dict: - return tuple(hidden_states) - return BaseModelOutput( - last_hidden_state=hidden_states, - hidden_states=None, - attentions=None, - ) - - -def ViTModel_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): - from transformers.models.vit.modeling_vit import BaseModelOutputWithPooling - - def pp_forward( - self, - pixel_values: Optional[torch.Tensor] = None, - bool_masked_pos: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - interpolate_pos_encoding: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.FloatTensor] = None, - ) -> Union[Tuple, BaseModelOutputWithPooling]: - r""" - bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*): - Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - logger = logging.get_logger(__name__) - - # Preprocess passed in arguments - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - - # Prepare head mask if needed - # 1.0 in head_mask indicate we keep the head - # attention_probs has shape bsz x n_heads x N x N - # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] - # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] - head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) - - if stage_manager.is_first_stage(): - if pixel_values is None: - raise ValueError("You have to specify pixel_values") - - # TODO(FoolPlayer): maybe have a cleaner way to cast the input (from `ImageProcessor` side?) - expected_dtype = self.embeddings.patch_embeddings.projection.weight.dtype - if pixel_values.dtype != expected_dtype: - pixel_values = pixel_values.to(expected_dtype) - - embedding_output = self.embeddings( - pixel_values, - bool_masked_pos=bool_masked_pos, - interpolate_pos_encoding=interpolate_pos_encoding, - ) - hidden_states = embedding_output - else: - assert ( - hidden_states is not None - ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" - - encoder_outputs = _encoder_forward( - encoder=self.encoder, - start_idx=stage_index[0], - end_idx=stage_index[1], - hidden_states=hidden_states, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - ) - if not stage_manager.is_last_stage(): - return {"hidden_states": encoder_outputs} - - sequence_output = encoder_outputs[0] - sequence_output = self.layernorm(sequence_output) - pooled_output = self.pooler(sequence_output) if self.pooler is not None else None - - if not return_dict: - head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,) - return head_outputs + encoder_outputs[1:] - - return BaseModelOutputWithPooling( - last_hidden_state=sequence_output, - pooler_output=pooled_output, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) - - return pp_forward - - -def ViTForImageClassification_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): - from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss - from transformers.models.vit.modeling_vit import ImageClassifierOutput - - def pp_forward( - self, - pixel_values: Optional[torch.Tensor] = None, - head_mask: Optional[torch.Tensor] = None, - labels: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - interpolate_pos_encoding: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.FloatTensor] = None, - ) -> Union[tuple, ImageClassifierOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): - Labels for computing the image classification/regression loss. Indices should be in `[0, ..., - config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If - `config.num_labels > 1` a classification loss is computed (Cross-Entropy). - """ - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if not stage_manager.is_first_stage(): - assert ( - hidden_states is not None - ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" - - outputs = self.vit( - pixel_values, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - interpolate_pos_encoding=interpolate_pos_encoding, - return_dict=return_dict, - hidden_states=hidden_states, - ) - - # not last stage, return hidden_states - if not stage_manager.is_last_stage(): - return outputs - else: - sequence_output = outputs[0] - - # last stage - logits = self.classifier(sequence_output[:, 0, :]) - - loss = None - if labels is not None: - # move labels to correct device to enable model parallelism - labels = labels.to(logits.device) - if self.config.problem_type is None: - if self.num_labels == 1: - self.config.problem_type = "regression" - elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int): - self.config.problem_type = "single_label_classification" - else: - self.config.problem_type = "multi_label_classification" - - if self.config.problem_type == "regression": - loss_fct = MSELoss() - if self.num_labels == 1: - loss = loss_fct(logits.squeeze(), labels.squeeze()) - else: - loss = loss_fct(logits, labels) - elif self.config.problem_type == "single_label_classification": - loss_fct = CrossEntropyLoss() - loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) - elif self.config.problem_type == "multi_label_classification": - loss_fct = BCEWithLogitsLoss() - loss = loss_fct(logits, labels) - - if not return_dict: - output = (logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output - - return ImageClassifierOutput( - loss=loss, - logits=logits, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - return pp_forward - - -def ViTForMaskedImageModeling_pipeline_forward(stage_manager: PipelineStageManager, stage_index: List[int]): - import math - - import torch.nn as nn - from transformers.models.vit.modeling_vit import ImageClassifierOutput, MaskedImageModelingOutput - - def pp_forward( - self, - pixel_values: Optional[torch.Tensor] = None, - bool_masked_pos: Optional[torch.BoolTensor] = None, - head_mask: Optional[torch.Tensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - interpolate_pos_encoding: Optional[bool] = None, - return_dict: Optional[bool] = None, - hidden_states: Optional[torch.FloatTensor] = None, - ) -> Union[tuple, ImageClassifierOutput]: - r""" - bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`): - Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). - - Returns: - - Examples: - ```python - >>> from transformers import AutoImageProcessor, ViTForMaskedImageModeling - >>> import torch - >>> from PIL import Image - >>> import requests - - >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" - >>> image = Image.open(requests.get(url, stream=True).raw) - - >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k") - >>> model = ViTForMaskedImageModeling.from_pretrained("google/vit-base-patch16-224-in21k") - - >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2 - >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values - >>> # create random boolean mask of shape (batch_size, num_patches) - >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool() - - >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos) - >>> loss, reconstructed_pixel_values = outputs.loss, outputs.reconstruction - >>> list(reconstructed_pixel_values.shape) - [1, 3, 224, 224] - ```""" - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if bool_masked_pos is not None and (self.config.patch_size != self.config.encoder_stride): - raise ValueError( - "When `bool_masked_pos` is provided, `patch_size` must be equal to `encoder_stride` to ensure that " - "the reconstructed image has the same dimensions as the input." - f"Got `patch_size` = {self.config.patch_size} and `encoder_stride` = {self.config.encoder_stride}." - ) - - if not stage_manager.is_first_stage(): - assert ( - hidden_states is not None - ), f"Current stage is {stage_manager.stage}, hidden_states should not be None" - - outputs = self.vit( - pixel_values, - bool_masked_pos=bool_masked_pos, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - interpolate_pos_encoding=interpolate_pos_encoding, - return_dict=return_dict, - hidden_states=hidden_states, - ) - if not stage_manager.is_last_stage(): - return outputs - else: - sequence_output = outputs[0] - - # Reshape to (batch_size, num_channels, height, width) - sequence_output = sequence_output[:, 1:] - batch_size, sequence_length, num_channels = sequence_output.shape - height = width = math.floor(sequence_length**0.5) - sequence_output = sequence_output.permute(0, 2, 1).reshape(batch_size, num_channels, height, width) - - # Reconstruct pixel values - reconstructed_pixel_values = self.decoder(sequence_output) - - masked_im_loss = None - if bool_masked_pos is not None: - size = self.config.image_size // self.config.patch_size - bool_masked_pos = bool_masked_pos.reshape(-1, size, size) - mask = ( - bool_masked_pos.repeat_interleave(self.config.patch_size, 1) - .repeat_interleave(self.config.patch_size, 2) - .unsqueeze(1) - .contiguous() - ) - reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none") - masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels - - if not return_dict: - output = (reconstructed_pixel_values,) + outputs[1:] - return ((masked_im_loss,) + output) if masked_im_loss is not None else output - - return MaskedImageModelingOutput( - loss=masked_im_loss, - reconstruction=reconstructed_pixel_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - ) - - return pp_forward - - -def get_vit_flash_self_attention_forward(): - from transformers.models.vit.modeling_vit import ViTSelfAttention - - def forward( - self: ViTSelfAttention, - hidden_states: torch.Tensor, - head_mask: Optional[torch.Tensor] = None, - output_attentions: bool = False, - ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: - assert head_mask is None, "head_mask is not supported for FlashAttention" - mixed_query_layer = self.query(hidden_states) - - key_layer = self.transpose_for_scores(self.key(hidden_states)) - value_layer = self.transpose_for_scores(self.value(hidden_states)) - query_layer = self.transpose_for_scores(mixed_query_layer) - - dropout_p = self.dropout_prob if self.training else 0.0 - context_layer = ColoAttention.attention(query_layer, key_layer, value_layer, dropout_p=dropout_p) - - context_layer = context_layer.permute(0, 2, 1, 3).contiguous() - new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) - context_layer = context_layer.view(new_context_layer_shape) - - outputs = (context_layer, None) if output_attentions else (context_layer,) - - return outputs - - return forward - - -def get_jit_fused_vit_output_forward(): - from transformers.models.vit.modeling_vit import ViTOutput - - def forward(self: ViTOutput, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor: - hidden_states = self.dense(hidden_states) - hidden_states = self.dropout_add(hidden_states, input_tensor, self.dropout.p, self.dropout.training) - return hidden_states - - return forward - - -def get_jit_fused_vit_intermediate_forward(): - from colossalai.kernel.jit.bias_gelu import GeLUFunction as JitGeLUFunction - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - hidden_states, bias = self.dense(hidden_states) - hidden_states = JitGeLUFunction.apply(hidden_states, bias) - - return hidden_states - - return forward +reexport_transformers_impl(globals(), __package__, "_vit") diff --git a/colossalai/shardformer/modeling/whisper.py b/colossalai/shardformer/modeling/whisper.py index 619bbc98e3a0..3285c6f3fde6 100644 --- a/colossalai/shardformer/modeling/whisper.py +++ b/colossalai/shardformer/modeling/whisper.py @@ -1,1193 +1,3 @@ -import logging -import random -from typing import List, Optional, Tuple, Union +from colossalai._compat import reexport_transformers_impl -import torch -from torch import nn -from torch.nn import CrossEntropyLoss -from transformers.modeling_attn_mask_utils import ( - _prepare_4d_causal_attention_mask, - _prepare_4d_causal_attention_mask_for_sdpa, -) -from transformers.modeling_outputs import ( - BaseModelOutput, - BaseModelOutputWithPastAndCrossAttentions, - Seq2SeqLMOutput, - Seq2SeqModelOutput, - SequenceClassifierOutput, -) -from transformers.models.whisper.modeling_whisper import ( - _HIDDEN_STATES_START_POSITION, - WhisperDecoder, - WhisperEncoder, - WhisperForAudioClassification, - WhisperForConditionalGeneration, - WhisperModel, - shift_tokens_right, -) -from transformers.utils import logging - -from colossalai.pipeline.stage_manager import PipelineStageManager -from colossalai.shardformer.layer import ColoAttention -from colossalai.shardformer.shard import ShardConfig - -logger = logging.get_logger(__name__) - - -def _get_attention_mask( - self: WhisperDecoder, - shard_config: ShardConfig, - hidden_states: torch.Tensor, - past_key_values_length: int, - attention_mask: Optional[torch.FloatTensor], - head_mask: Optional[torch.Tensor] = None, - output_attentions: bool = False, -): - batch_size, seq_length = hidden_states.shape[:2] - mask_seq_length = past_key_values_length + seq_length - if shard_config.enable_flash_attention: - attention_mask = ColoAttention.prepare_attn_kwargs( - (batch_size, 1, seq_length, mask_seq_length), - hidden_states.dtype, - hidden_states.device, - attention_mask, - is_causal=True, - ) - else: - input_shape = (batch_size, seq_length) - if self._use_flash_attention_2: - # 2d mask is passed through the layers - attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None - elif self._use_sdpa and head_mask is None and not output_attentions: - # output_attentions=True & head_mask can not be supported when using SDPA. - attention_mask = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, input_shape, hidden_states, past_key_values_length - ) - else: - # 4d mask is passed through the layers - attention_mask = _prepare_4d_causal_attention_mask( - attention_mask, input_shape, hidden_states, past_key_values_length - ) - return attention_mask - - -def get_whisper_flash_attention_forward(): - from transformers.models.whisper.modeling_whisper import WhisperAttention - - def forward( - self: WhisperAttention, - hidden_states: torch.Tensor, - key_value_states: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - attention_mask: Optional[dict] = None, - layer_head_mask: Optional[torch.Tensor] = None, - output_attentions: bool = False, - cache_position: Optional[torch.LongTensor] = None, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: - """Input shape: Batch x Time x Channel""" - assert layer_head_mask is None, "layer_head_mask is not supported for FlashAttention" - # for encoder, attention_mask is None - if attention_mask is None: - attention_mask = {} - # if key_value_states are provided this layer is used as a cross-attention layer - # for the decoder - is_cross_attention = key_value_states is not None - - bsz, tgt_len, _ = hidden_states.size() - - # get query proj - query_states = self.q_proj(hidden_states) - # get key, value proj - # `past_key_value[0].shape[2] == key_value_states.shape[1]` - # is checking that the `sequence_length` of the `past_key_value` is the same as - # the provided `key_value_states` to support prefix tuning - if ( - is_cross_attention - and past_key_value is not None - and past_key_value[0].shape[2] == key_value_states.shape[1] - ): - # reuse k,v, cross_attentions - key_states = past_key_value[0] - value_states = past_key_value[1] - elif is_cross_attention: - # cross_attentions - key_states = self._shape(self.k_proj(key_value_states), -1, bsz) - value_states = self._shape(self.v_proj(key_value_states), -1, bsz) - elif past_key_value is not None: - # reuse k, v, self_attention - key_states = self._shape(self.k_proj(hidden_states), -1, bsz) - value_states = self._shape(self.v_proj(hidden_states), -1, bsz) - key_states = torch.cat([past_key_value[0], key_states], dim=2) - value_states = torch.cat([past_key_value[1], value_states], dim=2) - else: - # self_attention - key_states = self._shape(self.k_proj(hidden_states), -1, bsz) - value_states = self._shape(self.v_proj(hidden_states), -1, bsz) - - if self.is_decoder: - # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states. - # Further calls to cross_attention layer can then reuse all cross-attention - # key/value_states (first "if" case) - # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of - # all previous decoder key/value_states. Further calls to uni-directional self-attention - # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case) - # if encoder bi-directional self-attention `past_key_value` is always `None` - past_key_value = (key_states, value_states) - - query_states = self._shape(query_states, tgt_len, bsz) - - dropout_p = self.dropout if self.training else 0.0 - attn_output = ColoAttention.attention( - query_states, - key_states, - value_states, - **attention_mask, - dropout_p=dropout_p, - scale=self.scaling, - ) - attn_output = attn_output.transpose(1, 2) - - # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be - # partitioned across GPUs when using tensor-parallelism. - attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim) - - attn_output = self.out_proj(attn_output) - - return attn_output, None, past_key_value - - return forward - - -def get_whisper_decoder_forward_for_flash_attention(shard_config: ShardConfig): - def forward( - self: WhisperDecoder, - input_ids=None, - attention_mask=None, - encoder_hidden_states=None, - head_mask=None, - cross_attn_head_mask=None, - past_key_values=None, - inputs_embeds=None, - position_ids=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - cache_position=None, - ): - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - # past_key_values_length - past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - attention_mask = _get_attention_mask(self, shard_config, inputs_embeds, past_key_values_length, attention_mask) - - # embed positions - if input_ids is not None: - positions = self.embed_positions( - input_ids, past_key_values_length=past_key_values_length, position_ids=position_ids - ) - else: - positions = self.embed_positions( - inputs_embeds, past_key_values_length=past_key_values_length, position_ids=position_ids - ) - - hidden_states = inputs_embeds + positions - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`..." - ) - use_cache = False - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None - next_decoder_cache = () if use_cache else None - - # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired - for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]): - if attn_mask is not None: - assert attn_mask.size()[0] == (len(self.layers)), ( - f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for" - f" {head_mask.size()[0]}." - ) - for idx, decoder_layer in enumerate(self.layers): - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) - if output_hidden_states: - all_hidden_states += (hidden_states,) - if self.training: - dropout_probability = torch.rand([]) - if dropout_probability < self.layerdrop: - continue - - past_key_value = past_key_values[idx] if past_key_values is not None else None - - if self.gradient_checkpointing and self.training: - - def create_custom_forward(module): - def custom_forward(*inputs): - # None for past_key_value - return module(*inputs, output_attentions, use_cache) - - return custom_forward - - layer_outputs = torch.utils.checkpoint.checkpoint( - create_custom_forward(decoder_layer), - hidden_states, - attention_mask, - encoder_hidden_states, - None, # encoder attention mask - head_mask[idx] if head_mask is not None else None, - (cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), - None, # past_key_value - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - cross_attn_layer_head_mask=( - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None - ), - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - ) - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache += (layer_outputs[3 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if encoder_hidden_states is not None: - all_cross_attentions += (layer_outputs[2],) - - hidden_states = self.layer_norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - next_cache, - all_hidden_states, - all_self_attns, - all_cross_attentions, - ] - if v is not None - ) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - cross_attentions=all_cross_attentions, - ) - - return forward - - -def get_jit_fused_whisper_encoder_layer_forward(): - from transformers.models.whisper.modeling_whisper import WhisperEncoderLayer - - def forward( - self: WhisperEncoderLayer, - hidden_states: torch.Tensor, - attention_mask: torch.Tensor, - layer_head_mask: torch.Tensor, - output_attentions: bool = False, - ) -> torch.Tensor: - """ - Args: - hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` - attention_mask (`torch.FloatTensor`): attention mask of size - `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. - layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size - `(encoder_attention_heads,)`. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - """ - residual = hidden_states - hidden_states = self.self_attn_layer_norm(hidden_states) - hidden_states, attn_weights, _ = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - layer_head_mask=layer_head_mask, - output_attentions=output_attentions, - ) - hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) - - residual = hidden_states - hidden_states = self.final_layer_norm(hidden_states) - hidden_states = self.activation_fn(self.fc1(hidden_states)) - hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) - hidden_states = self.fc2(hidden_states) - hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) - - if hidden_states.dtype == torch.float16 and ( - torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any() - ): - clamp_value = torch.finfo(hidden_states.dtype).max - 1000 - hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value) - - outputs = (hidden_states,) - - if output_attentions: - outputs += (attn_weights,) - - return outputs - - return forward - - -def get_jit_fused_whisper_decoder_layer_forward(): - from transformers.models.whisper.modeling_whisper import WhisperDecoderLayer - - def forward( - self: WhisperDecoderLayer, - hidden_states: torch.Tensor, - attention_mask: Optional[torch.Tensor] = None, - encoder_hidden_states: Optional[torch.Tensor] = None, - encoder_attention_mask: Optional[torch.Tensor] = None, - layer_head_mask: Optional[torch.Tensor] = None, - cross_attn_layer_head_mask: Optional[torch.Tensor] = None, - past_key_value: Optional[Tuple[torch.Tensor]] = None, - output_attentions: Optional[bool] = False, - use_cache: Optional[bool] = True, - ) -> torch.Tensor: - """ - Args: - hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` - attention_mask (`torch.FloatTensor`): attention mask of size - `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. - encoder_hidden_states (`torch.FloatTensor`): - cross attention input to the layer of shape `(batch, seq_len, embed_dim)` - encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size - `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. - layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size - `(encoder_attention_heads,)`. - cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of - size `(decoder_attention_heads,)`. - past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - """ - residual = hidden_states - hidden_states = self.self_attn_layer_norm(hidden_states) - - # Self Attention - # decoder uni-directional self-attention cached key/values tuple is at positions 1,2 - self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None - # add present self-attn cache to positions 1,2 of present_key_value tuple - hidden_states, self_attn_weights, present_key_value = self.self_attn( - hidden_states=hidden_states, - past_key_value=self_attn_past_key_value, - attention_mask=attention_mask, - layer_head_mask=layer_head_mask, - output_attentions=output_attentions, - ) - hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) - - # Cross-Attention Block - cross_attn_present_key_value = None - cross_attn_weights = None - if encoder_hidden_states is not None: - residual = hidden_states - hidden_states = self.encoder_attn_layer_norm(hidden_states) - - # cross_attn cached key/values tuple is at positions 3,4 of present_key_value tuple - cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None - hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn( - hidden_states=hidden_states, - key_value_states=encoder_hidden_states, - attention_mask=encoder_attention_mask, - layer_head_mask=cross_attn_layer_head_mask, - past_key_value=cross_attn_past_key_value, - output_attentions=output_attentions, - ) - hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) - - # add cross-attn to positions 3,4 of present_key_value tuple - present_key_value = present_key_value + cross_attn_present_key_value - - # Fully Connected - residual = hidden_states - hidden_states = self.final_layer_norm(hidden_states) - hidden_states = self.activation_fn(self.fc1(hidden_states)) - hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training) - hidden_states = self.fc2(hidden_states) - hidden_states = self.dropout_add(hidden_states, residual, self.dropout, self.training) - - outputs = (hidden_states,) - - if output_attentions: - outputs += (self_attn_weights, cross_attn_weights) - - if use_cache: - outputs += (present_key_value,) - - return outputs - - return forward - - -class WhisperPipelineForwards: - """ - This class serves as a micro library for forward function substitution of Llama models - under pipeline setting. - """ - - @staticmethod - def whisper_encoder_forward( - self: WhisperEncoder, - input_features, - attention_mask=None, - head_mask=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - encoder_states=None, - all_attentions=None, - stage_index: Optional[List[int]] = None, - decoder_starting_stage: Optional[int] = None, - shard_config: Optional[ShardConfig] = None, - ): - r""" - Args: - input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`): - Float values of mel features extracted from the raw speech waveform. Raw speech waveform can be - obtained by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a - `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into - `input_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding - and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] - attention_mask (`torch.Tensor`)`, *optional*): - Whisper does not support masking of the `input_features`, this argument is preserved for compatibility, - but it is not used. By default the silence in the input log mel spectrogram are ignored. - head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): - Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - output_hidden_states (`bool`, *optional*): - Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors - for more detail. - return_dict (`bool`, *optional*): - Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. - """ - logging.get_logger(__name__) - - stage = stage_manager.stage - at_first_stage = stage == 0 - at_last_stage = stage == decoder_starting_stage - 1 - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # Process inputs if at the first stage of encoder. - if at_first_stage: - inputs_embeds = nn.functional.gelu(self.conv1(input_features)) - inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds)) - - inputs_embeds = inputs_embeds.permute(0, 2, 1) - embed_pos = self.embed_positions.weight - - hidden_states = inputs_embeds + embed_pos - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - - encoder_states = () if output_hidden_states else None - all_attentions = () if output_attentions else None - - # check if head_mask has a correct number of layers specified if desired - if head_mask is not None: - assert head_mask.size()[0] == ( - len(self.layers) - ), f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}." - - else: - if hidden_states is None: - raise ValueError( - "hidden_states shouldn't be None for stages other than the first stage of encoder/decoder." - ) - - start_idx, end_idx = stage_index[0], stage_index[1] - - for idx in range(start_idx, end_idx): - encoder_layer = self.layers[idx] - - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) - dropout_probability = random.uniform(0, 1) - if self.training and (dropout_probability < self.layerdrop): # skip the layer - layer_outputs = (None, None) - else: - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - encoder_layer.__call__, - hidden_states, - None, - (head_mask[idx] if head_mask is not None else None), - output_attentions, - ) - else: - layer_outputs = encoder_layer( - hidden_states, - None, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - output_attentions=output_attentions, - ) - - hidden_states = layer_outputs[0] - - if output_attentions: - all_attentions = all_attentions + (layer_outputs[1],) - - if at_last_stage: - hidden_states = self.layer_norm(hidden_states) - if output_hidden_states: - encoder_states = encoder_states + (hidden_states,) - - if not return_dict: - return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None) - return BaseModelOutput( - last_hidden_state=hidden_states, - hidden_states=encoder_states, - attentions=all_attentions, - ) - - else: - return {"hidden_states": hidden_states, "head_mask": head_mask} - - @staticmethod - def whisper_decoder_forward( - self: WhisperDecoder, - input_ids=None, - attention_mask=None, - encoder_hidden_states=None, - head_mask=None, - cross_attn_head_mask=None, - past_key_values=None, - inputs_embeds=None, - position_ids=None, - use_cache=None, - output_attentions=None, - output_hidden_states=None, - return_dict=None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - decoder_starting_stage: Optional[int] = None, - shard_config: Optional[ShardConfig] = None, - ): - r""" - Args: - input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you - provide it. - - Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and - [`PreTrainedTokenizer.__call__`] for details. - - [What are input IDs?](../glossary#input-ids) - attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): - Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - - [What are attention masks?](../glossary#attention-mask) - encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): - Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention - of the decoder. - head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): - Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): - Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention - on hidden heads. Mask values selected in `[0, 1]`: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of - shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of - shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. - - Contains pre-computed hidden-states (key and values in the self-attention blocks and in the - cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. - - If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those - that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of - all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of - shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing - `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more - control over how to convert `input_ids` indices into associated vectors than the model's internal - embedding lookup matrix. - output_attentions (`bool`, *optional*): - Whether or not to return the attentions tensors of all attention layers. See `attentions` under - returned tensors for more detail. - output_hidden_states (`bool`, *optional*): - Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors - for more detail. - return_dict (`bool`, *optional*): - Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. - """ - logger = logging.get_logger(__name__) - stage = stage_manager.stage - at_first_stage = stage == decoder_starting_stage - at_last_stage = stage == stage_manager.num_stages - 1 - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None - next_decoder_cache = () if use_cache else None - - # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired - for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]): - if attn_mask is not None: - assert attn_mask.size()[0] == (len(self.layers)), ( - f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for" - f" {head_mask.size()[0]}." - ) - - # past_key_values_length - past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0 - - if at_first_stage: - # retrieve input_ids and inputs_embeds - if input_ids is not None and inputs_embeds is not None: - raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") - elif input_ids is not None: - input_shape = input_ids.size() - input_ids = input_ids.view(-1, input_shape[-1]) - elif inputs_embeds is not None: - input_shape = inputs_embeds.size()[:-1] - else: - raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - attention_mask = _get_attention_mask( - self, shard_config, inputs_embeds, past_key_values_length, attention_mask - ) - - # embed positions - if input_ids is not None: - positions = self.embed_positions( - input_ids, past_key_values_length=past_key_values_length, position_ids=position_ids - ) - else: - positions = self.embed_positions( - inputs_embeds, past_key_values_length=past_key_values_length, position_ids=position_ids - ) - - hidden_states = inputs_embeds + positions - hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training) - - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`..." - ) - use_cache = False - - else: - if hidden_states is None: - raise ValueError( - "hidden_states shouldn't be None for stages other than the first stage of encoder/decoder." - ) - input_shape = hidden_states.size()[:-1] - attention_mask = _get_attention_mask( - self, - shard_config, - hidden_states, - past_key_values_length, - attention_mask, - ) - - start_idx, end_idx = stage_index[0], stage_index[1] - - for idx in range(start_idx, end_idx): - # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description) - decoder_layer = self.layers[idx] - - if output_hidden_states: - all_hidden_states += (hidden_states,) - dropout_probability = random.uniform(0, 1) - if self.training and (dropout_probability < self.layerdrop): - continue - - past_key_value = past_key_values[idx] if past_key_values is not None else None - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - encoder_hidden_states, - None, # encoder attention mask - head_mask[idx] if head_mask is not None else None, - (cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None), - None, # past_key_value - output_attentions, - use_cache, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - encoder_hidden_states=encoder_hidden_states, - layer_head_mask=(head_mask[idx] if head_mask is not None else None), - cross_attn_layer_head_mask=( - cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None - ), - past_key_value=past_key_value, - output_attentions=output_attentions, - use_cache=use_cache, - ) - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache += (layer_outputs[3 if output_attentions else 1],) - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - if encoder_hidden_states is not None: - all_cross_attentions += (layer_outputs[2],) - - if at_last_stage: - hidden_states = self.layer_norm(hidden_states) - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - next_cache = next_decoder_cache if use_cache else None - if not return_dict: - return tuple( - v - for v in [ - hidden_states, - next_cache, - all_hidden_states, - all_self_attns, - all_cross_attentions, - ] - if v is not None - ) - return BaseModelOutputWithPastAndCrossAttentions( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - cross_attentions=all_cross_attentions, - ) - - else: - return { - "head_mask": head_mask, - "cross_attn_head_mask": cross_attn_head_mask, - "hidden_states": hidden_states, - } - - @staticmethod - def whisper_model_forward( - self: WhisperModel, - input_features: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.LongTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - decoder_head_mask: Optional[torch.Tensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, - decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - decoder_starting_stage: Optional[int] = None, - shard_config: Optional[ShardConfig] = None, - ): - r""" - Returns: - - Example: - ```python - >>> import torch - >>> from transformers import AutoFeatureExtractor, WhisperModel - >>> from datasets import load_dataset - - >>> model = WhisperModel.from_pretrained("openai/whisper-base") - >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base") - >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") - >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") - >>> input_features = inputs.input_features - >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id - >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state - >>> list(last_hidden_state.shape) - [1, 2, 512] - ```""" - # TODO: left the recording kv-value tensors as () or None type, this feature may be added in the future. - if past_key_values: - logger.warning_once("Non-empty past_key_values is not supported for pipeline models at the moment.") - past_key_values = None - if output_attentions: - logger.warning_once("output_attentions=True is not supported for pipeline models at the moment.") - output_attentions = False - if output_hidden_states: - logger.warning_once("output_hidden_states=True is not supported for pipeline models at the moment.") - output_hidden_states = False - if use_cache: - logger.warning_once("use_cache=True is not supported for pipeline models at the moment.") - use_cache = False - - logging.get_logger(__name__) - - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - in_decoder = stage_manager.stage >= decoder_starting_stage - if not in_decoder: - if encoder_outputs is None: - input_features = self._mask_input_features(input_features, attention_mask=attention_mask) - - encoder_outputs = WhisperPipelineForwards.whisper_encoder_forward( - self.encoder, - input_features, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - decoder_starting_stage=decoder_starting_stage, - ) - - if stage_manager.stage == decoder_starting_stage - 1: - # last stage of encoder - return {"encoder_hidden_states": encoder_outputs[0]} - else: - return encoder_outputs - # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput when return_dict=True - elif return_dict and not isinstance(encoder_outputs, BaseModelOutput): - encoder_outputs = BaseModelOutput( - last_hidden_state=encoder_outputs[0], - hidden_states=(encoder_outputs[1] if len(encoder_outputs) > 1 else None), - attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None, - ) - - at_last_decoder_stage = stage_manager.is_last_stage() - at_first_decoder_stage = stage_manager.stage == decoder_starting_stage - if encoder_outputs is not None: - encoder_hidden_states = encoder_outputs[0] - elif encoder_hidden_states is None: - raise ValueError("Non-empty encoder_hidden_states should be passed in at decoder stages.") - - if not at_first_decoder_stage and hidden_states is None: - raise ValueError("If not at the first layer of decoder, non-empty hidden_states must be provided.") - - # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn) - decoder_outputs = WhisperPipelineForwards.whisper_decoder_forward( - self.decoder, - input_ids=decoder_input_ids, - attention_mask=decoder_attention_mask, - encoder_hidden_states=encoder_hidden_states, - head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - past_key_values=past_key_values, - inputs_embeds=decoder_inputs_embeds, - position_ids=decoder_position_ids, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - decoder_starting_stage=decoder_starting_stage, - shard_config=shard_config, - ) - - # Directly return outputs of overloaded Whisper forward if not at last stage. - if not at_last_decoder_stage: - # encoder_hidden_states should be passed to the next stage - decoder_outputs["encoder_hidden_states"] = encoder_hidden_states - return decoder_outputs - - if not return_dict: - return decoder_outputs + encoder_outputs - - return Seq2SeqModelOutput( - last_hidden_state=decoder_outputs.last_hidden_state, - past_key_values=decoder_outputs.past_key_values, - decoder_hidden_states=decoder_outputs.hidden_states, - decoder_attentions=decoder_outputs.attentions, - cross_attentions=decoder_outputs.cross_attentions, - encoder_last_hidden_state=encoder_hidden_states, - ) - - @staticmethod - def whisper_for_conditional_generation_forward( - self: WhisperForConditionalGeneration, - input_features: Optional[torch.FloatTensor] = None, - attention_mask: Optional[torch.LongTensor] = None, - decoder_input_ids: Optional[torch.LongTensor] = None, - decoder_attention_mask: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - decoder_head_mask: Optional[torch.Tensor] = None, - cross_attn_head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - decoder_inputs_embeds: Optional[Tuple[torch.FloatTensor]] = None, - decoder_position_ids: Optional[Tuple[torch.LongTensor]] = None, - labels: Optional[torch.LongTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - encoder_hidden_states: Optional[torch.FloatTensor] = None, - stage_index: Optional[List[int]] = None, - decoder_starting_stage: Optional[int] = None, - shard_config: Optional[ShardConfig] = None, - ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` - or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is - only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - - Returns: - - Example: - - ```python - >>> import torch - >>> from transformers import AutoProcessor, WhisperForConditionalGeneration - >>> from datasets import load_dataset - - >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") - >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") - - >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") - - >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") - >>> input_features = inputs.input_features - - >>> generated_ids = model.generate(inputs=input_features) - - >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] - >>> transcription - ' Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.' - ```""" - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if labels is not None: - if decoder_input_ids is None and decoder_inputs_embeds is None: - decoder_input_ids = shift_tokens_right( - labels, self.config.pad_token_id, self.config.decoder_start_token_id - ) - in_decoder = stage_manager.stage >= decoder_starting_stage - at_last_decoder_stage = stage_manager.is_last_stage() - outputs = WhisperPipelineForwards.whisper_model_forward( - self.model, - input_features, - attention_mask=attention_mask, - decoder_input_ids=decoder_input_ids, - encoder_outputs=encoder_outputs, - decoder_attention_mask=decoder_attention_mask, - head_mask=head_mask, - decoder_head_mask=decoder_head_mask, - cross_attn_head_mask=cross_attn_head_mask, - past_key_values=past_key_values, - decoder_inputs_embeds=decoder_inputs_embeds, - decoder_position_ids=decoder_position_ids, - use_cache=use_cache, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - encoder_hidden_states=encoder_hidden_states, - stage_index=stage_index, - decoder_starting_stage=decoder_starting_stage, - shard_config=shard_config, - ) - if not in_decoder: - return outputs - - if not at_last_decoder_stage: - # encoder_hidden_states should be passed to the next stage - outputs["encoder_hidden_states"] = encoder_hidden_states - return outputs - - lm_logits = self.proj_out(outputs[0]) - - loss = None - if labels is not None: - loss_fct = CrossEntropyLoss() - # move labels to correct device to enable PP - labels = labels.to(lm_logits.device) - loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.reshape(-1)) - - if not return_dict: - output = (lm_logits,) + outputs[1:] - return ((loss,) + output) if loss is not None else output - - return Seq2SeqLMOutput( - loss=loss, - logits=lm_logits, - past_key_values=outputs.past_key_values, - decoder_hidden_states=outputs.decoder_hidden_states, - decoder_attentions=outputs.decoder_attentions, - cross_attentions=outputs.cross_attentions, - encoder_last_hidden_state=outputs.encoder_last_hidden_state, - encoder_hidden_states=outputs.encoder_hidden_states, - encoder_attentions=outputs.encoder_attentions, - ) - - @staticmethod - def whisper_for_audio_classification_forward( - self: WhisperForAudioClassification, - input_features: Optional[torch.LongTensor] = None, - head_mask: Optional[torch.Tensor] = None, - encoder_outputs: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, - labels: Optional[torch.LongTensor] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - stage_manager: Optional[PipelineStageManager] = None, - hidden_states: Optional[torch.FloatTensor] = None, - encoder_states=None, - all_attentions=None, - stage_index: Optional[List[int]] = None, - decoder_starting_stage: Optional[int] = None, - shard_config: Optional[ShardConfig] = None, - ): - r""" - This function is modified on the basis of transformers.models.whisper.modeling_whisper.WhisperForAudioClassification.forward. - Please refer to original code of transformers for more details. - """ - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - - if self.config.use_weighted_layer_sum: - output_hidden_states = True - elif output_hidden_states is None: - output_hidden_states = self.config.output_hidden_states - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - # audio_classification only holds encoder - encoder_outputs = WhisperPipelineForwards.whisper_encoder_forward( - self.encoder, - input_features, - head_mask=head_mask, - output_attentions=output_attentions, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - stage_manager=stage_manager, - hidden_states=hidden_states, - stage_index=stage_index, - decoder_starting_stage=decoder_starting_stage, - ) - - if not stage_manager.is_last_stage(): - return encoder_outputs - - if self.config.use_weighted_layer_sum: - hidden_states = encoder_outputs[_HIDDEN_STATES_START_POSITION] - hidden_states = torch.stack(hidden_states, dim=1) - norm_weights = nn.functional.softmax(self.layer_weights, dim=-1) - hidden_states = (hidden_states * norm_weights.view(-1, 1, 1)).sum(dim=1) - else: - hidden_states = encoder_outputs[0] - - hidden_states = self.projector(hidden_states) - pooled_output = hidden_states.mean(dim=1) - - logits = self.classifier(pooled_output) - - loss = None - - if labels is not None: - loss_fct = CrossEntropyLoss() - # move labels to correct device to enable PP - labels = labels.to(logits.device) - loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1)) - - if not return_dict: - output = (logits,) + encoder_outputs[1:] - return ((loss,) + output) if loss is not None else output - - return SequenceClassifierOutput( - loss=loss, - logits=logits, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) +reexport_transformers_impl(globals(), __package__, "_whisper") diff --git a/colossalai/shardformer/policies/_bert_tf4.py b/colossalai/shardformer/policies/_bert_tf4.py new file mode 100644 index 000000000000..fd4e020b0cb1 --- /dev/null +++ b/colossalai/shardformer/policies/_bert_tf4.py @@ -0,0 +1,825 @@ +import warnings +from functools import partial +from typing import Callable, Dict, List + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +import colossalai.shardformer.layer as col_nn + +from ..modeling.bert import ( + BertPipelineForwards, + bert_sequence_parallel_forward_fn, + get_bert_sequence_parallel_attention_forward, + get_jit_fused_bert_intermediate_forward, + get_jit_fused_bert_output_forward, + get_jit_fused_bert_self_output_forward, +) +from ..modeling.jit import get_jit_fused_dropout_add_func +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = [ + "BertPolicy", + "BertModelPolicy", + "BertForPreTrainingPolicy", + "BertLMHeadModelPolicy", + "BertForMaskedLMPolicy", + "BertForNextSentencePredictionPolicy", + "BertForSequenceClassificationPolicy", + "BertForTokenClassificationPolicy", + "BertForMultipleChoicePolicy", + "BertForQuestionAnsweringPolicy", +] + + +class BertPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.enable_bias_gelu_fused = self.shard_config.enable_jit_fused and self.model.config.hidden_act == "gelu" + return self.model + + def module_policy(self): + from transformers.models.bert.modeling_bert import ( + BertEmbeddings, + BertIntermediate, + BertLayer, + BertModel, + BertOutput, + BertSdpaSelfAttention, + BertSelfOutput, + ) + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = col_nn.VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = col_nn.PaddingEmbedding + + if self.shard_config.enable_fused_normalization: + norm_cls = col_nn.FusedLayerNorm + else: + norm_cls = col_nn.LayerNorm + + sp_mode = self.shard_config.sequence_parallelism_mode or None + assert sp_mode != "all_to_all", "all_to_all sequence parallelism is not supported for Bert" + if sp_mode == "ring": + warnings.warn( + f"For Bert, sequence parallelism is currently not support mode {sp_mode}, will set to be split_gather" + ) + sp_mode = "split_gather" + + sp_partial_derived = sp_mode == "split_gather" + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_sequence_parallelism: + # Fix the tgt_len size in bert sequence parallel attention forward. + self.append_or_create_method_replacement( + description={ + "forward": get_bert_sequence_parallel_attention_forward(self.shard_config), + }, + policy=policy, + target_key=BertSdpaSelfAttention, + ) + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + policy[BertLayer] = ModulePolicyDescription( + attribute_replacement={ + "attention.self.all_head_size": self.model.config.hidden_size + // self.shard_config.tensor_parallel_size, + "crossattention.self.all_head_size": self.model.config.hidden_size + // self.shard_config.tensor_parallel_size, + "attention.self.num_attention_heads": self.model.config.num_attention_heads + // self.shard_config.tensor_parallel_size, + "crossattention.self.num_attention_heads": self.model.config.num_attention_heads + // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.self.query", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.key", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.value", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=col_nn.Linear1D_Row, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "skip_bias_add": self.enable_bias_gelu_fused, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=col_nn.Linear1D_Row, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + ], + ) + + policy[BertEmbeddings] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + ] + ) + if self.enable_bias_gelu_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_intermediate_forward(), + }, + policy=policy, + target_key=BertIntermediate, + ) + + elif use_zbv: + policy[BertLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.self.query", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.key", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.value", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "skip_bias_add": self.enable_bias_gelu_fused, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + ], + ) + + policy[BertEmbeddings] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + ] + ) + if self.enable_bias_gelu_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_intermediate_forward(), + }, + policy=policy, + target_key=BertIntermediate, + ) + + if sp_mode == "split_gather": + self.append_or_create_method_replacement( + description={"forward": bert_sequence_parallel_forward_fn(self.shard_config)}, + policy=policy, + target_key=BertModel, + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="word_embeddings", + target_module=embedding_cls, + kwargs=( + { + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {} + ), + ) + ], + policy=policy, + target_key=BertEmbeddings, + ) + + # optimization configuration + # Handle bert layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="attention.output.LayerNorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="output.LayerNorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=BertLayer, + ) + # handle embedding layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="LayerNorm", + target_module=norm_cls, + ) + ], + policy=policy, + target_key=BertEmbeddings, + ) + + # use jit operator + if self.shard_config.enable_jit_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_self_output_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=BertSelfOutput, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_output_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=BertOutput, + ) + + return policy + + def add_lm_head_policy(self, base_policy): + from transformers.models.bert.modeling_bert import BertLMPredictionHead + + # optimize for tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="decoder", + target_module=col_nn.VocabParallelLMHead1D, + kwargs={ + "gather_output": True, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ), + policy=base_policy, + target_key=BertLMPredictionHead, + ) + else: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="decoder", + target_module=col_nn.PaddingLMHead, + kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, + ), + policy=base_policy, + target_key=BertLMPredictionHead, + ) + + # optimize with fused normalization + if self.shard_config.enable_fused_normalization: + # Handle bert lm prediction head + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="transform.LayerNorm", + target_module=col_nn.FusedLayerNorm, + ), + policy=base_policy, + target_key=BertLMPredictionHead, + ) + return base_policy + + def add_lm_prediction_policy(self, base_policy): + from transformers.models.bert.modeling_bert import BertLMPredictionHead + + method_replacement = { + "_save_to_state_dict": col_nn.ParallelModule._save_to_state_dict, + "_load_from_state_dict": col_nn.ParallelModule._load_from_state_dict, + } + self.append_or_create_method_replacement( + description=method_replacement, + policy=base_policy, + target_key=BertLMPredictionHead, + ) + return base_policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """ + If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy. + """ + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "BertModel": + module = self.model + else: + module = self.model.bert + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, + stage_manager=stage_manager, + shard_config=self.shard_config, + ) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=self.shard_config, + ) + } + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "BertModel": + module = self.model + else: + module = self.model.bert + stage_manager = self.pipeline_stage_manager + + held_layers = [] + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embeddings) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.encoder.layer[start_idx:end_idx]) + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(module.pooler) + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + if stage_manager.is_first_stage(): + held_layers.append(module.embeddings) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.encoder.layer[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.pooler) + + return held_layers + + +# BertModel +class BertModelPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + from transformers.models.bert.modeling_bert import BertModel + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertModel, + new_forward=BertPipelineForwards.bert_model_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in bert model""" + return [] + + +# BertForPreTraining +class BertForPreTrainingPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + policy = self.add_lm_prediction_policy(policy) + from transformers.models.bert.modeling_bert import BertForPreTraining + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForPreTraining, + new_forward=BertPipelineForwards.bert_for_pretraining_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage""" + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + model = self.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if id(model.bert.embeddings.word_embeddings.weight) == id(model.cls.predictions.decoder.weight): + # tie weights + return [ + { + 0: model.bert.embeddings.word_embeddings.weight, + self.pipeline_stage_manager.num_stages - 1: model.cls.predictions.decoder.weight, + } + ] + return [] + + +# BertLMHeadModel +class BertLMHeadModelPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + policy = self.add_lm_prediction_policy(policy) + from transformers.models.bert.modeling_bert import BertLMHeadModel + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertLMHeadModel, + new_forward=BertPipelineForwards.bert_lm_head_model_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + bert_model = self.model.bert + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if id(bert_model.embeddings.word_embeddings.weight) == id(self.model.cls.predictions.decoder.weight): + # tie weights + return [ + { + 0: bert_model.embeddings.word_embeddings.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.cls.predictions.decoder.weight, + } + ] + return [] + + +# BertForMaskedLM +class BertForMaskedLMPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + policy = self.add_lm_prediction_policy(policy) + from transformers.models.bert.modeling_bert import BertForMaskedLM + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForMaskedLM, + new_forward=BertPipelineForwards.bert_for_masked_lm_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + bert_model = self.model.bert + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if id(bert_model.embeddings.word_embeddings.weight) == id(self.model.cls.predictions.decoder.weight): + # tie weights + return [ + { + 0: bert_model.embeddings.word_embeddings.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.cls.predictions.decoder.weight, + } + ] + return [] + + +# BertForSequenceClassification +class BertForSequenceClassificationPolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForSequenceClassification + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ] + ) + } + policy.update(addon_module) + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForSequenceClassification, + new_forward=BertPipelineForwards.bert_for_sequence_classification_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.dropout) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +# BertForTokenClassification +class BertForTokenClassificationPolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForTokenClassification + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForTokenClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ] + ) + } + policy.update(addon_module) + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForTokenClassification, + new_forward=BertPipelineForwards.bert_for_token_classification_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.dropout) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +# BertForNextSentencePrediction +class BertForNextSentencePredictionPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + from transformers.models.bert.modeling_bert import BertForNextSentencePrediction + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForNextSentencePrediction, + new_forward=BertPipelineForwards.bert_for_next_sentence_prediction_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +# BertForMultipleChoice +class BertForMultipleChoicePolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForMultipleChoice + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForMultipleChoice: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ] + ) + } + policy.update(addon_module) + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForMultipleChoice, + new_forward=BertPipelineForwards.bert_for_multiple_choice_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.dropout) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +class BertForQuestionAnsweringPolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForQuestionAnswering + + policy = super().module_policy() + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForQuestionAnswering, + new_forward=BertPipelineForwards.bert_for_question_answering_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.qa_outputs) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] diff --git a/colossalai/shardformer/policies/_bert_tf5.py b/colossalai/shardformer/policies/_bert_tf5.py new file mode 100644 index 000000000000..608ae5cb7ce7 --- /dev/null +++ b/colossalai/shardformer/policies/_bert_tf5.py @@ -0,0 +1,832 @@ +import warnings +from functools import partial +from typing import Callable, Dict, List + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +import colossalai.shardformer.layer as col_nn + +from ..modeling.bert import ( + BertPipelineForwards, + bert_sequence_parallel_forward_fn, + get_bert_sequence_parallel_attention_forward, + get_jit_fused_bert_intermediate_forward, + get_jit_fused_bert_output_forward, + get_jit_fused_bert_self_output_forward, +) +from ..modeling.jit import get_jit_fused_dropout_add_func +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = [ + "BertPolicy", + "BertModelPolicy", + "BertForPreTrainingPolicy", + "BertLMHeadModelPolicy", + "BertForMaskedLMPolicy", + "BertForNextSentencePredictionPolicy", + "BertForSequenceClassificationPolicy", + "BertForTokenClassificationPolicy", + "BertForMultipleChoicePolicy", + "BertForQuestionAnsweringPolicy", +] + + +class BertPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.enable_bias_gelu_fused = self.shard_config.enable_jit_fused and self.model.config.hidden_act == "gelu" + return self.model + + def module_policy(self): + from transformers.models.bert.modeling_bert import ( + BertEmbeddings, + BertIntermediate, + BertLayer, + BertModel, + BertOutput, + BertSelfOutput, + ) + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = col_nn.VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = col_nn.PaddingEmbedding + + if self.shard_config.enable_fused_normalization: + norm_cls = col_nn.FusedLayerNorm + else: + norm_cls = col_nn.LayerNorm + + sp_mode = self.shard_config.sequence_parallelism_mode or None + assert sp_mode != "all_to_all", "all_to_all sequence parallelism is not supported for Bert" + if sp_mode == "ring": + warnings.warn( + f"For Bert, sequence parallelism is currently not support mode {sp_mode}, will set to be split_gather" + ) + sp_mode = "split_gather" + + sp_partial_derived = sp_mode == "split_gather" + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_sequence_parallelism: + # transformers v5 dropped ``BertSdpaSelfAttention``: ``modeling_bert`` now only ships + # ``BertSelfAttention`` / ``BertCrossAttention``, and self-attention vs cross-attention is + # a **class** difference rather than a constructor flag. The SP forward therefore targets + # ``BertSelfAttention`` -- the class ``BertAttention`` instantiates whenever + # ``is_cross_attention`` is False. (Cross-attention stays on v5's own forward: the policy + # does not shard ``crossattention.*`` either, so there is no gathered sequence to fix up.) + from transformers.models.bert.modeling_bert import BertSelfAttention + + # Fix the tgt_len size in bert sequence parallel attention forward. + self.append_or_create_method_replacement( + description={ + "forward": get_bert_sequence_parallel_attention_forward(self.shard_config), + }, + policy=policy, + target_key=BertSelfAttention, + ) + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + policy[BertLayer] = ModulePolicyDescription( + attribute_replacement={ + "attention.self.all_head_size": self.model.config.hidden_size + // self.shard_config.tensor_parallel_size, + "crossattention.self.all_head_size": self.model.config.hidden_size + // self.shard_config.tensor_parallel_size, + "attention.self.num_attention_heads": self.model.config.num_attention_heads + // self.shard_config.tensor_parallel_size, + "crossattention.self.num_attention_heads": self.model.config.num_attention_heads + // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.self.query", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.key", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.value", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=col_nn.Linear1D_Row, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=col_nn.Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "skip_bias_add": self.enable_bias_gelu_fused, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=col_nn.Linear1D_Row, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + ], + ) + + policy[BertEmbeddings] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + ] + ) + if self.enable_bias_gelu_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_intermediate_forward(), + }, + policy=policy, + target_key=BertIntermediate, + ) + + elif use_zbv: + policy[BertLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.self.query", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.key", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.value", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.self.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "skip_bias_add": self.enable_bias_gelu_fused, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + ], + ) + + policy[BertEmbeddings] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + ] + ) + if self.enable_bias_gelu_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_intermediate_forward(), + }, + policy=policy, + target_key=BertIntermediate, + ) + + if sp_mode == "split_gather": + self.append_or_create_method_replacement( + description={"forward": bert_sequence_parallel_forward_fn(self.shard_config)}, + policy=policy, + target_key=BertModel, + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="word_embeddings", + target_module=embedding_cls, + kwargs=( + { + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {} + ), + ) + ], + policy=policy, + target_key=BertEmbeddings, + ) + + # optimization configuration + # Handle bert layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="attention.output.LayerNorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="output.LayerNorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=BertLayer, + ) + # handle embedding layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="LayerNorm", + target_module=norm_cls, + ) + ], + policy=policy, + target_key=BertEmbeddings, + ) + + # use jit operator + if self.shard_config.enable_jit_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_self_output_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=BertSelfOutput, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_bert_output_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=BertOutput, + ) + + return policy + + def add_lm_head_policy(self, base_policy): + from transformers.models.bert.modeling_bert import BertLMPredictionHead + + # optimize for tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="decoder", + target_module=col_nn.VocabParallelLMHead1D, + kwargs={ + "gather_output": True, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ), + policy=base_policy, + target_key=BertLMPredictionHead, + ) + else: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="decoder", + target_module=col_nn.PaddingLMHead, + kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, + ), + policy=base_policy, + target_key=BertLMPredictionHead, + ) + + # optimize with fused normalization + if self.shard_config.enable_fused_normalization: + # Handle bert lm prediction head + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="transform.LayerNorm", + target_module=col_nn.FusedLayerNorm, + ), + policy=base_policy, + target_key=BertLMPredictionHead, + ) + return base_policy + + def add_lm_prediction_policy(self, base_policy): + from transformers.models.bert.modeling_bert import BertLMPredictionHead + + method_replacement = { + "_save_to_state_dict": col_nn.ParallelModule._save_to_state_dict, + "_load_from_state_dict": col_nn.ParallelModule._load_from_state_dict, + } + self.append_or_create_method_replacement( + description=method_replacement, + policy=base_policy, + target_key=BertLMPredictionHead, + ) + return base_policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """ + If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy. + """ + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "BertModel": + module = self.model + else: + module = self.model.bert + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, + stage_manager=stage_manager, + shard_config=self.shard_config, + ) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, + stage_manager=stage_manager, + stage_index=stage_index, + shard_config=self.shard_config, + ) + } + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "BertModel": + module = self.model + else: + module = self.model.bert + stage_manager = self.pipeline_stage_manager + + held_layers = [] + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embeddings) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.encoder.layer[start_idx:end_idx]) + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(module.pooler) + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + if stage_manager.is_first_stage(): + held_layers.append(module.embeddings) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.encoder.layer[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.pooler) + + return held_layers + + +# BertModel +class BertModelPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + from transformers.models.bert.modeling_bert import BertModel + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertModel, + new_forward=BertPipelineForwards.bert_model_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in bert model""" + return [] + + +# BertForPreTraining +class BertForPreTrainingPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + policy = self.add_lm_prediction_policy(policy) + from transformers.models.bert.modeling_bert import BertForPreTraining + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForPreTraining, + new_forward=BertPipelineForwards.bert_for_pretraining_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage""" + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + model = self.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if id(model.bert.embeddings.word_embeddings.weight) == id(model.cls.predictions.decoder.weight): + # tie weights + return [ + { + 0: model.bert.embeddings.word_embeddings.weight, + self.pipeline_stage_manager.num_stages - 1: model.cls.predictions.decoder.weight, + } + ] + return [] + + +# BertLMHeadModel +class BertLMHeadModelPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + policy = self.add_lm_prediction_policy(policy) + from transformers.models.bert.modeling_bert import BertLMHeadModel + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertLMHeadModel, + new_forward=BertPipelineForwards.bert_lm_head_model_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + bert_model = self.model.bert + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if id(bert_model.embeddings.word_embeddings.weight) == id(self.model.cls.predictions.decoder.weight): + # tie weights + return [ + { + 0: bert_model.embeddings.word_embeddings.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.cls.predictions.decoder.weight, + } + ] + return [] + + +# BertForMaskedLM +class BertForMaskedLMPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + policy = self.add_lm_prediction_policy(policy) + from transformers.models.bert.modeling_bert import BertForMaskedLM + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForMaskedLM, + new_forward=BertPipelineForwards.bert_for_masked_lm_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + bert_model = self.model.bert + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if id(bert_model.embeddings.word_embeddings.weight) == id(self.model.cls.predictions.decoder.weight): + # tie weights + return [ + { + 0: bert_model.embeddings.word_embeddings.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.cls.predictions.decoder.weight, + } + ] + return [] + + +# BertForSequenceClassification +class BertForSequenceClassificationPolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForSequenceClassification + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ] + ) + } + policy.update(addon_module) + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForSequenceClassification, + new_forward=BertPipelineForwards.bert_for_sequence_classification_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.dropout) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +# BertForTokenClassification +class BertForTokenClassificationPolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForTokenClassification + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForTokenClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ] + ) + } + policy.update(addon_module) + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForTokenClassification, + new_forward=BertPipelineForwards.bert_for_token_classification_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.dropout) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +# BertForNextSentencePrediction +class BertForNextSentencePredictionPolicy(BertPolicy): + def module_policy(self): + policy = super().module_policy() + from transformers.models.bert.modeling_bert import BertForNextSentencePrediction + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForNextSentencePrediction, + new_forward=BertPipelineForwards.bert_for_next_sentence_prediction_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.cls) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +# BertForMultipleChoice +class BertForMultipleChoicePolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForMultipleChoice + + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + addon_module = { + BertForMultipleChoice: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=col_nn.DropoutForParallelInput, + ) + ] + ) + } + policy.update(addon_module) + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForMultipleChoice, + new_forward=BertPipelineForwards.bert_for_multiple_choice_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.dropout) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] + + +class BertForQuestionAnsweringPolicy(BertPolicy): + def module_policy(self): + from transformers.models.bert.modeling_bert import BertForQuestionAnswering + + policy = super().module_policy() + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=BertForQuestionAnswering, + new_forward=BertPipelineForwards.bert_for_question_answering_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """ + get pipeline layers for current stage + """ + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.qa_outputs) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + # no shared params for sequence classification model + return [] diff --git a/colossalai/shardformer/policies/_deepseek_v3_tf4.py b/colossalai/shardformer/policies/_deepseek_v3_tf4.py new file mode 100644 index 000000000000..8b386bcaa6e8 --- /dev/null +++ b/colossalai/shardformer/policies/_deepseek_v3_tf4.py @@ -0,0 +1,164 @@ +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn + +from colossalai.shardformer.layer import FusedRMSNorm +from colossalai.shardformer.modeling.deepseek_v3 import ( + EpDeepseekV3MoE, + deepseek_v3_for_causal_lm_forward, + deepseek_v3_model_forward, +) +from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["DeepseekPolicy", "DeepseekForCausalLMPolicy"] + + +class DeepseekV3Policy(Policy): + def config_sanity_check(self): + assert not self.shard_config.enable_tensor_parallelism, "DeepSeekV3 does not support tensor parallelism" + assert not self.shard_config.enable_sequence_parallelism, "DeepSeekV3 does not support sequence parallelism" + if self.shard_config.pipeline_stage_manager: + assert not self.shard_config.pipeline_stage_manager.use_zbv, "DeepSeekV3 does not support ZBV" + + def preprocess(self): + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + + policy = {} + + # support gradient checkpointing + if self.shard_config.pipeline_stage_manager is None: + policy["DeepseekV3Model"] = ModulePolicyDescription( + method_replacement={"forward": deepseek_v3_model_forward} + ) + + if self.shard_config.expert_parallel_size > 1: + # expert parallel + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="mlp", + target_module=EpDeepseekV3MoE, + kwargs={ + "ep_group": self.shard_config.ep_group, + "moe_dp_group": self.shard_config.moe_dp_group, + }, + ) + ], + policy=policy, + target_key="DeepseekV3DecoderLayer", + ) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + # TODO: prevent casting to fp32 + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=FusedRMSNorm, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=FusedRMSNorm, + ), + ], + policy=policy, + target_key="DeepseekV3DecoderLayer", + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=FusedRMSNorm, + ), + policy=policy, + target_key="DeepseekV3Model", + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: str, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager: + num_layers = self.model.config.num_hidden_layers + stage_manager = self.pipeline_stage_manager + + layers_per_stage = stage_manager.distribute_layers(num_layers) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)} + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + return + + def get_held_layers(self) -> List[nn.Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + module = self.model + if module.__class__.__name__.startswith("PeftModel"): + module = module.get_base_model() + if module.__class__.__name__ != "DeepseekV3Model": + module = module.model + + stage_manager = self.pipeline_stage_manager + + held_layers = [] + + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + stage_manager.stage_indices = stage_indices + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + # for zbv, when is_first_stage (last fwd), we append norm + # for interleaved, when is_last_stage (last fwd), we also append norm + held_layers.append(module.norm) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + return held_layers + + +class DeepseekV3ModelPolicy(DeepseekV3Policy): + def module_policy(self): + policy = super().module_policy() + if self.shard_config.pipeline_stage_manager: + self.set_pipeline_forward("DeepseekV3Model", deepseek_v3_model_forward, policy) + return policy + + +class DeepseekV3ForCausalLMPolicy(DeepseekV3Policy): + def module_policy(self): + policy = super().module_policy() + if self.shard_config.pipeline_stage_manager: + self.set_pipeline_forward("DeepseekV3ForCausalLM", deepseek_v3_for_causal_lm_forward, policy) + return policy + + def get_held_layers(self): + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + elif stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers diff --git a/colossalai/shardformer/policies/_deepseek_v3_tf5.py b/colossalai/shardformer/policies/_deepseek_v3_tf5.py new file mode 100644 index 000000000000..f6f9361fddee --- /dev/null +++ b/colossalai/shardformer/policies/_deepseek_v3_tf5.py @@ -0,0 +1,171 @@ +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn + +from colossalai.shardformer.layer import FusedRMSNorm +from colossalai.shardformer.modeling.deepseek_v3 import ( + EpDeepseekV3MoE, + deepseek_v3_for_causal_lm_forward, + deepseek_v3_model_forward, +) +from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["DeepseekPolicy", "DeepseekForCausalLMPolicy"] + + +class DeepseekV3Policy(Policy): + def config_sanity_check(self): + assert not self.shard_config.enable_tensor_parallelism, "DeepSeekV3 does not support tensor parallelism" + assert not self.shard_config.enable_sequence_parallelism, "DeepSeekV3 does not support sequence parallelism" + if self.shard_config.pipeline_stage_manager: + assert not self.shard_config.pipeline_stage_manager.use_zbv, "DeepSeekV3 does not support ZBV" + + def preprocess(self): + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + + policy = {} + + # support gradient checkpointing + if self.shard_config.pipeline_stage_manager is None: + policy["DeepseekV3Model"] = ModulePolicyDescription( + method_replacement={"forward": deepseek_v3_model_forward} + ) + + if self.shard_config.expert_parallel_size > 1: + # expert parallel + # The v5 EP path uses three primitives (see docs/30 §4): slice fused expert parameters, + # mark their EP loading group, and run each local expert with fused matmuls. + # Unlike Mixtral, this policy replaces the whole `mlp`; dense layers before + # `first_k_dense_replace` are skipped by `from_native_module`. + # Pure EP (tp=1) is implemented; 4-GPU acceptance is pending. + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="mlp", + target_module=EpDeepseekV3MoE, + kwargs={ + "ep_group": self.shard_config.ep_group, + "moe_dp_group": self.shard_config.moe_dp_group, + }, + ) + ], + policy=policy, + target_key="DeepseekV3DecoderLayer", + ) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + # TODO: prevent casting to fp32 + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=FusedRMSNorm, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=FusedRMSNorm, + ), + ], + policy=policy, + target_key="DeepseekV3DecoderLayer", + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=FusedRMSNorm, + ), + policy=policy, + target_key="DeepseekV3Model", + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: str, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager: + num_layers = self.model.config.num_hidden_layers + stage_manager = self.pipeline_stage_manager + + layers_per_stage = stage_manager.distribute_layers(num_layers) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)} + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + return + + def get_held_layers(self) -> List[nn.Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + module = self.model + if module.__class__.__name__.startswith("PeftModel"): + module = module.get_base_model() + if module.__class__.__name__ != "DeepseekV3Model": + module = module.model + + stage_manager = self.pipeline_stage_manager + + # Every PP stage constructs its position embeddings from this buffer. + # Leaving it unheld lets the sharder release rotary_emb.inv_freq. + held_layers = [module.rotary_emb] + + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + stage_manager.stage_indices = stage_indices + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + # for zbv, when is_first_stage (last fwd), we append norm + # for interleaved, when is_last_stage (last fwd), we also append norm + held_layers.append(module.norm) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + return held_layers + + +class DeepseekV3ModelPolicy(DeepseekV3Policy): + def module_policy(self): + policy = super().module_policy() + if self.shard_config.pipeline_stage_manager: + self.set_pipeline_forward("DeepseekV3Model", deepseek_v3_model_forward, policy) + return policy + + +class DeepseekV3ForCausalLMPolicy(DeepseekV3Policy): + def module_policy(self): + policy = super().module_policy() + if self.shard_config.pipeline_stage_manager: + self.set_pipeline_forward("DeepseekV3ForCausalLM", deepseek_v3_for_causal_lm_forward, policy) + return policy + + def get_held_layers(self): + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + elif stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers diff --git a/colossalai/shardformer/policies/_llama_tf4.py b/colossalai/shardformer/policies/_llama_tf4.py new file mode 100644 index 000000000000..9ad63dd7f051 --- /dev/null +++ b/colossalai/shardformer/policies/_llama_tf4.py @@ -0,0 +1,538 @@ +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + PaddingLMHead, + RMSNorm, + VocabParallelEmbedding1D, + VocabParallelLMHead1D, +) + +from ..modeling.llama import LlamaPipelineForwards, get_llama_flash_attention_forward +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["LlamaPolicy", "LlamaForCausalLMPolicy", "LlamaForSequenceClassificationPolicy"] + + +class LlamaPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel + + policy = {} + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + + if self.shard_config.enable_fused_normalization: + norm_cls = FusedRMSNorm + else: + norm_cls = RMSNorm + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + if sp_mode == "ring_attn" and not self.is_causal: + raise ValueError("Ring attention is only meant for causal language modeling.") + + tp_size = self.shard_config.tensor_parallel_size + # Modified by SP and TP + num_q_heads = self.model.config.num_attention_heads + num_kv_heads = getattr(self.model.config, "num_key_value_heads", None) + + if sp_mode == "all_to_all": + num_q_heads //= sp_size + decoder_attribute_replacement = {"num_heads": num_q_heads} + if num_kv_heads: + num_kv_heads //= sp_size + decoder_attribute_replacement["num_key_value_heads"] = num_kv_heads + + policy[LlamaAttention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: + self.append_or_create_method_replacement( + description={ + "forward": get_llama_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=LlamaAttention, + ) + + if self.pipeline_stage_manager is None: + self.append_or_create_method_replacement( + description={ + "forward": partial( + LlamaPipelineForwards.llama_model_forward, + shard_config=self.shard_config, + ), + }, + policy=policy, + target_key=LlamaModel, + ) + # enable tp, replace layer to tp Linear1D_Col,Linear1D_Row, + if self.shard_config.enable_tensor_parallelism: + assert ( + num_q_heads % tp_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + if hasattr(self.model.config, "num_key_value_heads"): + assert ( + num_kv_heads >= tp_size and num_kv_heads % tp_size == 0 + ), f"The number of key_value heads must be divisible by, and must not be less than tensor parallel size." + num_q_heads //= tp_size + decoder_attribute_replacement = { + "self_attn.hidden_size": self.model.config.hidden_size // tp_size, + "self_attn.num_heads": num_q_heads, + } + if getattr(self.model.config, "num_key_value_heads", False): + num_kv_heads //= tp_size + decoder_attribute_replacement["self_attn.num_key_value_heads"] = num_kv_heads + + policy[LlamaDecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + # not enable tp, replace layer to LinearWithGradAccum + elif use_zbv: + policy[LlamaDecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=LlamaModel, + ) + + # optimization configuration + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=LlamaDecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=LlamaModel, + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "LlamaModel": + module = self.model + else: + module = self.model.model + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config + ) + } + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "LlamaModel": + module = self.model + else: + module = self.model.model + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.norm) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + + return held_layers + + +class LlamaModelPolicy(LlamaPolicy): + def module_policy(self): + policy = super().module_policy() + from transformers.models.llama.modeling_llama import LlamaModel + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=LlamaModel, new_forward=LlamaPipelineForwards.llama_model_forward, policy=policy + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in llama model""" + return [] + + +class LlamaForCausalLMPolicy(LlamaPolicy): + def module_policy(self): + from transformers import LlamaForCausalLM + + self.is_causal = True + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + # add a new item for causal lm + new_item = { + LlamaForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=VocabParallelLMHead1D, + kwargs={ + "gather_output": not self.shard_config.parallel_output, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ) + ], + ) + } + else: + new_item = { + LlamaForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=PaddingLMHead, + kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, + ) + ], + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=LlamaForCausalLM, new_forward=LlamaPipelineForwards.llama_for_causal_lm_forward, policy=policy + ) + elif self.shard_config.enable_tensor_parallelism or self.shard_config.enable_sequence_parallelism: + # Compute loss distributedly along the sequence dimension + new_item[LlamaForCausalLM].method_replacement = { + # "forward": get_lm_forward_with_dist_cross_entropy(self.shard_config) + "forward": partial(LlamaPipelineForwards.llama_for_causal_lm_forward, shard_config=self.shard_config) + } + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + if self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv: + return [] + llama_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(llama_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: llama_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class LlamaForSequenceClassificationPolicy(LlamaPolicy): + def module_policy(self): + from transformers import LlamaForSequenceClassification + + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + # enable tp, replace layer to tp Linear1D_Col,Linear1D_Row, + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + LlamaForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + # enable tp, replace layer to LinearWithGradAccum + elif use_zbv: + # add a new item for sequence classification + new_item = { + LlamaForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=LinearWithGradAccum, + kwargs=dict( + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + + # to be confirmed + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=LlamaForSequenceClassification, + new_forward=LlamaPipelineForwards.llama_for_sequence_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in llama for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_llama_tf5.py b/colossalai/shardformer/policies/_llama_tf5.py new file mode 100644 index 000000000000..a686c3ea8c07 --- /dev/null +++ b/colossalai/shardformer/policies/_llama_tf5.py @@ -0,0 +1,536 @@ +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + PaddingLMHead, + RMSNorm, + VocabParallelEmbedding1D, + VocabParallelLMHead1D, +) + +from ..modeling.llama import LlamaPipelineForwards, get_llama_flash_attention_forward +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["LlamaPolicy", "LlamaForCausalLMPolicy", "LlamaForSequenceClassificationPolicy"] + + +class LlamaPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel + + policy = {} + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + + if self.shard_config.enable_fused_normalization: + norm_cls = FusedRMSNorm + else: + norm_cls = RMSNorm + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + if sp_mode == "ring_attn" and not self.is_causal: + raise ValueError("Ring attention is only meant for causal language modeling.") + + tp_size = self.shard_config.tensor_parallel_size + # Modified by SP and TP + num_q_heads = self.model.config.num_attention_heads + num_kv_heads = getattr(self.model.config, "num_key_value_heads", None) + + if sp_mode == "all_to_all": + num_q_heads //= sp_size + decoder_attribute_replacement = {"num_heads": num_q_heads} + if num_kv_heads: + num_kv_heads //= sp_size + decoder_attribute_replacement["num_key_value_heads"] = num_kv_heads + + policy[LlamaAttention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: + self.append_or_create_method_replacement( + description={ + "forward": get_llama_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=LlamaAttention, + ) + + if self.pipeline_stage_manager is None: + self.append_or_create_method_replacement( + description={ + "forward": partial( + LlamaPipelineForwards.llama_model_forward, + shard_config=self.shard_config, + ), + }, + policy=policy, + target_key=LlamaModel, + ) + # enable tp, replace layer to tp Linear1D_Col,Linear1D_Row, + if self.shard_config.enable_tensor_parallelism: + assert ( + num_q_heads % tp_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + if hasattr(self.model.config, "num_key_value_heads"): + assert ( + num_kv_heads >= tp_size and num_kv_heads % tp_size == 0 + ), f"The number of key_value heads must be divisible by, and must not be less than tensor parallel size." + num_q_heads //= tp_size + # v5 reads `hidden_size` / `num_key_value_heads` off `self.config`, so only the local + # head count is written (the flash-attention forward in `modeling/llama.py` uses it). + decoder_attribute_replacement = { + "self_attn.num_heads": num_q_heads, + } + + policy[LlamaDecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + # not enable tp, replace layer to LinearWithGradAccum + elif use_zbv: + policy[LlamaDecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=LlamaModel, + ) + + # optimization configuration + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=LlamaDecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=LlamaModel, + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "LlamaModel": + module = self.model + else: + module = self.model.model + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config + ) + } + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "LlamaModel": + module = self.model + else: + module = self.model.model + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.norm) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + + return held_layers + + +class LlamaModelPolicy(LlamaPolicy): + def module_policy(self): + policy = super().module_policy() + from transformers.models.llama.modeling_llama import LlamaModel + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=LlamaModel, new_forward=LlamaPipelineForwards.llama_model_forward, policy=policy + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in llama model""" + return [] + + +class LlamaForCausalLMPolicy(LlamaPolicy): + def module_policy(self): + from transformers import LlamaForCausalLM + + self.is_causal = True + policy = super().module_policy() + + if self.shard_config.enable_tensor_parallelism: + # add a new item for causal lm + new_item = { + LlamaForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=VocabParallelLMHead1D, + kwargs={ + "gather_output": not self.shard_config.parallel_output, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ) + ], + ) + } + else: + new_item = { + LlamaForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=PaddingLMHead, + kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, + ) + ], + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=LlamaForCausalLM, new_forward=LlamaPipelineForwards.llama_for_causal_lm_forward, policy=policy + ) + elif self.shard_config.enable_tensor_parallelism or self.shard_config.enable_sequence_parallelism: + # Compute loss distributedly along the sequence dimension + new_item[LlamaForCausalLM].method_replacement = { + # "forward": get_lm_forward_with_dist_cross_entropy(self.shard_config) + "forward": partial(LlamaPipelineForwards.llama_for_causal_lm_forward, shard_config=self.shard_config) + } + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + if self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv: + return [] + llama_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(llama_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: llama_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class LlamaForSequenceClassificationPolicy(LlamaPolicy): + def module_policy(self): + from transformers import LlamaForSequenceClassification + + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + # enable tp, replace layer to tp Linear1D_Col,Linear1D_Row, + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + LlamaForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + # enable tp, replace layer to LinearWithGradAccum + elif use_zbv: + # add a new item for sequence classification + new_item = { + LlamaForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=LinearWithGradAccum, + kwargs=dict( + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + + # to be confirmed + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=LlamaForSequenceClassification, + new_forward=LlamaPipelineForwards.llama_for_sequence_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in llama for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_mixtral_tf4.py b/colossalai/shardformer/policies/_mixtral_tf4.py new file mode 100644 index 000000000000..a9584db9bc52 --- /dev/null +++ b/colossalai/shardformer/policies/_mixtral_tf4.py @@ -0,0 +1,491 @@ +import warnings +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module +from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM, MixtralModel + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + VocabParallelEmbedding1D, +) + +# from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col +# from colossalai.shardformer.layer.embedding import PaddingEmbedding, VocabParallelEmbedding1D +# from colossalai.shardformer.layer.linear import Linear1D_Row +from colossalai.shardformer.modeling.mixtral import ( + EPMixtralSparseMoeBlock, + MixtralPipelineForwards, + get_mixtral_flash_attention_forward, + get_mixtral_flash_attention_model_forward, +) +from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["MixtralPolicy", "MixtralForCausalLMPolicy"] + + +class MixtralPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.mixtral.modeling_mixtral import MixtralAttention, MixtralDecoderLayer, MixtralModel + + policy = {} + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + tp_size = self.shard_config.tensor_parallel_size + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + # modified for both SP and TP + num_q_heads = self.model.config.num_attention_heads + num_kv_heads = getattr(self.model.config, "num_key_value_heads", None) + + if sp_mode == "all_to_all": + num_q_heads //= sp_size + decoder_attribute_replacement = { + "num_heads": num_q_heads, + } + if getattr(self.model.config, "num_key_value_heads", False): + num_kv_heads //= sp_size + decoder_attribute_replacement["num_key_value_heads"] = num_kv_heads + + policy[MixtralAttention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + if self.shard_config.enable_sequence_parallelism: + if self.pipeline_stage_manager is not None: + # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism + # if both are enabled, one of them will be ignored + raise NotImplementedError("Sequence parallelism is not supported with pipeline parallelism.") + self.append_or_create_method_replacement( + description={ + "forward": get_mixtral_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=MixtralAttention, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_mixtral_flash_attention_model_forward( + self.shard_config, + sp_mode=sp_mode, + sp_size=sp_size, + sp_group=sp_group, + ), + }, + policy=policy, + target_key=MixtralModel, + ) + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + + if self.shard_config.enable_tensor_parallelism: + # tensor parallelism for non-moe params + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + assert ( + self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of key_value heads must be divisible by tensor parallel size." + num_q_heads //= tp_size + decoder_attribute_replacement = { + "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": num_q_heads, + } + if num_kv_heads: + num_kv_heads //= tp_size + decoder_attribute_replacement["self_attn.num_key_value_heads"] = num_kv_heads + + policy[MixtralDecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="block_sparse_moe.gate", + target_module=Linear1D_Col, + kwargs={ + "gather_output": True, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + elif use_zbv: + policy[MixtralDecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="block_sparse_moe.gate", + target_module=LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=MixtralModel, + ) + + if self.shard_config.ep_group: + # expert parallel + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="block_sparse_moe", + target_module=EPMixtralSparseMoeBlock, + kwargs={ + "ep_group": self.shard_config.ep_group, + "tp_group": self.shard_config.tensor_parallel_process_group, + "moe_dp_group": self.shard_config.moe_dp_group, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ) + ], + policy=policy, + target_key=MixtralDecoderLayer, + ) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=FusedRMSNorm, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=FusedRMSNorm, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=MixtralDecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=FusedRMSNorm, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=MixtralModel, + ) + + if self.shard_config.enable_flash_attention: + warnings.warn("Flash attention is natively supported in transformers, will ignore the flag.") + self.shard_config.enable_flash_attention = False + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager: + if self.shard_config.enable_sequence_parallelism: + # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism + # if both are enabled, one of them will be ignored + raise NotImplementedError("Pipeline parallelism is not supported with sequence parallelism.") + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "MixtralModel": + module = self.model + else: + module = self.model.model + + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)} + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + return + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "MixtralModel": + module = self.model + else: + module = self.model.model + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + stage_manager.stage_indices = stage_indices + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + # for zbv, when is_first_stage (last fwd), we append norm + # for interleaved, when is_last_stage (last fwd), we also append norm + held_layers.append(module.norm) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + return held_layers + + +class MixtralModelPolicy(MixtralPolicy): + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + policy = super().module_policy() + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=MixtralModel, + new_forward=MixtralPipelineForwards.mixtral_model_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in mixtral model""" + return [] + + +class MixtralForCausalLMPolicy(MixtralPolicy): + def module_policy(self): + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + # TODO: assign pg mesh from plugin to all modules + if self.shard_config.enable_tensor_parallelism: + # add a new item for causal lm + new_item = { + MixtralForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ], + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + MixtralForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=LinearWithGradAccum, + kwargs=dict( + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ], + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=MixtralForCausalLM, + new_forward=MixtralPipelineForwards.mixtral_for_causal_lm_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + elif stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + mixtral_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(mixtral_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: mixtral_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class MixtralForSequenceClassificationPolicy(MixtralPolicy): + def module_policy(self): + from transformers import MixtralForSequenceClassification + + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + MixtralForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + raise NotImplementedError + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in mixtral for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_mixtral_tf5.py b/colossalai/shardformer/policies/_mixtral_tf5.py new file mode 100644 index 000000000000..fc24d6132b7e --- /dev/null +++ b/colossalai/shardformer/policies/_mixtral_tf5.py @@ -0,0 +1,518 @@ +import warnings +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module +from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM, MixtralModel + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + VocabParallelEmbedding1D, +) + +# from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col +# from colossalai.shardformer.layer.embedding import PaddingEmbedding, VocabParallelEmbedding1D +# from colossalai.shardformer.layer.linear import Linear1D_Row +from colossalai.shardformer.modeling.mixtral import ( + EPMixtralSparseMoeBlock, + MixtralPipelineForwards, + MixtralTopKRouter1D, + MixtralTopKRouterWithGradAccum, + get_mixtral_flash_attention_forward, + get_mixtral_flash_attention_model_forward, +) +from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["MixtralPolicy", "MixtralForCausalLMPolicy"] + + +class MixtralPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.mixtral.modeling_mixtral import MixtralAttention, MixtralDecoderLayer, MixtralModel + + policy = {} + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + if sp_partial_derived: + # These modes keep the sequence split across the tp group, so the ranks of a tp group dispatch + # different tokens to the experts, while the tp-sharded experts reduce their outputs over that group. + raise NotImplementedError(f"Mixtral does not support sequence parallelism mode {sp_mode} yet.") + tp_size = self.shard_config.tensor_parallel_size + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + # modified for both SP and TP + num_q_heads = self.model.config.num_attention_heads + num_kv_heads = getattr(self.model.config, "num_key_value_heads", None) + + if sp_mode == "all_to_all": + num_q_heads //= sp_size + decoder_attribute_replacement = { + "num_heads": num_q_heads, + } + if getattr(self.model.config, "num_key_value_heads", False): + num_kv_heads //= sp_size + decoder_attribute_replacement["num_key_value_heads"] = num_kv_heads + + policy[MixtralAttention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + if self.shard_config.enable_sequence_parallelism: + if self.pipeline_stage_manager is not None: + # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism + # if both are enabled, one of them will be ignored + raise NotImplementedError("Sequence parallelism is not supported with pipeline parallelism.") + self.append_or_create_method_replacement( + description={ + "forward": get_mixtral_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=MixtralAttention, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_mixtral_flash_attention_model_forward( + self.shard_config, + sp_mode=sp_mode, + sp_size=sp_size, + sp_group=sp_group, + ), + }, + policy=policy, + target_key=MixtralModel, + ) + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + + if self.shard_config.enable_tensor_parallelism: + # tensor parallelism for non-moe params + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + assert ( + self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of key_value heads must be divisible by tensor parallel size." + num_q_heads //= tp_size + decoder_attribute_replacement = { + "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": num_q_heads, + } + if num_kv_heads: + num_kv_heads //= tp_size + decoder_attribute_replacement["self_attn.num_key_value_heads"] = num_kv_heads + + policy[MixtralDecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + # The SP attention forward multiplies `q_len` by `sp_size` and relies on the + # q/k/v projections gathering the full sequence (o_proj splitting it back), so + # they must carry `seq_parallel_mode` -- mirroring BERT's SP policy. + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + # v5 renamed this submodule from `block_sparse_moe` to `mlp`. + suffix="mlp.gate", + target_module=MixtralTopKRouter1D, + kwargs={ + "gather_output": True, + # SP shards the sequence, so TP gathering would pair local token j from + # different ranks. Keep the small router replicated and compute full logits + # locally in SP mode (see the early return in modeling). + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + elif use_zbv: + policy[MixtralDecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs={ + "seq_parallel_mode": sp_mode, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + # v5 renamed this submodule from `block_sparse_moe` to `mlp`; the gate returns + # a three-tuple, so `LinearWithGradAccum` cannot replace it directly. + suffix="mlp.gate", + target_module=MixtralTopKRouterWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=MixtralModel, + ) + + if self.shard_config.ep_group: + # expert parallel + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + # v5 renamed `block_sparse_moe` to `mlp` and fused `mlp.experts` into 3D + # `gate_up_proj` / `down_proj` parameters. The three EP primitives in + # `EPMixtralSparseMoeBlock` are implemented and pure EP is validated. + suffix="mlp", + target_module=EPMixtralSparseMoeBlock, + kwargs={ + "ep_group": self.shard_config.ep_group, + "tp_group": self.shard_config.tensor_parallel_process_group, + "moe_dp_group": self.shard_config.moe_dp_group, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ) + ], + policy=policy, + target_key=MixtralDecoderLayer, + ) + + # optimization configuration + if self.shard_config.enable_fused_normalization: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=FusedRMSNorm, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=FusedRMSNorm, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=MixtralDecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=FusedRMSNorm, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=MixtralModel, + ) + + if self.shard_config.enable_flash_attention: + warnings.warn("Flash attention is natively supported in transformers, will ignore the flag.") + self.shard_config.enable_flash_attention = False + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager: + if self.shard_config.enable_sequence_parallelism: + # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism + # if both are enabled, one of them will be ignored + raise NotImplementedError("Pipeline parallelism is not supported with sequence parallelism.") + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "MixtralModel": + module = self.model + else: + module = self.model.model + + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)} + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + return + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "MixtralModel": + module = self.model + else: + module = self.model.model + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + stage_manager.stage_indices = stage_indices + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + # for zbv, when is_first_stage (last fwd), we append norm + # for interleaved, when is_last_stage (last fwd), we also append norm + held_layers.append(module.norm) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + return held_layers + + +class MixtralModelPolicy(MixtralPolicy): + def __init__(self) -> None: + super().__init__() + + def module_policy(self): + policy = super().module_policy() + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=MixtralModel, + new_forward=MixtralPipelineForwards.mixtral_model_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in mixtral model""" + return [] + + +class MixtralForCausalLMPolicy(MixtralPolicy): + def module_policy(self): + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + # TODO: assign pg mesh from plugin to all modules + if self.shard_config.enable_tensor_parallelism: + # add a new item for causal lm + new_item = { + MixtralForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ], + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + MixtralForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=LinearWithGradAccum, + kwargs=dict( + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ], + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=MixtralForCausalLM, + new_forward=MixtralPipelineForwards.mixtral_for_causal_lm_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + elif stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + mixtral_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(mixtral_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: mixtral_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class MixtralForSequenceClassificationPolicy(MixtralPolicy): + def module_policy(self): + from transformers import MixtralForSequenceClassification + + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + MixtralForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + raise NotImplementedError + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in mixtral for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_qwen2_tf4.py b/colossalai/shardformer/policies/_qwen2_tf4.py new file mode 100644 index 000000000000..32f7d37b190f --- /dev/null +++ b/colossalai/shardformer/policies/_qwen2_tf4.py @@ -0,0 +1,556 @@ +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2Attention, + Qwen2DecoderLayer, + Qwen2ForCausalLM, + Qwen2ForSequenceClassification, + Qwen2Model, +) + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + RMSNorm, + VocabParallelEmbedding1D, + VocabParallelLMHead1D, +) + +from ..modeling.qwen2 import ( + Qwen2PipelineForwards, + get_lm_forward_with_dist_cross_entropy, + get_qwen2_flash_attention_forward, + get_qwen2_model_forward_for_flash_attn, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["Qwen2Policy", "Qwen2ForCausalLMPolicy", "Qwen2ForSequenceClassificationPolicy"] + + +class Qwen2Policy(Policy): + def __init__(self) -> None: + super().__init__() + import transformers + from packaging.version import Version + + assert Version(transformers.__version__) >= Version( + "4.39.1" + ), "The Qwen2 model should run on a transformers version of 4.39.1." + + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + if sp_mode == "all_to_all": + decoder_attribute_replacement = { + "num_heads": self.model.config.num_attention_heads // sp_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size + + policy[Qwen2Attention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + if hasattr(self.model.config, "num_key_value_heads"): + assert ( + self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of key_value heads must be divisible by tensor parallel size." + decoder_attribute_replacement = { + "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["self_attn.num_key_value_heads"] = ( + self.model.config.num_key_value_heads // self.shard_config.tensor_parallel_size + ) + + policy[Qwen2DecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + elif use_zbv: + policy[Qwen2DecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=Qwen2Model, + ) + + # optimization configuration + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=Qwen2DecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=Qwen2Model, + ) + + if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: + self.append_or_create_method_replacement( + description={ + "forward": get_qwen2_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=Qwen2Attention, + ) + if self.pipeline_stage_manager is None: + # replace qwen2 model forward method + self.append_or_create_method_replacement( + description={ + "forward": get_qwen2_model_forward_for_flash_attn( + self.shard_config, sp_mode, sp_size, sp_group + ), + }, + policy=policy, + target_key=Qwen2Model, + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "Qwen2Model": + module = self.model + else: + module = self.model.model + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config + ) + } + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "Qwen2Model": + module = self.model + else: + module = self.model.model + + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.norm) + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + + return held_layers + + +class Qwen2ModelPolicy(Qwen2Policy): + def module_policy(self): + policy = super().module_policy() + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen2Model, new_forward=Qwen2PipelineForwards.qwen2_model_forward, policy=policy + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen2 model""" + return [] + + +class Qwen2ForCausalLMPolicy(Qwen2Policy): + def module_policy(self): + policy = super().module_policy() + setattr(self.shard_config, "causal_lm", True) + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + # add a new item for casual lm + new_item = { + Qwen2ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=VocabParallelLMHead1D, + kwargs=dict( + gather_output=not self.shard_config.parallel_output, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + elif use_zbv: + # add a new item for casual lm + new_item = { + Qwen2ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=LinearWithGradAccum, + kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), + ), + SubModuleReplacementDescription( + suffix="lm_head", + target_module=VocabParallelLMHead1D, + kwargs={ + "gather_output": not self.shard_config.parallel_output, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ), + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen2ForCausalLM, new_forward=Qwen2PipelineForwards.qwen2_for_causal_lm_forward, policy=policy + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.lm_head) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + qwen2_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(qwen2_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: qwen2_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class Qwen2ForSequenceClassificationPolicy(Qwen2Policy): + def module_policy(self): + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + Qwen2ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + Qwen2ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=LinearWithGradAccum, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + # to be confirmed + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen2ForSequenceClassification, + new_forward=Qwen2PipelineForwards.qwen2_for_sequence_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.score) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen2 for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_qwen2_tf5.py b/colossalai/shardformer/policies/_qwen2_tf5.py new file mode 100644 index 000000000000..3c25628a19bf --- /dev/null +++ b/colossalai/shardformer/policies/_qwen2_tf5.py @@ -0,0 +1,555 @@ +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module +from transformers.models.qwen2.modeling_qwen2 import ( + Qwen2Attention, + Qwen2DecoderLayer, + Qwen2ForCausalLM, + Qwen2ForSequenceClassification, + Qwen2Model, +) + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + RMSNorm, + VocabParallelEmbedding1D, + VocabParallelLMHead1D, +) + +from ..modeling.qwen2 import ( + Qwen2PipelineForwards, + get_lm_forward_with_dist_cross_entropy, + get_qwen2_flash_attention_forward, + get_qwen2_model_forward_for_flash_attn, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["Qwen2Policy", "Qwen2ForCausalLMPolicy", "Qwen2ForSequenceClassificationPolicy"] + + +class Qwen2Policy(Policy): + def __init__(self) -> None: + super().__init__() + import transformers + from packaging.version import Version + + assert Version(transformers.__version__) >= Version( + "5.17.0" + ), "The Qwen2 model should run on a transformers version of 5.17.0." + + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + if sp_mode == "all_to_all": + decoder_attribute_replacement = { + "num_heads": self.model.config.num_attention_heads // sp_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size + + policy[Qwen2Attention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + if hasattr(self.model.config, "num_key_value_heads"): + assert ( + self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of key_value heads must be divisible by tensor parallel size." + decoder_attribute_replacement = { + "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["self_attn.num_key_value_heads"] = ( + self.model.config.num_key_value_heads // self.shard_config.tensor_parallel_size + ) + + policy[Qwen2DecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + elif use_zbv: + policy[Qwen2DecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=Qwen2Model, + ) + + # optimization configuration + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=Qwen2DecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=Qwen2Model, + ) + + if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: + self.append_or_create_method_replacement( + description={ + "forward": get_qwen2_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=Qwen2Attention, + ) + if self.pipeline_stage_manager is None: + # replace qwen2 model forward method + self.append_or_create_method_replacement( + description={ + "forward": get_qwen2_model_forward_for_flash_attn( + self.shard_config, sp_mode, sp_size, sp_group + ), + }, + policy=policy, + target_key=Qwen2Model, + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "Qwen2Model": + module = self.model + else: + module = self.model.model + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config + ) + } + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "Qwen2Model": + module = self.model + else: + module = self.model.model + + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.norm) + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + + return held_layers + + +class Qwen2ModelPolicy(Qwen2Policy): + def module_policy(self): + policy = super().module_policy() + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen2Model, new_forward=Qwen2PipelineForwards.qwen2_model_forward, policy=policy + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen2 model""" + return [] + + +class Qwen2ForCausalLMPolicy(Qwen2Policy): + def module_policy(self): + policy = super().module_policy() + setattr(self.shard_config, "causal_lm", True) + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + # add a new item for casual lm + new_item = { + Qwen2ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=VocabParallelLMHead1D, + kwargs=dict( + gather_output=not self.shard_config.parallel_output, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + elif use_zbv: + # add a new item for casual lm + new_item = { + Qwen2ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=LinearWithGradAccum, + kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), + ), + SubModuleReplacementDescription( + suffix="lm_head", + target_module=VocabParallelLMHead1D, + kwargs={ + "gather_output": not self.shard_config.parallel_output, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ), + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen2ForCausalLM, new_forward=Qwen2PipelineForwards.qwen2_for_causal_lm_forward, policy=policy + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.lm_head) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + qwen2_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(qwen2_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: qwen2_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class Qwen2ForSequenceClassificationPolicy(Qwen2Policy): + def module_policy(self): + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + Qwen2ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + Qwen2ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=LinearWithGradAccum, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + # to be confirmed + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen2ForSequenceClassification, + new_forward=Qwen2PipelineForwards.qwen2_for_sequence_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.score) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen2 for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_qwen3_tf4.py b/colossalai/shardformer/policies/_qwen3_tf4.py new file mode 100644 index 000000000000..e9cc9543278a --- /dev/null +++ b/colossalai/shardformer/policies/_qwen3_tf4.py @@ -0,0 +1,541 @@ +# Modifed from qwen2 policy +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.nn as nn +from torch import Tensor +from torch.nn import Module +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3Attention, + Qwen3DecoderLayer, + Qwen3ForCausalLM, + Qwen3ForSequenceClassification, + Qwen3Model, +) + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + RMSNorm, + VocabParallelEmbedding1D, +) + +from ..modeling.qwen3 import ( + Qwen3PipelineForwards, + get_lm_forward_with_dist_cross_entropy, + get_qwen3_flash_attention_forward, + get_qwen3_model_forward_for_flash_attn, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["Qwen3Policy", "Qwen3ForCausalLMPolicy", "Qwen3ForSequenceClassificationPolicy"] + + +class Qwen3Policy(Policy): + def __init__(self) -> None: + super().__init__() + import transformers + from packaging.version import Version + + assert Version(transformers.__version__) >= Version( + "4.51.0" + ), "The Qwen3 model should run on a transformers version of 4.51.0 or higher." + + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + if sp_mode == "all_to_all": + decoder_attribute_replacement = { + "num_heads": self.model.config.num_attention_heads // sp_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size + + policy[Qwen3Attention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + if hasattr(self.model.config, "num_key_value_heads"): + assert ( + self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of key_value heads must be divisible by tensor parallel size." + decoder_attribute_replacement = { + "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["self_attn.num_key_value_heads"] = ( + self.model.config.num_key_value_heads // self.shard_config.tensor_parallel_size + ) + + policy[Qwen3DecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + elif use_zbv: + policy[Qwen3DecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=Qwen3Model, + ) + + # optimization configuration + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=Qwen3DecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=Qwen3Model, + ) + + if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: + self.append_or_create_method_replacement( + description={ + "forward": get_qwen3_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=Qwen3Attention, + ) + if self.pipeline_stage_manager is None: + # replace qwen3 model forward method + self.append_or_create_method_replacement( + description={ + "forward": get_qwen3_model_forward_for_flash_attn( + self.shard_config, sp_mode, sp_size, sp_group + ), + }, + policy=policy, + target_key=Qwen3Model, + ) + + return policy + + def postprocess(self): + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "Qwen3Model": + module = self.model + else: + module = self.model.model + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config + ) + } + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "Qwen3Model": + module = self.model + else: + module = self.model.model + + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.norm) + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + + return held_layers + + +class Qwen3ModelPolicy(Qwen3Policy): + def module_policy(self): + policy = super().module_policy() + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=Qwen3Model, new_forward=Qwen3PipelineForwards.qwen3_model_forward, policy=policy + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen3 model""" + return [] + + +class Qwen3ForCausalLMPolicy(Qwen3Policy): + def module_policy(self): + policy = super().module_policy() + setattr(self.shard_config, "causal_lm", True) + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + # add a new item for casual lm + new_item = { + Qwen3ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=Linear1D_Col, + kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), + ) + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + elif use_zbv: + # add a new item for casual lm + new_item = { + Qwen3ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=LinearWithGradAccum, + kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), + ) + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen3ForCausalLM, new_forward=Qwen3PipelineForwards.qwen3_for_causal_lm_forward, policy=policy + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.lm_head) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + qwen3_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(qwen3_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: qwen3_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class Qwen3ForSequenceClassificationPolicy(Qwen3Policy): + def module_policy(self): + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + Qwen3ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + Qwen3ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=LinearWithGradAccum, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + # to be confirmed + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=Qwen3ForSequenceClassification, + new_forward=Qwen3PipelineForwards.qwen3_for_sequence_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.score) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen3 for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_qwen3_tf5.py b/colossalai/shardformer/policies/_qwen3_tf5.py new file mode 100644 index 000000000000..4c809602b6db --- /dev/null +++ b/colossalai/shardformer/policies/_qwen3_tf5.py @@ -0,0 +1,565 @@ +# Modifed from qwen2 policy +from functools import partial +from typing import Callable, Dict, List, Union + +import torch.distributed as dist +import torch.nn as nn +from torch import Tensor +from torch.nn import Module +from transformers.models.qwen3.modeling_qwen3 import ( + Qwen3Attention, + Qwen3DecoderLayer, + Qwen3ForCausalLM, + Qwen3ForSequenceClassification, + Qwen3Model, +) + +from colossalai.shardformer.layer import ( + FusedRMSNorm, + Linear1D_Col, + Linear1D_Row, + LinearWithGradAccum, + PaddingEmbedding, + RMSNorm, + VocabParallelEmbedding1D, +) + +from ..modeling.qwen3 import ( + Qwen3PipelineForwards, + get_lm_forward_with_dist_cross_entropy, + get_qwen3_flash_attention_forward, + get_qwen3_model_forward_for_flash_attn, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["Qwen3Policy", "Qwen3ForCausalLMPolicy", "Qwen3ForSequenceClassificationPolicy"] + + +class Qwen3Policy(Policy): + def __init__(self) -> None: + super().__init__() + import transformers + from packaging.version import Version + + assert Version(transformers.__version__) >= Version( + "5.17.0" + ), "The Qwen3 model should run on a transformers version of 5.17.0 or higher." + + def config_sanity_check(self): + pass + + def preprocess(self): + self.tie_weight = self.tie_weight_check() + self.origin_attn_implement = self.model.config._attn_implementation + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = PaddingEmbedding + norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm + + sp_mode = self.shard_config.sequence_parallelism_mode or None + sp_size = self.shard_config.sequence_parallel_size or None + sp_group = self.shard_config.sequence_parallel_process_group or None + sp_partial_derived = sp_mode in ["split_gather", "ring"] + if sp_mode == "all_to_all": + decoder_attribute_replacement = { + "num_heads": self.model.config.num_attention_heads // sp_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size + + policy[Qwen3Attention] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + ) + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + if hasattr(self.model.config, "num_key_value_heads"): + assert ( + self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of key_value heads must be divisible by tensor parallel size." + decoder_attribute_replacement = { + "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, + } + if getattr(self.model.config, "num_key_value_heads", False): + decoder_attribute_replacement["self_attn.num_key_value_heads"] = ( + self.model.config.num_key_value_heads // self.shard_config.tensor_parallel_size + ) + + policy[Qwen3DecoderLayer] = ModulePolicyDescription( + attribute_replacement=decoder_attribute_replacement, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=Linear1D_Col, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=Linear1D_Row, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + elif use_zbv: + policy[Qwen3DecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="self_attn.o_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.gate_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.up_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + SubModuleReplacementDescription( + suffix="mlp.down_proj", + target_module=LinearWithGradAccum, + kwargs=dict( + seq_parallel_mode=sp_mode, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + policy=policy, + target_key=Qwen3Model, + ) + + # optimization configuration + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="input_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + SubModuleReplacementDescription( + suffix="post_attention_layernorm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + ], + policy=policy, + target_key=Qwen3DecoderLayer, + ) + + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="norm", + target_module=norm_cls, + kwargs={"sp_partial_derived": sp_partial_derived}, + ), + policy=policy, + target_key=Qwen3Model, + ) + + if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: + self.append_or_create_method_replacement( + description={ + "forward": get_qwen3_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=Qwen3Attention, + ) + if self.pipeline_stage_manager is None: + # replace qwen3 model forward method + self.append_or_create_method_replacement( + description={ + "forward": get_qwen3_model_forward_for_flash_attn( + self.shard_config, sp_mode, sp_size, sp_group + ), + }, + policy=policy, + target_key=Qwen3Model, + ) + + return policy + + def postprocess(self): + if self.shard_config.enable_tensor_parallelism and self.shard_config.tensor_parallel_size > 1: + tp_group = self.shard_config.tensor_parallel_process_group + + def sync_head_grad(grad): + # Q/K norm weights are shared by the locally partitioned heads. + grad = grad.clone() + dist.all_reduce(grad, group=tp_group) + return grad + + for module in self.model.modules(): + if isinstance(module, Qwen3Attention): + for norm in (module.q_norm, module.k_norm): + weight = norm.weight + if ( + weight is not None + and weight.requires_grad + and not getattr(weight, "_qwen3_tp_grad_hook", False) + ): + weight.register_hook(sync_head_grad) + weight._qwen3_tp_grad_hook = True + return self.model + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if self.pipeline_stage_manager is None: + return + + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "Qwen3Model": + module = self.model + else: + module = self.model.model + + if stage_manager.is_interleave: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) + } + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = { + "forward": partial( + new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config + ) + } + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None + + if self.model.__class__.__name__ == "Qwen3Model": + module = self.model + else: + module = self.model.model + + stage_manager = self.pipeline_stage_manager + + held_layers = [] + held_layers.append(module.rotary_emb) + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embed_tokens) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.norm) + + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embed_tokens) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + if stage_manager.is_last_stage(): + held_layers.append(module.norm) + + return held_layers + + +class Qwen3ModelPolicy(Qwen3Policy): + def module_policy(self): + policy = super().module_policy() + + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=Qwen3Model, new_forward=Qwen3PipelineForwards.qwen3_model_forward, policy=policy + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + held_layers = super().get_held_layers() + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen3 model""" + return [] + + +class Qwen3ForCausalLMPolicy(Qwen3Policy): + def module_policy(self): + policy = super().module_policy() + setattr(self.shard_config, "causal_lm", True) + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + # add a new item for casual lm + new_item = { + Qwen3ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=not self.shard_config.parallel_output, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + elif use_zbv: + # add a new item for casual lm + new_item = { + Qwen3ForCausalLM: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="lm_head", + target_module=LinearWithGradAccum, + kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), + ) + ], + method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, + ) + } + policy.update(new_item) + + if self.pipeline_stage_manager: + # set None as default + self.set_pipeline_forward( + model_cls=Qwen3ForCausalLM, new_forward=Qwen3PipelineForwards.qwen3_for_causal_lm_forward, policy=policy + ) + + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.lm_head) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.lm_head) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + qwen3_model = self.model.model + if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: + if ( + id(qwen3_model.embed_tokens.weight) == id(self.model.lm_head.weight) + and self.pipeline_stage_manager.num_stages > 1 + ): + # tie weights + return [ + { + 0: qwen3_model.embed_tokens.weight, + self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, + } + ] + return [] + + +class Qwen3ForSequenceClassificationPolicy(Qwen3Policy): + def module_policy(self): + policy = super().module_policy() + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + if self.shard_config.enable_tensor_parallelism: + # add a new item for sequence classification + new_item = { + Qwen3ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + Qwen3ForSequenceClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="score", + target_module=LinearWithGradAccum, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + # to be confirmed + if self.pipeline_stage_manager: + self.set_pipeline_forward( + model_cls=Qwen3ForSequenceClassification, + new_forward=Qwen3PipelineForwards.qwen3_for_sequence_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[Module]: + """Get pipeline layers for current stage.""" + stage_manager = self.pipeline_stage_manager + held_layers = super().get_held_layers() + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.score) + else: + if stage_manager.is_last_stage(ignore_chunk=True): + held_layers.append(self.model.score) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + """No shared params in Qwen3 for sequence classification model""" + return [] diff --git a/colossalai/shardformer/policies/_vit_tf4.py b/colossalai/shardformer/policies/_vit_tf4.py new file mode 100644 index 000000000000..420ea286fd0a --- /dev/null +++ b/colossalai/shardformer/policies/_vit_tf4.py @@ -0,0 +1,434 @@ +import warnings +from typing import Callable, Dict, List, Union + +import torch.nn as nn + +import colossalai.shardformer.layer as col_nn +from colossalai.shardformer.layer import DropoutForReplicatedInput, Linear1D_Col + +from ..modeling.jit import get_jit_fused_dropout_add_func +from ..modeling.vit import ( + ViTForImageClassification_pipeline_forward, + ViTForMaskedImageModeling_pipeline_forward, + ViTModel_pipeline_forward, + get_jit_fused_vit_intermediate_forward, + get_jit_fused_vit_output_forward, + get_vit_flash_self_attention_forward, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["ViTPolicy", "ViTModelPolicy", "ViTForImageClassificationPolicy", "ViTForMaskedImageModelingPolicy"] + + +class ViTPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + self.enable_bias_gelu_fused = self.shard_config.enable_jit_fused and self.model.config.hidden_act == "gelu" + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.vit.modeling_vit import ( + ViTEmbeddings, + ViTIntermediate, + ViTLayer, + ViTOutput, + ViTSelfAttention, + ) + + policy = {} + + if self.shard_config.enable_sequence_parallelism: + self.shard_config.enable_sequence_parallelism = False + warnings.warn("Vit doesn't support sequence parallelism now, will ignore the sequence parallelism flag.") + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + policy[ViTEmbeddings] = ModulePolicyDescription( + attribute_replacement={}, + param_replacement=[], + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForReplicatedInput, + ) + ], + ) + + policy[ViTLayer] = ModulePolicyDescription( + attribute_replacement={ + "attention.attention.num_attention_heads": self.model.config.num_attention_heads + // self.shard_config.tensor_parallel_size, + "attention.attention.all_head_size": self.model.config.hidden_size + // self.shard_config.tensor_parallel_size, + }, + param_replacement=[], + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.attention.query", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.attention.key", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.attention.value", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=col_nn.Linear1D_Col, + kwargs={ + "skip_bias_add": self.enable_bias_gelu_fused, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + ], + ) + if self.enable_bias_gelu_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_vit_intermediate_forward(), + }, + policy=policy, + target_key=ViTIntermediate, + ) + elif use_zbv: + policy[ViTEmbeddings] = ModulePolicyDescription( + attribute_replacement={}, + param_replacement=[], + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="dropout", + target_module=DropoutForReplicatedInput, + ) + ], + ) + + policy[ViTLayer] = ModulePolicyDescription( + param_replacement=[], + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="attention.attention.query", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.attention.key", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.attention.value", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.attention.dropout", + target_module=col_nn.DropoutForParallelInput, + ), + SubModuleReplacementDescription( + suffix="attention.output.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="attention.output.dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + SubModuleReplacementDescription( + suffix="intermediate.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "skip_bias_add": self.enable_bias_gelu_fused, + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dense", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="output.dropout", + target_module=col_nn.DropoutForReplicatedInput, + ), + ], + ) + if self.enable_bias_gelu_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_vit_intermediate_forward(), + }, + policy=policy, + target_key=ViTIntermediate, + ) + # use flash attention + if self.shard_config.enable_flash_attention: + self.append_or_create_method_replacement( + description={ + "forward": get_vit_flash_self_attention_forward(), + }, + policy=policy, + target_key=ViTSelfAttention, + ) + + # use jit fused operator + if self.shard_config.enable_jit_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_vit_output_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=ViTOutput, + ) + + return policy + + def new_model_class(self): + return None + + def postprocess(self): + return self.model + + def get_held_layers(self) -> List[nn.Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + if self.model.__class__.__name__ == "ViTModel": + module = self.model + else: + module = self.model.vit + stage_manager = self.pipeline_stage_manager + + held_layers = [] + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embeddings) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.encoder.layer[start_idx:end_idx]) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + if stage_manager.is_first_stage(): + held_layers.append(module.embeddings) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.encoder.layer[start_idx:end_idx]) + return held_layers + + def set_pipeline_forward(self, model_cls: nn.Module, pipeline_forward: Callable, policy: Dict): + if self.pipeline_stage_manager: + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "ViTModel": + module = self.model + else: + module = self.model.vit + + layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = {"forward": pipeline_forward(stage_manager=stage_manager, stage_index=stage_index)} + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + +# ViTModel +class ViTModelPolicy(ViTPolicy): + def module_policy(self): + from transformers.models.vit.modeling_vit import ViTModel + + policy = super().module_policy() + + if self.shard_config.pipeline_stage_manager is not None: + self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + module = self.model + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.layernorm) + held_layers.append(module.pooler) + else: + if stage_manager.is_last_stage(): + held_layers.append(module.layernorm) + held_layers.append(module.pooler) + + return held_layers + + +# ViTForImageClassification +class ViTForImageClassificationPolicy(ViTPolicy): + def module_policy(self): + from transformers.models.vit.modeling_vit import ViTForImageClassification, ViTModel + + policy = super().module_policy() + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + new_item = { + ViTForImageClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="classifier", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + ViTForImageClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="classifier", + target_module=col_nn.LinearWithGradAccum, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + if self.shard_config.pipeline_stage_manager is not None: + self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) + self.set_pipeline_forward( + model_cls=ViTForImageClassification, + pipeline_forward=ViTForImageClassification_pipeline_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + module = self.model.vit + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.layernorm) + held_layers.append(self.model.classifier) + else: + if stage_manager.is_last_stage(): + held_layers.append(module.layernorm) + held_layers.append(self.model.classifier) + + return held_layers + + +# ViTForMaskedImageModeling +class ViTForMaskedImageModelingPolicy(ViTPolicy): + def module_policy(self): + from transformers.models.vit.modeling_vit import ViTForMaskedImageModeling, ViTModel + + policy = super().module_policy() + + if self.shard_config.pipeline_stage_manager is not None: + self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) + self.set_pipeline_forward( + model_cls=ViTForMaskedImageModeling, + pipeline_forward=ViTForMaskedImageModeling_pipeline_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + module = self.model.vit + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.layernorm) + held_layers.append(self.model.decoder) + else: + if stage_manager.is_last_stage(): + held_layers.append(module.layernorm) + held_layers.append(self.model.decoder) + + return held_layers diff --git a/colossalai/shardformer/policies/_vit_tf5.py b/colossalai/shardformer/policies/_vit_tf5.py new file mode 100644 index 000000000000..7b62f8ebf078 --- /dev/null +++ b/colossalai/shardformer/policies/_vit_tf5.py @@ -0,0 +1,284 @@ +import warnings +from typing import Callable, Dict, List, Union + +import torch.nn as nn + +import colossalai.shardformer.layer as col_nn +from colossalai.shardformer.layer import DropoutForReplicatedInput, Linear1D_Col + +from ..modeling.jit import get_jit_fused_dropout_add_func +from ..modeling.vit import ( + ViTForImageClassification_pipeline_forward, + ViTForMaskedImageModeling_pipeline_forward, + ViTModel_pipeline_forward, + get_jit_fused_vit_layer_forward, + get_vit_flash_self_attention_forward, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = ["ViTPolicy", "ViTModelPolicy", "ViTForImageClassificationPolicy", "ViTForMaskedImageModelingPolicy"] + + +class ViTPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + return self.model + + def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: + from transformers.models.vit.modeling_vit import ViTAttention, ViTEmbeddings, ViTLayer + + policy = {} + if self.shard_config.enable_sequence_parallelism: + self.shard_config.enable_sequence_parallelism = False + warnings.warn("ViT does not support sequence parallelism; the flag is disabled.") + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + tp = self.shard_config.enable_tensor_parallelism + if tp or use_zbv: + if tp: + assert self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 + policy[ViTEmbeddings] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription(suffix="dropout", target_module=DropoutForReplicatedInput) + ] + ) + replacements = [] + for suffix in ( + "attention.q_proj", + "attention.k_proj", + "attention.v_proj", + "attention.o_proj", + "mlp.fc1", + "mlp.fc2", + ): + column = suffix in ("attention.q_proj", "attention.k_proj", "attention.v_proj", "mlp.fc1") + target = (col_nn.Linear1D_Col if column else col_nn.Linear1D_Row) if tp else col_nn.LinearWithGradAccum + replacements.append( + SubModuleReplacementDescription( + suffix=suffix, + target_module=target, + kwargs={"fp8_communication": self.shard_config.fp8_communication, "use_zbv": use_zbv}, + ) + ) + replacements.append( + SubModuleReplacementDescription(suffix="dropout", target_module=col_nn.DropoutForReplicatedInput) + ) + policy[ViTLayer] = ModulePolicyDescription( + attribute_replacement=( + { + "attention.num_attention_heads": self.model.config.num_attention_heads + // self.shard_config.tensor_parallel_size + } + if tp + else {} + ), + sub_module_replacement=replacements, + ) + if self.shard_config.enable_flash_attention: + self.append_or_create_method_replacement( + description={"forward": get_vit_flash_self_attention_forward()}, + policy=policy, + target_key=ViTAttention, + ) + if self.shard_config.enable_jit_fused: + # HF's default gelu is exact erf GELU. The legacy bias_gelu kernel is a + # tanh approximation: retain the configured MLP activation and fuse + # dropout/residual at the new ViTLayer boundary instead. + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_vit_layer_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=ViTLayer, + ) + return policy + + def new_model_class(self): + return None + + def postprocess(self): + return self.model + + def get_held_layers(self) -> List[nn.Module]: + """Get pipeline layers for current stage.""" + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + if self.model.__class__.__name__ == "ViTModel": + module = self.model + else: + module = self.model.vit + stage_manager = self.pipeline_stage_manager + + held_layers = [] + if stage_manager.is_interleave: + assert stage_manager.num_model_chunks is not None + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_indices = stage_manager.get_stage_index(layers_per_stage) + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(module.embeddings) + for start_idx, end_idx in stage_indices: + held_layers.extend(module.layers[start_idx:end_idx]) + else: + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + if stage_manager.is_first_stage(): + held_layers.append(module.embeddings) + start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) + held_layers.extend(module.layers[start_idx:end_idx]) + return held_layers + + def set_pipeline_forward(self, model_cls: nn.Module, pipeline_forward: Callable, policy: Dict): + if self.pipeline_stage_manager: + stage_manager = self.pipeline_stage_manager + if self.model.__class__.__name__ == "ViTModel": + module = self.model + else: + module = self.model.vit + + layers_per_stage = stage_manager.distribute_layers(len(module.layers)) + stage_index = stage_manager.get_stage_index(layers_per_stage) + method_replacement = {"forward": pipeline_forward(stage_manager=stage_manager, stage_index=stage_index)} + self.append_or_create_method_replacement( + description=method_replacement, policy=policy, target_key=model_cls + ) + + +# ViTModel +class ViTModelPolicy(ViTPolicy): + def module_policy(self): + from transformers.models.vit.modeling_vit import ViTModel + + policy = super().module_policy() + + if self.shard_config.pipeline_stage_manager is not None: + self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + module = self.model + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.layernorm) + held_layers.append(module.pooler) + else: + if stage_manager.is_last_stage(): + held_layers.append(module.layernorm) + held_layers.append(module.pooler) + + return held_layers + + +# ViTForImageClassification +class ViTForImageClassificationPolicy(ViTPolicy): + def module_policy(self): + from transformers.models.vit.modeling_vit import ViTForImageClassification, ViTModel + + policy = super().module_policy() + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + if self.shard_config.enable_tensor_parallelism: + new_item = { + ViTForImageClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="classifier", + target_module=Linear1D_Col, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + elif use_zbv: + new_item = { + ViTForImageClassification: ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="classifier", + target_module=col_nn.LinearWithGradAccum, + kwargs=dict( + gather_output=True, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), + ) + ] + ) + } + policy.update(new_item) + if self.shard_config.pipeline_stage_manager is not None: + self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) + self.set_pipeline_forward( + model_cls=ViTForImageClassification, + pipeline_forward=ViTForImageClassification_pipeline_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + module = self.model.vit + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.layernorm) + held_layers.append(self.model.classifier) + else: + if stage_manager.is_last_stage(): + held_layers.append(module.layernorm) + held_layers.append(self.model.classifier) + + return held_layers + + +# ViTForMaskedImageModeling +class ViTForMaskedImageModelingPolicy(ViTPolicy): + def module_policy(self): + from transformers.models.vit.modeling_vit import ViTForMaskedImageModeling, ViTModel + + policy = super().module_policy() + + if self.shard_config.pipeline_stage_manager is not None: + self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) + self.set_pipeline_forward( + model_cls=ViTForMaskedImageModeling, + pipeline_forward=ViTForMaskedImageModeling_pipeline_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + + module = self.model.vit + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(module.layernorm) + held_layers.append(self.model.decoder) + else: + if stage_manager.is_last_stage(): + held_layers.append(module.layernorm) + held_layers.append(self.model.decoder) + + return held_layers diff --git a/colossalai/shardformer/policies/_whisper_tf4.py b/colossalai/shardformer/policies/_whisper_tf4.py new file mode 100644 index 000000000000..5d9b38e6fe98 --- /dev/null +++ b/colossalai/shardformer/policies/_whisper_tf4.py @@ -0,0 +1,845 @@ +import warnings +from functools import partial +from typing import Callable, Dict, List, Tuple + +import numpy as np +import torch.nn as nn +from torch import Tensor + +import colossalai.shardformer.layer as col_nn + +from ..modeling.jit import get_jit_fused_dropout_add_func +from ..modeling.whisper import ( + WhisperPipelineForwards, + get_jit_fused_whisper_decoder_layer_forward, + get_jit_fused_whisper_encoder_layer_forward, + get_whisper_decoder_forward_for_flash_attention, + get_whisper_flash_attention_forward, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = [ + "WhisperPolicy", + "WhisperModelPolicy", + "WhisperForConditionalGenerationPolicy", + "WhisperForAudioClassificationPolicy", +] + + +class WhisperPolicy(Policy): + def __init__(self) -> None: + super().__init__() + + def config_sanity_check(self): + pass + + def preprocess(self): + # reshape the embedding layer + r""" + Reshape the Embedding layer to make the embedding dimension divisible by world_size + """ + self.tie_weight = self.tie_weight_check() + return self.model + + def module_policy(self): + from transformers.models.whisper.modeling_whisper import ( + WhisperAttention, + WhisperDecoder, + WhisperDecoderLayer, + WhisperEncoder, + WhisperEncoderLayer, + WhisperFlashAttention2, + WhisperSdpaAttention, + ) + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = col_nn.VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = col_nn.PaddingEmbedding + + if self.shard_config.enable_fused_normalization: + norm_cls = col_nn.FusedLayerNorm + else: + norm_cls = col_nn.LayerNorm + + if self.shard_config.enable_sequence_parallelism: + self.shard_config.enable_sequence_parallelism = False + warnings.warn( + "Whisper doesn't support sequence parallelism now, will ignore the sequence parallelism flag." + ) + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + # TODO using the jit fused add_and_dropout affect the accuracy + if self.shard_config.enable_jit_fused: + self.shard_config.enable_jit_fused = False + warnings.warn("Whisper doesn't support jit fused operator now, will ignore the jit fused operator flag.") + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.encoder_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + policy[WhisperEncoderLayer] = ModulePolicyDescription( + attribute_replacement={ + "self_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.encoder_attention_heads + // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + policy[WhisperDecoderLayer] = ModulePolicyDescription( + attribute_replacement={ + "self_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.decoder_attention_heads + // self.shard_config.tensor_parallel_size, + "encoder_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, + "encoder_attn.num_heads": self.model.config.encoder_attention_heads + // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.q_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.k_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.v_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.out_proj", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + elif use_zbv: + policy[WhisperEncoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + policy[WhisperDecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.q_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.k_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.v_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.out_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + ], + policy=policy, + target_key=WhisperDecoder, + ) + + # optimization configuration + # Handle encoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="self_attn_layer_norm", + target_module=norm_cls, + ), + SubModuleReplacementDescription( + suffix="final_layer_norm", + target_module=norm_cls, + ), + ], + policy=policy, + target_key=WhisperEncoderLayer, + ) + + # Handle decoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="self_attn_layer_norm", + target_module=norm_cls, + ), + SubModuleReplacementDescription( + suffix="final_layer_norm", + target_module=norm_cls, + ), + ], + policy=policy, + target_key=WhisperDecoderLayer, + ) + + # handle encoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="layer_norm", + target_module=norm_cls, + ) + ], + policy=policy, + target_key=WhisperEncoder, + ) + + # handle decoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="layer_norm", + target_module=norm_cls, + ) + ], + policy=policy, + target_key=WhisperDecoder, + ) + + # enable flash attention + if self.shard_config.enable_flash_attention: + self.append_or_create_method_replacement( + description={ + "forward": get_whisper_flash_attention_forward(), + }, + policy=policy, + target_key=WhisperAttention, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_whisper_flash_attention_forward(), + }, + policy=policy, + target_key=WhisperFlashAttention2, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_whisper_flash_attention_forward(), + }, + policy=policy, + target_key=WhisperSdpaAttention, + ) + if not self.shard_config.pipeline_stage_manager: + self.append_or_create_method_replacement( + description={ + "forward": get_whisper_decoder_forward_for_flash_attention(self.shard_config), + }, + policy=policy, + target_key=WhisperDecoder, + ) + + # use jit fused operator + if self.shard_config.enable_jit_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_whisper_decoder_layer_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=WhisperDecoderLayer, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_whisper_encoder_layer_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=WhisperEncoderLayer, + ) + + return policy + + def add_lm_head_policy(self, base_policy): + from transformers.models.whisper.modeling_whisper import WhisperForConditionalGeneration + + # optimize for tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="proj_out", + target_module=col_nn.VocabParallelLMHead1D, + kwargs={ + "gather_output": True, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ), + policy=base_policy, + target_key=WhisperForConditionalGeneration, + ) + else: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="proj_out", + target_module=col_nn.PaddingLMHead, + kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, + ), + policy=base_policy, + target_key=WhisperForConditionalGeneration, + ) + + return base_policy + + def postprocess(self): + return self.model + + def distribute_whisper_layers( + self, num_encoder_layers: int, num_decoder_layers: int, num_stages: int + ) -> Tuple[List[int], int]: + """ + Distribute whisper layers into stages when pipeline parallel is used. + Return the layer distribution as a list and the starting stage of decoder. + If decoder doesn't exist, returned decoder starting stage is set to num_encoder_layers. + """ + stage_manager = self.pipeline_stage_manager + assert stage_manager is not None, "pipeline_stage_manager is None" + + # number of encoder layers must be a positive integer + if num_encoder_layers <= 0: + raise ValueError("The number of encoder layers for whisper must be a positive integer.") + + # number of layers should be large enough to fill in every stage + if num_encoder_layers + num_decoder_layers < num_stages: + raise ValueError("The total number of layers can't be smaller than number of stages.") + + # in the case of whisperEncoderModel, set decoder starting stage to num_stages since it doesn't exist + if num_decoder_layers == 0: + return stage_manager.distribute_layers(num_encoder_layers, num_stages), num_stages + + # the number of stages distributed between encoder and decoder is optimized in this way: + # num_encoder_stages = argmin(abs(num_encoder_layers / encoder_stages - num_decoder_layers / decoder_stages)) + # s.t. num_encoder_stages + num_decoder_stages = num_stages, num_encoder_stages >= 1, num_decoder_stages >= 1 + def objective(num_encoder_stages): + return abs(num_encoder_layers / num_encoder_stages - num_decoder_layers / (num_stages - num_encoder_stages)) + + num_encoder_stages = np.argmin([objective(i) for i in range(1, num_stages)]) + 1 + num_decoder_stages = num_stages - num_encoder_stages + + encoder_distribution = stage_manager.distribute_layers(num_encoder_layers, num_encoder_stages) + decoder_distribution = stage_manager.distribute_layers(num_decoder_layers, num_decoder_stages) + return encoder_distribution + decoder_distribution, num_encoder_stages + + def get_whisper_stage_index( + self, layers_per_stage: List[int], stage: int, decoder_starting_stage: int + ) -> Tuple[int, int]: + """ + Input the distribution of layers among stages, the current stage and the first stage of decoder. + Return the starting/ending idx of layers in encoder/decoder + """ + stage_manager = self.pipeline_stage_manager + assert stage_manager is not None, "pipeline_stage_manager is None" + + if stage < decoder_starting_stage: + return stage_manager.get_stage_index(layers_per_stage[:decoder_starting_stage], stage) + else: + return stage_manager.get_stage_index( + layers_per_stage[decoder_starting_stage:], + stage - decoder_starting_stage, + ) + + def get_held_layers(self) -> List[nn.Module]: + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + stage_manager = self.pipeline_stage_manager + + if self.model.__class__.__name__ == "WhisperModel": + model = self.model + elif self.model.__class__.__name__ == "WhisperForConditionalGeneration": + model = self.model.model + else: + model = None + + if model: + encoder = self.model.get_encoder() + decoder = self.model.get_decoder() + else: + # whisper for audio classification holds encoder only + encoder = self.model.encoder + decoder = None + + num_encoder_layers = len(encoder.layers) + if decoder: + num_decoder_layers = len(decoder.layers) + else: + num_decoder_layers = 0 + + held_layers = [] + if stage_manager.is_interleave: + layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + stage_indices = self.get_whisper_stage_index(layers_per_stage, stage_manager.stage, decoder_starting_stage) + + if stage_manager.stage < decoder_starting_stage: + # current stage is in whisper's encoder + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(encoder.embed_positions) + held_layers.append(encoder.conv1) + held_layers.append(encoder.conv2) + # interleaved: not use_zbv & stage_manager.stage == decoder_starting_stage - 1 + # zbv: use_zbv & stage_manager.stage == first stage + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and decoder_starting_stage - 1 + ): + held_layers.append(encoder.layer_norm) + for start_idx, end_idx in stage_indices: + held_layers.extend(encoder.layers[start_idx:end_idx]) + else: + # current stage is in whisper's decoder + # TODO:(Jianghai) We divide encoder and decoder layers into different parts here, + # the case encoder and decoder put in same stage should be add in the future. + if stage_manager.stage == decoder_starting_stage: + held_layers.append(decoder.embed_tokens) + held_layers.append(decoder.embed_positions) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(decoder.layer_norm) + for start_idx, end_idx in stage_indices: + held_layers.extend(encoder.layers[start_idx:end_idx]) + else: + layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + start_idx, end_idx = self.get_whisper_stage_index( + layers_per_stage, stage_manager.stage, decoder_starting_stage + ) + + if stage_manager.stage < decoder_starting_stage: + # current stage is in whisper's encoder + if stage_manager.is_first_stage(): + held_layers.append(encoder.embed_positions) + held_layers.append(encoder.conv1) + held_layers.append(encoder.conv2) + if stage_manager.stage == decoder_starting_stage - 1: + held_layers.append(encoder.layer_norm) + held_layers.extend(encoder.layers[start_idx:end_idx]) + else: + # current stage is in whisper's decoder + # TODO:(Jianghai) We divide encoder and decoder layers into different parts here, + # the case encoder and decoder put in same stage should be add in the future. + if stage_manager.stage == decoder_starting_stage: + held_layers.append(decoder.embed_tokens) + held_layers.append(decoder.embed_positions) + if stage_manager.is_last_stage(): + held_layers.append(decoder.layer_norm) + held_layers.extend(decoder.layers[start_idx:end_idx]) + return held_layers + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if not self.pipeline_stage_manager: + raise ValueError("set_pipeline_forward method can only be called when pipeline parallel is enabled.") + stage_manager = self.pipeline_stage_manager + + if self.model.__class__.__name__ == "WhisperModel": + model = self.model + elif self.model.__class__.__name__ == "WhisperForConditionalGeneration": + model = self.model.model + else: + model = None + + if model: + encoder = self.model.get_encoder() + decoder = self.model.get_decoder() + else: + encoder = self.model.encoder + decoder = None + + num_encoder_layers = len(encoder.layers) + if decoder: + num_decoder_layers = len(decoder.layers) + else: + num_decoder_layers = 0 + + layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + stage_index = self.get_whisper_stage_index(layers_per_stage, stage_manager.stage, decoder_starting_stage) + + method_replacement = { + "forward": partial( + new_forward, + stage_manager=stage_manager, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + shard_config=self.shard_config, + ) + } + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + +# WhisperModel +class WhisperModelPolicy(WhisperPolicy): + def module_policy(self): + from transformers import WhisperModel + + policy = super().module_policy() + + if self.pipeline_stage_manager is not None: + self.set_pipeline_forward( + model_cls=WhisperModel, + new_forward=WhisperPipelineForwards.whisper_model_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[nn.Module]: + return super().get_held_layers() + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + "no shared params in whisper model" + return [] + + +# WhisperForConditionalGeneration +class WhisperForConditionalGenerationPolicy(WhisperPolicy): + def module_policy(self): + from transformers import WhisperForConditionalGeneration + + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + + if self.pipeline_stage_manager is not None: + self.set_pipeline_forward( + model_cls=WhisperForConditionalGeneration, + new_forward=WhisperPipelineForwards.whisper_for_conditional_generation_forward, + policy=policy, + ) + return policy + + def postprocess(self): + return self.model + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.proj_out) + else: + if self.pipeline_stage_manager.is_last_stage(): + held_layers.append(self.model.proj_out) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + module = self.model + model = module.model + + if model: + encoder = self.model.get_encoder() + decoder = self.model.get_decoder() + else: + encoder = self.model.encoder + decoder = None + + num_encoder_layers = len(encoder.layers) + if decoder: + num_decoder_layers = len(decoder.layers) + else: + num_decoder_layers = 0 + + stage_manager = self.pipeline_stage_manager + if stage_manager is not None and stage_manager.num_stages > 1: + _, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + shared_params = [] + shared_embedding = {} + if id(module.proj_out) == id(model.decoder.embed_tokens): + shared_embedding[decoder_starting_stage] = model.decoder.embed_tokens + shared_embedding[stage_manager.num_stages - 1] = module.proj_out + if len(shared_embedding) > 0: + shared_params.append(shared_embedding) + return shared_params + return [] + + +# WhisperForAudioClassification +class WhisperForAudioClassificationPolicy(WhisperPolicy): + def module_policy(self): + from transformers import WhisperForAudioClassification + + policy = super().module_policy() + + if self.pipeline_stage_manager is not None: + self.set_pipeline_forward( + model_cls=WhisperForAudioClassification, + new_forward=WhisperPipelineForwards.whisper_for_audio_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.projector) + held_layers.append(self.model.classifier) + else: + if self.pipeline_stage_manager.is_last_stage(): + held_layers.append(self.model.projector) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + return [] diff --git a/colossalai/shardformer/policies/_whisper_tf5.py b/colossalai/shardformer/policies/_whisper_tf5.py new file mode 100644 index 000000000000..65b72dbb57d2 --- /dev/null +++ b/colossalai/shardformer/policies/_whisper_tf5.py @@ -0,0 +1,832 @@ +import warnings +from functools import partial +from typing import Callable, Dict, List, Tuple + +import numpy as np +import torch.nn as nn +from torch import Tensor + +import colossalai.shardformer.layer as col_nn + +from ..modeling.jit import get_jit_fused_dropout_add_func +from ..modeling.whisper import ( + WhisperPipelineForwards, + get_jit_fused_whisper_decoder_layer_forward, + get_jit_fused_whisper_encoder_layer_forward, + get_whisper_decoder_forward_for_flash_attention, + get_whisper_flash_attention_forward, +) +from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription + +__all__ = [ + "WhisperPolicy", + "WhisperModelPolicy", + "WhisperForConditionalGenerationPolicy", + "WhisperForAudioClassificationPolicy", +] + + +class WhisperPolicy(Policy): + def __init__(self) -> None: + super().__init__() + + def config_sanity_check(self): + pass + + def preprocess(self): + # reshape the embedding layer + r""" + Reshape the Embedding layer to make the embedding dimension divisible by world_size + """ + self.tie_weight = self.tie_weight_check() + return self.model + + def module_policy(self): + from transformers.models.whisper.modeling_whisper import ( + WhisperAttention, + WhisperDecoder, + WhisperDecoderLayer, + WhisperEncoder, + WhisperEncoderLayer, + ) + + policy = {} + + embedding_cls = None + if self.shard_config.enable_tensor_parallelism: + embedding_cls = col_nn.VocabParallelEmbedding1D + else: + if self.tie_weight: + embedding_cls = col_nn.PaddingEmbedding + + if self.shard_config.enable_fused_normalization: + norm_cls = col_nn.FusedLayerNorm + else: + norm_cls = col_nn.LayerNorm + + if self.shard_config.enable_sequence_parallelism: + self.shard_config.enable_sequence_parallelism = False + warnings.warn( + "Whisper doesn't support sequence parallelism now, will ignore the sequence parallelism flag." + ) + + use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv + + # TODO using the jit fused add_and_dropout affect the accuracy + if self.shard_config.enable_jit_fused: + self.shard_config.enable_jit_fused = False + warnings.warn("Whisper doesn't support jit fused operator now, will ignore the jit fused operator flag.") + + if self.shard_config.enable_tensor_parallelism: + assert ( + self.model.config.encoder_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), f"The number of attention heads must be divisible by tensor parallel size." + assert ( + self.model.config.decoder_attention_heads % self.shard_config.tensor_parallel_size == 0 + ), "Decoder attention heads must be divisible by tensor parallel size." + policy[WhisperEncoderLayer] = ModulePolicyDescription( + attribute_replacement={ + "self_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.encoder_attention_heads + // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + policy[WhisperDecoderLayer] = ModulePolicyDescription( + attribute_replacement={ + "self_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, + "self_attn.num_heads": self.model.config.decoder_attention_heads + // self.shard_config.tensor_parallel_size, + "encoder_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, + "encoder_attn.num_heads": self.model.config.decoder_attention_heads + // self.shard_config.tensor_parallel_size, + }, + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.q_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.k_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.v_proj", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.out_proj", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.Linear1D_Col, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.Linear1D_Row, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + elif use_zbv: + policy[WhisperEncoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + policy[WhisperDecoderLayer] = ModulePolicyDescription( + sub_module_replacement=[ + SubModuleReplacementDescription( + suffix="self_attn.q_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.k_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.v_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="self_attn.out_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.q_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.k_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.v_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="encoder_attn.out_proj", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc1", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + SubModuleReplacementDescription( + suffix="fc2", + target_module=col_nn.LinearWithGradAccum, + kwargs={ + "fp8_communication": self.shard_config.fp8_communication, + "use_zbv": use_zbv, + }, + ), + ], + ) + + if embedding_cls is not None: + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="embed_tokens", + target_module=embedding_cls, + kwargs=( + { + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + } + if self.shard_config.enable_tensor_parallelism + else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} + ), + ), + ], + policy=policy, + target_key=WhisperDecoder, + ) + + # optimization configuration + # Handle encoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="self_attn_layer_norm", + target_module=norm_cls, + ), + SubModuleReplacementDescription( + suffix="final_layer_norm", + target_module=norm_cls, + ), + ], + policy=policy, + target_key=WhisperEncoderLayer, + ) + + # Handle decoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="self_attn_layer_norm", + target_module=norm_cls, + ), + SubModuleReplacementDescription( + suffix="final_layer_norm", + target_module=norm_cls, + ), + ], + policy=policy, + target_key=WhisperDecoderLayer, + ) + + # handle encoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="layer_norm", + target_module=norm_cls, + ) + ], + policy=policy, + target_key=WhisperEncoder, + ) + + # handle decoder layer + self.append_or_create_submodule_replacement( + description=[ + SubModuleReplacementDescription( + suffix="layer_norm", + target_module=norm_cls, + ) + ], + policy=policy, + target_key=WhisperDecoder, + ) + + # enable flash attention + if self.shard_config.enable_flash_attention: + self.append_or_create_method_replacement( + description={ + "forward": get_whisper_flash_attention_forward(), + }, + policy=policy, + target_key=WhisperAttention, + ) + if not self.shard_config.pipeline_stage_manager: + self.append_or_create_method_replacement( + description={ + "forward": get_whisper_decoder_forward_for_flash_attention(self.shard_config), + }, + policy=policy, + target_key=WhisperDecoder, + ) + + # use jit fused operator + if self.shard_config.enable_jit_fused: + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_whisper_decoder_layer_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=WhisperDecoderLayer, + ) + self.append_or_create_method_replacement( + description={ + "forward": get_jit_fused_whisper_encoder_layer_forward(), + "dropout_add": get_jit_fused_dropout_add_func(), + }, + policy=policy, + target_key=WhisperEncoderLayer, + ) + + return policy + + def add_lm_head_policy(self, base_policy): + from transformers.models.whisper.modeling_whisper import WhisperForConditionalGeneration + + # optimize for tensor parallelism + if self.shard_config.enable_tensor_parallelism: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="proj_out", + target_module=col_nn.VocabParallelLMHead1D, + kwargs={ + "gather_output": True, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ), + policy=base_policy, + target_key=WhisperForConditionalGeneration, + ) + else: + self.append_or_create_submodule_replacement( + description=SubModuleReplacementDescription( + suffix="proj_out", + target_module=col_nn.PaddingLMHead, + kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, + ), + policy=base_policy, + target_key=WhisperForConditionalGeneration, + ) + + return base_policy + + def postprocess(self): + return self.model + + def distribute_whisper_layers( + self, num_encoder_layers: int, num_decoder_layers: int, num_stages: int + ) -> Tuple[List[int], int]: + """ + Distribute whisper layers into stages when pipeline parallel is used. + Return the layer distribution as a list and the starting stage of decoder. + If decoder doesn't exist, returned decoder starting stage is set to num_encoder_layers. + """ + stage_manager = self.pipeline_stage_manager + assert stage_manager is not None, "pipeline_stage_manager is None" + + # number of encoder layers must be a positive integer + if num_encoder_layers <= 0: + raise ValueError("The number of encoder layers for whisper must be a positive integer.") + + # number of layers should be large enough to fill in every stage + if num_encoder_layers + num_decoder_layers < num_stages: + raise ValueError("The total number of layers can't be smaller than number of stages.") + + # in the case of whisperEncoderModel, set decoder starting stage to num_stages since it doesn't exist + if num_decoder_layers == 0: + return stage_manager.distribute_layers(num_encoder_layers, num_stages), num_stages + + # the number of stages distributed between encoder and decoder is optimized in this way: + # num_encoder_stages = argmin(abs(num_encoder_layers / encoder_stages - num_decoder_layers / decoder_stages)) + # s.t. num_encoder_stages + num_decoder_stages = num_stages, num_encoder_stages >= 1, num_decoder_stages >= 1 + def objective(num_encoder_stages): + return abs(num_encoder_layers / num_encoder_stages - num_decoder_layers / (num_stages - num_encoder_stages)) + + num_encoder_stages = np.argmin([objective(i) for i in range(1, num_stages)]) + 1 + num_decoder_stages = num_stages - num_encoder_stages + + encoder_distribution = stage_manager.distribute_layers(num_encoder_layers, num_encoder_stages) + decoder_distribution = stage_manager.distribute_layers(num_decoder_layers, num_decoder_stages) + return encoder_distribution + decoder_distribution, num_encoder_stages + + def get_whisper_stage_index( + self, layers_per_stage: List[int], stage: int, decoder_starting_stage: int + ) -> Tuple[int, int]: + """ + Input the distribution of layers among stages, the current stage and the first stage of decoder. + Return the starting/ending idx of layers in encoder/decoder + """ + stage_manager = self.pipeline_stage_manager + assert stage_manager is not None, "pipeline_stage_manager is None" + + if stage < decoder_starting_stage: + return stage_manager.get_stage_index(layers_per_stage[:decoder_starting_stage], stage) + else: + return stage_manager.get_stage_index( + layers_per_stage[decoder_starting_stage:], + stage - decoder_starting_stage, + ) + + def get_held_layers(self) -> List[nn.Module]: + assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" + stage_manager = self.pipeline_stage_manager + + if self.model.__class__.__name__ == "WhisperModel": + model = self.model + elif self.model.__class__.__name__ == "WhisperForConditionalGeneration": + model = self.model.model + else: + model = None + + if model: + encoder = self.model.get_encoder() + decoder = self.model.get_decoder() + else: + # whisper for audio classification holds encoder only + encoder = self.model.encoder + decoder = None + + num_encoder_layers = len(encoder.layers) + if decoder: + num_decoder_layers = len(decoder.layers) + else: + num_decoder_layers = 0 + + held_layers = [] + if stage_manager.is_interleave: + layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + stage_indices = self.get_whisper_stage_index(layers_per_stage, stage_manager.stage, decoder_starting_stage) + + if stage_manager.stage < decoder_starting_stage: + # current stage is in whisper's encoder + if stage_manager.is_first_stage(ignore_chunk=True): + held_layers.append(encoder.embed_positions) + held_layers.append(encoder.conv1) + held_layers.append(encoder.conv2) + # interleaved: not use_zbv & stage_manager.stage == decoder_starting_stage - 1 + # zbv: use_zbv & stage_manager.stage == first stage + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and decoder_starting_stage - 1 + ): + held_layers.append(encoder.layer_norm) + for start_idx, end_idx in stage_indices: + held_layers.extend(encoder.layers[start_idx:end_idx]) + else: + # current stage is in whisper's decoder + # TODO:(Jianghai) We divide encoder and decoder layers into different parts here, + # the case encoder and decoder put in same stage should be add in the future. + if stage_manager.stage == decoder_starting_stage: + held_layers.append(decoder.embed_tokens) + held_layers.append(decoder.embed_positions) + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(decoder.layer_norm) + for start_idx, end_idx in stage_indices: + held_layers.extend(encoder.layers[start_idx:end_idx]) + else: + layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + start_idx, end_idx = self.get_whisper_stage_index( + layers_per_stage, stage_manager.stage, decoder_starting_stage + ) + + if stage_manager.stage < decoder_starting_stage: + # current stage is in whisper's encoder + if stage_manager.is_first_stage(): + held_layers.append(encoder.embed_positions) + held_layers.append(encoder.conv1) + held_layers.append(encoder.conv2) + if stage_manager.stage == decoder_starting_stage - 1: + held_layers.append(encoder.layer_norm) + held_layers.extend(encoder.layers[start_idx:end_idx]) + else: + # current stage is in whisper's decoder + # TODO:(Jianghai) We divide encoder and decoder layers into different parts here, + # the case encoder and decoder put in same stage should be add in the future. + if stage_manager.stage == decoder_starting_stage: + held_layers.append(decoder.embed_tokens) + held_layers.append(decoder.embed_positions) + if stage_manager.is_last_stage(): + held_layers.append(decoder.layer_norm) + held_layers.extend(decoder.layers[start_idx:end_idx]) + return held_layers + + def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: + """If under pipeline parallel setting, replacing the original forward method of huggingface + to customized forward method, and add this changing to policy.""" + if not self.pipeline_stage_manager: + raise ValueError("set_pipeline_forward method can only be called when pipeline parallel is enabled.") + stage_manager = self.pipeline_stage_manager + + if self.model.__class__.__name__ == "WhisperModel": + model = self.model + elif self.model.__class__.__name__ == "WhisperForConditionalGeneration": + model = self.model.model + else: + model = None + + if model: + encoder = self.model.get_encoder() + decoder = self.model.get_decoder() + else: + encoder = self.model.encoder + decoder = None + + num_encoder_layers = len(encoder.layers) + if decoder: + num_decoder_layers = len(decoder.layers) + else: + num_decoder_layers = 0 + + layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + stage_index = self.get_whisper_stage_index(layers_per_stage, stage_manager.stage, decoder_starting_stage) + + method_replacement = { + "forward": partial( + new_forward, + stage_manager=stage_manager, + stage_index=stage_index, + decoder_starting_stage=decoder_starting_stage, + shard_config=self.shard_config, + ) + } + self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) + + +# WhisperModel +class WhisperModelPolicy(WhisperPolicy): + def module_policy(self): + from transformers import WhisperModel + + policy = super().module_policy() + + if self.pipeline_stage_manager is not None: + self.set_pipeline_forward( + model_cls=WhisperModel, + new_forward=WhisperPipelineForwards.whisper_model_forward, + policy=policy, + ) + + return policy + + def get_held_layers(self) -> List[nn.Module]: + return super().get_held_layers() + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + "no shared params in whisper model" + return [] + + +# WhisperForConditionalGeneration +class WhisperForConditionalGenerationPolicy(WhisperPolicy): + def module_policy(self): + from transformers import WhisperForConditionalGeneration + + policy = super().module_policy() + policy = self.add_lm_head_policy(policy) + + if self.pipeline_stage_manager is not None: + self.set_pipeline_forward( + model_cls=WhisperForConditionalGeneration, + new_forward=WhisperPipelineForwards.whisper_for_conditional_generation_forward, + policy=policy, + ) + return policy + + def postprocess(self): + return self.model + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.proj_out) + else: + if self.pipeline_stage_manager.is_last_stage(): + held_layers.append(self.model.proj_out) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + module = self.model + model = module.model + + if model: + encoder = self.model.get_encoder() + decoder = self.model.get_decoder() + else: + encoder = self.model.encoder + decoder = None + + num_encoder_layers = len(encoder.layers) + if decoder: + num_decoder_layers = len(decoder.layers) + else: + num_decoder_layers = 0 + + stage_manager = self.pipeline_stage_manager + if stage_manager is not None and stage_manager.num_stages > 1: + _, decoder_starting_stage = self.distribute_whisper_layers( + num_encoder_layers, num_decoder_layers, stage_manager.num_stages + ) + shared_params = [] + shared_embedding = {} + if id(module.proj_out) == id(model.decoder.embed_tokens): + shared_embedding[decoder_starting_stage] = model.decoder.embed_tokens + shared_embedding[stage_manager.num_stages - 1] = module.proj_out + if len(shared_embedding) > 0: + shared_params.append(shared_embedding) + return shared_params + return [] + + +# WhisperForAudioClassification +class WhisperForAudioClassificationPolicy(WhisperPolicy): + def module_policy(self): + from transformers import WhisperForAudioClassification + + policy = super().module_policy() + + if self.pipeline_stage_manager is not None: + self.set_pipeline_forward( + model_cls=WhisperForAudioClassification, + new_forward=WhisperPipelineForwards.whisper_for_audio_classification_forward, + policy=policy, + ) + return policy + + def get_held_layers(self) -> List[nn.Module]: + held_layers = super().get_held_layers() + stage_manager = self.pipeline_stage_manager + if stage_manager.is_interleave: + if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( + not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) + ): + held_layers.append(self.model.projector) + held_layers.append(self.model.classifier) + else: + if self.pipeline_stage_manager.is_last_stage(): + held_layers.append(self.model.projector) + held_layers.append(self.model.classifier) + return held_layers + + def get_shared_params(self) -> List[Dict[int, Tensor]]: + return [] diff --git a/colossalai/shardformer/policies/auto_policy.py b/colossalai/shardformer/policies/auto_policy.py index 9cdc2f7f529e..3986bffd738b 100644 --- a/colossalai/shardformer/policies/auto_policy.py +++ b/colossalai/shardformer/policies/auto_policy.py @@ -180,6 +180,13 @@ class PolicyLocation: "transformers_modules.modeling_deepseek.DeepseekV3ForCausalLM": PolicyLocation( file_name="deepseek_v3", class_name="DeepseekV3ForCausalLMPolicy" ), + # HF5 native DeepSeek-V3; legacy remote-code registrations remain distinct. + "transformers.models.deepseek_v3.modeling_deepseek_v3.DeepseekV3Model": PolicyLocation( + file_name="deepseek_v3", class_name="DeepseekV3ModelPolicy" + ), + "transformers.models.deepseek_v3.modeling_deepseek_v3.DeepseekV3ForCausalLM": PolicyLocation( + file_name="deepseek_v3", class_name="DeepseekV3ForCausalLMPolicy" + ), # Falcon "transformers.models.falcon.modeling_falcon.FalconModel": PolicyLocation( file_name="falcon", class_name="FalconModelPolicy" @@ -250,6 +257,9 @@ def import_policy(policy_location: PolicyLocation) -> Policy: """ Dynamically import a Policy class based on the policy location. """ + from .._utils import check_hf5_policy + + check_hf5_policy(policy_location.file_name) module_name = f"colossalai.shardformer.policies.{policy_location.file_name}" module = importlib.import_module(module_name) return getattr(module, policy_location.class_name) diff --git a/colossalai/shardformer/policies/bert.py b/colossalai/shardformer/policies/bert.py index fd4e020b0cb1..068218025b87 100644 --- a/colossalai/shardformer/policies/bert.py +++ b/colossalai/shardformer/policies/bert.py @@ -1,825 +1,9 @@ -import warnings -from functools import partial -from typing import Callable, Dict, List +r"""BERT sharding policies, selected by the installed transformers major version. -import torch.nn as nn -from torch import Tensor -from torch.nn import Module +``_bert_tf5`` targets transformers v5; ``_bert_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -import colossalai.shardformer.layer as col_nn +from colossalai._compat import reexport_transformers_impl -from ..modeling.bert import ( - BertPipelineForwards, - bert_sequence_parallel_forward_fn, - get_bert_sequence_parallel_attention_forward, - get_jit_fused_bert_intermediate_forward, - get_jit_fused_bert_output_forward, - get_jit_fused_bert_self_output_forward, -) -from ..modeling.jit import get_jit_fused_dropout_add_func -from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription - -__all__ = [ - "BertPolicy", - "BertModelPolicy", - "BertForPreTrainingPolicy", - "BertLMHeadModelPolicy", - "BertForMaskedLMPolicy", - "BertForNextSentencePredictionPolicy", - "BertForSequenceClassificationPolicy", - "BertForTokenClassificationPolicy", - "BertForMultipleChoicePolicy", - "BertForQuestionAnsweringPolicy", -] - - -class BertPolicy(Policy): - def config_sanity_check(self): - pass - - def preprocess(self): - self.tie_weight = self.tie_weight_check() - self.enable_bias_gelu_fused = self.shard_config.enable_jit_fused and self.model.config.hidden_act == "gelu" - return self.model - - def module_policy(self): - from transformers.models.bert.modeling_bert import ( - BertEmbeddings, - BertIntermediate, - BertLayer, - BertModel, - BertOutput, - BertSdpaSelfAttention, - BertSelfOutput, - ) - - policy = {} - - embedding_cls = None - if self.shard_config.enable_tensor_parallelism: - embedding_cls = col_nn.VocabParallelEmbedding1D - else: - if self.tie_weight: - embedding_cls = col_nn.PaddingEmbedding - - if self.shard_config.enable_fused_normalization: - norm_cls = col_nn.FusedLayerNorm - else: - norm_cls = col_nn.LayerNorm - - sp_mode = self.shard_config.sequence_parallelism_mode or None - assert sp_mode != "all_to_all", "all_to_all sequence parallelism is not supported for Bert" - if sp_mode == "ring": - warnings.warn( - f"For Bert, sequence parallelism is currently not support mode {sp_mode}, will set to be split_gather" - ) - sp_mode = "split_gather" - - sp_partial_derived = sp_mode == "split_gather" - - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_sequence_parallelism: - # Fix the tgt_len size in bert sequence parallel attention forward. - self.append_or_create_method_replacement( - description={ - "forward": get_bert_sequence_parallel_attention_forward(self.shard_config), - }, - policy=policy, - target_key=BertSdpaSelfAttention, - ) - - if self.shard_config.enable_tensor_parallelism: - assert ( - self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of attention heads must be divisible by tensor parallel size." - policy[BertLayer] = ModulePolicyDescription( - attribute_replacement={ - "attention.self.all_head_size": self.model.config.hidden_size - // self.shard_config.tensor_parallel_size, - "crossattention.self.all_head_size": self.model.config.hidden_size - // self.shard_config.tensor_parallel_size, - "attention.self.num_attention_heads": self.model.config.num_attention_heads - // self.shard_config.tensor_parallel_size, - "crossattention.self.num_attention_heads": self.model.config.num_attention_heads - // self.shard_config.tensor_parallel_size, - }, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="attention.self.query", - target_module=col_nn.Linear1D_Col, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.self.key", - target_module=col_nn.Linear1D_Col, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.self.value", - target_module=col_nn.Linear1D_Col, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.self.dropout", - target_module=col_nn.DropoutForParallelInput, - ), - SubModuleReplacementDescription( - suffix="attention.output.dense", - target_module=col_nn.Linear1D_Row, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.output.dropout", - target_module=col_nn.DropoutForParallelInput, - ), - SubModuleReplacementDescription( - suffix="intermediate.dense", - target_module=col_nn.Linear1D_Col, - kwargs={ - "seq_parallel_mode": sp_mode, - "skip_bias_add": self.enable_bias_gelu_fused, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dense", - target_module=col_nn.Linear1D_Row, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dropout", - target_module=col_nn.DropoutForParallelInput, - ), - ], - ) - - policy[BertEmbeddings] = ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="dropout", - target_module=col_nn.DropoutForReplicatedInput, - ), - ] - ) - if self.enable_bias_gelu_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_bert_intermediate_forward(), - }, - policy=policy, - target_key=BertIntermediate, - ) - - elif use_zbv: - policy[BertLayer] = ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="attention.self.query", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.self.key", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.self.value", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.self.dropout", - target_module=col_nn.DropoutForParallelInput, - ), - SubModuleReplacementDescription( - suffix="attention.output.dense", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.output.dropout", - target_module=col_nn.DropoutForParallelInput, - ), - SubModuleReplacementDescription( - suffix="intermediate.dense", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "seq_parallel_mode": sp_mode, - "skip_bias_add": self.enable_bias_gelu_fused, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dense", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "seq_parallel_mode": sp_mode, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dropout", - target_module=col_nn.DropoutForParallelInput, - ), - ], - ) - - policy[BertEmbeddings] = ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="dropout", - target_module=col_nn.DropoutForReplicatedInput, - ), - ] - ) - if self.enable_bias_gelu_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_bert_intermediate_forward(), - }, - policy=policy, - target_key=BertIntermediate, - ) - - if sp_mode == "split_gather": - self.append_or_create_method_replacement( - description={"forward": bert_sequence_parallel_forward_fn(self.shard_config)}, - policy=policy, - target_key=BertModel, - ) - - if embedding_cls is not None: - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="word_embeddings", - target_module=embedding_cls, - kwargs=( - { - "fp8_communication": self.shard_config.fp8_communication, - } - if self.shard_config.enable_tensor_parallelism - else {} - ), - ) - ], - policy=policy, - target_key=BertEmbeddings, - ) - - # optimization configuration - # Handle bert layer - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="attention.output.LayerNorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - SubModuleReplacementDescription( - suffix="output.LayerNorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - ], - policy=policy, - target_key=BertLayer, - ) - # handle embedding layer - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="LayerNorm", - target_module=norm_cls, - ) - ], - policy=policy, - target_key=BertEmbeddings, - ) - - # use jit operator - if self.shard_config.enable_jit_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_bert_self_output_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=BertSelfOutput, - ) - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_bert_output_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=BertOutput, - ) - - return policy - - def add_lm_head_policy(self, base_policy): - from transformers.models.bert.modeling_bert import BertLMPredictionHead - - # optimize for tensor parallelism - if self.shard_config.enable_tensor_parallelism: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="decoder", - target_module=col_nn.VocabParallelLMHead1D, - kwargs={ - "gather_output": True, - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - }, - ), - policy=base_policy, - target_key=BertLMPredictionHead, - ) - else: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="decoder", - target_module=col_nn.PaddingLMHead, - kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, - ), - policy=base_policy, - target_key=BertLMPredictionHead, - ) - - # optimize with fused normalization - if self.shard_config.enable_fused_normalization: - # Handle bert lm prediction head - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="transform.LayerNorm", - target_module=col_nn.FusedLayerNorm, - ), - policy=base_policy, - target_key=BertLMPredictionHead, - ) - return base_policy - - def add_lm_prediction_policy(self, base_policy): - from transformers.models.bert.modeling_bert import BertLMPredictionHead - - method_replacement = { - "_save_to_state_dict": col_nn.ParallelModule._save_to_state_dict, - "_load_from_state_dict": col_nn.ParallelModule._load_from_state_dict, - } - self.append_or_create_method_replacement( - description=method_replacement, - policy=base_policy, - target_key=BertLMPredictionHead, - ) - return base_policy - - def postprocess(self): - return self.model - - def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: - """ - If under pipeline parallel setting, replacing the original forward method of huggingface - to customized forward method, and add this changing to policy. - """ - if self.pipeline_stage_manager is None: - return - - stage_manager = self.pipeline_stage_manager - if self.model.__class__.__name__ == "BertModel": - module = self.model - else: - module = self.model.bert - - if stage_manager.is_interleave: - layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) - stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial( - new_forward, - stage_manager=stage_manager, - shard_config=self.shard_config, - ) - } - - else: - layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) - stage_index = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial( - new_forward, - stage_manager=stage_manager, - stage_index=stage_index, - shard_config=self.shard_config, - ) - } - - self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - assert self.pipeline_stage_manager is not None - - if self.model.__class__.__name__ == "BertModel": - module = self.model - else: - module = self.model.bert - stage_manager = self.pipeline_stage_manager - - held_layers = [] - if stage_manager.is_interleave: - assert stage_manager.num_model_chunks is not None - layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) - stage_indices = stage_manager.get_stage_index(layers_per_stage) - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(module.embeddings) - for start_idx, end_idx in stage_indices: - held_layers.extend(module.encoder.layer[start_idx:end_idx]) - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(module.pooler) - - else: - layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) - if stage_manager.is_first_stage(): - held_layers.append(module.embeddings) - start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) - held_layers.extend(module.encoder.layer[start_idx:end_idx]) - if stage_manager.is_last_stage(): - held_layers.append(module.pooler) - - return held_layers - - -# BertModel -class BertModelPolicy(BertPolicy): - def module_policy(self): - policy = super().module_policy() - from transformers.models.bert.modeling_bert import BertModel - - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertModel, - new_forward=BertPipelineForwards.bert_model_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - held_layers = super().get_held_layers() - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in bert model""" - return [] - - -# BertForPreTraining -class BertForPreTrainingPolicy(BertPolicy): - def module_policy(self): - policy = super().module_policy() - policy = self.add_lm_head_policy(policy) - policy = self.add_lm_prediction_policy(policy) - from transformers.models.bert.modeling_bert import BertForPreTraining - - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertForPreTraining, - new_forward=BertPipelineForwards.bert_for_pretraining_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage""" - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.cls) - - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - model = self.model - if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: - if id(model.bert.embeddings.word_embeddings.weight) == id(model.cls.predictions.decoder.weight): - # tie weights - return [ - { - 0: model.bert.embeddings.word_embeddings.weight, - self.pipeline_stage_manager.num_stages - 1: model.cls.predictions.decoder.weight, - } - ] - return [] - - -# BertLMHeadModel -class BertLMHeadModelPolicy(BertPolicy): - def module_policy(self): - policy = super().module_policy() - policy = self.add_lm_head_policy(policy) - policy = self.add_lm_prediction_policy(policy) - from transformers.models.bert.modeling_bert import BertLMHeadModel - - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertLMHeadModel, - new_forward=BertPipelineForwards.bert_lm_head_model_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """ - get pipeline layers for current stage - """ - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.cls) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - bert_model = self.model.bert - if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: - if id(bert_model.embeddings.word_embeddings.weight) == id(self.model.cls.predictions.decoder.weight): - # tie weights - return [ - { - 0: bert_model.embeddings.word_embeddings.weight, - self.pipeline_stage_manager.num_stages - 1: self.model.cls.predictions.decoder.weight, - } - ] - return [] - - -# BertForMaskedLM -class BertForMaskedLMPolicy(BertPolicy): - def module_policy(self): - policy = super().module_policy() - policy = self.add_lm_head_policy(policy) - policy = self.add_lm_prediction_policy(policy) - from transformers.models.bert.modeling_bert import BertForMaskedLM - - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertForMaskedLM, - new_forward=BertPipelineForwards.bert_for_masked_lm_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """ - get pipeline layers for current stage - """ - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.cls) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - bert_model = self.model.bert - if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: - if id(bert_model.embeddings.word_embeddings.weight) == id(self.model.cls.predictions.decoder.weight): - # tie weights - return [ - { - 0: bert_model.embeddings.word_embeddings.weight, - self.pipeline_stage_manager.num_stages - 1: self.model.cls.predictions.decoder.weight, - } - ] - return [] - - -# BertForSequenceClassification -class BertForSequenceClassificationPolicy(BertPolicy): - def module_policy(self): - from transformers.models.bert.modeling_bert import BertForSequenceClassification - - policy = super().module_policy() - - if self.shard_config.enable_tensor_parallelism: - addon_module = { - BertForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="dropout", - target_module=col_nn.DropoutForParallelInput, - ) - ] - ) - } - policy.update(addon_module) - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertForSequenceClassification, - new_forward=BertPipelineForwards.bert_for_sequence_classification_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """ - get pipeline layers for current stage - """ - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.dropout) - held_layers.append(self.model.classifier) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - # no shared params for sequence classification model - return [] - - -# BertForTokenClassification -class BertForTokenClassificationPolicy(BertPolicy): - def module_policy(self): - from transformers.models.bert.modeling_bert import BertForTokenClassification - - policy = super().module_policy() - - if self.shard_config.enable_tensor_parallelism: - addon_module = { - BertForTokenClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="dropout", - target_module=col_nn.DropoutForParallelInput, - ) - ] - ) - } - policy.update(addon_module) - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertForTokenClassification, - new_forward=BertPipelineForwards.bert_for_token_classification_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """ - get pipeline layers for current stage - """ - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.dropout) - held_layers.append(self.model.classifier) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - # no shared params for sequence classification model - return [] - - -# BertForNextSentencePrediction -class BertForNextSentencePredictionPolicy(BertPolicy): - def module_policy(self): - policy = super().module_policy() - from transformers.models.bert.modeling_bert import BertForNextSentencePrediction - - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertForNextSentencePrediction, - new_forward=BertPipelineForwards.bert_for_next_sentence_prediction_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """ - get pipeline layers for current stage - """ - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.cls) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - # no shared params for sequence classification model - return [] - - -# BertForMultipleChoice -class BertForMultipleChoicePolicy(BertPolicy): - def module_policy(self): - from transformers.models.bert.modeling_bert import BertForMultipleChoice - - policy = super().module_policy() - - if self.shard_config.enable_tensor_parallelism: - addon_module = { - BertForMultipleChoice: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="dropout", - target_module=col_nn.DropoutForParallelInput, - ) - ] - ) - } - policy.update(addon_module) - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertForMultipleChoice, - new_forward=BertPipelineForwards.bert_for_multiple_choice_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """ - get pipeline layers for current stage - """ - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.dropout) - held_layers.append(self.model.classifier) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - # no shared params for sequence classification model - return [] - - -class BertForQuestionAnsweringPolicy(BertPolicy): - def module_policy(self): - from transformers.models.bert.modeling_bert import BertForQuestionAnswering - - policy = super().module_policy() - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=BertForQuestionAnswering, - new_forward=BertPipelineForwards.bert_for_question_answering_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """ - get pipeline layers for current stage - """ - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.qa_outputs) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - # no shared params for sequence classification model - return [] +reexport_transformers_impl(globals(), __package__, "_bert") diff --git a/colossalai/shardformer/policies/blip2.py b/colossalai/shardformer/policies/blip2.py index a2f582e9daa6..d5d3a072dfeb 100644 --- a/colossalai/shardformer/policies/blip2.py +++ b/colossalai/shardformer/policies/blip2.py @@ -684,6 +684,21 @@ def module_policy(self): target_key=Blip2QFormerOutput, ) + from .._utils import is_hf5 + + if is_hf5(): + removed_dropout_paths = { + "self_attn.dropout", + "attention.attention.dropout", + "crossattention.attention.dropout", + } + for description in policy.values(): + if description.sub_module_replacement is not None: + description.sub_module_replacement = [ + replacement + for replacement in description.sub_module_replacement + if replacement.suffix not in removed_dropout_paths + ] return policy def postprocess(self): diff --git a/colossalai/shardformer/policies/command.py b/colossalai/shardformer/policies/command.py index 8a510307ef5d..8ee7c3febaac 100644 --- a/colossalai/shardformer/policies/command.py +++ b/colossalai/shardformer/policies/command.py @@ -76,13 +76,16 @@ def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: attribute_replacement=decoder_attribute_replacement, ) - self.append_or_create_method_replacement( - description={ - "forward": get_command_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), - }, - policy=policy, - target_key=attn_cls, - ) + from .._utils import is_hf5 + + if not is_hf5(): + self.append_or_create_method_replacement( + description={ + "forward": get_command_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), + }, + policy=policy, + target_key=attn_cls, + ) if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: if self.pipeline_stage_manager is None: self.append_or_create_method_replacement( diff --git a/colossalai/shardformer/policies/deepseek_v3.py b/colossalai/shardformer/policies/deepseek_v3.py index 8b386bcaa6e8..2adae90feeac 100644 --- a/colossalai/shardformer/policies/deepseek_v3.py +++ b/colossalai/shardformer/policies/deepseek_v3.py @@ -1,164 +1,9 @@ -from functools import partial -from typing import Callable, Dict, List, Union +r"""DeepSeek-v3 sharding policies, selected by the installed transformers major version. -import torch.nn as nn +``_deepseek_v3_tf5`` targets transformers v5; ``_deepseek_v3_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -from colossalai.shardformer.layer import FusedRMSNorm -from colossalai.shardformer.modeling.deepseek_v3 import ( - EpDeepseekV3MoE, - deepseek_v3_for_causal_lm_forward, - deepseek_v3_model_forward, -) -from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription +from colossalai._compat import reexport_transformers_impl -__all__ = ["DeepseekPolicy", "DeepseekForCausalLMPolicy"] - - -class DeepseekV3Policy(Policy): - def config_sanity_check(self): - assert not self.shard_config.enable_tensor_parallelism, "DeepSeekV3 does not support tensor parallelism" - assert not self.shard_config.enable_sequence_parallelism, "DeepSeekV3 does not support sequence parallelism" - if self.shard_config.pipeline_stage_manager: - assert not self.shard_config.pipeline_stage_manager.use_zbv, "DeepSeekV3 does not support ZBV" - - def preprocess(self): - return self.model - - def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: - - policy = {} - - # support gradient checkpointing - if self.shard_config.pipeline_stage_manager is None: - policy["DeepseekV3Model"] = ModulePolicyDescription( - method_replacement={"forward": deepseek_v3_model_forward} - ) - - if self.shard_config.expert_parallel_size > 1: - # expert parallel - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="mlp", - target_module=EpDeepseekV3MoE, - kwargs={ - "ep_group": self.shard_config.ep_group, - "moe_dp_group": self.shard_config.moe_dp_group, - }, - ) - ], - policy=policy, - target_key="DeepseekV3DecoderLayer", - ) - - # optimization configuration - if self.shard_config.enable_fused_normalization: - # TODO: prevent casting to fp32 - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="input_layernorm", - target_module=FusedRMSNorm, - ), - SubModuleReplacementDescription( - suffix="post_attention_layernorm", - target_module=FusedRMSNorm, - ), - ], - policy=policy, - target_key="DeepseekV3DecoderLayer", - ) - - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="norm", - target_module=FusedRMSNorm, - ), - policy=policy, - target_key="DeepseekV3Model", - ) - - return policy - - def postprocess(self): - return self.model - - def set_pipeline_forward(self, model_cls: str, new_forward: Callable, policy: Dict) -> None: - """If under pipeline parallel setting, replacing the original forward method of huggingface - to customized forward method, and add this changing to policy.""" - if self.pipeline_stage_manager: - num_layers = self.model.config.num_hidden_layers - stage_manager = self.pipeline_stage_manager - - layers_per_stage = stage_manager.distribute_layers(num_layers) - stage_index = stage_manager.get_stage_index(layers_per_stage) - method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)} - self.append_or_create_method_replacement( - description=method_replacement, policy=policy, target_key=model_cls - ) - - return - - def get_held_layers(self) -> List[nn.Module]: - """Get pipeline layers for current stage.""" - assert self.pipeline_stage_manager is not None - - module = self.model - if module.__class__.__name__.startswith("PeftModel"): - module = module.get_base_model() - if module.__class__.__name__ != "DeepseekV3Model": - module = module.model - - stage_manager = self.pipeline_stage_manager - - held_layers = [] - - if stage_manager.is_interleave: - assert stage_manager.num_model_chunks is not None - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_indices = stage_manager.get_stage_index(layers_per_stage) - stage_manager.stage_indices = stage_indices - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(module.embed_tokens) - for start_idx, end_idx in stage_indices: - held_layers.extend(module.layers[start_idx:end_idx]) - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - # for zbv, when is_first_stage (last fwd), we append norm - # for interleaved, when is_last_stage (last fwd), we also append norm - held_layers.append(module.norm) - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - if stage_manager.is_first_stage(): - held_layers.append(module.embed_tokens) - start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) - held_layers.extend(module.layers[start_idx:end_idx]) - if stage_manager.is_last_stage(): - held_layers.append(module.norm) - return held_layers - - -class DeepseekV3ModelPolicy(DeepseekV3Policy): - def module_policy(self): - policy = super().module_policy() - if self.shard_config.pipeline_stage_manager: - self.set_pipeline_forward("DeepseekV3Model", deepseek_v3_model_forward, policy) - return policy - - -class DeepseekV3ForCausalLMPolicy(DeepseekV3Policy): - def module_policy(self): - policy = super().module_policy() - if self.shard_config.pipeline_stage_manager: - self.set_pipeline_forward("DeepseekV3ForCausalLM", deepseek_v3_for_causal_lm_forward, policy) - return policy - - def get_held_layers(self): - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(self.model.lm_head) - elif stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.lm_head) - return held_layers +reexport_transformers_impl(globals(), __package__, "_deepseek_v3") diff --git a/colossalai/shardformer/policies/llama.py b/colossalai/shardformer/policies/llama.py index 9ad63dd7f051..4b8b4f3d9acb 100644 --- a/colossalai/shardformer/policies/llama.py +++ b/colossalai/shardformer/policies/llama.py @@ -1,538 +1,9 @@ -from functools import partial -from typing import Callable, Dict, List, Union +r"""Llama sharding policies, selected by the installed transformers major version. -import torch.nn as nn -from torch import Tensor -from torch.nn import Module +``_llama_tf5`` targets transformers v5; ``_llama_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -from colossalai.shardformer.layer import ( - FusedRMSNorm, - Linear1D_Col, - Linear1D_Row, - LinearWithGradAccum, - PaddingEmbedding, - PaddingLMHead, - RMSNorm, - VocabParallelEmbedding1D, - VocabParallelLMHead1D, -) +from colossalai._compat import reexport_transformers_impl -from ..modeling.llama import LlamaPipelineForwards, get_llama_flash_attention_forward -from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription - -__all__ = ["LlamaPolicy", "LlamaForCausalLMPolicy", "LlamaForSequenceClassificationPolicy"] - - -class LlamaPolicy(Policy): - def config_sanity_check(self): - pass - - def preprocess(self): - self.tie_weight = self.tie_weight_check() - self.origin_attn_implement = self.model.config._attn_implementation - return self.model - - def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: - from transformers.models.llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaModel - - policy = {} - embedding_cls = None - if self.shard_config.enable_tensor_parallelism: - embedding_cls = VocabParallelEmbedding1D - else: - if self.tie_weight: - embedding_cls = PaddingEmbedding - - if self.shard_config.enable_fused_normalization: - norm_cls = FusedRMSNorm - else: - norm_cls = RMSNorm - - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - sp_mode = self.shard_config.sequence_parallelism_mode or None - sp_size = self.shard_config.sequence_parallel_size or None - sp_group = self.shard_config.sequence_parallel_process_group or None - sp_partial_derived = sp_mode in ["split_gather", "ring"] - if sp_mode == "ring_attn" and not self.is_causal: - raise ValueError("Ring attention is only meant for causal language modeling.") - - tp_size = self.shard_config.tensor_parallel_size - # Modified by SP and TP - num_q_heads = self.model.config.num_attention_heads - num_kv_heads = getattr(self.model.config, "num_key_value_heads", None) - - if sp_mode == "all_to_all": - num_q_heads //= sp_size - decoder_attribute_replacement = {"num_heads": num_q_heads} - if num_kv_heads: - num_kv_heads //= sp_size - decoder_attribute_replacement["num_key_value_heads"] = num_kv_heads - - policy[LlamaAttention] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - ) - if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: - self.append_or_create_method_replacement( - description={ - "forward": get_llama_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), - }, - policy=policy, - target_key=LlamaAttention, - ) - - if self.pipeline_stage_manager is None: - self.append_or_create_method_replacement( - description={ - "forward": partial( - LlamaPipelineForwards.llama_model_forward, - shard_config=self.shard_config, - ), - }, - policy=policy, - target_key=LlamaModel, - ) - # enable tp, replace layer to tp Linear1D_Col,Linear1D_Row, - if self.shard_config.enable_tensor_parallelism: - assert ( - num_q_heads % tp_size == 0 - ), f"The number of attention heads must be divisible by tensor parallel size." - if hasattr(self.model.config, "num_key_value_heads"): - assert ( - num_kv_heads >= tp_size and num_kv_heads % tp_size == 0 - ), f"The number of key_value heads must be divisible by, and must not be less than tensor parallel size." - num_q_heads //= tp_size - decoder_attribute_replacement = { - "self_attn.hidden_size": self.model.config.hidden_size // tp_size, - "self_attn.num_heads": num_q_heads, - } - if getattr(self.model.config, "num_key_value_heads", False): - num_kv_heads //= tp_size - decoder_attribute_replacement["self_attn.num_key_value_heads"] = num_kv_heads - - policy[LlamaDecoderLayer] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=Linear1D_Row, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.gate_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.up_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.down_proj", - target_module=Linear1D_Row, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - ], - ) - - # not enable tp, replace layer to LinearWithGradAccum - elif use_zbv: - policy[LlamaDecoderLayer] = ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.gate_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.up_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.down_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - ], - ) - - if embedding_cls is not None: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="embed_tokens", - target_module=embedding_cls, - kwargs=( - { - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - } - if self.shard_config.enable_tensor_parallelism - else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} - ), - ), - policy=policy, - target_key=LlamaModel, - ) - - # optimization configuration - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="input_layernorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - SubModuleReplacementDescription( - suffix="post_attention_layernorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - ], - policy=policy, - target_key=LlamaDecoderLayer, - ) - - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="norm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - policy=policy, - target_key=LlamaModel, - ) - - return policy - - def postprocess(self): - return self.model - - def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: - """If under pipeline parallel setting, replacing the original forward method of huggingface - to customized forward method, and add this changing to policy.""" - if self.pipeline_stage_manager is None: - return - - stage_manager = self.pipeline_stage_manager - if self.model.__class__.__name__ == "LlamaModel": - module = self.model - else: - module = self.model.model - - if stage_manager.is_interleave: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) - } - - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_index = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial( - new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config - ) - } - - self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - assert self.pipeline_stage_manager is not None - - if self.model.__class__.__name__ == "LlamaModel": - module = self.model - else: - module = self.model.model - stage_manager = self.pipeline_stage_manager - - held_layers = [] - held_layers.append(module.rotary_emb) - if stage_manager.is_interleave: - assert stage_manager.num_model_chunks is not None - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_indices = stage_manager.get_stage_index(layers_per_stage) - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(module.embed_tokens) - for start_idx, end_idx in stage_indices: - held_layers.extend(module.layers[start_idx:end_idx]) - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(module.norm) - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - if stage_manager.is_first_stage(): - held_layers.append(module.embed_tokens) - start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) - held_layers.extend(module.layers[start_idx:end_idx]) - if stage_manager.is_last_stage(): - held_layers.append(module.norm) - - return held_layers - - -class LlamaModelPolicy(LlamaPolicy): - def module_policy(self): - policy = super().module_policy() - from transformers.models.llama.modeling_llama import LlamaModel - - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=LlamaModel, new_forward=LlamaPipelineForwards.llama_model_forward, policy=policy - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - held_layers = super().get_held_layers() - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in llama model""" - return [] - - -class LlamaForCausalLMPolicy(LlamaPolicy): - def module_policy(self): - from transformers import LlamaForCausalLM - - self.is_causal = True - policy = super().module_policy() - - if self.shard_config.enable_tensor_parallelism: - # add a new item for causal lm - new_item = { - LlamaForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=VocabParallelLMHead1D, - kwargs={ - "gather_output": not self.shard_config.parallel_output, - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - }, - ) - ], - ) - } - else: - new_item = { - LlamaForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=PaddingLMHead, - kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, - ) - ], - ) - } - policy.update(new_item) - - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=LlamaForCausalLM, new_forward=LlamaPipelineForwards.llama_for_causal_lm_forward, policy=policy - ) - elif self.shard_config.enable_tensor_parallelism or self.shard_config.enable_sequence_parallelism: - # Compute loss distributedly along the sequence dimension - new_item[LlamaForCausalLM].method_replacement = { - # "forward": get_lm_forward_with_dist_cross_entropy(self.shard_config) - "forward": partial(LlamaPipelineForwards.llama_for_causal_lm_forward, shard_config=self.shard_config) - } - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.lm_head) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - if self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv: - return [] - llama_model = self.model.model - if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: - if ( - id(llama_model.embed_tokens.weight) == id(self.model.lm_head.weight) - and self.pipeline_stage_manager.num_stages > 1 - ): - # tie weights - return [ - { - 0: llama_model.embed_tokens.weight, - self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, - } - ] - return [] - - -class LlamaForSequenceClassificationPolicy(LlamaPolicy): - def module_policy(self): - from transformers import LlamaForSequenceClassification - - policy = super().module_policy() - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - # enable tp, replace layer to tp Linear1D_Col,Linear1D_Row, - if self.shard_config.enable_tensor_parallelism: - # add a new item for sequence classification - new_item = { - LlamaForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="score", - target_module=Linear1D_Col, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - # enable tp, replace layer to LinearWithGradAccum - elif use_zbv: - # add a new item for sequence classification - new_item = { - LlamaForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="score", - target_module=LinearWithGradAccum, - kwargs=dict( - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - - # to be confirmed - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=LlamaForSequenceClassification, - new_forward=LlamaPipelineForwards.llama_for_sequence_classification_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.score) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in llama for sequence classification model""" - return [] +reexport_transformers_impl(globals(), __package__, "_llama") diff --git a/colossalai/shardformer/policies/mixtral.py b/colossalai/shardformer/policies/mixtral.py index a9584db9bc52..a2fdd5dbf635 100644 --- a/colossalai/shardformer/policies/mixtral.py +++ b/colossalai/shardformer/policies/mixtral.py @@ -1,491 +1,9 @@ -import warnings -from functools import partial -from typing import Callable, Dict, List, Union +r"""Mixtral sharding policies, selected by the installed transformers major version. -import torch.nn as nn -from torch import Tensor -from torch.nn import Module -from transformers.models.mixtral.modeling_mixtral import MixtralForCausalLM, MixtralModel +``_mixtral_tf5`` targets transformers v5; ``_mixtral_tf4`` is the implementation for the pinned v4 +(``transformers==4.51.3``). See :func:`colossalai._compat.reexport_transformers_impl`. +""" -from colossalai.shardformer.layer import ( - FusedRMSNorm, - Linear1D_Col, - Linear1D_Row, - LinearWithGradAccum, - PaddingEmbedding, - VocabParallelEmbedding1D, -) +from colossalai._compat import reexport_transformers_impl -# from colossalai.shardformer.layer import FusedRMSNorm, Linear1D_Col -# from colossalai.shardformer.layer.embedding import PaddingEmbedding, VocabParallelEmbedding1D -# from colossalai.shardformer.layer.linear import Linear1D_Row -from colossalai.shardformer.modeling.mixtral import ( - EPMixtralSparseMoeBlock, - MixtralPipelineForwards, - get_mixtral_flash_attention_forward, - get_mixtral_flash_attention_model_forward, -) -from colossalai.shardformer.policies.base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription - -__all__ = ["MixtralPolicy", "MixtralForCausalLMPolicy"] - - -class MixtralPolicy(Policy): - def config_sanity_check(self): - pass - - def preprocess(self): - self.tie_weight = self.tie_weight_check() - self.origin_attn_implement = self.model.config._attn_implementation - return self.model - - def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: - from transformers.models.mixtral.modeling_mixtral import MixtralAttention, MixtralDecoderLayer, MixtralModel - - policy = {} - - sp_mode = self.shard_config.sequence_parallelism_mode or None - sp_size = self.shard_config.sequence_parallel_size or None - sp_group = self.shard_config.sequence_parallel_process_group or None - sp_partial_derived = sp_mode in ["split_gather", "ring"] - tp_size = self.shard_config.tensor_parallel_size - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - # modified for both SP and TP - num_q_heads = self.model.config.num_attention_heads - num_kv_heads = getattr(self.model.config, "num_key_value_heads", None) - - if sp_mode == "all_to_all": - num_q_heads //= sp_size - decoder_attribute_replacement = { - "num_heads": num_q_heads, - } - if getattr(self.model.config, "num_key_value_heads", False): - num_kv_heads //= sp_size - decoder_attribute_replacement["num_key_value_heads"] = num_kv_heads - - policy[MixtralAttention] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - ) - if self.shard_config.enable_sequence_parallelism: - if self.pipeline_stage_manager is not None: - # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism - # if both are enabled, one of them will be ignored - raise NotImplementedError("Sequence parallelism is not supported with pipeline parallelism.") - self.append_or_create_method_replacement( - description={ - "forward": get_mixtral_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), - }, - policy=policy, - target_key=MixtralAttention, - ) - self.append_or_create_method_replacement( - description={ - "forward": get_mixtral_flash_attention_model_forward( - self.shard_config, - sp_mode=sp_mode, - sp_size=sp_size, - sp_group=sp_group, - ), - }, - policy=policy, - target_key=MixtralModel, - ) - - embedding_cls = None - if self.shard_config.enable_tensor_parallelism: - embedding_cls = VocabParallelEmbedding1D - else: - if self.tie_weight: - embedding_cls = PaddingEmbedding - - if self.shard_config.enable_tensor_parallelism: - # tensor parallelism for non-moe params - assert ( - self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of attention heads must be divisible by tensor parallel size." - assert ( - self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of key_value heads must be divisible by tensor parallel size." - num_q_heads //= tp_size - decoder_attribute_replacement = { - "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, - "self_attn.num_heads": num_q_heads, - } - if num_kv_heads: - num_kv_heads //= tp_size - decoder_attribute_replacement["self_attn.num_key_value_heads"] = num_kv_heads - - policy[MixtralDecoderLayer] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="block_sparse_moe.gate", - target_module=Linear1D_Col, - kwargs={ - "gather_output": True, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - ], - ) - - elif use_zbv: - policy[MixtralDecoderLayer] = ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="block_sparse_moe.gate", - target_module=LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - ], - ) - if embedding_cls is not None: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="embed_tokens", - target_module=embedding_cls, - kwargs=( - { - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - } - if self.shard_config.enable_tensor_parallelism - else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} - ), - ), - policy=policy, - target_key=MixtralModel, - ) - - if self.shard_config.ep_group: - # expert parallel - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="block_sparse_moe", - target_module=EPMixtralSparseMoeBlock, - kwargs={ - "ep_group": self.shard_config.ep_group, - "tp_group": self.shard_config.tensor_parallel_process_group, - "moe_dp_group": self.shard_config.moe_dp_group, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ) - ], - policy=policy, - target_key=MixtralDecoderLayer, - ) - - # optimization configuration - if self.shard_config.enable_fused_normalization: - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="input_layernorm", - target_module=FusedRMSNorm, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - SubModuleReplacementDescription( - suffix="post_attention_layernorm", - target_module=FusedRMSNorm, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - ], - policy=policy, - target_key=MixtralDecoderLayer, - ) - - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="norm", - target_module=FusedRMSNorm, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - policy=policy, - target_key=MixtralModel, - ) - - if self.shard_config.enable_flash_attention: - warnings.warn("Flash attention is natively supported in transformers, will ignore the flag.") - self.shard_config.enable_flash_attention = False - - return policy - - def postprocess(self): - return self.model - - def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: - """If under pipeline parallel setting, replacing the original forward method of huggingface - to customized forward method, and add this changing to policy.""" - if self.pipeline_stage_manager: - if self.shard_config.enable_sequence_parallelism: - # NOTE: we are replacing model forward for both sequence parallelism and pipeline parallelism - # if both are enabled, one of them will be ignored - raise NotImplementedError("Pipeline parallelism is not supported with sequence parallelism.") - stage_manager = self.pipeline_stage_manager - if self.model.__class__.__name__ == "MixtralModel": - module = self.model - else: - module = self.model.model - - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_index = stage_manager.get_stage_index(layers_per_stage) - method_replacement = {"forward": partial(new_forward, stage_manager=stage_manager, stage_index=stage_index)} - self.append_or_create_method_replacement( - description=method_replacement, policy=policy, target_key=model_cls - ) - - return - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - assert self.pipeline_stage_manager is not None - - if self.model.__class__.__name__ == "MixtralModel": - module = self.model - else: - module = self.model.model - stage_manager = self.pipeline_stage_manager - - held_layers = [] - held_layers.append(module.rotary_emb) - if stage_manager.is_interleave: - assert stage_manager.num_model_chunks is not None - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_indices = stage_manager.get_stage_index(layers_per_stage) - stage_manager.stage_indices = stage_indices - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(module.embed_tokens) - for start_idx, end_idx in stage_indices: - held_layers.extend(module.layers[start_idx:end_idx]) - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - # for zbv, when is_first_stage (last fwd), we append norm - # for interleaved, when is_last_stage (last fwd), we also append norm - held_layers.append(module.norm) - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - if stage_manager.is_first_stage(): - held_layers.append(module.embed_tokens) - start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) - held_layers.extend(module.layers[start_idx:end_idx]) - if stage_manager.is_last_stage(): - held_layers.append(module.norm) - return held_layers - - -class MixtralModelPolicy(MixtralPolicy): - def __init__(self) -> None: - super().__init__() - - def module_policy(self): - policy = super().module_policy() - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=MixtralModel, - new_forward=MixtralPipelineForwards.mixtral_model_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - held_layers = super().get_held_layers() - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in mixtral model""" - return [] - - -class MixtralForCausalLMPolicy(MixtralPolicy): - def module_policy(self): - policy = super().module_policy() - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - # TODO: assign pg mesh from plugin to all modules - if self.shard_config.enable_tensor_parallelism: - # add a new item for causal lm - new_item = { - MixtralForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=Linear1D_Col, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ], - ) - } - policy.update(new_item) - elif use_zbv: - new_item = { - MixtralForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=LinearWithGradAccum, - kwargs=dict( - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ], - ) - } - policy.update(new_item) - - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=MixtralForCausalLM, - new_forward=MixtralPipelineForwards.mixtral_for_causal_lm_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(self.model.lm_head) - elif stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.lm_head) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - mixtral_model = self.model.model - if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: - if ( - id(mixtral_model.embed_tokens.weight) == id(self.model.lm_head.weight) - and self.pipeline_stage_manager.num_stages > 1 - ): - # tie weights - return [ - { - 0: mixtral_model.embed_tokens.weight, - self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, - } - ] - return [] - - -class MixtralForSequenceClassificationPolicy(MixtralPolicy): - def module_policy(self): - from transformers import MixtralForSequenceClassification - - policy = super().module_policy() - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_tensor_parallelism: - # add a new item for sequence classification - new_item = { - MixtralForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="score", - target_module=Linear1D_Col, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - - if self.pipeline_stage_manager: - raise NotImplementedError - - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.score) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in mixtral for sequence classification model""" - return [] +reexport_transformers_impl(globals(), __package__, "_mixtral") diff --git a/colossalai/shardformer/policies/opt.py b/colossalai/shardformer/policies/opt.py index 50742b850b24..bfec4ff6b87a 100644 --- a/colossalai/shardformer/policies/opt.py +++ b/colossalai/shardformer/policies/opt.py @@ -50,12 +50,19 @@ def preprocess(self): return self.model def module_policy(self): - from transformers.models.opt.modeling_opt import OPTAttention, OPTDecoder, OPTDecoderLayer, OptFlashAttention2 + from transformers.models.opt.modeling_opt import OPTAttention, OPTDecoder, OPTDecoderLayer - ATTN_IMPLEMENTATION = { - "eager": OPTAttention, - "flash_attention_2": OptFlashAttention2, - } + from .._utils import is_hf5 + + if is_hf5(): + ATTN_IMPLEMENTATION = dict.fromkeys(("eager", "sdpa", "flash_attention_2"), OPTAttention) + else: + from transformers.models.opt.modeling_opt import OptFlashAttention2 + + ATTN_IMPLEMENTATION = { + "eager": OPTAttention, + "flash_attention_2": OptFlashAttention2, + } policy = {} diff --git a/colossalai/shardformer/policies/qwen2.py b/colossalai/shardformer/policies/qwen2.py index 32f7d37b190f..258e7c565613 100644 --- a/colossalai/shardformer/policies/qwen2.py +++ b/colossalai/shardformer/policies/qwen2.py @@ -1,556 +1,3 @@ -from functools import partial -from typing import Callable, Dict, List, Union +from colossalai._compat import reexport_transformers_impl -import torch.nn as nn -from torch import Tensor -from torch.nn import Module -from transformers.models.qwen2.modeling_qwen2 import ( - Qwen2Attention, - Qwen2DecoderLayer, - Qwen2ForCausalLM, - Qwen2ForSequenceClassification, - Qwen2Model, -) - -from colossalai.shardformer.layer import ( - FusedRMSNorm, - Linear1D_Col, - Linear1D_Row, - LinearWithGradAccum, - PaddingEmbedding, - RMSNorm, - VocabParallelEmbedding1D, - VocabParallelLMHead1D, -) - -from ..modeling.qwen2 import ( - Qwen2PipelineForwards, - get_lm_forward_with_dist_cross_entropy, - get_qwen2_flash_attention_forward, - get_qwen2_model_forward_for_flash_attn, -) -from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription - -__all__ = ["Qwen2Policy", "Qwen2ForCausalLMPolicy", "Qwen2ForSequenceClassificationPolicy"] - - -class Qwen2Policy(Policy): - def __init__(self) -> None: - super().__init__() - import transformers - from packaging.version import Version - - assert Version(transformers.__version__) >= Version( - "4.39.1" - ), "The Qwen2 model should run on a transformers version of 4.39.1." - - def config_sanity_check(self): - pass - - def preprocess(self): - self.tie_weight = self.tie_weight_check() - self.origin_attn_implement = self.model.config._attn_implementation - return self.model - - def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: - - policy = {} - - embedding_cls = None - if self.shard_config.enable_tensor_parallelism: - embedding_cls = VocabParallelEmbedding1D - else: - if self.tie_weight: - embedding_cls = PaddingEmbedding - norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm - - sp_mode = self.shard_config.sequence_parallelism_mode or None - sp_size = self.shard_config.sequence_parallel_size or None - sp_group = self.shard_config.sequence_parallel_process_group or None - sp_partial_derived = sp_mode in ["split_gather", "ring"] - if sp_mode == "all_to_all": - decoder_attribute_replacement = { - "num_heads": self.model.config.num_attention_heads // sp_size, - } - if getattr(self.model.config, "num_key_value_heads", False): - decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size - - policy[Qwen2Attention] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - ) - - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_tensor_parallelism: - assert ( - self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of attention heads must be divisible by tensor parallel size." - if hasattr(self.model.config, "num_key_value_heads"): - assert ( - self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of key_value heads must be divisible by tensor parallel size." - decoder_attribute_replacement = { - "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, - "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, - } - if getattr(self.model.config, "num_key_value_heads", False): - decoder_attribute_replacement["self_attn.num_key_value_heads"] = ( - self.model.config.num_key_value_heads // self.shard_config.tensor_parallel_size - ) - - policy[Qwen2DecoderLayer] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=Linear1D_Row, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.gate_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.up_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.down_proj", - target_module=Linear1D_Row, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - ], - ) - elif use_zbv: - policy[Qwen2DecoderLayer] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.gate_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.up_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.down_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - ], - ) - - if embedding_cls is not None: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="embed_tokens", - target_module=embedding_cls, - kwargs=( - { - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - } - if self.shard_config.enable_tensor_parallelism - else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} - ), - ), - policy=policy, - target_key=Qwen2Model, - ) - - # optimization configuration - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="input_layernorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - SubModuleReplacementDescription( - suffix="post_attention_layernorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - ], - policy=policy, - target_key=Qwen2DecoderLayer, - ) - - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="norm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - policy=policy, - target_key=Qwen2Model, - ) - - if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: - self.append_or_create_method_replacement( - description={ - "forward": get_qwen2_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), - }, - policy=policy, - target_key=Qwen2Attention, - ) - if self.pipeline_stage_manager is None: - # replace qwen2 model forward method - self.append_or_create_method_replacement( - description={ - "forward": get_qwen2_model_forward_for_flash_attn( - self.shard_config, sp_mode, sp_size, sp_group - ), - }, - policy=policy, - target_key=Qwen2Model, - ) - - return policy - - def postprocess(self): - return self.model - - def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: - """If under pipeline parallel setting, replacing the original forward method of huggingface - to customized forward method, and add this changing to policy.""" - if self.pipeline_stage_manager is None: - return - - stage_manager = self.pipeline_stage_manager - if self.model.__class__.__name__ == "Qwen2Model": - module = self.model - else: - module = self.model.model - - if stage_manager.is_interleave: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) - } - - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_index = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial( - new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config - ) - } - self.append_or_create_method_replacement( - description=method_replacement, policy=policy, target_key=model_cls - ) - - self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - assert self.pipeline_stage_manager is not None - - if self.model.__class__.__name__ == "Qwen2Model": - module = self.model - else: - module = self.model.model - - stage_manager = self.pipeline_stage_manager - - held_layers = [] - held_layers.append(module.rotary_emb) - if stage_manager.is_interleave: - assert stage_manager.num_model_chunks is not None - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_indices = stage_manager.get_stage_index(layers_per_stage) - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(module.embed_tokens) - for start_idx, end_idx in stage_indices: - held_layers.extend(module.layers[start_idx:end_idx]) - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(module.norm) - - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - if stage_manager.is_first_stage(): - held_layers.append(module.embed_tokens) - start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) - held_layers.extend(module.layers[start_idx:end_idx]) - if stage_manager.is_last_stage(): - held_layers.append(module.norm) - - return held_layers - - -class Qwen2ModelPolicy(Qwen2Policy): - def module_policy(self): - policy = super().module_policy() - - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=Qwen2Model, new_forward=Qwen2PipelineForwards.qwen2_model_forward, policy=policy - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - held_layers = super().get_held_layers() - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in Qwen2 model""" - return [] - - -class Qwen2ForCausalLMPolicy(Qwen2Policy): - def module_policy(self): - policy = super().module_policy() - setattr(self.shard_config, "causal_lm", True) - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_tensor_parallelism: - # add a new item for casual lm - new_item = { - Qwen2ForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=VocabParallelLMHead1D, - kwargs=dict( - gather_output=not self.shard_config.parallel_output, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ], - method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, - ) - } - policy.update(new_item) - elif use_zbv: - # add a new item for casual lm - new_item = { - Qwen2ForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=LinearWithGradAccum, - kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), - ), - SubModuleReplacementDescription( - suffix="lm_head", - target_module=VocabParallelLMHead1D, - kwargs={ - "gather_output": not self.shard_config.parallel_output, - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - }, - ), - ], - method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, - ) - } - policy.update(new_item) - - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=Qwen2ForCausalLM, new_forward=Qwen2PipelineForwards.qwen2_for_causal_lm_forward, policy=policy - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.lm_head) - else: - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.lm_head) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - qwen2_model = self.model.model - if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: - if ( - id(qwen2_model.embed_tokens.weight) == id(self.model.lm_head.weight) - and self.pipeline_stage_manager.num_stages > 1 - ): - # tie weights - return [ - { - 0: qwen2_model.embed_tokens.weight, - self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, - } - ] - return [] - - -class Qwen2ForSequenceClassificationPolicy(Qwen2Policy): - def module_policy(self): - policy = super().module_policy() - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - if self.shard_config.enable_tensor_parallelism: - # add a new item for sequence classification - new_item = { - Qwen2ForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="score", - target_module=Linear1D_Col, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - elif use_zbv: - new_item = { - Qwen2ForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="score", - target_module=LinearWithGradAccum, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - # to be confirmed - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=Qwen2ForSequenceClassification, - new_forward=Qwen2PipelineForwards.qwen2_for_sequence_classification_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.score) - else: - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.score) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in Qwen2 for sequence classification model""" - return [] +reexport_transformers_impl(globals(), __package__, "_qwen2") diff --git a/colossalai/shardformer/policies/qwen3.py b/colossalai/shardformer/policies/qwen3.py index e9cc9543278a..19d46e525ca6 100644 --- a/colossalai/shardformer/policies/qwen3.py +++ b/colossalai/shardformer/policies/qwen3.py @@ -1,541 +1,3 @@ -# Modifed from qwen2 policy -from functools import partial -from typing import Callable, Dict, List, Union +from colossalai._compat import reexport_transformers_impl -import torch.nn as nn -from torch import Tensor -from torch.nn import Module -from transformers.models.qwen3.modeling_qwen3 import ( - Qwen3Attention, - Qwen3DecoderLayer, - Qwen3ForCausalLM, - Qwen3ForSequenceClassification, - Qwen3Model, -) - -from colossalai.shardformer.layer import ( - FusedRMSNorm, - Linear1D_Col, - Linear1D_Row, - LinearWithGradAccum, - PaddingEmbedding, - RMSNorm, - VocabParallelEmbedding1D, -) - -from ..modeling.qwen3 import ( - Qwen3PipelineForwards, - get_lm_forward_with_dist_cross_entropy, - get_qwen3_flash_attention_forward, - get_qwen3_model_forward_for_flash_attn, -) -from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription - -__all__ = ["Qwen3Policy", "Qwen3ForCausalLMPolicy", "Qwen3ForSequenceClassificationPolicy"] - - -class Qwen3Policy(Policy): - def __init__(self) -> None: - super().__init__() - import transformers - from packaging.version import Version - - assert Version(transformers.__version__) >= Version( - "4.51.0" - ), "The Qwen3 model should run on a transformers version of 4.51.0 or higher." - - def config_sanity_check(self): - pass - - def preprocess(self): - self.tie_weight = self.tie_weight_check() - self.origin_attn_implement = self.model.config._attn_implementation - return self.model - - def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: - - policy = {} - - embedding_cls = None - if self.shard_config.enable_tensor_parallelism: - embedding_cls = VocabParallelEmbedding1D - else: - if self.tie_weight: - embedding_cls = PaddingEmbedding - norm_cls = FusedRMSNorm if self.shard_config.enable_fused_normalization else RMSNorm - - sp_mode = self.shard_config.sequence_parallelism_mode or None - sp_size = self.shard_config.sequence_parallel_size or None - sp_group = self.shard_config.sequence_parallel_process_group or None - sp_partial_derived = sp_mode in ["split_gather", "ring"] - if sp_mode == "all_to_all": - decoder_attribute_replacement = { - "num_heads": self.model.config.num_attention_heads // sp_size, - } - if getattr(self.model.config, "num_key_value_heads", False): - decoder_attribute_replacement["num_key_value_heads"] = self.model.config.num_key_value_heads // sp_size - - policy[Qwen3Attention] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - ) - - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_tensor_parallelism: - assert ( - self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of attention heads must be divisible by tensor parallel size." - if hasattr(self.model.config, "num_key_value_heads"): - assert ( - self.model.config.num_key_value_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of key_value heads must be divisible by tensor parallel size." - decoder_attribute_replacement = { - "self_attn.hidden_size": self.model.config.hidden_size // self.shard_config.tensor_parallel_size, - "self_attn.num_heads": self.model.config.num_attention_heads // self.shard_config.tensor_parallel_size, - } - if getattr(self.model.config, "num_key_value_heads", False): - decoder_attribute_replacement["self_attn.num_key_value_heads"] = ( - self.model.config.num_key_value_heads // self.shard_config.tensor_parallel_size - ) - - policy[Qwen3DecoderLayer] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=Linear1D_Row, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.gate_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.up_proj", - target_module=Linear1D_Col, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.down_proj", - target_module=Linear1D_Row, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - ], - ) - elif use_zbv: - policy[Qwen3DecoderLayer] = ModulePolicyDescription( - attribute_replacement=decoder_attribute_replacement, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="self_attn.o_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.gate_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.up_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - SubModuleReplacementDescription( - suffix="mlp.down_proj", - target_module=LinearWithGradAccum, - kwargs=dict( - seq_parallel_mode=sp_mode, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ), - ], - ) - - if embedding_cls is not None: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="embed_tokens", - target_module=embedding_cls, - kwargs=( - { - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - } - if self.shard_config.enable_tensor_parallelism - else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} - ), - ), - policy=policy, - target_key=Qwen3Model, - ) - - # optimization configuration - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="input_layernorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - SubModuleReplacementDescription( - suffix="post_attention_layernorm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - ], - policy=policy, - target_key=Qwen3DecoderLayer, - ) - - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="norm", - target_module=norm_cls, - kwargs={"sp_partial_derived": sp_partial_derived}, - ), - policy=policy, - target_key=Qwen3Model, - ) - - if self.shard_config.enable_flash_attention or self.shard_config.enable_sequence_parallelism: - self.append_or_create_method_replacement( - description={ - "forward": get_qwen3_flash_attention_forward(self.shard_config, sp_mode, sp_size, sp_group), - }, - policy=policy, - target_key=Qwen3Attention, - ) - if self.pipeline_stage_manager is None: - # replace qwen3 model forward method - self.append_or_create_method_replacement( - description={ - "forward": get_qwen3_model_forward_for_flash_attn( - self.shard_config, sp_mode, sp_size, sp_group - ), - }, - policy=policy, - target_key=Qwen3Model, - ) - - return policy - - def postprocess(self): - return self.model - - def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: - """If under pipeline parallel setting, replacing the original forward method of huggingface - to customized forward method, and add this changing to policy.""" - if self.pipeline_stage_manager is None: - return - - stage_manager = self.pipeline_stage_manager - if self.model.__class__.__name__ == "Qwen3Model": - module = self.model - else: - module = self.model.model - - if stage_manager.is_interleave: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_manager.stage_indices = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial(new_forward, stage_manager=stage_manager, shard_config=self.shard_config) - } - - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_index = stage_manager.get_stage_index(layers_per_stage) - method_replacement = { - "forward": partial( - new_forward, stage_manager=stage_manager, stage_index=stage_index, shard_config=self.shard_config - ) - } - self.append_or_create_method_replacement( - description=method_replacement, policy=policy, target_key=model_cls - ) - - self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - assert self.pipeline_stage_manager is not None - - if self.model.__class__.__name__ == "Qwen3Model": - module = self.model - else: - module = self.model.model - - stage_manager = self.pipeline_stage_manager - - held_layers = [] - held_layers.append(module.rotary_emb) - if stage_manager.is_interleave: - assert stage_manager.num_model_chunks is not None - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - stage_indices = stage_manager.get_stage_index(layers_per_stage) - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(module.embed_tokens) - for start_idx, end_idx in stage_indices: - held_layers.extend(module.layers[start_idx:end_idx]) - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(module.norm) - - else: - layers_per_stage = stage_manager.distribute_layers(len(module.layers)) - if stage_manager.is_first_stage(): - held_layers.append(module.embed_tokens) - start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) - held_layers.extend(module.layers[start_idx:end_idx]) - if stage_manager.is_last_stage(): - held_layers.append(module.norm) - - return held_layers - - -class Qwen3ModelPolicy(Qwen3Policy): - def module_policy(self): - policy = super().module_policy() - - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=Qwen3Model, new_forward=Qwen3PipelineForwards.qwen3_model_forward, policy=policy - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - held_layers = super().get_held_layers() - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in Qwen3 model""" - return [] - - -class Qwen3ForCausalLMPolicy(Qwen3Policy): - def module_policy(self): - policy = super().module_policy() - setattr(self.shard_config, "causal_lm", True) - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_tensor_parallelism: - # add a new item for casual lm - new_item = { - Qwen3ForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=Linear1D_Col, - kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), - ) - ], - method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, - ) - } - policy.update(new_item) - elif use_zbv: - # add a new item for casual lm - new_item = { - Qwen3ForCausalLM: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="lm_head", - target_module=LinearWithGradAccum, - kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), - ) - ], - method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, - ) - } - policy.update(new_item) - - if self.pipeline_stage_manager: - # set None as default - self.set_pipeline_forward( - model_cls=Qwen3ForCausalLM, new_forward=Qwen3PipelineForwards.qwen3_for_causal_lm_forward, policy=policy - ) - - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.lm_head) - else: - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.lm_head) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - qwen3_model = self.model.model - if self.pipeline_stage_manager and self.pipeline_stage_manager.num_stages > 1: - if ( - id(qwen3_model.embed_tokens.weight) == id(self.model.lm_head.weight) - and self.pipeline_stage_manager.num_stages > 1 - ): - # tie weights - return [ - { - 0: qwen3_model.embed_tokens.weight, - self.pipeline_stage_manager.num_stages - 1: self.model.lm_head.weight, - } - ] - return [] - - -class Qwen3ForSequenceClassificationPolicy(Qwen3Policy): - def module_policy(self): - policy = super().module_policy() - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - if self.shard_config.enable_tensor_parallelism: - # add a new item for sequence classification - new_item = { - Qwen3ForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="score", - target_module=Linear1D_Col, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - elif use_zbv: - new_item = { - Qwen3ForSequenceClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="score", - target_module=LinearWithGradAccum, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - # to be confirmed - if self.pipeline_stage_manager: - self.set_pipeline_forward( - model_cls=Qwen3ForSequenceClassification, - new_forward=Qwen3PipelineForwards.qwen3_for_sequence_classification_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[Module]: - """Get pipeline layers for current stage.""" - stage_manager = self.pipeline_stage_manager - held_layers = super().get_held_layers() - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.score) - else: - if stage_manager.is_last_stage(ignore_chunk=True): - held_layers.append(self.model.score) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - """No shared params in Qwen3 for sequence classification model""" - return [] +reexport_transformers_impl(globals(), __package__, "_qwen3") diff --git a/colossalai/shardformer/policies/vit.py b/colossalai/shardformer/policies/vit.py index 420ea286fd0a..cf60db7af5e5 100644 --- a/colossalai/shardformer/policies/vit.py +++ b/colossalai/shardformer/policies/vit.py @@ -1,434 +1,3 @@ -import warnings -from typing import Callable, Dict, List, Union +from colossalai._compat import reexport_transformers_impl -import torch.nn as nn - -import colossalai.shardformer.layer as col_nn -from colossalai.shardformer.layer import DropoutForReplicatedInput, Linear1D_Col - -from ..modeling.jit import get_jit_fused_dropout_add_func -from ..modeling.vit import ( - ViTForImageClassification_pipeline_forward, - ViTForMaskedImageModeling_pipeline_forward, - ViTModel_pipeline_forward, - get_jit_fused_vit_intermediate_forward, - get_jit_fused_vit_output_forward, - get_vit_flash_self_attention_forward, -) -from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription - -__all__ = ["ViTPolicy", "ViTModelPolicy", "ViTForImageClassificationPolicy", "ViTForMaskedImageModelingPolicy"] - - -class ViTPolicy(Policy): - def config_sanity_check(self): - pass - - def preprocess(self): - self.enable_bias_gelu_fused = self.shard_config.enable_jit_fused and self.model.config.hidden_act == "gelu" - return self.model - - def module_policy(self) -> Dict[Union[str, nn.Module], ModulePolicyDescription]: - from transformers.models.vit.modeling_vit import ( - ViTEmbeddings, - ViTIntermediate, - ViTLayer, - ViTOutput, - ViTSelfAttention, - ) - - policy = {} - - if self.shard_config.enable_sequence_parallelism: - self.shard_config.enable_sequence_parallelism = False - warnings.warn("Vit doesn't support sequence parallelism now, will ignore the sequence parallelism flag.") - - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_tensor_parallelism: - assert ( - self.model.config.num_attention_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of attention heads must be divisible by tensor parallel size." - policy[ViTEmbeddings] = ModulePolicyDescription( - attribute_replacement={}, - param_replacement=[], - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="dropout", - target_module=DropoutForReplicatedInput, - ) - ], - ) - - policy[ViTLayer] = ModulePolicyDescription( - attribute_replacement={ - "attention.attention.num_attention_heads": self.model.config.num_attention_heads - // self.shard_config.tensor_parallel_size, - "attention.attention.all_head_size": self.model.config.hidden_size - // self.shard_config.tensor_parallel_size, - }, - param_replacement=[], - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="attention.attention.query", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.attention.key", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.attention.value", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.output.dense", - target_module=col_nn.Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.output.dropout", - target_module=col_nn.DropoutForReplicatedInput, - ), - SubModuleReplacementDescription( - suffix="intermediate.dense", - target_module=col_nn.Linear1D_Col, - kwargs={ - "skip_bias_add": self.enable_bias_gelu_fused, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dense", - target_module=col_nn.Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dropout", - target_module=col_nn.DropoutForReplicatedInput, - ), - ], - ) - if self.enable_bias_gelu_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_vit_intermediate_forward(), - }, - policy=policy, - target_key=ViTIntermediate, - ) - elif use_zbv: - policy[ViTEmbeddings] = ModulePolicyDescription( - attribute_replacement={}, - param_replacement=[], - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="dropout", - target_module=DropoutForReplicatedInput, - ) - ], - ) - - policy[ViTLayer] = ModulePolicyDescription( - param_replacement=[], - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="attention.attention.query", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.attention.key", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.attention.value", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.attention.dropout", - target_module=col_nn.DropoutForParallelInput, - ), - SubModuleReplacementDescription( - suffix="attention.output.dense", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="attention.output.dropout", - target_module=col_nn.DropoutForReplicatedInput, - ), - SubModuleReplacementDescription( - suffix="intermediate.dense", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "skip_bias_add": self.enable_bias_gelu_fused, - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dense", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="output.dropout", - target_module=col_nn.DropoutForReplicatedInput, - ), - ], - ) - if self.enable_bias_gelu_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_vit_intermediate_forward(), - }, - policy=policy, - target_key=ViTIntermediate, - ) - # use flash attention - if self.shard_config.enable_flash_attention: - self.append_or_create_method_replacement( - description={ - "forward": get_vit_flash_self_attention_forward(), - }, - policy=policy, - target_key=ViTSelfAttention, - ) - - # use jit fused operator - if self.shard_config.enable_jit_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_vit_output_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=ViTOutput, - ) - - return policy - - def new_model_class(self): - return None - - def postprocess(self): - return self.model - - def get_held_layers(self) -> List[nn.Module]: - """Get pipeline layers for current stage.""" - assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" - - if self.model.__class__.__name__ == "ViTModel": - module = self.model - else: - module = self.model.vit - stage_manager = self.pipeline_stage_manager - - held_layers = [] - if stage_manager.is_interleave: - assert stage_manager.num_model_chunks is not None - layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) - stage_indices = stage_manager.get_stage_index(layers_per_stage) - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(module.embeddings) - for start_idx, end_idx in stage_indices: - held_layers.extend(module.encoder.layer[start_idx:end_idx]) - else: - layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) - if stage_manager.is_first_stage(): - held_layers.append(module.embeddings) - start_idx, end_idx = stage_manager.get_stage_index(layers_per_stage) - held_layers.extend(module.encoder.layer[start_idx:end_idx]) - return held_layers - - def set_pipeline_forward(self, model_cls: nn.Module, pipeline_forward: Callable, policy: Dict): - if self.pipeline_stage_manager: - stage_manager = self.pipeline_stage_manager - if self.model.__class__.__name__ == "ViTModel": - module = self.model - else: - module = self.model.vit - - layers_per_stage = stage_manager.distribute_layers(len(module.encoder.layer)) - stage_index = stage_manager.get_stage_index(layers_per_stage) - method_replacement = {"forward": pipeline_forward(stage_manager=stage_manager, stage_index=stage_index)} - self.append_or_create_method_replacement( - description=method_replacement, policy=policy, target_key=model_cls - ) - - -# ViTModel -class ViTModelPolicy(ViTPolicy): - def module_policy(self): - from transformers.models.vit.modeling_vit import ViTModel - - policy = super().module_policy() - - if self.shard_config.pipeline_stage_manager is not None: - self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) - return policy - - def get_held_layers(self) -> List[nn.Module]: - held_layers = super().get_held_layers() - assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" - - module = self.model - stage_manager = self.pipeline_stage_manager - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(module.layernorm) - held_layers.append(module.pooler) - else: - if stage_manager.is_last_stage(): - held_layers.append(module.layernorm) - held_layers.append(module.pooler) - - return held_layers - - -# ViTForImageClassification -class ViTForImageClassificationPolicy(ViTPolicy): - def module_policy(self): - from transformers.models.vit.modeling_vit import ViTForImageClassification, ViTModel - - policy = super().module_policy() - - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - if self.shard_config.enable_tensor_parallelism: - new_item = { - ViTForImageClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="classifier", - target_module=Linear1D_Col, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - elif use_zbv: - new_item = { - ViTForImageClassification: ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="classifier", - target_module=col_nn.LinearWithGradAccum, - kwargs=dict( - gather_output=True, - fp8_communication=self.shard_config.fp8_communication, - use_zbv=use_zbv, - ), - ) - ] - ) - } - policy.update(new_item) - if self.shard_config.pipeline_stage_manager is not None: - self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) - self.set_pipeline_forward( - model_cls=ViTForImageClassification, - pipeline_forward=ViTForImageClassification_pipeline_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[nn.Module]: - held_layers = super().get_held_layers() - assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" - - module = self.model.vit - stage_manager = self.pipeline_stage_manager - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(module.layernorm) - held_layers.append(self.model.classifier) - else: - if stage_manager.is_last_stage(): - held_layers.append(module.layernorm) - held_layers.append(self.model.classifier) - - return held_layers - - -# ViTForMaskedImageModeling -class ViTForMaskedImageModelingPolicy(ViTPolicy): - def module_policy(self): - from transformers.models.vit.modeling_vit import ViTForMaskedImageModeling, ViTModel - - policy = super().module_policy() - - if self.shard_config.pipeline_stage_manager is not None: - self.set_pipeline_forward(model_cls=ViTModel, pipeline_forward=ViTModel_pipeline_forward, policy=policy) - self.set_pipeline_forward( - model_cls=ViTForMaskedImageModeling, - pipeline_forward=ViTForMaskedImageModeling_pipeline_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[nn.Module]: - held_layers = super().get_held_layers() - assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" - - module = self.model.vit - stage_manager = self.pipeline_stage_manager - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(module.layernorm) - held_layers.append(self.model.decoder) - else: - if stage_manager.is_last_stage(): - held_layers.append(module.layernorm) - held_layers.append(self.model.decoder) - - return held_layers +reexport_transformers_impl(globals(), __package__, "_vit") diff --git a/colossalai/shardformer/policies/whisper.py b/colossalai/shardformer/policies/whisper.py index 5d9b38e6fe98..3285c6f3fde6 100644 --- a/colossalai/shardformer/policies/whisper.py +++ b/colossalai/shardformer/policies/whisper.py @@ -1,845 +1,3 @@ -import warnings -from functools import partial -from typing import Callable, Dict, List, Tuple +from colossalai._compat import reexport_transformers_impl -import numpy as np -import torch.nn as nn -from torch import Tensor - -import colossalai.shardformer.layer as col_nn - -from ..modeling.jit import get_jit_fused_dropout_add_func -from ..modeling.whisper import ( - WhisperPipelineForwards, - get_jit_fused_whisper_decoder_layer_forward, - get_jit_fused_whisper_encoder_layer_forward, - get_whisper_decoder_forward_for_flash_attention, - get_whisper_flash_attention_forward, -) -from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription - -__all__ = [ - "WhisperPolicy", - "WhisperModelPolicy", - "WhisperForConditionalGenerationPolicy", - "WhisperForAudioClassificationPolicy", -] - - -class WhisperPolicy(Policy): - def __init__(self) -> None: - super().__init__() - - def config_sanity_check(self): - pass - - def preprocess(self): - # reshape the embedding layer - r""" - Reshape the Embedding layer to make the embedding dimension divisible by world_size - """ - self.tie_weight = self.tie_weight_check() - return self.model - - def module_policy(self): - from transformers.models.whisper.modeling_whisper import ( - WhisperAttention, - WhisperDecoder, - WhisperDecoderLayer, - WhisperEncoder, - WhisperEncoderLayer, - WhisperFlashAttention2, - WhisperSdpaAttention, - ) - - policy = {} - - embedding_cls = None - if self.shard_config.enable_tensor_parallelism: - embedding_cls = col_nn.VocabParallelEmbedding1D - else: - if self.tie_weight: - embedding_cls = col_nn.PaddingEmbedding - - if self.shard_config.enable_fused_normalization: - norm_cls = col_nn.FusedLayerNorm - else: - norm_cls = col_nn.LayerNorm - - if self.shard_config.enable_sequence_parallelism: - self.shard_config.enable_sequence_parallelism = False - warnings.warn( - "Whisper doesn't support sequence parallelism now, will ignore the sequence parallelism flag." - ) - - use_zbv = self.pipeline_stage_manager is not None and self.pipeline_stage_manager.use_zbv - - # TODO using the jit fused add_and_dropout affect the accuracy - if self.shard_config.enable_jit_fused: - self.shard_config.enable_jit_fused = False - warnings.warn("Whisper doesn't support jit fused operator now, will ignore the jit fused operator flag.") - - if self.shard_config.enable_tensor_parallelism: - assert ( - self.model.config.encoder_attention_heads % self.shard_config.tensor_parallel_size == 0 - ), f"The number of attention heads must be divisible by tensor parallel size." - policy[WhisperEncoderLayer] = ModulePolicyDescription( - attribute_replacement={ - "self_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, - "self_attn.num_heads": self.model.config.encoder_attention_heads - // self.shard_config.tensor_parallel_size, - }, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.out_proj", - target_module=col_nn.Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc1", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc2", - target_module=col_nn.Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - ], - ) - - policy[WhisperDecoderLayer] = ModulePolicyDescription( - attribute_replacement={ - "self_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, - "self_attn.num_heads": self.model.config.decoder_attention_heads - // self.shard_config.tensor_parallel_size, - "encoder_attn.embed_dim": self.model.config.d_model // self.shard_config.tensor_parallel_size, - "encoder_attn.num_heads": self.model.config.encoder_attention_heads - // self.shard_config.tensor_parallel_size, - }, - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.out_proj", - target_module=col_nn.Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.q_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.k_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.v_proj", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.out_proj", - target_module=col_nn.Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc1", - target_module=col_nn.Linear1D_Col, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc2", - target_module=col_nn.Linear1D_Row, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - ], - ) - elif use_zbv: - policy[WhisperEncoderLayer] = ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.out_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc1", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc2", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - ], - ) - - policy[WhisperDecoderLayer] = ModulePolicyDescription( - sub_module_replacement=[ - SubModuleReplacementDescription( - suffix="self_attn.q_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.k_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.v_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="self_attn.out_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.q_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.k_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.v_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="encoder_attn.out_proj", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc1", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - SubModuleReplacementDescription( - suffix="fc2", - target_module=col_nn.LinearWithGradAccum, - kwargs={ - "fp8_communication": self.shard_config.fp8_communication, - "use_zbv": use_zbv, - }, - ), - ], - ) - - if embedding_cls is not None: - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="embed_tokens", - target_module=embedding_cls, - kwargs=( - { - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - } - if self.shard_config.enable_tensor_parallelism - else {"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by} - ), - ), - ], - policy=policy, - target_key=WhisperDecoder, - ) - - # optimization configuration - # Handle encoder layer - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="self_attn_layer_norm", - target_module=norm_cls, - ), - SubModuleReplacementDescription( - suffix="final_layer_norm", - target_module=norm_cls, - ), - ], - policy=policy, - target_key=WhisperEncoderLayer, - ) - - # Handle decoder layer - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="self_attn_layer_norm", - target_module=norm_cls, - ), - SubModuleReplacementDescription( - suffix="final_layer_norm", - target_module=norm_cls, - ), - ], - policy=policy, - target_key=WhisperDecoderLayer, - ) - - # handle encoder layer - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="layer_norm", - target_module=norm_cls, - ) - ], - policy=policy, - target_key=WhisperEncoder, - ) - - # handle decoder layer - self.append_or_create_submodule_replacement( - description=[ - SubModuleReplacementDescription( - suffix="layer_norm", - target_module=norm_cls, - ) - ], - policy=policy, - target_key=WhisperDecoder, - ) - - # enable flash attention - if self.shard_config.enable_flash_attention: - self.append_or_create_method_replacement( - description={ - "forward": get_whisper_flash_attention_forward(), - }, - policy=policy, - target_key=WhisperAttention, - ) - self.append_or_create_method_replacement( - description={ - "forward": get_whisper_flash_attention_forward(), - }, - policy=policy, - target_key=WhisperFlashAttention2, - ) - self.append_or_create_method_replacement( - description={ - "forward": get_whisper_flash_attention_forward(), - }, - policy=policy, - target_key=WhisperSdpaAttention, - ) - if not self.shard_config.pipeline_stage_manager: - self.append_or_create_method_replacement( - description={ - "forward": get_whisper_decoder_forward_for_flash_attention(self.shard_config), - }, - policy=policy, - target_key=WhisperDecoder, - ) - - # use jit fused operator - if self.shard_config.enable_jit_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_whisper_decoder_layer_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=WhisperDecoderLayer, - ) - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_whisper_encoder_layer_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=WhisperEncoderLayer, - ) - - return policy - - def add_lm_head_policy(self, base_policy): - from transformers.models.whisper.modeling_whisper import WhisperForConditionalGeneration - - # optimize for tensor parallelism - if self.shard_config.enable_tensor_parallelism: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="proj_out", - target_module=col_nn.VocabParallelLMHead1D, - kwargs={ - "gather_output": True, - "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, - "fp8_communication": self.shard_config.fp8_communication, - }, - ), - policy=base_policy, - target_key=WhisperForConditionalGeneration, - ) - else: - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="proj_out", - target_module=col_nn.PaddingLMHead, - kwargs={"make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by}, - ), - policy=base_policy, - target_key=WhisperForConditionalGeneration, - ) - - return base_policy - - def postprocess(self): - return self.model - - def distribute_whisper_layers( - self, num_encoder_layers: int, num_decoder_layers: int, num_stages: int - ) -> Tuple[List[int], int]: - """ - Distribute whisper layers into stages when pipeline parallel is used. - Return the layer distribution as a list and the starting stage of decoder. - If decoder doesn't exist, returned decoder starting stage is set to num_encoder_layers. - """ - stage_manager = self.pipeline_stage_manager - assert stage_manager is not None, "pipeline_stage_manager is None" - - # number of encoder layers must be a positive integer - if num_encoder_layers <= 0: - raise ValueError("The number of encoder layers for whisper must be a positive integer.") - - # number of layers should be large enough to fill in every stage - if num_encoder_layers + num_decoder_layers < num_stages: - raise ValueError("The total number of layers can't be smaller than number of stages.") - - # in the case of whisperEncoderModel, set decoder starting stage to num_stages since it doesn't exist - if num_decoder_layers == 0: - return stage_manager.distribute_layers(num_encoder_layers, num_stages), num_stages - - # the number of stages distributed between encoder and decoder is optimized in this way: - # num_encoder_stages = argmin(abs(num_encoder_layers / encoder_stages - num_decoder_layers / decoder_stages)) - # s.t. num_encoder_stages + num_decoder_stages = num_stages, num_encoder_stages >= 1, num_decoder_stages >= 1 - def objective(num_encoder_stages): - return abs(num_encoder_layers / num_encoder_stages - num_decoder_layers / (num_stages - num_encoder_stages)) - - num_encoder_stages = np.argmin([objective(i) for i in range(1, num_stages)]) + 1 - num_decoder_stages = num_stages - num_encoder_stages - - encoder_distribution = stage_manager.distribute_layers(num_encoder_layers, num_encoder_stages) - decoder_distribution = stage_manager.distribute_layers(num_decoder_layers, num_decoder_stages) - return encoder_distribution + decoder_distribution, num_encoder_stages - - def get_whisper_stage_index( - self, layers_per_stage: List[int], stage: int, decoder_starting_stage: int - ) -> Tuple[int, int]: - """ - Input the distribution of layers among stages, the current stage and the first stage of decoder. - Return the starting/ending idx of layers in encoder/decoder - """ - stage_manager = self.pipeline_stage_manager - assert stage_manager is not None, "pipeline_stage_manager is None" - - if stage < decoder_starting_stage: - return stage_manager.get_stage_index(layers_per_stage[:decoder_starting_stage], stage) - else: - return stage_manager.get_stage_index( - layers_per_stage[decoder_starting_stage:], - stage - decoder_starting_stage, - ) - - def get_held_layers(self) -> List[nn.Module]: - assert self.pipeline_stage_manager is not None, "pipeline_stage_manager is None" - stage_manager = self.pipeline_stage_manager - - if self.model.__class__.__name__ == "WhisperModel": - model = self.model - elif self.model.__class__.__name__ == "WhisperForConditionalGeneration": - model = self.model.model - else: - model = None - - if model: - encoder = self.model.get_encoder() - decoder = self.model.get_decoder() - else: - # whisper for audio classification holds encoder only - encoder = self.model.encoder - decoder = None - - num_encoder_layers = len(encoder.layers) - if decoder: - num_decoder_layers = len(decoder.layers) - else: - num_decoder_layers = 0 - - held_layers = [] - if stage_manager.is_interleave: - layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( - num_encoder_layers, num_decoder_layers, stage_manager.num_stages - ) - stage_indices = self.get_whisper_stage_index(layers_per_stage, stage_manager.stage, decoder_starting_stage) - - if stage_manager.stage < decoder_starting_stage: - # current stage is in whisper's encoder - if stage_manager.is_first_stage(ignore_chunk=True): - held_layers.append(encoder.embed_positions) - held_layers.append(encoder.conv1) - held_layers.append(encoder.conv2) - # interleaved: not use_zbv & stage_manager.stage == decoder_starting_stage - 1 - # zbv: use_zbv & stage_manager.stage == first stage - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and decoder_starting_stage - 1 - ): - held_layers.append(encoder.layer_norm) - for start_idx, end_idx in stage_indices: - held_layers.extend(encoder.layers[start_idx:end_idx]) - else: - # current stage is in whisper's decoder - # TODO:(Jianghai) We divide encoder and decoder layers into different parts here, - # the case encoder and decoder put in same stage should be add in the future. - if stage_manager.stage == decoder_starting_stage: - held_layers.append(decoder.embed_tokens) - held_layers.append(decoder.embed_positions) - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(decoder.layer_norm) - for start_idx, end_idx in stage_indices: - held_layers.extend(encoder.layers[start_idx:end_idx]) - else: - layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( - num_encoder_layers, num_decoder_layers, stage_manager.num_stages - ) - start_idx, end_idx = self.get_whisper_stage_index( - layers_per_stage, stage_manager.stage, decoder_starting_stage - ) - - if stage_manager.stage < decoder_starting_stage: - # current stage is in whisper's encoder - if stage_manager.is_first_stage(): - held_layers.append(encoder.embed_positions) - held_layers.append(encoder.conv1) - held_layers.append(encoder.conv2) - if stage_manager.stage == decoder_starting_stage - 1: - held_layers.append(encoder.layer_norm) - held_layers.extend(encoder.layers[start_idx:end_idx]) - else: - # current stage is in whisper's decoder - # TODO:(Jianghai) We divide encoder and decoder layers into different parts here, - # the case encoder and decoder put in same stage should be add in the future. - if stage_manager.stage == decoder_starting_stage: - held_layers.append(decoder.embed_tokens) - held_layers.append(decoder.embed_positions) - if stage_manager.is_last_stage(): - held_layers.append(decoder.layer_norm) - held_layers.extend(decoder.layers[start_idx:end_idx]) - return held_layers - - def set_pipeline_forward(self, model_cls: nn.Module, new_forward: Callable, policy: Dict) -> None: - """If under pipeline parallel setting, replacing the original forward method of huggingface - to customized forward method, and add this changing to policy.""" - if not self.pipeline_stage_manager: - raise ValueError("set_pipeline_forward method can only be called when pipeline parallel is enabled.") - stage_manager = self.pipeline_stage_manager - - if self.model.__class__.__name__ == "WhisperModel": - model = self.model - elif self.model.__class__.__name__ == "WhisperForConditionalGeneration": - model = self.model.model - else: - model = None - - if model: - encoder = self.model.get_encoder() - decoder = self.model.get_decoder() - else: - encoder = self.model.encoder - decoder = None - - num_encoder_layers = len(encoder.layers) - if decoder: - num_decoder_layers = len(decoder.layers) - else: - num_decoder_layers = 0 - - layers_per_stage, decoder_starting_stage = self.distribute_whisper_layers( - num_encoder_layers, num_decoder_layers, stage_manager.num_stages - ) - stage_index = self.get_whisper_stage_index(layers_per_stage, stage_manager.stage, decoder_starting_stage) - - method_replacement = { - "forward": partial( - new_forward, - stage_manager=stage_manager, - stage_index=stage_index, - decoder_starting_stage=decoder_starting_stage, - shard_config=self.shard_config, - ) - } - self.append_or_create_method_replacement(description=method_replacement, policy=policy, target_key=model_cls) - - -# WhisperModel -class WhisperModelPolicy(WhisperPolicy): - def module_policy(self): - from transformers import WhisperModel - - policy = super().module_policy() - - if self.pipeline_stage_manager is not None: - self.set_pipeline_forward( - model_cls=WhisperModel, - new_forward=WhisperPipelineForwards.whisper_model_forward, - policy=policy, - ) - - return policy - - def get_held_layers(self) -> List[nn.Module]: - return super().get_held_layers() - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - "no shared params in whisper model" - return [] - - -# WhisperForConditionalGeneration -class WhisperForConditionalGenerationPolicy(WhisperPolicy): - def module_policy(self): - from transformers import WhisperForConditionalGeneration - - policy = super().module_policy() - policy = self.add_lm_head_policy(policy) - - if self.pipeline_stage_manager is not None: - self.set_pipeline_forward( - model_cls=WhisperForConditionalGeneration, - new_forward=WhisperPipelineForwards.whisper_for_conditional_generation_forward, - policy=policy, - ) - return policy - - def postprocess(self): - return self.model - - def get_held_layers(self) -> List[nn.Module]: - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.proj_out) - else: - if self.pipeline_stage_manager.is_last_stage(): - held_layers.append(self.model.proj_out) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - module = self.model - model = module.model - - if model: - encoder = self.model.get_encoder() - decoder = self.model.get_decoder() - else: - encoder = self.model.encoder - decoder = None - - num_encoder_layers = len(encoder.layers) - if decoder: - num_decoder_layers = len(decoder.layers) - else: - num_decoder_layers = 0 - - stage_manager = self.pipeline_stage_manager - if stage_manager is not None and stage_manager.num_stages > 1: - _, decoder_starting_stage = self.distribute_whisper_layers( - num_encoder_layers, num_decoder_layers, stage_manager.num_stages - ) - shared_params = [] - shared_embedding = {} - if id(module.proj_out) == id(model.decoder.embed_tokens): - shared_embedding[decoder_starting_stage] = model.decoder.embed_tokens - shared_embedding[stage_manager.num_stages - 1] = module.proj_out - if len(shared_embedding) > 0: - shared_params.append(shared_embedding) - return shared_params - return [] - - -# WhisperForAudioClassification -class WhisperForAudioClassificationPolicy(WhisperPolicy): - def module_policy(self): - from transformers import WhisperForAudioClassification - - policy = super().module_policy() - - if self.pipeline_stage_manager is not None: - self.set_pipeline_forward( - model_cls=WhisperForAudioClassification, - new_forward=WhisperPipelineForwards.whisper_for_audio_classification_forward, - policy=policy, - ) - return policy - - def get_held_layers(self) -> List[nn.Module]: - held_layers = super().get_held_layers() - stage_manager = self.pipeline_stage_manager - if stage_manager.is_interleave: - if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or ( - not stage_manager.use_zbv and stage_manager.is_last_stage(ignore_chunk=True) - ): - held_layers.append(self.model.projector) - held_layers.append(self.model.classifier) - else: - if self.pipeline_stage_manager.is_last_stage(): - held_layers.append(self.model.projector) - held_layers.append(self.model.classifier) - return held_layers - - def get_shared_params(self) -> List[Dict[int, Tensor]]: - return [] +reexport_transformers_impl(globals(), __package__, "_whisper") diff --git a/colossalai/shardformer/shard/sharder.py b/colossalai/shardformer/shard/sharder.py index ee2f1f405879..dfa374e65591 100644 --- a/colossalai/shardformer/shard/sharder.py +++ b/colossalai/shardformer/shard/sharder.py @@ -28,7 +28,28 @@ class ModelSharder(object): def __init__(self, model: nn.Module, policy: Policy, shard_config: ShardConfig = None) -> None: self.model = model self.shard_config = shard_config - self.policy = get_autopolicy(self.model) if policy is None else policy + if policy is None: + from .._utils import check_hf5_policy + from ..policies.auto_policy import _POLICY_LIST, _fullname + + location = _POLICY_LIST.get(_fullname(model)) + if location is not None: + check_hf5_policy(location.file_name, shard_config) + self.policy = get_autopolicy(self.model) + else: + from .._utils import check_hf5_policy + from ..policies.auto_policy import _POLICY_LIST + + policy_type = type(policy) + for location in _POLICY_LIST.values(): + if policy_type.__module__ in { + f"colossalai.shardformer.policies.{location.file_name}", + f"colossalai.shardformer.policies._{location.file_name}_tf4", + f"colossalai.shardformer.policies._{location.file_name}_tf5", + }: + check_hf5_policy(location.file_name, shard_config) + break + self.policy = policy def shard(self) -> List[Dict[int, Tensor]]: r""" diff --git a/colossalai/tensor/moe_tensor/api.py b/colossalai/tensor/moe_tensor/api.py index f52802d47384..a565ecf34eb1 100644 --- a/colossalai/tensor/moe_tensor/api.py +++ b/colossalai/tensor/moe_tensor/api.py @@ -32,6 +32,31 @@ def set_moe_tensor_ep_group(tensor: torch.Tensor, ep_group: ProcessGroup) -> Non tensor.__setattr__("ep_group", ep_group) +def set_moe_tensor_ep_fused(tensor: torch.Tensor) -> None: + """ + Mark the given moe tensor as fused experts: all experts are stacked along dim 0 and the + tensor is sharded along dim 0 by its expert parallel group, so every EP rank holds a slice + under the same parameter name. + + Args: + tensor (torch.Tensor): The tensor to be marked. + """ + tensor.__setattr__("ep_fused", True) + + +def is_moe_tensor_ep_fused(tensor: torch.Tensor) -> bool: + """ + Check whether the given tensor holds fused experts sharded along dim 0 by expert parallelism. + + Args: + tensor (torch.Tensor): The tensor to be checked. + + Returns: + bool: Whether the given tensor holds fused experts. + """ + return getattr(tensor, "ep_fused", False) + + def get_moe_info(ep_size: int, dp_size: int, pp_size: int, ep_inside: bool) -> MoeParallelInfo: """ Get moe info for the given tensor. diff --git a/colossalai/tensor/padded_tensor/api.py b/colossalai/tensor/padded_tensor/api.py index 5b66c016b399..b460a4e42229 100644 --- a/colossalai/tensor/padded_tensor/api.py +++ b/colossalai/tensor/padded_tensor/api.py @@ -44,11 +44,12 @@ def _hijack_back_detach_and_clone(ptensor: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: The hijacked tensor. """ - ptensor.detach = ptensor._unpad_detach - ptensor.clone = ptensor._unpad_clone - - delattr(ptensor, "_unpad_detach") - delattr(ptensor, "_unpad_clone") + if hasattr(ptensor, "_unpad_detach"): + ptensor.detach = ptensor._unpad_detach + delattr(ptensor, "_unpad_detach") + if hasattr(ptensor, "_unpad_clone"): + ptensor.clone = ptensor._unpad_clone + delattr(ptensor, "_unpad_clone") return ptensor diff --git a/docs/HF5-COMPATIBILITY.md b/docs/HF5-COMPATIBILITY.md new file mode 100644 index 000000000000..ea986fcc5363 --- /dev/null +++ b/docs/HF5-COMPATIBILITY.md @@ -0,0 +1,64 @@ +# Transformers 5 compatibility and migration boundaries + +Target tested in this checkout: Transformers **5.17.0**, PyTorch **2.11.0**. This table records scopes, not an assertion that the full distributed matrix passes. + +## D4 model-family inventory + +There are **20 registered policy families**. Excluding the eight core families (BERT, Llama, Qwen2, Qwen3, Mixtral, DeepseekV3, ViT, Whisper) leaves **12**, not 14: + +T5, GPT2, GPTJ, OPT, Bloom, SAM, BLIP2, ChatGLM2, Deepseek (legacy/non-V3), Falcon, Mistral, Command/Cohere. + +Do not count model heads, remote-code aliases, or Command/Cohere twice. + +## Support and gate matrix + +| Family | HF5 evidence | Current boundary | +| --- | --- | --- | +| OPT | Removed obsolete attention-class import; tiny eager/SDPA causal-LM forward, gradients and SGD update through actual ShardFormer pass | Native, nonparallel path only; TP/PP/SP, Flash/JIT/fused/fp8/all-optimization are gated pending migration/validation | +| Command/Cohere | StaticCache imports from cache_utils; preserves native HF5 attention instead of unconditional legacy replacement; eager/SDPA causal-LM training checks pass | Same restricted native nonparallel boundary; legacy custom attention/Cache/PP paths are not released | +| BLIP2 | Tiny eager/SDPA multimodal training with valid HF5 image placeholders passes; three obsolete dropout target descriptions removed in HF5 | Native nonparallel path only. Removing targets does **not** establish functional-dropout parallel RNG or TP/Flash/ZBV correctness | +| ChatGLM2 | Remote-code policy module imports, but model code revision/HF5 behavior not validated | HF5 built-in policy gated before lazy import; supply a validated custom policy or use a separately managed legacy environment | +| Deepseek, non-V3 | Legacy remote-code policy, not the native HF5 DeepseekV2 registry | HF5 built-in policy gated; native classes remain unregistered, do not claim V2 support | +| T5/GPT2/GPTJ/Bloom/Falcon/Mistral | Import and tiny descriptor/path audit passes; no full HF5 numerical/distributed acceptance | No new blanket family gate. Existing failures remain visible; old Cache and optional-path contracts require further work | +| SAM | CPU RNG read fixed; TP descriptor next reaches unsupported CPU manual_seed | Remaining platform blocker; not declared an HF5 model incompatibility | +| Eight core families | Four colleague families integrated with HF4/HF5 dispatch; non-EP CPU training checked; HF5 fused expert identity, optimizer ownership/update, optimizer checkpoint and EP gradient regressions executed remotely | The obsolete Mixtral/DeepseekV3 EP identity gate is removed after production four-GPU tests pass. Model-specific restrictions remain; full eight-GPU matrix and physical HF4 regression are not certified | + +The three newly validated families require HF5 **>=5.17.0** in this compatibility branch; earlier HF5 minors are not silently treated as tested. Later versions still need regression testing. HF4 branches remain present, but a real HF4 regression was not executed in this D4 pass. + +### Validated basic configuration + +```python +from colossalai.shardformer import ShardConfig, ShardFormer + +config = ShardConfig(enable_tensor_parallelism=False) +model, shared_params = ShardFormer(config).optimize(model) +``` + +For the new basic scopes, keep optional optimizations and PP/SP disabled. Training probes use FP32, dropout disabled and `use_cache=False`; they do not certify generation/cache, arbitrary dropout, BF16 or long-sequence behavior. BLIP2 HF5 inputs must include the configured image placeholder token repeated for `num_query_tokens`, with valid label masking. + +## Error and legacy guidance + +Unsupported requests report family, installed Transformers version, requested settings, support boundary and this document path. Automatic policy checks run before model mutation; remote-code gates also run before policy import. Explicit registered built-in policies are checked too. Arbitrary custom policies/subclasses remain the caller's responsibility. + +The repository currently declares `transformers==4.51.3`. If an unsupported HF5 path is required, use an **isolated legacy environment**, follow its original dependency declaration and test the desired model/configuration there. This is not a promise that every legacy family passes at that version. Do not downgrade the shared HF5 environment in place. + +## Testing and skip accounting + +- `tests/test_shardformer/test_hf5_compat/test_d4_compat.py`: actual basic training and negative gate tests; these are executed, not skipped. +- Existing distributed entrypoints for OPT/Command/BLIP2 and unvalidated remote-code ChatGLM2/Deepseek are explicitly skipped under HF5 with the same reason as the gate. Lazy test dependency loading prevents unrelated model-zoo imports before these skip decisions. +- Other families are not blanket skipped. Unexpected import/runtime failures remain failures. +- **SKIP is not PASS**. Four/eight-card coverage, full optimized modes, native generation/Cache and the remaining tail numerical matrix remain pending. + +## Integrated loading and pipeline boundaries + +- Offline lazy `from_pretrained` records a checkpoint path. Materialization constructs tensors; `Booster.boost` subsequently loads checkpoint weights and clears the path. These are tested separately, including sharded safetensors with tied parameters. +- Strict loading accepts an omitted tied alias only when the same Parameter has a checkpoint-loaded name; missing unrelated weights still fail. +- Padding alias routing applies to the approved **unsharded HybridParallel load** path. It validates all groups before mutation, verifies zero padding and checks logical values. It does not certify sharded alias routing or optimizer checkpoints. +- Lazy loading currently accepts an explicit `torch.dtype`; HF dtype strings/`auto` inference are not certified and remain explicit errors rather than silent guesses. +- Nested scheduler backward and chunk-local P2P metadata are covered by CPU contract tests. Real GPU P2P, full zero-bubble and dynamic-shape behavior require separate validation. + +### Legacy inference boundary + +HF5 Drafter legacy tuple-cache speculation/resume is not migrated: HF5 removed the legacy Cache conversion APIs. The two legacy-reference tests in `tests/test_infer/test_runtime_compatibility.py` are explicitly skipped under HF5; the nine active PP mask contract cases remain executed. This D1–D5 training migration does not certify legacy inference/cache behavior. + +Evidence: `logs/2026-10-12-d4-audit.md`, `logs/2026-10-12-d4-fix/`, `logs/2026-10-13-integration-review/`, and fresh remote validation in `logs/2026-10-20-integration/FINAL.md`. diff --git a/docs/HF5-POLICY-GATES-zh.md b/docs/HF5-POLICY-GATES-zh.md new file mode 100644 index 000000000000..76965bb77f30 --- /dev/null +++ b/docs/HF5-POLICY-GATES-zh.md @@ -0,0 +1,84 @@ +# HF5 长尾模型 policy 门禁说明（审阅稿） + +环境版本： PyTorch 2.11.0 / Transformers 5.17.0 + +## 1. 名单与计数 + +当前注册 **20 个模型族**。扣除八核心 BERT、Llama、Qwen2、Qwen3、Mixtral、DeepSeek-V3、ViT、Whisper，长尾包含是 **12 个**： + +T5、GPT2、GPTJ、OPT、Bloom、SAM、BLIP2、ChatGLM2、Deepseek（旧版、非 V3）、Falcon、Mistral、Command/Cohere。 + +## 2. 状态矩阵 + +| 模型族 | 当前状态 | 门禁/后续边界 | +|---|---|---| +| OPT | 修复 HF5 共享 attention 类选择；eager/SDPA 非并行 FP32 输出、loss、全参数梯度、SGD 更新通过 | TP/PP/SP、Flash/JIT/fused/fp8 等仍门禁；基础通过不代表生成、Cache、任意 dropout 或 BF16 通过 | +| Command/Cohere | StaticCache 稳定导入；HF5 基础路径保留原生 attention；eager/SDPA 非并行训练对照通过 | 旧自定义 attention/Cache 与优化/并行路径未放行 | +| BLIP2 | 修正三个已消失 dropout 目标；有效 image placeholder 的 eager/SDPA 非并行训练通过 | TP/Flash/ZBV、parallel dropout 数值/RNG 未验收，继续门禁 | +| ChatGLM2 | policy 可导入，remote-code revision/行为未验证 | HF5 内置 policy 在导入前拒绝；需要固定模型代码版本后单独验收 | +| 旧 Deepseek（非 V3） | 旧 remote-code 注册，不等于 HF5 原生 DeepseekV2 | HF5 内置 policy 门禁；原生 V2 未注册 | +| T5/GPT2/GPTJ/Bloom/Falcon/Mistral | 导入及小配置描述/必需目标路径探针通过 | 不 blanket 禁用或 skip；尚未完成 HF5 数值/并行验收，后续失败必须显露 | +| SAM | CPU get_rng_state 已修；随后被 CPU manual_seed 不支持阻断 | CPU 平台阻断，尚不能归因于 HF5 模型不兼容 | + +## 3. 当前门禁逻辑 + +### 3.1 长尾模型 + +- 当前门禁只针对 Transformers 5，依据适配测试结果制定。Transformers 4 保留原有运行路径，兼容性尚未重新测试。 +- ▎ OPT、Command、BLIP2 的基础适配目前只在 Transformers 5.17.0 上验证通过。低于 5.17.0 的 HF5 版本暂不放行；更高版本仍需测试，不能保证兼容。 +- 上述三族有以下请求时拒绝：`enable_tensor_parallelism`、`enable_sequence_parallelism`、`enable_flash_attention`、`enable_jit_fused`、`enable_fused_normalization`、`enable_all_optimization`、`fp8_communication` 或 pipeline stage manager。 + +### 3.2 入口 + +- `colossalai/shardformer/_utils.py`：`hf5_policy_error` / `check_hf5_policy`。 +- `policies/auto_policy.py`：remote-code 在 policy 懒导入前检查。 +- `shard/sharder.py`：自动/显式内置 policy 在模型修改前按配置检查；已覆盖公开入口及 `_tf4`/`_tf5` 内置实现路径。 +- 用户自定义 policy / 外部子类不误拦，正确性由调用者负责；未知模型仍保留原“未实现”错误。 +- `get_autopolicy` 只选择 policy 时没有完整 shard_config，配置限制最终以实际 ShardFormer 接入检查为准。 + +## 4. 基础可用配置与报错指引 + +```python +from colossalai.shardformer import ShardConfig, ShardFormer + +config = ShardConfig(enable_tensor_parallelism=False) +model, shared_params = ShardFormer(config).optimize(model) +``` + +OPT/Command/BLIP2 新基础测试为小配置 FP32、eager/SDPA、dropout 关闭、use_cache=False。BLIP2 需要正确 image_token_id 及 num_query_tokens 个占位 token，并正确 mask label。 + +报错应包含：模型族、实际 HF 版本、被拒设置、限制原因和文档路径。例如： + +```text +ShardFormer opt policy does not support this configuration with transformers 5.17.0: +... requested enable_tensor_parallelism. See docs/HF5-COMPATIBILITY.md. +``` + +旧版指引：仓库当前声明 `transformers==4.51.3`。若需要未支持 HF5 路径，使用**独立旧环境**并针对目标配置验证，不在共享 HF5 环境原地降级。旧声明不是所有模型/配置都通过的保证。 + +## 5. 测试、skip 与证据 + +- D4 正式专项：`tests/test_shardformer/test_hf5_compat/test_d4_compat.py`，**33 passed**，含真实基础训练、配置门禁、import 前拒绝、自定义 policy、未知模型、内置基础 policy 防绕过及核心不误拦。 +- D4 当时 Mac/Linux HF5 CPU 回归各 **173 passed、2 skipped**，不是当前最新全套数字。 +- 最新 V3 本地 CPU 全套 **206 passed、2 skipped**，见 `../logs/2026-10-14-d5-v3/full-cpu.log`；两个 skip 为需要 CUDA 的 head 专项，不等于 D4 长尾 skip。 +- 原分布式入口明确 **9 skipped**：OPT 2、Command 2、BLIP2 1、ChatGLM2 2、旧 Deepseek 2。它们覆盖当前被拒的路径，不记 PASS。 +- skip 标记放在旧包装器外层；测试依赖仅在实际执行时加载，避免不支持用例在 skip 前触发无关 model_zoo 导入。未受限制的其他用例不隐藏依赖/运行失败。 + +证据： + +- [D4 初查](../logs/2026-10-12-d4-audit.md) / [原始 JSON](../logs/2026-10-12-d4-audit.json) +- [D4 修复与验证](../logs/2026-10-12-d4-fix/README.md) +- `../logs/2026-10-12-d4-fix/final-focused-base-policy.log` +- `../logs/2026-10-12-d4-fix/linux-skips.log` +- [最新 D5 状态](../logs/2026-10-14-d5-v3/D5-matrix-status.md) + +这些是仓库内相对文件链接，尚未变成公开在线日志链接。 + +## 6. 解除门禁的标准 + +1. 固定模型族/任务头、HF 版本及 remote-code revision，明确支持配置。 +2. 建立同权重、同输入、同精度对照，调用真实 policy/ShardFormer；检查必要替换目标确实命中。 +3. 通过输出、loss、有限值、全参数梯度/更新、参数身份测试；Cache/生成、checkpoint 另立专项。 +4. 并行必须用相应真实 GPU/拓扑；CPU serial 或描述探针不能替代。 +5. 保留修复前 FAIL 与修复后 PASS，不放宽原阈值制造绿灯；独立算术控制与原生 HF 基准分开报告。 +6. 仅放行已验证配置，并同步修改 gate、非 skip 验收测试、skip reason、文档与日志；不能先删 gate 再把未运行矩阵标绿。 diff --git a/environment-hf5-server.yml b/environment-hf5-server.yml new file mode 100644 index 000000000000..4d1e9404e4d7 --- /dev/null +++ b/environment-hf5-server.yml @@ -0,0 +1,9 @@ +# Conda bootstrap ONLY; follow the pip installation steps in HF5-ENVIRONMENT.md. +# No macOS build strings or local prefix. Linux resolution is not yet verified. +name: colossalai-hf5-gpu +channels: + - conda-forge + - defaults +dependencies: + - python=3.11.16 + - pip diff --git a/requirements-hf5-cu128.txt b/requirements-hf5-cu128.txt new file mode 100644 index 000000000000..cdedc8b60495 --- /dev/null +++ b/requirements-hf5-cu128.txt @@ -0,0 +1,6 @@ +# Install separately BEFORE requirements-hf5.lock.txt on Linux NVIDIA. +# Keep this CUDA index isolated from the general dependency installation. +--index-url https://download.pytorch.org/whl/cu128 +torch==2.11.0 +torchvision==0.26.0 +torchaudio==2.11.0 diff --git a/requirements-hf5.lock.txt b/requirements-hf5.lock.txt new file mode 100644 index 000000000000..1323532c5814 --- /dev/null +++ b/requirements-hf5.lock.txt @@ -0,0 +1,59 @@ +# Captured from colossalai-hf5 on macOS arm64, Python 3.11.16, 2026-09-22. +# Exact installed Python package versions; NOT a complete Linux/CUDA lock. +# On Linux NVIDIA: install requirements-hf5-cu128.txt FIRST, then this file. +# ColossalAI is excluded: install the transferred working tree separately. +# See HF5-ENVIRONMENT.md for scope, known conflicts, and verification. +accelerate==1.15.0 +annotated-doc==0.0.5 +annotated-types==0.8.0 +anyio==4.15.1 +bitsandbytes==0.50.2 +certifi==2026.7.22 +click==8.5.0 +contexttimer==0.3.3 +einops==0.8.2 +filelock==4.0.1 +fsspec==2026.9.0 +galore-torch==1.0 +h11==0.16.0 +hf-xet==1.6.0 +httpcore==1.0.9 +httpx==0.28.1 +huggingface-hub==1.32.0 +idna==3.20 +iniconfig==2.3.0 +jinja2==3.1.6 +markdown-it-py==4.2.0 +markupsafe==3.0.3 +mdurl==0.1.2 +mpmath==1.3.0 +networkx==3.6.1 +ninja==1.13.2 +numpy==2.4.6 +packaging==26.3 +peft==0.19.1 +pillow==12.3.0 +pip==26.2.1 +pluggy==1.6.0 +psutil==7.2.2 +pydantic==2.13.5 +pydantic-core==2.46.5 +pygments==2.21.0 +pytest==8.4.2 +pyyaml==6.0.3 +regex==2026.9.10 +rich==15.0.0 +safetensors==0.8.0 +setuptools==81.0.0 +shellingham==1.5.4 +sympy==1.14.0 +tokenizers==0.23.2 +torch==2.11.0 +torchaudio==2.11.0 +torchvision==0.26.0 +tqdm==4.70.1 +transformers==5.17.0 +typer==0.27.2 +typing-extensions==4.16.0 +typing-inspection==0.4.4 +wheel==0.48.0 diff --git a/tests/kit/model_zoo/transformers/deepseek_v3.py b/tests/kit/model_zoo/transformers/deepseek_v3.py index ba15be3c8400..1efecddb83b8 100644 --- a/tests/kit/model_zoo/transformers/deepseek_v3.py +++ b/tests/kit/model_zoo/transformers/deepseek_v3.py @@ -40,30 +40,50 @@ def data_gen_for_lm(): loss_fn_for_lm = lambda x: x.loss -def init_deepseek(): +# Shared tiny shape for both DeepSeek-v3 implementations below. +TINY_DEEPSEEK_V3 = dict( + hidden_size=128, + intermediate_size=320, + kv_lora_rank=4, + moe_intermediate_size=32, + num_attention_heads=4, + num_experts_per_tok=4, + n_group=4, + num_hidden_layers=3, + num_key_value_heads=4, + first_k_dense_replace=1, + q_lora_rank=8, + torch_dtype="bfloat16", + n_routed_experts=16, + topk_group=2, + v_head_dim=32, + qk_nope_head_dim=32, + qk_rope_head_dim=32, + vocab_size=2048, +) + + +def remote_code_unsupported_reason(): + """Why the Hub remote code (`deepseek-ai/DeepSeek-V3`, `trust_remote_code=True`) cannot run here, or None. + + The repository's modeling_deepseek.py imports `is_torch_fx_available` at module level and uses the + legacy `DynamicCache` API; Transformers v5 removed both (docs/43 DS-10). Detect the APIs, not the version. + """ + from transformers.cache_utils import DynamicCache + from transformers.utils import import_utils + + missing = [name for name in ("is_torch_fx_available",) if not hasattr(import_utils, name)] + missing += [ + f"DynamicCache.{name}" for name in ("from_legacy_cache", "get_usable_length") if not hasattr(DynamicCache, name) + ] + if missing: + return f"DeepSeek-V3 Hub remote code needs APIs removed from this Transformers: {', '.join(missing)}" + return None - config = AutoConfig.from_pretrained( - "deepseek-ai/DeepSeek-V3", - hidden_size=128, - intermediate_size=320, - kv_lora_rank=4, - moe_intermediate_size=32, - num_attention_heads=4, - num_experts_per_tok=4, - n_group=4, - num_hidden_layers=3, - num_key_value_heads=4, - first_k_dense_replace=1, - q_lora_rank=8, - torch_dtype="bfloat16", - n_routed_experts=16, - topk_group=2, - v_head_dim=32, - qk_nope_head_dim=32, - qk_rope_head_dim=32, - trust_remote_code=True, - vocab_size=2048, - ) + +def init_deepseek(): + """DeepSeek-v3 from the Hub remote code (`transformers_modules.modeling_deepseek`).""" + config = AutoConfig.from_pretrained("deepseek-ai/DeepSeek-V3", trust_remote_code=True, **TINY_DEEPSEEK_V3) if hasattr(config, "pad_token_id"): config.pad_token_id = config.eos_token_id @@ -73,3 +93,12 @@ def init_deepseek(): if m.__class__.__name__ == "DeepseekV3MoE": m.moe_infer = MethodType(m.moe_infer.__wrapped__, m) return model + + +def init_deepseek_native(): + """DeepSeek-v3 built into Transformers (`transformers.models.deepseek_v3`), no Hub access needed.""" + from transformers import DeepseekV3Config, DeepseekV3ForCausalLM + + config = DeepseekV3Config(**TINY_DEEPSEEK_V3) + config.pad_token_id = config.eos_token_id + return DeepseekV3ForCausalLM(config) diff --git a/tests/test_checkpoint_io/test_hybrid_parallel_padding_alias.py b/tests/test_checkpoint_io/test_hybrid_parallel_padding_alias.py new file mode 100644 index 000000000000..5b7a60dbaaa1 --- /dev/null +++ b/tests/test_checkpoint_io/test_hybrid_parallel_padding_alias.py @@ -0,0 +1,39 @@ +import torch +from torch import nn + +from colossalai.checkpoint_io.hybrid_parallel_checkpoint_io import _route_padded_parameter_aliases +from colossalai.shardformer.layer.parallel_module import PaddingParallelModule + + +class _PaddedLinear(PaddingParallelModule): + def __init__(self, weight): + super().__init__(new_num_embeddings=4, old_num_embeddings=3, weight=weight) + + @staticmethod + def from_native_module(module, process_group=None): + raise NotImplementedError + + def forward(self, value): + return value @ self.weight.T + + +class _AliasedModel(nn.Module): + def __init__(self): + super().__init__() + weight = nn.Parameter(torch.zeros(3, 2)) + self.weight_alias = weight + self.parallel = _PaddedLinear(weight) + + +def test_load_padded_parameter_through_parallel_alias(): + model = _AliasedModel() + checkpoint = {"weight_alias": torch.arange(6, dtype=torch.float32).reshape(3, 2)} + + _route_padded_parameter_aliases(model, checkpoint) + result = model.load_state_dict(checkpoint, strict=False) + + assert result.missing_keys == ["weight_alias"] + assert not result.unexpected_keys + expected = torch.tensor([[0, 1], [2, 3], [4, 5], [0, 0]], dtype=torch.float32) + assert torch.equal(model.parallel.weight, expected) + assert model.weight_alias is model.parallel.weight diff --git a/tests/test_checkpoint_io/test_moe_checkpoint_ep_gather.py b/tests/test_checkpoint_io/test_moe_checkpoint_ep_gather.py new file mode 100644 index 000000000000..706831e949d0 --- /dev/null +++ b/tests/test_checkpoint_io/test_moe_checkpoint_ep_gather.py @@ -0,0 +1,132 @@ +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +from torch import nn + +from colossalai.checkpoint_io.moe_checkpoint import MoECheckpointIO +from colossalai.interface.model import ModelWrapper +from colossalai.tensor.moe_tensor.api import set_moe_tensor_ep_fused, set_moe_tensor_ep_group +from colossalai.testing.utils import spawn + +NUM_EXPERTS = 4 +EP_SIZE = 2 +NUM_LOCAL_EXPERTS = NUM_EXPERTS // EP_SIZE + + +class _PerExpertMoE(nn.Module): + """Experts kept as separate modules, as in the transformers v4 Mixtral / DeepSeek EP blocks. + + Each EP rank keeps only its own experts; the others are released, so parameter names differ across ranks. + """ + + def __init__(self, ep_rank: int): + super().__init__() + self.experts = nn.ModuleList(nn.Linear(2, 2, bias=False) for _ in range(NUM_EXPERTS)) + held = range(ep_rank * NUM_LOCAL_EXPERTS, (ep_rank + 1) * NUM_LOCAL_EXPERTS) + for idx, expert in enumerate(self.experts): + if idx in held: + expert.weight.data.fill_(float(idx)) + else: + expert.weight = None + + +class _FusedMoE(nn.Module): + """Experts stacked along dim 0 of one parameter, as in the transformers v5 Mixtral / DeepSeek-v3 EP blocks.""" + + def __init__(self, ep_rank: int): + super().__init__() + start = ep_rank * NUM_LOCAL_EXPERTS + self.gate_up_proj = nn.Parameter(_full_fused_weight()[start : start + NUM_LOCAL_EXPERTS].clone()) + + +def _full_fused_weight() -> torch.Tensor: + return torch.arange(NUM_EXPERTS * 2 * 2, dtype=torch.float32).reshape(NUM_EXPERTS, 2, 2) + + +def _collect(model: nn.Module) -> dict: + state = {} + for shard, _ in MoECheckpointIO._model_sharder(model): + state.update(shard) + return state + + +def _check_model_sharder(rank: int, world_size: int, port: int): + dist.init_process_group( + backend="gloo", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + try: + ep_group = dist.new_group(ranks=list(range(world_size))) + + # Per-expert parameters are only tagged with the EP group: each rank keeps its own experts. + per_expert = _PerExpertMoE(rank) + for p in per_expert.parameters(): + set_moe_tensor_ep_group(p, ep_group) + state = _collect(per_expert) + held = range(rank * NUM_LOCAL_EXPERTS, (rank + 1) * NUM_LOCAL_EXPERTS) + assert set(state) == {f"experts.{idx}.weight" for idx in held} + for idx in held: + assert torch.equal(state[f"experts.{idx}.weight"], torch.full((2, 2), float(idx))) + + # Fused parameters are gathered along the expert dimension to the EP root. + fused = _FusedMoE(rank) + set_moe_tensor_ep_group(fused.gate_up_proj, ep_group) + set_moe_tensor_ep_fused(fused.gate_up_proj) + state = _collect(fused) + if rank == 0: + assert torch.equal(state["gate_up_proj"], _full_fused_weight()) + else: + assert torch.equal(state["gate_up_proj"], fused.gate_up_proj.detach()) + finally: + dist.destroy_process_group() + + +def test_model_sharder_gathers_only_fused_experts(): + spawn(_check_model_sharder, nprocs=EP_SIZE) + + +@pytest.mark.parametrize("fused", [False, True]) +def test_non_root_ep_rank_writes_only_per_expert_checkpoints(monkeypatch, tmp_path, fused): + from colossalai.checkpoint_io import moe_checkpoint as moe_checkpoint_module + + module = _FusedMoE(ep_rank=1) if fused else _PerExpertMoE(ep_rank=1) + for p in module.parameters(): + set_moe_tensor_ep_group(p, None) + if fused: + set_moe_tensor_ep_fused(p) + + checkpoint_io = object.__new__(MoECheckpointIO) + checkpoint_io.moe_dp_rank = 0 + checkpoint_io.ep_rank = 1 + checkpoint_io.ep_size = EP_SIZE + checkpoint_io.tp_rank = 0 + checkpoint_io.sp_rank = 0 + checkpoint_io.pp_rank = 0 + checkpoint_io.pp_size = 1 + checkpoint_io.verbose = False + + class _Coordinator: + def is_master(self): + return False + + checkpoint_io.coordinator = _Coordinator() + observed = {} + + def fake_save_state_dict_shards(sharded_state_dict, checkpoint, index_file, base_filename, is_master, **kwargs): + observed["is_master"] = is_master + return 0 + + monkeypatch.setattr(MoECheckpointIO, "_model_sharder", staticmethod(lambda model, **kwargs: iter(()))) + monkeypatch.setattr(moe_checkpoint_module, "save_state_dict_shards", fake_save_state_dict_shards) + monkeypatch.setattr(moe_checkpoint_module.dist, "barrier", lambda: None) + + checkpoint_io.save_sharded_model(ModelWrapper(module), str(tmp_path / "checkpoint")) + + # Fused experts are gathered to the EP root, which writes them; per-expert parameters have + # different names on each EP rank, so every EP rank writes its own experts. + assert observed["is_master"] is (not fused) diff --git a/tests/test_checkpoint_io/test_save_config_file.py b/tests/test_checkpoint_io/test_save_config_file.py new file mode 100644 index 000000000000..5e30b8ac1dbf --- /dev/null +++ b/tests/test_checkpoint_io/test_save_config_file.py @@ -0,0 +1,35 @@ +import json + +import torch +from transformers import LlamaConfig, LlamaForCausalLM + +from colossalai.checkpoint_io.utils import save_config_file + + +def test_save_config_file_writes_model_config(tmp_path): + # Sharded checkpoints are reloaded with `from_pretrained(checkpoint_dir)`; without config.json + # Transformers falls back to the default (full-size) config instead of the saved model's. + config = LlamaConfig( + vocab_size=128, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + ) + model = LlamaForCausalLM(config).to(torch.bfloat16) + + save_config_file(model, str(tmp_path)) + + saved = json.loads((tmp_path / "config.json").read_text()) + assert saved["hidden_size"] == 32 + assert saved["num_hidden_layers"] == 2 + assert saved["architectures"] == ["LlamaForCausalLM"] + assert saved.get("dtype", saved.get("torch_dtype")) == "bfloat16" + assert (tmp_path / "generation_config.json").is_file() + + +def test_save_config_file_skips_non_huggingface_model(tmp_path): + save_config_file(torch.nn.Linear(2, 2), str(tmp_path)) + + assert not (tmp_path / "config.json").exists() diff --git a/tests/test_infer/test_runtime_compatibility.py b/tests/test_infer/test_runtime_compatibility.py index 349199255313..ac07dca16cb1 100644 --- a/tests/test_infer/test_runtime_compatibility.py +++ b/tests/test_infer/test_runtime_compatibility.py @@ -10,8 +10,8 @@ from transformers.models.llama.modeling_llama import LlamaAttention, LlamaRotaryEmbedding from transformers.models.mixtral.modeling_mixtral import MixtralAttention, MixtralConfig, MixtralRotaryEmbedding +from colossalai._compat import is_transformers_v5 from colossalai.booster.plugin import HybridParallelPlugin -from colossalai.inference.spec.drafter import Drafter from colossalai.shardformer.modeling.llama import LlamaPipelineForwards, get_llama_flash_attention_forward from colossalai.shardformer.modeling.mixtral import get_mixtral_flash_attention_forward from colossalai.shardformer.shard import ShardConfig @@ -43,7 +43,13 @@ def test_plugin_cleanup_does_not_hide_rank_error(): mesh.destroy_mesh_process_groups.assert_called_once() +@pytest.mark.skipif( + is_transformers_v5(), + reason="HF5 Drafter legacy tuple-cache inference is unsupported: cache conversion APIs are not migrated", +) def test_drafter_resumes_legacy_cache(): + from colossalai.inference.spec.drafter import Drafter + torch.manual_seed(123) config = LlamaConfig( vocab_size=64, @@ -91,6 +97,10 @@ def test_mixtral_sdpa_attention_layout(kv_heads): torch.testing.assert_close(actual, expected) +@pytest.mark.skipif( + is_transformers_v5(), + reason="HF4 reserved cache-mask column reference is unsupported in HF5 native eager attention", +) def test_llama_attention_slices_reserved_cache_mask_column(): config = LlamaConfig(hidden_size=32, num_attention_heads=4, num_key_value_heads=2) config._attn_implementation = "eager" @@ -110,7 +120,8 @@ def test_llama_attention_slices_reserved_cache_mask_column(): @pytest.mark.parametrize("mode", ["split_gather", "ring", "all_to_all"]) def test_llama_pipeline_mask_uses_full_sequence(mode, monkeypatch): - from colossalai.shardformer.modeling import llama as module + # Patch the implementation: the public module only re-exports version-specific objects. + module = importlib.import_module(LlamaPipelineForwards.__module__) model = LlamaModel( LlamaConfig( @@ -126,16 +137,20 @@ def test_llama_pipeline_mask_uses_full_sequence(mode, monkeypatch): def decoder(hidden, **kwargs): masks.append(kwargs["attention_mask"]) - return (hidden,) + return hidden if is_transformers_v5() else (hidden,) monkeypatch.setattr(model.layers[0], "forward", decoder) monkeypatch.setattr(module, "gather_sp_output", lambda hidden, config: hidden) stage = SimpleNamespace(is_first_stage=lambda: False, is_last_stage=lambda: True) + # This serial mask-contract test does not initialize process groups. A real + # SP ShardConfig would query distributed world sizes during construction. shard = SimpleNamespace( sequence_parallel_size=2, sequence_parallel_process_group=None, sequence_parallelism_mode=mode, + enable_sequence_parallelism=True, enable_flash_attention=False, + gradient_checkpoint_config=None, parallel_output=False, ) LlamaPipelineForwards.llama_model_forward( @@ -155,8 +170,10 @@ def decoder(hidden, **kwargs): @pytest.mark.parametrize("mode", ["split_gather", "ring", "all_to_all"]) def test_qwen_pipeline_mask_uses_full_sequence(family, mode, monkeypatch): transformers_module = importlib.import_module(f"transformers.models.{family}.modeling_{family}") - module = importlib.import_module(f"colossalai.shardformer.modeling.{family}") + public_module = importlib.import_module(f"colossalai.shardformer.modeling.{family}") prefix = "Qwen2" if family == "qwen2" else "Qwen3" + forwards = getattr(public_module, prefix + "PipelineForwards") + module = importlib.import_module(forwards.__module__) config = getattr(transformers_module, prefix + "Config")( vocab_size=64, hidden_size=32, @@ -167,14 +184,24 @@ def test_qwen_pipeline_mask_uses_full_sequence(family, mode, monkeypatch): ) model = getattr(transformers_module, prefix + "Model")(config) shapes = [] + masks = [] def prepare(shape, *args, **kwargs): shapes.append(shape) return {} - monkeypatch.setattr(module.ColoAttention, "prepare_attn_kwargs", prepare) + def decoder(hidden, *args, **kwargs): + if is_transformers_v5(): + masks.append(kwargs["attention_mask"]) + return hidden + return (hidden,) + + # HF4 creates ColoAttention kwargs directly. HF5 creates real additive HF + # masks by layer type; inspect the decoder mask instead of mocking an unused API. + if not is_transformers_v5(): + monkeypatch.setattr(module.ColoAttention, "prepare_attn_kwargs", prepare) monkeypatch.setattr(module, "gather_sp_output", lambda hidden, config: hidden) - monkeypatch.setattr(model.layers[0], "forward", lambda hidden, *args: (hidden,)) + monkeypatch.setattr(model.layers[0], "forward", decoder) stage = SimpleNamespace(is_first_stage=lambda: False, is_last_stage=lambda: True) shard = SimpleNamespace( sequence_parallel_size=2, @@ -182,6 +209,7 @@ def prepare(shape, *args, **kwargs): sequence_parallelism_mode=mode, enable_sequence_parallelism=True, enable_flash_attention=True, + gradient_checkpoint_config=None, parallel_output=False, ) forwards = getattr(module, prefix + "PipelineForwards") @@ -193,4 +221,10 @@ def prepare(shape, *args, **kwargs): stage_index=[0, 1], shard_config=shard, ) - assert shapes == [(2, 1, 8, 8)] + if is_transformers_v5(): + assert len(masks) == 1 + assert masks[0].shape == (2, 1, 8, 8) + assert masks[0][0, 0, 0, 1] < -1e30 + assert masks[0][0, 0, 7, 0] == 0 + else: + assert shapes == [(2, 1, 8, 8)] diff --git a/tests/test_lazy/test_hf5_regressions.py b/tests/test_lazy/test_hf5_regressions.py new file mode 100644 index 000000000000..f8132a74a9dc --- /dev/null +++ b/tests/test_lazy/test_hf5_regressions.py @@ -0,0 +1,97 @@ +"""Lazy construction regressions requiring no GPU or model downloads.""" + +import subprocess +import sys +from typing import Optional, get_args + +import pytest +import torch + +from colossalai.lazy import LazyInitContext, LazyTensor + + +@pytest.mark.parametrize("name,dtype", [("LongTensor", torch.int64), ("FloatTensor", torch.float32)]) +def test_legacy_constructor_annotations(name, dtype): + original = getattr(torch, name) + real_tensor = original([1, 2, 3]) + with LazyInitContext(): + constructor = getattr(torch, name) + assert isinstance(constructor, type) + assert isinstance(real_tensor, constructor) + assert get_args(constructor | None) == (original, type(None)) + assert get_args(None | constructor) == (type(None), original) + assert get_args(Optional[constructor]) + x = constructor([1, 2, 3]) + sized = constructor(3) + empty = constructor() + assert isinstance(x, LazyTensor) + assert getattr(torch, name) is original + torch.testing.assert_close(x.materialize(), torch.tensor([1, 2, 3], dtype=dtype)) + assert sized.materialize().shape == (3,) + assert empty.materialize().shape == (0,) + + +def test_copy_does_not_modify_source_history(): + with LazyInitContext(): + source = torch.arange(4, dtype=torch.float32) + a, b = torch.zeros(4), torch.zeros(4) + a.copy_(source) + b.copy_(source) + # Materializing either destination must not recurse through the source. + torch.testing.assert_close(b.materialize(), torch.arange(4, dtype=torch.float32)) + torch.testing.assert_close(a.materialize(), torch.arange(4, dtype=torch.float32)) + torch.testing.assert_close(source.materialize(), torch.arange(4, dtype=torch.float32)) + + +def test_inplace_repeated_operand_is_recorded_once(): + with LazyInitContext(): + x = torch.ones(4) + x.add_(x) + torch.testing.assert_close(x.materialize(), torch.full((4,), 2.0)) + + +def test_keyword_initializer(): + with LazyInitContext(): + x = torch.empty(4) + torch.nn.init.constant_(tensor=x, val=3) + torch.testing.assert_close(x.materialize(), torch.full((4,), 3.0)) + + +@pytest.mark.parametrize("family", ["qwen2", "qwen3"]) +def test_hf_cold_import_and_materialize(family): + # Each family starts with no HF modeling module in sys.modules. Loading it + # outside the lazy context would hide the annotation regression. + program = f""" +import copy +import torch +import transformers +from colossalai.lazy import LazyInitContext, LazyTensor +family = {family!r} +prefix = "Qwen2" if family == "qwen2" else "Qwen3" +config = getattr(transformers, prefix + "Config")(vocab_size=32, hidden_size=16, + intermediate_size=32, num_hidden_layers=1, num_attention_heads=2, + num_key_value_heads=2, head_dim=8) +with LazyInitContext(): + model = getattr(transformers, prefix + "ForCausalLM")(config) + duplicate = copy.deepcopy(model) +assert isinstance(next(model.parameters()), LazyTensor) +LazyInitContext.materialize(model) +LazyInitContext.materialize(duplicate) +for a, b in zip(model.parameters(), duplicate.parameters()): + torch.testing.assert_close(a, b, rtol=0, atol=0) +ref = getattr(transformers, prefix + "ForCausalLM")(config) +ref.load_state_dict(model.state_dict()) +for name, buffer in model.named_buffers(): + torch.testing.assert_close(buffer, dict(ref.named_buffers())[name], rtol=0, atol=0) +inputs = torch.tensor([[1, 2, 3, 4]]) +actual = model(input_ids=inputs, labels=inputs) +expected = ref(input_ids=inputs, labels=inputs) +torch.testing.assert_close(actual.loss, expected.loss) +actual.loss.backward() +expected.loss.backward() +for a, b in zip(model.parameters(), ref.parameters()): + if a.grad is not None: + torch.testing.assert_close(a.grad, b.grad) +""" + result = subprocess.run([sys.executable, "-c", program], capture_output=True, text=True, timeout=90) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/tests/test_moe/moe_utils.py b/tests/test_moe/moe_utils.py index dbcd28ab5939..be7c72abffda 100644 --- a/tests/test_moe/moe_utils.py +++ b/tests/test_moe/moe_utils.py @@ -8,6 +8,8 @@ import torch.distributed as dist from torch.utils._pytree import tree_map +from colossalai.tensor.moe_tensor.api import get_ep_rank, is_moe_tensor_ep_fused + def assert_loose_close(a, b, dtype: torch.dtype = torch.float32, name=""): assert loose_close(a, b, dtype), f"{name} not close {a.mean()} {b.mean()}" @@ -33,6 +35,20 @@ def loose_close(a, b, dtype: torch.dtype = torch.float32): return torch.allclose(a, b, rtol=rtol, atol=atol) +def ep_local_slice(full: torch.Tensor, param: torch.Tensor) -> torch.Tensor: + """Take the part of a full (unsharded) tensor that ``param`` holds on this rank. + + Transformers v5 fuses the experts of a MoE block into one ``[num_experts, ...]`` parameter, which expert + parallelism shards along dim 0, so a reference model's tensor has to be sliced before comparing. Parameters that + are not EP-fused are returned unchanged. + """ + if not is_moe_tensor_ep_fused(param) or full.shape == param.shape: + return full + num_local = param.shape[0] + start = get_ep_rank(param) * num_local + return full[start : start + num_local] + + def check_model_equal(model1, model2, dtype): assert set(model1.state_dict().keys()) == set(model2.state_dict().keys()) for i, ((name, p1), p2) in enumerate(zip(model1.named_parameters(), model2.parameters())): diff --git a/tests/test_moe/test_moe_ep_sp_grad_sync.py b/tests/test_moe/test_moe_ep_sp_grad_sync.py index d6bc7f6549f4..0f7fd522a519 100644 --- a/tests/test_moe/test_moe_ep_sp_grad_sync.py +++ b/tests/test_moe/test_moe_ep_sp_grad_sync.py @@ -15,7 +15,7 @@ from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.random import seed_all -from tests.test_moe.moe_utils import check_model_equal +from tests.test_moe.moe_utils import check_model_equal, ep_local_slice # Relative L2 error allowed between ZeRO grads (bf16) and an fp32 single-GPU reference. bf16 noise stays below 2e-2 # here, while without the sp reduction the expert grads are off by 0.8-1.0. @@ -54,6 +54,9 @@ def build_model() -> MixtralModel: num_local_experts=4, num_experts_per_tok=2, attn_implementation="sdpa", + # transformers v5 defaults to grouped_mm experts, which need 16-byte aligned strides that these tiny + # sizes do not have; v4 ignores this argument + experts_implementation="eager", ) return MixtralModel(config).to(torch.bfloat16).cuda() @@ -97,7 +100,7 @@ def run_grad_sync(zero_stage: int): grad = zero_optimizer.get_param_grad(param) grad = torch.zeros_like(param, dtype=torch.float) if grad is None else grad.float() ref = single_params[name].grad - ref = torch.zeros_like(grad) if ref is None else ref.float() + ref = torch.zeros_like(grad) if ref is None else ep_local_slice(ref.float(), param) err, ref_norm = (grad - ref).norm().item(), ref.norm().item() if err > GRAD_RTOL * ref_norm + 1e-6: errors.append(f"{name}: |zero - single| / |single| = {err / max(ref_norm, 1e-12):.3f}") diff --git a/tests/test_moe/test_moe_ep_zero_grad_norm.py b/tests/test_moe/test_moe_ep_zero_grad_norm.py index b66378833902..9713725ad0ee 100644 --- a/tests/test_moe/test_moe_ep_zero_grad_norm.py +++ b/tests/test_moe/test_moe_ep_zero_grad_norm.py @@ -20,24 +20,25 @@ NUM_EXPERTS = 8 +@parameterize("tp_size", [1, 2]) @parameterize("ep_size", [2, 4]) @parameterize("moe_only", [True, False]) -def run_grad_norm(ep_size: int, moe_only: bool): +def run_grad_norm(tp_size: int, ep_size: int, moe_only: bool): """The ZeRO grad norm must be the norm of the whole gradient and equal on every rank. Dense and expert grads live in two ZeRO stores, and every EP rank holds different experts, so the norm has to combine both stores and all EP ranks. Training only the MoE blocks makes experts carry most of the norm, as in a real MoE model; training everything mixes dense and expert grads. Routing is random, so every expert is used on - every moe_dp replica and the moe_dp scaling of #6464 cannot interfere. + every moe_dp replica and the moe_dp scaling of #6464 cannot interfere. With tp, the norm must also count the + tp-split params (fused MoE experts in transformers v5 are customized distributed tensors) on every tp rank. """ - if dist.get_world_size() % ep_size != 0: + if dist.get_world_size() % (ep_size * tp_size) != 0: return dtype = torch.bfloat16 - rank = dist.get_rank() plugin = MoeHybridParallelPlugin( pp_size=1, - tp_size=1, + tp_size=tp_size, ep_size=ep_size, zero_stage=1, precision="bf16", @@ -57,31 +58,35 @@ def run_grad_norm(ep_size: int, moe_only: bool): ) torch_model = MixtralModel(config).to(dtype).cuda() for name, param in torch_model.named_parameters(): - param.requires_grad_(not moe_only or "block_sparse_moe" in name) + # transformers v5 renamed `block_sparse_moe` to `mlp` + param.requires_grad_(not moe_only or "block_sparse_moe" in name or ".mlp." in name) - # single-GPU reference: same weights in fp32 on the concatenation of every rank's batch (= global mean loss) + # single-GPU reference: same weights in fp32 on the concatenation of every data replica's batch (= global mean loss) single_model = deepcopy(torch_model).float() zero_model = deepcopy(torch_model) zero_optimizer = torch.optim.SGD(zero_model.parameters(), lr=0.0) zero_model, zero_optimizer, _, _, _ = booster.boost(zero_model, zero_optimizer) - seed_all(1453 + rank) + # ranks of one tp group must see the same batch + seed_all(1453 + dist.get_rank(plugin.dp_group)) input_data = torch.rand(8, 16, HIDDEN_SIZE).cuda().to(dtype) zero_optimizer.backward(zero_model(inputs_embeds=input_data).last_hidden_state.mean()) zero_optimizer.step() # lr = 0: nothing moves, but this computes the norm used for clipping zero_norm = zero_optimizer.get_grad_norm() - all_inputs = [torch.empty_like(input_data) for _ in range(dist.get_world_size())] - dist.all_gather(all_inputs, input_data) + all_inputs = [torch.empty_like(input_data) for _ in range(dist.get_world_size(plugin.dp_group))] + dist.all_gather(all_inputs, input_data, group=plugin.dp_group) single_model(inputs_embeds=torch.cat(all_inputs).float()).last_hidden_state.mean().backward() grads = [p.grad.double() for p in single_model.parameters() if p.grad is not None] ref_norm = sum(g.pow(2).sum() for g in grads).sqrt().item() norms = [None] * dist.get_world_size() dist.all_gather_object(norms, zero_norm) - assert len(set(norms)) == 1, f"ep_size={ep_size} moe_only={moe_only}: grad norm differs across ranks: {norms}" + assert ( + len(set(norms)) == 1 + ), f"tp_size={tp_size} ep_size={ep_size} moe_only={moe_only}: grad norm differs across ranks: {norms}" assert abs(zero_norm - ref_norm) <= NORM_RTOL * ref_norm, ( - f"ep_size={ep_size} moe_only={moe_only}: ZeRO grad norm {zero_norm:.6g}, reference {ref_norm:.6g}, " + f"tp_size={tp_size} ep_size={ep_size} moe_only={moe_only}: ZeRO grad norm {zero_norm:.6g}, reference {ref_norm:.6g}, " f"ratio {zero_norm / ref_norm:.4f}" ) diff --git a/tests/test_moe/test_moe_ep_zero_grad_scale.py b/tests/test_moe/test_moe_ep_zero_grad_scale.py index 4f0317b54a70..99d48c4a1381 100644 --- a/tests/test_moe/test_moe_ep_zero_grad_scale.py +++ b/tests/test_moe/test_moe_ep_zero_grad_scale.py @@ -12,7 +12,7 @@ from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin from colossalai.testing import rerun_if_address_is_in_use, spawn from colossalai.testing.random import seed_all -from tests.test_moe.moe_utils import assert_loose_close +from tests.test_moe.moe_utils import assert_loose_close, ep_local_slice # Relative L2 error allowed against the reference grads. bf16 noise stays below 1e-2 here, while the old moe_dp # scaling was off by 0.33x-3x on the experts. The atol of assert_loose_close (4e-3) exceeds these tiny grads. @@ -60,7 +60,9 @@ def run_forced_routing(): torch_model = MixtralModel(config).to(dtype).cuda() with torch.no_grad(): for layer in torch_model.layers: - gate = layer.block_sparse_moe.gate.weight + # transformers v5 renamed `block_sparse_moe` to `mlp` + moe_block = layer.block_sparse_moe if hasattr(layer, "block_sparse_moe") else layer.mlp + gate = moe_block.gate.weight gate.zero_() gate[0, 0], gate[1, 0] = 10.0, -10.0 @@ -98,6 +100,7 @@ def run_forced_routing(): "single": single_params[n].grad, } grads = {k: torch.zeros_like(p, dtype=torch.float) if g is None else g.float() for k, g in grads.items()} + grads = {k: ep_local_slice(g, p) for k, g in grads.items()} for name, ref_name in (("zero", "ddp"), ("zero", "single"), ("ddp", "single")): grad, ref_grad = grads[name], grads[ref_name] err, ref = (grad - ref_grad).norm().item(), ref_grad.norm().item() diff --git a/tests/test_pipeline/test_schedule/test_integration_backward.py b/tests/test_pipeline/test_schedule/test_integration_backward.py new file mode 100644 index 000000000000..d1a0ba348a79 --- /dev/null +++ b/tests/test_pipeline/test_schedule/test_integration_backward.py @@ -0,0 +1,50 @@ +"""CPU scheduler backward contracts; no P2P transport is exercised here.""" + +import pytest +import torch + +from colossalai.pipeline.schedule.interleaved_pp import InterleavedSchedule +from colossalai.pipeline.schedule.one_f_one_b import OneForwardOneBackwardSchedule + + +class NativeOptimizer: + def __init__(self): + self.calls = 0 + + def backward(self, loss): + self.calls += 1 + loss.backward() + + def backward_by_grad(self, tensor, grad): + self.calls += 1 + torch.autograd.backward(tensor, grad) + + +@pytest.mark.parametrize("schedule", [OneForwardOneBackwardSchedule, InterleavedSchedule]) +def test_scheduler_nested_backward(schedule): + inputs = {"hidden": torch.ones(2, requires_grad=True), "router": (torch.ones(2, requires_grad=True),)} + outputs = {"hidden": inputs["hidden"] * 2, "router": (inputs["router"][0] * 3,)} + upstream = {"hidden": torch.ones(2), "router": (torch.ones(2),)} + optimizer = NativeOptimizer() + gradients = schedule.backward_step(None, optimizer, inputs, outputs, upstream) + assert optimizer.calls == 1 + torch.testing.assert_close(gradients["hidden"], torch.full((2,), 2.0)) + torch.testing.assert_close(gradients["router"][0], torch.full((2,), 3.0)) + + +@pytest.mark.parametrize("schedule", [OneForwardOneBackwardSchedule, InterleavedSchedule]) +def test_scheduler_no_backward_tensor(schedule): + optimizer = NativeOptimizer() + result = schedule.backward_step(None, optimizer, {"metadata": "x"}, {"metadata": "y"}, {}) + assert result is None + assert optimizer.calls == 0 + + +@pytest.mark.parametrize("schedule", [OneForwardOneBackwardSchedule, InterleavedSchedule]) +def test_scheduler_structure_error_before_optimizer(schedule): + optimizer = NativeOptimizer() + inputs = {"hidden": torch.ones(2, requires_grad=True)} + outputs = {"hidden": (inputs["hidden"] * 2,)} + with pytest.raises(ValueError, match="structures must match"): + schedule.backward_step(None, optimizer, inputs, outputs, {"hidden": torch.ones(2)}) + assert optimizer.calls == 0 diff --git a/tests/test_pipeline/test_schedule/test_interleaved_metadata_cache.py b/tests/test_pipeline/test_schedule/test_interleaved_metadata_cache.py new file mode 100644 index 000000000000..ca3e1aa4a758 --- /dev/null +++ b/tests/test_pipeline/test_schedule/test_interleaved_metadata_cache.py @@ -0,0 +1,186 @@ +from contextlib import contextmanager + +import torch + +from colossalai.pipeline.schedule.interleaved_pp import InterleavedSchedule + + +class _StageManager: + def __init__(self): + self.model_chunk_id = None + self.first_by_chunk = {} + self.last_by_chunk = {} + + @contextmanager + def switch_model_chunk_id(self, model_chunk_id): + self.model_chunk_id = model_chunk_id + yield + self.model_chunk_id = None + + def is_first_stage(self): + return self.first_by_chunk.get(self.model_chunk_id, False) + + def is_last_stage(self): + return self.last_by_chunk.get(self.model_chunk_id, False) + + +class _Communication: + def __init__(self): + self.forward_recv_metadata = [] + self.backward_recv_metadata = [] + self.forward_send_metadata = [] + self.backward_send_metadata = [] + self.forward_send_recv_metadata = [] + self.backward_send_recv_metadata = [] + + def recv_forward(self, _prev_rank, metadata_recv): + self.forward_recv_metadata.append(metadata_recv) + return torch.ones(1), [] + + def recv_backward(self, _next_rank, metadata_recv): + self.backward_recv_metadata.append(metadata_recv) + return torch.ones(1), [] + + def send_forward(self, _output_tensor, _next_rank, send_metadata): + self.forward_send_metadata.append(send_metadata) + return [] + + def send_backward(self, _input_tensor_grad, _prev_rank, send_metadata): + self.backward_send_metadata.append(send_metadata) + return [] + + def send_forward_recv_forward(self, _output_tensor, _is_send, _is_recv, send_metadata, metadata_recv, send_first): + del send_first + self.forward_send_recv_metadata.append((send_metadata, metadata_recv)) + return torch.ones(1), [] + + def send_backward_recv_backward( + self, _input_tensor_grad, _is_send, _is_recv, send_metadata, metadata_recv, send_first + ): + del send_first + self.backward_send_recv_metadata.append((send_metadata, metadata_recv)) + return torch.ones(1), [] + + +def test_interleaved_metadata_cache_is_chunk_local(): + schedule = InterleavedSchedule(_StageManager(), num_model_chunks=2, num_microbatch=2) + communication = _Communication() + schedule.comm = communication + + schedule.recv_forward(0) + schedule.recv_forward(1) + assert communication.forward_recv_metadata == [None, None] + assert schedule.tensor_metadata_recv[0] is not None + assert schedule.tensor_metadata_recv[1] is not None + schedule.recv_forward(0) + assert communication.forward_recv_metadata[2] is schedule.tensor_metadata_recv[0] + + schedule.recv_backward(0) + schedule.recv_backward(1) + assert communication.backward_recv_metadata == [None, None] + assert schedule.grad_metadata_recv[0] is not None + assert schedule.grad_metadata_recv[1] is not None + schedule.recv_backward(0) + assert communication.backward_recv_metadata[2] is schedule.grad_metadata_recv[0] + + tensor = torch.ones(1) + schedule.send_forward(0, tensor) + schedule.send_forward(1, tensor) + schedule.send_forward(0, tensor) + assert communication.forward_send_metadata == [True, True, False] + + schedule.send_backward(0, tensor) + schedule.send_backward(1, tensor) + schedule.send_backward(0, tensor) + assert communication.backward_send_metadata == [True, True, False] + + schedule = InterleavedSchedule(_StageManager(), num_model_chunks=2, num_microbatch=2) + communication = _Communication() + schedule.comm = communication + schedule.send_forward_recv_forward(0, 1, tensor) + schedule.send_forward_recv_forward(1, 0, tensor) + assert communication.forward_send_recv_metadata[0] == (True, None) + assert communication.forward_send_recv_metadata[1] == (True, None) + assert schedule.send_tensor_metadata == [False, False] + assert all(metadata is not None for metadata in schedule.tensor_metadata_recv) + + schedule.send_backward_recv_backward(0, 1, tensor) + schedule.send_backward_recv_backward(1, 0, tensor) + assert communication.backward_send_recv_metadata[0] == (True, None) + assert communication.backward_send_recv_metadata[1] == (True, None) + assert schedule.send_grad_metadata == [False, False] + assert all(metadata is not None for metadata in schedule.grad_metadata_recv) + + +def test_chunk_metadata_structure_and_no_send_state(): + stage_manager = _StageManager() + schedule = InterleavedSchedule(stage_manager, num_model_chunks=2, num_microbatch=2) + communication = _Communication() + communication.recv_forward = lambda rank, metadata_recv: ( + ( + {"hidden": torch.ones(2)} + if len(communication.forward_recv_metadata) == 1 + else {"hidden": torch.ones(3), "router": (torch.ones(1),)} + ), + [], + ) + original_recv = communication.recv_forward + + def record_recv(rank, metadata_recv): + communication.forward_recv_metadata.append(metadata_recv) + return original_recv(rank, metadata_recv) + + communication.recv_forward = record_recv + schedule.comm = communication + schedule.recv_forward(0) + schedule.recv_forward(1) + assert schedule.tensor_metadata_recv[0] != schedule.tensor_metadata_recv[1] + + # A combined op that does not send must not consume this chunk's first-send metadata state. + stage_manager.last_by_chunk[0] = True + schedule.send_forward_recv_forward(0, 1, torch.ones(1)) + assert schedule.send_tensor_metadata[0] is True + stage_manager.last_by_chunk[0] = False + schedule.send_forward_recv_forward(0, 1, torch.ones(1)) + assert schedule.send_tensor_metadata[0] is False + + +def test_metadata_cache_disabled_keeps_sending(): + schedule = InterleavedSchedule(_StageManager(), num_model_chunks=2, num_microbatch=2, enable_metadata_cache=False) + communication = _Communication() + schedule.comm = communication + tensor = torch.ones(1) + schedule.send_forward(0, tensor) + schedule.send_forward(0, tensor) + schedule.send_backward(1, tensor) + schedule.send_backward(1, tensor) + assert communication.forward_send_metadata == [True, True] + assert communication.backward_send_metadata == [True, True] + + +def test_batch_size_change_invalidates_all_chunks(): + schedule = InterleavedSchedule(_StageManager(), num_model_chunks=2, microbatch_size=2) + schedule.forward_only = True + schedule.last_batch_size = 4 + schedule.send_tensor_metadata = [False, False] + schedule.send_grad_metadata = [False, False] + schedule.tensor_metadata_recv = [object(), object()] + schedule.grad_metadata_recv = [object(), object()] + schedule.load_batch(iter([{"input": torch.ones(6, 3)}])) + assert schedule.enable_metadata_cache is False + assert schedule.send_tensor_metadata == [True, True] + assert schedule.send_grad_metadata == [True, True] + assert schedule.tensor_metadata_recv == [None, None] + assert schedule.grad_metadata_recv == [None, None] + + +def test_backward_no_send_keeps_first_metadata(): + manager = _StageManager() + manager.first_by_chunk[0] = True + schedule = InterleavedSchedule(manager, num_model_chunks=2, num_microbatch=2) + schedule.comm = _Communication() + schedule.send_backward_recv_backward(0, 1, torch.ones(1)) + assert schedule.send_grad_metadata[0] is True + manager.first_by_chunk[0] = False + schedule.send_backward_recv_backward(0, 1, torch.ones(1)) + assert schedule.send_grad_metadata[0] is False diff --git a/tests/test_pipeline/test_schedule/test_pipeline_schedule_utils.py b/tests/test_pipeline/test_schedule/test_pipeline_schedule_utils.py index 462355ee470b..b57497a370c8 100644 --- a/tests/test_pipeline/test_schedule/test_pipeline_schedule_utils.py +++ b/tests/test_pipeline/test_schedule/test_pipeline_schedule_utils.py @@ -1,6 +1,12 @@ import torch -from colossalai.pipeline.schedule._utils import get_batch_size, get_micro_batch, merge_batch +from colossalai.pipeline.schedule._utils import ( + collect_gradients, + flatten_tensors_and_grads, + get_batch_size, + get_micro_batch, + merge_batch, +) def test_get_batch_size(): @@ -45,3 +51,40 @@ def test_merge_batch(): merged = merge_batch([{"x": x[0:1], "y": y[0:1]}, {"x": x[1:2], "y": y[1:2]}]) assert torch.equal(merged["x"], x) assert torch.equal(merged["y"], y) + + +def test_collect_nested_gradients(): + hidden_states = torch.rand(2, requires_grad=True) + router_logits = (torch.rand(2, requires_grad=True), torch.rand(2, requires_grad=True)) + outputs = {"hidden_states": hidden_states, "past_router_logits": router_logits} + + (hidden_states.sum() + sum(logits.sum() for logits in router_logits)).backward() + + gradients = collect_gradients(outputs) + assert gradients is not None + pairs = flatten_tensors_and_grads(outputs, gradients) + assert len(pairs) == 3 + for tensor, grad in pairs: + assert grad is tensor.grad + + assert collect_gradients({"metadata": ("not a tensor",)}) is None + + +def test_nested_gradient_structure_mismatch(): + output = {"hidden": torch.ones(2), "router": (torch.ones(1),)} + gradients = {"hidden": torch.ones(2), "router": torch.ones(1)} + try: + flatten_tensors_and_grads(output, gradients) + except ValueError as error: + assert "structures must match" in str(error) + else: + raise AssertionError("Expected mismatched structures to fail") + + +def test_nested_gradient_filters_none_and_metadata(): + tensor = torch.ones(2, requires_grad=True) + (tensor * 2).sum().backward() + output = {"tensor": tensor, "metadata": "x", "unused": torch.ones(1)} + gradients = collect_gradients(output) + assert gradients == {"tensor": tensor.grad, "metadata": None, "unused": None} + assert flatten_tensors_and_grads(output, gradients) == [(tensor, tensor.grad)] diff --git a/tests/test_shardformer/test_hf5_compat/BF16-ACCEPTANCE.md b/tests/test_shardformer/test_hf5_compat/BF16-ACCEPTANCE.md new file mode 100644 index 000000000000..03e5ae900065 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/BF16-ACCEPTANCE.md @@ -0,0 +1,37 @@ +# HF5 BF16 验收约定 + +目标：PyTorch 2.11.0+cu128、Transformers 5.17.0。以下约定不代表全部组合已经通过。 + +## 两种结果分开报告 + +1. **原生 HF 基准**：同权重、同 BF16、同任务和优化器初始化；保留原断言结果，失败不删除、不用 skip/放宽阈值改写为通过。 +2. **独立计算路径控制**：单设备原生 PyTorch/HF 分块线性、分组 backward、明确 attention 后端。用于解释 TP/内核/微批舍入差异，不使用生产适配 forward 作为唯一参考，不自动替代原生基准验收。 + +## 每次执行必须记录 + +- Python/Torch/HF/FlashAttention 版本、源码哈希、GPU 编号、命令、退出码和逐 rank PASS/FAIL。 +- seed、初始化权重、精度、输入/labels/padding、cache、dropout、attention 实际 backend（不能仅看 config 名）、mask、scale/stride、微批、TP/PP/SP/ZeRO 拓扑。 +- 优化器超参数及 FP32 master 初值：BF16-rounded 初始权重的 FP32 副本，不混用未舍入的 reference master。 + +## 必检内容 + +- 输出形状、有限值、任务 logits/reconstruction/hidden state、loss dtype 和数值。 +- 全部可训练逻辑参数梯度及更新权重。缺失参数、requires_grad 不一致、单侧 grad=None、形状错误均失败。 +- 只有存在明确 partition inventory 才 gather 分片；词表 padding/tied aliases 单独记录与验证，不能靠自动截断掩盖形状错误。ZeRO 检查实际 working gradient buffer，不把清空 .grad 当成无需比较。 +- 同组复制参数（包括 Qwen3 q/k norm）梯度和更新一致。q/k norm 使用 SUM，不用平均；TP+SP 不重复归约。 +- `bf16_checks.py` 是非 ZeRO 全参数比较的公共小工具；当前不自动支持 padding/PP/ZeRO，调用方必须给出对应映射或明确标注未覆盖。 + +## 数值标准 + +- 默认 BF16 保持 `atol=rtol=0.005`；FP32 保持原测试阈值。 +- 同计算路径输出可额外记录逐位一致，但不要求所有不同内核/累积路径逐位一致。 +- 任何新增全参数检查超阈值，专项结论必须失败；原采样断言 PASS 另列，不能混成全参数通过。 +- 不仅报告最大误差，也报告超差元素数、均值/relative L2、参考幅度以及首次分歧层。更改阈值必须有独立证据和明确评审，不以测试变绿为理由。 + +## 最小回归与当前边界 + +- 固定小模型，至少两个 seed、两个更新步的关键专项；更广拓扑允许单步定向回归。 +- CPU/Gloo q/k norm 专项用于同步/累积机制；真实两卡 CUDA/NCCL 验证 TP、split/ring 和可执行 PP/all_to_all 路径。 +- TP×PP 或 TP×独立 SP 的四 rank 组合、八卡矩阵仍待最后验收，不能由 CPU/Gloo 模拟替代。 +- Whisper ConditionalGeneration 原独立控制按未补齐词表均分，生产 head 则按 51968 个物理行分片（逻辑词表 51866，TP2 每片25984）。固定真实 padding/分片边界及 grouped backward 后，seed42/123 的 logits/loss、全部 decoder 梯度边界精确一致，88 个参数梯度均通过；证据 `logs/2026-10-09-whisper-causal-193000/`。 +- `bf16_reference.grouped_column` 提供独立 zero-padding/F.linear/分组反向控制；必须显式传入真实 padding_multiple，不自动从候选层复制实现。完整 native-HF 默认矩阵仍不能据此改写为全部通过；Whisper 多步收敛/四八卡及更广输入范围尚未验收。 diff --git a/tests/test_shardformer/test_hf5_compat/bf16_checks.py b/tests/test_shardformer/test_hf5_compat/bf16_checks.py new file mode 100644 index 000000000000..a81b42ec8a10 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/bf16_checks.py @@ -0,0 +1,46 @@ +"""Strict logical-parameter comparisons for small HF5 acceptance tests.""" + +import torch +import torch.distributed as dist + + +def compare_tensor(reference, actual, *, atol=0.005, rtol=0.005): + assert reference.shape == actual.shape + assert torch.isfinite(reference).all() + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual.float(), reference.float(), atol=atol, rtol=rtol) + return (actual.float() - reference.float()).abs().max().item() + + +def compare_parameters(reference, actual, *, group=None, shard_dims=None, gradients=True): + """Compare every named parameter using an explicit partition inventory. + + ``shard_dims`` maps parameter names to their actual partition dimension. + A shape mismatch without this inventory is an error, not an inferred shard. + ZeRO-owned gradients require a separate working-buffer inventory; this + helper must not silently accept cleared .grad as full-gradient coverage. + """ + expected = dict(reference.named_parameters()) + candidate = dict(actual.named_parameters()) + assert expected.keys() == candidate.keys(), (expected.keys() - candidate.keys(), candidate.keys() - expected.keys()) + shard_dims = shard_dims or {} + assert shard_dims.keys() <= expected.keys() + errors = {} + for name, parameter in expected.items(): + other = candidate[name] + assert parameter.requires_grad == other.requires_grad, name + r = parameter.grad if gradients else parameter.detach() + a = other.grad if gradients else other.detach() + assert (r is None) == (a is None), name + if r is None: + continue + if name in shard_dims: + assert group is not None and dist.is_initialized(), name + parts = [torch.empty_like(a) for _ in range(dist.get_world_size(group))] + dist.all_gather(parts, a, group=group) + a = torch.cat(parts, dim=shard_dims[name]) + # Vocabulary padding must be accounted for explicitly by the caller. + assert r.shape == a.shape, name + errors[name] = compare_tensor(r, a) + assert errors, "No parameters compared" + return errors diff --git a/tests/test_shardformer/test_hf5_compat/bf16_reference.py b/tests/test_shardformer/test_hf5_compat/bf16_reference.py new file mode 100644 index 000000000000..8e16faa61ef3 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/bf16_reference.py @@ -0,0 +1,44 @@ +"""Independent PyTorch controls, not replacements for native-HF baselines.""" + +import torch +import torch.nn.functional as F + + +class GroupedColumn(torch.autograd.Function): + @staticmethod + def forward(ctx, inputs, weight, bias, parts, padding_multiple): + rows = weight.shape[0] + multiple = parts * padding_multiple + padded_rows = ((rows + multiple - 1) // multiple) * multiple + padded_weight = F.pad(weight, (0, 0, 0, padded_rows - rows)) + padded_bias = None if bias is None else F.pad(bias, (0, padded_rows - rows)) + ctx.save_for_backward(inputs, padded_weight) + ctx.has_bias = bias is not None + ctx.rows = rows + ctx.parts = parts + biases = padded_bias.chunk(parts) if padded_bias is not None else [None] * parts + output = torch.cat([F.linear(inputs, w, b) for w, b in zip(padded_weight.chunk(parts, 0), biases)], dim=-1) + return output[..., :rows] + + @staticmethod + def backward(ctx, grad_output): + inputs, weight = ctx.saved_tensors + grad_output = F.pad(grad_output, (0, weight.shape[0] - ctx.rows)) + slices = grad_output.chunk(ctx.parts, -1) + input_gradients = [g.matmul(w) for g, w in zip(slices, weight.chunk(ctx.parts, 0))] + grad_input = input_gradients[0] + for partial in input_gradients[1:]: + grad_input = grad_input + partial + flat_input = inputs.reshape(-1, inputs.shape[-1]) + grad_weight = torch.cat([g.reshape(-1, g.shape[-1]).t().matmul(flat_input) for g in slices], dim=0)[: ctx.rows] + grad_bias = grad_output.reshape(-1, grad_output.shape[-1]).sum(0)[: ctx.rows] if ctx.has_bias else None + return grad_input, grad_weight, grad_bias, None, None + + +def grouped_column(inputs, weight, bias=None, *, parts=2, padding_multiple=1): + """Match declared physical row partitions, including zero vocabulary padding. + + This is a numerical diagnostic: it performs no distributed communication, + uses no ColossalAI layer, and must be reported separately from native HF. + """ + return GroupedColumn.apply(inputs, weight, bias, parts, padding_multiple) diff --git a/tests/test_shardformer/test_hf5_compat/conftest.py b/tests/test_shardformer/test_hf5_compat/conftest.py new file mode 100644 index 000000000000..493156791422 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/conftest.py @@ -0,0 +1,16 @@ +import pytest +import torch + + +@pytest.fixture(autouse=True) +def deterministic_cpu(): + previous = torch.get_num_threads() + torch.set_num_threads(2) + torch.manual_seed(42) + yield + torch.set_num_threads(previous) + + +@pytest.fixture(params=["qwen2", "qwen3"]) +def family(request): + return request.param diff --git a/tests/test_shardformer/test_hf5_compat/test_bf16_checks.py b/tests/test_shardformer/test_hf5_compat/test_bf16_checks.py new file mode 100644 index 000000000000..ecfa0c877089 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_bf16_checks.py @@ -0,0 +1,31 @@ +import pytest +import torch + +from tests.test_shardformer.test_hf5_compat.bf16_checks import compare_parameters, compare_tensor + + +def test_finite_and_threshold(): + with pytest.raises(AssertionError): + compare_tensor(torch.ones(2), torch.tensor([1.0, float("nan")])) + with pytest.raises(AssertionError): + compare_tensor(torch.zeros(2), torch.tensor([0.0, 0.02])) + + +def test_missing_gradients_and_parameters(): + reference = torch.nn.Linear(2, 2) + actual = torch.nn.Linear(2, 2) + actual.load_state_dict(reference.state_dict()) + reference(torch.ones(1, 2)).sum().backward() + with pytest.raises(AssertionError): + compare_parameters(reference, actual) + actual(torch.ones(1, 2)).sum().backward() + assert len(compare_parameters(reference, actual)) == 2 + with pytest.raises(AssertionError): + compare_parameters(reference, torch.nn.Linear(2, 2, bias=False), gradients=False) + + +def test_unexplained_shape_mismatch(): + reference = torch.nn.Linear(2, 4, bias=False) + actual = torch.nn.Linear(2, 2, bias=False) + with pytest.raises(AssertionError): + compare_parameters(reference, actual, gradients=False) diff --git a/tests/test_shardformer/test_hf5_compat/test_bf16_padded_head_cuda.py b/tests/test_shardformer/test_hf5_compat/test_bf16_padded_head_cuda.py new file mode 100644 index 000000000000..9035f9c4bd70 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_bf16_padded_head_cuda.py @@ -0,0 +1,66 @@ +"""Two-rank physical vocabulary-head acceptance; not full-model coverage.""" + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from colossalai.shardformer.layer import VocabParallelLMHead1D +from tests.test_shardformer.test_hf5_compat.bf16_reference import grouped_column + + +def _worker(rank, rendezvous, seed): + torch.cuda.set_device(rank) + dist.init_process_group("nccl", init_method=rendezvous, rank=rank, world_size=2) + try: + torch.manual_seed(seed) + reference = torch.nn.Linear(256, 51866, bias=False, device="cuda", dtype=torch.bfloat16) + candidate = VocabParallelLMHead1D.from_native_module( + torch.nn.Linear(256, 51866, bias=False, device="cuda", dtype=torch.bfloat16), + process_group=dist.group.WORLD, + gather_output=True, + make_vocab_size_divisible_by=64, + ) + dist.broadcast(reference.weight.data, src=0) + padded = torch.nn.functional.pad(reference.weight.detach(), (0, 0, 0, 102)) + with torch.no_grad(): + candidate.weight.copy_(padded.chunk(2, 0)[rank]) + assert candidate.weight.shape == (25984, 256) + ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01) + optimizer = torch.optim.SGD(candidate.parameters(), lr=0.01) + for _ in range(2): + ref_optimizer.zero_grad() + optimizer.zero_grad() + x = torch.randn(1, 2, 256, device="cuda", dtype=torch.bfloat16, requires_grad=True) + dist.broadcast(x.data, src=0) + actual_x = x.detach().clone().requires_grad_(True) + expected = grouped_column(x, reference.weight, padding_multiple=64) + actual = candidate(actual_x) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + labels = torch.tensor([0, 1], device="cuda") + ref_loss = torch.nn.functional.cross_entropy(expected.reshape(-1, 51866).float(), labels) + loss = torch.nn.functional.cross_entropy(actual.reshape(-1, 51866).float(), labels) + torch.testing.assert_close(loss, ref_loss, atol=0, rtol=0) + ref_loss.backward() + loss.backward() + torch.testing.assert_close(actual_x.grad, x.grad, atol=0.005, rtol=0.005) + parts = [torch.empty_like(candidate.weight.grad) for _ in range(2)] + dist.all_gather(parts, candidate.weight.grad) + full_grad = torch.cat(parts, 0) + torch.testing.assert_close(full_grad[:51866], reference.weight.grad, atol=0.005, rtol=0.005) + assert torch.count_nonzero(full_grad[51866:]) == 0 + ref_optimizer.step() + optimizer.step() + parts = [torch.empty_like(candidate.weight) for _ in range(2)] + dist.all_gather(parts, candidate.weight.detach()) + full_weight = torch.cat(parts, 0) + torch.testing.assert_close(full_weight[:51866], reference.weight, atol=0.005, rtol=0.005) + assert torch.count_nonzero(full_weight[51866:]) == 0 + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires two real CUDA GPUs and NCCL") +@pytest.mark.parametrize("seed", [42, 123]) +def test_padded_vocabulary_head_cuda(tmp_path, seed): + mp.spawn(_worker, args=((tmp_path / "nccl-init").as_uri(), seed), nprocs=2, join=True) diff --git a/tests/test_shardformer/test_hf5_compat/test_bf16_reference.py b/tests/test_shardformer/test_hf5_compat/test_bf16_reference.py new file mode 100644 index 000000000000..b9b64b30314f --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_bf16_reference.py @@ -0,0 +1,28 @@ +import pytest +import torch +import torch.nn.functional as F + +from tests.test_shardformer.test_hf5_compat.bf16_reference import grouped_column + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("bias", [False, True]) +def test_physical_padded_column(dtype, bias): + inputs = torch.randn(2, 3, 8, dtype=dtype, requires_grad=True) + weight = torch.randn(13, 8, dtype=dtype, requires_grad=True) + offset = torch.randn(13, dtype=dtype, requires_grad=True) if bias else None + output = grouped_column(inputs, weight, offset, padding_multiple=4) + padded_weight = F.pad(weight, (0, 0, 0, 3)) + padded_bias = F.pad(offset, (0, 3)) if offset is not None else None + biases = padded_bias.chunk(2) if padded_bias is not None else (None, None) + expected = torch.cat([F.linear(inputs, w, b) for w, b in zip(padded_weight.chunk(2), biases)], -1)[..., :13] + torch.testing.assert_close(output, expected, atol=0, rtol=0) + upstream = torch.randn_like(output) + output.backward(upstream) + padded_grad = F.pad(upstream, (0, 3)) + gradients = [g.matmul(w) for g, w in zip(padded_grad.chunk(2, -1), padded_weight.chunk(2))] + torch.testing.assert_close(inputs.grad, gradients[0] + gradients[1], atol=0, rtol=0) + expected_weight = upstream.reshape(-1, 13).t().matmul(inputs.detach().reshape(-1, 8)) + torch.testing.assert_close(weight.grad, expected_weight, atol=0.005, rtol=0.005) + if offset is not None: + torch.testing.assert_close(offset.grad, padded_grad.reshape(-1, 16).sum(0)[:13], atol=0, rtol=0) diff --git a/tests/test_shardformer/test_hf5_compat/test_cpu_rng.py b/tests/test_shardformer/test_hf5_compat/test_cpu_rng.py new file mode 100644 index 000000000000..2168a6ae0b78 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_cpu_rng.py @@ -0,0 +1,16 @@ +import torch + +from colossalai.accelerator.cpu_accelerator import CpuAccelerator + + +def test_cpu_rng_device_and_restore(): + accelerator = CpuAccelerator() + initial = torch.get_rng_state() + try: + state = accelerator.get_rng_state() + torch.testing.assert_close(state, accelerator.get_rng_state(torch.device("cpu")), atol=0, rtol=0) + expected = torch.rand(8) + accelerator.set_rng_state(state) + torch.testing.assert_close(expected, torch.rand(8), atol=0, rtol=0) + finally: + torch.set_rng_state(initial) diff --git a/tests/test_shardformer/test_hf5_compat/test_d4_compat.py b/tests/test_shardformer/test_hf5_compat/test_d4_compat.py new file mode 100644 index 000000000000..10b305c38f70 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_d4_compat.py @@ -0,0 +1,207 @@ +"""D4 tiny CPU integration and early support-boundary errors.""" + +import copy +from types import SimpleNamespace + +import pytest + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("HF5 long-tail policy contracts", allow_module_level=True) + +import torch +import transformers as hf + +from colossalai.shardformer import ShardConfig, ShardFormer +from colossalai.shardformer._utils import check_hf5_policy, hf5_policy_error +from colossalai.shardformer.policies.auto_policy import PolicyLocation, import_policy +from tests.test_shardformer.test_hf5_compat.bf16_checks import compare_parameters + + +def model_and_data(family, backend): + if family == "opt": + config = hf.OPTConfig( + vocab_size=64, + hidden_size=32, + word_embed_proj_dim=32, + ffn_dim=48, + num_hidden_layers=1, + num_attention_heads=4, + max_position_embeddings=16, + dropout=0, + attention_dropout=0, + use_cache=False, + attn_implementation=backend, + ) + model = hf.OPTForCausalLM(config) + elif family == "blip2": + config = hf.Blip2Config( + vision_config=dict( + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + image_size=16, + patch_size=8, + attention_dropout=0, + ), + qformer_config=dict( + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + encoder_hidden_size=32, + hidden_dropout_prob=0, + attention_probs_dropout_prob=0, + ), + text_config=dict( + model_type="opt", + vocab_size=64, + hidden_size=32, + ffn_dim=48, + num_hidden_layers=1, + num_attention_heads=4, + word_embed_proj_dim=32, + dropout=0, + attention_dropout=0, + use_cache=False, + ), + num_query_tokens=2, + image_token_id=63, + ) + config._attn_implementation = backend + model = hf.Blip2ForConditionalGeneration(config).eval() + inputs = torch.tensor([[2, 3, 4, 5], [2, 6, 7, 8]]) + inputs[:, :2] = 63 + labels = inputs.clone() + labels[:, :2] = -100 + return model, dict(pixel_values=torch.randn(2, 3, 16, 16), input_ids=inputs, labels=labels) + else: + config = hf.CohereConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=16, + attention_dropout=0, + use_cache=False, + attn_implementation=backend, + ) + model = hf.CohereForCausalLM(config) + inputs = torch.tensor([[2, 3, 4, 5], [2, 6, 7, 8]]) + return model.eval(), dict(input_ids=inputs, labels=inputs.clone(), use_cache=False) + + +@pytest.mark.parametrize("family", ["opt", "command", "blip2"]) +@pytest.mark.parametrize("backend", ["eager", "sdpa"]) +def test_native_training_integration(family, backend): + reference, data = model_and_data(family, backend) + candidate = copy.deepcopy(reference) + candidate, _ = ShardFormer(ShardConfig(enable_tensor_parallelism=False)).optimize(candidate) + expected, actual = reference(**data), candidate(**data) + torch.testing.assert_close(actual.logits, expected.logits, atol=1e-5, rtol=1e-3) + torch.testing.assert_close(actual.loss, expected.loss, atol=1e-5, rtol=1e-3) + expected.loss.backward() + actual.loss.backward() + compare_parameters(reference, candidate) + torch.optim.SGD(reference.parameters(), lr=0.01).step() + torch.optim.SGD(candidate.parameters(), lr=0.01).step() + compare_parameters(reference, candidate, gradients=False) + + +@pytest.mark.parametrize("family", ["opt", "command", "blip2"]) +@pytest.mark.parametrize( + "flag", ["enable_tensor_parallelism", "enable_flash_attention", "enable_jit_fused", "pipeline_stage_manager"] +) +def test_unsupported_configuration_error(family, flag): + config = SimpleNamespace(**{flag: True}) + with pytest.raises(NotImplementedError, match="HF5-COMPATIBILITY") as exc: + check_hf5_policy(family, config) + assert family in str(exc.value) + assert hf.__version__ in str(exc.value) + + +@pytest.mark.parametrize("family", ["chatglm2", "deepseek"]) +def test_remote_code_gate_before_import(family, monkeypatch): + from colossalai.shardformer.policies import auto_policy + + def forbidden_import(*args, **kwargs): + pytest.fail("Denied policy must not be imported") + + monkeypatch.setattr(auto_policy.importlib, "import_module", forbidden_import) + with pytest.raises(NotImplementedError, match="remote-code"): + import_policy(PolicyLocation(family, "UnusedPolicy")) + + +def test_legacy_version_not_gated(monkeypatch): + from colossalai.shardformer import _utils + + monkeypatch.setattr(_utils, "is_hf5", lambda: False) + assert hf5_policy_error("command", SimpleNamespace(enable_tensor_parallelism=True)) is None + + +@pytest.mark.parametrize("family", ["bert", "llama", "qwen2", "qwen3", "mixtral", "deepseek_v3", "vit", "whisper"]) +def test_core_not_blocked_by_tail_gate(family): + assert hf5_policy_error(family, SimpleNamespace(enable_tensor_parallelism=True)) is None + + +def test_unknown_model_keeps_original_error(): + from colossalai.shardformer.policies.auto_policy import get_autopolicy + + with pytest.raises(NotImplementedError, match="Auto policy"): + get_autopolicy(torch.nn.Linear(2, 2)) + + +def test_blip2_obsolete_dropout_targets_removed(): + from colossalai.shardformer.policies.auto_policy import get_autopolicy + + model, _ = model_and_data("blip2", "eager") + policy = get_autopolicy(model) + config = ShardConfig(enable_tensor_parallelism=False) + config.enable_tensor_parallelism = True + config._tensor_parallel_size = 2 + policy.set_model(model) + policy.set_shard_config(config) + policy.preprocess() + descriptions = policy.module_policy() + removed = {"self_attn.dropout", "attention.attention.dropout", "crossattention.attention.dropout"} + assert all( + replacement.suffix not in removed + for description in descriptions.values() + for replacement in description.sub_module_replacement or [] + ) + + +def test_explicit_builtin_policy_gate(): + from colossalai.shardformer.policies.opt import OPTPolicy + + model, _ = model_and_data("opt", "eager") + config = ShardConfig(enable_tensor_parallelism=False, enable_jit_fused=True) + before = copy.deepcopy(model.state_dict()) + with pytest.raises(NotImplementedError, match="enable_jit_fused"): + ShardFormer(config).optimize(model, policy=OPTPolicy()) + for name, parameter in model.state_dict().items(): + torch.testing.assert_close(parameter, before[name], atol=0, rtol=0) + + +def test_custom_policy_not_gated(): + from colossalai.shardformer.policies.base_policy import Policy + + class CustomPolicy(Policy): + def config_sanity_check(self): + pass + + def preprocess(self): + return self.model + + def module_policy(self): + return {} + + def postprocess(self): + return self.model + + model, _ = model_and_data("opt", "eager") + ShardFormer(ShardConfig(enable_tensor_parallelism=False, enable_jit_fused=True)).optimize(model, CustomPolicy()) diff --git a/tests/test_shardformer/test_hf5_compat/test_d5_pp_regressions.py b/tests/test_shardformer/test_hf5_compat/test_d5_pp_regressions.py new file mode 100644 index 000000000000..0bde2d72d8e7 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_d5_pp_regressions.py @@ -0,0 +1,93 @@ +"""Serial PP contract regressions; real distributed acceptance runs separately.""" + +import copy +from types import SimpleNamespace + +import pytest + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("HF5 pipeline contracts", allow_module_level=True) + +import torch +from transformers import MixtralConfig, MixtralModel + +from colossalai.shardformer import ShardConfig +from colossalai.shardformer.modeling._mixtral_tf5 import MixtralPipelineForwards +from colossalai.shardformer.policies._deepseek_v3_tf5 import DeepseekV3ModelPolicy +from colossalai.shardformer.shard.utils import set_tensors_to_none +from tests.test_shardformer.test_hf5_compat.test_integrated_core import model_and_data + + +def stage(index, total=2): + return SimpleNamespace( + is_interleave=False, + use_zbv=False, + is_first_stage=lambda **kwargs: index == 0, + is_last_stage=lambda **kwargs: index == total - 1, + distribute_layers=lambda count: [count // total] * total, + get_stage_index=lambda counts: (sum(counts[:index]), sum(counts[: index + 1])), + ) + + +@pytest.mark.parametrize("backend", ["eager", "sdpa"]) +def test_mixtral_pp_causal_mask_with_flash_installed(monkeypatch, backend): + from colossalai.shardformer.modeling import _mixtral_tf5 as implementation + + monkeypatch.setattr(implementation, "is_flash_attn_2_available", lambda: True) + config = MixtralConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + use_cache=False, + attn_implementation=backend, + ) + model = MixtralModel(config).eval() + candidate = copy.deepcopy(model) + tokens = torch.tensor([[1, 3, 4, 5], [1, 6, 7, 8]]) + expected = model(input_ids=tokens).last_hidden_state + + def run(ids): + first = MixtralPipelineForwards.mixtral_model_forward( + candidate, input_ids=ids, stage_manager=stage(0), stage_index=(0, 1), use_cache=False + ) + return MixtralPipelineForwards.mixtral_model_forward( + candidate, hidden_states=first["hidden_states"], stage_manager=stage(1), stage_index=(1, 2), use_cache=False + ).last_hidden_state + + actual = run(tokens) + torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-3) + changed = tokens.clone() + changed[:, -1] = 10 + # Expert token grouping can change GEMM rounding even for an unchanged + # prefix; use the existing FP32 gradient-level bound, not bitwise equality. + torch.testing.assert_close(run(changed)[:, :-1], actual[:, :-1], atol=1e-6, rtol=1e-4) + expected.square().mean().backward() + actual.square().mean().backward() + for ref, value in zip(model.parameters(), candidate.parameters()): + assert (ref.grad is None) == (value.grad is None) + if ref.grad is not None: + torch.testing.assert_close(value.grad, ref.grad, atol=1e-6, rtol=1e-4) + + +@pytest.mark.parametrize("index", [0, 1]) +def test_deepseek_pp_preserves_rotary_buffer(index): + model, _ = model_and_data("deepseek_v3") + model = model.model + policy = DeepseekV3ModelPolicy() + policy.set_model(model) + config = ShardConfig(enable_tensor_parallelism=False) + config.pipeline_stage_manager = stage(index) + policy.set_shard_config(config) + held = policy.get_held_layers() + assert model.rotary_emb in held + set_tensors_to_none(model, exclude=set(held)) + assert model.rotary_emb.inv_freq is not None + cos, sin = model.rotary_emb(torch.zeros(1, 4, 32), position_ids=torch.arange(4).unsqueeze(0)) + assert torch.isfinite(cos).all() and torch.isfinite(sin).all() diff --git a/tests/test_shardformer/test_hf5_compat/test_integrated_core.py b/tests/test_shardformer/test_hf5_compat/test_integrated_core.py new file mode 100644 index 000000000000..c81d1e9a5efa --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_integrated_core.py @@ -0,0 +1,158 @@ +import copy +import importlib +from types import SimpleNamespace + +import pytest + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("HF5 core migration contracts", allow_module_level=True) + +import torch +import transformers as hf + +from colossalai.shardformer import ShardConfig, ShardFormer +from colossalai.shardformer._utils import check_hf5_policy +from colossalai.shardformer.policies.auto_policy import get_autopolicy + + +def model_and_data(family): + common = dict( + vocab_size=64, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=16, + use_cache=False, + attn_implementation="eager", + bos_token_id=1, + eos_token_id=2, + pad_token_id=0, + ) + if family == "bert": + model = hf.BertForMaskedLM( + hf.BertConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + max_position_embeddings=16, + hidden_dropout_prob=0, + attention_probs_dropout_prob=0, + attn_implementation="eager", + ) + ) + elif family == "llama": + model = hf.LlamaForCausalLM(hf.LlamaConfig(**common)) + elif family == "mixtral": + model = hf.MixtralForCausalLM(hf.MixtralConfig(**common, num_local_experts=4, num_experts_per_tok=2)) + else: + common["num_key_value_heads"] = 4 + model = hf.DeepseekV3ForCausalLM( + hf.DeepseekV3Config( + **common, + q_lora_rank=None, + kv_lora_rank=8, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=8, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + n_group=1, + topk_group=1, + moe_intermediate_size=16, + first_k_dense_replace=0, + ) + ) + tokens = torch.tensor([[1, 3, 4, 5], [1, 6, 7, 8]]) + return model.eval(), dict(input_ids=tokens, labels=tokens.clone()) + + +@pytest.mark.parametrize("seed", [42, 123]) +@pytest.mark.parametrize("family", ["bert", "llama", "mixtral", "deepseek_v3"]) +def test_integrated_non_ep_training(family, seed): + torch.manual_seed(seed) + reference, data = model_and_data(family) + candidate = copy.deepcopy(reference) + policy = get_autopolicy(candidate) + assert type(policy).__module__.endswith(f"_{family}_tf5") + assert importlib.import_module("colossalai.shardformer.policies." + family).__dict__[type(policy).__name__] is type( + policy + ) + candidate, _ = ShardFormer(ShardConfig(enable_tensor_parallelism=False)).optimize(candidate) + ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01) + optimizer = torch.optim.SGD(candidate.parameters(), lr=0.01) + params = dict(candidate.named_parameters()) + ref_params = dict(reference.named_parameters()) + assert params.keys() == ref_params.keys() + optimizer_ids = {id(param) for group in optimizer.param_groups for param in group["params"]} + assert optimizer_ids == {id(param) for param in candidate.parameters()} + for step in range(2): + ref_optimizer.zero_grad() + optimizer.zero_grad() + if step == 1: + reference.gradient_checkpointing_enable() + candidate.gradient_checkpointing_enable() + reference.train() + candidate.train() + expected, actual = reference(**data), candidate(**data) + assert torch.isfinite(actual.logits).all() + torch.testing.assert_close(actual.logits, expected.logits, atol=1e-5, rtol=1e-3) + torch.testing.assert_close(actual.loss, expected.loss, atol=1e-5, rtol=1e-3) + expected.loss.backward() + actual.loss.backward() + for name, parameter in ref_params.items(): + other = params[name] + assert (parameter.grad is None) == (other.grad is None), name + if parameter.grad is not None: + assert torch.isfinite(other.grad).all(), name + torch.testing.assert_close(other.grad, parameter.grad, atol=1e-6, rtol=1e-4) + ref_optimizer.step() + optimizer.step() + for name, parameter in ref_params.items(): + torch.testing.assert_close(params[name], parameter, atol=1e-6, rtol=1e-4) + + +@pytest.mark.parametrize("family", ["mixtral", "deepseek_v3"]) +def test_repaired_ep_policy_is_not_blanket_gated(family): + check_hf5_policy(family, SimpleNamespace(expert_parallel_size=2)) + + +@pytest.mark.parametrize("family", ["mixtral", "deepseek_v3"]) +def test_explicit_dispatched_ep_policy_passes_compatibility_check(family): + model, _ = model_and_data(family) + policy = get_autopolicy(model) + config = ShardConfig(enable_tensor_parallelism=False) + config._expert_parallel_size = 2 + before = {name: id(parameter) for name, parameter in model.named_parameters()} + # Policy selection must not mutate weights; actual EP training/checkpoint + # acceptance belongs to the distributed optimizer regressions. + check_hf5_policy(family, config) + assert type(policy).__module__.endswith(f"_{family}_tf5") + assert before == {name: id(parameter) for name, parameter in model.named_parameters()} + + +@pytest.mark.parametrize("v5", [False, True]) +def test_dispatch_selects_actual_implementation(monkeypatch, v5): + from colossalai import _compat + + calls = [] + exported = SimpleNamespace(public_name=object(), _private_name=object(), __name__="implementation") + monkeypatch.setattr(_compat, "is_transformers_v5", lambda: v5) + + def import_impl(name): + calls.append(name) + return exported + + monkeypatch.setattr(_compat, "import_module", import_impl) + namespace = {} + _compat.reexport_transformers_impl(namespace, "example", "_llama") + assert calls == ["example._llama_tf5" if v5 else "example._llama_tf4"] + assert namespace["public_name"] is exported.public_name + assert namespace["_private_name"] is exported._private_name + assert "__name__" not in namespace diff --git a/tests/test_shardformer/test_hf5_compat/test_integration_alias.py b/tests/test_shardformer/test_hf5_compat/test_integration_alias.py new file mode 100644 index 000000000000..cb0c074a5bf4 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_integration_alias.py @@ -0,0 +1,95 @@ +import pytest +import torch +from torch import nn + +from colossalai.checkpoint_io.hybrid_parallel_checkpoint_io import _route_padded_parameter_aliases +from colossalai.shardformer.layer.parallel_module import PaddingParallelModule + + +class PaddedLinear(PaddingParallelModule): + def __init__(self, weight): + super().__init__(new_num_embeddings=4, old_num_embeddings=3, weight=weight) + + @staticmethod + def from_native_module(module, process_group=None): + raise NotImplementedError + + +class AliasedModel(nn.Module): + def __init__(self): + super().__init__() + self.weight_alias = nn.Parameter(torch.zeros(3, 2)) + self.parallel = PaddedLinear(self.weight_alias) + + +def test_alias_route_and_identity(): + model = AliasedModel() + checkpoint = {"weight_alias": torch.arange(6).float().reshape(3, 2)} + _route_padded_parameter_aliases(model, checkpoint) + result = model.load_state_dict(checkpoint, strict=False) + assert result.missing_keys == ["weight_alias"] + assert not result.unexpected_keys + assert model.weight_alias is model.parallel.weight + torch.testing.assert_close(model.parallel.weight[:3], torch.arange(6).float().reshape(3, 2)) + assert torch.count_nonzero(model.parallel.weight[3:]) == 0 + + +def test_conflicting_alias_values(): + state = {"weight_alias": torch.ones(3, 2), "parallel.weight": torch.zeros(3, 2)} + with pytest.raises(ValueError, match="Conflicting checkpoint"): + _route_padded_parameter_aliases(AliasedModel(), state) + assert set(state) == {"weight_alias", "parallel.weight"} + + +def test_later_alias_conflict_does_not_mutate_checkpoint(): + model = nn.ModuleDict({"first": AliasedModel(), "second": AliasedModel()}) + state = { + "first.weight_alias": torch.ones(3, 2), + "second.weight_alias": torch.ones(3, 2), + "second.parallel.weight": torch.zeros(3, 2), + } + before = dict(state) + with pytest.raises(ValueError, match="Conflicting checkpoint"): + _route_padded_parameter_aliases(model, state) + assert state.keys() == before.keys() + for name in before: + assert state[name] is before[name] + + +def test_duplicate_module_aliases_are_routed(): + model = AliasedModel() + model.other_parallel = model.parallel + state = {"weight_alias": torch.ones(3, 2)} + _route_padded_parameter_aliases(model, state) + assert set(state) == {"parallel.weight", "other_parallel.weight"} + model.load_state_dict(state, strict=False) + assert model.parallel.weight is model.other_parallel.weight + + +def test_saver_aliases_with_padding_roundtrip(): + model = AliasedModel() + with torch.no_grad(): + model.weight_alias[:3].fill_(2) + state = model.state_dict() + assert state["weight_alias"].shape == (4, 2) + assert state["parallel.weight"].shape == (3, 2) + _route_padded_parameter_aliases(model, state) + target = AliasedModel() + target.load_state_dict(state, strict=False) + torch.testing.assert_close(target.parallel.weight, model.parallel.weight) + + +def test_nonzero_padding_is_not_silently_discarded(): + model = AliasedModel() + state = {"weight_alias": torch.ones(4, 2)} + with pytest.raises(ValueError, match="Nonzero padding"): + _route_padded_parameter_aliases(model, state) + assert set(state) == {"weight_alias"} + + +def test_strict_alias_missing_key_is_visible(): + model = AliasedModel() + state = {"weight_alias": torch.ones(3, 2)} + _route_padded_parameter_aliases(model, state) + with pytest.raises(RuntimeError, match="weight_alias"): + model.load_state_dict(state, strict=True) diff --git a/tests/test_shardformer/test_hf5_compat/test_integration_loading.py b/tests/test_shardformer/test_hf5_compat/test_integration_loading.py new file mode 100644 index 000000000000..84ff551f73a1 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_integration_loading.py @@ -0,0 +1,120 @@ +"""Offline integration of config saving and deferred lazy pretrained setup.""" + +import json + +import pytest +import torch +from transformers import OPTConfig, OPTForCausalLM + +from colossalai.checkpoint_io.utils import save_config_file +from colossalai.interface.pretrained import get_pretrained_path +from colossalai.lazy import LazyInitContext, LazyTensor + + +def small_model(): + return OPTForCausalLM( + OPTConfig( + vocab_size=64, + hidden_size=32, + word_embed_proj_dim=32, + ffn_dim=48, + num_hidden_layers=1, + num_attention_heads=4, + max_position_embeddings=16, + bos_token_id=1, + eos_token_id=2, + pad_token_id=0, + ) + ) + + +def test_config_roundtrip(tmp_path): + model = small_model().bfloat16() + save_config_file(model, str(tmp_path)) + saved = json.loads((tmp_path / "config.json").read_text()) + assert saved.get("dtype", saved.get("torch_dtype")) == "bfloat16" + config = OPTConfig.from_pretrained(tmp_path, local_files_only=True) + assert config.hidden_size == 32 + assert (tmp_path / "generation_config.json").exists() + + +def test_config_nonmaster_and_nonhf(tmp_path): + save_config_file(small_model(), str(tmp_path), is_master=False) + save_config_file(torch.nn.Linear(2, 2), str(tmp_path)) + assert not (tmp_path / "config.json").exists() + + +def test_lazy_pretrained_setup(tmp_path): + model = small_model() + model.save_pretrained(tmp_path) + initial_dtype = torch.get_default_dtype() + with LazyInitContext(): + candidate = OPTForCausalLM.from_pretrained(tmp_path, local_files_only=True, dtype=torch.bfloat16) + assert isinstance(next(candidate.parameters()), LazyTensor) + assert torch.get_default_dtype() == initial_dtype + assert next(candidate.parameters()).dtype == torch.bfloat16 + assert get_pretrained_path(candidate).endswith("model.safetensors") + # Lazy from_pretrained stores the checkpoint path; materialization alone + # does not load saved weights. Verify explicit subsequent weight loading. + LazyInitContext.materialize(candidate) + # HF5 may initialize selected norms in FP32; convert the fully materialized + # model explicitly before exercising this CPU BF16 comparison. + candidate.bfloat16() + candidate.load_state_dict(model.state_dict()) + inputs = torch.tensor([[1, 3, 4]]) + reference = model.bfloat16().eval() + torch.testing.assert_close(candidate(input_ids=inputs).logits, reference(input_ids=inputs).logits) + + +def test_booster_loads_deferred_checkpoint(tmp_path): + from colossalai.booster import Booster + + model = small_model().eval() + with torch.no_grad(): + for index, parameter in enumerate(model.parameters()): + parameter.fill_((index + 1) / 128) + model.save_pretrained(tmp_path, max_shard_size="4KB") + with LazyInitContext(): + candidate = OPTForCausalLM.from_pretrained(tmp_path, local_files_only=True) + assert get_pretrained_path(candidate).endswith("model.safetensors.index.json") + LazyInitContext.materialize(candidate) + candidate, *_ = Booster().boost(candidate) + assert get_pretrained_path(candidate) is None + expected = model.state_dict() + for name, tensor in candidate.state_dict().items(): + torch.testing.assert_close(tensor, expected[name], atol=0, rtol=0) + inputs = torch.tensor([[1, 3, 4]]) + torch.testing.assert_close(candidate(input_ids=inputs).logits, model(input_ids=inputs).logits) + + +def test_constructor_failure_restores_default_dtype(tmp_path, monkeypatch): + model = small_model() + model.save_pretrained(tmp_path) + original_dtype = torch.get_default_dtype() + + def fail_constructor(*args, **kwargs): + raise RuntimeError("deliberate constructor failure") + + monkeypatch.setattr(OPTForCausalLM, "__init__", fail_constructor) + with LazyInitContext(), pytest.raises(RuntimeError, match="deliberate constructor"): + OPTForCausalLM.from_pretrained(tmp_path, dtype=torch.float64, local_files_only=True) + assert torch.get_default_dtype() == original_dtype + + +def test_strict_booster_still_rejects_missing_untied_weight(tmp_path): + from colossalai.booster import Booster + + model = small_model() + model.config.tie_word_embeddings = False + model.lm_head = torch.nn.Linear(32, 64, bias=False) + state = model.state_dict() + state.pop("lm_head.weight") + torch.save(state, tmp_path / "incomplete.bin") + with pytest.raises(RuntimeError, match="lm_head.weight"): + Booster().load_model(model, str(tmp_path / "incomplete.bin"), strict=True) + + +def test_auth_conflict_does_not_expose_values(tmp_path): + with LazyInitContext(), pytest.raises(ValueError, match="Conflicting token") as exc: + OPTForCausalLM.from_pretrained(tmp_path, token="first-private-value", use_auth_token="second-private-value") + assert "private-value" not in str(exc.value) diff --git a/tests/test_shardformer/test_hf5_compat/test_qwen.py b/tests/test_shardformer/test_hf5_compat/test_qwen.py new file mode 100644 index 000000000000..75facb0cf70a --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_qwen.py @@ -0,0 +1,410 @@ +"""CPU regression tests of real ShardFormer Qwen forwards (no GPU kernels or collectives). + +The pipeline checks execute stage functions serially; they do not validate a +distributed pipeline schedule. Run this directory with --confcutdir to avoid the +repository-wide accelerator fixture on CPU-only development machines. +""" + +import copy +import importlib +from types import MethodType, SimpleNamespace + +import pytest +import torch +import transformers +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.17.0"): + pytest.skip("HF 5.17 contract", allow_module_level=True) + +from transformers.cache_utils import DynamicCache + +from colossalai.shardformer.shard import ShardConfig + +pytestmark = pytest.mark.skipif(Version(transformers.__version__) < Version("5.17.0"), reason="HF 5.17 contract") + + +def make_pair(family, *, head_dim=8, sliding=False, backend="eager", rope=False, dropout=0.0, head="ForCausalLM"): + prefix = "Qwen2" if family == "qwen2" else "Qwen3" + config = getattr(transformers, prefix + "Config")( + vocab_size=64, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + max_position_embeddings=64, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + attention_dropout=dropout, + attn_implementation=backend, + use_cache=False, + **({"head_dim": head_dim} if family == "qwen3" else {}), + **( + {"use_sliding_window": True, "sliding_window": 3, "layer_types": ["full_attention", "sliding_attention"]} + if sliding + else {} + ), + **({"rope_parameters": {"rope_type": "linear", "factor": 2.0, "rope_theta": 10000.0}} if rope else {}), + ) + reference = getattr(transformers, prefix + head)(config).eval() + return reference, copy.deepcopy(reference) + + +def local_module(family): + return importlib.import_module("colossalai.shardformer.modeling." + family) + + +def patch_model(model, family, *, attention=True): + module = local_module(family) + config = ShardConfig(enable_tensor_parallelism=False) + model.forward = MethodType(getattr(module, f"get_{family}_model_forward_for_flash_attn")(config), model) + if attention: + for layer in model.layers: + layer.self_attn.forward = MethodType( + getattr(module, f"get_{family}_flash_attention_forward")(config), layer.self_attn + ) + + +def inputs(padding="none"): + ids = torch.tensor([[3, 7, 9, 11, 13, 15], [4, 8, 10, 12, 14, 16]]) + if padding == "left": + ids[0, :2] = 0 + elif padding == "right": + ids[0, -2:] = 0 + mask = ids.ne(0).long() + positions = (mask.cumsum(-1) - 1).clamp_min(0) + labels = ids.masked_fill(mask == 0, -100) + return dict(input_ids=ids, attention_mask=mask, position_ids=positions, labels=labels, use_cache=False) + + +def close(actual, reference, *, atol=1e-5, rtol=1e-3): + assert torch.isfinite(actual).all() + assert actual.shape == reference.shape + torch.testing.assert_close(actual, reference, atol=atol, rtol=rtol) + + +@pytest.mark.parametrize( + "padding,backend,sliding,rope", + [ + ("none", "eager", False, False), + ("left", "eager", False, False), + ("right", "sdpa", False, False), + ("none", "eager", True, False), + ("left", "eager", False, True), + ], +) +def test_forward_loss_gradients_and_step(family, padding, backend, sliding, rope): + reference, actual = make_pair(family, backend=backend, sliding=sliding, rope=rope, head_dim=16) + patch_model(actual.model, family) + data = inputs(padding) + ref_out, out = reference(**data), actual(**data) + close(out.logits[data["attention_mask"].bool()], ref_out.logits[data["attention_mask"].bool()]) + close(out.loss, ref_out.loss) + ref_out.loss.backward() + out.loss.backward() + compared = 0 + for (name, param), (ref_name, ref_param) in zip(actual.named_parameters(), reference.named_parameters()): + assert name == ref_name + assert param.grad is not None and ref_param.grad is not None, name + close(param.grad, ref_param.grad, atol=1e-6, rtol=1e-4) + compared += 1 + assert compared > 0 + torch.optim.SGD(actual.parameters(), lr=0.01).step() + torch.optim.SGD(reference.parameters(), lr=0.01).step() + for param, ref_param in zip(actual.parameters(), reference.parameters()): + close(param, ref_param, atol=1e-4, rtol=1e-3) + + +@pytest.mark.parametrize("sliding", [False, True]) +def test_cache_prefill_and_decode(family, sliding): + reference, actual = make_pair(family, sliding=sliding, head_dim=16) + patch_model(actual.model, family) + ids = inputs()["input_ids"] + ids = torch.cat((ids, torch.tensor([[17], [18]])), dim=1) + caches = [DynamicCache(config=m.config) for m in (reference, actual)] + with torch.no_grad(): + for end in range(4, 8): + start = 0 if end == 4 else end - 1 + outputs = [ + m(input_ids=ids[:, start:end], attention_mask=torch.ones(2, end), past_key_values=cache, use_cache=True) + for m, cache in zip((reference, actual), caches) + ] + full = reference(input_ids=ids[:, :end], use_cache=False) + close(outputs[1].logits[:, -1], outputs[0].logits[:, -1]) + close(outputs[1].logits[:, -1], full.logits[:, -1]) + assert outputs[1].past_key_values is caches[1] + for layer_idx in range(2): + assert caches[1].get_seq_length(layer_idx) == caches[0].get_seq_length(layer_idx) == end + + +def test_automatic_cache_and_output_capture(family): + reference, actual = make_pair(family) + patch_model(actual.model, family) + kwargs = dict(input_ids=inputs()["input_ids"], use_cache=True, output_hidden_states=True, output_attentions=True) + ref_out, out = reference(**kwargs), actual(**kwargs) + assert isinstance(out.past_key_values, DynamicCache) + assert len(out.hidden_states) == len(ref_out.hidden_states) == 3 + assert len(out.attentions) == len(ref_out.attentions) == 2 + for a, b in zip(out.hidden_states + out.attentions, ref_out.hidden_states + ref_out.attentions): + close(a, b) + tup = actual.model(input_ids=inputs()["input_ids"], use_cache=False, return_dict=False) + assert isinstance(tup, tuple) and tup[0].shape == (2, 6, 32) + + +@pytest.mark.parametrize("reentrant", [False, True]) +def test_gradient_checkpointing(family, reentrant): + reference, actual = make_pair(family) + patch_model(actual.model, family) + actual.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": reentrant}) + actual.train() + reference.train() + ref_out, out = reference(**inputs()), actual(**inputs()) + close(out.loss, ref_out.loss) + ref_out.loss.backward() + out.loss.backward() + for a, b in zip(actual.parameters(), reference.parameters()): + assert a.grad is not None and b.grad is not None + close(a.grad, b.grad, atol=1e-6, rtol=1e-4) + + +def stage(index): + return SimpleNamespace( + stage=index, + num_stages=2, + is_interleave=False, + num_model_chunks=1, + is_first_stage=lambda: index == 0, + is_last_stage=lambda: index == 1, + ) + + +@pytest.mark.parametrize("checkpoint", [False, True]) +def test_pipeline_stage_functions_serially(family, checkpoint): + reference, actual = make_pair(family) + config = ShardConfig(enable_tensor_parallelism=False) + forward = getattr( + getattr(local_module(family), "Qwen" + family[-1] + "PipelineForwards"), f"{family}_model_forward" + ) + if checkpoint: + actual.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True}) + actual.train() + data = inputs("left") + data.pop("labels") + ref_out = reference.model(**data) + first = forward(actual.model, **data, stage_manager=stage(0), stage_index=[0, 1], shard_config=config) + out = forward( + actual.model, + hidden_states=first["hidden_states"], + attention_mask=data["attention_mask"], + use_cache=False, + stage_manager=stage(1), + stage_index=[1, 2], + shard_config=config, + position_ids=data["position_ids"], + ) + close(out.last_hidden_state, ref_out.last_hidden_state) + out.last_hidden_state.square().mean().backward() + ref_out.last_hidden_state.square().mean().backward() + for a, b in zip(actual.model.parameters(), reference.model.parameters()): + assert a.grad is not None and b.grad is not None + close(a.grad, b.grad, atol=1e-6, rtol=1e-4) + + +def test_attention_training_dropout(family): + reference, actual = make_pair(family, dropout=0.2) + patch_model(actual.model, family) + reference.train() + actual.train() + torch.manual_seed(123) + ref_out = reference(**inputs()) + torch.manual_seed(123) + out = actual(**inputs()) + close(out.loss, ref_out.loss) + + +def test_pipeline_sequence_classification_left_padding(family): + reference, actual = make_pair(family, head="ForSequenceClassification") + forward = getattr( + getattr(local_module(family), "Qwen" + family[-1] + "PipelineForwards"), + f"{family}_for_sequence_classification_forward", + ) + data = inputs("left") + data["labels"] = torch.tensor([0, 1]) + manager = SimpleNamespace(is_first_stage=lambda: True, is_last_stage=lambda: True) + out = forward( + actual, + **data, + stage_manager=manager, + stage_index=[0, 2], + shard_config=ShardConfig(enable_tensor_parallelism=False), + ) + ref_out = reference(**data) + close(out.logits, ref_out.logits) + close(out.loss, ref_out.loss) + + +@pytest.mark.parametrize("backend", ["eager", "sdpa"]) +def test_precomputed_mask_mapping(family, backend): + from transformers.masking_utils import create_causal_mask + + reference, actual = make_pair(family, backend=backend) + patch_model(actual.model, family) + ids = inputs()["input_ids"] + mask = create_causal_mask(reference.config, reference.model.embed_tokens(ids), None, None) + out = actual(input_ids=ids, attention_mask={"full_attention": mask}, use_cache=False) + ref_out = reference(input_ids=ids, use_cache=False) + close(out.logits, ref_out.logits) + + +def test_left_padding_chunked_decode(family): + reference, actual = make_pair(family) + patch_model(actual.model, family) + data = inputs("left") + cache = None + for start, end in ((0, 4), (4, 6)): + out = actual( + input_ids=data["input_ids"][:, start:end], + attention_mask=data["attention_mask"][:, :end], + position_ids=data["position_ids"][:, start:end], + past_key_values=cache, + use_cache=True, + ) + cache = out.past_key_values + ref_out = reference( + input_ids=data["input_ids"][:, :end], + attention_mask=data["attention_mask"][:, :end], + position_ids=data["position_ids"][:, :end], + use_cache=False, + ) + close(out.logits[:, -1], ref_out.logits[:, -1]) + assert cache.get_seq_length() == end + + +def test_pipeline_default_positions_and_checkpoint_selection(family): + from colossalai.shardformer.shard.grad_ckpt_config import PipelineGradientCheckpointConfig + + reference, actual = make_pair(family) + actual.train() + actual.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True}) + config = ShardConfig( + enable_tensor_parallelism=False, + gradient_checkpoint_config=PipelineGradientCheckpointConfig(num_ckpt_layers_per_stage=[1, 0]), + ) + calls = [0, 0] + + def count(index): + def hook(module, args): + calls[index] += 1 + + return hook + + handles = [layer.register_forward_pre_hook(count(i)) for i, layer in enumerate(actual.model.layers)] + forward = getattr( + getattr(local_module(family), "Qwen" + family[-1] + "PipelineForwards"), f"{family}_model_forward" + ) + first = forward( + actual.model, + input_ids=inputs()["input_ids"], + use_cache=False, + stage_manager=stage(0), + stage_index=[0, 1], + shard_config=config, + ) + out = forward( + actual.model, + hidden_states=first["hidden_states"], + use_cache=False, + stage_manager=stage(1), + stage_index=[1, 2], + shard_config=config, + ) + ref_out = reference.model(input_ids=inputs()["input_ids"], use_cache=False) + close(out.last_hidden_state, ref_out.last_hidden_state) + out.last_hidden_state.square().mean().backward() + assert calls == [2, 1] # Only the selected first layer is recomputed, exactly once. + for handle in handles: + handle.remove() + + +def test_lm_wrapper_generation_arguments(family): + reference, actual = make_pair(family) + patch_model(actual.model, family) + actual.forward = MethodType( + local_module(family).get_lm_forward_with_dist_cross_entropy(ShardConfig(enable_tensor_parallelism=False)), + actual, + ) + kwargs = dict(input_ids=inputs()["input_ids"], logits_to_keep=1, use_cache=True) + ref_out, out = reference(**kwargs), actual(**kwargs) + assert out.logits.shape == (2, 1, 64) + close(out.logits, ref_out.logits) + assert out.past_key_values.get_seq_length() == 6 + + +def test_zbv_policy_without_tensor_parallelism(family): + reference, _ = make_pair(family) + module = importlib.import_module("colossalai.shardformer.policies." + family) + policy = getattr(module, "Qwen" + family[-1] + "Policy")() + policy.set_model(reference) + policy.set_shard_config( + ShardConfig(enable_tensor_parallelism=False, pipeline_stage_manager=SimpleNamespace(use_zbv=True)) + ) + policy.preprocess() + descriptions = policy.module_policy() + replacements = descriptions[type(reference.model.layers[0])].sub_module_replacement + assert any(replacement.suffix == "self_attn.q_proj" for replacement in replacements) + + +def test_cache_length_helper(family): + from colossalai.shardformer._utils import get_hf_cache_seq_length + + reference, _ = make_pair(family) + cache = DynamicCache(config=reference.config) + assert get_hf_cache_seq_length(None) == get_hf_cache_seq_length(cache) == 0 + kv = torch.randn(2, 2, 4, 8) + for idx in range(2): + cache.update(kv, kv, idx) + assert get_hf_cache_seq_length(cache, idx) == 4 + cache.update(kv[:, :, :1], kv[:, :, :1], idx) + assert get_hf_cache_seq_length(cache, idx) == 5 + + +def test_flash_mask_contract(family): + from colossalai.shardformer._utils import get_hf_flash_attention_kwargs, prepare_hf_causal_masks + from colossalai.shardformer.layer import AttnMaskType + + reference, _ = make_pair(family, sliding=True) + data = inputs("left") + masks = prepare_hf_causal_masks( + reference.config, + reference.model.embed_tokens(data["input_ids"]), + data["attention_mask"], + None, + data["position_ids"], + eager=True, + ) + for layer_type, mask in masks.items(): + kwargs = get_hf_flash_attention_kwargs(mask) + assert kwargs["attention_mask_type"] == AttnMaskType.CUSTOM + # This checks the real mask bridge only; it never substitutes a fake GPU kernel. + close(kwargs["attention_mask"][0, :, 2:], mask[0, :, 2:]) + assert (kwargs["attention_mask"][0, :, :2] == 0).all() + assert mask[1, 0, 2, 3] < -1e20 # Future token stays masked. + if layer_type == "sliding_attention": + assert mask[1, 0, 5, 1] < -1e20 # Token outside the window stays masked. + + +def test_shardformer_cpu_policy_integration(family): + from colossalai.shardformer import ShardFormer + + reference, actual = make_pair(family) + actual, _ = ShardFormer(ShardConfig(enable_tensor_parallelism=False)).optimize(actual) + out, ref_out = actual(**inputs()), reference(**inputs()) + close(out.logits, ref_out.logits) + close(out.loss, ref_out.loss) + out.loss.backward() + ref_out.loss.backward() + for a, b in zip(actual.parameters(), reference.parameters()): + assert a.grad is not None and b.grad is not None + close(a.grad, b.grad, atol=1e-6, rtol=1e-4) diff --git a/tests/test_shardformer/test_hf5_compat/test_qwen3_norm_sync.py b/tests/test_shardformer/test_hf5_compat/test_qwen3_norm_sync.py new file mode 100644 index 000000000000..375ed11602a2 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_qwen3_norm_sync.py @@ -0,0 +1,110 @@ +"""Qwen3 head-shared norm gradients, independently checked on CPU/Gloo.""" + +from types import SimpleNamespace + +import pytest + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("HF5 Qwen3 norm contracts", allow_module_level=True) + +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from transformers import Qwen3Config, Qwen3Model + +from colossalai.shardformer.policies.qwen3 import Qwen3Policy + + +def _model(): + return Qwen3Model( + Qwen3Config( + vocab_size=32, + hidden_size=32, + intermediate_size=48, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + ) + ) + + +def _policy(model, group, size): + policy = Qwen3Policy() + policy.model = model + policy.shard_config = SimpleNamespace( + enable_tensor_parallelism=size > 1, + tensor_parallel_size=size, + tensor_parallel_process_group=group, + ) + return policy + + +def _worker(rank, rendezvous, dtype, seed): + torch.set_num_threads(1) + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=2) + try: + torch.manual_seed(seed) + reference = _model().to(dtype) + actual = _model().to(dtype) + actual.load_state_dict(reference.state_dict()) + policy = _policy(actual, dist.group.WORLD, 2) + policy.postprocess() + policy.postprocess() # Repeated installation must not multiply gradients. + ref_attention = reference.layers[0].self_attn + attention = actual.layers[0].self_attn + ref_optimizer = torch.optim.SGD(reference.parameters(), lr=0.01) + optimizer = torch.optim.SGD(actual.parameters(), lr=0.01) + for step in range(2): + if step == 1: + actual.gradient_checkpointing_enable() + policy.postprocess() + ref_optimizer.zero_grad() + optimizer.zero_grad() + for _ in range(2): + for name, heads in (("q_norm", 4), ("k_norm", 2)): + x = torch.randn(2, 3, heads, 8).to(dtype) + upstream = torch.randn_like(x) / 32 + expected = getattr(ref_attention, name)(x) + local_input = x.chunk(2, 2)[rank].contiguous() + if step == 1: + from torch.utils.checkpoint import checkpoint + + local = checkpoint(getattr(attention, name), local_input, use_reentrant=False) + else: + local = getattr(attention, name)(local_input) + expected.backward(upstream) + local.backward(upstream.chunk(2, 2)[rank].contiguous()) + for name in ("q_norm", "k_norm"): + r, a = getattr(ref_attention, name).weight, getattr(attention, name).weight + tolerance = 0.005 if dtype == torch.bfloat16 else 1e-5 + torch.testing.assert_close(a.grad, r.grad, atol=tolerance, rtol=tolerance) + replicas = [torch.empty_like(a.grad) for _ in range(2)] + dist.all_gather(replicas, a.grad) + torch.testing.assert_close(replicas[0], replicas[1], atol=0, rtol=0) + ref_optimizer.step() + optimizer.step() + for name in ("q_norm", "k_norm"): + torch.testing.assert_close( + getattr(attention, name).weight, getattr(ref_attention, name).weight, atol=tolerance, rtol=tolerance + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("seed", [42, 123]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_head_norm_gradient_sum(tmp_path, dtype, seed): + mp.spawn(_worker, args=((tmp_path / "gloo-init").as_uri(), dtype, seed), nprocs=2, join=True) + + +def test_singleton_and_frozen_norms(): + model = _model() + _policy(model, None, 1).postprocess() + assert not getattr(model.layers[0].self_attn.q_norm.weight, "_qwen3_tp_grad_hook", False) + for norm in (model.layers[0].self_attn.q_norm, model.layers[0].self_attn.k_norm): + norm.weight.requires_grad_(False) + _policy(model, None, 2).postprocess() + assert not getattr(model.layers[0].self_attn.q_norm.weight, "_qwen3_tp_grad_hook", False) diff --git a/tests/test_shardformer/test_hf5_compat/test_qwen_loss.py b/tests/test_shardformer/test_hf5_compat/test_qwen_loss.py new file mode 100644 index 000000000000..31e9003c1dda --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_qwen_loss.py @@ -0,0 +1,57 @@ +"""Qwen loss-wrapper checks with a real, single-rank CPU Gloo process group. + +This exercises dist_cross_entropy's non-TP path, not vocabulary-parallel CE. +""" + +from types import MethodType, SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import transformers +from packaging.version import Version + +if Version(transformers.__version__) < Version("5.17.0"): + pytest.skip("HF 5.17 contract", allow_module_level=True) + +from test_qwen import close, inputs, local_module, make_pair, patch_model + +from colossalai.shardformer.shard import ShardConfig + +pytestmark = pytest.mark.skipif(Version(transformers.__version__) < Version("5.17.0"), reason="HF 5.17 contract") + + +@pytest.fixture(scope="module") +def single_rank_gloo(tmp_path_factory): + if dist.is_initialized(): + pytest.skip("This isolated CPU test must own its single-rank process group") + rendezvous = tmp_path_factory.mktemp("qwen-gloo") / "rendezvous" + dist.init_process_group("gloo", init_method=rendezvous.as_uri(), rank=0, world_size=1) + yield + dist.destroy_process_group() + + +@pytest.mark.parametrize("pipeline,dtype", [(False, torch.float32), (True, torch.float32), (True, torch.bfloat16)]) +def test_lm_loss_wrapper(single_rank_gloo, family, pipeline, dtype): + reference, actual = make_pair(family) + reference, actual = reference.to(dtype), actual.to(dtype) + module = local_module(family) + config = ShardConfig(enable_tensor_parallelism=False) + if pipeline: + forward = getattr(getattr(module, "Qwen" + family[-1] + "PipelineForwards"), f"{family}_for_causal_lm_forward") + manager = SimpleNamespace(is_first_stage=lambda: True, is_last_stage=lambda: True) + kwargs = dict(stage_manager=manager, stage_index=[0, 2], shard_config=config) + else: + patch_model(actual.model, family) + forward = module.get_lm_forward_with_dist_cross_entropy(config) + kwargs = {} + actual.forward = MethodType(forward, actual) + ref_out, out = reference(**inputs("left")), actual(**inputs("left"), **kwargs) + assert out.loss.dtype == ref_out.loss.dtype == torch.float32 + close(out.logits, ref_out.logits) + close(out.loss, ref_out.loss) + ref_out.loss.backward() + out.loss.backward() + for a, b in zip(actual.parameters(), reference.parameters()): + assert a.grad is not None and b.grad is not None + close(a.grad, b.grad, atol=1e-6, rtol=1e-4) diff --git a/tests/test_shardformer/test_hf5_compat/test_vit_whisper.py b/tests/test_shardformer/test_hf5_compat/test_vit_whisper.py new file mode 100644 index 000000000000..a304fc92b728 --- /dev/null +++ b/tests/test_shardformer/test_hf5_compat/test_vit_whisper.py @@ -0,0 +1,674 @@ +"""HF 5.17 CPU contracts. Serial stages are NOT distributed PP validation. + +Flash tests replace only the GPU kernel with explicit CPU attention; the real +ColossalAI projection/cache/mask/return adapters run, but no GPU kernel runs. +""" + +import copy +import importlib +import subprocess +import sys +from types import MethodType, SimpleNamespace + +import pytest + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("HF5 ViT and Whisper contracts", allow_module_level=True) + +import torch +import transformers as hf + +from colossalai.shardformer.shard import ShardConfig + + +def local(family): + return importlib.import_module("colossalai.shardformer.modeling." + family) + + +def stage(index, total=2): + return SimpleNamespace( + stage=index, + num_stages=total, + is_interleave=False, + use_zbv=False, + is_first_stage=lambda **kw: index == 0, + is_last_stage=lambda **kw: index == total - 1, + ) + + +def pair(family, variant="base", backend="eager"): + if family == "vit": + config = hf.ViTConfig( + hidden_size=32, + intermediate_size=48, + num_hidden_layers=2, + num_attention_heads=4, + image_size=16, + patch_size=8, + encoder_stride=8, + num_labels=3, + attn_implementation=backend, + ) + cls = { + "base": hf.ViTModel, + "classification": hf.ViTForImageClassification, + "masked": hf.ViTForMaskedImageModeling, + }[variant] + else: + config = hf.WhisperConfig( + vocab_size=32, + d_model=32, + encoder_layers=2, + decoder_layers=2, + encoder_attention_heads=4, + decoder_attention_heads=2, + encoder_ffn_dim=48, + decoder_ffn_dim=48, + num_mel_bins=8, + max_source_positions=8, + max_target_positions=16, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + decoder_start_token_id=1, + num_labels=3, + classifier_proj_size=16, + use_cache=False, + apply_spec_augment=False, + attn_implementation=backend, + ) + cls = { + "base": hf.WhisperModel, + "generation": hf.WhisperForConditionalGeneration, + "classification": hf.WhisperForAudioClassification, + }[variant] + reference = cls(config).eval() + return reference, copy.deepcopy(reference) + + +def data(family, variant="base"): + if family == "vit": + result = {"pixel_values": torch.randn(2, 3, 16, 16)} + if variant == "classification": + result["labels"] = torch.tensor([0, 2]) + if variant == "masked": + result["bool_masked_pos"] = torch.tensor([[1, 0, 0, 1], [0, 1, 1, 0]], dtype=torch.bool) + else: + result = {"input_features": torch.randn(2, 8, 16)} + if variant != "classification": + result.update(decoder_input_ids=torch.tensor([[1, 4, 5, 6], [1, 7, 8, 9]]), use_cache=False) + if variant == "generation": + result["labels"] = torch.tensor([[4, 5, 6, 2], [7, 8, 9, 2]]) + if variant == "classification": + result["labels"] = torch.tensor([0, 2]) + return result + + +def close(a, b, atol=1e-5, rtol=1e-3): + assert a.shape == b.shape + assert torch.isfinite(a).all() + torch.testing.assert_close(a, b, atol=atol, rtol=rtol) + + +def compare_training(reference, actual, ref_out, out, record_property, logits_mask=None): + for key in ("last_hidden_state", "pooler_output", "logits", "reconstruction", "loss", "encoder_last_hidden_state"): + if key in ref_out: + value, expected = out[key], ref_out[key] + if key == "logits" and logits_mask is not None: + # Fully masked padding queries have no supervised target; the + # custom-mask kernel unblocks these rows. Compare valid tokens. + value, expected = value[logits_mask], expected[logits_mask] + close(value, expected) + record_property(key + "_max_abs_error", (value - expected).abs().max().item()) + + def objective(outputs): + if "loss" in outputs: + return outputs.loss + result = outputs.last_hidden_state.square().mean() + if "pooler_output" in outputs: + result = result + outputs.pooler_output.square().mean() + return result + + objective(ref_out).backward() + objective(out).backward() + compared = 0 + max_error = 0.0 + ref_params = dict(reference.named_parameters()) + for name, p in actual.named_parameters(): + r = ref_params[name] + assert (p.grad is None) == (r.grad is None), name + if p.grad is not None: + close(p.grad, r.grad, atol=1e-6, rtol=1e-4) + compared += 1 + max_error = max(max_error, (p.grad - r.grad).abs().max().item()) + assert compared > 10 + record_property("gradient_parameters_compared", compared) + record_property("gradient_max_abs_error", max_error) + torch.optim.SGD(reference.parameters(), lr=0.01).step() + torch.optim.SGD(actual.parameters(), lr=0.01).step() + for name, p in actual.named_parameters(): + close(p, ref_params[name], atol=1e-6, rtol=1e-4) + + +@pytest.fixture +def cpu_flash(monkeypatch): + from colossalai.shardformer.layer import ColoAttention + + calls = [] + + def attention(q, k, v, attention_mask=None, attention_mask_type=None, dropout_p=0.0, scale=None, **kw): + calls.append((tuple(q.shape), tuple(k.shape), dropout_p, scale)) + scores = q @ k.transpose(-2, -1) * (q.shape[-1] ** -0.5 if scale is None else scale) + if attention_mask is not None: + scores = scores + attention_mask + weights = scores.softmax(-1, dtype=torch.float32).to(q.dtype) + weights = torch.nn.functional.dropout(weights, p=dropout_p, training=dropout_p > 0) + return weights @ v + + monkeypatch.setattr(ColoAttention, "attention", staticmethod(attention)) + return calls + + +def patch_flash(model, family): + module = local(family) + if family == "vit": + from transformers.models.vit.modeling_vit import ViTAttention + + for child in model.modules(): + if isinstance(child, ViTAttention): + child.forward = MethodType(module.get_vit_flash_self_attention_forward(), child) + else: + from transformers.models.whisper.modeling_whisper import WhisperAttention, WhisperDecoder + + for child in model.modules(): + if isinstance(child, WhisperAttention): + child.forward = MethodType(module.get_whisper_flash_attention_forward(), child) + if isinstance(child, WhisperDecoder): + child.forward = MethodType( + module.get_whisper_decoder_forward_for_flash_attention( + ShardConfig(enable_tensor_parallelism=False, enable_flash_attention=True) + ), + child, + ) + + +@pytest.mark.parametrize("family", ["vit", "whisper"]) +def test_clean_process_import(family): + result = subprocess.run( + [ + sys.executable, + "-c", + f"import colossalai.shardformer.modeling.{family}; import colossalai.shardformer.policies.{family}", + ], + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr + + +@pytest.mark.parametrize( + "family,variant", + [ + ("vit", "base"), + ("vit", "classification"), + ("vit", "masked"), + ("whisper", "base"), + ("whisper", "generation"), + ("whisper", "classification"), + ], +) +@pytest.mark.parametrize("tp", [False, True]) +def test_policy_target_paths(family, variant, tp, record_property): + reference, _ = pair(family, variant) + pm = importlib.import_module("colossalai.shardformer.policies." + family) + policy = getattr(pm, type(reference).__name__ + "Policy")() + config = ShardConfig(enable_tensor_parallelism=False, enable_flash_attention=True) + config.enable_tensor_parallelism = tp + config._tensor_parallel_size = 2 if tp else 1 + policy.set_model(reference) + policy.set_shard_config(config) + policy.preprocess() + descriptions = policy.module_policy() + checked = 0 + for cls, desc in descriptions.items(): + instances = [m for m in reference.modules() if isinstance(m, cls)] + # AudioClassification has only an encoder; the shared policy also + # describes decoder types used by the other two Whisper variants. + if ( + family == "whisper" + and variant == "classification" + and cls.__name__ in ("WhisperDecoder", "WhisperDecoderLayer") + ): + assert not instances + continue + assert instances, cls + for m in instances: + for replacement in desc.sub_module_replacement or []: + assert m.get_submodule(replacement.suffix) is not None + checked += 1 + for path in desc.attribute_replacement or {}: + parent, _, attr = path.rpartition(".") + assert hasattr(m.get_submodule(parent) if parent else m, attr), path + assert checked > 0 or (family == "vit" and not tp) + record_property("replacement_targets", checked) + if family == "whisper" and tp and variant != "classification": + from transformers.models.whisper.modeling_whisper import WhisperDecoderLayer + + assert descriptions[WhisperDecoderLayer].attribute_replacement["encoder_attn.num_heads"] == 1 + + +@pytest.mark.parametrize("variant", ["base", "classification", "masked"]) +@pytest.mark.parametrize("checkpoint", [False, True]) +def test_vit_pipeline_serial(variant, checkpoint, record_property): + reference, actual = pair("vit", variant) + module = local("vit") + inputs = data("vit", variant) + if checkpoint: + actual.train() + reference.train() + actual.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + vit = actual if variant == "base" else actual.vit + result = None + for i in range(2): + manager = stage(i) + indices = [i, i + 1] + vit.forward = MethodType(module.ViTModel_pipeline_forward(manager, indices), vit) + extra = {} if i == 0 else {"hidden_states": result["hidden_states"]} + if variant == "base": + result = actual(**inputs, **extra) + else: + forward = getattr(module, type(actual).__name__ + "_pipeline_forward")(manager, indices) + result = forward(actual, **inputs, **extra) + compare_training(reference, actual, reference(**inputs), result, record_property) + + +@pytest.mark.parametrize("variant", ["base", "generation", "classification"]) +@pytest.mark.parametrize("checkpoint", [False, True]) +def test_whisper_pipeline_serial(variant, checkpoint, record_property): + reference, actual = pair("whisper", variant) + forwards = local("whisper").WhisperPipelineForwards + inputs = data("whisper", variant) + if checkpoint: + actual.train() + reference.train() + actual.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True}) + function = { + "base": forwards.whisper_model_forward, + "generation": forwards.whisper_for_conditional_generation_forward, + "classification": forwards.whisper_for_audio_classification_forward, + }[variant] + total = 2 if variant == "classification" else 4 + result = {} + for i in range(total): + result = function( + actual, + **inputs, + **result, + stage_manager=stage(i, total), + stage_index=[i % 2, i % 2 + 1], + decoder_starting_stage=2, + shard_config=ShardConfig(enable_tensor_parallelism=False), + ) + if i != total - 1: + assert isinstance(result, dict) + compare_training(reference, actual, reference(**inputs), result, record_property) + + +@pytest.mark.parametrize( + "family,variant", + [ + ("vit", "base"), + ("vit", "classification"), + ("vit", "masked"), + ("whisper", "base"), + ("whisper", "generation"), + ("whisper", "classification"), + ], +) +@pytest.mark.parametrize("backend", ["eager", "sdpa"]) +def test_flash_adapter_training_cpu_kernel(family, variant, backend, cpu_flash, record_property): + reference, actual = pair(family, variant, backend) + patch_flash(actual, family) + inputs = data(family, variant) + compare_training(reference, actual, reference(**inputs), actual(**inputs), record_property) + assert len(cpu_flash) >= 2 + + +@pytest.mark.parametrize("backend", ["eager", "sdpa"]) +@pytest.mark.parametrize("padding", ["none", "left", "right"]) +@pytest.mark.parametrize("jit", [False, True]) +def test_whisper_cache_decode(backend, padding, jit, cpu_flash, record_property): + reference, actual = pair("whisper", "generation", backend) + patch_flash(actual, "whisper") + if jit: + from colossalai.shardformer.modeling.jit import get_jit_fused_dropout_add_func + + for layer in actual.model.decoder.layers: + layer.forward = MethodType(local("whisper").get_jit_fused_whisper_decoder_layer_forward(), layer) + layer.dropout_add = MethodType(get_jit_fused_dropout_add_func(), layer) + features = torch.randn(2, 8, 16) + ids = torch.tensor([[1, 4, 5, 6, 7, 8, 9], [1, 9, 8, 7, 6, 5, 4]]) + mask = torch.ones_like(ids) + if padding == "left": + mask[0, :2] = 0 + if padding == "right": + mask[0, 2:4] = 0 + caches = [None, None] + counts = [0, 0] + handles = [] + for i, layer in enumerate(actual.model.decoder.layers): + + def counter(m, a, o, idx=i): + counts[idx] += 1 + + handles.append(layer.encoder_attn.k_proj.register_forward_hook(counter)) + first_cross = [] + with torch.no_grad(): + encoders = [m.get_encoder()(features) for m in (reference, actual)] + for end in range(4, 8): + start = 0 if end == 4 else end - 1 + outputs = [] + for j, m in enumerate((reference, actual)): + out = m( + encoder_outputs=encoders[j], + decoder_input_ids=ids[:, start:end], + decoder_attention_mask=mask[:, :end], + past_key_values=caches[j], + use_cache=True, + ) + caches[j] = out.past_key_values + outputs.append(out) + full = reference( + encoder_outputs=encoders[0], + decoder_input_ids=ids[:, :end], + decoder_attention_mask=mask[:, :end], + use_cache=False, + ) + valid = mask[:, end - 1].bool() + close(outputs[1].logits[valid, -1], outputs[0].logits[valid, -1]) + close(outputs[1].logits[valid, -1], full.logits[valid, -1]) + for i in range(2): + assert caches[1].self_attention_cache.get_seq_length(i) == end + assert caches[1].cross_attention_cache.get_seq_length(i) == 8 + cross = caches[1].cross_attention_cache.layers[i].keys + if end == 4: + first_cross.append(cross.clone()) + else: + torch.testing.assert_close(cross, first_cross[i], atol=0, rtol=0) + assert counts == [1, 1] + for h in handles: + h.remove() + record_property("self_cache_lengths", "4,5,6,7") + record_property("cross_projection_calls", str(counts)) + + +@pytest.mark.parametrize("family", ["vit", "whisper"]) +def test_jit_layer_contract(family, record_property): + from colossalai.shardformer.modeling.jit import get_jit_fused_dropout_add_func + + reference, actual = pair(family, "classification" if family == "vit" else "generation") + module = local(family) + for m in actual.modules(): + if family == "vit" and type(m).__name__ == "ViTLayer": + factory = module.get_jit_fused_vit_layer_forward + elif family == "whisper" and type(m).__name__ == "WhisperEncoderLayer": + factory = module.get_jit_fused_whisper_encoder_layer_forward + elif family == "whisper" and type(m).__name__ == "WhisperDecoderLayer": + factory = module.get_jit_fused_whisper_decoder_layer_forward + else: + continue + m.forward = MethodType(factory(), m) + m.dropout_add = MethodType(get_jit_fused_dropout_add_func(), m) + inputs = data(family, "classification" if family == "vit" else "generation") + compare_training(reference, actual, reference(**inputs), actual(**inputs), record_property) + + +@pytest.mark.parametrize("backend", ["eager", "sdpa"]) +def test_whisper_mask_helper(backend): + reference, _ = pair("whisper", backend=backend) + module = local("whisper") + hidden = torch.randn(2, 4, 32) + config = ShardConfig(enable_tensor_parallelism=False, enable_flash_attention=True) + mask = module._get_attention_mask(reference.decoder, config, hidden, 0, torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]])) + assert mask.shape == (2, 1, 4, 4) + assert mask[0, 0, 3, 3] < -1e20 + assert mask[1, 0, 1, 2] < -1e20 + assert mask[1, 0, 3, 0] == 0 + assert reference.config._attn_implementation == backend + + +def test_decoder_cache_allocation(): + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + + from colossalai.shardformer._utils import prepare_hf_encoder_decoder_cache + + ref, _ = pair("whisper") + ref.config.encoder_layers = 4 + cache = prepare_hf_encoder_decoder_cache(ref.config) + assert isinstance(cache, EncoderDecoderCache) + assert len(cache.self_attention_cache.layers) == 2 + assert ref.config.encoder_layers == 4 + assert prepare_hf_encoder_decoder_cache(ref.config, cache) is cache + assert isinstance(prepare_hf_encoder_decoder_cache(ref.config, DynamicCache()), EncoderDecoderCache) + with pytest.raises(TypeError): + prepare_hf_encoder_decoder_cache(ref.config, ()) + + +@pytest.mark.parametrize( + "family,variant", + [ + ("vit", "base"), + ("vit", "classification"), + ("vit", "masked"), + ("whisper", "base"), + ("whisper", "generation"), + ("whisper", "classification"), + ], +) +def test_shardformer_optimize_flash_cpu_kernel(family, variant, cpu_flash, record_property): + from colossalai.shardformer import ShardFormer + + reference, actual = pair(family, variant) + sharder = ShardFormer( + ShardConfig(enable_tensor_parallelism=False, enable_flash_attention=True, make_vocab_size_divisible_by=1) + ) + actual, _ = sharder.optimize(actual) + inputs = data(family, variant) + compare_training(reference, actual, reference(**inputs), actual(**inputs), record_property) + assert len(cpu_flash) >= 2 + + +@pytest.mark.parametrize("padding", ["left", "right"]) +def test_whisper_decoder_checkpoint_and_capture(padding, cpu_flash, record_property): + reference, actual = pair("whisper", "generation") + patch_flash(actual, "whisper") + reference.train() + actual.train() + actual.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": True}) + inputs = data("whisper", "generation") + mask = torch.ones(2, 4, dtype=torch.long) + mask[0, 0 if padding == "left" else -1] = 0 + inputs["decoder_attention_mask"] = mask + inputs["labels"] = inputs["labels"].masked_fill(mask == 0, -100) + # Explicitly requesting cache while checkpointing must not mutate a cache. + out = actual(**{**inputs, "use_cache": True}, output_hidden_states=True) + ref = reference(**inputs, output_hidden_states=True) + assert out.past_key_values is None + assert len(out.decoder_hidden_states) == len(ref.decoder_hidden_states) == 3 + compare_training(reference, actual, ref, out, record_property, logits_mask=mask.bool()) + + +def test_whisper_cross_mask_attention(cpu_flash): + ref, _ = pair("whisper") + reference = ref.decoder.layers[0].encoder_attn + actual = copy.deepcopy(reference) + actual.forward = MethodType(local("whisper").get_whisper_flash_attention_forward(), actual) + hidden = torch.randn(2, 3, 32) + audio = torch.randn(2, 8, 32) + mask = torch.zeros(2, 1, 3, 8) + mask[0, :, :, 5:] = torch.finfo(torch.float32).min + expected = reference(hidden, key_value_states=audio, attention_mask=mask)[0] + out = actual(hidden, key_value_states=audio, attention_mask=mask)[0] + close(out, expected) + audio[0, 5:] = 10000 + close(actual(hidden, key_value_states=audio, attention_mask=mask)[0], out) + + +@pytest.mark.parametrize("backend", ["eager", "sdpa"]) +def test_vit_padding_mask_and_tuple(backend, cpu_flash): + ref, actual = pair("vit", backend=backend) + patch_flash(actual, "vit") + pixels = torch.randn(2, 3, 16, 16) + mask = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]) + expected = ref(pixels, attention_mask=mask, return_dict=False) + out = actual(pixels, attention_mask=mask, return_dict=False) + assert isinstance(out, tuple) and len(out) == len(expected) + for a, b in zip(out, expected): + close(a, b) + module = local("vit") + state = module.ViTModel_pipeline_forward(stage(0), [0, 1])(actual, pixels, attention_mask=mask) + result = module.ViTModel_pipeline_forward(stage(1), [1, 2])( + actual, hidden_states=state["hidden_states"], attention_mask=mask, return_dict=False + ) + for a, b in zip(result, expected): + close(a, b) + + +def test_whisper_audio_weighted_pipeline(record_property): + reference, actual = pair("whisper", "classification") + # Construct with the feature enabled so HF creates trainable layer_weights. + config = copy.deepcopy(reference.config) + config.use_weighted_layer_sum = True + reference = hf.WhisperForAudioClassification(config).eval() + actual = copy.deepcopy(reference) + inputs = data("whisper", "classification") + fn = local("whisper").WhisperPipelineForwards.whisper_for_audio_classification_forward + state = fn(actual, **inputs, stage_manager=stage(0), stage_index=[0, 1], decoder_starting_stage=2) + assert len(state["encoder_states"]) == 1 + out = fn(actual, **inputs, **state, stage_manager=stage(1), stage_index=[1, 2], decoder_starting_stage=2) + compare_training(reference, actual, reference(**inputs), out, record_property) + + +@pytest.mark.parametrize("family", ["vit", "whisper"]) +def test_pipeline_held_layers(family): + from colossalai.pipeline.stage_manager import PipelineStageManager + + reference, _ = pair(family) + pm = importlib.import_module("colossalai.shardformer.policies." + family) + total = 2 if family == "vit" else 4 + held = [] + for i in range(total): + manager = stage(i, total) + manager.num_model_chunks = 1 + manager.num_layers_per_stage = None + manager.distribute_layers = MethodType(PipelineStageManager.distribute_layers, manager) + manager.get_stage_index = MethodType(PipelineStageManager.get_stage_index, manager) + config = ShardConfig(enable_tensor_parallelism=False, pipeline_stage_manager=manager) + policy = getattr(pm, type(reference).__name__ + "Policy")() + policy.set_model(reference) + policy.set_shard_config(config) + policy.preprocess() + policy.module_policy() + held.extend(policy.get_held_layers()) + layers = ( + list(reference.layers) if family == "vit" else list(reference.encoder.layers) + list(reference.decoder.layers) + ) + for layer in layers: + assert sum(x is layer for x in held) == 1 + + +def test_whisper_jit_stays_disabled(): + pm = importlib.import_module("colossalai.shardformer.policies.whisper") + ref, _ = pair("whisper") + policy = pm.WhisperModelPolicy() + config = ShardConfig(enable_tensor_parallelism=False, enable_jit_fused=True) + policy.set_model(ref) + policy.set_shard_config(config) + policy.preprocess() + with pytest.warns(UserWarning, match="doesn't support jit"): + policy.module_policy() + assert not config.enable_jit_fused + + +@pytest.mark.parametrize("variant", ["base", "generation"]) +def test_whisper_pipeline_tuple_and_default_cache(variant): + reference, actual = pair("whisper", variant) + actual.config.use_cache = True + inputs = data("whisper", variant) + inputs.pop("use_cache") + fn = getattr( + local("whisper").WhisperPipelineForwards, + "whisper_model_forward" if variant == "base" else "whisper_for_conditional_generation_forward", + ) + state = {} + for i in range(2): + state = fn( + actual, + **inputs, + **state, + return_dict=False, + stage_manager=stage(i), + stage_index=[0, 2], + decoder_starting_stage=1, + shard_config=ShardConfig(enable_tensor_parallelism=False), + ) + expected = reference(**inputs, use_cache=False, return_dict=False) + assert isinstance(state, tuple) and len(state) == len(expected) + for a, b in zip(state, expected): + close(a, b) + + +def test_vit_jit_policy_keeps_exact_activation(record_property): + from colossalai.shardformer import ShardFormer + + reference, actual = pair("vit", "classification") + # Large enough activations to expose replacing erf GELU by tanh GELU. + with torch.no_grad(): + for layer in reference.vit.layers: + layer.mlp.fc1.weight.mul_(10) + actual.load_state_dict(reference.state_dict()) + actual, _ = ShardFormer(ShardConfig(enable_tensor_parallelism=False, enable_jit_fused=True)).optimize(actual) + assert all(hasattr(layer, "dropout_add") for layer in actual.vit.layers) + inputs = data("vit", "classification") + compare_training(reference, actual, reference(**inputs), actual(**inputs), record_property) + + +def test_vit_interleaved_stage_functions_serially(record_property): + reference, _ = pair("vit") + config = copy.deepcopy(reference.config) + config.num_hidden_layers = 4 + reference = hf.ViTModel(config).eval() + actual = copy.deepcopy(reference) + pixels = torch.randn(2, 3, 16, 16) + state = {} + for logical in range(4): + physical = logical % 2 + chunk = logical // 2 + manager = stage(physical) + manager.is_interleave = True + manager.model_chunk_id = chunk + manager.is_first_stage = lambda **kw: logical == 0 + manager.is_last_stage = lambda **kw: logical == 3 + indices = [[physical, physical + 1], [physical + 2, physical + 3]] + state = local("vit").ViTModel_pipeline_forward(manager, indices)(actual, pixel_values=pixels, **state) + compare_training(reference, actual, reference(pixels), state, record_property) + + +@pytest.mark.parametrize("family", ["vit", "whisper"]) +def test_flash_training_dropout_argument(family, cpu_flash): + reference, actual = pair(family) + patch_flash(actual, family) + attn = actual.layers[0].attention if family == "vit" else actual.encoder.layers[0].self_attn + if family == "vit": + attn.attention_dropout = 0.2 + else: + attn.dropout = 0.2 + hidden = torch.randn(2, 4, 32) + attn.train() + attn(hidden) + assert cpu_flash[-1][2] == 0.2 + attn.eval() + attn(hidden) + assert cpu_flash[-1][2] == 0.0 diff --git a/tests/test_shardformer/test_layer/test_linear_1d.py b/tests/test_shardformer/test_layer/test_linear_1d.py index 773799c1cc09..2f1352473283 100644 --- a/tests/test_shardformer/test_layer/test_linear_1d.py +++ b/tests/test_shardformer/test_layer/test_linear_1d.py @@ -264,7 +264,7 @@ def check_linear_col_plus_row(lazy_init: bool, seq_parallel_mode: bool, overlap: @parameterize("lazy_init", [False, True]) -@parameterize("seq_parallel_mode", [None, "split_gather"]) +@parameterize("seq_parallel_mode", [None, "split_gather", "ring"]) @parameterize("overlap", [True]) def run_dist_linear_test(lazy_init, seq_parallel_mode, overlap): check_linear_1d_col(lazy_init, seq_parallel_mode, overlap) diff --git a/tests/test_shardformer/test_layer/test_ring_as_gather.py b/tests/test_shardformer/test_layer/test_ring_as_gather.py new file mode 100644 index 000000000000..0bd2372a2df3 --- /dev/null +++ b/tests/test_shardformer/test_layer/test_ring_as_gather.py @@ -0,0 +1,35 @@ +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from colossalai.shardformer.layer._operation import _ring_as_gather +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +def check_ring_as_gather(rank, world_size, port): + # gloo on CPU: the ring logic is independent of the device. + dist.init_process_group("gloo", init_method=f"tcp://localhost:{port}", rank=rank, world_size=world_size) + torch.manual_seed(0) + full = torch.randn(2, 4 * world_size, 8) + weight = torch.randn(6, 8) + local = full.chunk(world_size, dim=1)[rank].contiguous() + + output, gathered = _ring_as_gather(F.linear, {"input": local}, {"weight": weight}, process_group=dist.group.WORLD) + + assert torch.equal(output, F.linear(full, weight)) + # The backward of the ring linear computes the weight gradient from the gathered input, so the chunk + # used in every round must be recorded. 2, 3 and 4 ranks exercise the last round, the per-round + # record and the reuse of the send / recv buffers (a later irecv overwrites a recorded reference). + assert torch.equal(gathered["input"], full) + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", [2, 3, 4]) +@rerun_if_address_is_in_use() +def test_ring_as_gather(world_size): + spawn(check_ring_as_gather, world_size) + + +if __name__ == "__main__": + test_ring_as_gather(4) diff --git a/tests/test_shardformer/test_layer/test_ring_gather.py b/tests/test_shardformer/test_layer/test_ring_gather.py new file mode 100644 index 000000000000..91fc5179428b --- /dev/null +++ b/tests/test_shardformer/test_layer/test_ring_gather.py @@ -0,0 +1,42 @@ +"""CPU regression for the input snapshots retained by ring sequence parallelism.""" + +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +import torch.nn.functional as F + +from colossalai.shardformer.layer._operation import _ring_as_gather + + +def _check_ring_gather(rank, world_size, rendezvous): + torch.set_num_threads(1) + dist.init_process_group( + "gloo", init_method=rendezvous, rank=rank, world_size=world_size, timeout=timedelta(seconds=60) + ) + try: + # Distinct shards expose incorrect rank ordering and reused receive buffers. + chunks = [torch.arange(12, dtype=torch.float64).reshape(1, 3, 4) + 100 * r for r in range(world_size)] + local = chunks[rank] + # Dyadic values make this a bit-exact communication test, independent + # of the GEMM implementation's floating-point reduction order. + weight = torch.arange(20, dtype=torch.float64).reshape(5, 4) / 8 + output, saved = _ring_as_gather(F.linear, {"input": local}, {"weight": weight}, dist.group.WORLD) + expected_input = torch.cat(chunks, dim=1) + torch.testing.assert_close(output, F.linear(expected_input, weight), rtol=0, atol=0) + torch.testing.assert_close(saved["input"], expected_input, rtol=0, atol=0) + # Backward's weight gradient must use the same inputs as the forward. + grad = torch.arange(output.numel(), dtype=torch.float64).reshape_as(output) + actual_grad = grad.flatten(0, 1).T @ saved["input"].flatten(0, 1) + reference_weight = weight.clone().requires_grad_() + F.linear(expected_input, reference_weight).backward(grad) + torch.testing.assert_close(actual_grad, reference_weight.grad, rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", [2, 3, 4]) +def test_ring_gather_saved_inputs(world_size, tmp_path): + mp.spawn(_check_ring_gather, args=(world_size, (tmp_path / "rendezvous").as_uri()), nprocs=world_size) diff --git a/tests/test_shardformer/test_model/test_mixtral_checkpoint_loading.py b/tests/test_shardformer/test_model/test_mixtral_checkpoint_loading.py new file mode 100644 index 000000000000..b825230bdea1 --- /dev/null +++ b/tests/test_shardformer/test_model/test_mixtral_checkpoint_loading.py @@ -0,0 +1,255 @@ +from collections import OrderedDict +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +from torch import nn + +from colossalai._compat import is_transformers_v5 +from colossalai.checkpoint_io.index_file import CheckpointIndexFile +from colossalai.checkpoint_io.moe_checkpoint import MoECheckpointIO +from colossalai.interface.model import ModelWrapper +from colossalai.tensor.moe_tensor.api import ( + is_moe_tensor, + is_moe_tensor_ep_fused, + set_moe_tensor_ep_fused, + set_moe_tensor_ep_group, +) +from colossalai.testing.utils import spawn + +if not is_transformers_v5(): + pytest.skip("fused Mixtral experts only exist in transformers v5", allow_module_level=True) + +from colossalai.shardformer.modeling.mixtral import _mark_fused_expert_tp_shard, _register_fused_expert_checkpoint_hook + + +class _FusedExperts(nn.Module): + def __init__(self, num_experts: int, gate_up_width: int = 6, down_input_width: int = 3): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.empty(num_experts, gate_up_width, 4)) + self.down_proj = nn.Parameter(torch.empty(num_experts, 4, down_input_width)) + + +def test_load_full_fused_expert_checkpoint_for_local_ep_shard(): + full_gate_up = torch.arange(8 * 6 * 4, dtype=torch.float32).reshape(8, 6, 4) + full_down = torch.arange(8 * 4 * 3, dtype=torch.float32).reshape(8, 4, 3) + experts = _FusedExperts(num_experts=2) + _register_fused_expert_checkpoint_hook(experts, expert_start_idx=2, num_experts_per_ep=2, num_experts=8) + + state_dict = OrderedDict(gate_up_proj=full_gate_up, down_proj=full_down) + experts.load_state_dict(state_dict) + + assert torch.equal(experts.gate_up_proj, full_gate_up[2:4]) + assert torch.equal(experts.down_proj, full_down[2:4]) + + +def test_load_local_fused_expert_checkpoint_without_reslicing(): + local_gate_up = torch.randn(2, 6, 4) + local_down = torch.randn(2, 4, 3) + experts = _FusedExperts(num_experts=2) + _register_fused_expert_checkpoint_hook(experts, expert_start_idx=2, num_experts_per_ep=2, num_experts=8) + + experts.load_state_dict(OrderedDict(gate_up_proj=local_gate_up, down_proj=local_down)) + + assert torch.equal(experts.gate_up_proj, local_gate_up) + assert torch.equal(experts.down_proj, local_down) + + +def test_load_full_fused_expert_checkpoint_with_tp_shards(): + full_gate_up = torch.arange(8 * 8 * 4, dtype=torch.float32).reshape(8, 8, 4) + full_down = torch.arange(8 * 4 * 6, dtype=torch.float32).reshape(8, 4, 6) + experts = _FusedExperts(num_experts=2, gate_up_width=4, down_input_width=3) + _register_fused_expert_checkpoint_hook( + experts, expert_start_idx=2, num_experts_per_ep=2, num_experts=8, tp_rank=1, tp_size=2 + ) + + experts.load_state_dict(OrderedDict(gate_up_proj=full_gate_up, down_proj=full_down)) + + expected_gate_up = torch.cat((full_gate_up[2:4, 2:4], full_gate_up[2:4, 6:8]), dim=1) + assert torch.equal(experts.gate_up_proj, expected_gate_up) + assert torch.equal(experts.down_proj, full_down[2:4, :, 3:6]) + + +def _check_tp_expert_checkpoint_gather(rank: int, world_size: int, port: int): + dist.init_process_group( + backend="gloo", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + try: + ep_groups = [dist.new_group(ranks=[0]), dist.new_group(ranks=[1])] + full_gate_up = torch.arange(8 * 8 * 4, dtype=torch.float32).reshape(8, 8, 4) + full_down = torch.arange(8 * 4 * 6, dtype=torch.float32).reshape(8, 4, 6) + local_gate_up = torch.cat( + (full_gate_up[2:4, rank * 2 : (rank + 1) * 2], full_gate_up[2:4, 4 + rank * 2 : 6 + rank * 2]), + dim=1, + ) + local_down = full_down[2:4, :, rank * 3 : (rank + 1) * 3] + experts = _FusedExperts(num_experts=2, gate_up_width=4, down_input_width=3) + experts.gate_up_proj = nn.Parameter(local_gate_up.clone()) + experts.down_proj = nn.Parameter(local_down.clone()) + _mark_fused_expert_tp_shard(experts.gate_up_proj, "gate_up_proj", dist.group.WORLD, 2, 2, 8) + _mark_fused_expert_tp_shard(experts.down_proj, "down_proj", dist.group.WORLD, 2, 2, 8) + set_moe_tensor_ep_group(experts.gate_up_proj, ep_groups[rank]) + set_moe_tensor_ep_group(experts.down_proj, ep_groups[rank]) + assert is_moe_tensor(experts.gate_up_proj) and is_moe_tensor(experts.down_proj) + + saved_state = {} + for shard, _ in MoECheckpointIO._model_sharder(experts): + saved_state.update(shard) + assert torch.equal(saved_state["gate_up_proj"], full_gate_up[2:4]) + assert torch.equal(saved_state["down_proj"], full_down[2:4]) + assert torch.equal(experts.gate_up_proj.shard_fn(full_gate_up), local_gate_up) + assert torch.equal(experts.down_proj.shard_fn(full_down), local_down) + + reloaded = _FusedExperts(num_experts=2, gate_up_width=4, down_input_width=3) + _register_fused_expert_checkpoint_hook( + reloaded, expert_start_idx=2, num_experts_per_ep=2, num_experts=8, tp_rank=rank, tp_size=world_size + ) + reloaded.load_state_dict(saved_state) + assert torch.equal(reloaded.gate_up_proj, local_gate_up) + assert torch.equal(reloaded.down_proj, local_down) + finally: + dist.destroy_process_group() + + +def test_gather_and_load_tp_sharded_fused_expert_checkpoint(): + spawn(_check_tp_expert_checkpoint_gather, nprocs=2) + + +def _check_tp_ep_expert_checkpoint_gather(rank: int, world_size: int, port: int): + assert world_size == 4 + dist.init_process_group( + backend="gloo", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + try: + # Use a 2x2 EP/TP mesh: ranks in each TP group have different EP groups. + tp_groups = [dist.new_group(ranks=[0, 1]), dist.new_group(ranks=[2, 3])] + ep_groups = [dist.new_group(ranks=[0, 2]), dist.new_group(ranks=[1, 3])] + tp_rank = rank % 2 + ep_rank = rank // 2 + tp_group = tp_groups[ep_rank] + ep_group = ep_groups[tp_rank] + expert_start_idx = ep_rank * 4 + + full_gate_up = torch.arange(8 * 8 * 4, dtype=torch.float32).reshape(8, 8, 4) + full_down = torch.arange(8 * 4 * 6, dtype=torch.float32).reshape(8, 4, 6) + ep_gate_up = full_gate_up[expert_start_idx : expert_start_idx + 4] + ep_down = full_down[expert_start_idx : expert_start_idx + 4] + local_gate_up = torch.cat( + ( + ep_gate_up[:, tp_rank * 2 : (tp_rank + 1) * 2], + ep_gate_up[:, 4 + tp_rank * 2 : 6 + tp_rank * 2], + ), + dim=1, + ) + local_down = ep_down[:, :, tp_rank * 3 : (tp_rank + 1) * 3] + + experts = _FusedExperts(num_experts=4, gate_up_width=4, down_input_width=3) + experts.gate_up_proj = nn.Parameter(local_gate_up.clone()) + experts.down_proj = nn.Parameter(local_down.clone()) + _mark_fused_expert_tp_shard(experts.gate_up_proj, "gate_up_proj", tp_group, expert_start_idx, 4, 8) + _mark_fused_expert_tp_shard(experts.down_proj, "down_proj", tp_group, expert_start_idx, 4, 8) + set_moe_tensor_ep_group(experts.gate_up_proj, ep_group) + set_moe_tensor_ep_group(experts.down_proj, ep_group) + set_moe_tensor_ep_fused(experts.gate_up_proj) + set_moe_tensor_ep_fused(experts.down_proj) + assert is_moe_tensor(experts.gate_up_proj) and is_moe_tensor(experts.down_proj) + assert is_moe_tensor_ep_fused(experts.gate_up_proj) and is_moe_tensor_ep_fused(experts.down_proj) + assert dist.get_process_group_ranks(experts.gate_up_proj.ep_group) == dist.get_process_group_ranks(ep_group) + + saved_state = {} + for shard, _ in MoECheckpointIO._model_sharder(experts): + saved_state.update(shard) + if ep_rank == 0: + assert torch.equal(saved_state["gate_up_proj"], full_gate_up) + assert torch.equal(saved_state["down_proj"], full_down) + else: + assert torch.equal(saved_state["gate_up_proj"], ep_gate_up) + assert torch.equal(saved_state["down_proj"], ep_down) + + # The shared index points to the EP-root file, whose expert tensors are full. + checkpoint_state = { + "gate_up_proj": saved_state["gate_up_proj"] if rank == 0 else torch.empty_like(full_gate_up), + "down_proj": saved_state["down_proj"] if rank == 0 else torch.empty_like(full_down), + } + dist.broadcast(checkpoint_state["gate_up_proj"], src=0) + dist.broadcast(checkpoint_state["down_proj"], src=0) + + reloaded = _FusedExperts(num_experts=4, gate_up_width=4, down_input_width=3) + _register_fused_expert_checkpoint_hook( + reloaded, + expert_start_idx=expert_start_idx, + num_experts_per_ep=4, + num_experts=8, + tp_rank=tp_rank, + tp_size=2, + ) + reloaded.load_state_dict(checkpoint_state) + assert torch.equal(reloaded.gate_up_proj, local_gate_up) + assert torch.equal(reloaded.down_proj, local_down) + finally: + dist.destroy_process_group() + + +def test_gather_and_load_tp_ep_sharded_fused_expert_checkpoint(): + spawn(_check_tp_ep_expert_checkpoint_gather, nprocs=4) + + +def test_async_sharded_checkpoint_uses_moe_shards(monkeypatch, tmp_path): + from colossalai.checkpoint_io import moe_checkpoint as moe_checkpoint_module + from colossalai.checkpoint_io.utils import get_model_base_filenames + + model = ModelWrapper(nn.Linear(2, 2)) + checkpoint_io = object.__new__(MoECheckpointIO) + checkpoint_io.moe_dp_rank = 0 + checkpoint_io.ep_rank = 0 + checkpoint_io.tp_rank = 0 + checkpoint_io.sp_rank = 0 + checkpoint_io.ep_size = 1 + checkpoint_io.pp_size = 1 + checkpoint_io.verbose = False + checkpoint_io.async_writers = [] + checkpoint_io.pinned_state_dicts = {} + + class FakeAsyncWriter: + def sync_before_step(self): + pass + + def synchronize(self): + pass + + writer = FakeAsyncWriter() + observed = {} + + def fake_async_save_state_dict_shards(sharded_state_dict, checkpoint, index_file, base_filename, is_master): + shards = list(sharded_state_dict) + observed.update(checkpoint=checkpoint, base_filename=base_filename, is_master=is_master, shards=shards) + for idx, (shard, _) in enumerate(shards): + for name in shard: + index_file.append_weight_map(name, f"shard-{idx}.safetensors") + return sum(size for _, size in shards), [writer] + + monkeypatch.setattr(moe_checkpoint_module, "async_save_state_dict_shards", fake_async_save_state_dict_shards) + monkeypatch.setattr(moe_checkpoint_module.dist, "barrier", lambda: None) + monkeypatch.setattr(moe_checkpoint_module, "save_config_file", lambda *args, **kwargs: None) + + checkpoint = tmp_path / "checkpoint" + checkpoint_io.save_sharded_model(model, str(checkpoint), use_safetensors=True, use_async=True) + + _, index_filename = get_model_base_filenames(use_safetensors=True) + assert observed["checkpoint"] == str(checkpoint) + assert observed["base_filename"].endswith(".safetensors") + assert observed["is_master"] is True + assert len(observed["shards"]) == 1 + assert checkpoint_io.async_writers == [writer] + index_file = checkpoint / index_filename + assert index_file.is_file() + assert set(CheckpointIndexFile.from_file(index_file).weight_map) == {"weight", "bias"} diff --git a/tests/test_shardformer/test_model/test_mixtral_sp_attention_mask.py b/tests/test_shardformer/test_model/test_mixtral_sp_attention_mask.py new file mode 100644 index 000000000000..8a7ce10c580c --- /dev/null +++ b/tests/test_shardformer/test_model/test_mixtral_sp_attention_mask.py @@ -0,0 +1,137 @@ +from copy import deepcopy +from typing import Tuple + +import pytest +import torch +import torch.distributed as dist +from transformers.models.mixtral.configuration_mixtral import MixtralConfig +from transformers.models.mixtral.modeling_mixtral import MixtralModel + +import colossalai +from colossalai._compat import is_transformers_v5 +from colossalai.booster.booster import Booster +from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin +from colossalai.shardformer.layer.utils import Randomizer +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.testing.random import seed_all + +# head_dim is set explicitly, so num_attention_heads * head_dim (4 * 8) differs from hidden_size. +HIDDEN_SIZE, NUM_HEADS, HEAD_DIM = 24, 4, 8 +BATCH_SIZE, SEQ_LEN, NUM_PADDED = 2, 16, 5 +# Relative L2 error allowed between the gathered parallel output and the unsharded one, both in bf16 (unit roundoff +# 3.9e-3). A wrong mask or reshape changes whole rows of the output, i.e. errors of order 1. +OUTPUT_RTOL = 2e-2 + + +@parameterize( + "config", + [ + # (sp_mode, tp_size, sp_size) on 4 ranks; split_gather / ring are not supported yet + ("all_to_all", 1, 2), + ], +) +def run_sp_with_attention_mask(config: Tuple[str, int, int]): + """Mixtral with sequence parallelism must accept a padding attention mask, like the unsharded model. + + The SP attention / model forwards replaced by the policy only run with sequence parallelism, and the existing + Mixtral tests pass no attention mask. + """ + sp_mode, tp_size, sp_size = config + Randomizer.reset_index() + plugin = MoeHybridParallelPlugin( + pp_size=1, + tp_size=tp_size, + sp_size=sp_size, + ep_size=1, + # zero 0 would wrap the model in DDP, which the MoE plugin rejects when sp splits the dp group; ZeRO in turn + # needs fp16 / bf16 + zero_stage=1, + enable_sequence_parallelism=True, + sequence_parallelism_mode=sp_mode, + precision="bf16", + overlap_communication=False, + find_unused_parameters=True, + ) + booster = Booster(plugin=plugin) + + seed_all(10086) + model_config = MixtralConfig( + hidden_size=HIDDEN_SIZE, + intermediate_size=HIDDEN_SIZE * 2, + num_hidden_layers=2, + num_attention_heads=NUM_HEADS, + num_key_value_heads=NUM_HEADS, + head_dim=HEAD_DIM, + num_local_experts=4, + num_experts_per_tok=2, + attn_implementation="eager", + ) + torch_model = MixtralModel(model_config).to(torch.bfloat16).cuda() + parallel_model = deepcopy(torch_model) + optimizer = torch.optim.SGD(parallel_model.parameters(), lr=0.0) + parallel_model, optimizer, _, _, _ = booster.boost(parallel_model, optimizer) + + # every rank uses the same batch, so the gathered parallel output equals the unsharded one + seed_all(1453) + inputs_embeds = torch.rand(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) + attention_mask = torch.ones(BATCH_SIZE, SEQ_LEN, dtype=torch.long, device="cuda") + attention_mask[1, :NUM_PADDED] = 0 # left padding + + parallel_output = parallel_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask).last_hidden_state + optimizer.backward(parallel_output.mean()) + with torch.no_grad(): + torch_output = torch_model(inputs_embeds=inputs_embeds, attention_mask=attention_mask).last_hidden_state + + # sequence parallelism may leave each rank with its own part of the sequence + if parallel_output.shape[1] != SEQ_LEN: + parts = [torch.empty_like(parallel_output) for _ in range(dist.get_world_size(plugin.sp_group))] + dist.all_gather(parts, parallel_output.detach().contiguous(), group=plugin.sp_group) + parallel_output = torch.cat(parts, dim=1) + error = ((parallel_output.detach().float() - torch_output.float()).norm() / torch_output.float().norm()).item() + if dist.get_rank() == 0: + print(f"{sp_mode} tp={tp_size} sp={sp_size}: relative output error {error:.3e}", flush=True) + assert error <= OUTPUT_RTOL, f"{sp_mode} tp={tp_size} sp={sp_size}: relative output error {error:.3e}" + + +def run_unsupported_sp_mode(): + """split_gather / ring must be rejected when sharding, not hang in the MoE layer.""" + plugin = MoeHybridParallelPlugin( + pp_size=1, + tp_size=2, + ep_size=1, + zero_stage=1, + enable_sequence_parallelism=True, + sequence_parallelism_mode="split_gather", + precision="bf16", + ) + model_config = MixtralConfig( + hidden_size=HIDDEN_SIZE, + intermediate_size=HIDDEN_SIZE * 2, + num_hidden_layers=1, + num_attention_heads=NUM_HEADS, + num_key_value_heads=NUM_HEADS, + head_dim=HEAD_DIM, + num_local_experts=4, + num_experts_per_tok=2, + ) + model = MixtralModel(model_config).to(torch.bfloat16).cuda() + with pytest.raises(NotImplementedError, match="split_gather"): + Booster(plugin=plugin).boost(model, torch.optim.SGD(model.parameters(), lr=0.0)) + + +def run_dist(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + torch.cuda.set_device(rank) + run_sp_with_attention_mask() + if is_transformers_v5(): + run_unsupported_sp_mode() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_mixtral_sp_attention_mask(): + spawn(run_dist, 4) + + +if __name__ == "__main__": + test_mixtral_sp_attention_mask() diff --git a/tests/test_shardformer/test_model/test_moe_expert_param_identity.py b/tests/test_shardformer/test_model/test_moe_expert_param_identity.py new file mode 100644 index 000000000000..84bf16f224e2 --- /dev/null +++ b/tests/test_shardformer/test_model/test_moe_expert_param_identity.py @@ -0,0 +1,126 @@ +import socket +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("fused MoE experts only exist in transformers v5", allow_module_level=True) + +from transformers import DeepseekV3Config, MixtralConfig +from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3MoE +from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock + +from colossalai.shardformer.modeling.deepseek_v3 import EpDeepseekV3MoE +from colossalai.shardformer.modeling.mixtral import EPMixtralSparseMoeBlock +from colossalai.testing import rerun_if_address_is_in_use, spawn + + +@pytest.fixture +def single_rank_group(): + # Optimizers are created from the unsharded model and the MoE plugin only prunes optimizer params + # when PP/EP > 1, so the EP/TP expert slicing must keep the original Parameter objects. + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=0, world_size=1) + yield dist.group.WORLD + dist.destroy_process_group() + + +def _expert_params(module): + return module.experts.gate_up_proj, module.experts.down_proj + + +def test_mixtral_ep_keeps_expert_parameter_objects(single_rank_group): + config = MixtralConfig(hidden_size=32, intermediate_size=64, num_local_experts=4, num_experts_per_tok=2) + module = MixtralSparseMoeBlock(config) + before = _expert_params(module) + optimizer = torch.optim.SGD(module.parameters(), lr=1.0) + + module = EPMixtralSparseMoeBlock.from_native_module( + module, tp_group=single_rank_group, moe_dp_group=single_rank_group, ep_group=single_rank_group + ) + + after = _expert_params(module) + assert all(a is b for a, b in zip(after, before)) + optimizer_params = {id(p) for group in optimizer.param_groups for p in group["params"]} + assert all(id(p) in optimizer_params for p in after) + + +def _check_mixtral_tp_keeps_expert_parameter_objects(rank: int, world_size: int, port: int): + dist.init_process_group( + backend="gloo", + init_method=f"tcp://localhost:{port}", + rank=rank, + world_size=world_size, + timeout=timedelta(seconds=30), + ) + try: + own_group = [dist.new_group(ranks=[r]) for r in range(world_size)][rank] + torch.manual_seed(0) + config = MixtralConfig(hidden_size=32, intermediate_size=64, num_local_experts=4, num_experts_per_tok=2) + module = MixtralSparseMoeBlock(config) + # v5 allocates fused experts with torch.empty (weights are initialised by the full model), so + # fill them with known values before comparing the local shards. + with torch.no_grad(): + for param in _expert_params(module): + param.copy_(torch.arange(param.numel(), dtype=param.dtype).reshape(param.shape)) + full_gate_up = module.experts.gate_up_proj.detach().clone() + full_down = module.experts.down_proj.detach().clone() + before = _expert_params(module) + optimizer = torch.optim.SGD(module.parameters(), lr=1.0) + + # TP over the expert intermediate dimension; EP = 1 on each rank. + module = EPMixtralSparseMoeBlock.from_native_module( + module, tp_group=dist.group.WORLD, moe_dp_group=own_group, ep_group=own_group + ) + + after = _expert_params(module) + assert all(a is b for a, b in zip(after, before)) + optimizer_params = {id(p) for group in optimizer.param_groups for p in group["params"]} + assert all(id(p) in optimizer_params for p in after) + # The in-place storage swap must still hold this rank's gate/up/down shards. + part = 64 // world_size + expected_gate_up = torch.cat( + ( + full_gate_up[:, rank * part : (rank + 1) * part], + full_gate_up[:, 64 + rank * part : 64 + (rank + 1) * part], + ), + dim=1, + ) + assert torch.equal(module.experts.gate_up_proj.detach(), expected_gate_up) + assert torch.equal(module.experts.down_proj.detach(), full_down[:, :, rank * part : (rank + 1) * part]) + finally: + dist.destroy_process_group() + + +@rerun_if_address_is_in_use() +def test_mixtral_tp_keeps_expert_parameter_objects(): + spawn(_check_mixtral_tp_keeps_expert_parameter_objects, nprocs=2) + + +def test_deepseek_v3_ep_keeps_expert_parameter_objects(single_rank_group): + config = DeepseekV3Config( + hidden_size=64, + intermediate_size=96, + moe_intermediate_size=32, + n_routed_experts=8, + n_shared_experts=1, + num_experts_per_tok=2, + n_group=2, + topk_group=1, + ) + module = DeepseekV3MoE(config) + before = _expert_params(module) + optimizer = torch.optim.SGD(module.parameters(), lr=1.0) + + module = EpDeepseekV3MoE.from_native_module(module, moe_dp_group=single_rank_group, ep_group=single_rank_group) + + after = _expert_params(module) + assert all(a is b for a, b in zip(after, before)) + optimizer_params = {id(p) for group in optimizer.param_groups for p in group["params"]} + assert all(id(p) in optimizer_params for p in after) diff --git a/tests/test_shardformer/test_model/test_moe_optimizer_checkpoint.py b/tests/test_shardformer/test_model/test_moe_optimizer_checkpoint.py new file mode 100644 index 000000000000..43733b15c1a0 --- /dev/null +++ b/tests/test_shardformer/test_model/test_moe_optimizer_checkpoint.py @@ -0,0 +1,125 @@ +import os +import tempfile +from contextlib import nullcontext +from typing import Tuple + +import pytest +import torch +import torch.distributed as dist + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("fused MoE experts only exist in transformers v5", allow_module_level=True) + +import colossalai +from colossalai.booster.booster import Booster +from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.testing.random import seed_all +from tests.test_shardformer.test_model.test_moe_optimizer_coverage import NUM_TOKENS, _build + + +def _state_snapshot(optimizer): + """Tensor optimizer states keyed by working-parameter name order (ZeRO keys state by master params).""" + inner = optimizer.unwrap() if hasattr(optimizer, "unwrap") else optimizer.optim + return [ + {k: v.detach().clone() for k, v in inner.state[p].items() if isinstance(v, torch.Tensor)} + for group in inner.param_groups + for p in group["params"] + if p in inner.state + ] + + +def run_optimizer_checkpoint(config: Tuple): + """Sharded optimizer save/load must round-trip Adam state for fused EP/TP-sharded experts. + + The upstream MoE checkpoint test saves a momentum-free SGD, whose state is empty, so the expert + optimizer-state gather/scatter path is never exercised; Adam keeps exp_avg/exp_avg_sq per param. + """ + model_name, zero_stage, ep_size, tp_size = config + rank = dist.get_rank() + torch.cuda.set_device(rank) + plugin = MoeHybridParallelPlugin( + pp_size=1, + tp_size=tp_size, + ep_size=ep_size, + zero_stage=zero_stage, + precision="bf16", + initial_scale=1, + overlap_communication=False, + find_unused_parameters=True, + ) + booster = Booster(plugin=plugin) + seed_all(1234) + model = _build(model_name).to(torch.bfloat16).cuda() + hidden_size = model.config.hidden_size + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + model, optimizer, *_ = booster.boost(model, optimizer) + + # Ranks of one TP group must see the same inputs; otherwise their replicated (non TP-sharded) parameters get + # different gradients and optimizer states, while a checkpoint keeps one copy per TP group. + seed_all(4321 + dist.get_rank(plugin.dp_group)) + inputs = torch.rand(2, NUM_TOKENS // 2, hidden_size, device="cuda", dtype=torch.bfloat16) + loss = model(inputs_embeds=inputs).last_hidden_state.float().pow(2).mean() + booster.backward(loss, optimizer) + optimizer.step() + optimizer.zero_grad() + + expected = _state_snapshot(optimizer) + assert expected, "optimizer has no tensor state; the round trip would be vacuous" + + context = tempfile.TemporaryDirectory() if rank == 0 else nullcontext() + with context as tmp: + shared = [tmp] + dist.broadcast_object_list(shared, src=0) + optim_dir = os.path.join(shared[0], "optim") + booster.save_optimizer(optimizer, optim_dir, shard=True) + dist.barrier() + inner = optimizer.unwrap() if hasattr(optimizer, "unwrap") else optimizer.optim + for state in inner.state.values(): + for value in state.values(): + if isinstance(value, torch.Tensor) and value.dim() > 0: + value.zero_() + booster.load_optimizer(optimizer, optim_dir) + dist.barrier() + + loaded = _state_snapshot(optimizer) + assert len(loaded) == len(expected) + for index, (left, right) in enumerate(zip(expected, loaded)): + assert left.keys() == right.keys(), f"param {index}: state keys differ" + for key in left: + assert torch.equal(left[key], right[key].to(left[key].device)), f"param {index}: {key} not restored" + + +@parameterize( + "config", + [ + # (model, zero_stage, ep_size, tp_size) on 4 ranks + ("mixtral", 1, 4, 1), + ("mixtral", 1, 2, 2), + ("deepseek_v3", 1, 4, 1), + # ep_size 2 on 4 ranks gives moe_dp_size 2, so expert states are also sharded by Zero. + ("mixtral", 1, 2, 1), + ("deepseek_v3", 1, 2, 1), + # expert sizes whose Zero padding differs between moe_dp_size and the global dp size + ("mixtral_odd", 1, 2, 1), + ], +) +def run_all(config: Tuple): + run_optimizer_checkpoint(config) + + +def check_optimizer_checkpoint(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + run_all() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_moe_optimizer_checkpoint(): + spawn(check_optimizer_checkpoint, 4) + + +if __name__ == "__main__": + test_moe_optimizer_checkpoint() diff --git a/tests/test_shardformer/test_model/test_moe_optimizer_coverage.py b/tests/test_shardformer/test_model/test_moe_optimizer_coverage.py new file mode 100644 index 000000000000..5bfad76cbb1e --- /dev/null +++ b/tests/test_shardformer/test_model/test_moe_optimizer_coverage.py @@ -0,0 +1,157 @@ +from typing import Tuple + +import pytest +import torch +import torch.distributed as dist + +from colossalai._compat import is_transformers_v5 + +if not is_transformers_v5(): + pytest.skip("fused MoE experts only exist in transformers v5", allow_module_level=True) + +from transformers import DeepseekV3Config, MixtralConfig +from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3Model +from transformers.models.mixtral.modeling_mixtral import MixtralModel + +import colossalai +from colossalai.booster.booster import Booster +from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin +from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn +from colossalai.testing.random import seed_all + +NUM_TOKENS, HIDDEN = 64, 32 + + +def _build(model_name: str): + if model_name == "mixtral": + config = MixtralConfig( + hidden_size=HIDDEN, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_local_experts=4, + num_experts_per_tok=2, + attn_implementation="eager", + ) + return MixtralModel(config) + if model_name == "mixtral_odd": + # One expert per rank under ep_size 2 with odd sizes: the local fused experts have 1 x 34 x 33 = 1122 and + # 1 x 33 x 17 = 561 elements, so padding them for moe_dp_size 2 (0 and 1) differs from padding them for the + # global dp size 4 (2 and 3). + config = MixtralConfig( + hidden_size=33, + intermediate_size=17, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_local_experts=2, + num_experts_per_tok=2, + attn_implementation="eager", + ) + return MixtralModel(config) + config = DeepseekV3Config( + vocab_size=128, + hidden_size=HIDDEN, + intermediate_size=64, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=4, + n_shared_experts=1, + n_routed_experts=8, + n_group=2, + topk_group=1, + num_experts_per_tok=2, + first_k_dense_replace=1, + kv_lora_rank=16, + q_lora_rank=32, + qk_rope_head_dim=8, + qk_nope_head_dim=16, + v_head_dim=16, + attn_implementation="eager", + ) + return DeepseekV3Model(config) + + +def run_coverage(config: Tuple): + """Every trainable parameter of the boosted model must be owned by the optimizer and updated. + + Numerical parity tests can miss parameters that silently fall out of the optimizer (e.g. MoE + experts sharded into new Parameter objects) because small updates stay within loose tolerances. + This check is structural. The update is checked on the fp32 master copy when the wrapper keeps + one: a bf16 working weight near 1.0 (e.g. RMSNorm) can round a small update away. + """ + model_name, zero_stage, ep_size, tp_size = config + rank = dist.get_rank() + torch.cuda.set_device(rank) + plugin = MoeHybridParallelPlugin( + pp_size=1, + tp_size=tp_size, + ep_size=ep_size, + zero_stage=zero_stage, + precision="bf16", + initial_scale=1, + overlap_communication=False, + find_unused_parameters=True, + ) + booster = Booster(plugin=plugin) + seed_all(1234) + model = _build(model_name).to(torch.bfloat16).cuda() + optimizer = torch.optim.SGD(model.parameters(), lr=1e3) + model, optimizer, *_ = booster.boost(model, optimizer) + + # `embed_tokens` is unused because the test feeds `inputs_embeds`. + params = [ + (name, p) for name, p in model.unwrap().named_parameters() if p.requires_grad and "embed_tokens" not in name + ] + label = f"{model_name} zero={zero_stage} ep={ep_size} tp={tp_size}" + working_to_master = optimizer.get_working_to_master_map() + if working_to_master is None: + owned_ids = {id(p) for group in optimizer.param_groups for p in group["params"]} + tracked = {name: p for name, p in params if id(p) in owned_ids} + else: + tracked = {name: working_to_master[id(p)] for name, p in params if id(p) in working_to_master} + not_owned = [name for name, _ in params if name not in tracked] + assert not not_owned, f"{label}: not owned by the optimizer {not_owned}" + + before = {name: t.detach().float().clone() for name, t in tracked.items()} + seed_all(4321 + rank) + inputs = torch.rand(2, NUM_TOKENS // 2, HIDDEN, device="cuda", dtype=torch.bfloat16) + loss = model(inputs_embeds=inputs).last_hidden_state.float().pow(2).mean() + booster.backward(loss, optimizer) + optimizer.step() + + unchanged = [name for name, t in tracked.items() if torch.equal(t.detach().float(), before[name])] + assert not unchanged, f"{label}: not updated {unchanged}" + assert any(".experts." in name for name, _ in params) + + +@parameterize( + "config", + [ + # (model, zero_stage, ep_size, tp_size) on 4 ranks + ("mixtral", 1, 4, 1), + ("mixtral", 1, 2, 2), + ("mixtral", 0, 1, 2), + ("deepseek_v3", 1, 4, 1), + ], +) +def run_all(config: Tuple): + run_coverage(config) + + +def check_coverage(rank, world_size, port): + colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") + run_all() + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_moe_optimizer_coverage(): + spawn(check_coverage, 4) + + +if __name__ == "__main__": + test_moe_optimizer_coverage() diff --git a/tests/test_shardformer/test_model/test_shard_blip2.py b/tests/test_shardformer/test_model/test_shard_blip2.py index 04a8f57e9df6..e023be1a6b4d 100644 --- a/tests/test_shardformer/test_model/test_shard_blip2.py +++ b/tests/test_shardformer/test_model/test_shard_blip2.py @@ -1,8 +1,11 @@ +from types import SimpleNamespace + import pytest import torch import colossalai from colossalai.logging import disable_existing_loggers +from colossalai.shardformer._utils import hf5_policy_error from colossalai.testing import ( assert_hf_output_close, clear_cache_before_run, @@ -10,12 +13,17 @@ rerun_if_address_is_in_use, spawn, ) -from tests.kit.model_zoo import model_zoo -from tests.test_shardformer.test_model._utils import build_model, check_grad, run_forward + + +def _load_test_dependencies(): + global model_zoo, build_model, check_grad, run_forward + from tests.kit.model_zoo import model_zoo + from tests.test_shardformer.test_model._utils import build_model, check_grad, run_forward def check_forward_backward(org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn): # check forward + _load_test_dependencies() org_output, org_loss, shard_output, shard_loss = run_forward( org_model, sharded_model, data_gen_fn, output_transform_fn, loss_fn ) @@ -97,6 +105,7 @@ def run_blip2_test( def check_blip2(rank, world_size, port): + _load_test_dependencies() disable_existing_loggers() colossalai.launch( rank=rank, @@ -108,6 +117,10 @@ def check_blip2(rank, world_size, port): run_blip2_test() +@pytest.mark.skipif( + hf5_policy_error("blip2", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("blip2", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.dist @rerun_if_address_is_in_use() @clear_cache_before_run() diff --git a/tests/test_shardformer/test_model/test_shard_chatglm2.py b/tests/test_shardformer/test_model/test_shard_chatglm2.py index 17a8bf318976..6cc4bce3fef3 100644 --- a/tests/test_shardformer/test_model/test_shard_chatglm2.py +++ b/tests/test_shardformer/test_model/test_shard_chatglm2.py @@ -1,25 +1,35 @@ +from types import SimpleNamespace + import pytest import torch import colossalai from colossalai.logging import disable_existing_loggers +from colossalai.shardformer._utils import hf5_policy_error from colossalai.shardformer.layer.utils import Randomizer from colossalai.tensor.d_tensor.api import clear_layout_converter from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn -from tests.kit.model_zoo import model_zoo -from tests.test_shardformer.test_model._utils import ( - build_model_from_hybrid_plugin, - check_all_grad_tensors, - check_loss, - check_output_hidden_state, - check_weight, - get_grad_tensors_for_check, - run_forward_backward_with_hybrid_plugin, - unwrap_model, -) + + +def _load_test_dependencies(): + global model_zoo, build_model_from_hybrid_plugin, check_all_grad_tensors + global check_loss, check_output_hidden_state, check_weight + global get_grad_tensors_for_check, run_forward_backward_with_hybrid_plugin, unwrap_model + from tests.kit.model_zoo import model_zoo + from tests.test_shardformer.test_model._utils import ( + build_model_from_hybrid_plugin, + check_all_grad_tensors, + check_loss, + check_output_hidden_state, + check_weight, + get_grad_tensors_for_check, + run_forward_backward_with_hybrid_plugin, + unwrap_model, + ) def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config): + _load_test_dependencies() org_model, org_optimizer, sharded_model, sharded_optimizer, criterion, booster = build_model_from_hybrid_plugin( model_fn, loss_fn, test_config ) @@ -287,6 +297,7 @@ def run_chatglm_3d_test(test_config): def check_chatglm(rank, world_size, port): + _load_test_dependencies() disable_existing_loggers() colossalai.launch( rank=rank, @@ -299,6 +310,7 @@ def check_chatglm(rank, world_size, port): def check_chatglm_3d(rank, world_size, port): + _load_test_dependencies() disable_existing_loggers() colossalai.launch( rank=rank, @@ -310,6 +322,10 @@ def check_chatglm_3d(rank, world_size, port): run_chatglm_3d_test() +@pytest.mark.skipif( + hf5_policy_error("chatglm2", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("chatglm2", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.dist @rerun_if_address_is_in_use() @clear_cache_before_run() @@ -317,6 +333,10 @@ def test_chatglm(): spawn(check_chatglm, 4) +@pytest.mark.skipif( + hf5_policy_error("chatglm2", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("chatglm2", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.largedist @rerun_if_address_is_in_use() @clear_cache_before_run() diff --git a/tests/test_shardformer/test_model/test_shard_command.py b/tests/test_shardformer/test_model/test_shard_command.py index 4d156d84db82..5a2da52bee97 100644 --- a/tests/test_shardformer/test_model/test_shard_command.py +++ b/tests/test_shardformer/test_model/test_shard_command.py @@ -1,4 +1,5 @@ import os +from types import SimpleNamespace import pytest import torch @@ -8,25 +9,33 @@ import colossalai from colossalai.logging import disable_existing_loggers from colossalai.shardformer import PipelineGradientCheckpointConfig +from colossalai.shardformer._utils import hf5_policy_error from colossalai.shardformer.layer.utils import Randomizer from colossalai.tensor.d_tensor.api import clear_layout_converter from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn -from tests.kit.model_zoo import model_zoo -from tests.test_shardformer.test_model._utils import ( - build_model_from_hybrid_plugin, - check_all_grad_tensors, - check_loss, - check_output_hidden_state, - check_weight, - get_grad_tensors_for_check, - run_forward_backward_with_hybrid_plugin, - unwrap_model, -) os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" +def _load_test_dependencies(): + global model_zoo, build_model_from_hybrid_plugin, check_all_grad_tensors + global check_loss, check_output_hidden_state, check_weight + global get_grad_tensors_for_check, run_forward_backward_with_hybrid_plugin, unwrap_model + from tests.kit.model_zoo import model_zoo + from tests.test_shardformer.test_model._utils import ( + build_model_from_hybrid_plugin, + check_all_grad_tensors, + check_loss, + check_output_hidden_state, + check_weight, + get_grad_tensors_for_check, + run_forward_backward_with_hybrid_plugin, + unwrap_model, + ) + + def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config): + _load_test_dependencies() enable_gradient_checkpointing = test_config.pop("enable_gradient_checkpointing", False) org_model, org_optimizer, sharded_model, sharded_optimizer, criterion, booster = build_model_from_hybrid_plugin( model_fn, loss_fn, test_config @@ -343,17 +352,23 @@ def run_command_3d_test(test_config): def check_command(rank, world_size, port): + _load_test_dependencies() disable_existing_loggers() colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") run_command_test() def check_command_3d(rank, world_size, port): + _load_test_dependencies() disable_existing_loggers() colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") run_command_3d_test() +@pytest.mark.skipif( + hf5_policy_error("command", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("command", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.dist @rerun_if_address_is_in_use() @clear_cache_before_run() @@ -361,6 +376,10 @@ def test_command(): spawn(check_command, 4) +@pytest.mark.skipif( + hf5_policy_error("command", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("command", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.largedist @rerun_if_address_is_in_use() @clear_cache_before_run() diff --git a/tests/test_shardformer/test_model/test_shard_deepseek.py b/tests/test_shardformer/test_model/test_shard_deepseek.py index 99409074c527..588d7ff348cd 100644 --- a/tests/test_shardformer/test_model/test_shard_deepseek.py +++ b/tests/test_shardformer/test_model/test_shard_deepseek.py @@ -1,6 +1,7 @@ import os import shutil from copy import deepcopy +from types import SimpleNamespace from typing import Tuple import pytest @@ -12,6 +13,7 @@ import colossalai from colossalai.booster.booster import Booster from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin +from colossalai.shardformer._utils import hf5_policy_error from colossalai.shardformer.layer.utils import Randomizer from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn from colossalai.testing.random import seed_all @@ -222,6 +224,10 @@ def check_deepseek_3d(rank, world_size, port): run_deepseek_3d_test() +@pytest.mark.skipif( + hf5_policy_error("deepseek", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("deepseek", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.dist @pytest.mark.parametrize("world_size", [4]) @rerun_if_address_is_in_use() @@ -229,6 +235,10 @@ def test_deepseek(world_size): spawn(check_deepseek, world_size) +@pytest.mark.skipif( + hf5_policy_error("deepseek", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("deepseek", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.largedist @pytest.mark.parametrize("world_size", [8]) @rerun_if_address_is_in_use() diff --git a/tests/test_shardformer/test_model/test_shard_deepseek_v3.py b/tests/test_shardformer/test_model/test_shard_deepseek_v3.py index 52cb8365d03e..2a5a00e5218e 100644 --- a/tests/test_shardformer/test_model/test_shard_deepseek_v3.py +++ b/tests/test_shardformer/test_model/test_shard_deepseek_v3.py @@ -7,6 +7,7 @@ from torch.testing import assert_close import colossalai +from colossalai._compat import is_transformers_v5 from colossalai.booster.plugin import MoeHybridParallelPlugin from colossalai.booster.plugin.moe_hybrid_parallel_plugin import MoeHybridParallelPlugin from colossalai.testing import parameterize, rerun_if_address_is_in_use, spawn @@ -14,14 +15,20 @@ from tests.kit.model_zoo.transformers.deepseek_v3 import ( data_gen_for_lm, init_deepseek, + init_deepseek_native, loss_fn_for_lm, output_transform_fn, + remote_code_unsupported_reason, ) from tests.test_shardformer.test_model._utils import ( build_model_from_hybrid_plugin, run_forward_backward_with_hybrid_plugin, ) +# DeepSeek-v3 has two implementations with different classes, expert layouts and policy keys: +# the Hub remote code (`trust_remote_code=True`) and the model built into Transformers. +MODEL_FNS = {"remote_code": init_deepseek, "native": init_deepseek_native} + def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config): enable_gradient_checkpointing = test_config.pop("enable_gradient_checkpointing", False) @@ -41,21 +48,34 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, org_model, sharded_model, sharded_optimizer, data_gen_fn, output_transform_fn, criterion, booster ) - assert_close(org_loss, sharded_loss) + # The loss is an FP32 scalar from bf16 models; assert_close's FP32 default (rtol=1.3e-6) would demand + # bit-identical kernels, but EP and the Transformers implementation combine the routed experts in a + # different order. Measured difference is ~1.8e-5 relative, while a sharding error (e.g. mis-routed + # tokens) moves the loss by >1e-2, so 1e-3 separates rounding noise from real faults. + assert_close(org_loss, sharded_loss, rtol=1e-3, atol=1e-3) param_dict = {n: p for n, p in org_model.named_parameters()} + sharded_modules = dict(sharded_model.unwrap().named_modules()) for n, p in sharded_model.unwrap().named_parameters(): if n in param_dict: + reference_grad = param_dict[n].grad + if reference_grad is not None and p.shape != param_dict[n].shape: + # Fused experts (`gate_up_proj [E, 2I, H]`, `down_proj [E, H, I]`) keep only this EP rank's + # experts, so compare against the matching slice of the full reference gradient. + owner = sharded_modules[n.rsplit(".", 2)[0]] + reference_grad = reference_grad[ + owner.expert_start_idx : owner.expert_start_idx + owner.num_experts_per_ep + ] if booster.plugin.zero_stage == 0: grad = p.grad - target_grad = param_dict[n].grad + target_grad = reference_grad else: grad = sharded_optimizer.get_working_grad_by_param_id(id(p)) pg = sharded_optimizer.param_to_pg[p] - target_grad = param_dict[n].grad + target_grad = reference_grad if target_grad is None: continue - target_grad = target_grad.view(-1).chunk(dist.get_world_size(pg))[dist.get_rank(pg)] + target_grad = target_grad.reshape(-1).chunk(dist.get_world_size(pg))[dist.get_rank(pg)] assert_close(grad, target_grad, atol=5e-1, rtol=0) @@ -67,7 +87,7 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, (1, 2), ], ) -def run_deepseek_v3_test(config: Tuple[int, ...]): +def run_deepseek_v3_test(config: Tuple[int, ...], model_fn): zero_stage, ep_size = config plugin_config = dict( pp_size=1, @@ -80,7 +100,7 @@ def run_deepseek_v3_test(config: Tuple[int, ...]): ) check_forward_backward( - init_deepseek, + model_fn, data_gen_for_lm, output_transform_fn, loss_fn_for_lm, @@ -88,16 +108,27 @@ def run_deepseek_v3_test(config: Tuple[int, ...]): ) -def check_deepseek_v3(rank, world_size, port): +def check_deepseek_v3(rank, world_size, port, implementation): colossalai.launch(rank=rank, world_size=world_size, host="localhost", port=port, backend="nccl") - run_deepseek_v3_test() + run_deepseek_v3_test(model_fn=MODEL_FNS[implementation]) @pytest.mark.dist +@pytest.mark.skipif( + not is_transformers_v5(), reason="the v4 policy targets the remote-code model; native DeepSeek-v3 needs v5" +) @pytest.mark.parametrize("world_size", [4]) @rerun_if_address_is_in_use() def test_deepseek_v3(world_size): - spawn(check_deepseek_v3, world_size) + spawn(check_deepseek_v3, world_size, implementation="native") + + +@pytest.mark.dist +@pytest.mark.skipif(remote_code_unsupported_reason() is not None, reason=str(remote_code_unsupported_reason())) +@pytest.mark.parametrize("world_size", [4]) +@rerun_if_address_is_in_use() +def test_deepseek_v3_remote_code(world_size): + spawn(check_deepseek_v3, world_size, implementation="remote_code") if __name__ == "__main__": diff --git a/tests/test_shardformer/test_model/test_shard_llama.py b/tests/test_shardformer/test_model/test_shard_llama.py index 54e4ccde1a7a..0191e4b832a7 100644 --- a/tests/test_shardformer/test_model/test_shard_llama.py +++ b/tests/test_shardformer/test_model/test_shard_llama.py @@ -7,6 +7,7 @@ from torch.testing import assert_close import colossalai +from colossalai._compat import is_transformers_v5 from colossalai.logging import disable_existing_loggers from colossalai.pipeline.schedule.v_schedule import PipelineGraph from colossalai.shardformer import PipelineGradientCheckpointConfig @@ -30,6 +31,71 @@ os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" +@pytest.mark.skipif(not is_transformers_v5(), reason="`create_causal_mask` is only used by the v5 implementation") +def test_llama_pp_all_to_all_causal_mask_uses_full_query_length(): + from transformers.masking_utils import create_causal_mask + from transformers.models.llama.configuration_llama import LlamaConfig + + from colossalai.shardformer.modeling.llama import _get_causal_mask_input_embeddings + + hidden_states = torch.empty(4, 2, 8) + assert _get_causal_mask_input_embeddings(hidden_states, query_length=2, sequence_is_sharded=True) is hidden_states + assert _get_causal_mask_input_embeddings(hidden_states, query_length=4, sequence_is_sharded=False) is hidden_states + mask_inputs_embeds = _get_causal_mask_input_embeddings(hidden_states, query_length=4, sequence_is_sharded=True) + assert mask_inputs_embeds.shape == (4, 4, 8) + assert mask_inputs_embeds.stride(1) == 0 + assert mask_inputs_embeds.untyped_storage().data_ptr() == hidden_states.untyped_storage().data_ptr() + + config = LlamaConfig() + config._attn_implementation = "eager" + position_ids = torch.arange(4).unsqueeze(0).expand(4, -1) + mask = create_causal_mask( + config=config, + inputs_embeds=mask_inputs_embeds, + attention_mask=None, + past_key_values=None, + position_ids=position_ids, + allow_is_causal_skip=False, + ) + assert mask is not None and mask.shape == (4, 1, 4, 4) + + +@pytest.mark.skipif(not is_transformers_v5(), reason="`create_causal_mask` is only used by the v5 implementation") +@pytest.mark.parametrize("attn_implementation", ["sdpa", "flash_attention_2", "flex_attention"]) +def test_llama_explicit_softmax_mask_is_additive(attn_implementation): + from transformers.masking_utils import create_causal_mask + from transformers.models.llama.configuration_llama import LlamaConfig + + from colossalai.shardformer.modeling.llama import _EagerMaskConfig + + # Under these implementations `create_causal_mask` returns a boolean mask, None / a 2D padding + # mask, or a BlockMask; the explicit-softmax attention needs an additive float mask instead. + config = LlamaConfig(num_hidden_layers=1, hidden_size=32, intermediate_size=64, num_attention_heads=4) + config._attn_implementation = attn_implementation + mask_config = _EagerMaskConfig(config) + assert mask_config._attn_implementation == "eager" + assert config._attn_implementation == attn_implementation + # Every other field is read from the live config, so later changes are seen at once. + config.num_attention_heads = 8 + assert mask_config.num_attention_heads == 8 + config.num_attention_heads = 4 + + future = torch.ones(4, 4, dtype=torch.bool).triu(1) + for attention_mask in (torch.ones(2, 4, dtype=torch.long), torch.tensor([[1, 1, 1, 0], [1, 1, 1, 1]])): + mask = create_causal_mask( + config=mask_config, + inputs_embeds=torch.empty(2, 4, 32, dtype=torch.float16), + attention_mask=attention_mask, + past_key_values=None, + position_ids=torch.arange(4).unsqueeze(0).expand(2, -1), + allow_is_causal_skip=False, + ) + assert mask.dtype == torch.float16 and mask.shape == (2, 1, 4, 4) + assert (mask[:, :, future] == torch.finfo(torch.float16).min).all() + assert (mask[0, 0, :, 3] == torch.finfo(torch.float16).min).all() # the padded key + assert (mask[1, :, ~future] == 0).all() + + def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config): enable_gradient_checkpointing = test_config.pop("enable_gradient_checkpointing", False) org_model, org_optimizer, sharded_model, sharded_optimizer, criterion, booster = build_model_from_hybrid_plugin( @@ -228,6 +294,49 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, "precision": "fp16", "initial_scale": 1, }, + # SP + PP without flash attention: later stages build the eager causal mask from a sequence shard. + # v5 only: the v4 implementation (`_llama_tf4`, upstream as-is) builds that mask from the shard + # length and fails there with a 12 vs 24 size mismatch, leaving the other stage blocked in P2P. + *( + [ + { # Ulysess + PP + "tp_size": 1, + "pp_size": 2, + "sp_size": 2, + "num_microbatches": 2, + "enable_sequence_parallelism": True, + "sequence_parallelism_mode": "all_to_all", + "use_lazy_init": False, + "zero_stage": 1, + "precision": "fp16", + "initial_scale": 1, + }, + { # split_gather + PP + "tp_size": 2, + "pp_size": 2, + "num_microbatches": 2, + "enable_sequence_parallelism": True, + "sequence_parallelism_mode": "split_gather", + "use_lazy_init": False, + "zero_stage": 0, + "precision": "fp32", + "initial_scale": 1, + }, + { # ring + PP; zero_stage 0 so that the weight gradients are checked + "tp_size": 2, + "pp_size": 2, + "num_microbatches": 2, + "enable_sequence_parallelism": True, + "sequence_parallelism_mode": "ring", + "use_lazy_init": False, + "zero_stage": 0, + "precision": "fp32", + "initial_scale": 1, + }, + ] + if is_transformers_v5() + else [] + ), { "tp_size": 2, "pp_size": 1, diff --git a/tests/test_shardformer/test_model/test_shard_opt.py b/tests/test_shardformer/test_model/test_shard_opt.py index b7c77d20b807..e474e7c88572 100644 --- a/tests/test_shardformer/test_model/test_shard_opt.py +++ b/tests/test_shardformer/test_model/test_shard_opt.py @@ -1,29 +1,38 @@ import os +from types import SimpleNamespace import pytest import torch import colossalai from colossalai.logging import disable_existing_loggers +from colossalai.shardformer._utils import hf5_policy_error from colossalai.shardformer.layer.utils import Randomizer from colossalai.tensor.d_tensor.api import clear_layout_converter from colossalai.testing import clear_cache_before_run, parameterize, rerun_if_address_is_in_use, spawn -from tests.kit.model_zoo import model_zoo -from tests.test_shardformer.test_model._utils import ( - build_model_from_hybrid_plugin, - check_all_grad_tensors, - check_loss, - check_output_hidden_state, - check_weight, - get_grad_tensors_for_check, - run_forward_backward_with_hybrid_plugin, - unwrap_model, -) os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true" +def _load_test_dependencies(): + global model_zoo, build_model_from_hybrid_plugin, check_all_grad_tensors + global check_loss, check_output_hidden_state, check_weight + global get_grad_tensors_for_check, run_forward_backward_with_hybrid_plugin, unwrap_model + from tests.kit.model_zoo import model_zoo + from tests.test_shardformer.test_model._utils import ( + build_model_from_hybrid_plugin, + check_all_grad_tensors, + check_loss, + check_output_hidden_state, + check_weight, + get_grad_tensors_for_check, + run_forward_backward_with_hybrid_plugin, + unwrap_model, + ) + + def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config): + _load_test_dependencies() org_model, org_optimizer, sharded_model, sharded_optimizer, criterion, booster = build_model_from_hybrid_plugin( model_fn, loss_fn, test_config ) @@ -231,6 +240,7 @@ def run_opt_3d_test(test_config): def check_OPTModel(rank, world_size, port): + _load_test_dependencies() disable_existing_loggers() colossalai.launch( rank=rank, @@ -243,6 +253,7 @@ def check_OPTModel(rank, world_size, port): def check_opt_3d(rank, world_size, port): + _load_test_dependencies() disable_existing_loggers() colossalai.launch( rank=rank, @@ -254,6 +265,10 @@ def check_opt_3d(rank, world_size, port): run_opt_3d_test() +@pytest.mark.skipif( + hf5_policy_error("opt", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("opt", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.dist @rerun_if_address_is_in_use() @clear_cache_before_run() @@ -261,6 +276,10 @@ def test_OPTModel(): spawn(check_OPTModel, 4) +@pytest.mark.skipif( + hf5_policy_error("opt", SimpleNamespace(enable_tensor_parallelism=True)) is not None, + reason=hf5_policy_error("opt", SimpleNamespace(enable_tensor_parallelism=True)) or "", +) @pytest.mark.largedist @rerun_if_address_is_in_use() @clear_cache_before_run() diff --git a/tests/test_shardformer/test_model/test_shard_qwen3.py b/tests/test_shardformer/test_model/test_shard_qwen3.py index 9670a5999e8c..0118eb7c888a 100644 --- a/tests/test_shardformer/test_model/test_shard_qwen3.py +++ b/tests/test_shardformer/test_model/test_shard_qwen3.py @@ -1,3 +1,5 @@ +import copy + import pytest import torch import transformers @@ -55,6 +57,24 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, grads_to_check.update(col_layer_grads) grads_to_check.update(row_layer_grads) + norm_params = [] + if booster.plugin.zero_stage == 0: + reference_params = dict(qwen3_model.named_parameters()) + for name, param in shard_qwen3_model.named_parameters(): + if name.endswith(("q_norm.weight", "k_norm.weight")) and param is not None: + reference = reference_params[name] + assert (reference.grad is None) == (param.grad is None), name + if param.grad is None: + continue + atol, rtol = (1e-6, 1e-4) if test_config["precision"] == "fp32" else (5e-3, 5e-3) + grads_to_check[name] = { + "org_grad": reference.grad.detach().float().clone(), + "shard_grad": param.grad.detach().float().clone(), + "atol": atol, + "rtol": rtol, + } + norm_params.append((reference, param)) + # optimizer executes step org_optimizer.step() sharded_optimizer.step() @@ -81,9 +101,19 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, qwen3_model, shard_qwen3_model, col_layer_for_check, tp_group, atol=atol, rtol=rtol, dim=1, verbose=False ) + norm_atol, norm_rtol = (1e-3, 1e-3) if test_config["precision"] == "fp32" else (5e-3, 5e-3) + for reference, param in norm_params: + torch.testing.assert_close(param.float(), reference.float(), atol=norm_atol, rtol=norm_rtol) + # check grads check_all_grad_tensors(grads_to_check) + # NCCL group destruction must happen in the same order on every rank, + # rather than during rank-dependent collection of model/optimizer cycles. + torch.distributed.barrier() + booster.plugin.pg_mesh.destroy_mesh_process_groups() + del booster.plugin.pg_mesh + torch.distributed.barrier() torch.cuda.empty_cache() @@ -210,7 +240,7 @@ def run_qwen3_test(test_config): for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): try: - check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config) + check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, copy.deepcopy(test_config)) except Exception as e: print(f"Failed config: {test_config}") raise e @@ -259,7 +289,7 @@ def run_qwen3_3d_test(test_config): for name, (model_fn, data_gen_fn, output_transform_fn, loss_fn, _) in sub_model_zoo.items(): try: - check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, test_config) + check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, copy.deepcopy(test_config)) except Exception as e: print(f"Failed config: {test_config}") raise e diff --git a/tests/test_shardformer/test_model/test_shard_vit.py b/tests/test_shardformer/test_model/test_shard_vit.py index 99d6f7d2ccc1..903708bd7568 100644 --- a/tests/test_shardformer/test_model/test_shard_vit.py +++ b/tests/test_shardformer/test_model/test_shard_vit.py @@ -2,6 +2,7 @@ import torch import colossalai +from colossalai._compat import is_transformers_v5 from colossalai.logging import disable_existing_loggers from colossalai.shardformer.layer.utils import Randomizer from colossalai.tensor.d_tensor.api import clear_layout_converter @@ -36,8 +37,12 @@ def check_forward_backward(model_fn, data_gen_fn, output_transform_fn, loss_fn, shard_vit_model = unwrap_model(sharded_model, "ViTModel", "vit") # check grad - row_layer_for_check = ["encoder.layer[0].attention.attention.query", "embeddings.patch_embeddings.projection"] - col_layer_for_check = ["encoder.layer[0].attention.output.dense"] + if is_transformers_v5(): + row_layer_for_check = ["layers[0].attention.q_proj", "embeddings.patch_embeddings.projection"] + col_layer_for_check = ["layers[0].attention.o_proj"] + else: + row_layer_for_check = ["encoder.layer[0].attention.attention.query", "embeddings.patch_embeddings.projection"] + col_layer_for_check = ["encoder.layer[0].attention.output.dense"] # Save gradient tensors for comparison between the original model and the sharded model before optimizer step. grads_to_check = {}