41. ModelRunner 解析：vLLM的模型执行核心

# vllm/model_runner.py
import torch
from typing import Dict, List, Any, Optional
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer

class ModelRunner:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.model = None
        self.tokenizer = None
        self.compiled_model = None
        self.device = None
        self.dtype = None
        self.parallel_mode = None
        
        # 初始化组件
        self.memory_manager = MemoryManager(config)
        self.parallel_manager = ParallelManager(config)
        self.compile_optimizer = CompileOptimizer(config)
        
        # 设置设备和数据类型
        self._setup_device()
        self._setup_dtype()
    
    def _setup_device(self):
        """设置设备"""
        if torch.cuda.is_available():
            self.device = torch.device(f"cuda:{self.config.get('device', 0)}")
        else:
            self.device = torch.device("cpu")
    
    def _setup_dtype(self):
        """设置数据类型"""
        self.dtype = getattr(torch, self.config.get('dtype', 'float16'))
    
    def load_model(self):
        """加载模型"""
        from vllm.model_loader import ModelLoader
        
        model_loader = ModelLoader(self.config)
        loaded = model_loader.load()
        
        self.model = loaded['model']
        self.tokenizer = loaded['tokenizer']
        
        # 设置模型设备和数据类型
        self.model = self.model.to(self.device, dtype=self.dtype)
        
        # 初始化并行设置
        self.parallel_manager.setup_parallel(self.model)
        
        # 如果启用编译，编译模型
        if self.config.get('compile', False):
            self.compile_model()
    
    def compile_model(self):
        """编译模型"""
        self.compiled_model = self.compile_optimizer.compile(self.model)
    
    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> Dict[str, Any]:
        """模型前向传播"""
        # 选择使用编译后的模型或原始模型
        model = self.compiled_model if self.compiled_model is not None else self.model
        
        # 执行前向传播
        with torch.no_grad():
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        
        return {
            'logits': outputs.logits,
            'hidden_states': outputs.hidden_states if hasattr(outputs, 'hidden_states') else None
        }
    
    def generate(self, input_ids: torch.Tensor, **kwargs) -> Dict[str, Any]:
        """生成文本"""
        # 简化实现，实际会调用vLLM的生成逻辑
        return self.forward(input_ids, **kwargs)

# vllm/forward_pass.py
import torch
from typing import Dict, List, Any, Optional

class ForwardPass:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.memory_manager = MemoryManager(config)
        self.parallel_manager = ParallelManager(config)
    
    def setup_kv_cache(self, batch_size: int, seq_len: int, num_layers: int) -> Dict[str, Any]:
        """设置KVCache"""
        # 计算KVCache所需内存
        kv_cache_size = self._calculate_kv_cache_size(batch_size, seq_len, num_layers)
        
        # 分配内存
        kv_cache = self.memory_manager.allocate_memory(kv_cache_size)
        
        return {
            'kv_cache': kv_cache,
            'batch_size': batch_size,
            'seq_len': seq_len,
            'num_layers': num_layers
        }
    
    def _calculate_kv_cache_size(self, batch_size: int, seq_len: int, num_layers: int) -> int:
        """计算KVCache大小"""
        # 简化计算，实际会根据模型架构和数据类型计算
        head_dim = self.config.get('head_dim', 128)
        num_heads = self.config.get('num_heads', 32)
        kv_heads = self.config.get('kv_heads', num_heads)
        dtype_size = torch.finfo(self.config.get('dtype', torch.float16)).bits // 8
        
        # KVCache大小 = 2（K和V） * 层数 * 批次大小 * 头数 * 序列长度 * 头维度 * 数据类型大小
        return 2 * num_layers * batch_size * kv_heads * seq_len * head_dim * dtype_size
    
    def update_kv_cache(self, kv_cache: Dict[str, Any], new_seq_len: int) -> Dict[str, Any]:
        """更新KVCache大小"""
        batch_size = kv_cache['batch_size']
        num_layers = kv_cache['num_layers']
        
        # 计算新的KVCache大小
        new_kv_cache_size = self._calculate_kv_cache_size(batch_size, new_seq_len, num_layers)
        
        # 重新分配或扩展内存
        kv_cache['kv_cache'] = self.memory_manager.reallocate_memory(
            kv_cache['kv_cache'], new_kv_cache_size
        )
        kv_cache['seq_len'] = new_seq_len
        
        return kv_cache
    
    def execute(self, model: torch.nn.Module, input_ids: torch.Tensor, 
               attention_mask: Optional[torch.Tensor] = None, 
               kv_cache: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """执行前向传播"""
        with torch.no_grad():
            outputs = model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                use_cache=True if kv_cache is not None else False,
                past_key_values=kv_cache['kv_cache'] if kv_cache is not None else None
            )
        
        result = {
            'logits': outputs.logits
        }
        
        # 更新KVCache
        if kv_cache is not None:
            result['kv_cache'] = {
                'kv_cache': outputs.past_key_values,
                'batch_size': input_ids.shape[0],
                'seq_len': input_ids.shape[1] + (kv_cache['seq_len'] if kv_cache['seq_len'] > 0 else 0),
                'num_layers': kv_cache['num_layers']
            }
        
        return result

# vllm/compile_optimizer.py
import torch
from typing import Dict, Any, Optional
import os
import hashlib

class CompileOptimizer:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.compile_cache_dir = self.config.get('compile_cache_dir', './compile_cache')
        os.makedirs(self.compile_cache_dir, exist_ok=True)
    
    def _get_cache_key(self, model: torch.nn.Module) -> str:
        """生成编译缓存键"""
        # 简化实现，实际会考虑模型架构、配置、硬件等因素
        model_hash = hashlib.md5(str(model).encode()).hexdigest()
        config_hash = hashlib.md5(str(self.config).encode()).hexdigest()
        return f"{model_hash}_{config_hash}"
    
    def _load_from_cache(self, cache_key: str) -> Optional[torch.nn.Module]:
        """从缓存加载编译后的模型"""
        cache_path = os.path.join(self.compile_cache_dir, f"{cache_key}.pt")
        if os.path.exists(cache_path):
            try:
                return torch.load(cache_path)
            except Exception as e:
                print(f"Failed to load compiled model from cache: {e}")
        return None
    
    def _save_to_cache(self, cache_key: str, compiled_model: torch.nn.Module):
        """将编译后的模型保存到缓存"""
        cache_path = os.path.join(self.compile_cache_dir, f"{cache_key}.pt")
        try:
            torch.save(compiled_model, cache_path)
        except Exception as e:
            print(f"Failed to save compiled model to cache: {e}")
    
    def optimize_graph(self, model: torch.nn.Module) -> torch.nn.Module:
        """优化模型图"""
        # 简化实现，实际会使用PyTorch的图优化功能
        if hasattr(torch, 'fx'):
            from torch import fx
            # 使用FX进行图优化
            tracer = fx.Tracer()
            graph = tracer.trace(model)
            # 应用优化变换
            # ...
        return model
    
    def compile(self, model: torch.nn.Module) -> torch.nn.Module:
        """编译模型"""
        # 生成缓存键
        cache_key = self._get_cache_key(model)
        
        # 尝试从缓存加载
        compiled_model = self._load_from_cache(cache_key)
        if compiled_model is not None:
            print("Loaded compiled model from cache")
            return compiled_model
        
        print("Compiling model...")
        
        # 优化模型图
        optimized_model = self.optimize_graph(model)
        
        # 编译模型
        compile_config = {
            'backend': self.config.get('compile_backend', 'inductor'),
            'mode': self.config.get('compile_mode', 'max-autotune'),
        }
        
        compiled_model = torch.compile(optimized_model, **compile_config)
        
        # 保存到缓存
        self._save_to_cache(cache_key, compiled_model)
        
        print("Model compilation completed")
        return compiled_model

# vllm/parallel_manager.py
import torch
import torch.distributed as dist
from typing import Dict, Any, Optional

class ParallelManager:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.tensor_parallel_size = config.get('tensor_parallel_size', 1)
        self.pipeline_parallel_size = config.get('pipeline_parallel_size', 1)
        self.world_size = self.tensor_parallel_size * self.pipeline_parallel_size
        
        self.tensor_parallel_group = None
        self.pipeline_parallel_group = None
    
    def setup_parallel(self, model: torch.nn.Module):
        """设置并行策略"""
        if self.world_size <= 1:
            return model
        
        # 初始化分布式通信
        if not dist.is_initialized():
            dist.init_process_group(
                backend='nccl',
                rank=self.config.get('rank', 0),
                world_size=self.world_size
            )
        
        # 设置张量并行
        if self.tensor_parallel_size > 1:
            self._setup_tensor_parallel()
            model = self._apply_tensor_parallel(model)
        
        # 设置流水线并行
        if self.pipeline_parallel_size > 1:
            self._setup_pipeline_parallel()
            model = self._apply_pipeline_parallel(model)
        
        return model
    
    def _setup_tensor_parallel(self):
        """设置张量并行"""
        # 简化实现，实际会创建张量并行组
        self.tensor_parallel_group = dist.new_group(
            ranks=list(range(self.tensor_parallel_size))
        )
    
    def _setup_pipeline_parallel(self):
        """设置流水线并行"""
        # 简化实现，实际会创建流水线并行组
        self.pipeline_parallel_group = dist.new_group(
            ranks=list(range(0, self.world_size, self.tensor_parallel_size))
        )
    
    def _apply_tensor_parallel(self, model: torch.nn.Module) -> torch.nn.Module:
        """应用张量并行"""
        # 简化实现，实际会使用模型并行库（如Megatron-LM）进行张量并行
        if hasattr(model, 'parallelize'):
            model = model.parallelize(self.tensor_parallel_group)
        return model
    
    def _apply_pipeline_parallel(self, model: torch.nn.Module) -> torch.nn.Module:
        """应用流水线并行"""
        # 简化实现，实际会将模型划分为多个阶段
        # ...
        return model
    
    def execute_parallel(self, func: Callable, *args, **kwargs) -> Any:
        """并行执行函数"""
        if self.world_size <= 1:
            return func(*args, **kwargs)
        
        # 简化实现，实际会根据并行策略执行
        return func(*args, **kwargs)
    
    def gather_results(self, tensor: torch.Tensor) -> torch.Tensor:
        """收集并行执行的结果"""
        if self.world_size <= 1:
            return tensor
        
        # 简化实现，实际会根据并行策略收集结果
        gathered = [torch.empty_like(tensor) for _ in range(self.world_size)]
        dist.all_gather(gathered, tensor)
        return torch.cat(gathered, dim=0)

# vllm/memory_manager.py
import torch
from typing import Dict, Any, Optional

class MemoryManager:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.memory_pool = {}
        self.memory_stats = {
            'total_allocated': 0,
            'total_free': 0,
            'peak_allocated': 0,
            'fragmentation_rate': 0.0
        }
    
    def allocate_memory(self, size: int) -> Any:
        """分配内存"""
        # 简化实现，实际会使用更高效的内存分配策略
        if size == 0:
            return None
        
        # 检查内存池是否有合适的空闲内存
        for key in list(self.memory_pool.keys()):
            if self.memory_pool[key]['size'] >= size and not self.memory_pool[key]['in_use']:
                # 复用内存
                self.memory_pool[key]['in_use'] = True
                self.memory_stats['total_allocated'] += size
                if self.memory_stats['total_allocated'] > self.memory_stats['peak_allocated']:
                    self.memory_stats['peak_allocated'] = self.memory_stats['total_allocated']
                return self.memory_pool[key]['memory']
        
        # 分配新内存
        device = torch.device(f"cuda:{self.config.get('device', 0)}")
        memory = torch.empty(size, dtype=torch.uint8, device=device)
        
        # 添加到内存池
        self.memory_pool[id(memory)] = {
            'memory': memory,
            'size': size,
            'in_use': True
        }
        
        self.memory_stats['total_allocated'] += size
        if self.memory_stats['total_allocated'] > self.memory_stats['peak_allocated']:
            self.memory_stats['peak_allocated'] = self.memory_stats['total_allocated']
        
        return memory
    
    def free_memory(self, memory: Any):
        """释放内存"""
        if memory is None:
            return
        
        memory_id = id(memory)
        if memory_id in self.memory_pool:
            self.memory_pool[memory_id]['in_use'] = False
            self.memory_stats['total_allocated'] -= self.memory_pool[memory_id]['size']
            self.memory_stats['total_free'] += self.memory_pool[memory_id]['size']
    
    def reallocate_memory(self, memory: Any, new_size: int) -> Any:
        """重新分配内存"""
        if memory is None:
            return self.allocate_memory(new_size)
        
        # 释放旧内存
        self.free_memory(memory)
        
        # 分配新内存
        return self.allocate_memory(new_size)
    
    def reuse_memory(self, memory: Any, new_size: int) -> Any:
        """复用内存，如果大小合适"""
        if memory is None:
            return self.allocate_memory(new_size)
        
        memory_id = id(memory)
        if memory_id in self.memory_pool and self.memory_pool[memory_id]['size'] >= new_size:
            # 可以复用
            self.memory_pool[memory_id]['in_use'] = True
            self.memory_stats['total_allocated'] += new_size
            if self.memory_stats['total_allocated'] > self.memory_stats['peak_allocated']:
                self.memory_stats['peak_allocated'] = self.memory_stats['total_allocated']
            return memory
        else:
            # 重新分配
            return self.reallocate_memory(memory, new_size)
    
    def defragment_memory(self):
        """整理内存碎片"""
        # 简化实现，实际会将多个小的空闲内存块合并为大的块
        free_blocks = [
            block for block in self.memory_pool.values() 
            if not block['in_use']
        ]
        
        if not free_blocks:
            return
        
        # 按大小排序
        free_blocks.sort(key=lambda x: x['size'], reverse=True)
        
        # 合并空闲块（简化实现）
        total_free = sum(block['size'] for block in free_blocks)
        print(f"Defragmenting memory: {len(free_blocks)} free blocks, {total_free} bytes")
    
    def get_memory_stats(self) -> Dict[str, Any]:
        """获取内存统计信息"""
        # 计算碎片率（简化实现）
        free_blocks = [block for block in self.memory_pool.values() if not block['in_use']]
        if free_blocks:
            max_free_block = max(block['size'] for block in free_blocks)
            total_free = sum(block['size'] for block in free_blocks)
            self.memory_stats['fragmentation_rate'] = 1.0 - (max_free_block / total_free) if total_free > 0 else 0.0
        else:
            self.memory_stats['fragmentation_rate'] = 0.0
        
        return self.memory_stats

# vllm/model_adapters.py
from typing import Dict, Any, Callable
import torch
from transformers import AutoModelForCausalLM

class ModelAdapter:
    """模型适配器基类"""
    def __init__(self, config: Dict[str, Any]):
        self.config = config
    
    def load_model(self, model_path: str) -> torch.nn.Module:
        """加载模型"""
        return AutoModelForCausalLM.from_pretrained(
            model_path,
            trust_remote_code=self.config.get('trust_remote_code', False),
            torch_dtype=getattr(torch, self.config.get('dtype', 'float16')),
        )
    
    def forward(self, model: torch.nn.Module, input_ids: torch.Tensor, **kwargs) -> Dict[str, Any]:
        """模型前向传播"""
        outputs = model(input_ids=input_ids, **kwargs)
        return {
            'logits': outputs.logits
        }

class LLaMAAdapter(ModelAdapter):
    """LLaMA模型适配器"""
    def load_model(self, model_path: str) -> torch.nn.Module:
        # LLaMA特定的加载逻辑
        return super().load_model(model_path)

class MoEAdapter(ModelAdapter):
    """MoE模型适配器"""
    def forward(self, model: torch.nn.Module, input_ids: torch.Tensor, **kwargs) -> Dict[str, Any]:
        # MoE特定的前向传播逻辑
        outputs = model(input_ids=input_ids, **kwargs)
        return {
            'logits': outputs.logits,
            'router_logits': outputs.router_logits
        }

class MultiModalAdapter(ModelAdapter):
    """多模态模型适配器"""
    def forward(self, model: torch.nn.Module, input_ids: torch.Tensor, image_inputs: torch.Tensor, **kwargs) -> Dict[str, Any]:
        # 多模态特定的前向传播逻辑
        outputs = model(input_ids=input_ids, image_inputs=image_inputs, **kwargs)
        return {
            'logits': outputs.logits
        }

class ModelAdapterRegistry:
    """模型适配器注册表"""
    def __init__(self):
        self.adapters = {}
    
    def register(self, model_type: str, adapter_cls: Callable):
        """注册模型适配器"""
        self.adapters[model_type] = adapter_cls
    
    def get_adapter(self, model_type: str, config: Dict[str, Any]) -> ModelAdapter:
        """获取模型适配器"""
        if model_type in self.adapters:
            return self.adapters[model_type](config)
        return ModelAdapter(config)

# 全局注册表
model_adapter_registry = ModelAdapterRegistry()

# 注册适配器
model_adapter_registry.register('llama', LLaMAAdapter)
model_adapter_registry.register('moe', MoEAdapter)
model_adapter_registry.register('multimodal', MultiModalAdapter)

from vllm.model_runner import ModelRunner

# 配置
config = {
    'model': 'meta-llama/Llama-2-70b-hf',
    'dtype': 'float16',
    'compile': True,
    'compile_backend': 'inductor',
    'tensor_parallel_size': 8,
    'trust_remote_code': True,
    'gpu_memory_utilization': 0.9
}

# 初始化ModelRunner
model_runner = ModelRunner(config)

# 加载模型
model_runner.load_model()

# 准备输入
import torch
input_ids = torch.tensor([[1, 2, 3, 4, 5]], device='cuda:0')

# 前向传播
outputs = model_runner.forward(input_ids)
print(f"Logits shape: {outputs['logits'].shape}")

# 生成文本
generate_outputs = model_runner.generate(input_ids, max_new_tokens=100)
print(f"Generated tokens: {generate_outputs['logits'].shape}")