Blame for trainer/trainer_args.py

# SPDX-License-Identifier: Apache-2.0
# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/server_args.py
"""The arguments of Trainer Inference."""

import argparse
import dataclasses
from contextlib import contextmanager
from dataclasses import field
from enum import Enum
from typing import Any

from trainer.configs.configs import PreprocessConfig
from trainer.configs.pipelines.base import PipelineConfig, STA_Mode
from trainer.configs.utils import clean_cli_args
from trainer.logger import init_logger
from trainer.platforms import current_platform
from trainer.utils import FlexibleArgumentParser, StoreBoolean

logger = init_logger(__name__)


class ExecutionMode(str, Enum):
    """
    Enumeration for different pipeline modes.
    
    Inherits from str to allow string comparison for backward compatibility.
    """
    INFERENCE = "inference"
    PREPROCESS = "preprocess"
    FINETUNING = "finetuning"
    DISTILLATION = "distillation"

    @classmethod
    def from_string(cls, value: str) -> "ExecutionMode":
        """Convert string to ExecutionMode enum."""
        try:
            return cls(value.lower())
        except ValueError:
            raise ValueError(
                f"Invalid mode: {value}. Must be one of: {', '.join([m.value for m in cls])}"
            ) from None

    @classmethod
    def choices(cls) -> list[str]:
        """Get all available choices as strings for argparse."""
        return [mode.value for mode in cls]


class WorkloadType(str, Enum):
    """
    Enumeration for different workload types.
    
    Inherits from str to allow string comparison for backward compatibility.
    """
    I2V = "i2v"  # Image to Video
    T2V = "t2v"  # Text to Video
    T2I = "t2i"  # Text to Image
    I2I = "i2i"  # Image to Image

    @classmethod
    def from_string(cls, value: str) -> "WorkloadType":
        """Convert string to WorkloadType enum."""
        try:
            return cls(value.lower())
        except ValueError:
            raise ValueError(
                f"Invalid workload type: {value}. Must be one of: {', '.join([m.value for m in cls])}"
            ) from None

    @classmethod
    def choices(cls) -> list[str]:
        """Get all available choices as strings for argparse."""
        return [workload.value for workload in cls]


# args for trainer framework
@dataclasses.dataclass
class TrainerArgs:
    # Model and path configuration (for convenience)
    model_path: str

    # Running mode
    mode: ExecutionMode = ExecutionMode.INFERENCE

    # Workload type
    workload_type: WorkloadType = WorkloadType.T2V

    # Cache strategy
    cache_strategy: str = "none"

    # Distributed executor backend
    distributed_executor_backend: str = "mp"

    inference_mode: bool = True  # if False == training mode

    # HuggingFace specific parameters
    trust_remote_code: bool = False
    revision: str | None = None

    # Parallelism
    num_gpus: int = 1
    tp_size: int = -1
    sp_size: int = -1
    hsdp_replicate_dim: int = 1
    hsdp_shard_dim: int = -1
    dist_timeout: int | None = None  # timeout for torch.distributed

    pipeline_config: PipelineConfig = field(default_factory=PipelineConfig)
    preprocess_config: PreprocessConfig | None = None

    # LoRA parameters
    # (Wenxuan) prefer to keep it here instead of in pipeline config to not make it complicated.
    lora_path: str | None = None
    lora_nickname: str = "default"  # for swapping adapters in the pipeline
    # can restrict layers to adapt, e.g. ["q_proj"]
    # Will adapt only q, k, v, o by default.
    lora_target_modules: list[str] | None = None

    output_type: str = "pil"

    # CPU offload parameters
    dit_cpu_offload: bool = True
    use_fsdp_inference: bool = True
    text_encoder_cpu_offload: bool = True
    image_encoder_cpu_offload: bool = True
    vae_cpu_offload: bool = True
    pin_cpu_memory: bool = True

    # STA (Sliding Tile Attention) parameters
    mask_strategy_file_path: str | None = None
    STA_mode: STA_Mode = STA_Mode.STA_INFERENCE
    skip_time_steps: int = 15

    # Compilation
    enable_torch_compile: bool = False

    disable_autocast: bool = False

    # VSA parameters
    VSA_sparsity: float = 0.0  # inference/validation sparsity

    # Master port for distributed training/inference
    master_port: int | None = None

    # Stage verification
    enable_stage_verification: bool = True

    # Prompt text file for batch processing
    prompt_txt: str | None = None

    # model paths for correct deallocation
    model_paths: dict[str, str] = field(default_factory=dict)
    model_loaded: dict[str, bool] = field(default_factory=lambda: {
        "transformer": True,
        "vae": True,
    })

    cls_name: str | None = None
    load_from_dir: str | None = None
    module_name: str| None = None
    ar_action_load_from_dir: str | None = None

    # # DMD parameters
    # dmd_denoising_steps: List[int] | None = field(default=None)

    # MoE parameters used by Wan2.2
    boundary_ratio: float | None = None

    @property
    def training_mode(self) -> bool:
        return not self.inference_mode

    def __post_init__(self):
        self.check_trainer_args()

    @staticmethod
    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
        # Model and path configuration
        parser.add_argument(
            "--model-path",
            type=str,
            help=
            "The path of the model weights. This can be a local folder or a Hugging Face repo ID.",
        )
        parser.add_argument(
            "--model-dir",
            type=str,
            help="Directory containing StepVideo model",
        )

        # Running mode
        parser.add_argument(
            "--mode",
            type=str,
            choices=ExecutionMode.choices(),
            default=TrainerArgs.mode.value,
            help="The mode to run Trainer",
        )

        # Workload type
        parser.add_argument(
            "--workload-type",
            type=str,
            choices=WorkloadType.choices(),
            default=TrainerArgs.workload_type.value,
            help="The workload type",
        )

        # distributed_executor_backend
        parser.add_argument(
            "--distributed-executor-backend",
            type=str,
            choices=["mp"],
            default=TrainerArgs.distributed_executor_backend,
            help="The distributed executor backend to use",
        )

        parser.add_argument(
            "--inference-mode",
            action=StoreBoolean,
            default=TrainerArgs.inference_mode,
            help="Whether to use inference mode",
        )

        # HuggingFace specific parameters
        parser.add_argument(
            "--trust-remote-code",
            action=StoreBoolean,
            default=TrainerArgs.trust_remote_code,
            help="Trust remote code when loading HuggingFace models",
        )
        parser.add_argument(
            "--revision",
            type=str,
            default=TrainerArgs.revision,
            help=
            "The specific model version to use (can be a branch name, tag name, or commit id)",
        )

        # Parallelism
        parser.add_argument(
            "--num-gpus",
            type=int,
            default=TrainerArgs.num_gpus,
            help="The number of GPUs to use.",
        )
        parser.add_argument(
            "--tp-size",
            type=int,
            default=TrainerArgs.tp_size,
            help="The tensor parallelism size.",
        )
        parser.add_argument(
            "--sp-size",
            type=int,
            default=TrainerArgs.sp_size,
            help="The sequence parallelism size.",
        )
        parser.add_argument(
            "--hsdp-replicate-dim",
            type=int,
            default=TrainerArgs.hsdp_replicate_dim,
            help="The data parallelism size.",
        )
        parser.add_argument(
            "--hsdp-shard-dim",
            type=int,
            default=TrainerArgs.hsdp_shard_dim,
            help="The data parallelism shards.",
        )
        parser.add_argument(
            "--dist-timeout",
            type=int,
            default=TrainerArgs.dist_timeout,
            help="Set timeout for torch.distributed initialization.",
        )

        # Output type
        parser.add_argument(
            "--output-type",
            type=str,
            default=TrainerArgs.output_type,
            choices=["pil"],
            help="Output type for the generated video",
        )

        # Prompt text file for batch processing
        parser.add_argument(
            "--prompt-txt",
            type=str,
            default=TrainerArgs.prompt_txt,
            help=
            "Path to a text file containing prompts (one per line) for batch processing",
        )

        # STA (Sliding Tile Attention) parameters
        parser.add_argument(
            "--STA-mode",
            type=str,
            default=TrainerArgs.STA_mode.value,
            choices=[mode.value for mode in STA_Mode],
            help=
            "STA mode contains STA_inference, STA_searching, STA_tuning, STA_tuning_cfg, None",
        )
        parser.add_argument(
            "--skip-time-steps",
            type=int,
            default=TrainerArgs.skip_time_steps,
            help="Number of time steps to warmup (full attention) for STA",
        )
        parser.add_argument(
            "--mask-strategy-file-path",
            type=str,
            help="Path to mask strategy JSON file for STA",
        )
        parser.add_argument(
            "--enable-torch-compile",
            action=StoreBoolean,
            default=TrainerArgs.enable_torch_compile,
            help="Use torch.compile to speed up DiT inference." +
            "However, will likely cause precision drifts. See (https://github.com/pytorch/pytorch/issues/145213)",
        )

        parser.add_argument(
            "--dit-cpu-offload",
            action=StoreBoolean,
            help=
            "Use CPU offload for DiT inference. Enable if run out of memory with FSDP.",
        )
        parser.add_argument(
            "--use-fsdp-inference",
            action=StoreBoolean,
            help=
            "Use FSDP for inference by sharding the model weights. Latency is very low due to prefetch--enable if run out of memory.",
        )
        parser.add_argument(
            "--text-encoder-cpu-offload",
            action=StoreBoolean,
            help=
            "Use CPU offload for text encoder. Enable if run out of memory.",
        )
        parser.add_argument(
            "--image-encoder-cpu-offload",
            action=StoreBoolean,
            help=
            "Use CPU offload for image encoder. Enable if run out of memory.",
        )
        parser.add_argument(
            "--vae-cpu-offload",
            action=StoreBoolean,
            help="Use CPU offload for VAE. Enable if run out of memory.",
        )
        parser.add_argument(
            "--pin-cpu-memory",
            action=StoreBoolean,
            help=
            "Pin memory for CPU offload. Only added as a temp workaround if it throws \"CUDA error: invalid argument\". "
            "Should be enabled in almost all cases",
        )
        parser.add_argument(
            "--disable-autocast",
            action=StoreBoolean,
            help=
            "Disable autocast for denoising loop and vae decoding in pipeline sampling",
        )

        # VSA parameters
        parser.add_argument(
            "--VSA-sparsity",
            type=float,
            default=TrainerArgs.VSA_sparsity,
            help="Validation sparsity for VSA",
        )

        # Master port for distributed training/inference
        parser.add_argument(
            "--master-port",
            type=int,
            default=TrainerArgs.master_port,
            help="Master port for distributed training/inference",
        )

        # Stage verification
        parser.add_argument(
            "--enable-stage-verification",
            action=StoreBoolean,
            default=TrainerArgs.enable_stage_verification,
            help="Enable input/output verification for pipeline stages",
        )

        parser.add_argument(
            "--cls-name",
            type=str,
            help="model package name",
        )
        parser.add_argument(
            "--load-from-dir",
            type=str,
            help="ar model checkpoint directory",
        )
        parser.add_argument(
            "--ar-action-load-from-dir",
            type=str,
            help="ar action model checkpoint directory",
        )
        # Add pipeline configuration arguments
        PipelineConfig.add_cli_args(parser)

        # Add preprocessing configuration arguments
        PreprocessConfig.add_cli_args(parser)

        return parser

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace) -> "TrainerArgs":
        provided_args = clean_cli_args(args)
        # Get all fields from the dataclass
        attrs = [attr.name for attr in dataclasses.fields(cls)]

        # Create a dictionary of attribute values, with defaults for missing attributes
        kwargs: dict[str, Any] = {}
        for attr in attrs:
            if attr == 'pipeline_config':
                if "HunyuanTransformer" in provided_args["cls_name"]: continue
                pipeline_config = PipelineConfig.from_kwargs(provided_args)
                kwargs['pipeline_config'] = pipeline_config
            elif attr == 'preprocess_config':
                preprocess_config = PreprocessConfig.from_kwargs(provided_args)
                kwargs['preprocess_config'] = preprocess_config
            elif attr == 'mode':
                # Convert string to ExecutionMode enum
                mode_value = getattr(args, attr, TrainerArgs.mode.value)
                kwargs['mode'] = ExecutionMode.from_string(
                    mode_value) if isinstance(mode_value, str) else mode_value
            elif attr == 'workload_type':
                # Convert string to WorkloadType enum
                workload_type_value = getattr(args, 'workload_type',
                                              TrainerArgs.workload_type.value)
                kwargs['workload_type'] = WorkloadType.from_string(
                    workload_type_value) if isinstance(
                        workload_type_value, str) else workload_type_value
            # Use getattr with default value from the dataclass for potentially missing attributes
            else:
                # Get the field to check if it has a default_factory
                field = dataclasses.fields(cls)[next(
                    i for i, f in enumerate(dataclasses.fields(cls))
                    if f.name == attr)]
                if field.default_factory is not dataclasses.MISSING:
                    # Use the default_factory to create the default value
                    default_value = field.default_factory()
                else:
                    default_value = getattr(cls, attr, None)
                value = getattr(args, attr, default_value)
                kwargs[attr] = value  # type: ignore

        return cls(**kwargs)  # type: ignore

    @classmethod
    def from_kwargs(cls, **kwargs: Any) -> "TrainerArgs":
        # Convert mode string to enum if necessary
        if 'mode' in kwargs and isinstance(kwargs['mode'], str):
            kwargs['mode'] = ExecutionMode.from_string(kwargs['mode'])

        # Convert workload_type string to enum if necessary
        if 'workload_type' in kwargs and isinstance(kwargs['workload_type'],
                                                    str):
            kwargs['workload_type'] = WorkloadType.from_string(
                kwargs['workload_type'])

        kwargs['pipeline_config'] = PipelineConfig.from_kwargs(kwargs)
        kwargs['preprocess_config'] = PreprocessConfig.from_kwargs(kwargs)
        return cls(**kwargs)

    def check_trainer_args(self) -> None:
        """Validate inference arguments for consistency"""
        if current_platform.is_mps():
            self.use_fsdp_inference = False

        # Validate mode and inference_mode consistency
        assert isinstance(
            self.mode, ExecutionMode
        ), f"Mode must be an ExecutionMode enum, got {type(self.mode)}"
        assert self.mode in ExecutionMode.choices(
        ), f"Invalid execution mode: {self.mode}"

        # Validate workload type
        assert isinstance(
            self.workload_type, WorkloadType
        ), f"Workload type must be a WorkloadType enum, got {type(self.workload_type)}"
        assert self.workload_type in WorkloadType.choices(
        ), f"Invalid workload type: {self.workload_type}"

        if self.mode in [ExecutionMode.DISTILLATION, ExecutionMode.FINETUNING
                         ] and self.inference_mode:
            logger.warning(
                "Mode is 'training' but inference_mode is True. Setting inference_mode to False."
            )
            self.inference_mode = False
        elif self.mode in [ExecutionMode.INFERENCE, ExecutionMode.PREPROCESS
                           ] and not self.inference_mode:
            logger.warning(
                "Mode is '%s' but inference_mode is False. Setting inference_mode to True.",
                self.mode)
            self.inference_mode = True

        if not self.inference_mode:
            assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training"
            assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training"
            assert self.sp_size != -1, "sp_size must be set for training"

        if self.tp_size == -1:
            self.tp_size = 1
        if self.sp_size == -1:
            self.sp_size = self.num_gpus
        if self.hsdp_shard_dim == -1:
            self.hsdp_shard_dim = self.num_gpus

        assert self.sp_size <= self.num_gpus and self.num_gpus % self.sp_size == 0, "num_gpus must >= and be divisible by sp_size"
        assert self.hsdp_replicate_dim <= self.num_gpus and self.num_gpus % self.hsdp_replicate_dim == 0, "num_gpus must >= and be divisible by hsdp_replicate_dim"
        assert self.hsdp_shard_dim <= self.num_gpus and self.num_gpus % self.hsdp_shard_dim == 0, "num_gpus must >= and be divisible by hsdp_shard_dim"

        if self.num_gpus < max(self.tp_size, self.sp_size):
            self.num_gpus = max(self.tp_size, self.sp_size)

        if self.pipeline_config is None:
            raise ValueError("pipeline_config is not set in TrainerArgs")

        self.pipeline_config.check_pipeline_config()

        # Add preprocessing config validation if needed
        if self.mode == ExecutionMode.PREPROCESS:
            if self.preprocess_config is None:
                raise ValueError(
                    "preprocess_config is not set in TrainerArgs when mode is PREPROCESS"
                )
            if self.preprocess_config.model_path == "":
                self.preprocess_config.model_path = self.model_path
            if not self.pipeline_config.vae_config.load_encoder:
                self.pipeline_config.vae_config.load_encoder = True
            self.preprocess_config.check_preprocess_config()


_current_trainer_args = None


def prepare_trainer_args(argv: list[str]) -> TrainerArgs:
    """
    Prepare the inference arguments from the command line arguments.

    Args:
        argv: The command line arguments. Typically, it should be `sys.argv[1:]`
            to ensure compatibility with `parse_args` when no arguments are passed.

    Returns:
        The inference arguments.
    """
    parser = FlexibleArgumentParser()
    TrainerArgs.add_cli_args(parser)
    raw_args = parser.parse_args(argv)
    trainer_args = TrainerArgs.from_cli_args(raw_args)
    global _current_trainer_args
    _current_trainer_args = trainer_args
    return trainer_args


@contextmanager
def set_current_trainer_args(trainer_args: TrainerArgs):
    """
    Temporarily set the current trainer config.
    Used during model initialization.
    We save the current trainer config in a global variable,
    so that all modules can access it, e.g. custom ops
    can access the trainer config to determine how to dispatch.
    """
    global _current_trainer_args
    old_trainer_args = _current_trainer_args
    try:
        _current_trainer_args = trainer_args
        yield
    finally:
        _current_trainer_args = old_trainer_args


def get_current_trainer_args() -> TrainerArgs:
    if _current_trainer_args is None:
        # in ci, usually when we test custom ops/modules directly,
        # we don't set the trainer config. In that case, we set a default
        # config.
        # TODO(will): may need to handle this for CI.
        raise ValueError("Current trainer args is not set.")
    return _current_trainer_args


@dataclasses.dataclass
class TrainingArgs(TrainerArgs):
    """
    Training arguments. Inherits from TrainerArgs and adds training-specific
    arguments. If there are any conflicts, the training arguments will take
    precedence.
    """
    data_path: str = ""
    dataloader_num_workers: int = 0
    num_height: int = 0
    num_width: int = 0
    num_frames: int = 0

    train_batch_size: int = 0
    num_latent_t: int = 0
    group_frame: bool = False
    group_resolution: bool = False

    # text encoder & vae & diffusion model
    pretrained_model_name_or_path: str = ""
    dit_model_name_or_path: str = ""

    # diffusion setting
    ema_decay: float = 0.0
    ema_start_step: int = 0
    training_cfg_rate: float = 0.0
    precondition_outputs: bool = False

    # validation & logs
    validation_dataset_file: str = ""
    validation_preprocessed_path: str = ""
    validation_sampling_steps: str = ""
    validation_guidance_scale: str = ""
    validation_steps: float = 0.0
    log_validation: bool = False
    tracker_project_name: str = ""
    wandb_run_name: str = ""
    seed: int | None = None

    # output
    output_dir: str = ""
    checkpoints_total_limit: int = 0
    checkpointing_steps: int = 0
    resume_from_checkpoint: str = ""  # specify the checkpoint folder to resume from

    # optimizer & scheduler
    num_train_epochs: int = 0
    max_train_steps: int = 0
    gradient_accumulation_steps: int = 0
    learning_rate: float = 0.0
    scale_lr: bool = False
    lr_scheduler: str = "constant"
    lr_warmup_steps: int = 0
    max_grad_norm: float = 0.0
    enable_gradient_checkpointing_type: str | None = None
    selective_checkpointing: float = 0.0
    mixed_precision: str = ""
    train_sp_batch_size: int = 0
    fsdp_sharding_startegy: str = ""

    weighting_scheme: str = ""
    logit_mean: float = 0.0
    logit_std: float = 1.0
    mode_scale: float = 0.0

    num_euler_timesteps: int = 0
    lr_num_cycles: int = 0
    lr_power: float = 0.0
    min_lr_ratio: float = 0.5  # minimum learning rate ratio for cosine_with_min_lr scheduler
    not_apply_cfg_solver: bool = False
    distill_cfg: float = 0.0
    scheduler_type: str = ""
    linear_quadratic_threshold: float = 0.0
    linear_range: float = 0.0
    weight_decay: float = 0.0
    use_ema: bool = False
    multi_phased_distill_schedule: str = ""
    pred_decay_weight: float = 0.0
    pred_decay_type: str = ""
    hunyuan_teacher_disable_cfg: bool = False

    # master_weight_type
    master_weight_type: str = ""

    # VSA training decay parameters
    VSA_decay_rate: float = 0.01  # decay rate -> 0.02
    VSA_decay_interval_steps: int = 1  # decay interval steps -> 50

    # LoRA training parameters
    lora_rank: int | None = None
    lora_alpha: int | None = None
    lora_training: bool = False

    # camera dataset
    json_path: str = ""
    causal: bool = False
    window_frames: int = 9
    wandb_key: str = "f8f8f614d325bd04540f8517b142ca942b21017a"
    wandb_entity: str = "zhyzhy"
    action: bool = False
    use_memory: bool = False
    i2v_rate: float = 0.0
    train_time_shift: float = 1.0

    # distillation args
    generator_update_interval: int = 5
    min_timestep_ratio: float = 0.2
    max_timestep_ratio: float = 0.98
    real_score_guidance_scale: float = 3.5
    fake_score_learning_rate: float = 0.0  # separate learning rate for fake_score_transformer, if 0.0, use learning_rate
    fake_score_lr_scheduler: str = "constant"  # separate lr scheduler for fake_score_transformer, if not set, use lr_scheduler
    training_state_checkpointing_steps: int = 0  # for resuming training
    weight_only_checkpointing_steps: int = 0  # for inference
    log_visualization: bool = False
    # simulate generator forward to match inference
    simulate_generator_forward: bool = False

    @classmethod
    def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
        provided_args = clean_cli_args(args)
        # Get all fields from the dataclass
        attrs = [attr.name for attr in dataclasses.fields(cls)]
        logger.info(provided_args)
        # Create a dictionary of attribute values, with defaults for missing attributes
        kwargs: dict[str, Any] = {}
        for attr in attrs:
            if attr == 'pipeline_config':
                if "HunyuanTransformer" in provided_args["cls_name"]: continue
                pipeline_config = PipelineConfig.from_kwargs(provided_args)
                kwargs[attr] = pipeline_config
            elif attr == 'mode':
                # Convert string to ExecutionMode enum
                mode_value = getattr(args, attr, ExecutionMode.FINETUNING.value)
                kwargs[attr] = ExecutionMode.from_string(
                    mode_value) if isinstance(mode_value, str) else mode_value
            elif attr == 'workload_type':
                # Convert string to WorkloadType enum
                workload_type_value = getattr(args, 'workload_type',
                                              WorkloadType.T2V.value)
                kwargs[attr] = WorkloadType.from_string(
                    workload_type_value) if isinstance(
                        workload_type_value, str) else workload_type_value
            # Use getattr with default value from the dataclass for potentially missing attributes
            else:
                # Get the field to check its default value
                field = dataclasses.fields(cls)[next(
                    i for i, f in enumerate(dataclasses.fields(cls))
                    if f.name == attr)]

                # Check if the attribute is provided in args
                if hasattr(args, attr):
                    value = getattr(args, attr)
                else:
                    # Use the field's default value
                    if field.default_factory is not dataclasses.MISSING:
                        value = field.default_factory()
                    elif field.default is not dataclasses.MISSING:
                        value = field.default
                    else:
                        # No default value, use None
                        value = None

                kwargs[attr] = value

        return cls(**kwargs)  # type: ignore

    @staticmethod
    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
        parser.add_argument("--data-path",
                            type=str,
                            required=True,
                            help="Path to parquet files")
        parser.add_argument("--dataloader-num-workers",
                            type=int,
                            required=True,
                            help="Number of workers for dataloader")
        parser.add_argument("--num-height",
                            type=int,
                            required=True,
                            help="Number of heights")
        parser.add_argument("--num-width",
                            type=int,
                            required=True,
                            help="Number of widths")
        parser.add_argument("--num-frames",
                            type=int,
                            required=True,
                            help="Number of frames")

        # Training batch and model configuration
        parser.add_argument("--train-batch-size",
                            type=int,
                            required=True,
                            help="Training batch size")
        parser.add_argument("--num-latent-t",
                            type=int,
                            required=True,
                            help="Number of latent time steps")
        parser.add_argument("--group-frame",
                            action=StoreBoolean,
                            help="Whether to group frames during training")
        parser.add_argument("--group-resolution",
                            action=StoreBoolean,
                            help="Whether to group resolutions during training")

        # Model paths
        parser.add_argument("--pretrained-model-name-or-path",
                            type=str,
                            required=True,
                            help="Path to pretrained model or model name")
        parser.add_argument("--dit-model-name-or-path",
                            type=str,
                            required=False,
                            help="Path to DiT model or model name")
        parser.add_argument("--cache-dir",
                            type=str,
                            help="Directory to cache models")

        # Diffusion settings
        parser.add_argument("--ema-decay",
                            type=float,
                            default=0.999,
                            help="EMA decay rate")
        parser.add_argument("--ema-start-step",
                            type=int,
                            default=0,
                            help="Step to start EMA")
        parser.add_argument("--training-cfg-rate",
                            type=float,
                            help="Classifier-free guidance scale")
        parser.add_argument(
            "--precondition-outputs",
            action=StoreBoolean,
            help="Whether to precondition the outputs of the model")

        # Validation and logging
        parser.add_argument("--validation-dataset-file",
                            type=str,
                            help="Path to unprocessed validation dataset")
        parser.add_argument("--validation-preprocessed-path",
                            type=str,
                            help="Path to processed validation dataset")
        parser.add_argument("--validation-sampling-steps",
                            type=str,
                            help="Validation sampling steps")
        parser.add_argument("--validation-guidance-scale",
                            type=str,
                            help="Validation guidance scale")
        parser.add_argument("--validation-steps",
                            type=float,
                            help="Number of validation steps")
        parser.add_argument("--log-validation",
                            action=StoreBoolean,
                            help="Whether to log validation results")
        parser.add_argument("--tracker-project-name",
                            type=str,
                            help="Project name for tracking")
        parser.add_argument("--wandb-run-name",
                            type=str,
                            help="Run name for wandb")
        parser.add_argument("--seed",
                            type=int,
                            default=42,
                            help="Seed for deterministic training")

        # Output configuration
        parser.add_argument("--output-dir",
                            type=str,
                            required=True,
                            help="Output directory for checkpoints and logs")
        parser.add_argument("--checkpoints-total-limit",
                            type=int,
                            help="Maximum number of checkpoints to keep")
        parser.add_argument("--checkpointing-steps",
                            type=int,
                            help="Steps between checkpoints")
        parser.add_argument(
            "--training-state-checkpointing-steps",
            type=int,
            help=
            "Steps between training state checkpoints (for resuming training)")
        parser.add_argument(
            "--weight-only-checkpointing-steps",
            type=int,
            help="Steps between weight-only checkpoints (for inference)")
        parser.add_argument("--resume-from-checkpoint",
                            type=str,
                            help="Path to checkpoint to resume from")
        parser.add_argument("--logging-dir",
                            type=str,
                            help="Directory for logging")

        # Training configuration
        parser.add_argument("--num-train-epochs",
                            type=int,
                            help="Number of training epochs")
        parser.add_argument("--max-train-steps",
                            type=int,
                            help="Maximum number of training steps")
        parser.add_argument("--gradient-accumulation-steps",
                            type=int,
                            help="Number of steps to accumulate gradients")
        parser.add_argument("--learning-rate",
                            type=float,
                            required=True,
                            help="Learning rate")
        parser.add_argument("--scale-lr",
                            action=StoreBoolean,
                            help="Whether to scale learning rate")
        parser.add_argument("--lr-scheduler",
                            type=str,
                            default="constant",
                            help="Learning rate scheduler type")
        parser.add_argument("--lr-warmup-steps",
                            type=int,
                            default=10,
                            help="Number of warmup steps for learning rate")
        parser.add_argument("--max-grad-norm",
                            type=float,
                            help="Maximum gradient norm")
        parser.add_argument("--enable-gradient-checkpointing-type",
                            type=str,
                            choices=["full", "ops", "block_skip"],
                            default=None,
                            help="Gradient checkpointing type")
        parser.add_argument("--selective-checkpointing",
                            type=float,
                            help="Selective checkpointing threshold")
        parser.add_argument("--mixed-precision",
                            type=str,
                            help="Mixed precision training type")
        parser.add_argument("--train-sp-batch-size",
                            type=int,
                            help="Training spatial parallelism batch size")

        parser.add_argument("--fsdp-sharding-strategy",
                            type=str,
                            help="FSDP sharding strategy")

        parser.add_argument(
            "--weighting-scheme",
            type=str,
            default="uniform",
            choices=["sigma_sqrt", "logit_normal", "mode", "cosmap", "uniform"],
        )
        parser.add_argument(
            "--logit-mean",
            type=float,
            default=0.0,
            help="mean to use when using the `'logit_normal'` weighting scheme.",
        )
        parser.add_argument(
            "--logit-std",
            type=float,
            default=1.0,
            help="std to use when using the `'logit_normal'` weighting scheme.",
        )
        parser.add_argument(
            "--mode_scale",
            type=float,
            default=1.29,
            help=
            "Scale of mode weighting scheme. Only effective when using the `'mode'` as the `weighting_scheme`.",
        )

        # Additional training parameters
        parser.add_argument("--num-euler-timesteps",
                            type=int,
                            help="Number of Euler timesteps")
        parser.add_argument("--lr-num-cycles",
                            type=int,
                            help="Number of learning rate cycles")
        parser.add_argument("--lr-power",
                            type=float,
                            help="Learning rate power")
        parser.add_argument(
            "--min-lr-ratio",
            type=float,
            default=TrainingArgs.min_lr_ratio,
            help="Minimum learning rate ratio for cosine_with_min_lr scheduler")
        parser.add_argument("--not-apply-cfg-solver",
                            action=StoreBoolean,
                            help="Whether to not apply CFG solver")
        parser.add_argument("--distill-cfg",
                            type=float,
                            help="Distillation CFG scale")
        parser.add_argument("--scheduler-type", type=str, help="Scheduler type")
        parser.add_argument("--linear-quadratic-threshold",
                            type=float,
                            help="Linear quadratic threshold")
        parser.add_argument("--linear-range", type=float, help="Linear range")
        parser.add_argument("--weight-decay", type=float, help="Weight decay")
        parser.add_argument("--use-ema",
                            action=StoreBoolean,
                            help="Whether to use EMA")
        parser.add_argument("--multi-phased-distill-schedule",
                            type=str,
                            help="Multi-phased distillation schedule")
        parser.add_argument("--pred-decay-weight",
                            type=float,
                            help="Prediction decay weight")
        parser.add_argument("--pred-decay-type",
                            type=str,
                            help="Prediction decay type")
        parser.add_argument("--hunyuan-teacher-disable-cfg",
                            action=StoreBoolean,
                            help="Whether to disable CFG for Hunyuan teacher")
        parser.add_argument("--master-weight-type",
                            type=str,
                            help="Master weight type")

        # VSA parameters for training with dense to sparse adaption
        parser.add_argument(
            "--VSA-decay-rate",  # decay rate, how much sparsity you want to decay each step
            type=float,
            default=TrainingArgs.VSA_decay_rate,
            help="VSA decay rate")
        parser.add_argument(
            "--VSA-decay-interval-steps",  # how many steps for training with current sparsity
            type=int,
            default=TrainingArgs.VSA_decay_interval_steps,
            help="VSA decay interval steps")
        parser.add_argument("--lora-training",
                            action=StoreBoolean,
                            help="Whether to use LoRA training")
        parser.add_argument("--lora-rank", type=int, help="LoRA rank")
        parser.add_argument("--lora-alpha", type=int, help="LoRA alpha")

        # Distillation arguments
        parser.add_argument("--generator-update-interval",
                            type=int,
                            default=TrainingArgs.generator_update_interval,
                            help="Ratio of student updates to critic updates.")
        parser.add_argument("--min-timestep-ratio",
                            type=float,
                            default=TrainingArgs.min_timestep_ratio,
                            help="Minimum step ratio")
        parser.add_argument("--max-timestep-ratio",
                            type=float,
                            default=TrainingArgs.max_timestep_ratio,
                            help="Maximum step ratio")
        parser.add_argument("--real-score-guidance-scale",
                            type=float,
                            default=TrainingArgs.real_score_guidance_scale,
                            help="Teacher guidance scale")
        parser.add_argument("--fake-score-learning-rate",
                            type=float,
                            default=TrainingArgs.fake_score_learning_rate,
                            help="Learning rate for fake score transformer")
        parser.add_argument(
            "--fake-score-lr-scheduler",
            type=str,
            default=TrainingArgs.fake_score_lr_scheduler,
            help="Learning rate scheduler for fake score transformer")
        parser.add_argument("--log-visualization",
                            action=StoreBoolean,
                            help="Whether to log visualization")
        parser.add_argument(
            "--simulate-generator-forward",
            action=StoreBoolean,
            help="Whether to simulate generator forward to match inference")

        # camera
        parser.add_argument("--json-path",
                            type=str,
                            default="/all_data.json",
                            help="camera json path")
        parser.add_argument("--causal",
                            action=StoreBoolean,
                            help="training ar model")
        parser.add_argument("--window-frames",
                            type=int,
                            default=9,
                            help="window size")
        parser.add_argument("--wandb-key",
                            type=str,
                            default="",
                            help="key")
        parser.add_argument("--wandb-entity",
                            type=str,
                            default="",
                            help="entity")
        parser.add_argument("--action",
                            action=StoreBoolean,
                            help="training ar model")

        parser.add_argument("--use-memory",
                            action=StoreBoolean,
                            help="training mem model")
        parser.add_argument("--i2v-rate",
                            type=float,
                            default=0.0,
                            help="training rating for i2v")
        parser.add_argument("--train-time-shift",
                            type=float,
                            default=1.0,
                            help="training time shift")
        return parser


def parse_int_list(value: str) -> list[int]:
    if not value:
        return []
    return [int(x.strip()) for x in value.split(",")]