6f42444803
Replaces `enabled`/`ref_strategy` with unified `mode` parameter in ASFTStreamingConfig. Adds "auto", "seq", "batch", "hybrid", and "off" modes with automatic fallback logic. Implements seq_kv_cache streaming with KV cache reuse and batch microbatching support. Updates notebook defaults to use `mode="auto"` for optimal VRAM reduction. Adds comprehensive tests for mode routing, fallback behavior, and backward compatibility.
80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""CLI argument parsing tests for unsloth-cli.py."""
|
|
|
|
from pathlib import Path
|
|
import importlib.util
|
|
|
|
import pytest
|
|
|
|
|
|
def _load_cli_module():
|
|
root = Path(__file__).resolve().parents[1]
|
|
cli_path = root / "unsloth-cli.py"
|
|
spec = importlib.util.spec_from_file_location("unsloth_cli", cli_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_cli_defaults_asft():
|
|
cli = _load_cli_module()
|
|
parser = cli.build_parser()
|
|
args = parser.parse_args([])
|
|
|
|
assert args.asft is False
|
|
assert args.asft_mode == "asft"
|
|
assert args.kl_weight == 0.0
|
|
assert args.reference_policy == "disable_adapter"
|
|
assert args.asft_streaming == "off"
|
|
assert args.ref_microbatch_size is None
|
|
assert args.seq_chunk_size is None
|
|
|
|
|
|
def test_cli_asft_streaming_flag_defaults_auto():
|
|
cli = _load_cli_module()
|
|
parser = cli.build_parser()
|
|
args = parser.parse_args(["--asft_streaming"])
|
|
|
|
assert args.asft_streaming == "auto"
|
|
|
|
|
|
def test_cli_asft_streaming_value():
|
|
cli = _load_cli_module()
|
|
parser = cli.build_parser()
|
|
args = parser.parse_args(["--asft_streaming", "batch"])
|
|
|
|
assert args.asft_streaming == "batch"
|
|
|
|
|
|
def test_cli_asft_options_parsed():
|
|
cli = _load_cli_module()
|
|
parser = cli.build_parser()
|
|
args = parser.parse_args(
|
|
[
|
|
"--asft",
|
|
"--asft_mode",
|
|
"sft+kl",
|
|
"--kl_weight",
|
|
"0.2",
|
|
"--reference_policy",
|
|
"frozen_copy",
|
|
]
|
|
)
|
|
|
|
assert args.asft is True
|
|
assert args.asft_mode == "sft+kl"
|
|
assert args.kl_weight == pytest.approx(0.2)
|
|
assert args.reference_policy == "frozen_copy"
|