ezpz.configs⚓︎

@dataclass
class HfDataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    """

    data_path: Optional[str] = field(
        default=None,
        metadata={"help": "Path to the training data."},
    )
    dataset_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "The name of the dataset to use (via the datasets library)."
        },
    )
    dataset_config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "The configuration name of the dataset to use (via the datasets library)."
        },
    )
    train_split_str: Optional[str] = field(
        default=None,
        metadata={
            "help": "The split string to use for the train split (via the datasets library)."
        },
    )
    train_split_name: Optional[str] = field(
        default="train",
        metadata={
            "help": "The name of the train split to use (via the datasets library)."
        },
    )
    validation_split_name: Optional[str] = field(
        default="validation",
        metadata={
            "help": "The name of the validation split to use (via the datasets library)."
        },
    )
    validation_split_str: Optional[str] = field(
        default=None,
        metadata={
            "help": "The split string to use for the validation split (via the datasets library)."
        },
    )
    test_split_name: Optional[str] = field(
        default="test",
        metadata={
            "help": "The name of the test split to use (via the datasets library)."
        },
    )
    test_split_str: Optional[str] = field(
        default=None,
        metadata={
            "help": "The split string to use for the test split (via the datasets library)."
        },
    )
    train_file: Optional[str] = field(
        default=None,
        metadata={"help": "The input training data file (a text file)."},
    )
    validation_file: Optional[str] = field(
        default=None,
        metadata={
            "help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."
        },
    )
    max_train_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of training examples to this "
                "value if set."
            )
        },
    )
    max_eval_samples: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
                "value if set."
            )
        },
    )
    streaming: bool = field(
        default=False, metadata={"help": "Enable streaming mode"}
    )
    block_size: Optional[int] = field(
        default=None,
        metadata={
            "help": (
                "Optional input sequence length after tokenization. "
                "The training dataset will be truncated in block of this size for training. "
                "Default to the model max input length for single sentence inputs (take into account special tokens)."
            )
        },
    )
    overwrite_cache: bool = field(
        default=False,
        metadata={"help": "Overwrite the cached training and evaluation sets"},
    )
    validation_split_percentage: Optional[int] = field(
        default=5,
        metadata={
            "help": "The percentage of the train set used as validation set in case there's no validation split"
        },
    )
    preprocessing_num_workers: Optional[int] = field(
        default=None,
        metadata={
            "help": "The number of processes to use for the preprocessing."
        },
    )
    keep_linebreaks: bool = field(
        default=True,
        metadata={
            "help": "Whether to keep line breaks when using TXT files or not."
        },
    )

    def __post_init__(self):
        """Validate dataset arguments and ensure required files are present."""
        from transformers.utils.versions import require_version

        if self.streaming:
            require_version(
                "datasets>=2.0.0",
                "The streaming feature requires `datasets>=2.0.0`",
            )

        if (
            self.dataset_name is None
            and self.data_path is None
            and self.train_file is None
            and self.validation_file is None
        ):
            raise ValueError(
                "You must specify at least one of the following: "
                "a dataset name, a data path, a training file, or a validation file."
            )

        if self.train_file is not None:
            extension = self.train_file.split(".")[-1]
            assert extension in [
                "csv",
                "json",
                "txt",
            ], "`train_file` should be a csv, a json or a txt file."
        if self.validation_file is not None:
            extension = self.validation_file.split(".")[-1]
            assert extension in [
                "csv",
                "json",
                "txt",
            ], "`validation_file` should be a csv, a json or a txt file."

@dataclass
class HfModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
    """

    wandb_project_name: Optional[str] = field(  # type:ignore
        default=None,
        metadata={
            "help": (
                "The name of the wandb project to use. If not specified, will use the model name."
            )
        },
    )

    model_name_or_path: Optional[str] = field(  # type:ignore
        default=None,
        metadata={
            "help": (
                "The model checkpoint for weights initialization. Don't set if you want to train a model from scratch."
            )
        },
    )
    model_type: Optional[str | None] = field(
        default=None,
        metadata={
            "help": "If training from scratch, pass a model type from the list: "
            + ", ".join("https://huggingface.co/docs/transformers/en/models")
        },
    )
    config_overrides: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "Override some existing default config settings when a model is trained from scratch. Example: "
                "n_embd=10,resid_pdrop=0.2,scale_attn_weights=false,summary_type=cls_index"
            )
        },
    )
    config_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Pretrained config name or path if not the same as model_name"
        },
    )
    tokenizer_name: Optional[str] = field(
        default=None,
        metadata={
            "help": "Pretrained tokenizer name or path if not the same as model_name"
        },
    )
    cache_dir: Optional[str] = field(
        default=None,
        metadata={
            "help": "Where do you want to store the pretrained models downloaded from huggingface.co"
        },
    )
    use_fast_tokenizer: bool = field(
        default=True,
        metadata={
            "help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."
        },
    )
    model_revision: str = field(
        default="main",
        metadata={
            "help": "The specific model version to use (can be a branch name, tag name or commit id)."
        },
    )
    token: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "The token to use as HTTP bearer authorization for remote files. If not specified, will use the token "
                "generated when running `huggingface-cli login` (stored in `~/.huggingface`)."
            )
        },
    )
    trust_remote_code: bool = field(
        default=False,
        metadata={
            "help": (
                "Whether to trust the execution of code from datasets/models defined on the Hub."
                " This option should only be set to `True` for repositories you trust and in which you have read the"
                " code, as it will execute code present on the Hub on your local machine."
            )
        },
    )
    torch_dtype: Optional[str] = field(
        default=None,
        metadata={
            "help": (
                "Override the default `torch.dtype` and load the model under this dtype. If `auto` is passed, the "
                "dtype will be automatically derived from the model's weights."
            ),
            "choices": ["auto", "bfloat16", "float16", "float32"],
        },
    )
    low_cpu_mem_usage: bool = field(
        default=False,
        metadata={
            "help": (
                "It is an option to create the model as an empty shell, then only materialize its parameters when the pretrained weights are loaded. "
                "set True will benefit LLM loading time and RAM consumption."
            )
        },
    )

    def __post_init__(self):
        """Validate mutually exclusive Hugging Face model configuration options."""
        if self.config_overrides is not None and (
            self.config_name is not None or self.model_name_or_path is not None
        ):
            raise ValueError(
                "--config_overrides can't be used in combination with --config_name or --model_name_or_path"
            )

@dataclass
class TrainArgs:
    """Arguments for vision model training runs."""

    img_size: int
    """Input image resolution."""

    batch_size: int
    """Training batch size per device."""

    num_heads: int
    """Number of self-attention heads."""

    head_dim: int
    """Dimension per attention head."""

    depth: int
    """Number of transformer encoder layers."""

    patch_size: int
    """Size of each image patch."""

    dtype: str
    """Data type string (e.g. ``"bf16"``, ``"fp32"``)."""

    compile: bool
    """Enable ``torch.compile`` for the model."""

    attn_type: str
    """Attention backend (``"native"``, ``"sdpa"``, ``"flash"``, etc.)."""

    warmup: int | float
    """Learning rate warmup steps or fraction of total steps."""

    num_workers: int
    """Number of data loader workers."""

    max_iters: int
    """Maximum training iterations."""

    fsdp: Optional[bool] = False
    """Enable Fully Sharded Data Parallel wrapping."""

    format: Optional[str] = field(default_factory=str)
    """Output format string for logging."""

    cuda_sdpa_backend: Optional[str] = field(default_factory=str)
    """Override the CUDA SDPA backend selection."""

@dataclass(init=False)
@rich.repr.auto
class TrainConfig(BaseConfig):
    """High-level training options shared by ezpz scripts."""

    gas: int = 1
    """Gradient accumulation steps."""

    use_wandb: bool = False
    """Enable Weights & Biases logging."""

    seed: Optional[int] = None
    """Random seed for reproducibility. ``None`` means no explicit seeding."""

    port: Optional[str] = None
    """Port for the distributed rendezvous endpoint."""

    dtype: Optional[Any] = None
    """Default data type for model parameters (e.g. ``"bf16"``)."""

    load_from: Optional[str] = None
    """Path to a checkpoint to resume training from."""

    save_to: Optional[str] = None
    """Directory where checkpoints will be saved."""

    ds_config_path: Optional[str] = None
    """Path to a DeepSpeed JSON/YAML config file."""

    wandb_project_name: Optional[str] = None
    """W&B project name. Auto-set from ``WANDB_PROJECT`` when ``use_wandb`` is True."""

    ngpus: Optional[int] = None
    """Number of GPUs to use. ``None`` means use all available."""

    extras: dict[str, Any] = field(
        default_factory=dict, init=False, repr=False
    )
    """Extra keyword arguments not matching any declared field."""

    def __init__(self, **kwargs: Any) -> None:
        """Populate known fields while capturing any extras in ``self.extras``."""
        extras: dict[str, Any] = {}
        for name, field_def in self.__dataclass_fields__.items():
            if name == "extras":
                continue
            if name in kwargs:
                value = kwargs.pop(name)
            elif field_def.default is not MISSING:
                value = field_def.default
            elif field_def.default_factory is not MISSING:  # type: ignore[attr-defined]
                value = field_def.default_factory()  # type: ignore[misc]
            else:
                value = None
            setattr(self, name, value)

        for key, value in kwargs.items():
            setattr(self, key, value)
            extras[key] = value
        self.extras = extras
        self.__post_init__()

    def to_str(self) -> str:
        """Return a compact identifier for this config."""
        parts = [f"gas-{self.gas}"]
        if self.seed is not None:
            parts.append(f"seed-{self.seed}")
        return "_".join(parts)

    def __post_init__(self):
        """Validate configuration after initialisation."""
        if self.use_wandb and self.wandb_project_name is None:
            self.wandb_project_name = os.environ.get(
                "WANDB_PROJECT", os.environ.get("WB_PROJECT", "ezpz")
            )

@dataclass
class ViTConfig:
    """Configuration for a Vision Transformer model."""

    img_size: int = 224
    """Input image resolution (height and width)."""

    patch_size: int = 16
    """Size of each image patch."""

    depth: int = 12
    """Number of transformer encoder layers."""

    num_heads: int = 12
    """Number of self-attention heads."""

    hidden_dim: int = 768
    """Embedding / hidden dimension."""

    mlp_dim: int = 3072
    """Feed-forward network inner dimension."""

    dropout: float = 0.0
    """Dropout rate for embeddings and feed-forward layers."""

    attention_dropout: float = 0.0
    """Dropout rate applied to attention weights."""

    num_classes: int = 10
    """Number of output classification classes."""

    def __post_init__(self):
        self.seq_len = (self.img_size // self.patch_size) ** 2  # 196, default

@dataclass
class ZeroConfig:
    """Subset of DeepSpeed ZeRO options exposed via the ezpz CLI."""

    stage: int = 0
    """ZeRO optimisation stage (0–3)."""

    allgather_partitions: Optional[bool] = None
    """Use allgather for partition reconstruction."""

    allgather_bucket_size: int = int(5e8)
    """Number of elements allgathered at a time."""

    overlap_comm: Optional[bool] = None
    """Overlap gradient reduction with backward pass."""

    reduce_scatter: bool = True
    """Use reduce-scatter instead of allreduce for gradients."""

    reduce_bucket_size: int = int(5e8)
    """Number of elements reduced at a time."""

    contiguous_gradients: Optional[bool] = None
    """Copy gradients to a contiguous buffer as they are produced."""

    offload_param: Optional[dict] = None
    """Parameter offloading config (e.g. ``{"device": "cpu"}``)."""

    offload_optimizer: Optional[dict] = None
    """Optimizer state offloading config (e.g. ``{"device": "cpu"}``)."""

    stage3_max_live_parameters: int = int(1e9)
    """Max number of parameters resident per GPU before releasing (stage 3)."""

    stage3_max_reuse_distance: int = int(1e9)
    """Max reuse distance for parameters before re-fetching (stage 3)."""

    stage3_prefetch_bucket_size: int = int(5e8)
    """Prefetch bucket size for stage 3 parameter fetching."""

    stage3_param_persistence_threshold: int = int(1e6)
    """Parameters smaller than this stay on the GPU (stage 3)."""

    sub_group_size: Optional[int] = None
    """Number of parameters within a sub-group for ZeRO-3 partitioning."""

    elastic_checkpoint: Optional[dict] = None
    """Configuration for elastic checkpointing."""

    stage3_gather_16bit_weights_on_model_save: Optional[bool] = None
    """Gather full fp16 weights on save for portability (stage 3)."""

    ignore_unused_parameters: Optional[bool] = None
    """Silence errors about parameters not used in the forward pass."""

    round_robin_gradients: Optional[bool] = None
    """Distribute gradient partitions round-robin across ranks."""

    zero_hpz_partition_size: Optional[int] = None
    """Hierarchical partitioning group size (ZeRO++)."""

    zero_quantized_weights: Optional[bool] = None
    """Enable weight quantisation for communication (ZeRO++)."""

    zero_quantized_gradients: Optional[bool] = None
    """Enable gradient quantisation for communication (ZeRO++)."""

    log_trace_cache_warnings: Optional[bool] = None
    """Log warnings about trace cache misses."""

@dataclass
class timmViTConfig:
    """Configuration for a timm-style Vision Transformer (larger defaults)."""

    img_size: int = 224
    """Input image resolution."""

    batch_size: int = 128
    """Training batch size per device."""

    num_heads: int = 16
    """Number of self-attention heads."""

    head_dim: int = 64
    """Dimension per attention head."""

    depth: int = 24
    """Number of transformer encoder layers."""

    patch_size: int = 16
    """Size of each image patch."""

    hidden_dim: int = 1024
    """Embedding / hidden dimension."""

    mlp_dim: int = 4096
    """Feed-forward network inner dimension."""

    dropout: float = 0.0
    """Dropout rate for embeddings and feed-forward layers."""

    attention_dropout: float = 0.0
    """Dropout rate applied to attention weights."""

    num_classes: int = 1000
    """Number of output classification classes."""

    def __post_init__(self):
        self.seq_len = (self.img_size // self.patch_size) ** 2  # 196, default