scripts/check_env.py

"""Diagnose the training environment on the RunPod."""
import sys
print(f"Python: {sys.version}")

# Check transformers
try:
    import transformers
    print(f"Transformers: {transformers.__version__}")
    from transformers import TrainingArguments
    print("TrainingArguments: OK")
except Exception as e:
    print(f"Transformers error: {e}")

# Check TRL
try:
    import trl
    print(f"TRL: {trl.__version__}")
except Exception as e:
    print(f"TRL import error: {e}")

try:
    from trl import GRPOConfig, GRPOTrainer
    print("TRL GRPOTrainer: OK")
except Exception as e:
    print(f"TRL GRPOTrainer error: {e}")

# Check Unsloth GRPO
try:
    from unsloth.trainer import UnslothGRPOTrainer, UnslothGRPOConfig
    print("UnslothGRPOTrainer: OK")
except Exception as e:
    print(f"UnslothGRPOTrainer error: {e}")

# Check torch + CUDA
try:
    import torch
    print(f"Torch: {torch.__version__}, CUDA: {torch.cuda.is_available()}, GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A'}")
except Exception as e:
    print(f"Torch error: {e}")

# Check peft
try:
    import peft
    print(f"PEFT: {peft.__version__}")
except Exception as e:
    print(f"PEFT error: {e}")