feat: TPU compatibility

[monologg]

Details

• Pytorch XLA, TPU VM에서 정상 작동하도록 설정
• 실제 코드에서는 resize_embedding 도 같이 했었습니다. 다음 PR에서 추가 부탁드립니다.

python3 xla_spawn.py --num_cores 8 pretrain_language_model.py \
  --output_dir test_ckpt \
  --per_device_train_batch_size 32 \
  --gradient_accumulation_steps 4 \
  --num_train_epochs 100 \
  --learning_rate 1e-4 \
  --lr_scheduler_type linear \
  --warmup_ratio 0.01 \
  --save_strategy epoch \
  --save_total_limit 5 \
  --seed 42

Screenshot

Related Issue

close #7