feat(fsdp): expose fsdp's min_num_params for SIZE_BASED_WRAP in config .yml#3710
feat(fsdp): expose fsdp's min_num_params for SIZE_BASED_WRAP in config .yml#3710thad0ctor wants to merge 3 commits into
Conversation
Wire the FSDP_MIN_NUM_PARAMS env var from YAML so the size-based auto-wrap policy can be configured without setting the env var manually.
|
Complex PR? Review this PR in Change Stack to move by importance, not file order. 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 PR adds support for the ChangesFSDP min_num_params Support
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes 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: 2
🤖 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/utils/schemas/fsdp.py`:
- Around line 69-72: FSDPConfig's schema must validate that when
auto_wrap_policy == "SIZE_BASED_WRAP" the min_num_params field is set to a
positive integer; update the Pydantic schema in
src/axolotl/utils/schemas/fsdp.py by adding a validator (e.g., a `@root_validator`
or `@validator` referencing FSDPConfig, auto_wrap_policy, and min_num_params) that
raises a ValueError if auto_wrap_policy == "SIZE_BASED_WRAP" and (min_num_params
is None or min_num_params <= 0), ensuring invalid combinations are rejected at
schema validation time.
In `@src/axolotl/utils/trainer.py`:
- Around line 617-618: The current truthy check on
cfg.fsdp_config.min_num_params will skip explicit falsy values like 0; change
the condition in trainer.py to test for None explicitly (use "is not None" on
cfg.fsdp_config.min_num_params) so FSDP_MIN_NUM_PARAMS is exported whenever the
config provides a value (including 0), and keep the assignment to
os.environ["FSDP_MIN_NUM_PARAMS"] = str(cfg.fsdp_config.min_num_params) intact.
🪄 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: a72c8dc3-6bed-4398-8a77-b40ac8c479eb
📒 Files selected for processing (3)
docs/multi-gpu.qmdsrc/axolotl/utils/schemas/fsdp.pysrc/axolotl/utils/trainer.py
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
The new FSDPConfig validator requires min_num_params when auto_wrap_policy is SIZE_BASED_WRAP, so the existing migration fixture now needs the field to pass validation.
1 participant
Description
The
SIZE_BASED_WRAPauto-wrap policy is selectable viafsdp_config.auto_wrap_policy, but there was no way to set its parameter threshold from the config, it was only configurable through env variable. This adds amin_num_paramsfield toFSDPConfigand wires it through to theFSDP_MIN_NUM_PARAMSenv var for Accelerate.Motivation and Context
FSDP only honors a size threshold when
FSDP_MIN_NUM_PARAMSis set (> 0); otherwise it reads the default of0andset_auto_wrap_policy()silently drops the policy toNoneunless manually setting ENV variables.Axolotl's
setup_fsdp_envs()exported every other FSDP knob but never this one, so selectingauto_wrap_policy: SIZE_BASED_WRAPresulted in no modules being wrapped.How has this been tested?
setup_fsdp_envs()exportsFSDP_MIN_NUM_PARAMSwhenmin_num_paramsis set, and leaves it unset otherwise (no behavior change for existing configs).min_num_params=100000000andset_auto_wrap_policy(model)yields an active size-based policy bound to that threshold; with the field unset, the pluginreads
0and the policy resolves toNone(confirming the field is required).transformers==5.9.0,accelerate==1.13.0, local venv.AI Usage Disclaimer
Opus 4.8 assisted with diagnosing and fix
Types of changes
Summary by CodeRabbit
New Features
min_num_paramsconfiguration in FSDP settings to control the minimum parameter count threshold for automatic module wrapping during distributed training.Documentation
min_num_paramsis required when theSIZE_BASED_WRAPauto-wrap policy is enabled; otherwise the policy will be ignored.