github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Offload ARA with the layer if doing layer offloading. Add support to offload the LoRA. Still needs optimizer support

Browse Source
This commit is contained in:
Jaret Burkett
2025-10-21 06:03:27 -06:00
parent 76ce757e0c
commit 0d8a33dc16
5 changed files with 37 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/memory_management/manager_modules.py
View File

@@ -583,6 +583,8 @@ class LinearLayerMemoryManager(BaseLayerMemoryManager):
self.module.ara_lora_ref().org_forward = _mm_forward
else:
self.module.forward = _mm_forward
self.module._memory_management_device = self.manager.process_device
class ConvLayerMemoryManager(BaseLayerMemoryManager):
@@ -638,3 +640,5 @@ class ConvLayerMemoryManager(BaseLayerMemoryManager):
self.module.ara_lora_ref().org_forward = _mm_forward
else:
self.module.forward = _mm_forward
self.module._memory_management_device = self.manager.process_device