memory-optimization

# 内存优化 关于 CPU 卸载、梯度检查点、内存分析以及使用 bitsandbytes 进行高级内存节省策略的完整指南。 ## 概述 用于适配大模型的内存优化技术： - **量化**：50-75% 节省（见其他文档） - **CPU 卸载**：将权重移至 CPU/磁盘 - **梯度检查点**：以计算换取内存 - **优化器策略**：8-bit、分页优化器 - **混合精度**：FP16/BF16 训练 ## CPU 卸载 ### 基础 CPU 卸载 在不使用时将模型部分移至 CPU RAM。 python from transformers import AutoModelForCausalLM, BitsAndBytesConfig import torch config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16 ) model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-2-70b-hf", quantization_config=config, device_map="auto", # 自动设备放置 max_memory={0: "40GB", "cpu": "100GB"} # 40GB GPU, 100GB CPU ) **工作原理**： - 权重存储在 CPU 上 - 仅在计算需要时移至 GPU - 由 `accelerate` 自动管理 **权衡**：速度慢约 5-10 倍，但能运行更大模型 ### 多 GPU 卸载 跨多个 GPU + CPU 分布： python model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-2-405b-hf", quantization_config=config, device_map="auto", max_memory={ 0: "70GB", # GPU 0 1: "70GB", # GPU 1 2: "70GB", # GPU 2 3: "70GB", # GPU 3 "cpu": "200GB" # CPU RAM } ) **结果**：405B 模型（4-bit = ~200GB）可运行在 4×80GB GPU + CPU 上 ### 磁盘卸载 针对连 CPU RAM 都放不下的模型： python model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-2-405b-hf", quantization_config=config, device_map="auto", offload_folder="./offload", # 磁盘卸载目录 offload_state_dict=True, max_memory={0: "40GB", "cpu": "50GB"} ) **权衡**：极慢（慢约 100 倍），但可行 ### 手动设备映射 用于精确控制： python device_map = { "model.embed_tokens": 0, # GPU 0 "model.layers.0": 0, "model.layers.1": 0, # ... "model.layers.40": 1, # GPU 1 "model.layers.41": 1, # ... "model.layers.79": "cpu", # CPU "model.norm": "cpu", "lm_head": "cpu" } model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-2-70b-hf", quantization_config=config, device