Thank you for such a comprehensive and detailed breakdown! I really appreciate the deep dive.
I just finished running some A/B tests based on your suggestions, and here are the actual results on my setup (RTX 5080 16GB, PyTorch 2.14, CUDA 13):
torch.compile: I tested eager vs. mode="default". In my WSL environment with custom checkpointing, it caused graph breaks and was actually slightly slower than eager mode.
Gradient Checkpointing: I tried reducing the checkpointed layers, but since my baseline VRAM is already at 15.0/16.0 GB, any reduction immediately caused an OOM (Out of Memory). So I had to keep all 24 layers checkpointed.
Fused Loss (Liger): I tried to implement Liger’s fused linear cross entropy. Unfortunately, PyTorch 2.14 + CUDA 13 seems to have a deep compatibility issue with Liger. It threw CUBLAS_STATUS_INTERNAL_ERROR and illegal memory access errors, even in the nightly build.
Current Status: I reverted to the native F.cross_entropy and it’s running perfectly stable at ~9,000 tokens/s with peak VRAM.
Thank you again for the thorough guide! It helped me rule out all the dead ends and find the most stable configuration for my hardware.