Support LLaMA training in NEOM format

[youth123]

PR Category

Train

PR Types

New Features

PR Description

• Supports loading and saving checkpoints in nemo zarr format
• Supports train packed seqs
• Fix the issue where wandb finalization cannot find the latest_checkpointed_iteration file
• Fix lora can not support layernorm weight load & not support nemo zarr

The checkpoint file format is as follows：
load zarr format:
-context
-weights
-module.decoder.xxx._extra_state
-module.decoder.xxx.weight
-optimizer.state.fp32_param.xxx.weight
-optimizer.state.fp32_param.xxx.weight.sync
common.pt
meatadata.json

save zarr format：
-iter_xxx
-module.decoder.xxx._extra_state
-module.decoder.xxx.weight
-optimizer.state.fp32_param.xxx.weight
-optimizer.state.fp32_param.xxx.weight.sync
common.pt
meatadata.json
latest_checkpointed_iteration.txt

The comparison of nemo and flagscale under different distributed strategies is as follows：
lora:

• dp2:

• tp2:

• pp2:

full:

• dp2：

• tp2:

• pp2:

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}