Add optional Axolotl MoRA/ReMoRA integration#3647
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:
📝 WalkthroughWalkthroughThis pull request adds MoRA (Mixture of Rank Adapters) and ReMoRA (Restart MoRA) support to Axolotl through a plugin-driven architecture. The changes extend the plugin system with adapter capability contracts, introduce MoRA configuration models and a plugin implementation, refactor adapter loading to centralize PEFT config construction, update schema validation to support plugin-registered adapters, and integrate MoRA into DPO/RL and ReLoRA training pipelines. ChangesMoRA/ReMoRA Adapter Plugin Implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
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)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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: 4
🧹 Nitpick comments (4)
src/axolotl/utils/schemas/validation.py (1)
1482-1483: 💤 Low valueRedundant
mora_relora_stepguard — this check is unreachable
MoraConfig.validate_relora(in bothargs.pyandpeft.py) already raisesValueError("mora.use_relora requires mora.use_relora_step")when validating the nested model. In Pydantic v2, nested model validators run before the parent'smodel_validator(mode="after"), so ifuse_relora=Trueanduse_relora_step=None, the exception is thrown beforecheck_relorais ever entered. The duplicate raise at line 1482-1483 is dead code.♻️ Proposed simplification
if mora_use_relora: - if not mora_relora_step: - raise ValueError("mora.use_relora requires mora.use_relora_step") self.relora = True if not self.jagged_restart_steps:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/utils/schemas/validation.py` around lines 1482 - 1483, Remove the redundant guard that raises ValueError("mora.use_relora requires mora.use_relora_step") inside the parent validator (check_relora) because MoraConfig.validate_relora (in args.py and peft.py) already performs this validation for the nested model and Pydantic v2 runs nested validators before the parent's model_validator(mode="after"); locate the check in the check_relora (or parent model_validator) and delete the unreachable if-not mora_relora_step block so the error only originates from MoraConfig.validate_relora.tests/utils/schemas/validation/test_mora.py (1)
31-51: ⚡ Quick winMissing test for conflicting
mora.use_relora_stepandjagged_restart_steps
check_relora(validation.py line 1487-1490) raises when both fields are set but disagree. This new branch has no test coverage. Consider adding:def test_remora_step_conflicts_with_jagged_restart_steps(self, min_base_cfg): cfg = min_base_cfg | DictDefault({ "adapter": "mora", "plugins": ["axolotl.integrations.mora.MoraPlugin"], "jagged_restart_steps": 1000, "mora": { "use_mora": True, "mora_type": 6, "use_relora": True, "use_relora_step": 2000, # conflicts with jagged_restart_steps=1000 }, }) prepare_plugins(cfg) with pytest.raises(ValueError, match="mora.use_relora_step must match jagged_restart_steps"): validate_config(cfg)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/utils/schemas/validation/test_mora.py` around lines 31 - 51, Add a new test function (e.g., test_remora_step_conflicts_with_jagged_restart_steps) that builds a config from min_base_cfg | DictDefault where "jagged_restart_steps" is set to 1000 while mora.use_relora_step is set to 2000 (and other mora fields same as existing test), call prepare_plugins(cfg) and then assert that validate_config(cfg) raises a ValueError by using pytest.raises(ValueError, match="mora.use_relora_step must match jagged_restart_steps"); this will exercise the check_relora branch in validation.py that errors when mora.use_relora_step conflicts with jagged_restart_steps.src/axolotl/utils/schemas/peft.py (1)
266-302: ⚡ Quick winDuplicate
MoraConfig— this block is already defined insrc/axolotl/integrations/mora/args.pyThis is the mirror of the issue flagged in
args.py. The entireMoraConfigclass here (fields, constraints, andvalidate_relora) is identical to the one inargs.py. Resolve by importing instead:♻️ Proposed fix
+from axolotl.integrations.mora.args import MoraConfig # canonical definition -class MoraConfig(BaseModel): - """MoRA / ReMoRA configuration subset.""" - - use_mora: bool = Field(...) - mora_type: int = Field(...) - use_relora: bool = Field(...) - use_relora_step: int | None = Field(...) - - `@model_validator`(mode="after") - def validate_relora(self): ...🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/utils/schemas/peft.py` around lines 266 - 302, Remove the duplicate MoraConfig class definition and instead import the canonical MoraConfig from the existing integrations module; specifically delete the MoraConfig class and replace it with an import of MoraConfig from integrations.mora.args, ensuring any local references still refer to MoraConfig and that validate_relora remains the same via the imported class.src/axolotl/integrations/mora/args.py (1)
8-48: ⚡ Quick win
MoraConfigis duplicated verbatim insrc/axolotl/utils/schemas/peft.py
src/axolotl/utils/schemas/peft.py(lines 266–302) defines an independentMoraConfigwith the same four fields (use_mora,mora_type,use_relora,use_relora_step) and an identicalvalidate_reloramodel-validator. Both copies will silently diverge the moment one is updated without touching the other.
peft.pyshould import from this canonical definition rather than re-declare it:♻️ Proposed fix in
src/axolotl/utils/schemas/peft.py+from axolotl.integrations.mora.args import MoraConfig # re-use canonical definition -class MoraConfig(BaseModel): - ... # remove the duplicate block entirely🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/integrations/mora/args.py` around lines 8 - 48, Remove the duplicate MoraConfig declaration in peft.py and instead import the canonical MoraConfig class (which contains the validate_relora method) from the module that defines it; replace the local class definition with a single import, update any local references to use the imported MoraConfig, and ensure any module exports or type references still expose MoraConfig as before.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/axolotl/loaders/adapter.py`:
- Around line 429-462: The load_mora path currently calls get_peft_model(model,
lora_config, ...) without the ParamWrapper dropout safeguard; add the same check
used in load_lora: detect if PEFT will auto-convert to target_parameters and
lora_config.lora_dropout (or cfg.lora_dropout) != 0, and if so set
lora_config.lora_dropout = 0 (or otherwise disable dropout) before calling
get_peft_model to avoid ParamWrapper runtime failures on fused-MoE models; make
the change around the branch where model = get_peft_model(model, lora_config,
**model_kwargs) and reference get_peft_model, lora_config, cfg.lora_dropout, and
ParamWrapper/target_parameters in the implementation.
In `@src/axolotl/utils/schemas/validation.py`:
- Around line 1492-1495: Run the ruff formatter on the modified file to fix CI
failures: format src/axolotl/utils/schemas/validation.py (the block around the
ReLoRA checks — the if not self.jagged_restart_steps and adapter membership
check) by running `ruff format src/axolotl/utils/schemas/validation.py` or
applying the same ruff-format changes so the file matches the project's style
rules and the CI stops reporting reformatting diffs.
In `@tests/integrations/test_mora.py`:
- Around line 81-84: The test file fails ruff-format due to long lines in the
block that sets up the fake PEFT objects; run the formatter (e.g., "ruff format
tests/integrations/test_mora.py") or manually break the long lines so they
conform to line-length rules for the three statements that set monkeypatch
attributes for adapter_module._peft_supports_mora, adapter_module.LoraConfig,
and adapter_module.get_peft_model (and the fake_model assignment); ensure each
monkeypatch.setattr call is on its own wrapped line and re-run ruff to verify
formatting passes.
- Around line 47-57: The test test_load_mora_raises_when_peft_missing_support is
relying on the runtime PEFT implementation; mock the internal guard
_peft_supports_mora to return False so the test is deterministic: in the test
before calling load_mora(model, cfg, config_only=True) patch or monkeypatch the
module-level function _peft_supports_mora to a lambda or stub that returns False
(reference _peft_supports_mora and load_mora to locate the code) so the
ImportError branch is exercised regardless of the environment.
---
Nitpick comments:
In `@src/axolotl/integrations/mora/args.py`:
- Around line 8-48: Remove the duplicate MoraConfig declaration in peft.py and
instead import the canonical MoraConfig class (which contains the
validate_relora method) from the module that defines it; replace the local class
definition with a single import, update any local references to use the imported
MoraConfig, and ensure any module exports or type references still expose
MoraConfig as before.
In `@src/axolotl/utils/schemas/peft.py`:
- Around line 266-302: Remove the duplicate MoraConfig class definition and
instead import the canonical MoraConfig from the existing integrations module;
specifically delete the MoraConfig class and replace it with an import of
MoraConfig from integrations.mora.args, ensuring any local references still
refer to MoraConfig and that validate_relora remains the same via the imported
class.
In `@src/axolotl/utils/schemas/validation.py`:
- Around line 1482-1483: Remove the redundant guard that raises
ValueError("mora.use_relora requires mora.use_relora_step") inside the parent
validator (check_relora) because MoraConfig.validate_relora (in args.py and
peft.py) already performs this validation for the nested model and Pydantic v2
runs nested validators before the parent's model_validator(mode="after"); locate
the check in the check_relora (or parent model_validator) and delete the
unreachable if-not mora_relora_step block so the error only originates from
MoraConfig.validate_relora.
In `@tests/utils/schemas/validation/test_mora.py`:
- Around line 31-51: Add a new test function (e.g.,
test_remora_step_conflicts_with_jagged_restart_steps) that builds a config from
min_base_cfg | DictDefault where "jagged_restart_steps" is set to 1000 while
mora.use_relora_step is set to 2000 (and other mora fields same as existing
test), call prepare_plugins(cfg) and then assert that validate_config(cfg)
raises a ValueError by using pytest.raises(ValueError,
match="mora.use_relora_step must match jagged_restart_steps"); this will
exercise the check_relora branch in validation.py that errors when
mora.use_relora_step conflicts with jagged_restart_steps.
🪄 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: 9f85ef28-aa77-4b32-9443-32535f20e661
📒 Files selected for processing (10)
src/axolotl/integrations/mora/__init__.pysrc/axolotl/integrations/mora/args.pysrc/axolotl/integrations/mora/plugin.pysrc/axolotl/loaders/adapter.pysrc/axolotl/loaders/model.pysrc/axolotl/train.pysrc/axolotl/utils/schemas/peft.pysrc/axolotl/utils/schemas/validation.pytests/integrations/test_mora.pytests/utils/schemas/validation/test_mora.py
| _patch_peft_clippable_linear() | ||
| lora_config = _build_peft_lora_config(model, cfg, adapter_kind="mora") | ||
|
|
||
| if config_only: | ||
| return None, lora_config | ||
|
|
||
| rank = int(os.environ.get("LOCAL_RANK", 0)) | ||
|
|
||
| if ( | ||
| cfg.fsdp_config | ||
| and cfg.adapter | ||
| and cfg.fsdp_config.cpu_ram_efficient_loading | ||
| and rank != 0 | ||
| ): | ||
| setup_quantized_meta_for_peft(model) | ||
|
|
||
| model_kwargs: Any = {} | ||
| if cfg.peft_autocast_adapter_dtype is not None: | ||
| model_kwargs["autocast_adapter_dtype"] = cfg.peft_autocast_adapter_dtype | ||
|
|
||
| if cfg.lora_model_dir: | ||
| LOG.debug("Loading pretrained PEFT - MoRA") | ||
| if cfg.lora_on_cpu: | ||
| model_kwargs["max_memory"] = {"cpu": "256GiB"} | ||
| model_kwargs["device_map"] = {"": "cpu"} | ||
| model = PeftModel.from_pretrained( | ||
| model, | ||
| cfg.lora_model_dir, | ||
| is_trainable=(not inference), | ||
| **model_kwargs, | ||
| ) | ||
| else: | ||
| model = get_peft_model(model, lora_config, **model_kwargs) | ||
|
|
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify parity gap between load_lora and load_mora around ParamWrapper/dropout guard.
rg -n -C4 'def load_lora|def load_mora|_peft_will_auto_convert_target_params|_patch_peft_param_wrapper_dropout' src/axolotl/loaders/adapter.pyRepository: axolotl-ai-cloud/axolotl
Length of output: 2003
🏁 Script executed:
sed -n '423,510p' src/axolotl/loaders/adapter.pyRepository: axolotl-ai-cloud/axolotl
Length of output: 2978
🏁 Script executed:
rg -i 'mora.*dropout|mora.*param.*wrapper|mora.*target.*param|mora.*fused' src/axolotl/loaders/adapter.pyRepository: axolotl-ai-cloud/axolotl
Length of output: 50
🏁 Script executed:
rg -n 'def _build_peft_lora_config' src/axolotl/loaders/adapter.py -A 50Repository: axolotl-ai-cloud/axolotl
Length of output: 2392
🏁 Script executed:
rg -n 'def _build_mora_config_kwargs|def _build_lora_config_kwargs' src/axolotl/loaders/adapter.py -A 20Repository: axolotl-ai-cloud/axolotl
Length of output: 2078
🏁 Script executed:
rg -n '_peft_will_auto_convert_target_params' src/axolotl/loaders/adapter.py -B 2 -A 25Repository: axolotl-ai-cloud/axolotl
Length of output: 2405
Mirror the ParamWrapper dropout safeguard in load_mora.
load_lora protects get_peft_model(...) when PEFT auto-converts to target_parameters and lora_dropout != 0, but load_mora skips that check. This can reintroduce the same runtime failure on fused-MoE models under MoRA.
🔧 Proposed parity fix
def load_mora(
model: PreTrainedModel,
cfg: DictDefault,
inference: bool = False,
config_only: bool = False,
) -> tuple[PreTrainedModel | PeftModel | PeftMixedModel | None, PeftConfig | None]:
_patch_peft_clippable_linear()
lora_config = _build_peft_lora_config(model, cfg, adapter_kind="mora")
if config_only:
return None, lora_config
+
+ if getattr(
+ lora_config, "lora_dropout", 0
+ ) and _peft_will_auto_convert_target_params(model, lora_config):
+ LOG.warning(
+ "lora_dropout=%s requested but PEFT will wrap this model's fused "
+ "MoE expert parameters with ParamWrapper, which cannot apply "
+ "dropout (the 3D einsum can't factor dropout out of "
+ "lora_B(lora_A(dropout(x)))). Dropout will still be applied to "
+ "non-expert LoRA layers (e.g. attention), and expert LoRA layers "
+ "will use nn.Identity for the dropout slot.",
+ lora_config.lora_dropout,
+ )
+ _patch_peft_param_wrapper_dropout()
rank = int(os.environ.get("LOCAL_RANK", 0))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/axolotl/loaders/adapter.py` around lines 429 - 462, The load_mora path
currently calls get_peft_model(model, lora_config, ...) without the ParamWrapper
dropout safeguard; add the same check used in load_lora: detect if PEFT will
auto-convert to target_parameters and lora_config.lora_dropout (or
cfg.lora_dropout) != 0, and if so set lora_config.lora_dropout = 0 (or otherwise
disable dropout) before calling get_peft_model to avoid ParamWrapper runtime
failures on fused-MoE models; make the change around the branch where model =
get_peft_model(model, lora_config, **model_kwargs) and reference get_peft_model,
lora_config, cfg.lora_dropout, and ParamWrapper/target_parameters in the
implementation.
| def test_load_mora_raises_when_peft_missing_support(self): | ||
| model = torch.nn.Linear(4, 4) | ||
| cfg = DictDefault( | ||
| { | ||
| "adapter": "mora", | ||
| "mora": {"use_mora": True, "mora_type": 6}, | ||
| } | ||
| ) | ||
|
|
||
| with pytest.raises(ImportError, match="MoRA support"): | ||
| load_mora(model, cfg, config_only=True) |
There was a problem hiding this comment.
Test relies on the environment having vanilla PEFT — will break on a MoRA-capable build
test_load_mora_raises_when_peft_missing_support calls load_mora without patching _peft_supports_mora. It implicitly relies on the installed PEFT (0.19.1) lacking use_mora on LoraConfig. If the test is ever run in an environment where a MoRA fork of PEFT is installed, _peft_supports_mora() returns True, the ImportError is never raised, and the test silently passes/fails for the wrong reason.
Explicitly mock the guard to make the test deterministic:
🔒 Proposed fix
- def test_load_mora_raises_when_peft_missing_support(self):
+ def test_load_mora_raises_when_peft_missing_support(self, monkeypatch):
model = torch.nn.Linear(4, 4)
cfg = DictDefault(
{
"adapter": "mora",
"mora": {"use_mora": True, "mora_type": 6},
}
)
+ monkeypatch.setattr(adapter_module, "_peft_supports_mora", lambda: False)
with pytest.raises(ImportError, match="MoRA support"):
load_mora(model, cfg, config_only=True)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/integrations/test_mora.py` around lines 47 - 57, The test
test_load_mora_raises_when_peft_missing_support is relying on the runtime PEFT
implementation; mock the internal guard _peft_supports_mora to return False so
the test is deterministic: in the test before calling load_mora(model, cfg,
config_only=True) patch or monkeypatch the module-level function
_peft_supports_mora to a lambda or stub that returns False (reference
_peft_supports_mora and load_mora to locate the code) so the ImportError branch
is exercised regardless of the environment.
| fake_model = SimpleNamespace(print_trainable_parameters=Mock()) | ||
| monkeypatch.setattr(adapter_module, "_peft_supports_mora", lambda: True) | ||
| monkeypatch.setattr(adapter_module, "LoraConfig", FakeLoraConfig) | ||
| monkeypatch.setattr(adapter_module, "get_peft_model", Mock(return_value=fake_model)) |
There was a problem hiding this comment.
Fix ruff-format failures — CI is blocked
ruff-format reports this file needs reformatting (lines 81-83 are likely too long). Run ruff format tests/integrations/test_mora.py locally.
🧰 Tools
🪛 GitHub Actions: lint / 0_pre-commit.txt
[error] 81-83: pre-commit hook 'ruff-format' failed because files were modified (2 files reformatted).
🪛 GitHub Actions: lint / pre-commit
[error] 81-83: pre-commit hook ruff-format failed; 2 files were reformatted by this hook (formatting needed).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/integrations/test_mora.py` around lines 81 - 84, The test file fails
ruff-format due to long lines in the block that sets up the fake PEFT objects;
run the formatter (e.g., "ruff format tests/integrations/test_mora.py") or
manually break the long lines so they conform to line-length rules for the three
statements that set monkeypatch attributes for
adapter_module._peft_supports_mora, adapter_module.LoraConfig, and
adapter_module.get_peft_model (and the fake_model assignment); ensure each
monkeypatch.setattr call is on its own wrapped line and re-run ruff to verify
formatting passes.
|
📖 Documentation Preview: https://69fffb277bef4de37b0c8bac--resonant-treacle-0fd729.netlify.app Deployed on Netlify from commit 1f95e8a |
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Co-authored-by: Axolotl Swarm <no-reply@axolotl.ai>
Co-authored-by: Axolotl Swarm <no-reply@axolotl.ai>
winglian
force-pushed
the
features/mora-remora-integration
branch
from
c4bce7d to
cd32c83
Compare
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
src/axolotl/loaders/adapter.py (1)
388-390:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAvoid mutating the model in
config_onlymode.Line 388-Line 390 calls
enable_input_require_grads()even whenconfig_only=True. That breaks the expected “config-only” no-side-effects behavior.Suggested fix
- if hasattr(model, "enable_input_require_grads"): + if not config_only and hasattr(model, "enable_input_require_grads"): model.enable_input_require_grads()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/loaders/adapter.py` around lines 388 - 390, The code currently calls model.enable_input_require_grads() even when config_only=True, causing side effects; update the check so the method is only invoked when not in config-only mode (e.g., wrap the call with if not config_only: model.enable_input_require_grads()). Ensure this guard is applied alongside the existing adapter logic (the lines handling adapter in ["lora", "qlora"]) so config-only paths do not mutate the model.
🧹 Nitpick comments (7)
src/axolotl/integrations/base.py (2)
118-128: 💤 Low valueConfirm intent of
(None, PeftConfig)return shape.The signature allows the model element of the returned tuple to be
None(tuple[PreTrainedModel | PeftModel | PeftMixedModel | None, PeftConfig | None] | None), which lets a plugin signal "config-only success" by returning(None, cfg). That distinction (vs. plainNonemeaning "skip plugin") is subtle and worth a one-line note in the docstring so future plugin authors don't conflate the two.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/integrations/base.py` around lines 118 - 128, The load_adapter return type can be either None (plugin skipped) or a tuple where the first element may be None and the second a PeftConfig to signal a "config-only" success; update the load_adapter docstring to explicitly state that returning None means "no plugin applied/skip", while returning (None, cfg) means the plugin succeeded but only provided configuration (no model loaded), and include the exact tuple shape (None, PeftConfig) and brief guidance for plugin authors to return the correct form when they only want to supply config.
463-482: 💤 Low valueRepeated dict rebuilds in capability lookups.
adapter_capabilities()constructs a fresh dict on every invocation, andget_adapter_capability,supports_adapter, andadapter_supports_reloraeach call it again. Validators (e.g.,LoraConfig.validate_adapter,check_relora) can call these in sequence, so for each validation pass the registry is iterated multiple times. With small plugin counts this is harmless, but a tiny memoization (or a private_adapter_capabilitiesdict updated onregister) would make capability queries O(1).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/integrations/base.py` around lines 463 - 482, The adapter_capabilities() method rebuilds the dict on every call causing repeated iteration when get_adapter_capability, supports_adapter, and adapter_supports_relora call it; instead introduce a cached private dict (e.g., _adapter_capabilities) that is populated/updated when plugins are registered/unregistered (update it inside the register/unregister methods that touch self.plugins) and have adapter_capabilities() return that cached dict (or lazily build it once and invalidate on register changes); then change get_adapter_capability, supports_adapter, and adapter_supports_relora to read from the cached _adapter_capabilities to make capability lookups O(1).tests/integrations/mora/test_mora.py (1)
28-30: ⚡ Quick winUse
register()(or a teardown fixture) instead of mutatingpluginsdirectly.Assigning into
PluginManager.get_instance().plugins[...]bypassesregister()and persists across tests becausePluginManageris a process-wide singleton. Suggest either callingPluginManager.get_instance().register("axolotl.integrations.mora.MoraPlugin")or, better, an autouse fixture that snapshotspluginsbefore the test and restores it after — this avoids cross-test order dependencies with other suites that assume a clean registry.♻️ Example fixture-based cleanup
`@pytest.fixture`(autouse=True) def _reset_plugin_registry(): pm = PluginManager.get_instance() snapshot = dict(pm.plugins) try: yield finally: pm.plugins.clear() pm.plugins.update(snapshot)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integrations/mora/test_mora.py` around lines 28 - 30, Replace direct mutation of PluginManager.get_instance().plugins in the test with a proper registration call or add an autouse teardown fixture: call PluginManager.get_instance().register("axolotl.integrations.mora.MoraPlugin") (or instantiate mora_plugin.MoraPlugin() and pass it to register) instead of assigning into plugins, or add a pytest fixture (e.g., _reset_plugin_registry) that snapshots pm.plugins (pm = PluginManager.get_instance()), yields, then restores the snapshot in the finally block to ensure the singleton registry is not mutated across tests.src/axolotl/utils/config/__init__.py (2)
331-384: ⚡ Quick winDocument the prepare_plugins → validate_config ordering.
plugin_manager.normalize_config_input(cfg)andplugin_manager.validate_config(validated_cfg)are no-ops unless plugins were already registered viaprepare_plugins(cfg). The schema-side checks added in this PR (LoraConfig.validate_adapterandcheck_relora) also depend on that registration. A caller that invokesvalidate_configwithout first callingprepare_pluginswill silently pass through plugin-only adapters and only fail later. A short docstring note (or an explicitprepare_plugins(cfg)call here, idempotent on the singleton) would prevent that footgun.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/utils/config/__init__.py` around lines 331 - 384, The plugin validation depends on plugins being registered via prepare_plugins(cfg) so calling plugin_manager.normalize_config_input(cfg) and plugin_manager.validate_config(validated_cfg) may be no-ops if prepare_plugins wasn't run; fix by invoking prepare_plugins(cfg) (idempotent on the PluginManager singleton) before normalize_config_input and before validate_config, or add a short docstring note explaining callers must call prepare_plugins(cfg) first—update references around PluginManager.get_instance(), plugin_manager.normalize_config_input, plugin_manager.validate_config, and the prepare_plugins symbol to ensure plugin registration happens before schema checks like LoraConfig.validate_adapter and check_relora.
368-384: 💤 Low valuePlugin
validate_configmutations bypass the Pydantic schema.
validated_cfgis built from a Pydantic dump and then handed to plugins, which can freely mutate it before it's returned. Any keys a plugin adds or rewrites here will not be re-validated againstAxolotlInputConfig/AxolotlConfigWCapabilities. That is presumably intentional (cheap escape hatch), but worth calling out so that plugin authors don't rely on this for fields they want type-checked. Consider re-running Pydantic validation, or restricting pluginvalidate_configto read-only assertions, to make the contract explicit.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/utils/config/__init__.py` around lines 368 - 384, The plugin_manager.validate_config call receives and may mutate validated_cfg (constructed from AxolotlInputConfig/AxolotlConfigWCapabilities.model_dump), allowing plugins to introduce fields that bypass Pydantic checking; to fix, after plugin_manager.validate_config(validated_cfg) re-run Pydantic validation by parsing the (possibly mutated) dict back into the appropriate model (e.g. AxolotlConfigWCapabilities or AxolotlInputConfig) via the model constructor or model_validate/parse_obj to enforce schema/types, and then return the re-validated model_dump (or alternatively restrict plugin_manager.validate_config to only perform read-only assertions and document that behavior); locate the logic around validated_cfg and plugin_manager.validate_config to implement the revalidation step or tighten the plugin contract.src/axolotl/utils/schemas/peft.py (1)
41-46: 💤 Low valueDocument built-in adapter values in the JSON schema description.
Loosening
adapterfromLiteral["lora", "qlora", "llama-adapter"]tostrremoves the enumeration that downstream tooling (docs, schema-aware editors) used to surface valid values. Consider naming the built-ins explicitly in the description so users still discover them without inspecting code.- "description": "If you want to use a built-in or plugin adapter, or leave blank to train all parameters in original model" + "description": "If you want to use a built-in adapter (lora, qlora, llama-adapter) or a plugin-registered adapter, or leave blank to train all parameters in the original model"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/utils/schemas/peft.py` around lines 41 - 46, The json_schema for the adapter Field lost the enumerated built-in values when you relaxed the type; update the json_schema_extra description on adapter (the adapter: str | None = Field(...) declaration) to explicitly list the supported built-in adapter names (for example: "lora", "qlora", "llama-adapter") and clarify that leaving it blank trains all parameters, so downstream tooling and schema-aware editors can surface valid choices without needing to inspect code.src/axolotl/loaders/adapter.py (1)
176-176: ⚡ Quick winPreserve deterministic
target_modulesordering.Using
list(set(...))at Line 176 makes ordering non-deterministic across runs, which can hurt reproducibility/debugging of emitted PEFT configs.Suggested fix
- lora_target_modules = list(set(lora_target_modules_as_list + linear_names)) + lora_target_modules = list( + dict.fromkeys(lora_target_modules_as_list + linear_names) + )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/loaders/adapter.py` at line 176, The assignment to lora_target_modules uses list(set(lora_target_modules_as_list + linear_names)) which makes ordering non-deterministic; change it to produce a deterministic, order-preserving list of unique module names by concatenating lora_target_modules_as_list and linear_names and then removing duplicates while preserving first-seen order (e.g., via an ordered-unique pattern like using a seen set or dict.fromkeys on the concatenated sequence) so the resulting lora_target_modules has stable ordering across runs.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/integrations/test_adapter_plugin_registry.py`:
- Line 36: Test mutates the global PluginManager singleton by assigning
PluginManager.get_instance().plugins["fake"] = FakeAdapterPlugin() and never
restores it; change the test to isolate this mutation by saving the original
plugins mapping or the original value for the "fake" key, then set the fake
plugin, and finally restore the previous state in a teardown/try-finally (or use
a test fixture/monkeypatch to inject and auto-restore) so
PluginManager.get_instance().plugins is returned to its original contents after
the test completes.
---
Duplicate comments:
In `@src/axolotl/loaders/adapter.py`:
- Around line 388-390: The code currently calls
model.enable_input_require_grads() even when config_only=True, causing side
effects; update the check so the method is only invoked when not in config-only
mode (e.g., wrap the call with if not config_only:
model.enable_input_require_grads()). Ensure this guard is applied alongside the
existing adapter logic (the lines handling adapter in ["lora", "qlora"]) so
config-only paths do not mutate the model.
---
Nitpick comments:
In `@src/axolotl/integrations/base.py`:
- Around line 118-128: The load_adapter return type can be either None (plugin
skipped) or a tuple where the first element may be None and the second a
PeftConfig to signal a "config-only" success; update the load_adapter docstring
to explicitly state that returning None means "no plugin applied/skip", while
returning (None, cfg) means the plugin succeeded but only provided configuration
(no model loaded), and include the exact tuple shape (None, PeftConfig) and
brief guidance for plugin authors to return the correct form when they only want
to supply config.
- Around line 463-482: The adapter_capabilities() method rebuilds the dict on
every call causing repeated iteration when get_adapter_capability,
supports_adapter, and adapter_supports_relora call it; instead introduce a
cached private dict (e.g., _adapter_capabilities) that is populated/updated when
plugins are registered/unregistered (update it inside the register/unregister
methods that touch self.plugins) and have adapter_capabilities() return that
cached dict (or lazily build it once and invalidate on register changes); then
change get_adapter_capability, supports_adapter, and adapter_supports_relora to
read from the cached _adapter_capabilities to make capability lookups O(1).
In `@src/axolotl/loaders/adapter.py`:
- Line 176: The assignment to lora_target_modules uses
list(set(lora_target_modules_as_list + linear_names)) which makes ordering
non-deterministic; change it to produce a deterministic, order-preserving list
of unique module names by concatenating lora_target_modules_as_list and
linear_names and then removing duplicates while preserving first-seen order
(e.g., via an ordered-unique pattern like using a seen set or dict.fromkeys on
the concatenated sequence) so the resulting lora_target_modules has stable
ordering across runs.
In `@src/axolotl/utils/config/__init__.py`:
- Around line 331-384: The plugin validation depends on plugins being registered
via prepare_plugins(cfg) so calling plugin_manager.normalize_config_input(cfg)
and plugin_manager.validate_config(validated_cfg) may be no-ops if
prepare_plugins wasn't run; fix by invoking prepare_plugins(cfg) (idempotent on
the PluginManager singleton) before normalize_config_input and before
validate_config, or add a short docstring note explaining callers must call
prepare_plugins(cfg) first—update references around
PluginManager.get_instance(), plugin_manager.normalize_config_input,
plugin_manager.validate_config, and the prepare_plugins symbol to ensure plugin
registration happens before schema checks like LoraConfig.validate_adapter and
check_relora.
- Around line 368-384: The plugin_manager.validate_config call receives and may
mutate validated_cfg (constructed from
AxolotlInputConfig/AxolotlConfigWCapabilities.model_dump), allowing plugins to
introduce fields that bypass Pydantic checking; to fix, after
plugin_manager.validate_config(validated_cfg) re-run Pydantic validation by
parsing the (possibly mutated) dict back into the appropriate model (e.g.
AxolotlConfigWCapabilities or AxolotlInputConfig) via the model constructor or
model_validate/parse_obj to enforce schema/types, and then return the
re-validated model_dump (or alternatively restrict
plugin_manager.validate_config to only perform read-only assertions and document
that behavior); locate the logic around validated_cfg and
plugin_manager.validate_config to implement the revalidation step or tighten the
plugin contract.
In `@src/axolotl/utils/schemas/peft.py`:
- Around line 41-46: The json_schema for the adapter Field lost the enumerated
built-in values when you relaxed the type; update the json_schema_extra
description on adapter (the adapter: str | None = Field(...) declaration) to
explicitly list the supported built-in adapter names (for example: "lora",
"qlora", "llama-adapter") and clarify that leaving it blank trains all
parameters, so downstream tooling and schema-aware editors can surface valid
choices without needing to inspect code.
In `@tests/integrations/mora/test_mora.py`:
- Around line 28-30: Replace direct mutation of
PluginManager.get_instance().plugins in the test with a proper registration call
or add an autouse teardown fixture: call
PluginManager.get_instance().register("axolotl.integrations.mora.MoraPlugin")
(or instantiate mora_plugin.MoraPlugin() and pass it to register) instead of
assigning into plugins, or add a pytest fixture (e.g., _reset_plugin_registry)
that snapshots pm.plugins (pm = PluginManager.get_instance()), yields, then
restores the snapshot in the finally block to ensure the singleton registry is
not mutated across tests.
🪄 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: 3745e7d0-a7d3-400a-972a-d55043eb54ff
📒 Files selected for processing (13)
src/axolotl/integrations/base.pysrc/axolotl/integrations/mora/__init__.pysrc/axolotl/integrations/mora/args.pysrc/axolotl/integrations/mora/plugin.pysrc/axolotl/loaders/adapter.pysrc/axolotl/loaders/model.pysrc/axolotl/train.pysrc/axolotl/utils/config/__init__.pysrc/axolotl/utils/schemas/peft.pysrc/axolotl/utils/schemas/validation.pytests/integrations/mora/test_mora.pytests/integrations/test_adapter_plugin_registry.pytests/utils/schemas/validation/mora/test_mora_validation.py
✅ Files skipped from review due to trivial changes (1)
- src/axolotl/integrations/mora/init.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/axolotl/integrations/mora/args.py
| "lora_dropout": 0.0, | ||
| } | ||
| ) | ||
| PluginManager.get_instance().plugins["fake"] = FakeAdapterPlugin() |
There was a problem hiding this comment.
Isolate plugin registry mutations to avoid cross-test leakage.
Line 36 and Line 61 write into a process-wide singleton plugin map and never restore it. This can make later tests order-dependent/flaky.
Suggested fix
- def test_lora_like_plugin_adapter_contributes_peft_kwargs(self, monkeypatch):
+ def test_lora_like_plugin_adapter_contributes_peft_kwargs(self, monkeypatch):
model = torch.nn.Linear(4, 4)
@@
- PluginManager.get_instance().plugins["fake"] = FakeAdapterPlugin()
+ plugin_manager = PluginManager.get_instance()
+ monkeypatch.setitem(plugin_manager.plugins, "fake", FakeAdapterPlugin())
@@
- def test_relora_accepts_plugin_adapter_capability(self, min_base_cfg):
- PluginManager.get_instance().plugins["fake"] = FakeAdapterPlugin()
+ def test_relora_accepts_plugin_adapter_capability(self, min_base_cfg, monkeypatch):
+ plugin_manager = PluginManager.get_instance()
+ monkeypatch.setitem(plugin_manager.plugins, "fake", FakeAdapterPlugin())Also applies to: 61-61
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/integrations/test_adapter_plugin_registry.py` at line 36, Test mutates
the global PluginManager singleton by assigning
PluginManager.get_instance().plugins["fake"] = FakeAdapterPlugin() and never
restores it; change the test to isolate this mutation by saving the original
plugins mapping or the original value for the "fake" key, then set the fake
plugin, and finally restore the previous state in a teardown/try-finally (or use
a test fixture/monkeypatch to inject and auto-restore) so
PluginManager.get_instance().plugins is returned to its original contents after
the test completes.
1 participant
Summary by CodeRabbit
New Features
Refactor
Tests