feat: add Mano manifold optimizer (mano_adamw)

- implement Mano (v2) with axis-rotating tangent projection and manifold normalization, replacing Newton-Schulz iteration
- composite ManoAdamW reuses partition_optimizer_parameters and composite helpers
- register mano_adamw in OptimizerFactory, export Mano and ManoAdamW
- add --mano_momentum and --mano_nesterov CLI options in Optimizer group
- add mano_adamw hyperparameters branch in train.py
- document mano_adamw in params.md
- add tests for single-step projection, axis alternation, factory registration, closure, and resume
This commit is contained in:
2026-08-01 08:51:08 +08:00
parent 6c76c16480
commit 6db276f37a
6 changed files with 367 additions and 2 deletions
+23
View File
@@ -219,6 +219,19 @@ _START_METHODS = ["spawn", "fork", "forkserver"]
group="Optimizer",
help="Muon LR adjustment strategy.",
)
@opt(
"--mano_momentum",
type=float,
default=0.95,
group="Optimizer",
help="Mano momentum factor.",
)
@opt(
"--mano_nesterov/--no-mano_nesterov",
default=True,
group="Optimizer",
help="Mano Nesterov momentum.",
)
@opt(
"--random_seed",
type=int,
@@ -674,6 +687,8 @@ def train(
"nesterov": kwargs.pop("muon_nesterov", True),
"ns_steps": kwargs.pop("muon_ns_steps", 5),
"adjust_lr_fn": kwargs.pop("muon_adjust_lr", "match_rms_adamw"),
"mano_momentum": kwargs.pop("mano_momentum", 0.95),
"mano_nesterov": kwargs.pop("mano_nesterov", True),
}
optimizer_fn = partial(
create_optimizer,
@@ -695,6 +710,14 @@ def train(
optimizer_hyperparameters.update(
{"nadamw_betas": [0.9, 0.999], "nadamw_eps": 1e-8, "nora_eps": 1e-10}
)
elif optimizer_name == "mano_adamw":
optimizer_hyperparameters = {
key: optimizer_kwargs[key]
for key in ("lr", "weight_decay", "mano_momentum", "mano_nesterov")
}
optimizer_hyperparameters.update(
{"adamw_betas": [0.9, 0.95], "adamw_eps": 1e-8}
)
else:
optimizer_hyperparameters = {
key: optimizer_kwargs[key]