refactor: enhance configuration and modularity in aimx-hydra-lightnin…

skills/aimx-hydra-lightning-builder/SKILL.md

@@ -54,13 +54,23 @@ Never edit, format, sync dependencies, generate files, or run mutation/codegen c
 
 Read `references/architecture.md` before scaffold or migration work.
 
-- `configs/<task>.yaml` composes `data`, `datamodule`, `model`, `plmodule`, `trainer`, `callbacks`, `logger`, `paths`, `accelerate`, and optional `experiment`.
+- `configs/<task>.yaml` composes `datamodule`, `model`, `plmodule`, `trainer`, `callbacks`, `logger`, `paths`, `accelerate`, `opt`, and `experiment`.
+- Keep `configs/<task>.yaml` as the baseline and select experiment deltas with `experiment=<name>`, where `configs/experiment/<name>.yaml` uses Hydra `override` defaults and parameter overrides.
 - `src/train.py` seeds, instantiates configured objects, logs hyperparameters, and calls `trainer.fit/validate/test`.
 - `BaseLitModule` owns `cfg`, `cfg.model` instantiation, optimizer/scheduler, compile/SDPA options, and shared trace helpers.
 - Task modules own batch parsing, loss, metrics, and prediction/evaluation outputs.
-- DataModules own splits, dataloaders, sampler/collate policy, and data preparation boundaries.
+- DataModules own splits, dataloaders, sampler/collate policy, and data preparation boundaries. Prefer dataset samples and batches as pytrees so task modules can evolve without positional tuple churn.
 - Aim trace uses Lightning loggers for scalars and explicit `experiment.track(...)` for images/distributions.
 
+## Design Principles
+
+- Keep high cohesion inside modules and low coupling across modules.
+- Let config define how an experiment runs; let code define what the domain operation means.
+- Keep inheritance trees shallow and explicit. Prefer composition through Hydra-configured modules when behavior varies.
+- Keep baseline defaults separate from experiment deltas. Experiments should override choices and parameters, not copy whole config trees.
+- Keep optimizer and scheduler policy in `opt`; experiments override `opt` values instead of hiding optimizer settings under `model` or `trainer`.
+- Use domain adapters for domain-specific behavior. Shared bases define contracts and common mechanics; child adapters implement radar, satellite, vision-frame, tabular, sequence, or other domain semantics.
+
 ## References
 
 - `references/architecture.md`: core relationships and file layout.

skills/aimx-hydra-lightning-builder/assets/template-repo/README.md

@@ -7,13 +7,18 @@ Hydra + Lightning + Aim template for Aimx AutoResearch.
 ```bash
 uv sync
 uv run python src/train.py trainer.fast_dev_run=true trainer.logger=false
+uv run python src/train.py experiment=exp trainer.fast_dev_run=true trainer.logger=false
 uv run pytest
 ```
 
+Use `experiment=<name>` to apply a file from `configs/experiment/<name>.yaml`.
+Experiment yaml files should override config groups and values such as
+`model`, `datamodule`, `trainer`, `opt`, `accelerate`, and `logger`.
+
 Enable Aim logging by leaving `trainer.logger=true` and using `logger=aim`.
 
 ```bash
-uv run python src/train.py
+uv run python src/train.py experiment=exp
 aimx query params "run.hash != ''" --repo .
 aimx query metrics "metric.name != ''" --repo .
 aimx query metrics "metric.name == 'acc'" --repo . --json

skills/aimx-hydra-lightning-builder/assets/template-repo/configs/accelerate/default.yaml

@@ -1,3 +1,4 @@
 compile: false
 precision: "32-true"
 fp32_matmul_precision: "highest"
+sdpa: ["efficient", "flash", "math"]

skills/aimx-hydra-lightning-builder/assets/template-repo/configs/experiment/exp.yaml

@@ -0,0 +1,39 @@
+# @package _global_
+
+# Run with:
+# uv run python src/train.py experiment=exp
+
+defaults:
+  - override /datamodule: dummy
+  - override /model: mlp
+  - override /plmodule: classifier
+  - override /callbacks: default
+  - override /trainer: default
+  - override /opt: default
+  - override /accelerate: default
+  - override /logger: aim
+
+task_name: train_exp
+tags: ["exp", "{{ preset }}"]
+
+seed: 42
+
+autoresearch:
+  experiment_name: "{{ project_name }}-exp"
+
+trainer:
+  max_epochs: 2
+  gradient_clip_val: 0.5
+
+datamodule:
+  batch_size: 32
+
+model:
+  hidden_dim: 32
+
+opt:
+  optimizer:
+    lr: 0.002
+
+accelerate:
+  compile: false

skills/aimx-hydra-lightning-builder/assets/template-repo/configs/opt/default.yaml

@@ -2,3 +2,9 @@ optimizer:
   _target_: torch.optim.AdamW
   lr: 0.001
   weight_decay: 0.0
+
+scheduler:
+  _target_: torch.optim.lr_scheduler.CosineAnnealingLR
+  _partial_: True
+  T_max: 10
+  eta_min: 0

skills/aimx-hydra-lightning-builder/assets/template-repo/configs/train.yaml

@@ -10,6 +10,7 @@ defaults:
   - logger: aim
   - opt: default
   - accelerate: default
+  - experiment: null
 
 task_name: train
 tags: ["dev", "{{ preset }}"]

skills/aimx-hydra-lightning-builder/assets/template-repo/src/__package__/datamodules/dummy.py

@@ -2,7 +2,22 @@
 
 import torch
 from lightning import LightningDataModule
-from torch.utils.data import DataLoader, TensorDataset, random_split
+from torch.utils.data import DataLoader, Dataset, random_split
+
+
+class PytreeClassificationDataset(Dataset):
+    def __init__(self, x: torch.Tensor, y: torch.Tensor) -> None:
+        self.x = x
+        self.y = y
+
+    def __len__(self) -> int:
+        return int(self.x.shape[0])
+
+    def __getitem__(self, index: int) -> dict[str, dict[str, torch.Tensor]]:
+        return {
+            "input": {"x": self.x[index]},
+            "target": {"label": self.y[index]},
+        }
 
 
 class RandomClassificationDataModule(LightningDataModule):
@@ -25,7 +40,7 @@ def setup(self, stage: str | None = None) -> None:
         x = torch.randn(int(self.hparams.num_samples), int(self.hparams.num_features), generator=generator)
         weights = torch.randn(int(self.hparams.num_features), int(self.hparams.num_classes), generator=generator)
         y = torch.argmax(x @ weights, dim=1)
-        dataset = TensorDataset(x, y)
+        dataset = PytreeClassificationDataset(x, y)
         train_len = max(1, int(0.8 * len(dataset)))
         val_len = len(dataset) - train_len
         self.train_dataset, self.val_dataset = random_split(dataset, [train_len, val_len], generator=generator)

skills/aimx-hydra-lightning-builder/assets/template-repo/src/__package__/plmodules/__init__.py

@@ -1,3 +1,93 @@
+from __future__ import annotations
+
+import hydra
+import lightning as L
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from omegaconf import DictConfig
+
+
+class BaseLitModule(L.LightningModule):
+    def __init__(self, cfg: DictConfig) -> None:
+        super().__init__()
+
+        self.save_hyperparameters(logger=False)
+        self.cfg = cfg
+        self.net = hydra.utils.instantiate(cfg.model)
+        self._net_compiled = False
+
+        sdpa_map = {
+            "cudnn": SDPBackend.CUDNN_ATTENTION,
+            "math": SDPBackend.MATH,
+            "efficient": SDPBackend.EFFICIENT_ATTENTION,
+            "flash": SDPBackend.FLASH_ATTENTION,
+        }
+
+        self.sdpa_backends = [sdpa_map[backend] for backend in self.cfg.accelerate.get("sdpa", ["math"])]
+
+    def forward(self, *args, **kwargs):
+        return self._model_forward(*args, **kwargs)
+
+    def _model_forward(self, *args, **kwargs):
+        with sdpa_kernel(self.sdpa_backends):
+            return self.net(*args, **kwargs)
+
+    def setup(self, stage: str) -> None:
+        if self.cfg.accelerate.compile and stage == "fit" and hasattr(torch, "compile") and not self._net_compiled:
+            self.net = torch.compile(self.net)
+            self._net_compiled = True
+
+    def get_lr_scheduler(self, optimizer):
+        scheduler = hydra.utils.instantiate(self.cfg.opt.scheduler)(optimizer=optimizer)
+        kwargs = {
+            key: value for key, value in self.cfg.opt.items() if key not in ["optimizer", "scheduler"]
+        }
+        return {
+            "scheduler": scheduler,
+            **kwargs,
+        }
+
+    def get_optimizer(self):
+        if self.cfg.opt.optimizer._target_ == "torch.optim.AdamW":
+            optimizer = hydra.utils.instantiate(
+                self.cfg.opt.optimizer,
+                params=filter(lambda p: p.requires_grad, self.net.parameters()),
+            )
+        elif self.cfg.opt.optimizer._target_ == "colossalai.nn.optimizer.HybridAdam":
+            optimizer = hydra.utils.instantiate(
+                self.cfg.opt.optimizer,
+                model_params=filter(lambda p: p.requires_grad, self.net.parameters()),
+            )
+        else:
+            optimizer = hydra.utils.instantiate(
+                self.cfg.opt.optimizer,
+                params=filter(lambda p: p.requires_grad, self.net.parameters()),
+            )
+        return optimizer
+
+    def configure_optimizers(self):
+        optimizer = self.get_optimizer()
+        if not self.cfg.opt.get("scheduler"):
+            return optimizer
+
+        lr_scheduler = self.get_lr_scheduler(optimizer)
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": lr_scheduler,
+        }
+
+    def _aim_experiments(self):
+        for logger in self.loggers:
+            experiment = getattr(logger, "experiment", None)
+            if experiment is not None and hasattr(experiment, "track"):
+                yield experiment
+
+    def _instantiate_metric(self, name: str, defaults: dict[str, dict[str, object]]):
+        metrics_cfg = self.cfg.get("metrics", {})
+        metric_cfg = metrics_cfg[name] if name in metrics_cfg else defaults[name]
+        return hydra.utils.instantiate(metric_cfg)
+
+
 from {{ package_name }}.plmodules.classifier import ClassificationModule
 
-__all__ = ["ClassificationModule"]
+__all__ = ["BaseLitModule", "ClassificationModule"]

skills/aimx-hydra-lightning-builder/assets/template-repo/src/__package__/plmodules/classifier.py

@@ -1,38 +1,39 @@
 from __future__ import annotations
 
-import hydra
-import lightning as L
 import torch
 import torch.nn.functional as F
 from omegaconf import DictConfig
 
+from {{ package_name }}.plmodules import BaseLitModule
 
-class ClassificationModule(L.LightningModule):
+
+class ClassificationModule(BaseLitModule):
     def __init__(self, cfg: DictConfig) -> None:
-        super().__init__()
-        self.save_hyperparameters(logger=False)
-        self.cfg = cfg
-        self.net = hydra.utils.instantiate(cfg.model)
+        super().__init__(cfg)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.net(x)
+    def _parse_batch(self, batch: dict[str, dict[str, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:
+        return batch["input"]["x"], batch["target"]["label"]
 
-    def _shared_step(self, batch, mode: str) -> torch.Tensor:
-        x, y = batch
+    def _shared_step(self, batch, mode: str) -> dict[str, torch.Tensor]:
+        x, y = self._parse_batch(batch)
         logits = self(x)
-        loss = F.cross_entropy(logits, y)
         preds = torch.argmax(logits, dim=1)
-        acc = (preds == y).float().mean()
-        on_step = mode == "train"
-        self.log(f"{mode}/loss", loss, on_step=on_step, on_epoch=True, prog_bar=True)
-        self.log(f"{mode}/acc", acc, on_step=on_step, on_epoch=True, prog_bar=True)
-        return loss
+        res = {
+            "y_hat": preds,
+            "y": y,
+        }
+        if mode in ["train", "val"]:
+            loss = F.cross_entropy(logits, y)
+            acc = (preds == y).float().mean()
+            on_step = mode == "train"
+            self.log(f"{mode}/loss", loss, on_step=on_step, on_epoch=True, prog_bar=True)
+            self.log(f"{mode}/acc", acc, on_step=on_step, on_epoch=True, prog_bar=True)
+            res["loss"] = loss
+        return res
 
     def training_step(self, batch, batch_idx: int) -> torch.Tensor:
-        return self._shared_step(batch, "train")
-
-    def validation_step(self, batch, batch_idx: int) -> None:
-        self._shared_step(batch, "val")
+        res = self._shared_step(batch, "train")
+        return res["loss"]
 
-    def configure_optimizers(self):
-        return hydra.utils.instantiate(self.cfg.opt.optimizer, params=self.parameters())
+    def validation_step(self, batch, batch_idx: int) -> dict[str, torch.Tensor]:
+        return self._shared_step(batch, "val")

skills/aimx-hydra-lightning-builder/assets/template-repo/tests/test_fast_dev_run.py

@@ -3,12 +3,16 @@
 import subprocess
 import sys
 
+import pytest
 
-def test_fast_dev_run() -> None:
+
+@pytest.mark.parametrize("overrides", [(), ("experiment=exp",)])
+def test_fast_dev_run(overrides: tuple[str, ...]) -> None:
     result = subprocess.run(
         [
             sys.executable,
             "src/train.py",
+            *overrides,
             "trainer.fast_dev_run=true",
             "trainer.logger=false",
             "trainer.enable_progress_bar=false",
@@ -18,3 +22,10 @@ def test_fast_dev_run() -> None:
         text=True,
     )
     assert result.returncode == 0, result.stderr
+
+
+def test_plmodule_exports() -> None:
+    from {{ package_name }}.plmodules import BaseLitModule, ClassificationModule
+
+    assert BaseLitModule.__name__ == "BaseLitModule"
+    assert ClassificationModule.__name__ == "ClassificationModule"