TPU v5e에서 PyTorch를 사용하여 Llama 3 학습
이 튜토리얼에서는 WikiText 데이터 세트를 사용하여 TPU v5e에서 PyTorch/XLA를 사용하여 Llama-3-8B 모델을 학습시키는 방법을 설명합니다. 모델 세부정보는 Meta-Llama-3-8B를 참고하세요.
Llama-3-8B 모델은 Hugging Face 플랫폼에 호스팅됩니다.
Meta-Llama-3-8B에는 두 가지 버전이 있습니다. 하나는 Transformer와 함께 사용하기 위한 버전이고 다른 하나는 원래 Llama 3 코드베이스와 함께 사용하기 위한 버전입니다. 이 튜토리얼에서는 다음과 같은 이유로 Transformers 버전을 사용합니다.
Hugging Face 생태계와 원활하게 통합: 이를 통해 모델을 더 쉽게 미세 조정하고, 사전 빌드된 파이프라인을 사용하고, 방대한 데이터 세트 및 도구 모음에 액세스할 수 있습니다.
유연성 및 맞춤설정 가능: Transformers 버전은 모델을 미세 조정하고 배포하는 데 유용한 유연성과 맞춤설정 옵션을 제공합니다.
커뮤니티 지원 제공: Hugging Face 커뮤니티는 Transformers 모델 사용에 관한 광범위한 문서, 튜토리얼, 지원을 제공합니다.
Transformers에 관한 자세한 내용은 Hugging Face Transformers 문서를 참고하세요.
가중치 및 토큰 생성기를 다운로드하는 등 Meta-Llama-3-8B 모델에 액세스하고 이를 사용하려면 Hugging Face 사용자 액세스 토큰이 필요합니다. 토큰은 다음을 제공합니다.
인증 및 승인: 액세스 토큰은 사용자 인증 정보 역할을 하며 Hugging Face 서버가 모델 리소스에 대한 액세스를 승인할 수 있도록 합니다. 이렇게 하면 승인된 사용자만 모델을 다운로드하고 사용할 수 있습니다.
보안: Hugging Face는 액세스 토큰을 사용하여 모델을 보호하고 무단 액세스 또는 오용을 방지합니다.
이 튜토리얼에서 액세스 토큰을 만들고 사용하는 방법에 관한 자세한 내용은 모델 실행을 참고하세요. 액세스 토큰 생성 및 사용에 관한 포괄적인 정보는 사용자 액세스 토큰에 관한 Hugging Face 문서를 참고하세요.
Hugging Face의 Llama 3 8B 모델에 액세스하려면 권한도 필요합니다. 이 권한을 얻으려면 Hugging Face의 Meta-Llama-3-8B 모델로 이동하여 액세스를 요청하세요.
TPU v5litepod-16 프로비저닝 준비
이 튜토리얼은 다음 Cloud TPU 환경 변수를 사용하여 테스트되었습니다. 가속기 유형, 영역, 런타임 버전이 호환되는 한 다른 변수를 사용하여 TPU를 프로비저닝할 수 있습니다.
예를 들어 이 튜토리얼에서는 europe-west4-b
이 영역으로 사용됩니다. 실행 중인 TPU 버전 (가속기 유형)(이 튜토리얼에서는 v5litepod-16)을 지원하는 다른 영역을 사용할 수 있습니다.
다음 TPU VM 환경 변수를 설정합니다.
export TPU_NAME=queued-resources-node-id #The TPU name is the queued resource node-id export PROJECT_ID=your-project-id export ACCELERATOR_TYPE=v5litepod-16 export ZONE=europe-west4-b export RUNTIME_VERSION=v2-alpha-tpuv5-lite export QUEUED_RESOURCE_ID=queued-resource-id export VALID_UNTIL_DURATION=1d
Hugging Face에서 Meta-Llama-3-8B 모델에 액세스할 수 있으면 튜토리얼을 실행할 TPU 환경을 준비합니다.
Cloud TPU 환경 설정 가이드를 따라 Cloud TPU를 사용할 수 있는 적절한 액세스 권한이 있는지 확인합니다.
TPU VM의 서비스 ID를 만듭니다.
gcloud alpha compute tpus tpu-vm service-identity create --zone=zone
TPU 서비스 계정을 만들고 Google Cloud 서비스에 대한 액세스 권한을 부여합니다.
서비스 계정은 Google Cloud TPU 서비스가 다른 Google Cloud서비스에 액세스하도록 허용합니다. 사용자 관리형 서비스 계정이 권장됩니다. Google Cloud 콘솔 또는
gcloud
명령어를 통해 서비스 계정을 만들 수 있습니다.gcloud
명령줄 도구를 사용하여 서비스 계정을 만듭니다.gcloud iam service-accounts create your-service-account-name \ --description="your-sa-description" \ --display-name="your-sa-display-name" export SERVICE_ACCOUNT_NAME=your-service-account-name
Google Cloud 콘솔에서 서비스 계정을 만듭니다.
- Google Cloud Console의 서비스 계정 페이지로 이동합니다.
- 서비스 계정 만들기를 클릭합니다.
- 서비스 계정 이름을 입력합니다.
- (선택사항) 서비스 계정에 대한 설명을 입력합니다.
- 만들기를 클릭하고 계속 진행합니다.
- 서비스 계정에 부여할 역할을 선택합니다.
- 계속을 클릭합니다.
- (선택사항) 서비스 계정을 관리할 수 있는 사용자 또는 그룹을 지정합니다.
- 완료를 클릭하여 서비스 계정 만들기를 마칩니다.
서비스 계정을 만든 후 다음 단계에 따라 서비스 계정 역할을 부여합니다.
다음 역할이 필요합니다.
- TPU 관리자: TPU를 만드는 데 필요합니다.
- 스토리지 관리자: Cloud Storage에 액세스하는 데 필요합니다.
- 로그 작성자
- 모니터링 측정항목 작성자: Cloud Monitoring에 측정항목을 쓰는 데 필요합니다.
사용자에게 IAM 역할을 할당하려면 관리자가 나에게
roles/resourcemanager.projectIamAdmin
역할을 부여해야 합니다. 프로젝트 IAM 관리자roles/resourcemanager.projectIamAdmin
역할이 있는 사용자도 이 역할을 부여할 수 있습니다.다음
gcloud
명령어를 사용하여 서비스 계정 역할을 추가합니다.gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/tpu.admin gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/storage.admin gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/logging.logWriter gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/monitoring.metricWriter
Google Cloud 콘솔을 사용하여 역할을 할당할 수도 있습니다.
Google Cloud 콘솔에서 다음 역할을 선택합니다.
- 서비스 계정을 선택하고 주 구성원 추가를 클릭합니다.
- 새 주 구성원 필드에 서비스 계정의 이메일 주소를 입력합니다.
- 역할 선택 드롭다운에서 역할을 검색하여 (예: 스토리지 관리자) 선택합니다.
- 저장을 클릭합니다.
Google Cloud 로 인증하고 Google Cloud CLI의 기본 프로젝트 및 ZONE을 구성합니다.
gcloud auth login gcloud config set project PROJECT_ID gcloud config set compute/zone ZONE
용량 확보
TPU 용량을 확보할 준비가 되면 할당량 페이지에서 Cloud Quotas 시스템에 관해 알아보세요. 용량 확보에 관해 추가 질문이 있는 경우 Cloud TPU 영업팀 또는 계정팀에 문의하세요.
Cloud TPU 환경 프로비저닝
TPU VM은 GKE, GKE 및 XPK 또는 대기열에 추가된 리소스로 프로비저닝할 수 있습니다.
기본 요건
- 이 튜토리얼은 Python 3.10 이상으로 테스트되었습니다.
- 프로젝트에Google Cloud 프로젝트 내에서 액세스할 수 있는 최대 칩 수를 지정하는
TPUS_PER_TPU_FAMILY
할당량이 충분한지 확인합니다. - 프로젝트에 다음에 충분한 TPU 할당량이 있는지 확인합니다.
- TPU VM 할당량
- IP 주소 할당량
- 하이퍼디스크 균형 할당량
- 사용자 프로젝트 권한
- XPK와 함께 GKE를 사용하는 경우 XPK를 실행하는 데 필요한 권한은 사용자 또는 서비스 계정의 Cloud 콘솔 권한을 참고하세요.
TPU v5litepod-16 프로비저닝
TPU VM을 만듭니다.
gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \ --node-id=${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --accelerator-type=${ACCELERATOR_TYPE} \ --runtime-version=${RUNTIME_VERSION} \ --service-account=${SERVICE_ACCOUNT_NAME} \ --spot
TPU가
ACTIVE
상태인지 확인합니다.gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
TPU가 활성화되면 (ACTIVE
) 다음과 유사한 출력이 표시됩니다.
createTime: '2025-02-28T21:16:08.053492925Z'
name: projects/my-project/locations/zone/queuedResources/tpu-name-zone
spot: {}
state:
state: ACTIVE
tpu:
nodeSpec:
- node:
acceleratorType: v5litepod-16
networkConfig:
enableExternalIps: true
network: default
queuedResource: projects/19672137403/locations/zone/queuedResources/qr-name
runtimeVersion: v2-alpha-tpuv5-lite
schedulingConfig: {}
my-service-account@your-project-id.iam.gserviceaccount.com
email: 19672137854-compute@developer.iam.gserviceaccount.com
shieldedInstanceConfig: {}
nodeId: tpu-name
parent: projects/19672137403/locations/zone
설치
Hugging Face Transformers의 pytorch-tpu/transformers
포크 및 종속 항목을 설치합니다. 이 튜토리얼은 다음 종속 항목 버전으로 테스트되었습니다.
torch
: 2.6.0과 호환됨torch_xla[tpu]
: 2.6.0과 호환됨jax
: 0.4.38jaxlib
: 0.4.38
프레임워크 소프트웨어 및 종속 항목 설치
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone -b flash_attention https://github.com/pytorch-tpu/transformers.git sudo apt install python3.10-venv python -m venv /home/$USER/venv/ source ~/venv/bin/activate cd transformers pip3 install --user -e . pip3 install datasets pip3 install evaluate pip3 install scikit-learn pip3 install accelerate pip install torch~=2.6.0 torch_xla[tpu]~=2.6.0 -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html pip install jax==0.4.38 jaxlib==0.4.38 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html'
설치가 완료되면 다음과 유사한 출력이 표시됩니다.
Collecting jax==0.4.38
Downloading jax-0.4.38-py3-none-any.whl (2.1 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.1/2.1 MB 18.0 MB/s eta 0:00:00
Collecting jaxlib==0.4.38
Downloading jaxlib-0.4.38-cp310-cp310-manylinux2014_x86_64.whl (85.0 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 85.0/85.0 MB 10.1 MB/s eta 0:00:00
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0
Collecting opt-einsum
Downloading opt_einsum-3.4.0-py3-none-any.whl (71 kB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 71.9/71.9 KB 186.4 kB/s eta 0:00:00
Requirement already satisfied: numpy>=1.24 in /home/your-username/.local/lib/python3.10/site-packages (from jax==0.4.38) (2.2.3)
Requirement already satisfied: scipy>=1.10 in /home/your-username/.local/lib/python3.10/site-packages (from jax==0.4.38) (1.15.2)
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0
Collecting ml-dtypes>=0.2.0
Downloading ml_dtypes-0.5.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.7 MB)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.7/4.7 MB 13.8 MB/s eta 0:00:00
Installing collected packages: opt-einsum, ml-dtypes, jaxlib, jax
Successfully installed jax-0.4.38 jaxlib-0.4.38 ml-dtypes-0.5.1 opt-einsum-3.4.0
모델 구성 설정
다음 섹션의 학습 명령어인 모델 실행은 두 개의 JSON 구성 파일을 사용하여 모델 매개변수와 FSDP (전체 샤딩 데이터 병렬) 구성을 정의합니다. FSDP 샤딩은 학습 중에 모델 가중치가 더 큰 배치 크기에 맞게 조정되는 데 사용됩니다. 소형 모델로 학습하는 경우 데이터 병렬 처리를 사용하고 각 기기에서 가중치를 복제하는 것으로 충분할 수 있습니다. PyTorch/XLA에서 기기 간에 텐서를 샤딩하는 방법에 관한 자세한 내용은 PyTorch/XLA SPMD 사용자 가이드를 참고하세요.
이 명령어는 Llama3-8B의 모델 매개변수 구성 파일을 만듭니다. 다른 모델의 경우 Hugging Face에서 구성을 찾습니다. 예를 보려면 Llama2-7B 구성을 참고하세요.
cat > llama-config.json <<EOF { "architectures": [ "LlamaForCausalLM" ], "attention_bias": false, "attention_dropout": 0.0, "bos_token_id": 128000, "eos_token_id": 128001, "hidden_act": "silu", "hidden_size": 4096, "initializer_range": 0.02, "intermediate_size": 14336, "max_position_embeddings": 8192, "model_type": "llama", "num_attention_heads": 32, "num_hidden_layers": 32, "num_key_value_heads": 8, "pretraining_tp": 1, "rms_norm_eps": 1e-05, "rope_scaling": null, "rope_theta": 500000.0, "tie_word_embeddings": false, "torch_dtype": "bfloat16", "transformers_version": "4.40.0.dev0", "use_cache": false, "vocab_size": 128256 } EOF
FSDP 구성 파일을 만듭니다.
cat > fsdp-config.json <<EOF { "fsdp_transformer_layer_cls_to_wrap": [ "LlamaDecoderLayer" ], "xla": true, "xla_fsdp_v2": true, "xla_fsdp_grad_ckpt": true } EOF
FSDP에 관한 자세한 내용은 FSDPv2를 참고하세요.
다음 명령어를 사용하여 구성 파일을 TPU VM에 업로드합니다.
ssh-add ~/.ssh/google_compute_engine #Setup SSH Key in the SSH agent. gcloud alpha compute tpus tpu-vm scp llama-config.json fsdp-config.json $TPU_NAME:. \ --worker=all \ --project=${PROJECT_ID} \ --zone=${ZONE}
이 명령어는 다음과 유사한 출력을 생성합니다.
Using scp batch size of 4.Attempting to SCP into 1 nodes with a total of 4 workers. SCP: Attempting to connect to worker 0... SCP: Attempting to connect to worker 1... SCP: Attempting to connect to worker 2... SCP: Attempting to connect to worker 3... llama-config.json 100% 707 4.1KB/s 00:00 llama-config.json 100% 707 4.0KB/s 00:00 llama-config.json 100% 707 4.1KB/s 00:00 llama-config.json 100% 707 4.1KB/s 00:00 fsdp-config.json 100% 156 0.9KB/s 00:00 fsdp-config.json 100% 156 0.9KB/s 00:00 fsdp-config.json 100% 156 0.9KB/s 00:00 fsdp-config.json 100% 156 0.9KB/s 00:00
모델 실행
이전 섹션에서 만든 구성 파일을 사용하여 run_clm.py
스크립트를 실행하여 WikiText 데이터 세트에서 Llama 3 8B 모델을 학습합니다. 학습 스크립트는 TPU v5litepod-16에서 실행되는 데 약 10분 정도 걸립니다.
아직 Hugging Face 토큰이 없다면 새 Hugging Face 토큰을 생성합니다.
- 내 프로필 > 설정 > 액세스 토큰을 클릭합니다.
- 새 토큰을 선택합니다.
- 원하는 이름과 Read 이상의 역할을 지정합니다.
- 토큰 생성을 선택합니다.
Hugging Face 토큰을 사용하여 다음 명령어를 통해 TPU VM에서 Hugging Face에 로그인합니다.
huggingface-cli login
토큰 변수를 이전 단계에서 Hugging Face에서 생성된 토큰으로 바꿉니다.gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' pip install -U "huggingface_hub[cli]" export PATH="/home/$USER/.local/bin/:$PATH" huggingface-cli login --token hf_abcxyzEFg'
이 명령어를 사용하면 Hugging Face에 로그인하고 현재 활성 토큰을 표시할 수 있습니다.
모델 학습을 실행합니다.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' source ~/venv/bin/activate export PJRT_DEVICE=TPU export XLA_USE_SPMD=1 export ENABLE_PJRT_COMPATIBILITY=true # Optional variables for debugging: export XLA_IR_DEBUG=1 export XLA_HLO_DEBUG=1 export PROFILE_EPOCH=0 export PROFILE_STEP=3 export PROFILE_DURATION_MS=100000 # Set PROFILE_LOGDIR to a local VM path or gs://my-bucket/profile_path export PROFILE_LOGDIR=your-bucket/profile_path python3 transformers/examples/pytorch/language-modeling/run_clm.py \ --dataset_name wikitext \ --dataset_config_name wikitext-2-raw-v1 \ --per_device_train_batch_size 16 \ --do_train \ --output_dir /home/$USER/tmp/test-clm \ --overwrite_output_dir \ --config_name /home/$USER/llama-config.json \ --cache_dir /home/$USER/cache \ --tokenizer_name meta-llama/Meta-Llama-3-8B \ --block_size 8192 \ --optim adafactor \ --save_strategy no \ --logging_strategy no \ --fsdp "full_shard" \ --fsdp_config /home/$USER/fsdp-config.json \ --torch_dtype bfloat16 \ --dataloader_drop_last yes \ --flash_attention \ --max_steps 20'
학습 단계는 약 10분 정도 소요됩니다. 학습이 끝나갈 때 다음과 유사한 메시지가 표시됩니다.
[INFO|trainer.py:2053] 2025-03-18 22:05:02,536 >> ***** Running training *****
[INFO|trainer.py:2054] 2025-03-18 22:05:02,536 >> Num examples = 272
[INFO|trainer.py:2055] 2025-03-18 22:05:02,536 >> Num Epochs = 2
[INFO|trainer.py:2056] 2025-03-18 22:05:02,536 >> Instantaneous batch size per device = 16
[INFO|trainer.py:2059] 2025-03-18 22:05:02,536 >> Total train batch size (w. parallel, distributed & accumulation) = 16
[INFO|trainer.py:2060] 2025-03-18 22:05:02,536 >> Gradient Accumulation steps = 1
[INFO|trainer.py:2061] 2025-03-18 22:05:02,536 >> Total optimization steps = 20
[INFO|trainer.py:2062] 2025-03-18 22:05:02,537 >> Number of trainable parameters = 8,030,261,248
0%| | 0/20 [00:00<?, ?it/s][INFO|trainer.py:2143] 2025-03-18 22:05:02,540 >> Profiling server started: <_XLAC.profiler.ProfilerServer object at 0x7f01bdcb6770>
5%|▌ | 1/20 [00:07<02:29, 7.86s/it]/home/your-username/venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1810: UserWarning: For backward hooks to be called, module output should be a Tensor or a tuple of Tensors but received <class 'transformers.modeling_outputs.CausalLMOutputWithPast'>
5%|▌ | 1/20 [00:07<02:29, 7.89s/it]Compilation at Step 0, time: 213.83555555343628
/home/your-username/venv/lib/python3.10/site-packages/torch/nn/modules/module.py:1810:
10%|█ | 2/20 [03:43<38:57, 129.87s/it]Compilation at Step 0, time: 213.12156581878662
/home/your-username/venv/lib/python3.10/site-packages/torch/nn/modules/module.py:"
10%|█ | 2/20 [03:40<38:29, 128.30s/it]Compilation at Step 1, time: 224.5414960384369
15%|█▌ | 3/20 [07:22<48:31, 171.24s/it]Compilation at Step 1, time: 226.23664164543152
15%|█▌ | 3/20 [07:26<48:56, 172.73s/it]Compilation at Step 1, time: 226.9180543422699
Compilation at Step 1, time: 224.3874273300171
20%|██ | 4/20 [07:23<27:45, 104.10s/it]Starting to trace for 100000 ms. Remaining attempt(s): 2
2025-03-18 22:12:32.104419: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 847930 nanoseconds and will start immediately.
Starting to trace for 100000 ms. Remaining attempt(s): 2
2025-03-18 22:12:32.104373: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 763960 nanoseconds and will start immediately.
Starting to trace for 100000 ms. Remaining attempt(s): 2
2025-03-18 22:12:32.104538: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 854020 nanoseconds and will start immediately.
2025-03-18 22:12:32.104347: W external/tsl/tsl/profiler/lib/profiler_session.cc:109] Profiling is late by 761070 nanoseconds and will start immediately.
Starting to trace for 100000 ms. Remaining attempt(s): 2
85%|████████▌ | 17/20 [07:55<00:06, 2.26s/it]Compilation at Step -1, time: 3.676558494567871
Compilation at Step -1, time: 3.447533130645752
Compilation at Step -1, time: 3.5890843868255615
Compilation at Step -1, time: 3.4956483840942383
100%|██████████| 20/20 [11:39<00:00, 35.14s/it][INFO|trainer.py:2350] 2025-03-18 22:16:42,476 >>
Training completed. Do not forget to share your model on huggingface.co/models =)
100%|██████████| 20/20 [11:47<00:00, 35.23s/it][INFO|trainer.py:2350] 2025-03-18 22:16:43,239 >>
Training completed. Do not forget to share your model on huggingface.co/models =)
삭제
학습이 완료되면 다음 단계에 따라 대기열에 추가된 리소스와 TPU VM을 삭제합니다. 이렇게 하면 TPU VM 사용에 대한 청구가 중단됩니다.
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --force \ --async