Merge branch 'main' into nightly

This commit is contained in:
Daniel Han 2025-01-08 12:42:18 -08:00
commit 9811d45059
2 changed files with 7 additions and 5 deletions

View file

@ -39,7 +39,7 @@ triton = [
"triton @ https://github.com/woct0rdho/triton-windows/releases/download/v3.1.0-windows.post5/triton-3.1.0-cp312-cp312-win_amd64.whl ; python_version=='3.12' and platform_system == 'Windows'",
]
huggingface = [
"unsloth_zoo>=2024.12.7",
"unsloth_zoo>=2025.1.1",
"packaging",
"tyro",
"transformers>=4.46.1,!=4.47.0",
@ -285,7 +285,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3",
]
colab-new = [
"unsloth_zoo>=2024.12.7",
"unsloth_zoo>=2025.1.1",
"packaging",
"tyro",
"transformers>=4.46.1,!=4.47.0",

View file

@ -48,9 +48,11 @@ pass
# 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"
"roundup_power2_divisions:[32:256,64:128,256:64,>:32]"
# [TODO] Check why some GPUs don't work
# "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: