Fix: adding magistral fsdp config, fixing not eval with test_datasets, handle mllama attention#2789
Conversation
WalkthroughA new YAML configuration was added for QLoRA fine-tuning with FSDP on the Magistral-Small model. The logic for disabling evaluation in training arguments was changed to require both no evaluation dataset and zero validation set size. The LoRA attention kernel monkeypatch now supports the "mllama" model type. A usage note was added for cut_cross_entropy integration regarding VLM training compatibility. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Config
participant Trainer
participant Model
User->>Config: Provide magistral-small-fsdp-qlora.yaml
Config->>Trainer: Load training, LoRA, FSDP, and eval params
Trainer->>Model: Initialize with QLoRA, FSDP, and tokenizer
Trainer->>Trainer: Check eval dataset and val set size
Trainer->>Model: Train and evaluate as configured
Assessment against linked issues
Assessment against linked issues: Out-of-scope changes
Possibly related PRs
Suggested labels
Suggested reviewers
Poem
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed 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)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![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)
examples/magistral/magistral-small-fsdp-qlora.yaml (1)
17-18: Floating fraction forval_set_sizerelies on downstream support
val_set_size: 0.1is convenient, but not every Axolotl code-path treats floats and ints identically (see change inbase.py).
Make sure the dataset-splitting utility you use converts this fraction into an absolute size; otherwise evaluation will be silently skipped.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
examples/magistral/magistral-small-fsdp-qlora.yaml(1 hunks)src/axolotl/core/builders/base.py(1 hunks)src/axolotl/monkeypatch/lora_kernels.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (6)
- GitHub Check: PyTest (3.11, 2.5.1)
- GitHub Check: PyTest from Source Dist (3.11, 2.6.0)
- GitHub Check: PyTest from Source Dist (3.11, 2.7.1)
- 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.1)
🔇 Additional comments (1)
src/axolotl/monkeypatch/lora_kernels.py (1)
148-152:Details
❓ Verification inconclusive
Verify the new ‘mllama’ special-case import
transformersmainline currently exposesMllamaAttention, notMllamaTextSelfAttention.
If the class name or module path is wrong the patch will raise at import time and break all LoRA kernel initialisation.Please run a quick check (or pin a compatible transformers commit) before merge.
🏁 Script executed:
#!/bin/bash # Attempt to import the proposed class – will exit non-zero on failure. python - <<'PY' import importlib, sys, traceback try: mod = importlib.import_module("transformers.models.mllama.modeling_mllama") getattr(mod, "MllamaTextSelfAttention") print("✅ Import succeeded") except Exception: traceback.print_exc() sys.exit(1) PYLength of output: 1431
Verify the new ‘mllama’ special-case import
The automated import check couldn’t run because the
transformerslibrary isn’t installed in this environment. Please confirm in your local setup that:
transformers.models.mllama.modeling_mllama.MllamaTextSelfAttentionis present in the version(s) you support- You’ve pinned or documented a compatible
transformersrelease to avoid breakage at import time
| if not self.eval_dataset and self.cfg.val_set_size == 0: | ||
| # do not eval if no eval_dataset and val_set_size=0 | ||
| training_args_kwargs["eval_strategy"] = "no" |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Condition should use truth-value rather than strict equality
val_set_size can be supplied as None, 0, or 0.0.
Using the strict comparison == 0 means that the “no-eval” shortcut will be skipped when the field is None, even though that represents no validation split just as well.
- if not self.eval_dataset and self.cfg.val_set_size == 0:
+ # disable evaluation only when BOTH no explicit eval dataset
+ # and the user provided an equivalent “no-split” value
+ if self.eval_dataset is None and not self.cfg.val_set_size:This guards against None and negative / falsy values in one shot, while still allowing fractions such as 0.1.
📝 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.
| if not self.eval_dataset and self.cfg.val_set_size == 0: | |
| # do not eval if no eval_dataset and val_set_size=0 | |
| training_args_kwargs["eval_strategy"] = "no" | |
| # disable evaluation only when BOTH no explicit eval dataset | |
| # and the user provided an equivalent “no-split” value | |
| if self.eval_dataset is None and not self.cfg.val_set_size: | |
| # do not eval if no eval_dataset and val_set_size=0 | |
| training_args_kwargs["eval_strategy"] = "no" |
🤖 Prompt for AI Agents
In src/axolotl/core/builders/base.py around lines 383 to 385, replace the strict
equality check 'self.cfg.val_set_size == 0' with a truth-value check like 'not
self.cfg.val_set_size' to correctly handle cases where val_set_size is None, 0,
or 0.0. This change ensures the "no-eval" shortcut triggers appropriately for
all falsy values representing no validation split, while still allowing
fractional values such as 0.1.
Codecov ReportAttention: Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
src/axolotl/integrations/cut_cross_entropy/README.md (2)
27-27: Grammar: add missing article "an" for clarity
The sentence reads:NOTE: If you are training a VLM model, please use older version of Axolotl...
It would be clearer as:
NOTE: If you are training a VLM model, please use an older version of Axolotl...🧰 Tools
🪛 LanguageTool
[uncategorized] ~27-~27: Possible missing article found.
Context: ...ou are training a VLM model, please use older version of Axolotl as upstream has appl...(AI_HYDRA_LEO_MISSING_AN)
29-33: Clarify context for repository commands
Thegit checkoutandpip3 installsteps assume the user is in the Axolotl repo root. Consider adding a note, for example:+Run these commands from the root of the Axolotl repository: ```bash git checkout 787880215b3ab32ccaf81c1b2e9588c6f3e6e764 pip3 install --no-build-isolation -e .</blockquote></details> </blockquote></details> <details> <summary>📜 Review details</summary> **Configuration used: CodeRabbit UI** **Review profile: CHILL** **Plan: Pro** <details> <summary>📥 Commits</summary> Reviewing files that changed from the base of the PR and between 7a21195f455b31d5d3e29515a1c929940fc159d3 and 5df2688355b2a7d076b0caf308082b760dad32b9. </details> <details> <summary>📒 Files selected for processing (1)</summary> * `src/axolotl/integrations/cut_cross_entropy/README.md` (1 hunks) </details> <details> <summary>🧰 Additional context used</summary> <details> <summary>🪛 LanguageTool</summary> <details> <summary>src/axolotl/integrations/cut_cross_entropy/README.md</summary> [uncategorized] ~27-~27: Possible missing article found. Context: ...ou are training a VLM model, please use older version of Axolotl as upstream has appl... (AI_HYDRA_LEO_MISSING_AN) </details> </details> </details> <details> <summary>⏰ Context from checks skipped due to timeout of 90000ms (9)</summary> * GitHub Check: PyTest (3.11, 2.6.0) * GitHub Check: pre-commit * GitHub Check: PyTest (3.11, 2.7.1) * GitHub Check: PyTest from Source Dist (3.11, 2.5.1) * GitHub Check: PyTest (3.11, 2.5.1) * GitHub Check: PyTest from Source Dist (3.11, 2.7.1) * GitHub Check: PyTest from Source Dist (3.11, 2.6.0) * GitHub Check: pre-commit * GitHub Check: preview </details> </details> <!-- This is an auto-generated comment by CodeRabbit for review status -->
2 participants
Description
Motivation and Context
How has this been tested?
Screenshots (if appropriate)
Types of changes
Social Handles (Optional)
Summary by CodeRabbit
New Features
Bug Fixes
Refactor
Documentation