Changes Made

Checkpoint System:

• Added thread-safe locking for checkpoint writes to prevent corruption
• Required optimizer_state_dict validation for checkpoint integrity
• Stopped background writers in test cleanups to prevent resource leaks
• Added cache_vae_dtype config field with fp32/bf16 tradeoff options

Dataset & Caching:

• Migrated dataset caching from pickle to safetensors format for faster/safer storage
• Added cache_vae_dtype VAE precision control (fp32 safe, bf16 faster/lower VRAM)
• Added seed parameter for reproducible dataset shuffling and deterministic sampling

Performance Optimizations:

• Added TF32 matmul and cuDNN benchmark for CUDA hardware
• Added torch.compile support for UNet with dynamic shapes
• Enabled pin_memory on CUDA for faster data transfer
• Added VRAM monitoring throughout training phases
• Reset peak memory stats per epoch for accurate telemetry

Hardware Guardrails:

• Added GPU-first guard verification test (skips on CUDA-enabled machines)
• Added mixed_precision=‘no’ VRAM usage warnings for <16GB GPUs
• Better handling of CPU RNG states after map_location=“cuda” loads
• GPU memory cleanup after model offloading

Testing:

• Fix test_checkpoint.py to shutdown background writers
• Add GPU-first guard test that requires non-CUDA environment

Commit