| """Strict execution settings kept separate from Yunmo-Next model geometry. |
| |
| The type records how one training process executes an already-defined model. |
| It intentionally has no context-to-stage inference and no production route |
| table; stage recipes must select an engine explicitly. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import asdict, dataclass, fields |
| from typing import Any, Mapping |
|
|
|
|
| @dataclass(frozen=True) |
| class TrainingEngineConfig: |
| """Stage-independent training execution schema.""" |
|
|
| feed_forward_chunk_tokens: int = 0 |
| feed_forward_serial_backward: bool = False |
| recurrent_segment_tokens: int = 0 |
| recurrent_packed_documents: bool = False |
| recurrent_official_forward_serial_backward: bool = False |
| activation_checkpointing: bool = True |
| activation_checkpoint_mode: str = "non_reentrant" |
| cross_entropy_backend: str = "liger" |
|
|
| def __post_init__(self) -> None: |
| for name in ("feed_forward_chunk_tokens", "recurrent_segment_tokens"): |
| value = getattr(self, name) |
| if type(value) is not int: |
| raise TypeError(f"{name} must be int, got {type(value).__name__}") |
| if value < 0: |
| raise ValueError(f"{name} must be non-negative") |
| for name in ( |
| "feed_forward_serial_backward", |
| "recurrent_packed_documents", |
| "recurrent_official_forward_serial_backward", |
| "activation_checkpointing", |
| ): |
| value = getattr(self, name) |
| if type(value) is not bool: |
| raise TypeError(f"{name} must be bool, got {type(value).__name__}") |
| for name in ("activation_checkpoint_mode", "cross_entropy_backend"): |
| value = getattr(self, name) |
| if type(value) is not str: |
| raise TypeError(f"{name} must be str, got {type(value).__name__}") |
|
|
| if ( |
| self.feed_forward_serial_backward |
| and self.feed_forward_chunk_tokens == 0 |
| ): |
| raise ValueError( |
| "serial feed-forward backward requires a positive chunk size" |
| ) |
| if ( |
| self.recurrent_segment_tokens |
| and self.recurrent_segment_tokens % 64 |
| ): |
| raise ValueError("recurrent segment size must be a multiple of 64") |
| if self.recurrent_packed_documents and self.recurrent_segment_tokens == 0: |
| raise ValueError( |
| "packed recurrent execution requires a positive segment size" |
| ) |
| if ( |
| self.recurrent_official_forward_serial_backward |
| and not self.recurrent_packed_documents |
| ): |
| raise ValueError( |
| "official-forward serial-backward recurrent execution " |
| "requires packed documents" |
| ) |
| if self.activation_checkpoint_mode != "non_reentrant": |
| raise ValueError( |
| "Yunmo-Next training uses non-reentrant activation checkpointing" |
| ) |
| if self.cross_entropy_backend not in {"liger", "pytorch"}: |
| raise ValueError( |
| "cross_entropy_backend must be liger or pytorch" |
| ) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return asdict(self) |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "TrainingEngineConfig": |
| if not isinstance(value, Mapping): |
| raise TypeError("TrainingEngineConfig must be loaded from a mapping") |
| expected = {field.name for field in fields(cls)} |
| actual = set(value) |
| if actual != expected: |
| raise ValueError( |
| "TrainingEngineConfig field set changed: " |
| f"missing={sorted(expected - actual)}, " |
| f"extra={sorted(actual - expected)}" |
| ) |
| return cls(**dict(value)) |
|
|
|
|
| TRAINING_ENGINE_CONFIG_FIELDS = tuple( |
| field.name for field in fields(TrainingEngineConfig) |
| ) |
|
|
