env_report.py
1 import platform 2 import sys 3 4 5 def print_env_report(): 6 """ 7 8 Prints a report of the environment. Useful for debugging and reproducibility. 9 Usage: 10 ``` 11 python -m liger_kernel.env_report 12 ``` 13 14 """ 15 print("Environment Report:") 16 print("-------------------") 17 print(f"Operating System: {platform.platform()}") 18 print(f"Python version: {sys.version.split()[0]}") 19 20 try: 21 import torch 22 23 print(f"PyTorch version: {torch.__version__}") 24 cuda_version = ( 25 torch.version.cuda if torch.cuda.is_available() else "Not available" 26 ) 27 print(f"CUDA version: {cuda_version}") 28 except ImportError: 29 print("PyTorch: Not installed") 30 print("CUDA version: Unable to query") 31 32 try: 33 import triton 34 35 print(f"Triton version: {triton.__version__}") 36 except ImportError: 37 print("Triton: Not installed") 38 39 try: 40 import transformers 41 42 print(f"Transformers version: {transformers.__version__}") 43 except ImportError: 44 print("Transformers: Not installed") 45 46 47 if __name__ == "__main__": 48 print_env_report()