384 lines
12 KiB
Python
384 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""CPU-only unit tests for :func:`load_run_config`."""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
from fastvideo.train.utils.config import RunConfig, load_run_config
|
|
from fastvideo.train.utils.training_config import TrainingConfig
|
|
|
|
|
|
def _write_yaml(tmp_path: Path, data: dict[str, Any]) -> str:
|
|
path = tmp_path / "run.yaml"
|
|
path.write_text(yaml.safe_dump(data), encoding="utf-8")
|
|
return str(path)
|
|
|
|
|
|
def _minimal_yaml() -> dict[str, Any]:
|
|
return {
|
|
"models": {
|
|
"student": {
|
|
"_target_": "fastvideo.train.models.wan.WanModel",
|
|
"init_from": "fake/model",
|
|
},
|
|
},
|
|
"method": {
|
|
"_target_": "fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod",
|
|
},
|
|
"training": {},
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Happy path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_minimal_yaml_loads_happy_path(tmp_path: Path) -> None:
|
|
cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
|
|
|
|
assert isinstance(cfg, RunConfig)
|
|
assert isinstance(cfg.training, TrainingConfig)
|
|
assert cfg.method["_target_"] == ("fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod")
|
|
assert "student" in cfg.models
|
|
assert cfg.callbacks == {}
|
|
# raw retains the original YAML dict for downstream logging.
|
|
assert "models" in cfg.raw and "method" in cfg.raw
|
|
|
|
|
|
def test_minimal_yaml_applies_all_defaults(tmp_path: Path) -> None:
|
|
cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
|
|
t = cfg.training
|
|
|
|
assert t.distributed.num_gpus == 1
|
|
assert t.distributed.tp_size == 1
|
|
assert t.distributed.sp_size == 1
|
|
assert t.distributed.hsdp_replicate_dim == 1
|
|
assert t.distributed.pin_cpu_memory is False
|
|
|
|
assert t.data.train_batch_size == 1
|
|
assert t.data.dataloader_num_workers == 0
|
|
assert t.data.training_cfg_rate == 0.0
|
|
assert t.data.seed == 0
|
|
|
|
assert t.optimizer.learning_rate == 0.0
|
|
assert t.optimizer.betas == (0.9, 0.999)
|
|
assert t.optimizer.weight_decay == 0.0
|
|
assert t.optimizer.lr_scheduler == "constant"
|
|
assert t.optimizer.min_lr_ratio == 0.5
|
|
|
|
assert t.loop.max_train_steps == 0
|
|
assert t.loop.gradient_accumulation_steps == 1
|
|
|
|
assert t.checkpoint.output_dir == ""
|
|
assert t.checkpoint.checkpoints_total_limit == 0
|
|
|
|
assert t.tracker.trackers == []
|
|
assert t.tracker.project_name == "fastvideo"
|
|
assert t.tracker.run_name == ""
|
|
|
|
assert t.model.weighting_scheme == "uniform"
|
|
assert t.model.precondition_outputs is False
|
|
assert t.model.moba_config == {}
|
|
|
|
assert t.dit_precision == "fp32"
|
|
assert t.vsa_sparsity == 0.0
|
|
assert t.pipeline_config is None
|
|
|
|
|
|
def test_full_yaml_populates_all_training_fields(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
data["training"] = {
|
|
"distributed": {
|
|
"num_gpus": 4,
|
|
"tp_size": 2,
|
|
"sp_size": 2,
|
|
"hsdp_replicate_dim": 2,
|
|
"hsdp_shard_dim": 2,
|
|
"pin_cpu_memory": True,
|
|
},
|
|
"data": {
|
|
"data_path": "/some/path",
|
|
"train_batch_size": 2,
|
|
"dataloader_num_workers": 4,
|
|
"training_cfg_rate": 0.1,
|
|
"seed": 42,
|
|
"num_height": 256,
|
|
"num_width": 512,
|
|
"num_latent_t": 8,
|
|
"num_frames": 33,
|
|
},
|
|
"optimizer": {
|
|
"learning_rate": 1e-4,
|
|
"betas": [0.9, 0.95],
|
|
"weight_decay": 0.01,
|
|
"lr_scheduler": "cosine",
|
|
"lr_warmup_steps": 100,
|
|
"min_lr_ratio": 0.1,
|
|
},
|
|
"loop": {
|
|
"max_train_steps": 1000,
|
|
"gradient_accumulation_steps": 4,
|
|
},
|
|
"checkpoint": {
|
|
"output_dir": "/out",
|
|
"training_state_checkpointing_steps": 50,
|
|
"checkpoints_total_limit": 3,
|
|
},
|
|
"tracker": {
|
|
"trackers": ["wandb"],
|
|
"project_name": "myproj",
|
|
"run_name": "myrun",
|
|
},
|
|
"vsa": {
|
|
"sparsity": 0.5
|
|
},
|
|
"model": {
|
|
"weighting_scheme": "logit_normal",
|
|
"logit_mean": 0.5,
|
|
"logit_std": 1.5,
|
|
"precondition_outputs": True,
|
|
},
|
|
"dit_precision": "bf16",
|
|
}
|
|
cfg = load_run_config(_write_yaml(tmp_path, data))
|
|
t = cfg.training
|
|
|
|
assert t.distributed.num_gpus == 4
|
|
assert t.distributed.tp_size == 2
|
|
assert t.distributed.pin_cpu_memory is True
|
|
|
|
assert t.data.train_batch_size == 2
|
|
assert t.data.data_path == "/some/path"
|
|
assert t.data.num_frames == 33
|
|
assert t.data.seed == 42
|
|
|
|
assert t.optimizer.learning_rate == pytest.approx(1e-4)
|
|
assert t.optimizer.betas == (0.9, 0.95)
|
|
assert t.optimizer.lr_scheduler == "cosine"
|
|
assert t.optimizer.lr_warmup_steps == 100
|
|
assert t.optimizer.min_lr_ratio == pytest.approx(0.1)
|
|
|
|
assert t.loop.max_train_steps == 1000
|
|
assert t.loop.gradient_accumulation_steps == 4
|
|
|
|
assert t.checkpoint.output_dir == "/out"
|
|
assert t.checkpoint.checkpoints_total_limit == 3
|
|
|
|
assert t.tracker.trackers == ["wandb"]
|
|
assert t.tracker.project_name == "myproj"
|
|
|
|
assert t.vsa_sparsity == pytest.approx(0.5)
|
|
assert t.model.weighting_scheme == "logit_normal"
|
|
assert t.model.precondition_outputs is True
|
|
assert t.dit_precision == "bf16"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Schema validation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_missing_models_raises(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
del data["models"]
|
|
with pytest.raises(ValueError, match="models"):
|
|
load_run_config(_write_yaml(tmp_path, data))
|
|
|
|
|
|
def test_missing_method_raises(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
del data["method"]
|
|
with pytest.raises(ValueError, match="method"):
|
|
load_run_config(_write_yaml(tmp_path, data))
|
|
|
|
|
|
def test_missing_training_raises(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
del data["training"]
|
|
with pytest.raises(ValueError, match="training"):
|
|
load_run_config(_write_yaml(tmp_path, data))
|
|
|
|
|
|
def test_models_role_without_target_raises(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
data["models"]["student"] = {"init_from": "fake/model"}
|
|
with pytest.raises(ValueError, match="_target_"):
|
|
load_run_config(_write_yaml(tmp_path, data))
|
|
|
|
|
|
def test_method_without_target_raises(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
data["method"] = {"some_param": 5}
|
|
with pytest.raises(ValueError, match="_target_"):
|
|
load_run_config(_write_yaml(tmp_path, data))
|
|
|
|
|
|
def test_missing_config_file_raises(tmp_path: Path) -> None:
|
|
with pytest.raises(FileNotFoundError):
|
|
load_run_config(str(tmp_path / "does_not_exist.yaml"))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Special parsing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("betas_value, expected", [
|
|
([0.8, 0.9], (0.8, 0.9)),
|
|
("0.9,0.999", (0.9, 0.999)),
|
|
])
|
|
def test_betas_parses_list_and_string_forms(
|
|
tmp_path: Path,
|
|
betas_value: Any,
|
|
expected: tuple[float, float],
|
|
) -> None:
|
|
data = _minimal_yaml()
|
|
data["training"] = {"optimizer": {"betas": betas_value}}
|
|
cfg = load_run_config(_write_yaml(tmp_path, data))
|
|
assert cfg.training.optimizer.betas == expected
|
|
|
|
|
|
def test_pipeline_quant_config_resolves_for_load_and_export(tmp_path: Path) -> None:
|
|
from fastvideo.layers.quantization.nvfp4_qat_train_config import (
|
|
NVFP4QATTrainConfig, )
|
|
from fastvideo.train.entrypoint.dcp_to_diffusers import (
|
|
_run_config_from_raw, )
|
|
|
|
data = _minimal_yaml()
|
|
data["models"]["student"]["init_from"] = ("FastVideo/LTX2-Distilled-Diffusers")
|
|
data["pipeline"] = {"dit_config": {"quant_config": "nvfp4_qat_train"}}
|
|
|
|
cfg = load_run_config(_write_yaml(tmp_path, data))
|
|
assert isinstance(cfg.training.pipeline_config.dit_config.quant_config, NVFP4QATTrainConfig)
|
|
|
|
export_cfg = _run_config_from_raw(cfg.raw)
|
|
assert isinstance(
|
|
export_cfg.training.pipeline_config.dit_config.quant_config,
|
|
NVFP4QATTrainConfig,
|
|
)
|
|
|
|
|
|
def test_pipeline_quant_config_rejects_unknown_name(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
data["models"]["student"]["init_from"] = ("FastVideo/LTX2-Distilled-Diffusers")
|
|
data["pipeline"] = {"dit_config": {"quant_config": "not_a_quantization_method"}}
|
|
|
|
with pytest.raises(ValueError, match="Invalid quantization method"):
|
|
load_run_config(_write_yaml(tmp_path, data))
|
|
|
|
|
|
def test_data_path_mapping_parses_repeat_counts(tmp_path: Path) -> None:
|
|
# Config loading should preserve structured multi-dataset paths so the
|
|
# dataset layer can interpret repeat counts later.
|
|
data = _minimal_yaml()
|
|
data["training"] = {
|
|
"data": {
|
|
"data_path": {
|
|
"data/path1": 1,
|
|
"data/path2": 2,
|
|
}
|
|
}
|
|
}
|
|
|
|
cfg = load_run_config(_write_yaml(tmp_path, data))
|
|
|
|
assert cfg.training.data.data_path == {
|
|
"data/path1": 1,
|
|
"data/path2": 2,
|
|
}
|
|
|
|
|
|
def test_dotted_override_replaces_mapping_data_path(tmp_path: Path) -> None:
|
|
# A dict-valued data_path is a single leaf for overrides: a scalar
|
|
# --training.data.data_path replaces the whole mapping.
|
|
data = _minimal_yaml()
|
|
data["training"] = {
|
|
"data": {
|
|
"data_path": {
|
|
"data/path1": 1,
|
|
"data/path2": 2,
|
|
}
|
|
}
|
|
}
|
|
|
|
cfg = load_run_config(
|
|
_write_yaml(tmp_path, data),
|
|
overrides=["--training.data.data_path=data/only"],
|
|
)
|
|
|
|
assert cfg.training.data.data_path == "data/only"
|
|
|
|
|
|
def test_dotted_overrides_apply_with_type_coercion(tmp_path: Path) -> None:
|
|
path = _write_yaml(tmp_path, _minimal_yaml())
|
|
overrides = [
|
|
"--training.distributed.num_gpus=4",
|
|
"--training.optimizer.learning_rate=1e-3",
|
|
"--training.distributed.pin_cpu_memory=true",
|
|
"--training.tracker.project_name=overridden",
|
|
]
|
|
cfg = load_run_config(path, overrides=overrides)
|
|
|
|
assert cfg.training.distributed.num_gpus == 4
|
|
assert cfg.training.optimizer.learning_rate == pytest.approx(1e-3)
|
|
assert cfg.training.distributed.pin_cpu_memory is True
|
|
assert cfg.training.tracker.project_name == "overridden"
|
|
|
|
|
|
def test_dotted_overrides_accept_separate_value_token(tmp_path: Path) -> None:
|
|
path = _write_yaml(tmp_path, _minimal_yaml())
|
|
cfg = load_run_config(
|
|
path,
|
|
overrides=["--training.distributed.num_gpus", "8"],
|
|
)
|
|
assert cfg.training.distributed.num_gpus == 8
|
|
|
|
|
|
def test_overrides_create_intermediate_keys(tmp_path: Path) -> None:
|
|
"""Overrides into a nested key absent from YAML should still apply."""
|
|
data = _minimal_yaml()
|
|
# No `training.checkpoint` block in the minimal YAML.
|
|
path = _write_yaml(tmp_path, data)
|
|
cfg = load_run_config(
|
|
path,
|
|
overrides=["--training.checkpoint.checkpoints_total_limit=5"],
|
|
)
|
|
assert cfg.training.checkpoint.checkpoints_total_limit == 5
|
|
|
|
|
|
def test_hsdp_shard_dim_defaults_to_num_gpus(tmp_path: Path) -> None:
|
|
"""When unset, hsdp_shard_dim and sp_size fall back to num_gpus."""
|
|
data = _minimal_yaml()
|
|
data["training"] = {"distributed": {"num_gpus": 4}}
|
|
cfg = load_run_config(_write_yaml(tmp_path, data))
|
|
assert cfg.training.distributed.hsdp_shard_dim == 4
|
|
assert cfg.training.distributed.sp_size == 4
|
|
|
|
|
|
def test_model_path_derived_from_student_init_from(tmp_path: Path) -> None:
|
|
cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
|
|
assert cfg.training.model_path == "fake/model"
|
|
|
|
|
|
def test_callbacks_default_to_empty_when_absent(tmp_path: Path) -> None:
|
|
cfg = load_run_config(_write_yaml(tmp_path, _minimal_yaml()))
|
|
assert cfg.callbacks == {}
|
|
|
|
|
|
def test_callbacks_passed_through_when_present(tmp_path: Path) -> None:
|
|
data = _minimal_yaml()
|
|
data["callbacks"] = {
|
|
"grad_clip": {
|
|
"_target_": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback",
|
|
"max_grad_norm": 1.0,
|
|
},
|
|
}
|
|
cfg = load_run_config(_write_yaml(tmp_path, data))
|
|
assert "grad_clip" in cfg.callbacks
|
|
assert cfg.callbacks["grad_clip"]["max_grad_norm"] == 1.0
|