训练 GPT-neox (1) ----环境配置

本章节仅 针对于训练环境配置进行记录与说明，对于训练中的工具使用，demo，优化，将在后续内容讨论。

1. EC2 的选择

对于小规模的GPT-neox，博主最开始是以省钱为目的出发的，因此选择的是 Nvidia T4 GPU，即 AWS g4dn系列的机器，即可支持自定义小规模的transformer模型。本篇章后续的所有依赖，也都基于 AWS g4dn的机器。（注意T4不支持 Flash-Attention 加速，Flash-Attention 仅支持 Ampere 架构)。 在部署环境测试的时候，可以选择最小的g4dn.xlarge，后续训练有需要，再进行扩容。

2.OS的选择

与 Llama3部署环境 相同，不建议大家 使用不含 驱动的 AWS linux，因为需要解决的冲突较多，博主继续选择了与部署Llama3 一样的AMI： Deep Learning OSS Nvidia Driver AMI GPU TensorFlow 2.16 (Amazon Linux 2) 20240509

3. python 依赖安装

注意！！！与Llama3部署环境 不同，由于训练时需要使用 apex 框架，因此对cuda和torch的版本有要求，默认pip 安装的 torch 版本 为 2.3.0，版本过高，如果使用该版本，则需要重新安装 cuda。因此博主使用匹配 当前os的torch 版本。

pip install torch torchvision torchaudio accelerate -i https://pypi.tuna.tsinghua.edu.cn/simple

除此以外，需要 安装python开发的工具包，以后续编译 apex。

sudo yum install python3-devel

4. apex 安装

Nvidia apex 是一个强大的工具包，可以进行混合精度训练与分布式训练等，显著的提高计算效率，GPT-neox也依赖于 apex。

需要注意，安装 apex 请通过官网下载 https://github.com/NVIDIA/apex

不要 使用 pip install（因为pip装上的并不是你以为的apex）。

a) clone repo到服务器 git clone https://github.com/NVIDIA/apex

b) 修改 setup.py 中，对 cuda 版本与 torch.version.cuda 版本的校验逻辑，注释 setup.py 的第 178 行。不然则会raise版本不匹配的错误，编译时退出。

c) 根据 “apex/README.md” 中指导安装 apex pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./

4. GPT-neox 依赖与环境修改

在配置你自己的模型并进行训练前，需要对gpt neox中的极小一部分 src code进行修改，以便后续编译。

a) clone repo 到服务器 https://github.com/EleutherAI/gpt-neox

b) 安装依赖 pip install -r requirements/requirements.txt， 这里不要使用国内的镜像源，因为requirement里面需要下载 DeepSpeed 的repo。如果出现网络问题，多次尝试即可。（泪目，国内开发者伤不起）

c) 修改 gpt-neox/megatron/data 中的 Makefile

我们需要 data 文件下的 src 帮我们处理自己的训练数据，需要用到里面的helpers.cpp。而当前 EC2环境下，python3-config 并不是我们需要的 python3.10版本的解释器配置，python-config才是，因此替换python3-config为python-config。

修改后可以 make 尝试编译，然后使用python3 import helpers，如果不报错，则修改成功。

5. 训练数据准备并 Pretokenize

训练的数据集可以是任何数据集，但是需要自己进行处理，process代码已经在 gpt neox的repo里面了。需要注意 a 和 b 步骤需要 连接 huggingface，而CN EC2是不行的，所以博主是在本地获取后，scp 到了EC2上。为啥我们要 pre tokenize 数据呢，还不是因为 CN 的环境不能连接 hugging face。

a) 获取训练数据，博主使用了 wiki的数据进行测试。可以直接使用以下代码获取。

from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-103-raw-v1")

with open("wikitext_train.txt", "w") as f:
    for line in dataset["train"]:
        f.write(line["text"] + "\n")

with open("wikitext_validation.txt", "w") as f:
    for line in dataset["validation"]:
        f.write(line["text"] + "\n")

with open("wikitext_test.txt", "w") as f:
    for line in dataset["test"]:
        f.write(line["text"] + "\n")

b）获取词汇表文件和合并文件

这取决于你使用什么 Tokenizer，博主使用 GPT2Tokenizer，可以直接使用以下代码获取 vocab.json 和 merges.txt。

from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")

tokenizer.save_vocabulary(".")

c) Pretokenize 数据，README里有哦

python3 tools/datasets/preprocess_data.py --input ./tiashi/wikitext_train.txt --output-prefix ./tiashi/mydataset --vocab ./tiashi/vocab.json --merge-file ./tiashi/merges.txt --dataset-impl mmap --tokenizer-type GPT2BPETokenizer --append-eod

会生成 mydataset_text_document.bin 与 mydataset_text_document.idx 两个文件，在后续 模型的config中，使用

"data-path": "./tiashi/mydataset_text_document"

6. 配置你的模型

模型的参数配置部分后续会单独开一个article进行讲解，本章只对环境和流程进行详述，给出博主使用的一个17M 参数的用例：

{
  "pipe_parallel_size": 1,
  "model_parallel_size": 1,

  # model settings
  "num_layers": 6,
  "hidden_size": 512,
  "num_attention_heads": 8,
  "seq_length": 2048,
  "max_position_embeddings": 2048,
  "pos_emb": "rotary",
  "no_weight_tying": true,
  "gpt_j_residual": false,
  "output_layer_parallelism": "column",

    #"attention_config": [[["flash"], 6]],

  "scaled_upper_triang_masked_softmax_fusion": false,
  "bias_gelu_fusion": false,
  "rope_fusion": false,
  "layernorm_fusion": false,
  
  "vocab_file": "./tiashi/vocab.json",
  "save": "./tiashi/tom-19M_checkpoints",
  "load": "./tiashi/tom-19M_checkpoints",

  "merge_file": "./tiashi/merges.txt", 
  "data-path": "./tiashi/mydataset_text_document",
  # init methods
  "init_method": "small_init",
  "output_layer_init_method": "wang_init",

  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001,
      "betas": [0.9, 0.95],
      "eps": 1.0e-8,
    }
  },
  "min_lr": 0.0001,

  # for all zero_optimization options, see https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training
  "zero_optimization": {
    "stage": 1,
    "allgather_partitions": True,
    "allgather_bucket_size": 500000000,
    "overlap_comm": True,
    "reduce_scatter": True,
    "reduce_bucket_size": 500000000,
    "contiguous_gradients": True,
  },

  "train_micro_batch_size_per_gpu": 20, #32,
  "gradient_accumulation_steps": 1,
  "data_impl": "mmap",
  "num_workers": 1,

  # activation checkpointing
  "checkpoint_activations": true,
  "checkpoint_num_layers": 1,
  "partition_activations": true,
  "synchronize_each_layer": true,

  # regularization
  "gradient_clipping": 1.0,
  "weight_decay": 0.1,
  "hidden_dropout": 0,
  "attention_dropout": 0,

  # precision settings
  "fp16": {
    "fp16": true,
    "enabled": true,
    "loss_scale": 0,
    "loss_scale_window": 1000,
    "initial_scale_power": 12,
    "hysteresis": 2,
    "min_loss_scale": 1,
  },

  "train_iters": 143000,
  "lr_decay_iters": 143000,
  "distributed_backend": "nccl",
  "lr_decay_style": "cosine",
  "warmup": 0.01,
  "checkpoint_factor": 1000,
  "eval_interval": 500000,
  "eval_iters": 10,

  "log_interval": 50,
  "steps_per_print": 10,
  "wall_clock_breakdown": true,
  
  "tensorboard-dir": "./tensorboard",
  "log_dir": "./logs",

  # additional deepspeed args not specified above
  "deepspeed_extra_args": {
    "comms_logger": {
        "enabled": true,
        "verbose": true,
        "prof_all": true,
        "debug": false
    },
  }

}

7. 训练你的模型

直接跑 nohup python3 deepy.py train.py --conf_dir configs tom-19M.yml > runlog.log 2>&1 &。

可以在 日志中看到 模型的训练情况。博主使用了 tensorboard 对性能进行可视化，关于模型训练过程的调优和 辅助工具的使用，同样会在后续单独写article进行讨论。

8. 调用你的模型

模型在训练结束以后，会存于 checkpoint里面，可以选择最新的一个 checkpoint文件夹，作为模型调用的文件目录。需要注意的是，这里的模型训练好都是 .pt 的文件，而模型config 则是我们之前 的 yml 文件，也会在 checkpoint文件夹中保存一份。

其中，模型状态文件为：mp_rank_00_model_states.pt，zero_pp_rank_{}_mp_rank_00_optim_states.pt 是优化器状态文件，推理阶段仅需要 mp_rank_00_model_states.pt即可。

这里的config.json是博主自己后续生成的，其实不需要。

调用模型的代码如下：

import torch
from transformers import GPTNeoXConfig, GPTNeoXForCausalLM, GPT2Tokenizer
import yaml

# 指定模型配置文件路径
config_path = "/home/ec2-user/gpt-neox/tiashi/tom-small_checkpoints/global_step30000/configs/tom-small.yml"

# 将 YAML 文件加载为字典
with open(config_path, 'r') as file:
    config_data = yaml.safe_load(file)

config = GPTNeoXConfig(**config_data)

# 创建模型实例
model = GPTNeoXForCausalLM(config)
model_path = "/home/ec2-user/gpt-neox/tiashi/tom-small_checkpoints/global_step30000/"
# 加载模型权重
state_dict = torch.load(f'{model_path}mp_rank_00_model_states.pt', map_location=torch.device('cuda'))
model.load_state_dict(state_dict['module'], strict=False)

# 将模型设为评估模式
model.eval()

#from transformers import GPTNeoXTokenizer

# 加载GPT-NeoX的tokenizer (假设你已经训练并保存了一个tokenizer)
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
#tokenizer = GPTNeoXTokenizer.from_pretrained('gpt2')

# 输入示例
input_text = "Game"
inputs = tokenizer(input_text, return_tensors="pt")

# 生成输出
with torch.no_grad():
    outputs = model.generate(inputs.input_ids, max_length=50)

# 解码输出
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

运行脚本，即可得到模型的输出。评测模型的性能，如何使用harness框架，同样将在后续article中再进行详述。

附录

环境配置中常见的错误见：gpt-neox via CN AWS EC2 部署环境问题汇总，后续会逐步完善。

祝大家生活愉快。