Torch 一些使用的小Tips

Posted by MakiNaruto on Tue, May 17, 2022

为什么使用 contiguous()?

确保切片后的张量在内存中是连续存储的。这是因为切片操作可能会导致张量不再连续存储,而后续的计算(比如计算loss)需要连续存储的张量。

示例:

 1import torch
 2
 3# 创建一个连续的张量
 4x = torch.randn(3, 4)
 5print(x.is_contiguous())  # 输出 True
 6
 7# 转置操作会使张量不连续
 8y = x.T
 9print(y.is_contiguous())  # 输出 False
10
11# 将不连续张量转换为连续张量
12z = y.contiguous()
13print(z.is_contiguous())  # 输出 True

模型加载

限制哪些卡被使用.

1os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
2os.environ["CUDA_VISIBLE_DEVICES"] = "0,3"

1. 使用的卡, 自动分配

1AutoModel.from_pretrained(glm, trust_remote_code=True, device_map='auto')

2. 分别指定哪些层加载到哪个卡上

 11. 加载完模型后查看都有哪些层, 例如
 2print(model.hf_device_map)
 3
 42. 自定义加载显卡
 5device_map = {
 6    'transformer.embedding': 0, 
 7    'transformer.rotary_pos_emb': 0, 
 8    'transformer.encoder.layers.0': 0, 
 9    'transformer.encoder.layers.1': 0, 
10    'transformer.encoder.layers.2': 0, 
11    ...
12    'transformer.encoder.layers.27': 1, 
13    'transformer.encoder.final_layernorm': 1, 
14    'transformer.output_layer': 1
15    }
16
173. 加载模型
18AutoModel.from_pretrained(model, trust_remote_code=True, device_map=device_map)