Skip to content

TORCH_28_COMPATIBILITY.md

Latest commit

 

History

History
87 lines (60 loc) · 2.63 KB

TORCH_28_COMPATIBILITY.md

File metadata and controls

87 lines (60 loc) · 2.63 KB
Raw

PyTorch 2.8.0 + torch-xla 2.8 Compatibility Update

Summary

This document summarizes the changes made to ensure the Diff-VocSep TPU code is fully compatible with torch-xla 2.8.x.

Key Changes Made

1. Fixed Deprecated API Usage

File: validator_tpu.py

  • Issue: xm.xrt_world_size() is deprecated in torch-xla 2.8
  • Fix: Replaced with xm.xla_device_world_size()
  • Impact: Ensures compatibility with current torch-xla API

2. PJRT Compatibility Improvements

Files: train_tpu.py, inference_tpu.py, validator_tpu.py

  • Enhancement: Added PJRT environment variable setup
  • Code: 
    os.environ.setdefault("PJRT_DEVICE", "TPU")
os.environ.setdefault("XLA_USE_SPMD", "1")
  • Impact: Better compatibility with torch-xla 2.8 PJRT runtime

3. Lazy Import Pattern

File: train_tpu.py

  • Change: Converted direct torch-xla imports to lazy imports
  • Before: import torch_xla.core.xla_model as xm
  • After: xm = importlib.import_module("torch_xla.core.xla_model")
  • Benefit: Prevents import issues in non-TPU environments

4. Documentation Updates

Files: README.md, Diff_VocSep_Train_TPU.ipynb

  • Added: torch-xla 2.8+ compatibility notes
  • Added: torchrun usage recommendations
  • Example: torchrun --nproc_per_node=8 train_tpu.py --config config.yaml

Compatibility Test

A comprehensive test script test_torch_28_compatibility.py was created to validate:

✅ PJRT environment setup
✅ PyTorch 2.8.0 feature compatibility
✅ torch-xla lazy import patterns
✅ Autocast with XLA device_type
✅ Distributed training functions

Migration Notes

For Users Upgrading to torch-xla 2.8:

  1. Environment Setup: PJRT environment variables are now set automatically
  2. Training: Consider using torchrun instead of the built-in spawn mechanism
  3. Compatibility: All existing configs and scripts should work without changes

Recommended Setup:

# Install PyTorch 2.8
pip install "torch==2.8.*" "torchvision==0.23.*" "torchaudio==2.8.*"

# Install torch-xla 2.8 for TPU
pip install torch-xla[tpu]==2.8 -f https://storage.googleapis.com/libtpu-releases/index.html

# Run training (either method works)
python train_tpu.py --config config.yaml
# OR
torchrun --nproc_per_node=8 train_tpu.py --config config.yaml

Backward Compatibility

All changes maintain backward compatibility with older torch-xla versions while adding support for 2.8.x features.

Testing

Run the compatibility test to verify your environment:

python test_torch_28_compatibility.py

Expected output: 🎉 All compatibility tests passed!