LLM预训练

网上大量预训练代码都是封装了trainer-deepspeed后的结果，看了也不了解其中所用技术的优化点在哪。本文从最基础的训练过程开始，层层加码并对比。

基础版本

1.代码

from transformers import AutoModel,AutoTokenizer
from torch.utils.data import Dataset, DataLoader
# 模型加载
model_path = "xxx/glm2" # 你的模型路径
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True,torch_dtype=torch.float32,device_map="auto")

# 使用一半参数测试
for k,v in model.named_parameters():
    no_grad_layer = list(range(0,27,2))
    no_grad_layer = [str(i) for i in no_grad_layer]
    name_list = k.split(".")
    if len(name_list) >= 4 and name_list[3] in no_grad_layer:
        v.requires_grad=False

# dataset
example = ["text1","text2","text3"]
class GenerateDataset(Dataset):
    def __init__(self, example):
        super(GenerateDataset, self).__init__()
        self.example = example

    def __getitem__(self, item):
        return self.example[item]

    def __len__(self):
        return len(self.example)
ds = GenerateDataset(example)

# dataloader+coll_fn
def pretrain_fn(context):
    input_ids = []
    label_ids = []
    max_len = 2000
    context_ids = tokenizer(context,add_special_tokens=False)["input_ids"]
    target_ids = tokenizer(context,add_special_tokens=False)["input_ids"]
    for c_ids, t_ids in zip(context_ids, target_ids):
        length = len(c_ids)
        if length >= max_len:
            c_ids = c_ids[:max_len - 1] + [2]
            t_ids = t_ids[:max_len - 1] + [2]
        else:
            c_ids = c_ids + [2] * (max_len - length)
            t_ids = t_ids + [-100] * (max_len - length)
        input_ids.append(c_ids)
        label_ids.append(t_ids)
    return {"input_ids":torch.LongTensor(input_ids),"labels":torch.LongTensor(label_ids)}
    
dl = DataLoader(ds, batch_size=1,collate_fn=pretrain_fn)
# 优化器
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
# 训练
for batch in tqdm(dl):
    optimizer.zero_grad()
    out=model(input_ids=batch["input_ids"], labels=batch["labels"])
    loss = out.loss
    print(loss)
    loss.backward()
    optimizer.step()
    # print(loss.item)

2.有几个注意事项：

1.torch_dtype=torch.float32，如果用了fp16，模型参数和损失直接变为nan，训练过程出错。用accelerate和trainer会报ValueError: Attempting to unscale FP16 gradients.

  0%|                                                  | 0/1615 [00:00<?, ?it/s]tensor(11.4844, dtype=torch.float16, grad_fn=<ToCopyBackward0>)
  0%|                                        | 1/1615 [00:04<1:54:59,  4.27s/it]tensor(nan, dtype=torch.float16, grad_fn=<ToCopyBackward0>)
  0%|                                          | 2/1615 [00:04<50:53,  1.89s/it]tensor(nan, dtype=torch.float16, grad_fn=<ToCopyBackward0>)

ValueError: Attempting to unscale FP16 gradients.

2.优化器SGD/Adamw,6B模型全参数训练时候，如果用Adamw会爆显存（batch=1,文本长度20），用SGD则不会。实际llm没看到用SGD做优化器，这里只是提一嘴。

Accelerate 版本

# 加载model optimizer dataloader 后，加入以下代码
from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp16")
device = accelerator.device
model, optimizer, dl = accelerator.prepare(model, optimizer, dl)
# 原来的loss.backward()替换为accelerator.backward(loss)
# loss.backward()
accelerator.backward(loss)

Trainer 版本

原代码的dataloader optimizer 和训练代码都不用了

train_args = TrainingArguments(save_strategy="epoch",
                               log_level="debug",
                               output_dir="./saved_checkpoint/",
                               per_device_train_batch_size=4,  
                               gradient_accumulation_steps=1,  # 梯度累积步数
                               num_train_epochs=1,
                               learning_rate=1e-5,
                               fp16=True, # 是否使用fp16
                               logging_steps=1,
                               warmup_steps=50,
                               optim="adamw_hf",  # 指定优化器
                               gradient_checkpointing=True, # 梯度检查点
                               )
trainer = Trainer(model=model,train_dataset=ds,args=train_args, data_collator=pretrain_fn)
trainer.train()

需要注意的时：如果gradient_checkpointing=True，则需要加上下面的代码

model.enable_input_require_grads()
# 或者 直接修改模型embedding参数的requires_grad
for k,v in model.named_parameters():
	if k == "transformer.embedding.word_embeddings.weight":
		v.requires_grad = True

Trainer + DeepSpeed

修改train_args, 其他不变

train_args = TrainingArguments(save_strategy="steps",
                               save_steps=40,
                               log_level="debug",
                               output_dir="./test_lora/",
                               per_device_train_batch_size=2,
                               gradient_accumulation_steps=2,
                               num_train_epochs=10,
                               learning_rate=1e-6,
                               fp16=True,
                               logging_steps=1,
                               warmup_steps=10,
                               optim="adamw_hf",
                               gradient_checkpointing=True,
                               deepspeed="./notebooks/deepspeed-stage2.json" # 加入ds路径
                               )

stage-2参数

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 20,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

stage-3参数

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 20,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

启动命令

deepspeed pretrain.py --seepspeed ./notebooks/deepspeed-stage2.json