export PROJECT=<project_name>
export REGION=<region>
export ZONE=<tpu_vm_instance_zone>
export VPC=<vpc_name>
export SUBNET=<vpc_subnet>
export TPUVM=<tpu_vm_instance_name>
export TYPE=v3-8
export IMAGE=tpu-vm-pt-2.0

gcloud compute tpus tpu-vm create ${TPUVM} \
    --zone=${ZONE} \
    --accelerator-type=${TYPE} \
    --version=${IMAGE} \
    --network=${VPC} \
    --subnetwork="projects/${PROJECT}/regions/${REGION}/subnetworks/${SUBNET}" \
    --internal-ips

gcloud alpha compute tpus tpu-vm ssh ${TPUVM} --zone=${ZONE} --tunnel-through-iap

sudo apt update -y && sudo apt upgrade -y
pip install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-nightly-cp38-cp38-linux_x86_64.whl
pip install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-nightly-cp38-cp38-linux_x86_64.whl
pip install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torchvision-nightly-cp38-cp38-linux_x86_64.whl

cd $HOME
# git clone -b v4.31-release https://github.com/huggingface/transformers.git
git clone https://github.com/huggingface/transformers.git
cd transformers
# For Python 3.8
pip install -e .
pip install datasets evaluate scikit-learn accelerate py7zr

# huggingface-cli login --token <YOUR_HF_TOKEN>
# llama2_fsdp_config.json
{
    "fsdp_transformer_layer_cls_to_wrap": [
        "LlamaDecoderLayer"
    ],
    "xla": true,
    "xla_fsdp_settings": {
        "compute_dtype": "bfloat16",
        "shard_param_on_dim_0": true,
        "pin_layout_in_collective_ops": true
    },
   "xla_fsdp_grad_ckpt": true
}

export PJRT_DEVICE=TPU
nohup python3 -u examples/pytorch/xla_spawn.py --num_cores 8 examples/pytorch/language-modeling/run_clm.py \
    --model_name_or_path "meta-llama/Llama-2-7b-hf" \
    --num_train_epochs 3 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 6 \
    --per_device_eval_batch_size 6 \
    --do_train \
    --do_eval \
    --output_dir /tmp/llama-2-7b-hf-ft-xla \
    --overwrite_output_dir \
    --cache_dir /tmp \
    --block_size 2048 \
    --optim adafactor \
    --save_strategy no \
    --logging_strategy no \
    --gradient_checkpointing \
    --fsdp "full_shard" \
    --fsdp_config ~/llama2_fsdp_config.json > run.log 2>&1 &

from transformers import AutoTokenizer, LlamaTokenizer
import transformers
import torch
model = "~/llama-2-7b-hf-ft-xla"
tokenizer = AutoTokenizer.from_pretrained(model, use_auth_token=True)
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
sequences = pipeline(
    ['I liked "Breaking Bad" and "Band of Brothers". Do you have any recommendations of other shows I might like?\n'] * 1,
    do_sample=True,
    top_k=10,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    max_length=200,
)
print(sequences)

RuntimeError: probability tensor contains either `inf`, `nan` or element < 0

[{'generated_text': 'I liked "Breaking Bad" and "Band of Brothers". Do you have any recommendations of other shows I might like?\n contends Creator smiling reminiscentoffset prophets contends contends Sheffield contends wetlandslocked maximizing maximizing WIratorct continuity=- ...'}]

Some weights of the model checkpoint at /home/hzchen/scripts/llm/gpt-ft-test were not used when initializing GPT2LMHeadModel: [<FSDP_LAYERS_OMITTED...>]
- This IS expected if you are initializing GPT2LMHeadModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing GPT2LMHeadModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of GPT2LMHeadModel were not initialized from the model checkpoint at /home/hzchen/scripts/llm/gpt-ft-test and are newly initialized: [<LAYERS_OMITTED...>]
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

# fsdp_config.yaml
compute_environment: LOCAL_MACHINE
debug: false
distributed_type: FSDP
downcast_bf16: 'no'
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_backward_prefetch_policy: NO_PREFETCH
  fsdp_cpu_ram_efficient_loading: true
  fsdp_forward_prefetch: false
  fsdp_offload_params: false
  fsdp_sharding_strategy: 1
  fsdp_state_dict_type: FULL_STATE_DICT
  fsdp_sync_module_states: true
  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
  fsdp_use_orig_params: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 8
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

nohup accelerate launch --config_file ~/fsdp_config.yaml examples/pytorch/language-modeling/run_clm.py \
    --model_name_or_path "meta-llama/Llama-2-7b-hf" \
    --num_train_epochs 3 \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 2 \
    --do_train \
    --do_eval \
    --output_dir /tmp/llama-2-7b-hf-ft-cuda \
    --overwrite_output_dir \
    --cache_dir /tmp \
    --block_size 2048 \
    --optim adafactor \
    --save_strategy no \
    --logging_strategy no \
    --gradient_checkpointing > run.log 2>&1 &