@TomMonkeyMan

@TomMonkeyMan

训练 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 部署环境问题汇总，后续会逐步完善。

祝大家生活愉快。

【GPT neox training on-hands (1) 】配置与训练 Training GPT neox on CN AWS EC2