Add Colab notebook for Side-Step LoRA/LoKR training on ACE-Step v1.5 …

acestep/training_v2/adapter_resume.py

@@ -0,0 +1,61 @@
+"""Restore saved LoRA adapter weights into a live decoder when resuming.
+
+PEFT's ``save_pretrained`` strips the adapter name from tensor keys
+(``...lora_A.weight``), while a live ``PeftModel`` names them
+``...lora_A.default.weight``.  A plain ``load_state_dict(strict=False)``
+therefore matches nothing and silently resumes from an untrained adapter.
+This module loads through ``set_peft_model_state_dict`` (which restores the
+adapter name) and refuses to continue when no saved tensor was applied.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class AdapterResumeError(RuntimeError):
+    """Raised when checkpoint adapter weights cannot be applied to the model."""
+
+
+def load_adapter_state_dict(decoder: Any, state_dict: Dict[str, torch.Tensor]) -> int:
+    """Apply saved adapter weights to *decoder*.
+
+    Args:
+        decoder: The training decoder, usually a PEFT ``PeftModel``.
+        state_dict: Tensors read from ``adapter_model.safetensors``.
+
+    Returns:
+        Number of saved tensors that were applied to the decoder.
+
+    Raises:
+        AdapterResumeError: If the file is empty or no tensor matches the decoder.
+    """
+    if not state_dict:
+        raise AdapterResumeError("Checkpoint adapter file contains no tensors")
+
+    if hasattr(decoder, "peft_config"):
+        from peft import set_peft_model_state_dict
+
+        result = set_peft_model_state_dict(decoder, state_dict)
+    else:
+        result = decoder.load_state_dict(state_dict, strict=False)
+
+    unexpected = list(getattr(result, "unexpected_keys", None) or [])
+    loaded = len(state_dict) - len(unexpected)
+    if loaded <= 0:
+        raise AdapterResumeError(
+            "None of the checkpoint adapter weights match the model "
+            "(check rank, alpha and target modules are identical to the original run)"
+        )
+    if unexpected:
+        logger.warning(
+            "[Side-Step] %d checkpoint tensors did not match the model, e.g. %s",
+            len(unexpected),
+            unexpected[0],
+        )
+    return loaded

acestep/training_v2/adapter_resume_test.py

@@ -0,0 +1,99 @@
+"""Tests for acestep.training_v2.adapter_resume.load_adapter_state_dict.
+
+Uses a tiny real PEFT LoRA model so the key naming of ``save_pretrained``
+(adapter name stripped) is exercised exactly as in a Side-Step checkpoint.
+"""
+
+import os
+import tempfile
+import unittest
+
+try:
+    import torch
+    import torch.nn as nn
+    from peft import LoraConfig, get_peft_model
+    from safetensors.torch import load_file
+
+    _HAS_DEPS = True
+except ImportError:  # pragma: no cover - depends on the environment
+    _HAS_DEPS = False
+
+from acestep.training_v2.adapter_resume import AdapterResumeError, load_adapter_state_dict
+
+
+def _make_peft_model(seed: int) -> "nn.Module":
+    """Return a small PEFT LoRA model with non-zero, seed-dependent adapter weights."""
+
+    class Block(nn.Module):
+        """Minimal block with attention- and MLP-like projections."""
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.q_proj = nn.Linear(8, 8)
+            self.gate_proj = nn.Linear(8, 8)
+
+        def forward(self, x):
+            return self.gate_proj(self.q_proj(x))
+
+    torch.manual_seed(0)
+    model = get_peft_model(
+        Block(), LoraConfig(r=4, lora_alpha=8, target_modules=["q_proj", "gate_proj"])
+    )
+    torch.manual_seed(seed)
+    for name, param in model.named_parameters():
+        if "lora_" in name:
+            nn.init.normal_(param)
+    return model
+
+
+def _lora_tensors(model: "nn.Module") -> dict:
+    """Return a copy of the adapter tensors keyed by parameter name."""
+    return {n: p.detach().clone() for n, p in model.named_parameters() if "lora_" in n}
+
+
+@unittest.skipUnless(_HAS_DEPS, "torch, peft and safetensors are required")
+class LoadAdapterStateDictTest(unittest.TestCase):
+    """Resume must restore the trained adapter, never silently skip it."""
+
+    def setUp(self) -> None:
+        self.trained = _make_peft_model(seed=1)
+        self.tmp = tempfile.TemporaryDirectory()
+        self.trained.save_pretrained(self.tmp.name)
+        self.saved = load_file(os.path.join(self.tmp.name, "adapter_model.safetensors"))
+
+    def tearDown(self) -> None:
+        self.tmp.cleanup()
+
+    def test_restores_peft_weights_saved_by_save_pretrained(self) -> None:
+        fresh = _make_peft_model(seed=2)
+        loaded = load_adapter_state_dict(fresh, self.saved)
+        self.assertEqual(loaded, len(self.saved))
+        expected = _lora_tensors(self.trained)
+        for name, tensor in _lora_tensors(fresh).items():
+            self.assertTrue(torch.equal(tensor, expected[name]), name)
+
+    def test_plain_strict_false_load_matches_nothing(self) -> None:
+        """Regression: the previous loader applied none of the saved tensors."""
+        fresh = _make_peft_model(seed=2)
+        result = fresh.load_state_dict(self.saved, strict=False)
+        self.assertEqual(len(result.unexpected_keys), len(self.saved))
+
+    def test_raises_when_no_tensor_matches(self) -> None:
+        fresh = _make_peft_model(seed=2)
+        wrong = {"base_model.model.other.lora_A.weight": torch.zeros(4, 8)}
+        with self.assertRaises(AdapterResumeError):
+            load_adapter_state_dict(fresh, wrong)
+
+    def test_raises_on_empty_file(self) -> None:
+        with self.assertRaises(AdapterResumeError):
+            load_adapter_state_dict(_make_peft_model(seed=2), {})
+
+    def test_non_peft_module_uses_plain_state_dict(self) -> None:
+        plain = nn.Linear(4, 4)
+        weights = {"weight": torch.ones(4, 4), "bias": torch.zeros(4)}
+        self.assertEqual(load_adapter_state_dict(plain, weights), 2)
+        self.assertTrue(torch.equal(plain.weight, torch.ones(4, 4)))
+
+
+if __name__ == "__main__":
+    unittest.main()

acestep/training_v2/cli/args.py

@@ -19,6 +19,10 @@
     "turbo": "acestep-v15-turbo",
     "base": "acestep-v15-base",
     "sft": "acestep-v15-sft",
+    # XL (4B DiT) variants -- same names as model_loader._VARIANT_DIR
+    "xl_turbo": "acestep-v15-xl-turbo",
+    "xl_base": "acestep-v15-xl-base",
+    "xl_sft": "acestep-v15-xl-sft",
 }
 
 
@@ -326,6 +330,15 @@ def _add_fixed_args(parser: argparse.ArgumentParser) -> None:
     """Add arguments specific to the fixed subcommand."""
     g = parser.add_argument_group("Corrected training")
     g.add_argument("--cfg-ratio", type=float, default=0.15, help="CFG dropout probability (default: 0.15)")
+    g.add_argument(
+        "--timestep-mu",
+        type=float,
+        default=None,
+        help=(
+            "Override the logit-normal timestep mean (default: model config, -0.4 for base). "
+            "Lower values (e.g. -0.8) train more on low-noise steps, i.e. fine detail"
+        ),
+    )
 
 
 

acestep/training_v2/cli/config_builder.py

@@ -87,6 +87,12 @@ def build_configs(args: argparse.Namespace) -> Tuple[AdapterConfig, TrainingConf
         timestep_mu = defaults.get("timestep_mu", timestep_mu)
         timestep_sigma = defaults.get("timestep_sigma", timestep_sigma)
 
+    # -- Explicit override (e.g. focus training on low-noise detail) --------
+    timestep_mu_override = getattr(args, "timestep_mu", None)
+    if timestep_mu_override is not None:
+        logger.info("[Side-Step] timestep_mu overridden: %s -> %s", timestep_mu, timestep_mu_override)
+        timestep_mu = timestep_mu_override
+
     # -- GPU info -----------------------------------------------------------
     gpu_info = detect_gpu(
         requested_device=args.device,

acestep/training_v2/cli/config_builder_timestep_test.py

@@ -0,0 +1,60 @@
+"""Tests for the --timestep-mu override in build_configs."""
+
+import json
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from acestep.training_v2.cli.args import build_root_parser
+from acestep.training_v2.cli.config_builder import build_configs
+
+
+def _parse(ckpt_dir: str, *extra: str):
+    """Parse a minimal `fixed` command line pointing at *ckpt_dir*."""
+    argv = [
+        "fixed", "--checkpoint-dir", ckpt_dir, "--model-variant", "base",
+        "--dataset-dir", ckpt_dir, "--output-dir", ckpt_dir, *extra,
+    ]
+    return build_root_parser().parse_args(argv)
+
+
+_FAKE_GPU = SimpleNamespace(device="cpu", precision="fp32")
+
+
+@patch("acestep.training_v2.cli.config_builder.detect_gpu", return_value=_FAKE_GPU)
+class TimestepMuOverrideTest(unittest.TestCase):
+    """The model config value is used unless --timestep-mu is given."""
+
+    def setUp(self) -> None:
+        self.tmp = tempfile.TemporaryDirectory()
+        model_dir = Path(self.tmp.name) / "acestep-v15-base"
+        model_dir.mkdir()
+        (model_dir / "config.json").write_text(
+            json.dumps({"timestep_mu": -0.4, "timestep_sigma": 1.0})
+        )
+
+    def tearDown(self) -> None:
+        self.tmp.cleanup()
+
+    def test_defaults_to_model_config(self, _gpu) -> None:
+        _, cfg = build_configs(_parse(self.tmp.name))
+        self.assertEqual(cfg.timestep_mu, -0.4)
+
+    def test_override_replaces_model_config(self, _gpu) -> None:
+        _, cfg = build_configs(_parse(self.tmp.name, "--timestep-mu", "-0.8"))
+        self.assertEqual(cfg.timestep_mu, -0.8)
+        self.assertEqual(cfg.timestep_sigma, 1.0)
+
+    def test_vanilla_subcommand_unaffected(self, _gpu) -> None:
+        argv = [
+            "vanilla", "--checkpoint-dir", self.tmp.name, "--model-variant", "base",
+            "--dataset-dir", self.tmp.name, "--output-dir", self.tmp.name,
+        ]
+        _, cfg = build_configs(build_root_parser().parse_args(argv))
+        self.assertEqual(cfg.timestep_mu, -0.4)
+
+
+if __name__ == "__main__":
+    unittest.main()

acestep/training_v2/cli/variant_dir_map_test.py

@@ -0,0 +1,56 @@
+"""Tests that CLI variant aliases (including XL) resolve to checkpoint folders."""
+
+import json
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+
+from acestep.training_v2.cli.args import VARIANT_DIR_MAP
+from acestep.training_v2.cli.config_builder import _resolve_model_config_path
+from acestep.training_v2.cli.validation import validate_paths
+from acestep.training_v2.model_loader import _VARIANT_DIR
+
+
+class VariantDirMapTest(unittest.TestCase):
+    """--model-variant xl_base must find acestep-v15-xl-base like base does."""
+
+    def setUp(self) -> None:
+        self.tmp = tempfile.TemporaryDirectory()
+        self.root = Path(self.tmp.name)
+        for name in ("acestep-v15-base", "acestep-v15-xl-base"):
+            (self.root / name).mkdir()
+            (self.root / name / "config.json").write_text(json.dumps({"timestep_mu": -0.4}))
+
+    def tearDown(self) -> None:
+        self.tmp.cleanup()
+
+    def _validate(self, variant: str) -> SimpleNamespace:
+        args = SimpleNamespace(
+            checkpoint_dir=str(self.root), model_variant=variant,
+            dataset_dir=None, resume_from=None,
+        )
+        self.assertTrue(validate_paths(args), variant)
+        return args
+
+    def test_xl_base_alias_resolves(self) -> None:
+        args = self._validate("xl_base")
+        self.assertEqual(args.model_dir, self.root / "acestep-v15-xl-base")
+        self.assertEqual(
+            _resolve_model_config_path(self.root, "xl_base"),
+            self.root / "acestep-v15-xl-base" / "config.json",
+        )
+
+    def test_base_alias_unchanged(self) -> None:
+        self.assertEqual(self._validate("base").model_dir, self.root / "acestep-v15-base")
+
+    def test_literal_folder_name_still_works(self) -> None:
+        args = self._validate("acestep-v15-xl-base")
+        self.assertEqual(args.model_dir, self.root / "acestep-v15-xl-base")
+
+    def test_cli_map_matches_model_loader(self) -> None:
+        self.assertEqual(VARIANT_DIR_MAP, _VARIANT_DIR)
+
+
+if __name__ == "__main__":
+    unittest.main()

acestep/training_v2/trainer_fixed.py

@@ -29,6 +29,7 @@
 from acestep.training.data_module import PreprocessedDataModule
 
 # V2 modules
+from acestep.training_v2.adapter_resume import AdapterResumeError
 from acestep.training_v2.configs import TrainingConfigV2
 from acestep.training_v2.tensorboard_utils import TrainingLogger
 from acestep.training_v2.ui import TrainingUpdate
@@ -390,6 +391,10 @@ def _train_fabric(
                 )
                 if resumed is not None:
                     start_epoch, global_step = resumed
+            except AdapterResumeError:
+                # Optimizer/scheduler state may already be loaded: never
+                # silently continue with an untrained adapter.
+                raise
             except Exception as exc:
                 logger.exception("Failed to load checkpoint")
                 if is_main:

acestep/training_v2/trainer_helpers.py

@@ -24,6 +24,7 @@
     save_lokr_weights,
     load_lokr_weights,
 )
+from acestep.training_v2.adapter_resume import load_adapter_state_dict
 from acestep.training_v2.ui import TrainingUpdate
 
 logger = logging.getLogger(__name__)
@@ -395,11 +396,14 @@ def resume_checkpoint(
             decoder = module.model.decoder
             if hasattr(decoder, "_forward_module"):
                 decoder = decoder._forward_module
-            decoder.load_state_dict(state_dict, strict=False)
+            n_loaded = load_adapter_state_dict(decoder, state_dict)
 
             start_epoch = ckpt_info["epoch"]
             g_step = ckpt_info["global_step"]
-            parts = [f"[OK] Resumed from epoch {start_epoch}, step {g_step}"]
+            parts = [
+                f"[OK] Resumed from epoch {start_epoch}, step {g_step}",
+                f"adapter OK ({n_loaded} tensors)",
+            ]
             if ckpt_info["loaded_optimizer"]:
                 parts.append("optimizer OK")
             if ckpt_info["loaded_scheduler"]: