4.1 GPU显存分配与碎片整理策略

摘要

GPU显存管理是深度学习推理和训练的核心挑战之一。有效的显存分配策略和碎片整理技术能够显著提升硬件利用率、降低显存开销、提高系统稳定性。本章将深入剖析GPU显存管理的底层机制、碎片产生的原因、显存分配算法、碎片整理技术，以及在各种场景下的优化实践，为构建高效、稳定的深度学习系统提供全面的技术指导。

1. GPU显存管理基础

1.1 GPU显存架构概述

现代GPU显存系统主要由以下几个层次组成：

CPU内存 ──┐
          │ PCIe
GPU显存 ──┤
          │ GPU内部总线
          ├── 显存控制器
          ├── 显存分区 (Memory Banks)
          ├── 显存池 (Memory Pool)
          └── 显存缓存 (Cache)
​

import torch
import torch.nn as nn
import pynvml
import psutil

class GPUMemoryInfo:
    def __init__(self, device_id=0):
        self.device_id = device_id
        self._init_nVML()
    
    def _init_nVML(self):
        """初始化NVML库"""
        try:
            pynvml.nvmlInit()
            self.handle = pynvml.nvmlDeviceGetHandleByIndex(self.device_id)
        except:
            self.handle = None
    
    def get_memory_info(self):
        """获取显存信息"""
        if self.handle is None:
            return self._get_pytorch_memory_info()
        
        try:
            # NVML获取显存信息
            memory_info = pynvml.nvmlDeviceGetMemoryInfo(self.handle)
            
            return {
                'total': memory_info.total / (1024**3),  # GB
                'used': memory_info.used / (1024**3),    # GB
                'free': memory_info.free / (1024**3),     # GB
                'utilization': memory_info.used / memory_info.total * 100
            }
        except:
            return self._get_pytorch_memory_info()
    
    def _get_pytorch_memory_info(self):
        """通过PyTorch获取显存信息"""
        if torch.cuda.is_available():
            total_memory = torch.cuda.get_device_properties(self.device_id).total_memory
            used_memory = torch.cuda.memory_allocated(self.device_id)
            free_memory = total_memory - used_memory
            
            return {
                'total': total_memory / (1024**3),
                'used': used_memory / (1024**3),
                'free': free_memory / (1024**3),
                'utilization': used_memory / total_memory * 100
            }
        else:
            return None
    
    def get_memory_banks_info(self):
        """获取显存分区信息"""
        if self.handle is None:
            return None
        
        try:
            # 获取显存分区信息
            memory_clock = pynvml.nvmlDeviceGetMemoryClock(self.handle)
            bus_width = pynvml.nvmlDeviceGetPcieThroughput(self.handle)
            
            # 计算理论显存带宽
            bandwidth = memory_clock * bus_width / 8 / 1024 / 1024  # GB/s
            
            return {
                'memory_clock': memory_clock,
                'bus_width': bus_width,
                'bandwidth': bandwidth,
                'memory_banks': 8  # 常见GPU的显存分区数
            }
        except:
            return None
​

1.2 显存分配机制

PyTorch的显存分配采用了分层策略：

class MemoryAllocator:
    def __init__(self, device='cuda'):
        self.device = device
        self.allocated_blocks = {}
        self.fragmentation_stats = {}
        
    def allocate_tensor(self, size, dtype=torch.float32):
        """分配张量"""
        # 计算实际分配大小（考虑对齐）
        element_size = self._get_dtype_size(dtype)
        aligned_size = self._align_allocation(size * element_size)
        
        # 分配内存
        tensor = torch.empty(size, dtype=dtype, device=self.device)
        
        # 记录分配信息
        block_id = id(tensor)
        self.allocated_blocks[block_id] = {
            'tensor': tensor,
            'size': aligned_size,
            'dtype': dtype,
            'allocation_time': time.time()
        }
        
        return tensor
    
    def _get_dtype_size(self, dtype):
        """获取数据类型大小"""
        dtype_sizes = {
            torch.float32: 4,
            torch.float16: 2,
            torch.int8: 1,
            torch.int32: 4,
            torch.int64: 8
        }
        return dtype_sizes.get(dtype, 4)
    
    def _align_allocation(self, size, alignment=512):
        """对齐内存分配"""
        return (size + alignment - 1) // alignment * alignment
    
    def get_memory_usage(self):
        """获取内存使用情况"""
        total_allocated = sum(block['size'] for block in self.allocated_blocks.values())
        total_blocks = len(self.allocated_blocks)
        
        return {
            'total_allocated': total_allocated,
            'total_blocks': total_blocks,
            'average_block_size': total_allocated / total_blocks if total_blocks > 0 else 0
        }
​

1.3 显存碎片的概念与成因

显存碎片是指显存中存在大量不连续的小块内存，导致无法满足大块内存的分配需求。

class FragmentationAnalyzer:
    def __init__(self, allocator):
        self.allocator = allocator
        
    def analyze_fragmentation(self):
        """分析碎片化情况"""
        allocated_blocks = self.allocator.allocated_blocks
        total_memory = self._get_total_memory()
        
        if not allocated_blocks:
            return {'fragmentation_ratio': 0, 'largest_free_block': total_memory}
        
        # 计算已分配内存
        allocated_sizes = [block['size'] for block in allocated_blocks.values()]
        total_allocated = sum(allocated_sizes)
        
        # 模拟最大连续空闲块
        largest_free_block = self._estimate_largest_free_block(allocated_sizes, total_memory)
        
        # 计算碎片率
        fragmentation_ratio = 1 - (largest_free_block / (total_memory - total_allocated))
        
        return {
            'fragmentation_ratio': fragmentation_ratio,
            'largest_free_block': largest_free_block,
            'total_memory': total_memory,
            'total_allocated': total_allocated,
            'fragmentation_level': self._classify_fragmentation_level(fragmentation_ratio)
        }
    
    def _estimate_largest_free_block(self, allocated_sizes, total_memory):
        """估计最大空闲块大小"""
        # 简化估算：假设碎片主要由小块分配造成
        avg_allocation_size = sum(allocated_sizes) / len(allocated_sizes)
        fragmentation_overhead = len(allocated_sizes) * avg_allocation_size * 0.1  # 10%开销
        largest_free = (total_memory - sum(allocated_sizes)) - fragmentation_overhead
        return max(0, largest_free)
    
    def _classify_fragmentation_level(self, ratio):
        """分类碎片化程度"""
        if ratio < 0.1:
            return '低'
        elif ratio < 0.3:
            return '中'
        elif ratio < 0.6:
            return '高'
        else:
            return '严重'
    
    def _get_total_memory(self):
        """获取总显存"""
        try:
            return torch.cuda.get_device_properties(0).total_memory
        except:
            return 8 * 1024**3  # 默认8GB
​

2. 显存分配算法

2.1 首次适应算法（First Fit）

class FirstFitAllocator:
    def __init__(self, total_memory):
        self.total_memory = total_memory
        self.free_blocks = [(0, total_memory)]  # (start, size)
        self.allocated_blocks = {}
        
    def allocate(self, size, alignment=512):
        """首次适应分配"""
        # 对齐需求大小
        aligned_size = (size + alignment - 1) // alignment * alignment
        
        # 查找第一个合适的空闲块
        for i, (start, block_size) in enumerate(self.free_blocks):
            if block_size >= aligned_size:
                # 分配内存
                allocated_start = start
                allocated_end = start + aligned_size
                
                # 更新空闲块
                remaining_size = block_size - aligned_size
                if remaining_size > 0:
                    self.free_blocks[i] = (allocated_end, remaining_size)
                else:
                    del self.free_blocks[i]
                
                # 记录分配
                block_id = id((allocated_start, aligned_size))
                self.allocated_blocks[block_id] = {
                    'start': allocated_start,
                    'size': aligned_size,
                    'actual_size': size
                }
                
                return allocated_start, aligned_size
        
        # 分配失败
        return None, None
    
    def deallocate(self, start, size):
        """释放内存"""
        # 查找对应的分配块
        block_to_remove = None
        for block_id, block in self.allocated_blocks.items():
            if block['start'] == start:
                block_to_remove = block_id
                break
        
        if block_to_remove is not None:
            # 从已分配中移除
            allocated_size = self.allocated_blocks[block_to_remove]['size']
            del self.allocated_blocks[block_to_remove]
            
            # 插入空闲块列表
            self._insert_free_block(start, allocated_size)
            
            # 合并相邻空闲块
            self._merge_adjacent_blocks()
    
    def _insert_free_block(self, start, size):
        """插入空闲块"""
        # 按起始地址排序插入
        for i, (current_start, current_size) in enumerate(self.free_blocks):
            if start < current_start:
                self.free_blocks.insert(i, (start, size))
                return
        
        self.free_blocks.append((start, size))
    
    def _merge_adjacent_blocks(self):
        """合并相邻空闲块"""
        if not self.free_blocks:
            return
        
        # 按起始地址排序
        self.free_blocks.sort()
        
        merged_blocks = []
        current_start, current_size = self.free_blocks[0]
        
        for start, size in self.free_blocks[1:]:
            if start == current_start + current_size:
                # 相邻块，合并
                current_size += size
            else:
                # 不相邻，保存当前块
                merged_blocks.append((current_start, current_size))
                current_start, current_size = start, size
        
        merged_blocks.append((current_start, current_size))
        self.free_blocks = merged_blocks
​

2.2 最佳适应算法（Best Fit）

class BestFitAllocator:
    def __init__(self, total_memory):
        self.total_memory = total_memory
        self.free_blocks = [(0, total_memory)]
        self.allocated_blocks = {}
    
    def allocate(self, size, alignment=512):
        """最佳适应分配"""
        aligned_size = (size + alignment - 1) // alignment * alignment
        
        best_fit_index = -1
        best_fit_size = float('inf')
        
        # 查找最适合的块（最小的大于等于需求大小的块）
        for i, (start, block_size) in enumerate(self.free_blocks):
            if block_size >= aligned_size and block_size < best_fit_size:
                best_fit_index = i
                best_fit_size = block_size
        
        if best_fit_index != -1:
            # 找到最佳适配
            start, block_size = self.free_blocks[best_fit_index]
            allocated_start = start
            allocated_size = aligned_size
            
            # 更新空闲块
            remaining_size = block_size - allocated_size
            if remaining_size > 0:
                self.free_blocks[best_fit_index] = (start + allocated_size, remaining_size)
            else:
                del self.free_blocks[best_fit_index]
            
            # 记录分配
            block_id = id((allocated_start, allocated_size))
            self.allocated_blocks[block_id] = {
                'start': allocated_start,
                'size': allocated_size,
                'actual_size': size
            }
            
            return allocated_start, allocated_size
        
        return None, None
​

2.3 伙伴系统算法

class BuddySystemAllocator:
    def __init__(self, total_memory):
        self.total_memory = total_memory
        self.max_order = self._calculate_max_order(total_memory)
        self.memory_blocks = {}
        
        # 初始化：只有一个最大块
        self.memory_blocks[0] = {
            'order': self.max_order,
            'size': total_memory,
            'allocated': False,
            'buddy': None
        }
    
    def _calculate_max_order(self, size):
        """计算最大阶数"""
        order = 0
        while (1 << order) < size:
            order += 1
        return order
    
    def allocate(self, size):
        """伙伴系统分配"""
        required_order = self._calculate_order(size)
        
        if required_order > self.max_order:
            return None
        
        # 查找合适大小的块
        block = self._find_block(required_order)
        
        if block is None:
            # 没有合适大小的块，尝试分裂更大的块
            block = self._split_blocks(required_order)
        
        if block is not None:
            block['allocated'] = True
            return block
        
        return None
    
    def _calculate_order(self, size):
        """计算需求的阶数"""
        order = 0
        while (1 << order) < size:
            order += 1
        return order
    
    def _find_block(self, order):
        """查找指定阶数的块"""
        for block_id, block in self.memory_blocks.items():
            if (not block['allocated'] and 
                block['order'] == order):
                return block
        return None
    
    def _split_blocks(self, target_order):
        """分裂块"""
        current_order = target_order + 1
        
        while current_order <= self.max_order:
            # 查找当前阶数的块
            block = self._find_block(current_order)
            
            if block is not None:
                # 分裂成两个伙伴块
                new_order = current_order - 1
                new_size = 1 << new_order
                
                # 创建伙伴块
                buddy_id = len(self.memory_blocks)
                buddy_block = {
                    'order': new_order,
                    'size': new_size,
                    'allocated': False,
                    'buddy': block_id
                }
                
                # 更新原块
                block['order'] = new_order
                block['size'] = new_size
                block['buddy'] = buddy_id
                
                # 添加伙伴块
                self.memory_blocks[buddy_id] = buddy_block
                
                # 如果目标阶数匹配，返回其中一个块
                if new_order == target_order:
                    return block
                
                # 继续分裂
                current_order = new_order
            else:
                # 当前阶数没有可用块，尝试更大阶数
                current_order += 1
        
        return None
​

2.4 分级内存分配器

class TieredMemoryAllocator:
    def __init__(self, device='cuda'):
        self.device = device
        self.tiers = {
            'small': {'min_size': 1, 'max_size': 1024, 'allocator': FirstFitAllocator(8*1024**3)},
            'medium': {'min_size': 1024, 'max_size': 1024*1024, 'allocator': BestFitAllocator(4*1024**3)},
            'large': {'min_size': 1024*1024, 'max_size': float('inf'), 'allocator': BuddySystemAllocator(16*1024**3)}
        }
        
        self.allocation_stats = {
            'small': {'count': 0, 'total_size': 0},
            'medium': {'count': 0, 'total_size': 0},
            'large': {'count': 0, 'total_size': 0}
        }
    
    def allocate(self, size):
        """分级分配"""
        # 根据大小选择合适的分配器
        if size <= self.tiers['small']['max_size']:
            tier = 'small'
        elif size <= self.tiers['medium']['max_size']:
            tier = 'medium'
        else:
            tier = 'large'
        
        # 使用对应分配器
        start, allocated_size = self.tiers[tier]['allocator'].allocate(size)
        
        if start is not None:
            # 记录分配统计
            self.allocation_stats[tier]['count'] += 1
            self.allocation_stats[tier]['total_size'] += allocated_size
            
            return start, allocated_size, tier
        
        return None, None, None
    
    def deallocate(self, start, size, tier):
        """释放内存"""
        if tier in self.tiers:
            self.tiers[tier]['allocator'].deallocate(start, size)
    
    def get_fragmentation_stats(self):
        """获取碎片统计"""
        stats = {}
        for tier_name, tier_info in self.tiers.items():
            allocator = tier_info['allocator']
            if hasattr(allocator, 'analyze_fragmentation'):
                stats[tier_name] = allocator.analyze_fragmentation()
        
        return stats
    
    def get_allocation_stats(self):
        """获取分配统计"""
        return self.allocation_stats.copy()
​

3. 显存碎片整理技术

3.1 内存压缩算法

class MemoryCompactor:
    def __init__(self, total_memory):
        self.total_memory = total_memory
        self.allocated_blocks = {}
        self.free_blocks = [(0, total_memory)]
        
    def compact_memory(self):
        """内存压缩"""
        if not self.allocated_blocks:
            return
        
        # 按地址排序已分配块
        sorted_blocks = sorted(self.allocated_blocks.values(), 
                             key=lambda x: x['start'])
        
        # 重新分配内存
        new_start = 0
        for block in sorted_blocks:
            old_start = block['start']
​