#export CUDA_HOME=/usr/lib/nvidia-cuda-toolkit/bin && \
#export PATH=/usr/lib/nvidia-cuda-toolkit/bin:$PATH && \
#export LD_LIBRARY_PATH=/usr/lib/cuda:$LD_LIBRARY_PATH && \
#export CUDA_VISIBLE_DEVICES=0,1,2 && \
#export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True && \
#export TORCH_USE_CUDA_DSA=1
time python3 run_clm.py \
    --model_name_or_path sberbank-ai/rugpt3small_based_on_gpt2 \
    --train_file corpus.txt \
    --per_device_train_batch_size 2 \
    --block_size 2048 \
    --dataset_config_name plain_text \
    --do_train \
    --gradient_accumulation_steps 4 \
    --gradient_checkpointing True \
    --bf16 True \
    --optim adamw_torch \
    --weight_decay 0.1 \
    --num_train_epochs 10 \
    --max_steps 50000 \
    --save_steps 20 \
    --save_total_limit 5 \
    --output_dir models/nero-ysss