frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refine Allocator Config error message friendly (#165288)

Browse Source 
* __->__ #165288
Pull Request resolved: https://github.com/pytorch/pytorch/pull/165288
Approved by: https://github.com/albanD
This commit is contained in:
Yu, Guangye
2025-10-17 17:16:41 +00:00
committed by PyTorch MergeBot
parent 5d62b63a76
commit 4888ed440e
4 changed files with 23 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
c10/core/AllocatorConfig.cpp
View File

@ -46,7 +46,7 @@ size_t AcceleratorAllocatorConfig::roundup_power2_divisions(size_t size) {
63 - llvm::countLeadingZeros(kRoundUpPowerOfTwoStart); 63 - llvm::countLeadingZeros(kRoundUpPowerOfTwoStart);
const size_t interval_end = const size_t interval_end =
63 - llvm::countLeadingZeros(kRoundUpPowerOfTwoEnd); 63 - llvm::countLeadingZeros(kRoundUpPowerOfTwoEnd);
TORCH_CHECK( TORCH_CHECK_VALUE(
interval_end - interval_start == kRoundUpPowerOfTwoIntervals, interval_end - interval_start == kRoundUpPowerOfTwoIntervals,
"kRoundUpPowerOfTwoIntervals mismatch"); "kRoundUpPowerOfTwoIntervals mismatch");
@ -65,7 +65,7 @@ size_t AcceleratorAllocatorConfig::parseMaxSplitSize(
std::numeric_limits<size_t>::max() / kMB; std::numeric_limits<size_t>::max() / kMB;
size_t val_env = tokenizer.toSizeT(++i); size_t val_env = tokenizer.toSizeT(++i);
TORCH_CHECK( TORCH_CHECK_VALUE(
val_env >= min_allowed_split_size_mb, val_env >= min_allowed_split_size_mb,
"CachingAllocator option max_split_size_mb too small, must be >= ", "CachingAllocator option max_split_size_mb too small, must be >= ",
min_allowed_split_size_mb); min_allowed_split_size_mb);
@ -84,7 +84,7 @@ size_t AcceleratorAllocatorConfig::parseMaxNonSplitRoundingSize(
std::numeric_limits<size_t>::max() / kMB; std::numeric_limits<size_t>::max() / kMB;
size_t val_env = tokenizer.toSizeT(++i); size_t val_env = tokenizer.toSizeT(++i);
TORCH_CHECK( TORCH_CHECK_VALUE(
val_env >= min_allowed_split_size_mb, val_env >= min_allowed_split_size_mb,
"CachingAllocator option max_non_split_rounding_mb too small, must be >= ", "CachingAllocator option max_non_split_rounding_mb too small, must be >= ",
min_allowed_split_size_mb); min_allowed_split_size_mb);
@ -99,7 +99,7 @@ size_t AcceleratorAllocatorConfig::parseGarbageCollectionThreshold(
size_t i) { size_t i) {
tokenizer.checkToken(++i, ":"); tokenizer.checkToken(++i, ":");
double val_env = tokenizer.toDouble(++i); double val_env = tokenizer.toDouble(++i);
TORCH_CHECK( TORCH_CHECK_VALUE(
val_env > 0 && val_env < 1.0, val_env > 0 && val_env < 1.0,
"garbage_collect_threshold is invalid, set it in (0.0, 1.0)"); "garbage_collect_threshold is invalid, set it in (0.0, 1.0)");
garbage_collection_threshold_ = val_env; garbage_collection_threshold_ = val_env;
@ -120,7 +120,7 @@ size_t AcceleratorAllocatorConfig::parseRoundUpPower2Divisions(
size_t value_index = i; size_t value_index = i;
tokenizer.checkToken(++i, ":"); tokenizer.checkToken(++i, ":");
size_t value = tokenizer.toSizeT(++i); size_t value = tokenizer.toSizeT(++i);
TORCH_CHECK( TORCH_CHECK_VALUE(
value == 0 || llvm::isPowerOf2_64(value), value == 0 || llvm::isPowerOf2_64(value),
"For roundups, the divisions has to be power of 2 or 0 to disable roundup "); "For roundups, the divisions has to be power of 2 or 0 to disable roundup ");
@ -133,7 +133,7 @@ size_t AcceleratorAllocatorConfig::parseRoundUpPower2Divisions(
value); value);
} else { } else {
size_t boundary = tokenizer.toSizeT(value_index); size_t boundary = tokenizer.toSizeT(value_index);
TORCH_CHECK( TORCH_CHECK_VALUE(
llvm::isPowerOf2_64(boundary), llvm::isPowerOf2_64(boundary),
"For roundups, the intervals have to be power of 2 "); "For roundups, the intervals have to be power of 2 ");
@ -163,7 +163,7 @@ size_t AcceleratorAllocatorConfig::parseRoundUpPower2Divisions(
"Expected closing bracket ']' in ConfigTokenizer but reached end of config"); "Expected closing bracket ']' in ConfigTokenizer but reached end of config");
} else { // Keep this for backwards compatibility } else { // Keep this for backwards compatibility
size_t value = tokenizer.toSizeT(i); size_t value = tokenizer.toSizeT(i);
TORCH_CHECK( TORCH_CHECK_VALUE(
llvm::isPowerOf2_64(value), llvm::isPowerOf2_64(value),
"For roundups, the divisions has to be power of 2 "); "For roundups, the divisions has to be power of 2 ");
std::fill( std::fill(
@ -223,7 +223,7 @@ void AcceleratorAllocatorConfig::parseArgs(const std::string& env) {
// If a device-specific configuration parser hook is registered, it will // If a device-specific configuration parser hook is registered, it will
// check if the key is unrecognized. // check if the key is unrecognized.
if (device_config_parser_hook_) { if (device_config_parser_hook_) {
TORCH_CHECK( TORCH_CHECK_VALUE(
getKeys().find(key) != getKeys().end(), getKeys().find(key) != getKeys().end(),
"Unrecognized key '", "Unrecognized key '",
key, key,

4
c10/core/AllocatorConfig.h
View File

@ -76,7 +76,7 @@ class ConfigTokenizer {
} else if (token == "False") { } else if (token == "False") {
return false; return false;
} else { } else {
TORCH_CHECK( TORCH_CHECK_VALUE(
false, false,
"Expected 'True' or 'False' at index ", "Expected 'True' or 'False' at index ",
i, i,
@ -253,7 +253,7 @@ class C10_API AcceleratorAllocatorConfig {
device_config_parser_hook_ = std::move(hook); device_config_parser_hook_ = std::move(hook);
auto& mutable_keys = getMutableKeys(); auto& mutable_keys = getMutableKeys();
for (auto& key : keys) { for (auto& key : keys) {
TORCH_CHECK( TORCH_CHECK_VALUE(
mutable_keys.insert(key).second, mutable_keys.insert(key).second,
"Duplicated key '", "Duplicated key '",
key, key,

13
c10/cuda/CUDAAllocatorConfig.cpp
View File

@ -20,7 +20,7 @@ size_t CUDAAllocatorConfig::parseAllocatorConfig(
tokenizer.checkToken(++i, ":"); tokenizer.checkToken(++i, ":");
i++; // Move to the value after the colon i++; // Move to the value after the colon
#ifdef USE_ROCM #ifdef USE_ROCM
TORCH_CHECK( TORCH_CHECK_VALUE(
((tokenizer[i] == "native") || (tokenizer[i] == PYTORCH_TOKEN1) || ((tokenizer[i] == "native") || (tokenizer[i] == PYTORCH_TOKEN1) ||
(tokenizer[i] == PYTORCH_TOKEN2)), (tokenizer[i] == PYTORCH_TOKEN2)),
"Unknown allocator backend, " "Unknown allocator backend, "
@ -36,7 +36,7 @@ size_t CUDAAllocatorConfig::parseAllocatorConfig(
" != ", " != ",
get()->name()); get()->name());
#else // USE_ROCM #else // USE_ROCM
TORCH_CHECK( TORCH_CHECK_VALUE(
((tokenizer[i] == "native") || (tokenizer[i] == PYTORCH_TOKEN1)), ((tokenizer[i] == "native") || (tokenizer[i] == PYTORCH_TOKEN1)),
"Unknown allocator backend, " "Unknown allocator backend, "
"options are native and " PYTORCH_TOKEN1); "options are native and " PYTORCH_TOKEN1);
@ -109,7 +109,7 @@ void CUDAAllocatorConfig::parseArgs(const std::string& env) {
} else { } else {
const auto& keys = const auto& keys =
c10::CachingAllocator::AcceleratorAllocatorConfig::getKeys(); c10::CachingAllocator::AcceleratorAllocatorConfig::getKeys();
TORCH_CHECK( TORCH_CHECK_VALUE(
keys.find(key) != keys.end(), keys.find(key) != keys.end(),
"Unrecognized key '", "Unrecognized key '",
key, key,
@ -151,12 +151,12 @@ size_t CUDAAllocatorConfig::parsePinnedNumRegisterThreads(
size_t i) { size_t i) {
tokenizer.checkToken(++i, ":"); tokenizer.checkToken(++i, ":");
size_t val2 = tokenizer.toSizeT(++i); size_t val2 = tokenizer.toSizeT(++i);
TORCH_CHECK( TORCH_CHECK_VALUE(
llvm::isPowerOf2_64(val2), llvm::isPowerOf2_64(val2),
"Number of register threads has to be power of 2, got ", "Number of register threads has to be power of 2, got ",
val2); val2);
auto maxThreads = CUDAAllocatorConfig::pinned_max_register_threads(); auto maxThreads = CUDAAllocatorConfig::pinned_max_register_threads();
TORCH_CHECK( TORCH_CHECK_VALUE(
val2 <= maxThreads, val2 <= maxThreads,
"Number of register threads should be less than or equal to ", "Number of register threads should be less than or equal to ",
maxThreads, maxThreads,
@ -171,7 +171,8 @@ size_t CUDAAllocatorConfig::parsePinnedReserveSegmentSize(
size_t i) { size_t i) {
tokenizer.checkToken(++i, ":"); tokenizer.checkToken(++i, ":");
size_t val2 = tokenizer.toSizeT(++i); size_t val2 = tokenizer.toSizeT(++i);
TORCH_CHECK(val2 > 0, "Pinned reserve segment size has to be greater than 0"); TORCH_CHECK_VALUE(
val2 > 0, "Pinned reserve segment size has to be greater than 0");
m_pinned_reserve_segment_size_mb = val2; m_pinned_reserve_segment_size_mb = val2;
return i; return i;
} }

12
test/test_cuda.py
View File

@ -4582,21 +4582,21 @@ class TestCudaMallocAsync(TestCase):
reg_mem = torch.cuda.memory_stats()[key_allocated] reg_mem = torch.cuda.memory_stats()[key_allocated]
self.assertEqual(reg_mem - start_mem, nbytes) self.assertEqual(reg_mem - start_mem, nbytes)
with self.assertRaises(RuntimeError): with self.assertRaises(ValueError):
torch.cuda.memory._set_allocator_settings("foo:1,bar:2") torch.cuda.memory._set_allocator_settings("foo:1,bar:2")
with self.assertRaises(RuntimeError): with self.assertRaises(ValueError):
torch.cuda.memory._set_allocator_settings( torch.cuda.memory._set_allocator_settings(
"garbage_collection_threshold:1.2" "garbage_collection_threshold:1.2"
) )
with self.assertRaises(RuntimeError): with self.assertRaises(ValueError):
torch.cuda.memory._set_allocator_settings("max_split_size_mb:2") torch.cuda.memory._set_allocator_settings("max_split_size_mb:2")
with self.assertRaises(RuntimeError): with self.assertRaises(ValueError):
torch.cuda.memory._set_allocator_settings("release_lock_on_cudamalloc:none") torch.cuda.memory._set_allocator_settings("release_lock_on_cudamalloc:none")
with self.assertRaises(RuntimeError): with self.assertRaises(ValueError):
torch.cuda.memory._set_allocator_settings( torch.cuda.memory._set_allocator_settings(
"pinned_use_cuda_host_register:none" "pinned_use_cuda_host_register:none"
) )
@ -4606,7 +4606,7 @@ class TestCudaMallocAsync(TestCase):
"pinned_num_register_threads:none" "pinned_num_register_threads:none"
) )
with self.assertRaises(RuntimeError): with self.assertRaises(ValueError):
torch.cuda.memory._set_allocator_settings( torch.cuda.memory._set_allocator_settings(
"pinned_num_register_threads:1024" "pinned_num_register_threads:1024"
) )