[torch_xla] Fix NaN training loss

[tengyifei]

Fixes #9

Fixed by switching to the Adafactor optimizer, which is what we have been using for model optimization in the past several months.

Tested:

  XLA_IR_DEBUG=1 XLA_HLO_DEBUG=1 python3
  torchprime/torch_xla_models/train.py \
      torchprime/torch_xla_models/configs/run.json

  NUM_SLICES=1 TPU_TYPE=v6e-256 launcher/run_xpk.sh \
     torchprime/torch_xla_models/train.py \
     --dataset_name wikitext \
     --dataset_config_name 'wikitext-2-raw-v1' \
     --output_dir /tmp \
     --cache_dir /tmp \
     --global_batch_size 256 \
     --logging_steps 10 \
     --max_steps 15 \
     --profile_step 5 \
     --model_id 'meta-llama/Meta-Llama-3-8B' \
     --tokenizer_name 'meta-llama/Meta-Llama-3-8B' \
     --block_size 8192 \
     --fsdp full_shard \
     --fsdp_config torchprime/torch_xla_models/configs/fsdp_config.json

[bhavya01]

I think that the AdamW optimizer should also work well with the models. That's what MaxText uses.

Okay to submit this for now but can we open a separate issue that training loss is NaN with AdamW optimizer.

[tengyifei]

SGTM