Star860
1352
代码
代码
Issues1692
Pull Requests2223
项目讨论
Wiki
分析
Star860
1352
已合并
[feat] support custom segment_size for expandable_segments #26372
XDaoHong创建于 2025年11月10日
[feat] support custom segment_size for expandable_segments #26372
已合并
XDaoHong创建于 2025年11月10日
从XDaoHong/pytorchv2.7.1合入到Ascend/pytorchv2.7.1
共 2 个文件变更+52-2
Atest/npu/test_custom_segment_size.py+20-0
@@ -0,0 +1,20 @@
1+import os
2+import gc
3+ 
4+import torch
5+from torch.testing._internal.common_utils import TestCase, run_tests
6+import torch_npu
7+ 
8+os.environ["PYTORCH_NPU_ALLOC_CONF"] = "expandable_segments:True,segment_size_mb:128"
9+ 
10+ 
11+class Test_expandable_segments(TestCase):
12+ def test_empty_virt_addr_cache(self):
13+ gc.collect()
14+ torch_npu.npu.empty_cache()
15+ x = torch.empty((2, 1024, 1024), device="npu", dtype=torch.float32)
16+ self.assertEqual(torch_npu.npu.memory_reserved(), 128 * 1024 * 1024)
17+ 
18+ 
19+if __name__ == '__main__':
20+ run_tests()
Mtorch_npu/csrc/core/npu/NPUCachingAllocator.cpp+32-2
@@ -95,6 +95,9 @@ constexpr size_t kRoundLarge = 2097152; // round up large allocs t
95constexpr size_t kAlignRoundLarge = 16384; // round up large allocs to 16 KB95constexpr size_t kAlignRoundLarge = 16384; // round up large allocs to 16 KB
96constexpr size_t kSmallPoolVirAddrSize = 2147483648; // 2 GB96constexpr size_t kSmallPoolVirAddrSize = 2147483648; // 2 GB
97constexpr size_t kLargePoolVirAddrSize = 10737418240; // 10 GB97constexpr size_t kLargePoolVirAddrSize = 10737418240; // 10 GB
98+constexpr size_t kMB = 1024 * 1024; // 1 MB
99+constexpr size_t k20MB = 20; // 20 MB for segmemt_size
100+constexpr size_t k512MB = 512; // 512 MB for segmemt_size
98const std::string kMinCannVersion = "8.1.RC1"; // minimum cann version which supports 1g mem 8.1.RC1101const std::string kMinCannVersion = "8.1.RC1"; // minimum cann version which supports 1g mem 8.1.RC1
99const std::string kMinDriverVersion = "25.0.RC1"; // minimum driver version which supports 1g mem 25.0.RC1102const std::string kMinDriverVersion = "25.0.RC1"; // minimum driver version which supports 1g mem 25.0.RC1
100const std::string kCannModule = "CANN"; // cann module name103const std::string kCannModule = "CANN"; // cann module name
@@ -879,6 +882,11 @@ public:
879 return instance().m_page_size_1g;882 return instance().m_page_size_1g;
880 }883 }
881 884 
885+ static size_t segment_size_mb()
886+ {
887+ return instance().m_segment_size_mb;
888+ }
889+ 
882 static CachingAllocatorConfig &instance()890 static CachingAllocatorConfig &instance()
883 {891 {
884 static CachingAllocatorConfig *s_instance = ([]() {892 static CachingAllocatorConfig *s_instance = ([]() {
@@ -899,12 +907,14 @@ private:
899 bool set_expandable_segments_flag = false;907 bool set_expandable_segments_flag = false;
900 size_t m_base_addr_aligned_size = kAlignRoundLarge;908 size_t m_base_addr_aligned_size = kAlignRoundLarge;
901 bool m_page_size_1g = false; // 新增1G页配置标志909 bool m_page_size_1g = false; // 新增1G页配置标志
910+ size_t m_segment_size_mb;
902 911 
903 CachingAllocatorConfig()912 CachingAllocatorConfig()
904 : m_max_split_size(std::numeric_limits<size_t>::max()),913 : m_max_split_size(std::numeric_limits<size_t>::max()),
905 m_garbage_collection_threshold(0),914 m_garbage_collection_threshold(0),
906 m_expandable_segments(false),915 m_expandable_segments(false),
907- m_base_addr_aligned_size(kAlignRoundLarge)916+ m_base_addr_aligned_size(kAlignRoundLarge),
917+ m_segment_size_mb(0)
908 {}918 {}
909 919 
910 void lexArgs(const char *env, std::vector<std::string> &config);920 void lexArgs(const char *env, std::vector<std::string> &config);
@@ -914,6 +924,7 @@ private:
914 size_t parseExpandableSegments(const std::vector<std::string> &config, size_t i);924 size_t parseExpandableSegments(const std::vector<std::string> &config, size_t i);
915 size_t parseAddrAlignSize(const std::vector<std::string> &config, size_t i);925 size_t parseAddrAlignSize(const std::vector<std::string> &config, size_t i);
916 size_t parsePageSize(const std::vector<std::string> &config, size_t i);926 size_t parsePageSize(const std::vector<std::string> &config, size_t i);
927+ size_t parseSegmentSizeMb(const std::vector<std::string> &config, size_t i);
917};928};
918 929 
919void CachingAllocatorConfig::lexArgs(const char *env, std::vector<std::string> &config)930void CachingAllocatorConfig::lexArgs(const char *env, std::vector<std::string> &config)
@@ -1031,6 +1042,21 @@ size_t CachingAllocatorConfig::parsePageSize(const std::vector<std::string> &con
1031 return i + 2; // 返回最后处理的索引位置1042 return i + 2; // 返回最后处理的索引位置
1032}1043}
1033 1044 
1045+size_t CachingAllocatorConfig::parseSegmentSizeMb(const std::vector<std::string> &config, size_t i)
1046+{
1047+ consumeToken(config, ++i, ':');
1048+ if (++i < config.size()) {
1049+ size_t val = static_cast<size_t>(stoi(config[i]));
1050+ TORCH_CHECK(val >= k20MB && val <= k512MB,
1051+ "CachingAllocator option segment_size_mb error, must be [20, 512], dtype is int",
1052+ OPS_ERROR(ErrCode::VALUE));
1053+ m_segment_size_mb = val * kMB;
1054+ } else {
1055+ TORCH_CHECK(false, "Error, expecting segment_size_mb value", OPS_ERROR(ErrCode::VALUE));
1056+ }
1057+ return i;
1058+}
1059+ 
1034void CachingAllocatorConfig::parseArgs(const char *env, std::set<std::string> supported_settings)1060void CachingAllocatorConfig::parseArgs(const char *env, std::set<std::string> supported_settings)
1035{1061{
1036 // If empty, set the default values1062 // If empty, set the default values
@@ -1062,6 +1088,8 @@ void CachingAllocatorConfig::parseArgs(const char *env, std::set<std::string> su
1062 i = parseAddrAlignSize(config, i);1088 i = parseAddrAlignSize(config, i);
1063 } else if (config[i] == "page_size") {1089 } else if (config[i] == "page_size") {
1064 i = parsePageSize(config, i);1090 i = parsePageSize(config, i);
1091+ } else if (config[i] == "segment_size_mb") {
1092+ i = parseSegmentSizeMb(config, i);
1065 } else {1093 } else {
1066 TORCH_CHECK(false, "Unrecognized CachingAllocator option: ", config[i], PTA_ERROR(ErrCode::PARAM));1094 TORCH_CHECK(false, "Unrecognized CachingAllocator option: ", config[i], PTA_ERROR(ErrCode::PARAM));
1067 }1095 }
@@ -2310,7 +2338,9 @@ private:
2310 return c;2338 return c;
2311 }2339 }
2312 }2340 }
2313- auto segment_size = pool->is_small ? kSmallBuffer : kLargeBuffer;2341+ auto custom_segment_size = CachingAllocatorConfig::segment_size_mb();
2342+ auto segment_size = pool->is_small ?
2343+ kSmallBuffer : (custom_segment_size > 0 ? custom_segment_size : kLargeBuffer);
2314 // 此处申请虚拟内存，segment_size是页大小，实际虚拟内存巨大2344 // 此处申请虚拟内存，segment_size是页大小，实际虚拟内存巨大
2315 if (IsMallocPage1GMem(pool->is_small)) {2345 if (IsMallocPage1GMem(pool->is_small)) {
2316 segment_size = kExtraLargeBuffer;2346 segment_size = kExtraLargeBuffer;