Baichuan2

# 运行如下命令，生成当前机器的RANK_TABLE_FILE的json文件
python ./mindformers/tools/hccl_tools.py --device_num "[0,8)"

{
    "version": "1.0",
    "server_count": "1",
    "server_list": [
        {
            "server_id": "xx.xx.xx.xx",
            "device": [
                {"device_id": "0","device_ip": "192.1.27.6","rank_id": "0"},
                {"device_id": "1","device_ip": "192.2.27.6","rank_id": "1"},
                {"device_id": "2","device_ip": "192.3.27.6","rank_id": "2"},
                {"device_id": "3","device_ip": "192.4.27.6","rank_id": "3"},
                {"device_id": "4","device_ip": "192.1.27.7","rank_id": "4"},
                {"device_id": "5","device_ip": "192.2.27.7","rank_id": "5"},
                {"device_id": "6","device_ip": "192.3.27.7","rank_id": "6"},
                {"device_id": "7","device_ip": "192.4.27.7","rank_id": "7"}],
             "host_nic_ip": "reserve"
        }
    ],
    "status": "completed"
}

# 运行如下命令，生成当前机器的RANK_TABLE_FILE的json文件
python ./mindformers/tools/hccl_tools.py --device_num "[0,8)" --server_ip xx.xx.xx.xx

# 运行如下命令，合并每个机器的RANK_TABLE_FILE文件。
python ./mindformers/tools/merge_hccl.py hccl*.json

{
    "version": "1.0",
    "server_count": "2",
    "server_list": [
        {
            "server_id": "xx.xx.xx.xx",
            "device": [
                {
                    "device_id": "0", "device_ip": "192.168.0.0", "rank_id": "0"
                },
                {
                    "device_id": "1", "device_ip": "192.168.1.0", "rank_id": "1"
                },
                {
                    "device_id": "2", "device_ip": "192.168.2.0", "rank_id": "2"
                },
                {
                    "device_id": "3", "device_ip": "192.168.3.0", "rank_id": "3"
                },
                {
                    "device_id": "4", "device_ip": "192.168.0.1", "rank_id": "4"
                },
                {
                    "device_id": "5", "device_ip": "192.168.1.1", "rank_id": "5"
                },
                {
                    "device_id": "6", "device_ip": "192.168.2.1", "rank_id": "6"
                },
                {
                    "device_id": "7", "device_ip": "192.168.3.1", "rank_id": "7"
                }
            ],
            "host_nic_ip": "reserve"
        },
        {
            "server_id": "xx.xx.xx.xx",
            "device": [
                {
                    "device_id": "0", "device_ip": "192.168.0.1", "rank_id": "8"
                },
                {
                    "device_id": "1", "device_ip": "192.168.1.1", "rank_id": "9"
                },
                {
                    "device_id": "2", "device_ip": "192.168.2.1", "rank_id": "10"
                },
                {
                    "device_id": "3", "device_ip": "192.168.3.1", "rank_id": "11"
                },
                {
                    "device_id": "4", "device_ip": "192.168.0.2", "rank_id": "12"
                },
                {
                    "device_id": "5", "device_ip": "192.168.1.2", "rank_id": "13"
                },
                {
                    "device_id": "6", "device_ip": "192.168.2.2", "rank_id": "14"
                },
                {
                    "device_id": "7", "device_ip": "192.168.3.2", "rank_id": "15"
                }
            ],
            "host_nic_ip": "reserve"
        }
    ],
    "status": "completed"
}

# 脚本路径：research/baichuan2/belle_preprocess.py
python research/baichuan2/belle_preprocess.py \
--input_glob /{path}/belle_chat_ramdon_10k.json \
--model_file /{path}/tokenizer.model \
--output_file /{path}/belle_512.mindrecord \
--seq_length 512

# 参数说明
input_glob: 输入数据集路径
model_file: 词表文件路径
output_file: 输出数据集的路径和名称
seq_length: 生成数据集的序列长度

pip install torch==2.0.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install transformers==4.30.2 -i https://pypi.tuna.tsinghua.edu.cn/simple

python ./research/baichuan/convert_weight.py --torch_ckpt_dir TORCH_CKPT_DIR --mindspore_ckpt_path MS_CKPT_NAME

# 参数说明
TORCH_CKPT_DIR: huggingface权重保存目录路径
mindspore_ckpt_path: 权重保存文件名，保存为TORCH_CKPT_DIR/OUTPUT_NAME, 也可以指定为自定义保存路径

load_checkpoint: 'model_dir/xxx.ckpt'  # 使用完整权重路径
auto_trans_ckpt: True           # 打开自动权重转换
use_parallel: True
run_mode: 'finetune'
# dataset
train_dataset: &train_dataset
  data_loader:
    type: MindDataset
    dataset_dir: "dataset_dir"  # 配置训练数据集文件夹路径
    shuffle: True
  input_columns: ["input_ids", "labels"]
# 指令微调时（如belle数据集），input_columns: ["input_ids", "labels"]
# 继续预训练时（如wikitext2数据集），input_columns: ["input_ids"]

# 8卡分布式策略配置
parallel_config:
  data_parallel: 8
  model_parallel: 1
  pipeline_stage: 1
  micro_batch_num: 1
  vocab_emb_dp: True
  gradient_aggregation_group: 4

cd mindformers/research
bash run_singlenode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_7b_910b.yaml \
--load_checkpoint model_dir/xxx.ckpt \
--auto_trans_ckpt True \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [0,8] 8

# 参数说明
config: 配置文件路径
load_checkpoint: 权重文件路径
auto_trans_ckpt: 自动权重转换开关
run_mode: 运行模式，微调时设置为finetune
train_data: 训练数据集文件夹路径

output_dir: './output'          # 默认路径，若需要自动权重转换，请配置为共享盘输出路径
load_checkpoint: 'model_dir'    # 使用分布式权重，传入路径为离线切分完成的文件夹路径，包含rank_0-rank*
auto_trans_ckpt: False          # 关闭自动权重转换，若需要自动权重转换，则改为True
use_parallel: True
run_mode: 'finetune'
# dataset
train_dataset: &train_dataset
  data_loader:
    type: MindDataset
    dataset_dir: "dataset_dir"  # 配置训练数据集文件夹路径
    shuffle: True
  input_columns: ["input_ids", "labels"]
# 指令微调时（如belle数据集），input_columns: ["input_ids", "labels"]
# 继续预训练时（如wikitext2数据集），input_columns: ["input_ids"]

# 16卡分布式策略配置
parallel_config:
  data_parallel: 2
  model_parallel: 4
  pipeline_stage: 2
  micro_batch_num: 8
  vocab_emb_dp: True
  gradient_aggregation_group: 4

# node 1
cd mindformers/research
bash run_multinode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_7b.yaml \
--load_checkpoint model_dir \
--auto_trans_ckpt False \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [0,8] 16

# node 2
cd mindformers/research
bash run_multinode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_7b.yaml \
--load_checkpoint model_dir \
--auto_trans_ckpt False \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [8,16] 16

# 参数说明
config: 配置文件路径
load_checkpoint: 权重文件夹路径，权重按照'model_dir/rank_x/xxx.ckpt'格式存放
auto_trans_ckpt: 自动权重转换开关
run_mode: 运行模式，微调时设置为finetune
train_data: 训练数据集文件夹路径

load_checkpoint: 'model_dir'    # 使用完整权重路径或分别加载预训练、lora权重
auto_trans_ckpt: True           # 打开自动权重转换
use_parallel: True
run_mode: 'finetune'
# dataset
train_dataset: &train_dataset
  data_loader:
    type: MindDataset
    dataset_dir: "dataset_dir"  # 配置训练数据集文件夹路径
    shuffle: True
  input_columns: ["input_ids", "labels"]
# 指令微调时（如belle数据集），input_columns: ["input_ids", "labels"]
# 继续预训练时（如wikitext2数据集），input_columns: ["input_ids"]

# 8卡分布式策略配置
parallel_config:
  data_parallel: 8
  model_parallel: 1
  pipeline_stage: 1
  micro_batch_num: 1
  vocab_emb_dp: True
  gradient_aggregation_group: 4

# 权重保存配置参考
parallel:
  strategy_ckpt_config:
    save_file: "./ckpt_strategy.ckpt"
    only_trainable_params: False # 配置为False，保存完整权重

# callbacks
callbacks:
  - type: CheckpointMointor
    save_checkpoint_steps: 500
    save_trainable_params: True  #若需要单独保存lora参数的训练权重，可将此参数置为True

# model增加pet_config
model:
  model_config:
    pet_config:
      pet_type: lora
      # configuration of lora
      lora_rank: 1
      lora_alpha: 32
      lora_dropout: 0.1
      target_modules: '.*wq|.*wk|.*wv'

cd mindformers/research
bash run_singlenode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_7b_lora_910b.yaml \
--load_checkpoint model_dir \
--auto_trans_ckpt True \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [0,8] 8

# 参数说明
config: 配置文件路径
load_checkpoint: 权重文件路径
auto_trans_ckpt: 自动权重转换开关
run_mode: 运行模式，微调时设置为finetune
train_data: 训练数据集文件夹路径

# research/baichuan2/run_baichuan2_7b_lora_910b.yaml
max_device_memory: "31GB"    # Atlas 800将最大内存改为31GB即可

load_checkpoint: 'path/to/Baichuan2_7B_Chat.ckpt'   # 填写权重路径
auto_trans_ckpt: False                              # 关闭自动权重转换
use_past: True                                      # 使用增量推理
vocab_file: 'path/to/tokenizer.model'               # 配置词表路径
use_parallel: False                                 # 关闭并行模式

cd research
# 推理命令中参数会覆盖yaml文件中的相同参数
python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_7b.yaml \
--run_mode predict \
--use_parallel False \
--load_checkpoint path/to/Baichuan2_7B_Chat.ckpt \
--auto_trans_ckpt False \
--predict_data "<reserved_106>你是谁？<reserved_107>"

# output: [{'text_generation_text': ['<reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问']}]

# run_baichuan2_pipeline.py
from mindspore import context
from mindformers.pipeline import pipeline
from mindformers.models import LlamaConfig
from mindformers import MindFormerConfig

from baichuan2_7b import Baichuan7BV2ForCausalLM
from baichuan2_13b import Baichuan13BV2ForCausalLM
from baichuan2_tokenizer import Baichuan2Tokenizer

model_dict = {
    "baichuan2_7b": Baichuan7BV2ForCausalLM,
    "baichuan2_13b": Baichuan13BV2ForCausalLM,
}

context.set_context(device_id=0, mode=0)

inputs = ["<reserved_106>你是谁？<reserved_107>",
          "<reserved_106>《静夜思》作者是？<reserved_107>",
          "<reserved_106>白日依山尽，下一句是？<reserved_107>"]

# init model
baichuan2_config_path = "path/to/run_baichuan2_7b.yaml"
baichuan2_config = MindFormerConfig(baichuan2_config_path)

baichuan2_config.model.model_config.batch_size = 1
baichuan2_model_config = LlamaConfig(**baichuan2_config.model.model_config)
model_name = baichuan2_config.trainer.model_name
baichuan2_network = model_dict[model_name](
    config=baichuan2_model_config
)

# init tokenizer
tokenizer = Baichuan2Tokenizer(
    vocab_file=baichuan2_config.processor.tokenizer.vocab_file
)

# init and run pipeline
pipeline_task = pipeline(task="text_generation", model=baichuan2_network, tokenizer=tokenizer)
outputs = pipeline_task(inputs,
                        do_sample=False,
                        top_k=1,
                        top_p=1.0,
                        repetition_penalty=1.05,
                        temperature=1.0,
                        max_length=64)
for output in outputs:
    print(output)

load_checkpoint: 'path/to/Baichuan2-7B-Chat.ckpt'             # 填写权重绝对路径
auto_trans_ckpt: False                                        # 关闭自动权重转换
use_past: True                                                # 使用增量推理
vocab_file: 'path/to/tokenizer.model'                         # 配置词表路径
use_parallel: False                                           # 关闭并行模式

python baichuan2/run_baichuan2_pipeline.py

# 推理输出
# {'text_generation_text': ['<reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问']}
# {'text_generation_text': ['<reserved_106>《静夜思》作者是？<reserved_107>《静夜思》的作者是唐代著名诗人李白。']}
# {'text_generation_text': ['<reserved_106>白日依山尽，下一句是？<reserved_107>黄河入海流。']}

# run_baichuan2_generate.py
from mindspore import context
from mindformers.pipeline import pipeline
from mindformers.models import LlamaConfig
from mindformers import MindFormerConfig

from baichuan2_7b import Baichuan7BV2ForCausalLM
from baichuan2_13b import Baichuan13BV2ForCausalLM
from baichuan2_tokenizer import Baichuan2Tokenizer

model_dict = {
    "baichuan2_7b": Baichuan7BV2ForCausalLM,
    "baichuan2_13b": Baichuan13BV2ForCausalLM,
}

inputs = ["<reserved_106>你是谁？<reserved_107>",
          "<reserved_106>《静夜思》作者是？<reserved_107>",
          "<reserved_106>白日依山尽，下一句是？<reserved_107>"]
batch_size = len(inputs)

# init model
baichuan2_config_path = "path/to/run_baichuan2_7b.yaml"
baichuan2_config = MindFormerConfig(baichuan2_config_path)

baichuan2_config.model.model_config.batch_size = batch_size
baichuan2_model_config = LlamaConfig(**baichuan2_config.model.model_config)
model_name = baichuan2_config.trainer.model_name
baichuan2_network = model_dict[model_name](
    config=baichuan2_model_config
)

# init tokenizer
tokenizer = Baichuan2Tokenizer(
    vocab_file=baichuan2_config.processor.tokenizer.vocab_file
)

# predict using generate
inputs_ids = tokenizer(inputs, max_length=64, padding="max_length")["input_ids"]
outputs = baichuan2_network.generate(inputs_ids,
                                     do_sample=False,
                                     top_k=1,
                                     top_p=1.0,
                                     repetition_penalty=1.05,
                                     temperature=1.0,
                                     max_length=64)
for output in outputs:
    print(tokenizer.decode(output))

load_checkpoint: 'path/to/Baichuan2-7B-Chat.ckpt'             # 填写权重绝对路径
auto_trans_ckpt: False                                        # 关闭自动权重转换
use_past: True                                                # 使用增量推理
vocab_file: 'path/to/tokenizer.model'                         # 配置词表路径
use_parallel: False                                           # 关闭并行模式

python baichuan2/run_baichuan2_generate.py

# 推理输出
# <reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问</s>
# <reserved_106>《静夜思》作者是？<reserved_107>《静夜思》的作者是唐代著名诗人李白。</s>
# <reserved_106>白日依山尽，下一句是？<reserved_107>黄河入海流。</s>

cd research/baichuan2
python run_baichuan2_chat.py \
--config run_baichuan2_7b_910b.yaml \
--use_parallel False \
--load_checkpoint '/path/to/Baichuan2-7B-Chat.ckpt' \
--auto_trans_ckpt False \
--vocab_file '/path/to/tokenizer.model'

infer_baichuan2_7b_dir
    ├── run_baichuan2_7b_910b.yaml    # 推理模型的配置文件
    └── Baichuan2-7B-Chat.ckpt        # 推理模型的权重文件

infer:
  prefill_model_path: "/path/to/baichuan2_7b_export/baichuan2_7b_prefill.mindir"
  increment_model_path: "/path/to/baichuan2_7b_export/baichuan2_7b_inc.mindir"
  infer_seq_length: 512
  model_type: mindir

# 参数说明：
prefill_model_path: 全量图路径
increment_model_path: 增量图路径
infer_seq_length: 推理序列长度
model_type: 推理模型类型

# 注意与export.py中的batch_size设置保持一致，如下所示，否则用导出的MindIR图进行单卡推理可能会出现out of memory的问题
# 由于配置了prefill_model_path和increment_model_path两个路径，需要导出增量图，因此在模型配置中打开增量开关，如下所示
# 使用后处理加速需要配置is_sample_acceleration开关，注意与推理脚本中的设置保持一致，如下所示
model:
  model_config:
    batch_size: 1                    # 单batch推理设置为1，多batch推理设置为相应的batch数
    use_past: True
    is_sample_acceleration: False    # 后处理加速开关，当前baichuan2模型暂不支持，设置为False

python mindformers/tools/export.py --model_dir /path/to/infer_baichuan2_7b_dir

python run_infer_main.py \
--device_id 0 \
--model_name baichuan2_7b \
--seq_length 512 \                            # 注意与export导出时的推理序列长度保持一致
--tokenizer_path path/to/tokenizer.model \    # 不设置时，以from_pretrained的方式自动加载tokenizer（research模型不支持）
--prefill_model_path /path/to/baichuan2_7b_export/baichuan2_7b_prefill.mindir \
--increment_model_path /path/to/baichuan2_7b_export/baichuan2_7b_inc.mindir \
--config_path /path/to/910b_ge_default.cfg \
--do_sample False \
--top_k 1 \
--top_p 1.0 \
--repetition_penalty 1.0 \
--temperature 1.0 \
--max_length 512 \
--is_sample_acceleration False \              # 后处理加速开关，当前baichuan2模型暂不支持，设置为False
--add_special_tokens False \
--stream False

# 参数说明
device_id: 设备物理ID
model_name: 模型名称
seq_length: 推理序列长度
tokenizer_path: 模型tokenizer路径
prefill_model_path: 全量图路径
increment_model_path: 增量图路径
config_path: GE配置文件路径
do_sample: 是否对候选id进行采样
top_k: 选择top_k个token id作为候选
top_p: 将累积概率小于top_k的token id作为候选
repetition_penalty: 生成单词的惩罚因子，设置为1时不打开
temperature: 温度系数，用来调整下个token的概率
max_length: 能够生成的最大语句长度
is_sample_acceleration: 后处理加速开关
add_special_tokens: 对输入token化时是否添加特殊字符
stream: 是否采用流式结果返回
prompt: 输入中加入prompt的内容，Baichuan2可以选择不设置，按默认的prompt进行推理

# output
# ['<reserved_106>解释一下“温故而知新”<reserved_107>“温故而知新”是一个中国古代成语，出自《论语·为政》。它的意思是通过回顾和了解过去的事情，可以从中获得新的知识和启示。这个成语强调了学习和思考的重要性，鼓励人们在不断积累知识的过程中，不断地回顾和总结，从而实现自我提升和成长。']

[ascend_context]
provider=ge

[ge_session_options]
ge.externalWeight=1
ge.exec.atomicCleanPolicy=1
ge.event=notify
ge.exec.staticMemoryPolicy=2
ge.exec.formatMode=1
ge.exec.precision_mode=must_keep_origin_dtype

# 参数说明
# provider=ge：采用GE接口
# ge.externalWeight=1：将网络中Const/Constant节点的权重保存在单独的文件中
# ge.exec.atomicCleanPolicy=1：不集中清理网络中atomic算子占用的内存
# ge.exec.staticMemoryPolicy=2：网络运行使用动态扩展内存方式
# ge.exec.precision_mode=must_keep_origin_dtype：选择算子精度模式

load_checkpoint: 'model_dir/xxx.ckpt'  # 使用完整权重路径
auto_trans_ckpt: True           # 打开自动权重转换
use_parallel: True
run_mode: 'finetune'
# dataset
train_dataset: &train_dataset
  data_loader:
    type: MindDataset
    dataset_dir: "dataset_dir"  # 配置训练数据集文件夹路径
    shuffle: True
  input_columns: ["input_ids", "labels"]
# 指令微调时（如belle数据集），input_columns: ["input_ids", "labels"]
# 继续预训练时（如wikitext2数据集），input_columns: ["input_ids"]

# 8卡分布式策略配置
parallel_config:
  data_parallel: 1
  model_parallel: 2
  pipeline_stage: 4
  micro_batch_num: 16
  vocab_emb_dp: True
  gradient_aggregation_group: 4

cd mindformers/research
bash run_singlenode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_13b_910b.yaml \
--load_checkpoint model_dir/xxx.ckpt \
--auto_trans_ckpt True \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [0,8] 8

# 参数说明
config: 配置文件路径
load_checkpoint: 权重文件路径
auto_trans_ckpt: 自动权重转换开关
run_mode: 运行模式，微调时设置为finetune
train_data: 训练数据集文件夹路径

output_dir: './output'          # 默认路径，若需要自动权重转换，请配置为共享盘输出路径
load_checkpoint: 'model_dir'    # 使用分布式权重，传入路径为离线切分完成的文件夹路径，包含rank_0-rank*
auto_trans_ckpt: False          # 关闭自动权重转换，若需要自动权重转换，则改为True
use_parallel: True
run_mode: 'finetune'
# dataset
train_dataset: &train_dataset
  data_loader:
    type: MindDataset
    dataset_dir: "dataset_dir"  # 配置训练数据集文件夹路径
    shuffle: True
  input_columns: ["input_ids", "labels"]
# 指令微调时（如belle数据集），input_columns: ["input_ids", "labels"]
# 继续预训练时（如wikitext2数据集），input_columns: ["input_ids"]

# 16卡分布式策略配置
parallel_config:
  data_parallel: 1
  model_parallel: 8
  pipeline_stage: 2
  micro_batch_num: 16
  vocab_emb_dp: True
  gradient_aggregation_group: 4

# node 1
cd mindformers/research
bash run_multinode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_13b.yaml \
--load_checkpoint model_dir \
--auto_trans_ckpt False \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [0,8] 16

# node 2
cd mindformers/research
bash run_multinode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_13b.yaml \
--load_checkpoint model_dir \
--auto_trans_ckpt False \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [8,16] 16

# 参数说明
config: 配置文件路径
load_checkpoint: 权重文件夹路径，权重按照'model_dir/rank_x/xxx.ckpt'格式存放
auto_trans_ckpt: 自动权重转换开关
run_mode: 运行模式，微调时设置为finetune
train_data: 训练数据集文件夹路径

load_checkpoint: 'model_dir'    # 使用完整权重路径或分别加载预训练、lora权重
auto_trans_ckpt: True           # 打开自动权重转换
use_parallel: True
run_mode: 'finetune'
# dataset
train_dataset: &train_dataset
  data_loader:
    type: MindDataset
    dataset_dir: "dataset_dir"  # 配置训练数据集文件夹路径
    shuffle: True
  input_columns: ["input_ids", "labels"]
# 指令微调时（如belle数据集），input_columns: ["input_ids", "labels"]
# 继续预训练时（如wikitext2数据集），input_columns: ["input_ids"]

# 8卡分布式策略配置
parallel_config:
  data_parallel: 8
  model_parallel: 1
  pipeline_stage: 1
  micro_batch_num: 1
  vocab_emb_dp: True
  gradient_aggregation_group: 4

# 权重保存配置参考
parallel:
  strategy_ckpt_config:
    save_file: "./ckpt_strategy.ckpt"
    only_trainable_params: False # 配置为False，保存完整权重

# callbacks
callbacks:
  - type: CheckpointMointor
    save_checkpoint_steps: 500
    save_trainable_params: True  #若需要单独保存lora参数的训练权重，可将此参数置为True

# model增加pet_config
model:
  model_config:
    pet_config:
      pet_type: lora
      # configuration of lora
      lora_rank: 1
      lora_alpha: 32
      lora_dropout: 0.1
      target_modules: '.*wq|.*wk|.*wv'

cd mindformers/research
bash run_singlenode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_13b_lora_910b.yaml \
--load_checkpoint model_dir \
--auto_trans_ckpt True \
--use_parallel True \
--run_mode finetune \
--train_data dataset_dir" \
RANK_TABLE_FILE [0,8] 8

# 参数说明
config: 配置文件路径
load_checkpoint: 权重文件路径
auto_trans_ckpt: 自动权重转换开关
run_mode: 运行模式，微调时设置为finetune
train_data: 训练数据集文件夹路径

# research/baichuan2/run_baichuan2_13b_lora_910b.yaml
max_device_memory: "31GB"    # Atlas 800将最大内存改为31GB即可

# Atlas 800的8卡分布式策略配置
parallel_config:
  data_parallel: 2
  model_parallel: 1
  pipeline_stage: 4
  micro_batch_num: 4
  vocab_emb_dp: True
  gradient_aggregation_group: 4

python mindformers/tools/hccl_tools.py --device_num [0,2]

load_checkpoint: 'path/to/Baichuan2_13B_Chat.ckpt'  # 填写权重路径
auto_trans_ckpt: False                              # 关闭自动权重转换
use_past: True                                      # 使用增量推理
vocab_file: 'path/to/tokenizer.model'               # 配置词表路径
use_parallel: False                                 # 关闭并行模式

cd research
# 推理命令中参数会覆盖yaml文件中的相同参数
python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_13b_910b.yaml \
--run_mode predict \
--use_parallel False \
--load_checkpoint path/to/Baichuan2_13B_Chat.ckpt \
--auto_trans_ckpt False \
--predict_data <reserved_106>你是谁？<reserved_107>

# output: [{'text_generation_text': ['<reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问']}]

load_checkpoint: 'model_dir/xxx.ckpt'    # 使用完整权重路径
auto_trans_ckpt: True                    # 打开自动权重转换
use_past: True                           # 使用增量推理
use_parallel: True                       # 使用并行模式
vocab_file: 'path/to/tokenizer.model'    # 配置词表路径

# 分布式配置
parallel_config:
  data_parallel: 1
  model_parallel: 2
  pipeline_stage: 1
  micro_batch_num: 1
  vocab_emb_dp: True
  gradient_aggregation_group: 4

cd research
# 推理命令中参数会覆盖yaml文件中的相同参数
bash ./run_singlenode.sh \
"python baichuan2/run_baichuan2.py \
--config baichuan2/run_baichuan2_13b.yaml \
--run_mode predict \
--use_parallel True \
--load_checkpoint model_dir/xxx.ckpt \
--auto_trans_ckpt True \
--predict_data <reserved_106>你是谁？<reserved_107>" RANK_TABLE_FILE [0,2] 2

# output: [{'text_generation_text': ['<reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问']}]

# run_baichuan2_pipeline.py
import mindspore as ms
from mindspore import Model
from mindspore import Tensor
from mindspore.common import initializer as init

from mindformers import MindFormerConfig
from mindformers.pipeline import pipeline
from mindformers.models import LlamaConfig
from mindformers.pet import get_pet_model, LoraConfig
from mindformers.trainer.utils import transform_and_load_checkpoint
from mindformers.core.context import build_context
from mindformers.core.parallel_config import build_parallel_config

from baichuan2_7b import Baichuan7BV2ForCausalLM
from baichuan2_13b import Baichuan13BV2ForCausalLM
from baichuan2_tokenizer import Baichuan2Tokenizer

model_dict = {
    "baichuan2_7b": Baichuan7BV2ForCausalLM,
    "baichuan2_7b_lora": Baichuan7BV2ForCausalLM,
    "baichuan2_13b": Baichuan13BV2ForCausalLM,
    "baichuan2_13b_lora": Baichuan13BV2ForCausalLM
}

inputs = ["<reserved_106>你是谁？<reserved_107>",
          "<reserved_106>《静夜思》作者是？<reserved_107>",
          "<reserved_106>白日依山尽，下一句是？<reserved_107>"]

# init config，默认使用Atlas 800配置文件
baichuan2_config_path = "path/to/run_baichuan2_13b.yaml"
baichuan2_config = MindFormerConfig(baichuan2_config_path)

# init context
build_context(baichuan2_config)
build_parallel_config(baichuan2_config)

# init model
baichuan2_config.model.model_config.parallel_config = baichuan2_config.parallel_config
baichuan2_config.model.model_config.batch_size = 1
baichuan2_model_config = LlamaConfig(**baichuan2_config.model.model_config)
baichuan2_model_config.checkpoint_name_or_path = None
model_name = baichuan2_config.trainer.model_name
baichuan2_network = model_dict[model_name](
    config=baichuan2_model_config
)

if baichuan2_config.model.model_config.pet_config:
    print("----------------Init lora params----------------")
    pet_config = LoraConfig(
        lora_rank=baichuan2_config.model.model_config.pet_config.lora_rank,
        lora_alpha=baichuan2_config.model.model_config.pet_config.lora_alpha,
        lora_dropout=baichuan2_config.model.model_config.pet_config.lora_dropout,
        target_modules=baichuan2_config.model.model_config.pet_config.target_modules
    )
    baichuan2_network = get_pet_model(baichuan2_network, pet_config)

baichuan2_model = Model(baichuan2_network)

# load checkpoint
if baichuan2_config.load_checkpoint:
    print("----------------Transform and load checkpoint----------------")
    seq_length = baichuan2_config.model.model_config.seq_length
    infer_data = Tensor(shape=(1, seq_length), dtype=ms.int32, init=init.One())
    transform_and_load_checkpoint(baichuan2_config, baichuan2_model, baichuan2_network, infer_data, do_predict=True)

# init tokenizer
tokenizer = Baichuan2Tokenizer(
    vocab_file=baichuan2_config.processor.tokenizer.vocab_file
)

# init and run pipeline
pipeline_task = pipeline(task="text_generation", model=baichuan2_model, tokenizer=tokenizer)
outputs = pipeline_task(inputs,
                        do_sample=False,
                        top_k=1,
                        top_p=1.0,
                        repetition_penalty=1.0,
                        temperature=1.0,
                        max_length=64)
for output in outputs:
    print(output)

# 使用完整权重
load_checkpoint: 'path/to/Baichuan2-13B-Chat.ckpt'            # 填写权重绝对路径
auto_trans_ckpt: False                                        # 关闭自动权重转换
use_past: True                                                # 使用增量推理
vocab_file: 'path/to/tokenizer.model'                         # 配置词表路径
use_parallel: False                                           # 关闭并行模式

# 需要将分布式权重转换为完整权重
load_checkpoint: 'model_dir'              # 分布式权重文件夹路径，权重存放格式为"model_dir/rank_x/xxx.ckpt"
src_strategy_path_or_dir: 'strategy_path' # 使用分布式权重且需要权重转换，填写分布式策略文件路径
auto_trans_ckpt: True                     # 打开自动权重转换

python baichuan2/run_baichuan2_pipeline.py

# 推理输出
# {'text_generation_text': [<reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问</s>]}
# {'text_generation_text': [<reserved_106>《静夜思》作者是？<reserved_107>《静夜思》的作者是唐代著名诗人李白。这是一首描绘夜晚思乡之情的诗篇，表达了作者对故乡的思念和对亲人的牵挂之情。</s>]}
# {'text_generation_text': [<reserved_106>白日依山尽，下一句是？<reserved_107>黄河入海流。</s>]}

# 如果使用完整权重进行多卡推理，需要将权重转换为分布式权重
load_checkpoint: 'path/to/Baichuan2-13B-Chat.ckpt'   # 填写权重绝对路径
src_strategy_path_or_dir: ''                         # 使用完整权重，不需要填写策略文件路径
auto_trans_ckpt: True                                # 打开自动权重转换
use_parallel: True                                   # 使用并行模式
use_past: True                                       # 使用增量推理
vocab_file: 'path/to/tokenizer.model'                # 配置词表路径

# 分布式配置
parallel_config:
  data_parallel: 1
  model_parallel: 2
  pipeline_stage: 1
  micro_batch_num: 1
  vocab_emb_dp: True
  gradient_aggregation_group: 4

# 1.需要将分布式权重转换为推理对应的分布式权重
load_checkpoint: 'model_dir'              # 分布式权重文件夹路径，权重存放格式为"model_dir/rank_x/xxx.ckpt"
src_strategy_path_or_dir: 'strategy_path' # 使用分布式权重且需要权重转换，填写分布式策略文件路径
auto_trans_ckpt: True                     # 打开自动权重转换

# 2.直接使用推理对应的分布式权重
load_checkpoint: 'model_dir'              # 使用分布式权重，权重存放格式为"model_dir/rank_x/xxx.ckpt"
src_strategy_path_or_dir: ''              # 不进行权重转换，不需要填写策略文件路径
auto_trans_ckpt: False                    # 关闭自动权重转换

cd research
bash run_singlenode.sh "python baichuan2/run_baichuan2_pipeline.py" RANK_TABLE_FILE [0,2] 2

# 推理输出
# {'text_generation_text': [<reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问</s>]}
# {'text_generation_text': [<reserved_106>《静夜思》作者是？<reserved_107>《静夜思》的作者是唐代著名诗人李白。这是一首描绘夜晚思乡之情的诗篇，表达了作者对故乡的思念和对亲人的牵挂之情。</s>]}
# {'text_generation_text': [<reserved_106>白日依山尽，下一句是？<reserved_107>黄河入海流。</s>]}

# run_baichuan2_generate.py
import mindspore as ms
from mindspore import Model
from mindspore import Tensor
from mindspore.common import initializer as init

from mindformers import MindFormerConfig
from mindformers.models import LlamaConfig
from mindformers.pet import get_pet_model, LoraConfig
from mindformers.trainer.utils import transform_and_load_checkpoint
from mindformers.core.context import build_context
from mindformers.core.parallel_config import build_parallel_config

from baichuan2_7b import Baichuan7BV2ForCausalLM
from baichuan2_13b import Baichuan13BV2ForCausalLM
from baichuan2_tokenizer import Baichuan2Tokenizer

model_dict = {
    "baichuan2_7b": Baichuan7BV2ForCausalLM,
    "baichuan2_7b_lora": Baichuan7BV2ForCausalLM,
    "baichuan2_13b": Baichuan13BV2ForCausalLM,
    "baichuan2_13b_lora": Baichuan13BV2ForCausalLM
}

inputs = ["<reserved_106>你是谁？<reserved_107>",
          "<reserved_106>《静夜思》作者是？<reserved_107>",
          "<reserved_106>白日依山尽，下一句是？<reserved_107>"]
batch_size = len(inputs)

# init config，默认使用Atlas 800配置文件
baichuan2_config_path = "path/to/run_baichuan2_13b.yaml"
baichuan2_config = MindFormerConfig(baichuan2_config_path)

# init context
build_context(baichuan2_config)
build_parallel_config(baichuan2_config)

# init model
baichuan2_config.model.model_config.parallel_config = baichuan2_config.parallel_config
baichuan2_config.model.model_config.batch_size = batch_size
baichuan2_model_config = LlamaConfig(**baichuan2_config.model.model_config)
model_name = baichuan2_config.trainer.model_name
baichuan2_network = model_dict[model_name](
    config=baichuan2_model_config
)

if baichuan2_config.model.model_config.pet_config:
    print("----------------Init lora params----------------")
    pet_config = LoraConfig(
        lora_rank=baichuan2_config.model.model_config.pet_config.lora_rank,
        lora_alpha=baichuan2_config.model.model_config.pet_config.lora_alpha,
        lora_dropout=baichuan2_config.model.model_config.pet_config.lora_dropout,
        target_modules=baichuan2_config.model.model_config.pet_config.target_modules
    )
    baichuan2_network = get_pet_model(baichuan2_network, pet_config)

baichuan2_model = Model(baichuan2_network)

# load checkpoint
if baichuan2_config.load_checkpoint:
    print("----------------Transform and load checkpoint----------------")
    seq_length = baichuan2_config.model.model_config.seq_length
    infer_data = Tensor(shape=(batch_size, seq_length), dtype=ms.int32, init=init.One())
    transform_and_load_checkpoint(baichuan2_config, baichuan2_model, baichuan2_network, infer_data, do_predict=True)

# init tokenizer
tokenizer = Baichuan2Tokenizer(
    vocab_file=baichuan2_config.processor.tokenizer.vocab_file
)

# predict using generate
inputs_ids = tokenizer(inputs, max_length=64, padding="max_length")["input_ids"]
outputs = baichuan2_network.generate(inputs_ids,
                                     do_sample=False,
                                     top_k=1,
                                     top_p=1.0,
                                     repetition_penalty=1.0,
                                     temperature=1.0,
                                     max_length=64)
for output in outputs:
    print(tokenizer.decode(output))

# 使用完整权重
load_checkpoint: 'path/to/Baichuan2-13B-Chat.ckpt'            # 填写权重绝对路径
auto_trans_ckpt: False                                        # 关闭自动权重转换
use_past: True                                                # 使用增量推理
vocab_file: 'path/to/tokenizer.model'                         # 配置词表路径
use_parallel: False                                           # 关闭并行模式

# 需要将分布式权重转换为完整权重
load_checkpoint: 'model_dir'              # 分布式权重文件夹路径，权重存放格式为"model_dir/rank_x/xxx.ckpt"
src_strategy_path_or_dir: 'strategy_path' # 使用分布式权重且需要权重转换，填写分布式策略文件路径
auto_trans_ckpt: True                     # 打开自动权重转换

python baichuan2/run_baichuan2_generate.py

# 推理输出
# <reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问</s>
# <reserved_106>《静夜思》作者是？<reserved_107>《静夜思》的作者是唐代著名诗人李白。这是一首描绘夜晚思乡之情的诗篇，表达了作者对故乡的思念和对亲人的牵挂之情。</s>
# <reserved_106>白日依山尽，下一句是？<reserved_107>黄河入海流。</s>

# 如果使用完整权重进行多卡推理，需要将权重转换为分布式权重
load_checkpoint: 'path/to/Baichuan2-13B-Chat.ckpt'    # 填写权重绝对路径
src_strategy_path_or_dir: ''                          # 使用完整权重，不需要填写策略文件路径
auto_trans_ckpt: True                                 # 打开自动权重转换
use_parallel: True                                    # 使用并行模式
use_past: True                                        # 使用增量推理
vocab_file: 'path/to/tokenizer.model'                 # 配置词表路径

# 分布式配置
parallel_config:
  data_parallel: 1
  model_parallel: 2
  pipeline_stage: 1
  micro_batch_num: 1
  vocab_emb_dp: True
  gradient_aggregation_group: 4

# 1.需要将分布式权重转换为推理对应的分布式权重
load_checkpoint: 'model_dir'              # 分布式权重文件夹路径，权重存放格式为"model_dir/rank_x/xxx.ckpt"
src_strategy_path_or_dir: 'strategy_path' # 使用分布式权重且需要权重转换，填写分布式策略文件路径
auto_trans_ckpt: True                     # 打开自动权重转换

# 2.直接使用推理对应的分布式权重
load_checkpoint: 'model_dir'              # 使用分布式权重，权重存放格式为"model_dir/rank_x/xxx.ckpt"
src_strategy_path_or_dir: ''              # 不进行权重转换，不需要填写策略文件路径
auto_trans_ckpt: False                    # 关闭自动权重转换

cd research
bash run_singlenode.sh "python baichuan2/run_baichuan2_generate.py" RANK_TABLE_FILE [0,2] 2

# 推理输出
# <reserved_106>你是谁？<reserved_107>我是百川大模型，是由百川智能的工程师们创造的大语言模型，我可以和人类进行自然交流、解答问题、协助创作，帮助大众轻松、普惠的获得世界知识和专业服务。如果你有任何问题，可以随时向我提问</s>
# <reserved_106>《静夜思》作者是？<reserved_107>《静夜思》的作者是唐代著名诗人李白。这是一首描绘夜晚思乡之情的诗篇，表达了作者对故乡的思念和对亲人的牵挂之情。</s>
# <reserved_106>白日依山尽，下一句是？<reserved_107>黄河入海流。</s>

cd research/baichuan2
python run_baichuan2_chat.py \
--config run_baichuan2_13b_910b.yaml \
--use_parallel False \
--load_checkpoint '/path/to/Baichuan2-13B-Chat.ckpt' \
--auto_trans_ckpt False \
--vocab_file '/path/to/tokenizer.model'

# model config
model:
  model_config:
    seq_length: 512
    checkpoint_name_or_path: "/path/to/Baichuan2-13B-Chat.ckpt"
    use_past: True              # 开启增量推理
    is_dynamic: False           # 使用双动态推理时设置为True
    use_kvcache_op: True        # 是否使用kvcache融合算子，推荐设置为True
    is_flexible_shape: False    # 是否固定kvcache大小为bs*seq
    use_rope_slice: False       # 是否使用RoPE位置编码slice

python run_baichuan2.py \
--config export_baichuan2_13b.yaml \
--run_mode export \
--use_parallel False \
--device_id 0

[ascend_context]
plugin_custom_ops=All
provider=ge

[ge_session_options]
ge.exec.formatMode=1
ge.exec.atomicCleanPolicy=1
ge.exec.staticMemoryPolicy=2
ge.exec.precision_mode=must_keep_origin_dtype

# 参数说明
# provider=ge：采用GE接口
# ge.externalWeight=1：将网络中Const/Constant节点的权重保存在单独的文件中
# ge.exec.atomicCleanPolicy=1：不集中清理网络中atomic算子占用的内存
# ge.exec.staticMemoryPolicy=2：网络运行使用动态扩展内存方式
# ge.exec.precision_mode=must_keep_origin_dtype：选择算子精度模式

[ascend_context]
plugin_custom_ops=All
provider=ge

[ge_session_options]
ge.exec.formatMode=1
ge.exec.atomicCleanPolicy=1
ge.exec.staticMemoryPolicy=2
ge.exec.precision_mode=must_keep_origin_dtype

[ge_graph_options]
ge.inputShape=batch_index:-1;batch_valid_length:-1;tokens:-1,1;zactivate_len:-1
ge.dynamicDims=1,1,1,256;2,2,2,256;4,4,4,256;1,1,1,512
ge.dynamicNodeType=1

# 参数说明
# ge.inputShape：设置参数动态输入，-1表示动态入参
# ge.dynamicDims：设置实际推理的batch size和activate length，与ge.inputShape中-1的位置依次对应

python run_infer_main.py \
--device_id 0 \
--model_name baichuan2_13b \
--seq_length 2048 \                           # 注意静态推理时需要与export导出的推理序列长度保持一致
--tokenizer_path path/to/tokenizer.model \    # 不设置时，以from_pretrained的方式自动加载tokenizer（research模型不支持）
--prefill_model_path /path/to/minder_full_checkpoint/rank_0_graph.mindir \
--increment_model_path /path/to/minder_inc_checkpoint/rank_0_graph.mindir \
--config_path /path/to/910b_ge_default_ctx.ini \
--do_sample False \
--top_k 1 \
--top_p 1.0 \
--repetition_penalty 1.0 \
--temperature 1.0 \
--max_length 2048 \
--is_sample_acceleration False \              # 后处理加速开关，当前baichuan2模型暂不支持，设置为False
--add_special_tokens False \
--dynamic False

# 参数说明
device_id: 设备物理ID
model_name: 模型名称
seq_length: 推理序列长度
tokenizer_path: 模型tokenizer路径
prefill_model_path: 全量图路径
increment_model_path: 增量图路径
config_path: GE配置文件路径
do_sample: 是否对候选id进行采样
top_k: 选择top_k个token id作为候选
top_p: 将累积概率小于top_k的token id作为候选
repetition_penalty: 生成单词的惩罚因子，设置为1时不打开
temperature: 温度系数，用来调整下个token的概率
max_length: 能够生成的最大语句长度
is_sample_acceleration: 后处理加速开关
add_special_tokens: 对输入token化时是否添加特殊字符
dynamic: 是否采用双动态推理
prompt: 输入中加入prompt的内容，Baichuan2可以选择不设置，按默认的prompt进行推理

--config_path /path/to/910b_ge_default_ctx.ini，/path/to/910b_ge_default_inc.ini
--dynamic True

context:
  runtime_num_threads=1 # 新增配置

parallel_config:
  data_parallel: 8
  model_parallel: 1
  pipeline_stage: 2 # pipeline_stage设置为机器节点数量，比如有2台机器，则为2
  micro_batch_num: 8 # micro_batch_num必须是大于等于pipeline_stage

export MS_MEMORY_POOL_RECYCLE=1
export GE_NOT_CUT=1