feat: add SCAO (Sparse Curvature-Aware Adaptive Optimizer) support#3624
feat: add SCAO (Sparse Curvature-Aware Adaptive Optimizer) support#3624whispering3 wants to merge 13 commits into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds support for the SCAO custom optimizer by extending the Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/axolotl/utils/schemas/enums.py (1)
98-98: Nit: trailing whitespace.There's a trailing space after
"scao"on this line. Most pre-commit/linter configs (e.g.trailing-whitespace) will flag or auto-fix this.✂️ Proposed fix
- scao = "scao" + scao = "scao"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/axolotl/utils/schemas/enums.py` at line 98, The line defining the enum entry scao has a trailing whitespace after the string literal (scao = "scao" ), remove the trailing space so the assignment is exactly scao = "scao" to satisfy linters; update the enum entry in enums.py where the symbol scao is defined to remove the extra space character.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/axolotl/core/builders/base.py`:
- Around line 384-395: The ImportError raised when importing SCAO should
suppress the internal traceback by chaining it with from None; in the block
where you check self.cfg.optimizer == "scao" and try to from scao.optimizer
import SCAO, change the re-raised ImportError to use "raise ImportError('SCAO
optimizer not found. Please install it with \'pip install scao\'') from None" so
the user-facing install hint replaces the internal traceback.
---
Nitpick comments:
In `@src/axolotl/utils/schemas/enums.py`:
- Line 98: The line defining the enum entry scao has a trailing whitespace after
the string literal (scao = "scao" ), remove the trailing space so the assignment
is exactly scao = "scao" to satisfy linters; update the enum entry in enums.py
where the symbol scao is defined to remove the extra space character.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 1138560b-1f2c-495c-903f-aa69882c7f8b
📒 Files selected for processing (2)
src/axolotl/core/builders/base.pysrc/axolotl/utils/schemas/enums.py
| elif self.cfg.optimizer == "scao": | ||
| try: | ||
| from scao.optimizer import SCAO | ||
| except ImportError: | ||
| raise ImportError( | ||
| "SCAO optimizer not found. Please install it with 'pip install scao'" | ||
| ) | ||
| optimizer_cls = SCAO | ||
| beta1 = training_args_kwargs.get("adam_beta1", 0.9) | ||
| beta2 = training_args_kwargs.get("adam_beta2", 0.999) | ||
| adam_kwargs["betas"] = (beta1, beta2) | ||
| optimizer_kwargs.update(adam_kwargs) |
There was a problem hiding this comment.
Chain the re-raised ImportError (Ruff B904).
Ruff flags raise ... from err/from None inside an except block to avoid obscuring the original traceback. Since the intent here is to replace the internal traceback with a user-facing install hint, from None is appropriate and matches the style used elsewhere for optional-dependency imports.
🛠️ Proposed fix
elif self.cfg.optimizer == "scao":
try:
from scao.optimizer import SCAO
- except ImportError:
+ except ImportError as err:
raise ImportError(
"SCAO optimizer not found. Please install it with 'pip install scao'"
- )
+ ) from err
optimizer_cls = SCAO
beta1 = training_args_kwargs.get("adam_beta1", 0.9)
beta2 = training_args_kwargs.get("adam_beta2", 0.999)
adam_kwargs["betas"] = (beta1, beta2)
optimizer_kwargs.update(adam_kwargs)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| elif self.cfg.optimizer == "scao": | |
| try: | |
| from scao.optimizer import SCAO | |
| except ImportError: | |
| raise ImportError( | |
| "SCAO optimizer not found. Please install it with 'pip install scao'" | |
| ) | |
| optimizer_cls = SCAO | |
| beta1 = training_args_kwargs.get("adam_beta1", 0.9) | |
| beta2 = training_args_kwargs.get("adam_beta2", 0.999) | |
| adam_kwargs["betas"] = (beta1, beta2) | |
| optimizer_kwargs.update(adam_kwargs) | |
| elif self.cfg.optimizer == "scao": | |
| try: | |
| from scao.optimizer import SCAO | |
| except ImportError as err: | |
| raise ImportError( | |
| "SCAO optimizer not found. Please install it with 'pip install scao'" | |
| ) from err | |
| optimizer_cls = SCAO | |
| beta1 = training_args_kwargs.get("adam_beta1", 0.9) | |
| beta2 = training_args_kwargs.get("adam_beta2", 0.999) | |
| adam_kwargs["betas"] = (beta1, beta2) | |
| optimizer_kwargs.update(adam_kwargs) |
🧰 Tools
🪛 Ruff (0.15.11)
[warning] 388-390: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/axolotl/core/builders/base.py` around lines 384 - 395, The ImportError
raised when importing SCAO should suppress the internal traceback by chaining it
with from None; in the block where you check self.cfg.optimizer == "scao" and
try to from scao.optimizer import SCAO, change the re-raised ImportError to use
"raise ImportError('SCAO optimizer not found. Please install it with \'pip
install scao\'') from None" so the user-facing install hint replaces the
internal traceback.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
|
Hi! I just updated the branch to sync with main, and all CI checks are green. Let me know if you need any adjustments to the code or documentation before review. Thanks! |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Added E2E coverage for the SCAO optimizer in tests/e2e/test_optimizers.py (test_scao), following the existing optimizer test pattern. CI is currently waiting for maintainer approval to run workflows. |
# Conflicts: # pyproject.toml # src/axolotl/utils/schemas/enums.py
2 participants
Description
Following the discussion in Issue #3616, this PR adds native support for SCAO (Sparse Curvature-Aware Adaptive Optimizer).
As requested, the optimizer math itself is fully isolated in an external, pip-installable package (
scao). This PR simply adds the routing logic in the Trainer Builder to allow users to setoptimizer: scaoin their.ymlconfig files.You can check the PyPI package here: https://pypi.org/project/scao/
Benchmark / Rationale
SCAO is designed to bring 2nd-order convergence to QLoRA setups without the typical memory or speed penalties. Because QLoRA is heavily memory-bound, SCAO computes curvature essentially for free.
T4 GPU / Qwen-2.5-3B (First 100 Steps):
Changes
"scao"handling in_configure_custom_optimizer.try/exceptto guide users to runpip install scao.Testing
Tested locally with a Qwen-2.5-3B config. No OOM issues, parameters inject correctly.
Summary by CodeRabbit
scaooptimizer. Users can now configure and use this optimizer option in their training configurations.