multi_node.sh

torchrun --nnodes 2 --nproc_per_node 32 --master_addr algo-1 --master_port 7777 --node_rank 0 train_llama.py \
 --model_id "meta-llama/Llama-2-70b-hf" \
 --lr 5e-5 \
 --per_device_train_batch_size 16 \
 --bf16 True \
 --epochs 3

run.sh

torchrun --nproc_per_node=32 train_llama.py \
 --model_id "meta-llama/Llama-2-7b-hf" \
 --lr 5e-5 \
 --per_device_train_batch_size 16 \
 --bf16 True \
 --epochs 3

train_llama.py

from transformers import AutoTokenizer, AutoModelForCausalLM
from optimum.neuron import NeuronTrainer, NeuronTrainingArguments

# Prepare and tokenize dataset
# ....

# Load Llama model
llama = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")

# Define Hyperparameters
training_args = TrainingArguments(...)

# Create Trainer instance
trainer = Trainer(
    model=llama,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)

# Start training
trainer.train()