tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add FSDP v2 swap memory support + QLoRA compatibility fixes (#3167)

Browse Source 
Co-authored-by: salman <salman.mohammadi@outlook.com>
This commit is contained in:
Grant Holmes (Ren)
2025-09-26 04:23:59 -05:00
committed by GitHub
parent 7fa8ac40cd
commit 850c1a5f8d
6 changed files with 71 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
tests/utils/schemas/validation/test_fsdp.py
View File

@@ -61,12 +61,50 @@ class TestFSDPValidation:
},
fsdp_version=2,
)
validated_cfg = validate_config(cfg)
assert validated_cfg.fsdp_version == 2
assert validated_cfg.fsdp_config.cpu_ram_efficient_loading is True
def test_fsdp2_cpu_offload_pin_memory_requires_offload_params(self, min_base_cfg):
cfg = min_base_cfg | DictDefault(
fsdp_config={
"cpu_offload_pin_memory": False,
"offload_params": False,
},
fsdp_version=2,
)
with pytest.raises(
ValueError,
match="FSDP2 does not support load_in_8bit or load_in_4bit with cpu_ram_efficient_loading.",
match="disabling cpu_offload_pin_memory requires enabling offload_params",
):
validate_config(cfg)
def test_fsdp1_cpu_offload_pin_memory_not_supported(self, min_base_cfg):
cfg = min_base_cfg | DictDefault(
fsdp_config={
"cpu_offload_pin_memory": False,
"offload_params": True,
},
fsdp_version=1,
)
with pytest.raises(
ValueError,
match="FSDP1 does not support disabling cpu_offload_pin_memory, please set `fsdp_version` to 2",
):
validate_config(cfg)
def test_fsdp2_cpu_offload_pin_memory_w_offload_params(self, min_base_cfg):
cfg = min_base_cfg | DictDefault(
fsdp_config={
"cpu_offload_pin_memory": False,
"offload_params": True,
},
fsdp_version=2,
)
validated_cfg = validate_config(cfg)
assert validated_cfg.fsdp_config.cpu_offload_pin_memory is False
assert validated_cfg.fsdp_config.offload_params is True
def test_fsdp_prefixes_removed(self, min_base_cfg):
cfg = min_base_cfg | DictDefault(
fsdp_config={