Skip to content

Fix: lora kernel pre-patch applied despite post-patch not applied#2772

Merged
NanoCode012 merged 6 commits into
mainfrom
fix/lora-kernel-patch
Jun 14, 2025
Merged

Fix: lora kernel pre-patch applied despite post-patch not applied#2772
NanoCode012 merged 6 commits into
mainfrom
fix/lora-kernel-patch

Conversation

@NanoCode012

@NanoCode012 NanoCode012 commented Jun 9, 2025

Copy link
Copy Markdown
Collaborator

Description

Fix #2770

Motivation and Context

How has this been tested?

Added test and manual run before/after to ensure error appears and got fixed

Screenshots (if appropriate)

Types of changes

Social Handles (Optional)

Summary by CodeRabbit

  • Bug Fixes

    • Improved handling of LoRA dropout by ensuring kernel patching is skipped when dropout is non-zero, with a warning logged for clarity.
  • Tests

    • Added a new test to verify that self-attention kernel patching does not occur when LoRA dropout is set above zero.

@coderabbitai

coderabbitai Bot commented Jun 9, 2025

Copy link
Copy Markdown
Contributor

"""

Walkthrough

The changes introduce a precondition in the _apply_self_attention_lora_patch method to prevent patching when LoRA dropout is non-zero, add a helper function get_layers for consistent model layer retrieval, and provide a new test verifying that patching is skipped when dropout is set. Imports were updated to support the new test logic.

Changes

File(s) Change Summary
src/axolotl/monkeypatch/lora_kernels.py Added get_layers helper function; refactored apply_lora_kernel_patches to use get_layers.
src/axolotl/loaders/patch_manager.py Modified _apply_self_attention_lora_patch to skip patching if lora_dropout is non-zero and log a warning.
tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py Added test test_kernel_training_integration_dropout_non_zero to verify patching is skipped when LoRA dropout > 0; updated imports.

Sequence Diagram(s)

sequenceDiagram
    participant Test as Test Function
    participant Config as LoRA Config
    participant Patch as patch_self_attn_lora
    participant Model as Model
    participant Apply as apply_lora_kernel_patches

    Test->>Patch: Call with Config (lora_dropout > 0)
    Patch->>Patch: Check lora_dropout
    Patch-->>Test: Return without patching

    Test->>Model: Load model and tokenizer
    Test->>Apply: Call apply_lora_kernel_patches(Model)
    Apply->>Model: Retrieve layers via get_layers
    Apply->>Model: Attempt to patch (skipped if dropout > 0)
    Apply-->>Test: Model layers remain unpatched
Loading

Assessment against linked issues

Objective Addressed Explanation
Prevent AttributeError by ensuring patching is skipped when LoRA dropout is non-zero (#2770)
Add test to verify patching is not applied when LoRA dropout is non-zero (#2770)
Centralize layer retrieval logic for patching functions (#2770)

Assessment against linked issues: Out-of-scope changes

No out-of-scope changes found.

Possibly related PRs

  • Lora kernels fix #2732: Fixes and reorganizes self-attention LoRA patch application and updates model loading for PeftModelForCausalLM; closely related to patching logic and layer handling.

Suggested reviewers

  • winglian

Poem

In the code where kernels patch and play,
A rabbit hops to check the way—
If dropout’s high, no patch will land,
The layers stay just as they stand.
With tests in tow, the bunny cheers,
“No errors here, just bug-free gears!”
🐇✨
"""


📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 7f1329d and e4af596.

📒 Files selected for processing (3)
  • src/axolotl/loaders/patch_manager.py (1 hunks)
  • src/axolotl/monkeypatch/lora_kernels.py (2 hunks)
  • tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py (2 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
  • src/axolotl/loaders/patch_manager.py
  • src/axolotl/monkeypatch/lora_kernels.py
  • tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py
✨ Finishing Touches
  • 📝 Generate Docstrings

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.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Explain this complex logic.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai explain this code block.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and explain its main purpose.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

@NanoCode012

Copy link
Copy Markdown
Collaborator Author

Ran test and passed locally but other test failing. Not sure if it's some isolation issue.

model_name = 'Qwen/Qwen3-30B-A3B', attention_cls = <class 'transformers.models.qwen3_moe.modeling_qwen3_moe.Qwen3MoeAttention'>

    @pytest.mark.parametrize(
        "model_name,attention_cls",
        [
            ("HuggingFaceTB/SmolLM2-135M", LlamaAttention),
            ("Qwen/Qwen3-30B-A3B", Qwen3MoeAttention),
        ],
    )
    def test_attention_patching_integration(model_name, attention_cls):
        """Test attention patching in integration context."""
        cfg = DictDefault({"base_model": model_name})
    
        # Store the original implementation
        original_forward = getattr(attention_cls, "forward")
    
        # Apply patch
        patch_self_attn_lora(cfg)
    
        # Get the new forward method
        patched_forward = attention_cls.forward
    
        # Check the forward method was replaced
>       assert original_forward is not patched_forward
E       assert <function Qwen3MoeAttention.forward at 0x7a56d1cb1800> is not <function Qwen3MoeAttention.forward at 0x7a56d1cb1800>

tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py:100: AssertionError
---------------------------------------------------------------- Captured stdout call ----------------------------------------------------------------[2025-06-09 19:04:32,491] [WARNING] [axolotl.monkeypatch.lora_kernels.patch_self_attn_lora:185] [PID:5540] [RANK:0] Cannot patch self-attention - requires no dropout

 tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py ⨯✓✓✓✓✓✓✓✓✓✓✓✓✓                                                            100% ██████████============================================================== short test summary info ===============================================================FAILED tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py::test_attention_patching_integration[HuggingFaceTB/SmolLM2-135M-LlamaAttention] - assert <function LlamaAttention.forward at 0x7a56d1c6b560> is not <function LlamaAttention.forward at 0x7a56d1c6b560>
FAILED tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py::test_attention_patching_integration[Qwen/Qwen3-30B-A3B-Qwen3MoeAttention] - assert <function Qwen3MoeAttention.forward at 0x7a56d1cb1800> is not <function Qwen3MoeAttention.forward at 0x7a56d1cb1800>

Results (45.62s):
      13 passed
       2 failed
         - tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py:79 test_attention_patching_integration[HuggingFaceTB/SmolLM2-135M-LlamaAttention]
         - tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py:79 test_attention_patching_integration[Qwen/Qwen3-30B-A3B-Qwen3MoeAttention]
         - ```

@codecov

codecov Bot commented Jun 9, 2025

Copy link
Copy Markdown

Codecov Report

Attention: Patch coverage is 72.72727% with 3 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/axolotl/loaders/patch_manager.py 50.00% 2 Missing ⚠️
src/axolotl/monkeypatch/lora_kernels.py 85.71% 1 Missing ⚠️

📢 Thoughts on this report? Let us know!

@winglian winglian requested a review from djsaunde June 9, 2025 19:32
Comment thread src/axolotl/monkeypatch/lora_kernels.py Outdated
Comment on lines +181 to +187
# Only patch if conditions are met
can_patch = cfg.lora_dropout == 0

if not can_patch:
LOG.warning("Cannot patch self-attention - requires no dropout")
return

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

move this logic out of this function that does the actual patching, into the caller (def _apply_self_attention_lora_patch) and that will fix the test issue since the test is making sure that it patches.

@djsaunde djsaunde left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks for catching this!

Comment thread src/axolotl/monkeypatch/lora_kernels.py Outdated
# Only patch if conditions are met
can_patch = cfg.lora_dropout == 0
can_patch = (
cfg.lora_dropout == 0 if hasattr(cfg, "lora_dropout") else True

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

since it's a DictDefault, isn't hasattr always true?

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

🧹 Nitpick comments (1)
src/axolotl/loaders/patch_manager.py (1)

169-179: LGTM! The precondition check effectively addresses the reported bug.

The added guard clause correctly prevents self-attention LoRA patching when dropout is enabled, which aligns with the test failures and warning message observed in the PR comments. This should resolve the issue where pre-patch was applied despite post-patch not being applied.

Consider enhancing type safety for the dropout comparison:

-            can_patch = (
-                self.cfg.lora_dropout == 0
-                if hasattr(self.cfg, "lora_dropout")
-                else True
-            )  # default to True if lora_dropout is not set
+            can_patch = True  # default to True if lora_dropout is not set
+            if hasattr(self.cfg, "lora_dropout"):
+                dropout_val = getattr(self.cfg, "lora_dropout", 0)
+                can_patch = isinstance(dropout_val, (int, float)) and dropout_val == 0

Note: This change may impact backward compatibility for configurations that previously had non-zero dropout values and relied on patching being applied anyway.

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c2bc57c and 7f1329d.

📒 Files selected for processing (2)
  • src/axolotl/loaders/patch_manager.py (1 hunks)
  • src/axolotl/monkeypatch/lora_kernels.py (2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
  • src/axolotl/monkeypatch/lora_kernels.py
⏰ Context from checks skipped due to timeout of 90000ms (8)
  • GitHub Check: PyTest from Source Dist (3.11, 2.7.0)
  • GitHub Check: PyTest from Source Dist (3.11, 2.6.0)
  • GitHub Check: PyTest from Source Dist (3.11, 2.5.1)
  • GitHub Check: PyTest (3.11, 2.6.0)
  • GitHub Check: PyTest (3.11, 2.7.0)
  • GitHub Check: PyTest (3.11, 2.5.1)
  • GitHub Check: pre-commit
  • GitHub Check: pre-commit

Comment thread src/axolotl/loaders/patch_manager.py
@NanoCode012 NanoCode012 force-pushed the fix/lora-kernel-patch branch from 7f1329d to e4af596 Compare June 13, 2025 18:52
@NanoCode012 NanoCode012 merged commit 21388cf into main Jun 14, 2025
17 checks passed
@NanoCode012 NanoCode012 deleted the fix/lora-kernel-patch branch June 14, 2025 18:54
@coderabbitai coderabbitai Bot mentioned this pull request Sep 25, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

apply_lora_kernel_patches can be skipped but patch_self_attn_lora can rely on it breaking forward

3 participants