TPU로 LLM 학습하기 (feat. TRC, MaxText)

gcloud alpha compute tpus queued-resources create $QUEUED_RESOURCE_ID \
    --node-id $NODE_ID \
    --project $PROJECT_ID \
    --zone $REGION \
    --accelerator-type $NODE_TYPE \
    --runtime-version tpu-ubuntu2204-base \
    --spot \ // 스팟 노드를 만들고자 한다면 추가.
    --scopes=https://www.googleapis.com/auth/devstorage.read_only,https://www.googleapis.com/auth/logging.write,https://www.googleapis.com/auth/monitoring.write,https://www.googleapis.com/auth/pubsub,https://www.googleapis.com/auth/service.management.readonly,https://www.googleapis.com/auth/servicecontrol,https://www.googleapis.com/auth/trace.append,https://www.googleapis.com/auth/devstorage.full_control

gcloud alpha compute tpus tpu-vm ssh $node --zone=$region --project=$PROJECT_ID --worker=0

gcloud alpha compute tpus tpu-vm ssh $node \
    --zone=$region  \
    --project=$PROJECT_ID \
    --worker=all \
    --command="screen -L -Logfile logfile.txt -d -m bash -c \"bash <(curl -sL $path)\""

git clone AI-Hypercomputer/maxtext

cd maxtext

pip install -r ./requirements.txt

# huggingface hub에서 스트리밍하는 예시다.
dataset_type: hf
hf_path: 'allenai/c4'  # for using https://huggingface.co/datasets/allenai/c4
hf_data_dir: 'en'
hf_train_files: ''
# set eval_interval > 0 to use the specified eval dataset, otherwise, only metrics on the train set will be calculated.
eval_interval: 10000
hf_eval_split: 'validation'
hf_eval_files: ''
# for HF pipeline, tokenizer_path can be a path in HuggingFace Hub, 
# or a local path containing tokenizer in a format supported by transformers.AutoTokenizer
tokenizer_path: 'google-t5/t5-large'  # for using https://huggingface.co/google-t5/t5-large
hf_access_token: ''  # provide token if using gated dataset or tokenizer

# GCS에 있는 데이터셋을 사용하는 예시다.
dataset_type: hf
hf_path: 'parquet'  # or json, arrow, etc.
hf_data_dir: ''
hf_train_files: 'gs://<bucket>/<folder>/*-train-*.parquet'   # match the train files
# set eval_interval > 0 to use the specified eval dataset. Otherwise, only metrics on the train set will be calculated.
eval_interval: 10000
hf_eval_split: ''
hf_eval_files: 'gs://<bucket>/<folder>/*-validation-*.parquet'  # match the val files
# for HF pipeline, tokenizer_path can be a path in HuggingFace Hub, 
# or a local path containing tokenizer in a format supported by transformers.AutoTokenizer
tokenizer_path: 'google-t5/t5-large'  # for using https://huggingface.co/google-t5/t5-large

per_device_batch_size: 16 # 기기 당 배치 크기.
gradient_accumulation_steps: 2 # 그래디언트 누적 스텝.
ici_fsdp_parallelism: 16 # FSDP 병렬성. ICI는 한 슬라이스 내에 있는 칩들 간의 네트워크를 의미한다.
ici_tensor_parallelism: 1 # TP 병렬성.
remat_policy: minimal # remat 정책. minimal, full 등이 있다. 적을수록 성능은 좋아지지만 메모리 사용량이 늘어난다.

model_name: llama3-8b # 모델 아키텍처.
steps: 10000 # 총 스텝 수.
checkpoint_period: 1000 # 체크포인트 저장 주기.
eval_interval: 1000 # eval 주기.

learning_rate: 6e-4 # 학습률.

async_checkpointing: true # 체크포인트 비동기 저장.

python MaxText/train.py MaxText/configs/base.yml \
    base_output_directory=/path/to/output \
    tokenizer_path=google-t5/t5-large \
    per_device_batch_size=16 \
    gradient_accumulation_steps=2 \
    ici_fsdp_parallelism=16 \
    ici_tensor_parallelism=1 \
    remat_policy=minimal \
    async_checkpointing=true \
    model_name=llama3-8b \
    steps=10000 \
    checkpoint_period=1000 \
    eval_interval=1000 \
    learning_rate=6e-4 \
    dataset_type=hf \
    hf_path=allenai/c4 \
    hf_data_dir=en \
    hf_train_files= \
    hf_eval_split=validation \
    hf_eval_files= \
    hf_access_token=YOUR_HF_ACCESS_TOKEN \
    enable_goodput_recording=false \ // IAM에서 서비스 계정에 대한 권한을 부여해야 한다. 나는 그냥 false로 두었다. 
    monitor_goodput=false \ // 위와 같음