869 lines
34 KiB
Python
869 lines
34 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Training run config (``_target_`` based YAML)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import yaml
|
|
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.train.utils.training_config import (
|
|
CheckpointConfig,
|
|
CrossSampleSequencePackingConfig,
|
|
DataConfig,
|
|
TrainingConfig,
|
|
DistributedConfig,
|
|
ModelTrainingConfig,
|
|
OptimizerConfig,
|
|
PerformanceConfig,
|
|
TrackerConfig,
|
|
TrainingLoopConfig,
|
|
)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_TRAINING_DIT_ARCH_OVERRIDE_KEYS = ("local_attn_size", "sink_size")
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RunConfig:
|
|
"""Parsed run config loaded from YAML."""
|
|
|
|
models: dict[str, dict[str, Any]]
|
|
method: dict[str, Any]
|
|
training: TrainingConfig
|
|
callbacks: dict[str, dict[str, Any]]
|
|
raw: dict[str, Any]
|
|
|
|
def resolved_config(self) -> dict[str, Any]:
|
|
"""Return a fully-resolved config dict with defaults.
|
|
|
|
Suitable for logging to W&B so that every parameter
|
|
(including defaults) is visible.
|
|
"""
|
|
import dataclasses
|
|
|
|
def _safe_asdict(obj: Any) -> Any:
|
|
if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
|
|
return {
|
|
f.name: _safe_asdict(getattr(obj, f.name))
|
|
for f in dataclasses.fields(obj) if not callable(getattr(obj, f.name))
|
|
}
|
|
if isinstance(obj, dict):
|
|
return {k: _safe_asdict(v) for k, v in obj.items()}
|
|
if isinstance(obj, list | tuple):
|
|
return type(obj)(_safe_asdict(v) for v in obj)
|
|
return obj
|
|
|
|
resolved: dict[str, Any] = {}
|
|
resolved["models"] = dict(self.models)
|
|
resolved["method"] = dict(self.method)
|
|
resolved["training"] = _safe_asdict(self.training)
|
|
resolved["callbacks"] = dict(self.callbacks)
|
|
return resolved
|
|
|
|
|
|
# ---- parsing helpers (kept for use by methods) ----
|
|
|
|
|
|
def _resolve_existing_file(path: str) -> str:
|
|
if not path:
|
|
return path
|
|
expanded = os.path.expanduser(path)
|
|
resolved = Path(expanded).resolve()
|
|
if not resolved.exists():
|
|
raise FileNotFoundError(f"Config file not found: {resolved}")
|
|
if not resolved.is_file():
|
|
raise ValueError(f"Expected a file path, got: {resolved}")
|
|
return str(resolved)
|
|
|
|
|
|
def _require_mapping(raw: Any, *, where: str) -> dict[str, Any]:
|
|
if not isinstance(raw, dict):
|
|
raise ValueError(f"Expected mapping at {where}, "
|
|
f"got {type(raw).__name__}")
|
|
return raw
|
|
|
|
|
|
def _require_str(raw: Any, *, where: str) -> str:
|
|
if not isinstance(raw, str) or not raw.strip():
|
|
raise ValueError(f"Expected non-empty string at {where}")
|
|
return raw
|
|
|
|
|
|
def get_optional_int(mapping: dict[str, Any], key: str, *, where: str) -> int | None:
|
|
raw = mapping.get(key)
|
|
if raw is None:
|
|
return None
|
|
if isinstance(raw, bool):
|
|
raise ValueError(f"Expected int at {where}, got bool")
|
|
if isinstance(raw, int):
|
|
return int(raw)
|
|
if isinstance(raw, float) and raw.is_integer():
|
|
return int(raw)
|
|
if isinstance(raw, str) and raw.strip():
|
|
return int(raw)
|
|
raise ValueError(f"Expected int at {where}, "
|
|
f"got {type(raw).__name__}")
|
|
|
|
|
|
def get_optional_float(mapping: dict[str, Any], key: str, *, where: str) -> float | None:
|
|
raw = mapping.get(key)
|
|
if raw is None:
|
|
return None
|
|
if isinstance(raw, bool):
|
|
raise ValueError(f"Expected float at {where}, got bool")
|
|
if isinstance(raw, int | float):
|
|
return float(raw)
|
|
if isinstance(raw, str) and raw.strip():
|
|
return float(raw)
|
|
raise ValueError(f"Expected float at {where}, "
|
|
f"got {type(raw).__name__}")
|
|
|
|
|
|
def parse_betas(raw: Any, *, where: str) -> tuple[float, float]:
|
|
if raw is None:
|
|
raise ValueError(f"Missing betas for {where}")
|
|
if isinstance(raw, tuple | list) and len(raw) == 2:
|
|
return float(raw[0]), float(raw[1])
|
|
if isinstance(raw, str):
|
|
parts = [p.strip() for p in raw.split(",") if p.strip()]
|
|
if len(parts) != 2:
|
|
raise ValueError(f"Expected betas as 'b1,b2' at {where}, "
|
|
f"got {raw!r}")
|
|
return float(parts[0]), float(parts[1])
|
|
raise ValueError(f"Expected betas as 'b1,b2' at {where}, "
|
|
f"got {type(raw).__name__}")
|
|
|
|
|
|
# ---- config convenience helpers ----
|
|
|
|
|
|
def require_positive_int(
|
|
mapping: dict[str, Any],
|
|
key: str,
|
|
*,
|
|
default: int | None = None,
|
|
where: str | None = None,
|
|
) -> int:
|
|
"""Read an int that must be > 0."""
|
|
loc = where or key
|
|
raw = mapping.get(key)
|
|
if raw is None:
|
|
if default is not None:
|
|
return default
|
|
raise ValueError(f"Missing required key {loc!r}")
|
|
val = get_optional_int(mapping, key, where=loc)
|
|
if val is None or val <= 0:
|
|
raise ValueError(f"{loc} must be a positive integer, got {raw!r}")
|
|
return val
|
|
|
|
|
|
def require_non_negative_int(
|
|
mapping: dict[str, Any],
|
|
key: str,
|
|
*,
|
|
default: int | None = None,
|
|
where: str | None = None,
|
|
) -> int:
|
|
"""Read an int that must be >= 0."""
|
|
loc = where or key
|
|
raw = mapping.get(key)
|
|
if raw is None:
|
|
if default is not None:
|
|
return default
|
|
raise ValueError(f"Missing required key {loc!r}")
|
|
val = get_optional_int(mapping, key, where=loc)
|
|
if val is None or val < 0:
|
|
raise ValueError(f"{loc} must be a non-negative integer, "
|
|
f"got {raw!r}")
|
|
return val
|
|
|
|
|
|
def require_non_negative_float(
|
|
mapping: dict[str, Any],
|
|
key: str,
|
|
*,
|
|
default: float | None = None,
|
|
where: str | None = None,
|
|
) -> float:
|
|
"""Read a float that must be >= 0."""
|
|
loc = where or key
|
|
raw = mapping.get(key)
|
|
if raw is None:
|
|
if default is not None:
|
|
return default
|
|
raise ValueError(f"Missing required key {loc!r}")
|
|
val = get_optional_float(mapping, key, where=loc)
|
|
if val is None or val < 0.0:
|
|
raise ValueError(f"{loc} must be a non-negative float, "
|
|
f"got {raw!r}")
|
|
return val
|
|
|
|
|
|
def require_choice(
|
|
mapping: dict[str, Any],
|
|
key: str,
|
|
choices: set[str] | frozenset[str],
|
|
*,
|
|
default: str | None = None,
|
|
where: str | None = None,
|
|
) -> str:
|
|
"""Read a string that must be one of *choices*."""
|
|
loc = where or key
|
|
raw = mapping.get(key)
|
|
if raw is None:
|
|
if default is not None:
|
|
if default not in choices:
|
|
raise ValueError(f"Default {default!r} not in {choices}")
|
|
return default
|
|
raise ValueError(f"Missing required key {loc!r}")
|
|
if not isinstance(raw, str) or not raw.strip():
|
|
raise ValueError(f"{loc} must be a non-empty string, "
|
|
f"got {type(raw).__name__}")
|
|
val = raw.strip().lower()
|
|
if val not in choices:
|
|
raise ValueError(f"{loc} must be one of {sorted(choices)}, "
|
|
f"got {raw!r}")
|
|
return val
|
|
|
|
|
|
def require_bool(
|
|
mapping: dict[str, Any],
|
|
key: str,
|
|
*,
|
|
default: bool | None = None,
|
|
where: str | None = None,
|
|
) -> bool:
|
|
"""Read a bool value."""
|
|
loc = where or key
|
|
raw = mapping.get(key)
|
|
if raw is None:
|
|
if default is not None:
|
|
return default
|
|
raise ValueError(f"Missing required key {loc!r}")
|
|
if not isinstance(raw, bool):
|
|
raise ValueError(f"{loc} must be a bool, "
|
|
f"got {type(raw).__name__}")
|
|
return raw
|
|
|
|
|
|
def _parse_pipeline_config(
|
|
cfg: dict[str, Any],
|
|
*,
|
|
models: dict[str, dict[str, Any]],
|
|
) -> Any:
|
|
"""Resolve PipelineConfig from the ``pipeline:`` YAML key."""
|
|
from fastvideo.configs.pipelines.base import PipelineConfig
|
|
|
|
pipeline_raw = cfg.get("pipeline")
|
|
if pipeline_raw is None:
|
|
return None
|
|
|
|
pipeline_raw, dit_arch_overrides = _split_training_dit_arch_overrides(pipeline_raw)
|
|
|
|
# Derive model_path from models.student.init_from —
|
|
# needed by PipelineConfig.from_kwargs.
|
|
model_path: str | None = None
|
|
student_cfg = models.get("student")
|
|
if student_cfg is not None:
|
|
init_from = student_cfg.get("init_from")
|
|
if init_from is not None:
|
|
model_path = str(init_from)
|
|
|
|
kwargs: dict[str, Any] = {"pipeline_config": pipeline_raw}
|
|
if model_path is not None:
|
|
kwargs["model_path"] = model_path
|
|
|
|
if isinstance(pipeline_raw, str):
|
|
kwargs["pipeline_config"] = _resolve_existing_file(pipeline_raw)
|
|
|
|
pipeline_config = PipelineConfig.from_kwargs(kwargs)
|
|
_apply_training_dit_arch_overrides(pipeline_config, dit_arch_overrides)
|
|
_resolve_dit_quant_config(pipeline_config)
|
|
return pipeline_config
|
|
|
|
|
|
def _resolve_dit_quant_config(pipeline_config: Any) -> None:
|
|
"""Resolve a string ``pipeline.dit_config.quant_config`` (e.g.
|
|
``nvfp4_qat_train``) into its registered QuantizationConfig instance.
|
|
|
|
Model construction calls ``quant_config.get_quant_method()`` inside
|
|
``LinearBase.__init__``, so a bare YAML string would crash there.
|
|
This must live in ``_parse_pipeline_config`` (not ``load_run_config``)
|
|
because ``dcp_to_diffusers`` rebuilds models from a checkpoint's raw
|
|
config by calling ``_parse_pipeline_config`` directly.
|
|
"""
|
|
dit_config = getattr(pipeline_config, "dit_config", None)
|
|
quant = getattr(dit_config, "quant_config", None)
|
|
if isinstance(quant, str):
|
|
from fastvideo.layers.quantization import (
|
|
get_quantization_config, )
|
|
dit_config.quant_config = get_quantization_config(quant)()
|
|
|
|
|
|
def _split_training_dit_arch_overrides(pipeline_raw: Any) -> tuple[Any, dict[str, Any]]:
|
|
"""Remove train-only DiT arch overrides before generic config parsing."""
|
|
if not isinstance(pipeline_raw, dict) or not isinstance(pipeline_raw.get("dit_config"), dict):
|
|
return pipeline_raw, {}
|
|
|
|
dit_raw = pipeline_raw["dit_config"]
|
|
overrides = {key: dit_raw[key] for key in _TRAINING_DIT_ARCH_OVERRIDE_KEYS if key in dit_raw}
|
|
if not overrides:
|
|
return pipeline_raw, {}
|
|
|
|
pipeline_raw = dict(pipeline_raw)
|
|
pipeline_raw["dit_config"] = {key: value for key, value in dit_raw.items() if key not in overrides}
|
|
return pipeline_raw, overrides
|
|
|
|
|
|
def _apply_training_dit_arch_overrides(pipeline_config: Any, overrides: dict[str, Any]) -> None:
|
|
"""Apply train-only DiT arch overrides after PipelineConfig construction."""
|
|
if not overrides:
|
|
return
|
|
|
|
arch_config = pipeline_config.dit_config.arch_config
|
|
for key, value in overrides.items():
|
|
if not hasattr(arch_config, key):
|
|
raise ValueError(f"pipeline.dit_config.{key} is not a valid field for {type(arch_config).__name__}")
|
|
setattr(arch_config, key, value)
|
|
if hasattr(arch_config, "__post_init__"):
|
|
arch_config.__post_init__()
|
|
|
|
|
|
def _build_training_config(
|
|
t: dict[str, Any],
|
|
*,
|
|
models: dict[str, dict[str, Any]],
|
|
pipeline_config: Any,
|
|
) -> TrainingConfig:
|
|
"""Build ``TrainingConfig`` from the nested ``training`` YAML mapping.
|
|
|
|
``preprocessed_data_type`` selects the dataloader's tensor contract:
|
|
``t2v`` carries video and text, ``t2va`` carries synchronized video,
|
|
audio, and text, ``text_only`` carries prompt conditioning, and
|
|
``mmaudio_features`` carries cached MMAudio V2A features.
|
|
"""
|
|
d = dict(t.get("distributed", {}) or {})
|
|
da = dict(t.get("data", {}) or {})
|
|
o = dict(t.get("optimizer", {}) or {})
|
|
lo = dict(t.get("loop", {}) or {})
|
|
ck = dict(t.get("checkpoint", {}) or {})
|
|
tr = dict(t.get("tracker", {}) or {})
|
|
perf = dict(t.get("performance", {}) or {})
|
|
vs = dict(t.get("vsa", {}) or {})
|
|
m = dict(t.get("model", {}) or {})
|
|
|
|
num_gpus = int(d.get("num_gpus", 1) or 1)
|
|
|
|
betas_raw = o.get("betas", "0.9,0.999")
|
|
betas = parse_betas(betas_raw, where="training.optimizer.betas")
|
|
|
|
model_path = str(t.get("model_path", "") or "")
|
|
if not model_path:
|
|
student_cfg = models.get("student")
|
|
if student_cfg is not None:
|
|
init_from = student_cfg.get("init_from")
|
|
if init_from is not None:
|
|
model_path = str(init_from)
|
|
|
|
raw_data_path = da.get("data_path", "") or ""
|
|
data_path: str | list[str] | dict[str, int]
|
|
if isinstance(raw_data_path, dict):
|
|
data_path = {str(path): int(repeat) for path, repeat in raw_data_path.items()}
|
|
elif isinstance(raw_data_path, list | tuple):
|
|
data_path = [str(path) for path in raw_data_path]
|
|
else:
|
|
data_path = str(raw_data_path)
|
|
|
|
preprocessed_data_type = str(da.get("preprocessed_data_type", "t2v") or "t2v").strip().lower()
|
|
if preprocessed_data_type not in {"t2v", "t2va", "ref2va", "text_only", "mmaudio_features"}:
|
|
raise ValueError("training.data.preprocessed_data_type must be one of "
|
|
"{'t2v', 't2va', 'ref2va', 'text_only', 'mmaudio_features'}, got "
|
|
f"{preprocessed_data_type!r}")
|
|
|
|
raw_max_presentation_tokens = da.get("max_presentation_tokens")
|
|
if raw_max_presentation_tokens is None:
|
|
max_presentation_tokens = None
|
|
elif isinstance(raw_max_presentation_tokens,
|
|
bool) or not isinstance(raw_max_presentation_tokens, int) or raw_max_presentation_tokens <= 0:
|
|
raise ValueError("training.data.max_presentation_tokens must be a positive integer, got "
|
|
f"{raw_max_presentation_tokens!r}")
|
|
else:
|
|
max_presentation_tokens = raw_max_presentation_tokens
|
|
|
|
packing_where = "training.data.cross_sample_sequence_packing"
|
|
packing_raw_value = da.get("cross_sample_sequence_packing", {})
|
|
packing_raw = _require_mapping(packing_raw_value, where=packing_where)
|
|
unknown_packing_keys = set(packing_raw) - {
|
|
"enabled",
|
|
"max_joint_tokens",
|
|
"max_samples_per_pack",
|
|
}
|
|
if unknown_packing_keys:
|
|
raise ValueError(f"Unknown key(s) at {packing_where}: {sorted(unknown_packing_keys)}")
|
|
packing_enabled = require_bool(
|
|
packing_raw,
|
|
"enabled",
|
|
default=False,
|
|
where=f"{packing_where}.enabled",
|
|
)
|
|
|
|
def _strict_optional_positive_packing_int(key: str) -> int | None:
|
|
value = packing_raw.get(key)
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
|
|
raise ValueError(f"{packing_where}.{key} must be a positive YAML integer, got {value!r}")
|
|
return value
|
|
|
|
packing_max_joint_tokens = _strict_optional_positive_packing_int("max_joint_tokens")
|
|
packing_max_samples_per_pack = _strict_optional_positive_packing_int("max_samples_per_pack")
|
|
native_shape_bucketing = require_bool(
|
|
da,
|
|
"native_shape_bucketing",
|
|
default=False,
|
|
where="training.data.native_shape_bucketing",
|
|
)
|
|
train_batch_size = int(da.get("train_batch_size", 1) or 1)
|
|
from fastvideo.dataset.sequence_capacity import validate_fixed_sequence_capacities
|
|
|
|
capacities = validate_fixed_sequence_capacities(da.get("fixed_sequence_capacities"))
|
|
if capacities:
|
|
if not native_shape_bucketing:
|
|
raise ValueError("training.data.fixed_sequence_capacities requires native_shape_bucketing: true")
|
|
raw_train_batch_size = da.get("train_batch_size", 1)
|
|
if (isinstance(raw_train_batch_size, bool) or not isinstance(raw_train_batch_size, int)
|
|
or raw_train_batch_size != 1):
|
|
raise ValueError(
|
|
"training.data.fixed_sequence_capacities requires train_batch_size to be the YAML integer 1")
|
|
if preprocessed_data_type not in {"text_only", "t2va", "ref2va"}:
|
|
raise ValueError("training.data.fixed_sequence_capacities supports only text_only, t2va and ref2va")
|
|
if packing_enabled and packing_max_joint_tokens is not None and packing_max_joint_tokens > capacities[-1]:
|
|
raise ValueError("cross_sample_sequence_packing.max_joint_tokens exceeds largest fixed_sequence_capacities "
|
|
"value")
|
|
if packing_enabled:
|
|
if packing_max_joint_tokens is None:
|
|
raise ValueError(f"{packing_where}.max_joint_tokens is required when packing is enabled")
|
|
if not native_shape_bucketing:
|
|
raise ValueError(f"{packing_where} requires training.data.native_shape_bucketing: true")
|
|
raw_train_batch_size = da.get("train_batch_size", 1)
|
|
if isinstance(raw_train_batch_size, bool) or not isinstance(raw_train_batch_size, int):
|
|
raise ValueError(f"{packing_where} requires training.data.train_batch_size to be the YAML integer 1")
|
|
if train_batch_size != 1:
|
|
raise ValueError(f"{packing_where} requires training.data.train_batch_size: 1")
|
|
if preprocessed_data_type not in {"text_only", "t2va"}:
|
|
raise ValueError(f"{packing_where} currently supports only text_only and t2va data, got "
|
|
f"{preprocessed_data_type!r}")
|
|
|
|
vsa_tile_size = int(vs.get("tile_size", 256) or 256)
|
|
if vsa_tile_size not in (64, 128, 256):
|
|
raise ValueError(f"training.vsa.tile_size must be 64, 128 or 256, got {vsa_tile_size!r}")
|
|
vsa_ref_keep_rate = get_optional_float(vs, "ref_keep_rate", where="training.vsa.ref_keep_rate")
|
|
if vsa_ref_keep_rate is not None and not 0.0 < vsa_ref_keep_rate <= 1.0:
|
|
raise ValueError("training.vsa.ref_keep_rate must be in (0, 1], got "
|
|
f"{vsa_ref_keep_rate!r}")
|
|
|
|
save_inference_checkpoint_on_validation = require_bool(
|
|
ck,
|
|
"save_inference_checkpoint_on_validation",
|
|
default=False,
|
|
where="training.checkpoint.save_inference_checkpoint_on_validation",
|
|
)
|
|
inference_checkpoint_role = str(ck.get("inference_checkpoint_role", "student") or "").strip()
|
|
if save_inference_checkpoint_on_validation and not inference_checkpoint_role:
|
|
raise ValueError("training.checkpoint.inference_checkpoint_role must be non-empty "
|
|
"when validation inference checkpointing is enabled")
|
|
inference_checkpoint_dtype = str(ck.get("inference_checkpoint_dtype", "bfloat16") or "bfloat16").strip().lower()
|
|
if inference_checkpoint_dtype not in {"bfloat16", "float16", "float32"}:
|
|
raise ValueError("training.checkpoint.inference_checkpoint_dtype must be one of "
|
|
"['bfloat16', 'float16', 'float32'], got "
|
|
f"{inference_checkpoint_dtype!r}")
|
|
inference_base_model_revision = ck.get("inference_base_model_revision")
|
|
if inference_base_model_revision is not None:
|
|
from fastvideo.configs.pipelines.minimax_h3 import parse_base_model_revision
|
|
parse_base_model_revision(inference_base_model_revision)
|
|
require_complete_training_checkpoint = require_bool(
|
|
ck,
|
|
"require_complete_training_checkpoint",
|
|
default=False,
|
|
where="training.checkpoint.require_complete_training_checkpoint",
|
|
)
|
|
|
|
peak_tflops = get_optional_float(
|
|
perf,
|
|
"peak_tflops_per_gpu",
|
|
where="training.performance.peak_tflops_per_gpu",
|
|
)
|
|
if peak_tflops is not None and peak_tflops <= 0.0:
|
|
raise ValueError("training.performance.peak_tflops_per_gpu must be "
|
|
f"> 0, got {peak_tflops}")
|
|
|
|
from fastvideo.train.utils.distributed_strategy import (
|
|
normalize_distributed_strategy, )
|
|
|
|
return TrainingConfig(
|
|
distributed=DistributedConfig(
|
|
strategy=normalize_distributed_strategy(str(d.get("strategy", "fsdp") or "fsdp")),
|
|
num_gpus=num_gpus,
|
|
tp_size=int(d.get("tp_size", 1) or 1),
|
|
sp_size=int(d.get("sp_size", num_gpus) or num_gpus),
|
|
hsdp_replicate_dim=int(d.get("hsdp_replicate_dim", 1) or 1),
|
|
hsdp_shard_dim=int(d.get("hsdp_shard_dim", num_gpus) or num_gpus),
|
|
pin_cpu_memory=bool(d.get("pin_cpu_memory", False)),
|
|
),
|
|
data=DataConfig(
|
|
data_path=data_path,
|
|
preprocessed_data_type=preprocessed_data_type,
|
|
train_batch_size=int(da.get("train_batch_size", 1) or 1),
|
|
dataloader_num_workers=int(da.get("dataloader_num_workers", 0) or 0),
|
|
training_cfg_rate=float(da.get("training_cfg_rate", 0.0) or 0.0),
|
|
seed=int(da.get("seed", 0) or 0),
|
|
num_height=int(da.get("num_height", 0) or 0),
|
|
num_width=int(da.get("num_width", 0) or 0),
|
|
num_latent_t=int(da.get("num_latent_t", 0) or 0),
|
|
num_frames=int(da.get("num_frames", 0) or 0),
|
|
native_shape_bucketing=native_shape_bucketing,
|
|
require_existing_parquet_cache=da.get("require_existing_parquet_cache", False),
|
|
fixed_sequence_capacities=list(capacities) if capacities else None,
|
|
max_presentation_tokens=max_presentation_tokens,
|
|
cross_sample_sequence_packing=CrossSampleSequencePackingConfig(
|
|
enabled=packing_enabled,
|
|
max_joint_tokens=packing_max_joint_tokens,
|
|
max_samples_per_pack=packing_max_samples_per_pack,
|
|
),
|
|
),
|
|
optimizer=OptimizerConfig(
|
|
learning_rate=float(o.get("learning_rate", 0.0) or 0.0),
|
|
betas=betas,
|
|
eps=float(o.get("eps", 1e-8) or 1e-8),
|
|
weight_decay=float(o.get("weight_decay", 0.0) or 0.0),
|
|
lr_scheduler=str(o.get("lr_scheduler", "constant") or "constant"),
|
|
lr_warmup_steps=int(o.get("lr_warmup_steps", 0) or 0),
|
|
lr_num_cycles=int(o.get("lr_num_cycles", 0) or 0),
|
|
lr_power=float(o.get("lr_power", 0.0) or 0.0),
|
|
min_lr_ratio=float(o.get("min_lr_ratio", 0.5) or 0.5),
|
|
lr_milestones=tuple(int(step) for step in (o.get("lr_milestones", []) or [])),
|
|
lr_gamma=float(o.get("lr_gamma", 0.1) or 0.1),
|
|
fused=bool(o.get("fused", False)),
|
|
),
|
|
loop=TrainingLoopConfig(
|
|
max_train_steps=int(lo.get("max_train_steps", 0) or 0),
|
|
gradient_accumulation_steps=int(lo.get("gradient_accumulation_steps", 1) or 1),
|
|
fsdp_defer_backward_finalize=require_bool(
|
|
lo,
|
|
"fsdp_defer_backward_finalize",
|
|
default=False,
|
|
where="training.loop.fsdp_defer_backward_finalize",
|
|
),
|
|
),
|
|
checkpoint=CheckpointConfig(
|
|
output_dir=str(ck.get("output_dir", "") or ""),
|
|
resume_from_checkpoint=str(ck.get("resume_from_checkpoint", "") or ""),
|
|
save_inference_checkpoint_on_validation=save_inference_checkpoint_on_validation,
|
|
inference_checkpoint_role=inference_checkpoint_role or "student",
|
|
inference_checkpoint_dtype=inference_checkpoint_dtype,
|
|
inference_base_model_revision=inference_base_model_revision,
|
|
inference_checkpoints_total_limit=ck.get("inference_checkpoints_total_limit", 0),
|
|
training_state_checkpointing_steps=int(ck.get("training_state_checkpointing_steps", 0) or 0),
|
|
require_complete_training_checkpoint=require_complete_training_checkpoint,
|
|
checkpoints_total_limit=int(ck.get("checkpoints_total_limit", 0) or 0),
|
|
checkpointing_start_step=int(ck.get("checkpointing_start_step", 0) or 0),
|
|
reset_lr_on_resume=bool(ck.get("reset_lr_on_resume", False)),
|
|
),
|
|
tracker=TrackerConfig(
|
|
trackers=list(tr.get("trackers", []) or []),
|
|
entity=str(tr.get("entity", "") or ""),
|
|
project_name=str(tr.get("project_name", "fastvideo") or "fastvideo"),
|
|
run_name=str(tr.get("run_name", "") or ""),
|
|
),
|
|
performance=PerformanceConfig(
|
|
enabled=bool(perf.get("enabled", True)),
|
|
peak_tflops_per_gpu=peak_tflops,
|
|
),
|
|
vsa_sparsity=float(vs.get("sparsity", 0.0) or 0.0),
|
|
vsa_ref_keep_rate=vsa_ref_keep_rate,
|
|
vsa_tile_size=vsa_tile_size,
|
|
vsa_cache_tile_buf=bool(vs.get("cache_tile_buf", False) or False),
|
|
model=ModelTrainingConfig(
|
|
weighting_scheme=str(m.get("weighting_scheme", "uniform") or "uniform"),
|
|
logit_mean=float(m.get("logit_mean", 0.0) or 0.0),
|
|
logit_std=float(m.get("logit_std", 1.0) or 1.0),
|
|
mode_scale=float(m.get("mode_scale", 1.0) or 1.0),
|
|
precondition_outputs=bool(m.get("precondition_outputs", False)),
|
|
moba_config=dict(m.get("moba_config", {}) or {}),
|
|
enable_gradient_checkpointing_type=(m.get("enable_gradient_checkpointing_type")),
|
|
allow_low_precision_master_weights=bool(m.get("allow_low_precision_master_weights", False)),
|
|
minimax_h3_activation_offload_min_tokens=int(m.get("minimax_h3_activation_offload_min_tokens", 0)),
|
|
minimax_h3_student_compile=m.get("minimax_h3_student_compile", "off"),
|
|
minimax_h3_compile_limit_fallback=m.get("minimax_h3_compile_limit_fallback", False),
|
|
minimax_h3_accumulated_recompile_limit=m.get("minimax_h3_accumulated_recompile_limit", 1024),
|
|
enable_torch_compile=bool(m.get("enable_torch_compile", False)),
|
|
compile_train_fn=bool(m.get("compile_train_fn", False)),
|
|
torch_compile_kwargs=dict(m.get("torch_compile_kwargs", {}) or {}),
|
|
),
|
|
pipeline_config=pipeline_config,
|
|
model_path=model_path,
|
|
dit_precision=str(t.get("dit_precision", "fp32") or "fp32"),
|
|
)
|
|
|
|
|
|
def _parse_cli_overrides(overrides: list[str], ) -> dict[str, Any]:
|
|
"""Parse ``--dotted.key value`` CLI overrides.
|
|
|
|
Returns a flat dict mapping dotted keys to parsed
|
|
values. Values are auto-cast: ``true``/``false`` to
|
|
bool, integers/floats to numbers, and YAML list
|
|
literals (``[a, b]``) to lists.
|
|
"""
|
|
result: dict[str, Any] = {}
|
|
i = 0
|
|
while i < len(overrides):
|
|
arg = overrides[i]
|
|
if not arg.startswith("--"):
|
|
raise ValueError(f"Expected --dotted.key, got {arg!r}")
|
|
key = arg[2:] # strip leading --
|
|
if "=" in key:
|
|
key, raw_val = key.split("=", 1)
|
|
else:
|
|
i += 1
|
|
if i >= len(overrides):
|
|
raise ValueError(f"Missing value for override {arg!r}")
|
|
raw_val = overrides[i]
|
|
result[key] = _cast_value(raw_val)
|
|
i += 1
|
|
return result
|
|
|
|
|
|
def _cast_value(raw: str) -> Any:
|
|
"""Best-effort cast a CLI string to a Python value."""
|
|
if raw.lower() == "true":
|
|
return True
|
|
if raw.lower() == "false":
|
|
return False
|
|
if raw.lower() in ("none", "null"):
|
|
return None
|
|
# Try int
|
|
try:
|
|
return int(raw)
|
|
except ValueError:
|
|
pass
|
|
# Try float
|
|
try:
|
|
return float(raw)
|
|
except ValueError:
|
|
pass
|
|
# Try YAML list literal like [1, 2]
|
|
if raw.startswith("[") and raw.endswith("]"):
|
|
try:
|
|
return yaml.safe_load(raw)
|
|
except yaml.YAMLError:
|
|
pass
|
|
return raw
|
|
|
|
|
|
def _apply_overrides(
|
|
cfg: dict[str, Any],
|
|
overrides: dict[str, Any],
|
|
) -> None:
|
|
"""Apply dotted-key overrides to a nested dict."""
|
|
for dotted_key, value in overrides.items():
|
|
parts = dotted_key.split(".")
|
|
d = cfg
|
|
for part in parts[:-1]:
|
|
if part not in d or not isinstance(d[part], dict):
|
|
d[part] = {}
|
|
d = d[part]
|
|
d[parts[-1]] = value
|
|
|
|
|
|
def load_run_config(
|
|
path: str,
|
|
overrides: list[str] | None = None,
|
|
) -> RunConfig:
|
|
"""Load a run config from YAML.
|
|
|
|
Expected top-level keys: ``models``, ``method``,
|
|
``training`` (nested), and optionally ``callbacks``
|
|
and ``pipeline``.
|
|
|
|
Args:
|
|
path: Path to the YAML config file.
|
|
overrides: Optional list of CLI override tokens,
|
|
e.g. ``["--training.distributed.num_gpus", "4"]``.
|
|
Dotted keys map to nested YAML paths.
|
|
"""
|
|
path = _resolve_existing_file(path)
|
|
with open(path, encoding="utf-8") as f:
|
|
raw = yaml.safe_load(f)
|
|
cfg = _require_mapping(raw, where=path)
|
|
|
|
# Apply CLI overrides before building typed config.
|
|
if overrides:
|
|
parsed = _parse_cli_overrides(overrides)
|
|
_apply_overrides(cfg, parsed)
|
|
logger.info("Applied CLI overrides: %s", parsed)
|
|
|
|
# --- models ---
|
|
models_raw = _require_mapping(cfg.get("models"), where="models")
|
|
models: dict[str, dict[str, Any]] = {}
|
|
for role, model_cfg_raw in models_raw.items():
|
|
role_str = _require_str(role, where="models.<role>")
|
|
model_cfg = _require_mapping(model_cfg_raw, where=f"models.{role_str}")
|
|
if "_target_" not in model_cfg:
|
|
raise ValueError(f"models.{role_str} must have a "
|
|
"'_target_' key")
|
|
models[role_str] = dict(model_cfg)
|
|
|
|
# --- method ---
|
|
method_raw = _require_mapping(cfg.get("method"), where="method")
|
|
if "_target_" not in method_raw:
|
|
raise ValueError("method must have a '_target_' key")
|
|
method = dict(method_raw)
|
|
|
|
# --- callbacks ---
|
|
callbacks_raw = cfg.get("callbacks", None)
|
|
if callbacks_raw is None:
|
|
callbacks: dict[str, dict[str, Any]] = {}
|
|
else:
|
|
callbacks = _require_mapping(callbacks_raw, where="callbacks")
|
|
|
|
# --- pipeline config ---
|
|
pipeline_config = _parse_pipeline_config(cfg, models=models)
|
|
|
|
# --- training config ---
|
|
training_raw = _require_mapping(cfg.get("training"), where="training")
|
|
t = dict(training_raw)
|
|
training = _build_training_config(t, models=models, pipeline_config=pipeline_config)
|
|
|
|
return RunConfig(
|
|
models=models,
|
|
method=method,
|
|
training=training,
|
|
callbacks=callbacks,
|
|
raw=cfg,
|
|
)
|
|
|
|
|
|
def parse_cfg_uncond(method_config: dict[str, Any]) -> dict[str, Any] | None:
|
|
"""Validate ``method.cfg_uncond``, the unconditional-branch policy per channel.
|
|
|
|
Shared by DMD2 and the PDD family, whose teachers apply the same
|
|
classifier-free guidance contract (``{"text": "zero"}`` for MiniMax-H3).
|
|
``None`` (unset) means the model's default negative-prompt conditioning.
|
|
"""
|
|
raw = method_config.get("cfg_uncond")
|
|
if raw is None:
|
|
return None
|
|
if not isinstance(raw, dict):
|
|
raise ValueError("method_config.cfg_uncond must be a dict "
|
|
f"when set, got {type(raw).__name__}")
|
|
|
|
cfg: dict[str, Any] = dict(raw)
|
|
|
|
on_missing_raw = cfg.get("on_missing", "error")
|
|
if on_missing_raw is None:
|
|
on_missing_raw = "error"
|
|
if not isinstance(on_missing_raw, str):
|
|
raise ValueError("method_config.cfg_uncond.on_missing must "
|
|
"be a string, got "
|
|
f"{type(on_missing_raw).__name__}")
|
|
on_missing = on_missing_raw.strip().lower()
|
|
if on_missing not in {"error", "ignore"}:
|
|
raise ValueError("method_config.cfg_uncond.on_missing must "
|
|
"be one of {error, ignore}, got "
|
|
f"{on_missing_raw!r}")
|
|
cfg["on_missing"] = on_missing
|
|
|
|
for channel, policy_raw in list(cfg.items()):
|
|
if channel == "on_missing":
|
|
continue
|
|
if policy_raw is None:
|
|
continue
|
|
if not isinstance(policy_raw, str):
|
|
raise ValueError("method_config.cfg_uncond values must "
|
|
"be strings, got "
|
|
f"{channel}="
|
|
f"{type(policy_raw).__name__}")
|
|
policy = policy_raw.strip().lower()
|
|
allowed = {"keep", "zero", "drop"}
|
|
if channel == "text":
|
|
allowed = {*allowed, "negative_prompt"}
|
|
if policy not in allowed:
|
|
raise ValueError("method_config.cfg_uncond values must "
|
|
"be one of "
|
|
f"{sorted(allowed)}, got "
|
|
f"{channel}={policy_raw!r}")
|
|
cfg[channel] = policy
|
|
|
|
return cfg
|
|
|
|
|
|
def parse_rollout_mode(method_config: dict[str, Any], *, where: str = "method.rollout_mode") -> str:
|
|
"""Parse how a distillation method obtains the student's input point.
|
|
|
|
``simulate`` starts from fresh noise and lets the student create its own
|
|
trajectory, so it can run on text-only data; ``data_latent`` perturbs
|
|
preprocessed VAE latents at a sampled point. Returns the canonical name.
|
|
"""
|
|
raw = method_config.get("rollout_mode")
|
|
if raw is None:
|
|
raise ValueError(f"{where} must be set: 'simulate' (the student's own trajectory from fresh noise) "
|
|
"or 'data_latent' (noised real latents)")
|
|
if not isinstance(raw, str):
|
|
raise ValueError(f"{where} must be a string, got {type(raw).__name__}")
|
|
mode = raw.strip().lower()
|
|
if mode in ("simulate", "sim"):
|
|
return "simulate"
|
|
if mode in ("data_latent", "data", "vae_latent"):
|
|
return "data_latent"
|
|
raise ValueError(f"{where} must be one of {{simulate, data_latent}}, got {raw!r}")
|
|
|
|
|
|
def validate_text_only_rollout(preprocessed_data_type: Any, rollout_mode: str) -> None:
|
|
"""Text-only data carries no latents, so only simulated rollouts can consume it."""
|
|
if str(preprocessed_data_type).strip().lower() == "text_only" and rollout_mode != "simulate":
|
|
raise ValueError("training.data.preprocessed_data_type='text_only' requires "
|
|
"method.rollout_mode='simulate'; data_latent rollout requires vae_latent data.")
|
|
|
|
|
|
def uses_negative_prompt_conditioning(cfg_uncond: dict[str, Any] | None) -> bool:
|
|
"""Whether the unconditional branch needs an encoded negative prompt."""
|
|
if cfg_uncond is None:
|
|
return True
|
|
text_policy = cfg_uncond.get("text", None)
|
|
if text_policy is None:
|
|
return True
|
|
return str(text_policy).strip().lower() == "negative_prompt"
|
|
|
|
|
|
def parse_torch_dtype(
|
|
mapping: dict[str, Any],
|
|
key: str,
|
|
*,
|
|
default: str,
|
|
where: str | None = None,
|
|
) -> Any:
|
|
"""Read a floating-point torch dtype by short (``fp32``/``bf16``) or torch name."""
|
|
import torch
|
|
|
|
from fastvideo.utils import PRECISION_TO_TYPE
|
|
|
|
loc = where or key
|
|
raw = mapping.get(key, default)
|
|
if isinstance(raw, torch.dtype):
|
|
dtype = raw
|
|
else:
|
|
if not isinstance(raw, str) or not raw.strip():
|
|
raise ValueError(f"{loc} must name a torch dtype, got {raw!r}")
|
|
name = raw.strip().lower().removeprefix("torch.")
|
|
dtype = PRECISION_TO_TYPE.get(name) or getattr(torch, name, None)
|
|
if not isinstance(dtype, torch.dtype) or not dtype.is_floating_point:
|
|
raise ValueError(f"{loc} must name a floating-point torch dtype (fp32, bf16, fp16, float64, ...), "
|
|
f"got {raw!r}")
|
|
return dtype
|