diff --git a/unsloth/__init__.py b/unsloth/__init__.py index afd255dc35..90d2a63519 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -55,7 +55,12 @@ else: pass # Reduce VRAM usage by reducing fragmentation -os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,roundup_power2_divisions:[64:128,256:64,>:32]" +# And optimize pinning of memory +os.environ["PYTORCH_CUDA_ALLOC_CONF"] = \ + "expandable_segments:True,"\ + "roundup_power2_divisions:[32:256,64:128,256:64,>:32],"\ + "pinned_use_cuda_host_register:True,"\ + "pinned_num_register_threads:8" # Hugging Face Hub faster downloads if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: