求助：Colab TPUv2运行时因xla_latency_hiding_scheduler_rerun崩溃，T5模型训练失败

我尝试用wikitext-103-v1数据集训练google-t5/t5-base模型做通用语言处理，但无论怎么折腾，都没法在Google Colab Pro TPUv2运行时正常启动训练脚本。试过升级、降级、卸载重装各种依赖来保证兼容性，但每次解决一个问题就会冒出新问题，现在彻底卡壳了。

环境信息 #

TPUv2运行时
Python 3.10.12
Ubuntu 22.04.4 LTS
虚拟环境

遇到的问题序列 #

1. 初始导入错误 #

一开始执行脚本就遇到torch_xla导入失败：

ImportError                               Traceback (most recent call last)
<ipython-input-8-5baa930b08e9> in <cell line: 2>()
      1 import torch
----> 2 import torch_xla.core.xla_model as xm
      3 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
      4 from datasets import load_dataset
      5 

/usr/local/lib/python3.10/dist-packages/torch_xla/__init__.py in <module>
     18   sys.setdlopenflags(flags)
     19 
---> 20 import _XLAC
     21 from ._internal import tpu
     22 from .version import __version__

ImportError: /usr/local/lib/python3.10/dist-packages/_XLAC.cpython-310-x86_64-linux-gnu.so: undefined symbol: _ZN2at23SavedTensorDefaultHooks11set_tracingEb

2. 依赖版本匹配错误 #

调整PyTorch和torch_xla版本后，又遇到torchvision版本找不到的问题：

ERROR: Ignored the following yanked versions: 0.1.6, 0.1.7, 0.1.8, 0.1.9, 0.2.0, 0.2.1, 0.2.2, 0.2.2.post2, 0.2.2.post3, 0.15.0
ERROR: Could not find a version that satisfies the requirement torchvision0.15.1+cu118 (from versions: 0.12.0, 0.13.0, 0.13.1, 0.14.0, 0.14.1, 0.15.1, 0.15.2, 0.16.0, 0.16.1, 0.16.2, 0.17.0, 0.17.1, 0.17.2, 0.18.0, 0.18.1, 0.19.0)
ERROR: No matching distribution found for torchvision0.15.1+cu118

3. 最终运行时崩溃 #

解决依赖安装问题后，运行脚本先出现以下警告：

/usr/local/lib/python3.10/dist-packages/torch_xla/__init__.py:202: UserWarning: `tensorflow` can conflict with `torch-xla`. Prefer `tensorflow-cpu` when using PyTorch/XLA. To silence this warning, `pip uninstall -y tensorflow && pip install tensorflow-cpu`. If you are in a notebook environment such as Colab or Kaggle, restart your notebook runtime afterwards.
warnings.warn(

随后运行时直接崩溃重启，日志显示：

Aug 26, 2024, 3:26:40 PM    WARNING WARNING:root:kernel xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx restarted
Aug 26, 2024, 3:26:40 PM    INFO    KernelRestarter: restarting kernel (1/5), keep random ports
Aug 26, 2024, 3:26:38 PM    WARNING ERROR: Unknown command line flag 'xla_latency_hiding_scheduler_rerun'

已尝试的操作 #

以下命令并非同时执行：

• !pip uninstall cloud-tpu-client -y
• !pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 -f https://storage.googleapis.com/libtpu-releases/index.html（修正原命令中的拼写错误）
• !pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
• !pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.5.0.dev-cp310-cp310-linux_x86_64.whl' -f https://storage.googleapis.com/libtpu-releases/index.html
• !pip install torch~=2.4.0 torch_xla~=2.4.0 https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla_cuda_plugin-2.4.0-py3-none-any.whl
• !pip3 install datasets

当前依赖配置 #

执行!grep -E 'torch|torch_xla|numpy|scipy|protobuf|six'得到：

numpy                        1.26.4
protobuf                     3.20.3
scipy                        1.13.1
six                          1.16.0
torch                        2.4.0
torch-xla                    2.4.0
torch_xla_cuda_plugin        2.4.0
torchaudio                   2.4.0+cpu
torchvision                  0.19.0+cpu

当前训练脚本 #

%env source t5_env/bin/activate

from google.colab import drive
drive.mount('/content/drive')

import os
os.environ["HUGGINGFACE_TOKEN"] = "####" 
import torch
import torch_xla
import torch_xla.core.xla_model as xm
device = xm.xla_device()

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TrainingArguments, Trainer

# Set environment variable to avoid potential tokenizers parallelism issues
os.environ['TOKENIZERS_PARALLELISM'] = 'false'

# Load tokenizer and model
model_name = "google/t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=os.environ["HUGGINGFACE_TOKEN"])
model = AutoModelForSeq2SeqLM.from_pretrained(model_name, use_auth_token=os.environ["HUGGINGFACE_TOKEN"]).to(device)

# Load the dataset
dataset = load_dataset("wikitext", "wikitext-103-v1")

def preprocess_function(examples):
    inputs = examples['text']
    targets = examples['text']

    model_inputs = tokenizer(inputs, truncation=True, padding='max_length', return_tensors="pt")
    labels = tokenizer(targets, truncation=True, padding='max_length', return_tensors="pt")
    model_inputs['labels'] = labels['input_ids']
    return model_inputs

tokenized_datasets = dataset.map(preprocess_function, batched=True)
tokenized_datasets.save_to_disk("/content/drive/MyDrive/Colab Notebooks/t5base/tokenized_datasets")

# Define training arguments
training_args = TrainingArguments(
    output_dir="/content/drive/MyDrive/Colab Notebooks/t5base/checkpoints",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    save_steps=1000,  
    logging_steps=100,
    eval_steps=500,
    evaluation_strategy="steps",
    learning_rate=2e-5,
    weight_decay=0.01,
    fp16=True,
)

# Model and trainer
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) 

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
)

# Training loop
from tqdm.auto import tqdm
with tqdm(total=len(trainer.get_train_dataloader())) as pbar:
    try:
        trainer.train()
        trainer.save_model("/content/drive/MyDrive/Colab Notebooks/t5base/final_model")
        print("Model trained and saved!")
    except Exception as e:
        print(f"Training error: {e}")

我肯定漏掉了一些简单的关键点，但实在找不到解决方向，恳请帮忙排查问题。

内容的提问来源于stack exchange，提问作者therealtheruss