已合并
[fix]multi_stream_lazy_reclaim coredump error #42370
xuyun15创建于 7月22日
[fix]multi_stream_lazy_reclaim coredump error #42370
已合并
共 2 个文件变更+78-13
| @@ -1,5 +1,8 @@ | |||
| 1 | import os | 1 | import os |
| 2 | import math | 2 | import math |
| 3 | +import subprocess | ||
| 4 | +import sys | ||
| 5 | +import textwrap | ||
| 3 | import torch | 6 | import torch |
| 4 | import torch_npu | 7 | import torch_npu |
| 5 | 8 | ||
| @@ -96,6 +99,48 @@ class TestAllocator(TestCase): | |||
| 96 | set_config = False | 99 | set_config = False |
| 97 | self.assertTrue(set_config) | 100 | self.assertTrue(set_config) |
| 98 | 101 | ||
| 102 | + def test_multi_stream_lazy_reclaim_trigger_event(self): | ||
| 103 | + code = textwrap.dedent("""\ | ||
| 104 | + import os | ||
| 105 | + os.environ["PYTORCH_NPU_ALLOC_CONF"] = "multi_stream_lazy_reclaim:True" | ||
| 106 | + | ||
| 107 | + import time | ||
| 108 | + import torch | ||
| 109 | + import torch_npu | ||
| 110 | + | ||
| 111 | + max_event_lazy_num = 512 | ||
| 112 | + shared_stream = torch.npu.Stream() | ||
| 113 | + x_list = [] | ||
| 114 | + for i in range(max_event_lazy_num): | ||
| 115 | + x_list.append(torch.empty(16, 16, device="npu", dtype=torch.bfloat16)) | ||
| 116 | + x_list[i].record_stream(shared_stream) | ||
| 117 | + | ||
| 118 | + x = torch.empty(16, 16, device="npu", dtype=torch.bfloat16) | ||
| 119 | + x.record_stream(shared_stream) | ||
| 120 | + | ||
| 121 | + with torch.npu.stream(shared_stream): | ||
| 122 | + y = x + 0.1 | ||
| 123 | + | ||
| 124 | + del x_list | ||
| 125 | + del x | ||
| 126 | + | ||
| 127 | + time.sleep(0.1) | ||
| 128 | + dumb = torch.empty(16, 16, device="npu", dtype=torch.bfloat16) | ||
| 129 | + del dumb | ||
| 130 | + """) | ||
| 131 | + | ||
| 132 | + result = subprocess.run( | ||
| 133 | + [sys.executable, "-c", code], | ||
| 134 | + capture_output=True, | ||
| 135 | + text=True, | ||
| 136 | + ) | ||
| 137 | + | ||
| 138 | + assert result.returncode == 0, ( | ||
| 139 | + f"Subprocess failed with return code {result.returncode}.\\n" | ||
| 140 | + f"stdout: {result.stdout}\\n" | ||
| 141 | + f"stderr: {result.stderr}" | ||
| 142 | + ) | ||
| 143 | + | ||
| 99 | 144 | ||
| 100 | if __name__ == '__main__': | 145 | if __name__ == '__main__': |
| 101 | run_tests() | 146 | run_tests() |
| @@ -1222,25 +1222,45 @@ public: | |||
| 1222 | TORCH_NPU_MEMORY_LOGD("Rounded size: %zu, alloc size: %zu, using %s pool on device %d", | 1222 | TORCH_NPU_MEMORY_LOGD("Rounded size: %zu, alloc size: %zu, using %s pool on device %d", |
| 1223 | size, alloc_size, pool.is_small ? "small" : "large", device); | 1223 | size, alloc_size, pool.is_small ? "small" : "large", device); |
| 1224 | 1224 | ||
| 1225 | + // Lazy reclaim: bound the pending-event backlog BEFORE selecting a block. | ||
| 1226 | + // The original code ran process_events() AFTER get_free_block(): the block | ||
| 1227 | + // returned by get_free_block() is removed from the pool but not yet marked | ||
| 1228 | + // allocated (that happens later in alloc_found_block()), so it still looks | ||
| 1229 | + // free (allocated == false, event_count == 0, empty stream_uses). If | ||
| 1230 | + // process_events() then free_block()'d an adjacent reaped block, | ||
| 1231 | + // try_merge_blocks() would coalesce and delete the block we were about to | ||
| 1232 | + // hand out -> allocator corruption under large backlogs (sum > kLazyQuerySize). | ||
| 1233 | + // Reap first, then select, so process_events() never runs while an | ||
| 1234 | + // un-finalized block is held. Call-count is kept identical to the original | ||
| 1235 | + // via the `reaped` guard, so the lazy fast path (skip when sum <= threshold | ||
| 1236 | + // and a block is found) is preserved. | ||
| 1237 | + const bool lazy_reclaim = | ||
| 1238 | + CachingAllocatorConfig::multi_stream_lazy_reclaim() && C10_LIKELY(captures_underway.empty()); | ||
| 1239 | + bool reaped = false; | ||
| 1240 | + if (lazy_reclaim) { | ||
| 1241 | + size_t sum = 0; | ||
| 1242 | + for (auto it = npu_events.begin(); it != npu_events.end(); ++it) { | ||
| 1243 | + sum += it->second.size(); | ||
| 1244 | + } | ||
| 1245 | + if (sum > kLazyQuerySize) { | ||
| 1246 | + process_events(context); | ||
| 1247 | + reaped = true; | ||
| 1248 | + } | ||
| 1249 | + } | ||
| 1250 | + | ||
| 1225 | // First, try to get a block from the existing pool. | 1251 | // First, try to get a block from the existing pool. |
| 1226 | bool block_found = | 1252 | bool block_found = |
| 1227 | // Search pool | 1253 | // Search pool |
| 1228 | get_free_block(params) || | 1254 | get_free_block(params) || |
| 1229 | // Trigger callbacks and retry search | 1255 | // Trigger callbacks and retry search |
| 1230 | (trigger_free_memory_callbacks(params) && get_free_block(params)); | 1256 | (trigger_free_memory_callbacks(params) && get_free_block(params)); |
| 1231 | - if (CachingAllocatorConfig::multi_stream_lazy_reclaim() && C10_LIKELY(captures_underway.empty())) { | 1257 | + |
| 1232 | - // Lazy process events and free memory | 1258 | + // Out of blocks: it is safe to reap now (no un-finalized block is held yet). |
| 1233 | - size_t sum = 0; | 1259 | + // Guard with `reaped` so we never process events twice in one malloc. |
| 1234 | - for (auto it = npu_events.begin(); it != npu_events.end(); ++it) { | 1260 | + if (lazy_reclaim && !block_found && !reaped) { |
| 1235 | - sum += it->second.size(); | 1261 | + process_events(context); |
| 1236 | - } | 1262 | + block_found = get_free_block(params); |
| 1237 | - if (!block_found || sum > kLazyQuerySize) { | 1263 | + } |
| 1238 | - process_events(context); | ||
| 1239 | - } | ||
| 1240 | - if (!block_found) { | ||
| 1241 | - block_found = get_free_block(params); | ||
| 1242 | - } | ||
| 1243 | - } | ||
| 1244 | // Can't reuse an existing block; try to get a new one. | 1264 | // Can't reuse an existing block; try to get a new one. |
| 1245 | if (!block_found) { | 1265 | if (!block_found) { |
| 1246 | TORCH_NPU_MEMORY_LOGD("No existing block found on device %d, attempting to allocate new block", device); | 1266 | TORCH_NPU_MEMORY_LOGD("No existing block found on device %d, attempting to allocate new block", device); |