4.2-CPU Offload(1)

4.2 CPU Offload与Pipeline并行

摘要

CPU Offload与Pipeline并行是应对GPU显存限制的关键技术。通过将部分计算任务卸载到CPU内存、构建计算流水线，可以实现大模型的高效推理和训练。本章将深入剖析CPU Offload的实现机制、Pipeline并行的调度策略、内存优化技术，以及在实际应用中的性能优化方法，为构建高效的大模型系统提供全面的技术指导。

1. CPU Offload基础理论

1.1 CPU Offload的基本概念

CPU Offload是指将原本在GPU上运行的部分计算任务转移到CPU上执行，以缓解GPU显存压力的技术。

import torch
import torch.nn as nn
import torch.nn.functional as F
import threading
import queue
import time
from typing import Dict, List, Optional, Tuple

class CPUOffloadConfig:
    """CPU Offload配置"""
    def __init__(self, 
                 offload_threshold: int = 1024*1024*1024,  # 1GB
                 cpu_workers: int = 4,
                 batch_size: int = 32,
                 use_pipeline: bool = True):
        self.offload_threshold = offload_threshold
        self.cpu_workers = cpu_workers
        self.batch_size = batch_size
        self.use_pipeline = use_pipeline
​

1.2 CPU与GPU的协同计算

class CPUGPUCoordinator:
    def __init__(self, model, config: CPUOffloadConfig):
        self.model = model
        self.config = config
        self.device = next(model.parameters()).device
        self.cpu_device = torch.device('cpu')
        
        # 任务队列
        self.task_queue = queue.Queue()
        self.result_queue = queue.Queue()
        
        # 工作线程
        self.workers = []
        self._start_workers()
        
    def _start_workers(self):
        """启动工作线程"""
        for i in range(self.config.cpu_workers):
            worker = threading.Thread(
                target=self._worker_loop,
                daemon=True,
                name=f"CPU-Worker-{i}"
            )
            worker.start()
            self.workers.append(worker)
    
    def _worker_loop(self):
        """工作线程循环"""
        while True:
            # 获取任务
            task = self.task_queue.get()
            
            if task is None:  # 停止信号
                break
            
            try:
                # 执行任务
                result = self._execute_task(task)
                
                # 返回结果
                self.result_queue.put(result)
                
            except Exception as e:
                self.result_queue.put(('error', str(e)))
            
            finally:
                self.task_queue.task_done()
    
    def _execute_task(self, task):
        """执行任务"""
        task_type, args = task
        
        if task_type == 'forward':
            # CPU前向传播
            data = args['data']
            return self._cpu_forward(data)
        elif task_type == 'backward':
            # CPU反向传播
            loss = args['loss']
            return self._cpu_backward(loss)
        elif task_type == 'offload':
            # CPU卸载任务
            tensor = args['tensor']
            return self._offload_to_cpu(tensor)
        else:
            raise ValueError(f"Unknown task type: {task_type}")
    
    def _cpu_forward(self, data):
        """CPU前向传播"""
        # 将数据移动到CPU
        data_cpu = data.to(self.cpu_device)
        
        # 在CPU上执行前向传播
        with torch.no_grad():
            output = self.model(data_cpu)
        
        return output
    
    def _cpu_backward(self, loss):
        """CPU反向传播"""
        # 在CPU上执行反向传播
        loss.backward()
        return None
    
    def _offload_to_cpu(self, tensor):
        """卸载张量到CPU"""
        return tensor.to(self.cpu_device)
​

1.3 内存管理策略

class MemoryManager:
    def __init__(self, config: CPUOffloadConfig):
        self.config = config
        self.gpu_memory_used = 0
        self.cpu_memory_used = 0
        self.gpu_memory_limit = torch.cuda.get_device_properties(0).total_memory
        self.cpu_memory_limit = psutil.virtual_memory().total
        
        # 内存池
        self.gpu_pool = {}
        self.cpu_pool = {}
        
    def should_offload(self, tensor: torch.Tensor) -> bool:
        """判断是否需要卸载到CPU"""
        tensor_size = tensor.numel() * tensor.element_size()
        
        # 检查GPU内存使用情况
        gpu_usage = self.gpu_memory_used + tensor_size
        if gpu_usage > self.gpu_memory_limit * 0.9:  # 90%阈值
            return True
        
        # 检查GPU碎片化情况
        if self._check_gpu_fragmentation():
            return True
        
        # 检查单个张量大小
        if tensor_size > self.config.offload_threshold:
            return True
        
        return False
    
    def offload_tensor(self, tensor: torch.Tensor) -> torch.Tensor:
        """卸载张量到CPU"""
        if self.should_offload(tensor):
            tensor_cpu = tensor.cpu()
            self.gpu_memory_used -= tensor.numel() * tensor.element_size()
            self.cpu_memory_used += tensor_cpu.numel() * tensor_cpu.element_size()
            return tensor_cpu
        else:
            return tensor
    
    def _check_gpu_fragmentation(self) -> bool:
        """检查GPU碎片化"""
        # 简化实现：检查GPU内存块数量
        allocated = torch.cuda.memory_allocated()
        reserved = torch.cuda.memory_reserved()
        
        # 如果reserved比allocated大很多，说明有碎片
        fragmentation_ratio = (reserved - allocated) / reserved
        return fragmentation_ratio > 0.2  # 20%阈值
    
    def get_memory_stats(self) -> Dict:
        """获取内存统计"""
        return {
            'gpu_memory_used': self.gpu_memory_used,
            'gpu_memory_limit': self.gpu_memory_limit,
            'gpu_utilization': self.gpu_memory_used / self.gpu_memory_limit,
            'cpu_memory_used': self.cpu_memory_used,
            'cpu_memory_limit': self.cpu_memory_limit,
            'cpu_utilization': self.cpu_memory_used / self.cpu_memory_limit
        }
​

2. Pipeline并行技术

2.1 Pipeline并行基础

class PipelineParallelModel(nn.Module):
    """Pipeline并行模型"""
    def __init__(self, layers: List[nn.Module], config: CPUOffloadConfig):
        super().__init__()
        self.layers = nn.ModuleList(layers)
        self.config = config
        self.pipeline_manager = PipelineManager(len(layers), config)
        
    def forward(self, x):
        """前向传播"""
        return self.pipeline_manager.forward(x, self.layers)
    
    def backward(self, loss):
        """反向传播"""
        return self.pipeline_manager.backward(loss, self.layers)
​

2.2 Pipeline调度器

class PipelineManager:
    def __init__(self, num_layers: int, config: CPUOffloadConfig):
        self.num_layers = num_layers
        self.config = config
        self.current_batch = 0
        self.micro_batch_size = config.batch_size
        self.current_micro_batch = 0
        
        # Pipeline状态
        self.stage_outputs = {}
        self.stage_inputs = {}
        self.gradient_accumulation = {}
        
        # 同步锁
        self.locks = [threading.Lock() for _ in range(num_layers)]
        
    def forward(self, x, layers):
        """Pipeline前向传播"""
        # 准备微批次
        micro_batches = self._split_micro_batches(x)
        
        results = []
        
        for micro_batch in micro_batches:
            result = self._process_micro_batch(micro_batch, layers)
            results.append(result)
        
        # 合并结果
        return torch.cat(results, dim=0)
    
    def _split_micro_batches(self, x):
        """分割为微批次"""
        batch_size = x.size(0)
        num_micro_batches = (batch_size + self.micro_batch_size - 1) // self.micro_batch_size
        
        micro_batches = []
        for i in range(num_micro_batches):
            start = i * self.micro_batch_size
            end = min((i + 1) * self.micro_batch_size, batch_size)
            micro_batches.append(x[start:end])
        
        return micro_batches
    
    def _process_micro_batch(self, micro_batch, layers):
        """处理单个微批次"""
        # 1. 第一层
        with self.locks[0]:
            stage0_output = layers[0](micro_batch)
            self.stage_outputs[0] = stage0_output
            
            # 如果不是最后阶段，准备下一阶段的输入
            if self.num_layers > 1:
                self.stage_inputs[1] = stage0_output
        
        # 2. 中间层
        for stage in range(1, self.num_layers - 1):
            with self.locks[stage]:
                # 等待前一阶段完成
                while stage not in self.stage_outputs:
                    time.sleep(0.001)
                
                # 执行当前阶段
                stage_input = self.stage_outputs[stage-1]
                stage_output = layers[stage](stage_input)
                self.stage_outputs[stage] = stage_output
                
                # 准备下一阶段输入
                if stage + 1 < self.num_layers:
                    self.stage_inputs[stage + 1] = stage_output
                
                # 清除前一阶段输出
                del self.stage_outputs[stage-1]
        
        # 3. 最后阶段
        with self.locks[self.num_layers - 1]:
            if self.num_layers == 1:
                final_output = layers[0](micro_batch)
            else:
                final_output = layers[self.num_layers - 1](self.stage_outputs[self.num_layers - 2])
            
            # 清除所有中间输出
            self.stage_outputs.clear()
        
        return final_output
​

2.3 1F1B调度策略

class PipelineScheduler1F1B:
    """1F1B (One Forward, One Backward) Pipeline调度器"""
    def __init__(self, num_stages: int, micro_batch_size: int):
        self.num_stages = num_stages
        self.micro_batch_size = micro_batch_size
        
        # 调度状态
        self.current_micro_batch = 0
        self.current_stage = 0
        
        # 同步队列
        self.forward_queue = queue.Queue()
        self.backward_queue = queue.Queue()
        
        # 阶段状态
        self.stage_states = ['idle'] * num_stages
        
    def schedule_forward(self, micro_batch):
        """调度前向传播"""
        # 添加到调度队列
        self.forward_queue.put((micro_batch, 'forward'))
        
        # 开始调度
        self._start_schedule()
    
    def schedule_backward(self, gradient):
        """调度反向传播"""
        # 添加到调度队列
        self.backward_queue.put((gradient, 'backward'))
        
        # 开始调度
        self._start_schedule()
    
    def _start_schedule(self):
        """开始调度"""
        while not self.forward_queue.empty() or not self.backward_queue.empty():
            # 查找空闲阶段
            free_stage = self._find_free_stage()
            
            if free_stage is not None:
                # 执行调度
                if self._should_forward():
                    self._execute_forward(free_stage)
                else:
                    self._execute_backward(free_stage)
            else:
                # 没有空闲阶段，等待
                time.sleep(0.001)
    
    def _find_free_stage(self):
        """查找空闲阶段"""
        for i in range(self.num_stages):
            if self.stage_states[i] == 'idle':
                return i
        return None
    
    def _should_forward(self):
        """判断是否应该执行前向传播"""
        if self.forward_queue.empty():
            return False
        
        # 检查是否有足够的空间执行前向传播
        forward_count = sum(1 for state in self.stage_states if state == 'forward')
        backward_count = sum(1 for state in self.stage_states if state == 'backward')
        
        # 保持1F1B平衡
        return forward_count <= backward_count + 1
    
    def _execute_forward(self, stage):
        """执行前向传播"""
        self.stage_states[stage] = 'forward'
        
        # 从队列获取任务
        micro_batch, task_type = self.forward_queue.get()
        
        try:
            # 执行前向传播
            result = self._forward_stage(stage, micro_batch)
            
            # 完成后更新状态
            self.stage_states[stage] = 'forward_complete'
            
        except Exception as e:
            print(f"Forward error in stage {stage}: {e}")
            self.stage_states[stage] = 'idle'
    
    def _execute_backward(self, stage):
        """执行反向传播"""
        self.stage_states[stage] = 'backward'
        
        # 从队列获取任务
        gradient, task_type = self.backward_queue.get()
        
        try:
            # 执行反向传播
            self._backward_stage(stage, gradient)
            
            # 完成后更新状态
            self.stage_states[stage] = 'backward_complete'
            
        except Exception as e:
            print(f"Backward error in stage {stage}: {e}")
            self.stage_states[stage] = 'idle'
​

3. CPU Offload实现细节

3.1 异步CPU计算

class AsyncCPUExecutor:
    """异步CPU执行器"""
    def __init__(self, num_workers: int = 4):
        self.num_workers = num_workers
        self.task_queue = queue.Queue()
        self.result_queue = queue.Queue()
        self.workers = []
        
        # 启动工作线程
        self._start_workers()
    
    def _start_workers(self):
        """启动工作线程"""
        for i in range(self.num_workers):
            worker = threading.Thread(
                target=self._worker_loop,
                daemon=True,
                name=f"CPU-Executor-{i}"
            )
            worker.start()
            self.workers.append(worker)
    
    def _worker_loop(self):
        """工作线程循环"""
        while True:
            task = self.task_queue.get()
            
            if task is None:  # 停止信号
                break
            
            try:
                task_type, args, callback = task
                
                # 执行任务
                result = self._execute_task(task_type, args)
                
                # 调用回调
                if callback:
                    callback(result)
                    
            except Exception as e:
                print(f"Task execution failed: {e}")
            
            finally:
                self.task_queue.task_done()
    
    def _execute_task(self, task_type: str, args: dict):
        """执行任务"""
        if task_type == 'forward':
            return self._cpu_forward(args['model'], args['data'])
        elif task_type == 'backward':
            return self._cpu_backward(args['loss'])
        elif task_type == 'matmul':
            return self._cpu_matmul(args['A'], args['B'])
        else:
            raise ValueError(f"Unknown task type: {task_type}")
    
    def _cpu_forward(self, model, data):
        """CPU前向传播"""
        model.eval()  # 确保模型在评估模式
        with torch.no_grad():
            output = model(data)
        return output
    
    def _cpu_backward(self, loss):
        """CPU反向传播"""
        loss.backward(retain_graph=True)
        return None
    
    def _cpu_matmul(self, A, B):
        """CPU矩阵乘法"""
        return torch.matmul(A, B)
    
    def submit_task(self, task_type: str, args: dict, callback=None):
        """提交任务"""
        task = (task_type, args, callback)
        self.task_queue.put(task)
        return task
    
    def shutdown(self):
        """关闭执行器"""
        for _ in range(self.num_workers):
            self.task_queue.put(None)
        
        for worker in self.workers:
            worker.join()
​

3.2 内存传输优化

class MemoryTransferOptimizer:
    """内存传输优化器"""
    def __init__(self):
        self.transfer_stats = {
            'gpu_to_cpu_count': 0,
            'cpu_to_gpu_count': 0,
            'gpu_to_cpu_bytes': 0,
            'cpu_to_gpu_bytes': 0,
            'total_time': 0
        }
        
        # 传输队列
        self.transfer_queue = queue.Queue()
        self.transfer_workers = []
        
        # 启动传输工作线程
        self._start_transfer_workers()
    
    def _start_transfer_workers(self):
        """启动传输工作线程"""
        for i in range(2):  # 2个传输工作线程
            worker = threading.Thread(
                target=self._transfer_worker_loop,
                daemon=True,
                name=f"Transfer-Worker-{i}"
            )
            worker.start()
            self.transfer_workers.append(worker)
    
    def _transfer_worker_loop(self):
        """传输工作线程循环"""
        while True:
            transfer_task = self.transfer_queue.get()
            
            if transfer_task is None:  # 停止信号
                break
            
            try:
                self._execute_transfer(transfer_task)
                
            except Exception as e:
                print(f"Transfer failed: {e}")
​