#python3 -m torch.distributed.launch -m debugpy llama_pretrain.py \

torchrun  --nproc_per_node $nproc_per_node  --nnodes $nnodes  --node_rank $node_rank  --master_addr $master_addr  --master_port $master_port llama_pretrain/llama_pretrain.py \
    --model_type llama \
    --config_overrides "num_attention_heads=32,num_hidden_layers=1,num_key_value_heads=1" \
    --tokenizer_name ../llama_tokenizer \
    --train_file llama_pretrain/data/pretrain_data.txt \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --bf16 True \
    --overwrite_output_dir \
    --do_train \
    --do_eval \
    --logging_strategy steps \
    --logging_steps 10 \
    --output_dir llama_pretrain/tmp \
    --save_strategy no \
    --num_train_epochs 1

    # --deepspeed llama_pretrain/ds_config.json \
