File size: 4,317 Bytes
954544e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | from __future__ import annotations
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any
import yaml
@dataclass(frozen=True)
class ModelConfig:
base_model: str
insertion_after: int = 15
source_layers: tuple[int, ...] = (12, 13, 14, 15)
max_loops: int = 8
active_loops: int = 1
initial_loop_scale: float = 0.01
attention_implementation: str = "sdpa"
def validate(self) -> None:
if not self.base_model:
raise ValueError("model.base_model is required")
if self.insertion_after < 0:
raise ValueError("model.insertion_after must be non-negative")
if not self.source_layers:
raise ValueError("model.source_layers must not be empty")
if len(set(self.source_layers)) != len(self.source_layers):
raise ValueError("model.source_layers must be unique")
if min(self.source_layers) < 0:
raise ValueError("model.source_layers must be non-negative")
if self.max_loops < 1:
raise ValueError("model.max_loops must be at least 1")
if not 1 <= self.active_loops <= self.max_loops:
raise ValueError("model.active_loops must be within [1, max_loops]")
if not 0.0 <= self.initial_loop_scale < 1.0:
raise ValueError("model.initial_loop_scale must be within [0, 1)")
@dataclass(frozen=True)
class TrainingConfig:
data_path: str
output_dir: str
sequence_length: int = 512
micro_batch_size: int = 1
gradient_accumulation_steps: int = 4
max_steps: int = 20
learning_rate_core: float = 1.0e-5
learning_rate_scales: float = 1.0e-3
weight_decay: float = 0.1
warmup_steps: int = 2
max_grad_norm: float = 1.0
log_every: int = 1
save_every: int = 10
seed: int = 1337
gradient_checkpointing: bool = True
keep_last_checkpoints: int = 2
reasoning_loss_weight: float = 0.35
answer_loss_weight: float = 1.0
def validate(self) -> None:
if not self.data_path:
raise ValueError("training.data_path is required")
if not self.output_dir:
raise ValueError("training.output_dir is required")
positive_ints = {
"sequence_length": self.sequence_length,
"micro_batch_size": self.micro_batch_size,
"gradient_accumulation_steps": self.gradient_accumulation_steps,
"max_steps": self.max_steps,
"log_every": self.log_every,
"save_every": self.save_every,
"keep_last_checkpoints": self.keep_last_checkpoints,
}
for name, value in positive_ints.items():
if value < 1:
raise ValueError(f"training.{name} must be at least 1")
if self.warmup_steps < 0:
raise ValueError("training.warmup_steps must be non-negative")
if self.learning_rate_core <= 0 or self.learning_rate_scales <= 0:
raise ValueError("training learning rates must be positive")
if self.max_grad_norm <= 0:
raise ValueError("training.max_grad_norm must be positive")
if self.reasoning_loss_weight <= 0 or self.answer_loss_weight <= 0:
raise ValueError("training thinking loss weights must be positive")
@dataclass(frozen=True)
class DotRDConfig:
model: ModelConfig
training: TrainingConfig
metadata: dict[str, Any] = field(default_factory=dict)
def validate(self) -> None:
self.model.validate()
self.training.validate()
def to_dict(self) -> dict[str, Any]:
return asdict(self)
def load_config(path: str | Path) -> DotRDConfig:
config_path = Path(path)
payload = yaml.safe_load(config_path.read_text(encoding="utf-8"))
if not isinstance(payload, dict):
raise ValueError("configuration root must be an object")
model_payload = dict(payload.get("model") or {})
training_payload = dict(payload.get("training") or {})
if "source_layers" in model_payload:
model_payload["source_layers"] = tuple(model_payload["source_layers"])
config = DotRDConfig(
model=ModelConfig(**model_payload),
training=TrainingConfig(**training_payload),
metadata=dict(payload.get("metadata") or {}),
)
config.validate()
return config
|