Update Wall-X to 1.1.0 (#104)

wall_x/config/hyperparams_config.py

@@ -0,0 +1,129 @@
+"""Training hyperparameter config dataclasses."""
+
+from dataclasses import dataclass, field
+from typing import List, Optional, Tuple
+
+from .registry import register_optimizer_config, register_scheduler_config
+
+
+@dataclass
+class LRGroupConfig:
+    """Named optimizer LR group matched by parameter-name substrings."""
+
+    name: str
+    lr: float
+    include: List[str] = field(default_factory=list)
+    fail_on_empty: bool = True
+
+
+@dataclass
+class OptimizerConfig:
+    """Base optimizer config. ``optimizer_type`` selects a registered subclass."""
+
+    optimizer_type: str = "adamw"
+    learning_rate: float = 1e-4
+    max_grad_norm: float = 1.0
+    enable_grad_clip: bool = True
+    # Named parameter groups with independent learning rates. Unmatched
+    # trainable parameters remain in the base group using ``learning_rate``.
+    lr_groups: Optional[List[LRGroupConfig]] = None
+    # Optional action-expert LR split.
+    train_action_expert_only: bool = False
+    action_expert_learning_rate: Optional[float] = None
+    action_lr_keywords: Optional[List[str]] = None
+
+
+@register_optimizer_config("adamw")
+@dataclass
+class AdamWConfig(OptimizerConfig):
+    """AdamW optimizer config."""
+
+    optimizer_type: str = "adamw"
+    betas: Tuple[float, float] = (0.9, 0.98)
+    weight_decay: float = 1e-8
+    eps: float = 1e-8
+    fused: bool = True
+    foreach: Optional[bool] = None
+
+
+@register_optimizer_config("dmuon")
+@dataclass
+class DMuonConfig(OptimizerConfig):
+    """DMuon optimizer config."""
+
+    optimizer_type: str = "dmuon"
+    muon_lr: float = 0.02
+    momentum: float = 0.95
+    ns_steps: int = 5
+    muon_weight_decay: float = 0.0
+    adamw_lr: float = 1e-3
+    adamw_betas: Tuple[float, float] = (0.9, 0.999)
+    adamw_weight_decay: float = 0.01
+    adamw_eps: float = 1e-8
+    ns_backend: str = "gram"
+    ns_coefficients: str = "default"
+    nesterov: bool = True
+
+
+@dataclass
+class SchedulerConfig:
+    """Base scheduler config. ``scheduler_type`` selects a registered subclass."""
+
+    scheduler_type: str = "cosine"
+    # Optional training-step cap. When > 0, the trainer sets
+    # loss_guard_should_stop=True once global_step >= num_training_steps,
+    # regardless of scheduler type. Cosine reads it for its own decay
+    # horizon; constant / step schedulers use it only for the stop signal.
+    num_training_steps: int = 0
+
+
+@register_scheduler_config("cosine")
+@dataclass
+class CosineSchedulerConfig(SchedulerConfig):
+    """Cosine annealing with warmup."""
+
+    scheduler_type: str = "cosine"
+    num_warmup_steps: int = 0
+    num_training_steps: int = 0
+    min_lr: Optional[float] = None  # None means 0.1 * learning_rate at runtime.
+
+
+@register_scheduler_config("constant")
+@dataclass
+class ConstantSchedulerConfig(SchedulerConfig):
+    """Constant learning rate with no decay."""
+
+    scheduler_type: str = "constant"
+
+
+@register_scheduler_config("step")
+@dataclass
+class StepSchedulerConfig(SchedulerConfig):
+    """Step decay scheduler."""
+
+    scheduler_type: str = "step"
+    step_size: int = 10000
+    gamma: float = 0.1
+
+
+@dataclass
+class TrainHyperParams:
+    num_epoch: int = 1
+    batch_size_per_gpu: int = 1
+    gradient_accumulation_steps: int = 1
+    seed: int = 42
+    optimizer: OptimizerConfig = field(default_factory=AdamWConfig)
+    scheduler: SchedulerConfig = field(default_factory=CosineSchedulerConfig)
+
+
+__all__ = [
+    "AdamWConfig",
+    "ConstantSchedulerConfig",
+    "CosineSchedulerConfig",
+    "DMuonConfig",
+    "LRGroupConfig",
+    "OptimizerConfig",
+    "SchedulerConfig",
+    "StepSchedulerConfig",
+    "TrainHyperParams",
+]