抢占式Cloud TPU v3-32训练语言模型的数据管理难题

针对只读PD挂载的锁文件权限问题 #

• 把数据集缓存和锁文件路径改到TPU本地可写目录：抢占式TPU的/tmp或/mnt/disks/local目录是可写的，且空间足够。在脚本开头添加环境变量配置：

  import os
  os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_cache"
  os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_transformers_cache"
  

  或者在调用load_dataset时直接指定cache_dir参数：

  dataset = load_dataset("your_dataset", cache_dir="/tmp/hf_cache")
  

  这样脚本创建锁文件时会指向可写的本地目录，避免只读PD的权限错误。

针对GCS流式缓存空间不足问题 #

• 切换缓存目录到TPU本地大容量存储：TPU v3-32的本地SSD有384GB可用空间，完全容纳250GB缓存。修改TFDA的缓存路径配置，在脚本中设置：

  from datasets import load_dataset
  dataset = load_dataset(
      "your_dataset",
      streaming=True,
      cache_dir="/mnt/disks/local/tfda_cache"
  )
  

• 提前预处理数据集并上传到GCS：在其他可挂载读写PD的VM上完成数据预处理（比如分词、转成Arrow格式），将处理好的数据集直接上传到GCS，然后脚本直接从GCS加载预处理后的数据集，无需流式缓存，节省本地空间。

额外优化建议 #

• 检查脚本中的锁文件逻辑：如果脚本有硬编码的锁文件路径，手动修改为本地可写路径，比如将创建锁文件的代码中的路径替换为/tmp/dataset.lock。
• 利用TPU本地存储的临时特性：本地存储重启后会清空，但训练过程中只要不重启TPU，缓存就可以正常使用；如果TPU被抢占，重新启动时再重新加载缓存即可，抢占式TPU本身就需要做好重启恢复的准备。

内容的提问来源于stack exchange，提问作者A.Najafi