Baichuan 2#

Baichuan 2 is the new generation of large-scale open-source language models launched by Baichuan Intelligence inc.It is trained on a high-quality corpus with 2.6 trillion tokens and has achieved the best performance in authoritative Chinese and English benchmarks of the same size.

Note

You need to manually installbitsandbytes in a NeMo container in order to use the Baichuan 2 model.

We provide pre-defined recipes for finetuning a Baichuan2 7B model using NeMo 2.0 andNeMo-Run.These recipes configure arun.Partial for one of thenemo.collections.llm api functions introduced in NeMo 2.0.The recipes are hosted inbaichuan2_7b.

NeMo 2.0 Finetuning Recipes#

Note

The finetuning recipes use theSquadDataModule for thedata argument. You can replace theSquadDataModule with your custom dataset.

To import the HF model and convert to NeMo 2.0 format, run the following command (this only needs to be done once)

fromnemo.collectionsimportllmif__name__=="__main__":llm.import_ckpt(model=llm.Baichuan2Model(llm.Baichuan2Config7B()),source='hf://baichuan-inc/Baichuan2-7B-Base')

By default, the non-instruct version of the model is loaded. To load a different model, setfinetune.resume.restore_config.path=nemo://<hfmodelid> orfinetune.resume.restore_config.path=<localmodelpath>

We provide an example below on how to invoke the default recipe and override the data argument:

fromnemo.collectionsimportllmrecipe=llm.baichuan2_7b.finetune_recipe(name="baichuan2_7b_finetuning",dir=f"/path/to/checkpoints",num_nodes=1,num_gpus_per_node=8,peft_scheme='lora',# 'lora', 'none'packed_sequence=False,)# # To override the data argument# dataloader = a_function_that_configures_your_custom_dataset(#     gbs=gbs,#     mbs=mbs,#     seq_length=recipe.model.config.seq_length,# )# recipe.data = dataloader

By default, the finetuning recipe will run LoRA finetuning with LoRA applied to all linear layers in the language model.To finetune the entire model without LoRA, setpeft_scheme='none' in the recipe argument.

To finetune with sequence packing for a higher throughput, setpacked_sequence=True. Note that you may need totune the global batch size in order to achieve similar convergence.

Note

The configuration in the recipes is done using the NeMo-Runrun.Config andrun.Partial configuration objects. Please review the NeMo-Rundocumentation to learn more about its configuration and execution system.

Once you have your final configuration ready, you can execute it on any of the NeMo-Run supported executors. The simplest is the local executor, which just runs the pretraining locally in a separate process. You can use it as follows:

importnemo_runasrunrun.run(recipe,executor=run.LocalExecutor())

Additionally, you can also run it directly in the same Python process as follows:

run.run(recipe,direct=True)