这里测试的是开源模型是 meta-llama/Llama-3.1-405B-Instruct，属于开源模型里面最大的。
1. 默认方法
代码如下，这个方法可以加载起来模型，但是在推理的时候会非常慢。
|
1 2 3 4 5 6 7 8 9 10 11 12 |
import torch from transformers import AutoModelForCausalLM MODEL_ID = "meta-llama/Llama-3.1-405B-Instruct" print(f"Load Model {MODEL_ID} ... ") model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map="auto", trust_remote_code=True, torch_dtype=torch.bfloat16 ) |
2. 使用 BitsAndBytesConfig
这个方法可以把更多的权重加载到GPU里面，这样速度会快很多，下面是8个24G的GPU内存的内存映射代码。
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
import torch from transformers import AutoModelForCausalLM, BitsAndBytesConfig # 检测可用的GPU数量 NUM_GPUS = torch.cuda.device_count() print(f"NUM_GPUS: {NUM_GPUS}") # 定义Transformer模型的层数 NUM_TRANS_LAYERS = 126 MODEL_ID = "meta-llama/Llama-3.1-405B-Instruct" def create_device_map(): device_map = { 'model.embed_tokens': 0, 'model.norm': 0, 'model.rotary_emb': 0, 'lm_head': 0 } # 根据 GPU 数量配置设备映射 if NUM_GPUS > 0: for start, end, gpu_id in [(0, 6, 0), (6, 18, 1), (18, 30, 2), (30, 42, 3), (42, 54, 4), (54, 66, 5), (66, 78, 6), (78, 90, 7)]: for i in range(start, end): device_map[f'model.layers.{i}'] = gpu_id for i in range(90, NUM_TRANS_LAYERS): device_map[f'model.layers.{i}'] = "cpu" # 使用 CPU 处理剩余层 return device_map device_map = create_device_map() if NUM_GPUS > 0 else None # 模型加载 print(f"Loading Model {MODEL_ID} ...") model = AutoModelForCausalLM.from_pretrained( MODEL_ID, device_map=device_map, trust_remote_code=True, torch_dtype=torch.bfloat16, quantization_config=BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, llm_int8_enable_fp32_cpu_offload=True # 启用 CPU 辅助 ) ) |