trending/miles
1
0
Fork 0
mirror of https://github.com/radixark/miles.git synced 2026-10-02 07:14:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

bugfix on deepep init without offloading

Browse Source
This commit is contained in:
Zilin Zhu
2025-08-28 19:32:01 -07:00
parent 552ba47f5d
commit 0eff732497
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
slime/backends/megatron_utils/__init__.py
+4 -2
View File
@@ -8,10 +8,12 @@ try:
old_init = deep_ep.Buffer.__init__
def new_init(self, *args, **kwargs):
torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(False)
if torch_memory_saver._impl is not None:
torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(False)
old_init(self, *args, **kwargs)
torch.cuda.synchronize()
torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(True)
if torch_memory_saver._impl is not None:
torch_memory_saver._impl._binary_wrapper.cdll.tms_set_interesting_region(True)
deep_ep.Buffer.__init__ = new_init
except ImportError: