Update Wall-X to 1.1.0 (#104)
This commit is contained in:
Starrick Liu
@@ -0,0 +1,62 @@
|
||||
"""Typed training configuration system."""
|
||||
|
||||
from .data_config import DataConfig, LeRobotDataConfig
|
||||
from .hyperparams_config import (
|
||||
AdamWConfig,
|
||||
ConstantSchedulerConfig,
|
||||
CosineSchedulerConfig,
|
||||
DMuonConfig,
|
||||
OptimizerConfig,
|
||||
SchedulerConfig,
|
||||
StepSchedulerConfig,
|
||||
TrainHyperParams,
|
||||
)
|
||||
from .infra_config import (
|
||||
CheckpointConfig,
|
||||
DebugConfig,
|
||||
DistributedConfig,
|
||||
LoggingConfig,
|
||||
)
|
||||
from .loader import load_config, save_config
|
||||
from .model_config import ModelConfig, QActModelConfig
|
||||
from .registry import (
|
||||
register_data_config,
|
||||
register_model_config,
|
||||
register_optimizer_config,
|
||||
register_scheduler_config,
|
||||
)
|
||||
from .task_config import TaskConfig
|
||||
from .train_config import TrainConfig
|
||||
|
||||
__all__ = [
|
||||
# Core
|
||||
"TrainConfig",
|
||||
"TaskConfig",
|
||||
# Model configs
|
||||
"ModelConfig",
|
||||
"QActModelConfig",
|
||||
# Data configs
|
||||
"DataConfig",
|
||||
"LeRobotDataConfig",
|
||||
# Hyperparams
|
||||
"TrainHyperParams",
|
||||
"OptimizerConfig",
|
||||
"AdamWConfig",
|
||||
"DMuonConfig",
|
||||
"SchedulerConfig",
|
||||
"CosineSchedulerConfig",
|
||||
"ConstantSchedulerConfig",
|
||||
"StepSchedulerConfig",
|
||||
# Infra configs
|
||||
"DistributedConfig",
|
||||
"LoggingConfig",
|
||||
"CheckpointConfig",
|
||||
"DebugConfig",
|
||||
# Functions
|
||||
"load_config",
|
||||
"save_config",
|
||||
"register_data_config",
|
||||
"register_model_config",
|
||||
"register_optimizer_config",
|
||||
"register_scheduler_config",
|
||||
]
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Public data config dataclasses.
|
||||
|
||||
Only data backends shipped in the public package should define config classes
|
||||
here. Internal backends register their dataclasses from their own packages via
|
||||
``wall_x.config.registry.register_data_config``.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from .registry import register_data_config
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataConfig:
|
||||
"""Base fields shared by data backends.
|
||||
|
||||
``normalizer_config`` may contain:
|
||||
- ``min_key``: stats key for the minimum value.
|
||||
- ``delta_key``: stats key for the value range.
|
||||
- ``customized_action_statistic_dof``: explicit action-stats JSON path.
|
||||
"""
|
||||
|
||||
dataset_type: str = "lerobot"
|
||||
resolution: Dict[str, int] = field(
|
||||
default_factory=lambda: {
|
||||
"face_view": 256,
|
||||
"left_wrist_view": 256,
|
||||
"right_wrist_view": 256,
|
||||
}
|
||||
)
|
||||
train_test_split: float = 0.95
|
||||
normalizer_config: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
@register_data_config("lerobot")
|
||||
@dataclass
|
||||
class LeRobotDataConfig(DataConfig):
|
||||
"""LeRobot data config.
|
||||
|
||||
``lerobot_config`` is expected to contain fields such as ``repo_id`` and
|
||||
``root`` for a HuggingFace LeRobot dataset. ``norm_stats_path`` points to
|
||||
explicit action normalizer stats; the core package does not bundle private
|
||||
defaults.
|
||||
"""
|
||||
|
||||
dataset_type: str = "lerobot"
|
||||
lerobot_config: Optional[Dict[str, Any]] = None
|
||||
key_mappings: Optional[Dict[str, Any]] = None
|
||||
norm_stats_path: Optional[str] = None
|
||||
priority_order: Optional[Dict[str, float]] = None
|
||||
camera_name_mapping: Optional[Dict[str, str]] = None
|
||||
num_workers: int = 4
|
||||
action_tokenizer_path: Optional[str] = None
|
||||
use_fast_tokenizer: bool = False
|
||||
padding_side: str = "left"
|
||||
noise_scheduler: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DataConfig",
|
||||
"LeRobotDataConfig",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -0,0 +1,98 @@
|
||||
"""Infrastructure config: distributed runtime, logging, checkpoints, and debug."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class DistributedConfig:
|
||||
# FSDP
|
||||
use_fsdp: bool = False
|
||||
fsdp_sharding_strategy: str = "full_shard"
|
||||
fsdp_backward_prefetch: str = "backward_pre"
|
||||
fsdp_cpu_offload: bool = False
|
||||
fsdp_use_orig_params: bool = True
|
||||
fsdp_limit_all_gathers: bool = True
|
||||
fsdp_forward_prefetch: bool = False
|
||||
fsdp_sync_module_states: bool = True
|
||||
fsdp_save_policy: str = "full"
|
||||
fsdp_hsdp_replicate_size: Optional[int] = None
|
||||
# Mixed precision
|
||||
use_mixed_precision: bool = True
|
||||
bf16: bool = True
|
||||
fsdp_reduce_dtype: str = "bf16"
|
||||
use_amp: bool = False
|
||||
use_gradient_checkpointing: bool = False
|
||||
use_gradient_checkpointing_offload: bool = False
|
||||
use_selective_recompute: bool = False
|
||||
# DDP fallback
|
||||
find_unused_parameters: bool = False
|
||||
broadcast_buffers: bool = True
|
||||
bucket_cap_mb: int = 25
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoggingConfig:
|
||||
log_name: str = "exp"
|
||||
log_project: str = "wallx"
|
||||
log_entity: Optional[str] = None
|
||||
use_wandb: bool = True
|
||||
wandb_offline: bool = False
|
||||
log_interval: int = 1
|
||||
save_interval: int = 1000
|
||||
val_interval: int = 4000
|
||||
epoch_save_interval: int = 1
|
||||
gc_interval_steps: int = 1000
|
||||
ignore_until_interval: int = 0
|
||||
# Rolling-window for smoothing per-step training metrics displayed on the
|
||||
# console (and reused by tqdm). 1 = raw per-step (historical behavior).
|
||||
# 10 = DZ-style 10-step rolling average - diffusion losses are dominated
|
||||
# by timestep-sampling noise per step; the smoothing only changes display,
|
||||
# not training. Independent of log_interval (which buffers for wandb).
|
||||
loss_log_smooth_window: int = 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointConfig:
|
||||
"""Checkpoint save and resume options.
|
||||
|
||||
``resume_from`` is shorthand for setting every component-specific resume
|
||||
path to the same checkpoint. Component-specific fields take precedence.
|
||||
"""
|
||||
|
||||
save_path: str = "./ckpt"
|
||||
validate_first: bool = False
|
||||
# Shorthand path for all components.
|
||||
resume_from: Optional[str] = None
|
||||
# Component-specific resume paths.
|
||||
resume_model: Optional[str] = None
|
||||
resume_optimizer: Optional[str] = None
|
||||
resume_scheduler: Optional[str] = None
|
||||
resume_ema: Optional[str] = None
|
||||
resume_rng: Optional[str] = None
|
||||
resume_data: Optional[str] = None
|
||||
resume_step: Optional[str] = None
|
||||
|
||||
def get_resume_path(self, component: str) -> Optional[str]:
|
||||
"""Return the resume path for one checkpoint component."""
|
||||
specific = getattr(self, f"resume_{component}", None)
|
||||
if specific is not None:
|
||||
return specific
|
||||
return self.resume_from
|
||||
|
||||
|
||||
@dataclass
|
||||
class DebugConfig:
|
||||
profile: bool = False
|
||||
profile_save_path: str = "./profile"
|
||||
profile_wait_iters: int = 1
|
||||
profile_warmup_iters: int = 1
|
||||
profile_active_iters: int = 3
|
||||
show_time_details: bool = False
|
||||
visualize_sample: bool = False
|
||||
save_debug_batch_path: Optional[str] = None
|
||||
nvtx: bool = False
|
||||
# Formula MFU uses the local FLOPs estimate and measured step time.
|
||||
enable_mfu: bool = False
|
||||
# Optional FLOPs profiling runs an extra forward at step 0.
|
||||
enable_mfu_profile: bool = False
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Wall-X config loader."""
|
||||
|
||||
import dataclasses
|
||||
import os
|
||||
import shutil
|
||||
from typing import Any, Type, TypeVar
|
||||
|
||||
import yaml
|
||||
|
||||
from .data_config import LeRobotDataConfig
|
||||
from .hyperparams_config import (
|
||||
AdamWConfig,
|
||||
LRGroupConfig,
|
||||
OptimizerConfig,
|
||||
SchedulerConfig,
|
||||
TrainHyperParams,
|
||||
)
|
||||
from .infra_config import (
|
||||
CheckpointConfig,
|
||||
DebugConfig,
|
||||
DistributedConfig,
|
||||
LoggingConfig,
|
||||
)
|
||||
from .model_config import ModelConfig, QActModelConfig
|
||||
from .registry import (
|
||||
get_data_config,
|
||||
get_model_config,
|
||||
get_optimizer_config,
|
||||
get_scheduler_config,
|
||||
registered_data_configs,
|
||||
registered_model_configs,
|
||||
registered_optimizer_configs,
|
||||
registered_scheduler_configs,
|
||||
)
|
||||
from .task_config import TaskConfig
|
||||
from .train_config import TrainConfig
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
_CONFIG_PLUGINS_LOADED = False
|
||||
|
||||
|
||||
class _TrainConfigSafeLoader(yaml.SafeLoader):
|
||||
pass
|
||||
|
||||
|
||||
def _construct_python_tuple(loader: yaml.SafeLoader, node: yaml.Node) -> tuple:
|
||||
return tuple(loader.construct_sequence(node))
|
||||
|
||||
|
||||
_TrainConfigSafeLoader.add_constructor(
|
||||
"tag:yaml.org,2002:python/tuple", _construct_python_tuple
|
||||
)
|
||||
|
||||
|
||||
def _ensure_config_plugins_loaded() -> None:
|
||||
"""Load optional internal config plugins when they are present."""
|
||||
global _CONFIG_PLUGINS_LOADED
|
||||
if _CONFIG_PLUGINS_LOADED:
|
||||
return
|
||||
_CONFIG_PLUGINS_LOADED = True
|
||||
try:
|
||||
from .internal_plugins import register_internal_config_plugins
|
||||
except ImportError:
|
||||
return
|
||||
register_internal_config_plugins()
|
||||
|
||||
|
||||
def load_config(config_path: str, cli_args: Any = None) -> TrainConfig:
|
||||
"""Load a ``TrainConfig`` from a YAML file."""
|
||||
_ensure_config_plugins_loaded()
|
||||
|
||||
with open(config_path, "r") as f:
|
||||
raw = yaml.load(f, Loader=_TrainConfigSafeLoader)
|
||||
|
||||
if raw is None:
|
||||
raise ValueError(f"Config file is empty: {config_path}")
|
||||
|
||||
model_type = raw.get("model_type")
|
||||
if model_type is None:
|
||||
raise ValueError(f"Missing required field 'model_type' in {config_path}")
|
||||
|
||||
config = TrainConfig(
|
||||
model_type=model_type,
|
||||
task=_build_dataclass(TaskConfig, raw.get("task", {})),
|
||||
model=_build_model_config(model_type, raw.get("model", {})),
|
||||
data=_build_data_config(raw.get("data", {})),
|
||||
hyperparams=_build_hyperparams(raw.get("hyperparams", {})),
|
||||
distributed=_build_dataclass(DistributedConfig, raw.get("distributed", {})),
|
||||
logging=_build_dataclass(LoggingConfig, raw.get("logging", {})),
|
||||
checkpoint=_build_dataclass(CheckpointConfig, raw.get("checkpoint", {})),
|
||||
debug=_build_dataclass(DebugConfig, raw.get("debug", {})),
|
||||
_raw_data=raw.get("data", {}),
|
||||
_raw_yaml=raw,
|
||||
dataset_path=raw.get("dataset_path"),
|
||||
)
|
||||
|
||||
if cli_args is not None:
|
||||
_apply_cli_overrides(config, cli_args)
|
||||
|
||||
_validate(config)
|
||||
|
||||
# Register the active data backend now that cfg is fully resolved
|
||||
# (post-CLI-override, post-validate). This is the single source of
|
||||
# truth for "which dataset backend is this run using".
|
||||
from wall_x.data._registry import _set_data_backend
|
||||
|
||||
_set_data_backend(config.data.dataset_type)
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def save_config(config: TrainConfig, save_dir: str) -> str:
|
||||
"""Save ``TrainConfig`` to ``config.yml`` under ``save_dir``."""
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
config_path = os.path.join(save_dir, "config.yml")
|
||||
|
||||
data = _sanitize_for_yaml(dataclasses.asdict(config))
|
||||
with open(config_path, "w") as f:
|
||||
yaml.dump(
|
||||
data, f, default_flow_style=False, allow_unicode=True, sort_keys=False
|
||||
)
|
||||
|
||||
dataset_config_path = getattr(config.data, "dataset_config_path", None)
|
||||
if dataset_config_path and os.path.exists(dataset_config_path):
|
||||
dst = os.path.join(save_dir, "dataset_config.yml")
|
||||
shutil.copy(dataset_config_path, dst)
|
||||
|
||||
return config_path
|
||||
|
||||
|
||||
def _sanitize_for_yaml(obj: Any) -> Any:
|
||||
"""Convert dataclass output to YAML-safe containers."""
|
||||
if isinstance(obj, dict):
|
||||
return {k: _sanitize_for_yaml(v) for k, v in obj.items()}
|
||||
elif isinstance(obj, (list, tuple)):
|
||||
return [_sanitize_for_yaml(v) for v in obj]
|
||||
return obj
|
||||
|
||||
|
||||
def _build_dataclass(cls: Type[T], raw: dict) -> T:
|
||||
"""Build a dataclass from a dict, ignoring unknown keys."""
|
||||
if not raw:
|
||||
return cls()
|
||||
|
||||
field_names = {f.name for f in dataclasses.fields(cls)}
|
||||
field_types = {f.name: f.type for f in dataclasses.fields(cls)}
|
||||
filtered = {}
|
||||
|
||||
for k, v in raw.items():
|
||||
if k not in field_names:
|
||||
continue
|
||||
ft = field_types[k]
|
||||
if (
|
||||
isinstance(ft, type)
|
||||
and dataclasses.is_dataclass(ft)
|
||||
and isinstance(v, dict)
|
||||
):
|
||||
filtered[k] = _build_dataclass(ft, v)
|
||||
else:
|
||||
filtered[k] = v
|
||||
|
||||
return cls(**filtered)
|
||||
|
||||
|
||||
def _build_model_config(model_type: str, raw: dict) -> ModelConfig:
|
||||
"""Build the registered model config for ``model_type``."""
|
||||
cls = get_model_config(model_type)
|
||||
if cls is None:
|
||||
raise ValueError(
|
||||
f"Unknown model_type: '{model_type}'. "
|
||||
f"Supported: {registered_model_configs()}"
|
||||
)
|
||||
return _build_dataclass(cls, raw)
|
||||
|
||||
|
||||
def _build_data_config(raw: dict):
|
||||
"""Build the registered data config for ``dataset_type``."""
|
||||
if not raw:
|
||||
return LeRobotDataConfig()
|
||||
|
||||
dataset_type = raw.get("dataset_type", "lerobot")
|
||||
cls = get_data_config(dataset_type)
|
||||
if cls is None:
|
||||
raise ValueError(
|
||||
f"Unknown dataset_type: '{dataset_type}'. "
|
||||
f"Supported: {registered_data_configs()}"
|
||||
)
|
||||
|
||||
return _build_dataclass(cls, raw)
|
||||
|
||||
|
||||
def _build_optimizer_config(raw: dict) -> OptimizerConfig:
|
||||
"""Build the registered optimizer config for ``optimizer_type``."""
|
||||
if not raw:
|
||||
return AdamWConfig()
|
||||
|
||||
raw = dict(raw)
|
||||
optimizer_type = raw.get("optimizer_type", "adamw")
|
||||
cls = get_optimizer_config(optimizer_type)
|
||||
if cls is None:
|
||||
raise ValueError(
|
||||
f"Unknown optimizer_type: '{optimizer_type}'. "
|
||||
f"Supported: {registered_optimizer_configs()}"
|
||||
)
|
||||
|
||||
if "betas" in raw and isinstance(raw["betas"], list):
|
||||
raw["betas"] = tuple(raw["betas"])
|
||||
if "adamw_betas" in raw and isinstance(raw["adamw_betas"], list):
|
||||
raw["adamw_betas"] = tuple(raw["adamw_betas"])
|
||||
if raw.get("lr_groups") is not None:
|
||||
raw["lr_groups"] = [
|
||||
_build_dataclass(LRGroupConfig, group) for group in raw["lr_groups"]
|
||||
]
|
||||
|
||||
return _build_dataclass(cls, raw)
|
||||
|
||||
|
||||
def _build_scheduler_config(raw: dict) -> SchedulerConfig:
|
||||
"""Build the registered scheduler config for ``scheduler_type``."""
|
||||
if not raw:
|
||||
cls = get_scheduler_config("cosine")
|
||||
if cls is None:
|
||||
raise ValueError("Scheduler config 'cosine' is not registered")
|
||||
return cls()
|
||||
|
||||
scheduler_type = raw.get("scheduler_type", "cosine")
|
||||
cls = get_scheduler_config(scheduler_type)
|
||||
if cls is None:
|
||||
raise ValueError(
|
||||
f"Unknown scheduler_type: '{scheduler_type}'. "
|
||||
f"Supported: {registered_scheduler_configs()}"
|
||||
)
|
||||
return _build_dataclass(cls, raw)
|
||||
|
||||
|
||||
def _build_hyperparams(raw: dict) -> TrainHyperParams:
|
||||
"""Build ``TrainHyperParams`` with polymorphic optimizer/scheduler config."""
|
||||
if not raw:
|
||||
return TrainHyperParams()
|
||||
|
||||
raw = dict(raw) # shallow copy to avoid mutating caller's dict
|
||||
optimizer_raw = raw.pop("optimizer", {})
|
||||
scheduler_raw = raw.pop("scheduler", {})
|
||||
|
||||
optimizer = _build_optimizer_config(optimizer_raw)
|
||||
scheduler = _build_scheduler_config(scheduler_raw)
|
||||
|
||||
hp = _build_dataclass(TrainHyperParams, raw)
|
||||
hp.optimizer = optimizer
|
||||
hp.scheduler = scheduler
|
||||
return hp
|
||||
|
||||
|
||||
def _apply_cli_overrides(config: TrainConfig, args: Any) -> None:
|
||||
"""Apply CLI overrides to the loaded config."""
|
||||
if getattr(args, "fsdp_sharding_strategy", None) is not None:
|
||||
config.distributed.fsdp_sharding_strategy = args.fsdp_sharding_strategy
|
||||
|
||||
if getattr(args, "debug", False):
|
||||
config.logging.log_name = "debug"
|
||||
config.logging.log_project = "debug"
|
||||
config.checkpoint.save_path = "./ckpt/debug"
|
||||
|
||||
if getattr(args, "visualize", False):
|
||||
config.debug.visualize_sample = True
|
||||
|
||||
if getattr(args, "wandb_offline", None) is not None:
|
||||
config.logging.wandb_offline = args.wandb_offline in ("true", "True", "1")
|
||||
|
||||
|
||||
def _validate(config: TrainConfig) -> None:
|
||||
"""Validate required fields and model-specific config."""
|
||||
if not config.model_type:
|
||||
raise ValueError("model_type is required")
|
||||
|
||||
if not config.task.dof_config:
|
||||
raise ValueError("task.dof_config is required (cannot be empty)")
|
||||
|
||||
if isinstance(config.model, QActModelConfig):
|
||||
model = config.model
|
||||
if not model.config_path:
|
||||
raise ValueError("model.config_path is required for QAct models")
|
||||
if not model.processor_path:
|
||||
raise ValueError("model.processor_path is required for QAct models")
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Base model config dataclasses.
|
||||
|
||||
Optional model variants register their config dataclasses from their own
|
||||
packages via ``wall_x.config.registry.register_model_config``.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .registry import register_model_config
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""Base model architecture fields."""
|
||||
|
||||
use_ema: bool = False
|
||||
attn_implementation: Optional[str] = None
|
||||
attn_deterministic: Optional[bool] = None
|
||||
ar_loss_weight: float = 1.0
|
||||
|
||||
|
||||
@register_model_config("qwen2_5")
|
||||
@dataclass
|
||||
class QActModelConfig(ModelConfig):
|
||||
"""QAct model config for the Qwen2.5 VLA path."""
|
||||
|
||||
config_path: str = ""
|
||||
processor_path: str = ""
|
||||
pretrained_path: Optional[str] = None
|
||||
backbone: str = "qwen2_5"
|
||||
action_tokenizer_type: Optional[str] = None
|
||||
action_tokenizer_path: Optional[str] = None
|
||||
action_tokenizer_checkpoint_path: Optional[str] = None
|
||||
action_tokenizer_config_dir: Optional[str] = None
|
||||
new_special_tokens: Optional[List[str]] = None
|
||||
flow_loss_weight: float = 1.0
|
||||
enable_customized_robot_config: bool = False
|
||||
customized_robot_config: Optional[Dict[str, Any]] = None
|
||||
|
||||
|
||||
__all__ = ["ModelConfig", "QActModelConfig"]
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Runtime registries for typed config variants.
|
||||
|
||||
Core Wall-X keeps only base config types in ``wall_x.config``. Optional
|
||||
datasets, model families, and optimizer variants register their dataclasses
|
||||
from their own packages, so trimmed distributions do not leave dead config
|
||||
entries behind.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
_DATA_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||||
_MODEL_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||||
_OPTIMIZER_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||||
_SCHEDULER_CONFIG_REGISTRY: dict[str, type[Any]] = {}
|
||||
|
||||
|
||||
def _register(
|
||||
registry: dict[str, type[Any]], names: tuple[str, ...], cls: type[T]
|
||||
) -> type[T]:
|
||||
for name in names:
|
||||
if not name:
|
||||
raise ValueError("Config registry name must be non-empty")
|
||||
existing = registry.get(name)
|
||||
if existing is not None and existing is not cls:
|
||||
raise ValueError(
|
||||
f"Config name {name!r} is already registered by "
|
||||
f"{existing.__module__}.{existing.__name__}"
|
||||
)
|
||||
registry[name] = cls
|
||||
return cls
|
||||
|
||||
|
||||
def register_data_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||||
return lambda cls: _register(_DATA_CONFIG_REGISTRY, names, cls)
|
||||
|
||||
|
||||
def register_model_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||||
return lambda cls: _register(_MODEL_CONFIG_REGISTRY, names, cls)
|
||||
|
||||
|
||||
def register_optimizer_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||||
return lambda cls: _register(_OPTIMIZER_CONFIG_REGISTRY, names, cls)
|
||||
|
||||
|
||||
def register_scheduler_config(*names: str) -> Callable[[type[T]], type[T]]:
|
||||
return lambda cls: _register(_SCHEDULER_CONFIG_REGISTRY, names, cls)
|
||||
|
||||
|
||||
def get_data_config(name: str) -> type[Any] | None:
|
||||
return _DATA_CONFIG_REGISTRY.get(name)
|
||||
|
||||
|
||||
def get_model_config(name: str) -> type[Any] | None:
|
||||
return _MODEL_CONFIG_REGISTRY.get(name)
|
||||
|
||||
|
||||
def get_optimizer_config(name: str) -> type[Any] | None:
|
||||
return _OPTIMIZER_CONFIG_REGISTRY.get(name)
|
||||
|
||||
|
||||
def get_scheduler_config(name: str) -> type[Any] | None:
|
||||
return _SCHEDULER_CONFIG_REGISTRY.get(name)
|
||||
|
||||
|
||||
def registered_data_configs() -> list[str]:
|
||||
return sorted(_DATA_CONFIG_REGISTRY)
|
||||
|
||||
|
||||
def registered_model_configs() -> list[str]:
|
||||
return sorted(_MODEL_CONFIG_REGISTRY)
|
||||
|
||||
|
||||
def registered_optimizer_configs() -> list[str]:
|
||||
return sorted(_OPTIMIZER_CONFIG_REGISTRY)
|
||||
|
||||
|
||||
def registered_scheduler_configs() -> list[str]:
|
||||
return sorted(_SCHEDULER_CONFIG_REGISTRY)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"get_data_config",
|
||||
"get_model_config",
|
||||
"get_optimizer_config",
|
||||
"get_scheduler_config",
|
||||
"register_data_config",
|
||||
"register_model_config",
|
||||
"register_optimizer_config",
|
||||
"register_scheduler_config",
|
||||
"registered_data_configs",
|
||||
"registered_model_configs",
|
||||
"registered_optimizer_configs",
|
||||
"registered_scheduler_configs",
|
||||
]
|
||||
@@ -0,0 +1,17 @@
|
||||
"""Task config shared by model and data code."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class TaskConfig:
|
||||
"""Robot task definition shared by model construction and data slicing."""
|
||||
|
||||
dof_config: Dict[str, int] = field(default_factory=dict)
|
||||
agent_pos_config: Dict[str, int] = field(default_factory=dict)
|
||||
ar_dof_config: Optional[Dict[str, int]] = None
|
||||
action_horizon: int = 32
|
||||
action_horizon_flow: int = 32
|
||||
noise_scheduler: Optional[Dict[str, Any]] = None
|
||||
use_state_string_representation: bool = False
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Top-level training config."""
|
||||
|
||||
import dataclasses
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict
|
||||
|
||||
from .data_config import DataConfig
|
||||
from .hyperparams_config import TrainHyperParams
|
||||
from .infra_config import (
|
||||
CheckpointConfig,
|
||||
DebugConfig,
|
||||
DistributedConfig,
|
||||
LoggingConfig,
|
||||
)
|
||||
from .model_config import ModelConfig
|
||||
from .task_config import TaskConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainConfig:
|
||||
"""Top-level Wall-X training config."""
|
||||
|
||||
model_type: str = "qwen2_5"
|
||||
task: TaskConfig = field(default_factory=TaskConfig)
|
||||
model: ModelConfig = field(default_factory=ModelConfig)
|
||||
data: DataConfig = field(default_factory=DataConfig)
|
||||
hyperparams: TrainHyperParams = field(default_factory=TrainHyperParams)
|
||||
distributed: DistributedConfig = field(default_factory=DistributedConfig)
|
||||
logging: LoggingConfig = field(default_factory=LoggingConfig)
|
||||
checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
|
||||
debug: DebugConfig = field(default_factory=DebugConfig)
|
||||
# Raw YAML sections preserved verbatim for backend APIs that read fields
|
||||
# not captured by the typed DataConfig dataclass.
|
||||
_raw_data: Dict[str, Any] = field(default_factory=dict)
|
||||
# Full raw YAML dict for backend-specific compatibility paths.
|
||||
_raw_yaml: Dict[str, Any] = field(default_factory=dict)
|
||||
# dataset_path lives at top level in YAML, consumed by data loaders directly
|
||||
dataset_path: Any = None
|
||||
|
||||
@property
|
||||
def action_dim(self) -> int:
|
||||
return sum(self.task.dof_config.values())
|
||||
|
||||
@property
|
||||
def propri_dim(self) -> int:
|
||||
return sum(self.task.agent_pos_config.values())
|
||||
|
||||
def build_data_loader_dict(self) -> Dict[str, Any]:
|
||||
"""Build the raw dict consumed by backend-specific compatibility paths.
|
||||
|
||||
Merges ``_raw_data`` (verbatim YAML ``data:`` section) with task
|
||||
fields (dof_config, action_horizon, etc.) and hyperparams
|
||||
(batch_size). Backend compatibility APIs may read fields that the
|
||||
typed DataConfig dataclass does not carry.
|
||||
|
||||
This keeps legacy flat configs working while the main config surface
|
||||
stays typed.
|
||||
"""
|
||||
# Start with the raw YAML data section, then add typed task defaults.
|
||||
data_dict = dict(self._raw_data)
|
||||
task_dict = dataclasses.asdict(self.task)
|
||||
for key in (
|
||||
"dof_config",
|
||||
"agent_pos_config",
|
||||
"action_horizon",
|
||||
"action_horizon_flow",
|
||||
"ar_dof_config",
|
||||
"use_state_string_representation",
|
||||
):
|
||||
if key in task_dict and task_dict[key] is not None:
|
||||
data_dict.setdefault(key, task_dict[key])
|
||||
data_dict.setdefault("batch_size_per_gpu", self.hyperparams.batch_size_per_gpu)
|
||||
data_dict.setdefault("batch_size", self.hyperparams.batch_size_per_gpu)
|
||||
result: Dict[str, Any] = {
|
||||
"model_type": self.model_type,
|
||||
"data": data_dict,
|
||||
**task_dict,
|
||||
}
|
||||
if self.dataset_path is not None:
|
||||
result["dataset_path"] = self.dataset_path
|
||||
return result
|
||||
Reference in New Issue
Block a user