423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-16 20:00:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor mllm param logic

Browse Source
This commit is contained in:
hiyouga
2025-01-10 15:41:54 +00:00
parent 5ffd8ad192
commit f6f630a1c9
10 changed files with 198 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
tests/model/model_utils/test_checkpointing.py
View File

@@ -14,6 +14,7 @@
import os
import pytest
import torch
from llamafactory.extras.misc import get_current_device
@@ -39,16 +40,11 @@ TRAIN_ARGS = {
}
def test_checkpointing_enable():
model = load_train_model(disable_gradient_checkpointing=False, **TRAIN_ARGS)
@pytest.mark.parametrize("disable_gradient_checkpointing", [False, True])
def test_vanilla_checkpointing(disable_gradient_checkpointing: bool):
model = load_train_model(disable_gradient_checkpointing=disable_gradient_checkpointing, **TRAIN_ARGS)
for module in filter(lambda m: hasattr(m, "gradient_checkpointing"), model.modules()):
assert getattr(module, "gradient_checkpointing") is True
def test_checkpointing_disable():
model = load_train_model(disable_gradient_checkpointing=True, **TRAIN_ARGS)
for module in filter(lambda m: hasattr(m, "gradient_checkpointing"), model.modules()):
assert getattr(module, "gradient_checkpointing") is False
assert getattr(module, "gradient_checkpointing") != disable_gradient_checkpointing
def test_unsloth_gradient_checkpointing():