109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Loader-side adapters for proven direct-checkpoint mmap layouts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from functools import partial
|
|
from typing import Protocol
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DirectMmapTensorPolicy:
|
|
"""Loader-owned proof for one tensor's direct mmap behavior."""
|
|
|
|
allow_custom_loader: bool = False
|
|
transform: Callable[[torch.Tensor], torch.Tensor] | None = None
|
|
|
|
|
|
class DirectMmapAdapter(Protocol):
|
|
def policy_for(
|
|
self,
|
|
runtime_name: str,
|
|
target: torch.Tensor,
|
|
) -> DirectMmapTensorPolicy | None: ...
|
|
|
|
|
|
class _MiniMaxH3DirectMmapAdapter:
|
|
"""TP1 runtime-layout contract already enforced by MiniMax-H3's loader."""
|
|
|
|
def __init__(self, pipeline: nn.Module) -> None:
|
|
self.pipeline = pipeline
|
|
|
|
def policy_for(
|
|
self,
|
|
runtime_name: str,
|
|
target: torch.Tensor,
|
|
) -> DirectMmapTensorPolicy:
|
|
del target
|
|
transform = None
|
|
suffix = ".qkv_proj.weight"
|
|
if runtime_name.endswith(suffix):
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
MiniMaxH3Attention,
|
|
_reorder_grouped_qkv_to_qkv,
|
|
)
|
|
|
|
attention_path = runtime_name[: -len(suffix)]
|
|
attention = self.pipeline.get_submodule(attention_path)
|
|
if not isinstance(attention, MiniMaxH3Attention):
|
|
raise ValueError(
|
|
f"MiniMax-H3 direct-mmap adapter expected attention at {attention_path!r}, "
|
|
f"got {type(attention).__name__}"
|
|
)
|
|
transform = partial(
|
|
_reorder_grouped_qkv_to_qkv,
|
|
num_query_groups=attention.total_num_heads,
|
|
heads_per_group=1,
|
|
head_dim=attention.head_dim,
|
|
)
|
|
return DirectMmapTensorPolicy(
|
|
allow_custom_loader=True,
|
|
transform=transform,
|
|
)
|
|
|
|
|
|
class _Cosmos3DirectMmapAdapter:
|
|
"""Preserve Cosmos3's existing TP1 direct-mmap load contract."""
|
|
|
|
def policy_for(
|
|
self,
|
|
runtime_name: str,
|
|
target: torch.Tensor,
|
|
) -> DirectMmapTensorPolicy:
|
|
del runtime_name, target
|
|
# Cosmos3 has no packed module mapping or TP1 semantic transform after
|
|
# its checkpoint-key remap. Its custom vLLM loaders reduce to exact
|
|
# copies at TP1; topology and metadata are checked by plan preflight.
|
|
return DirectMmapTensorPolicy(allow_custom_loader=True)
|
|
|
|
|
|
def get_direct_mmap_adapter(pipeline: nn.Module) -> DirectMmapAdapter | None:
|
|
"""Return a loader-owned adapter without requiring a pipeline flag."""
|
|
from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
|
|
MiniMaxH3DiTModel,
|
|
)
|
|
|
|
if any(isinstance(module, MiniMaxH3DiTModel) for module in pipeline.modules()):
|
|
return _MiniMaxH3DirectMmapAdapter(pipeline)
|
|
|
|
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
|
|
Cosmos3VFMTransformer,
|
|
)
|
|
|
|
if any(isinstance(module, Cosmos3VFMTransformer) for module in pipeline.modules()):
|
|
return _Cosmos3DirectMmapAdapter()
|
|
return None
|
|
|
|
|
|
__all__ = [
|
|
"DirectMmapAdapter",
|
|
"DirectMmapTensorPolicy",
|
|
"get_direct_mmap_adapter",
|
|
]
|