Skip to content

[FSDP] Validate activation checkpointing auto-wrap policy - #4140

Draft
sdjasj wants to merge 1 commit into
huggingface:mainfrom
sdjasj:agent/fix-fsdp-ac-no-wrap
Draft

[FSDP] Validate activation checkpointing auto-wrap policy#4140
sdjasj wants to merge 1 commit into
huggingface:mainfrom
sdjasj:agent/fix-fsdp-ac-no-wrap

Conversation

@sdjasj

@sdjasj sdjasj commented Jul 29, 2026

Copy link
Copy Markdown

What does this PR do?

Rejects FSDP activation checkpointing configurations that use NO_WRAP and raises an actionable ValueError during plugin initialization.

Activation checkpointing selects modules through the configured auto-wrap policy. NO_WRAP is normalized to None, so the previous code failed later in model preparation with an opaque TypeError: 'NoneType' object is not callable.

The regression coverage exercises both environment-variable configuration and direct FullyShardedDataParallelPlugin construction for FSDP1 and FSDP2.

Fixes #3822

Validation

python -m pytest -q \
  tests/fsdp/test_fsdp.py::FSDPPluginIntegration::test_activation_checkpointing_requires_auto_wrap_policy \
  tests/fsdp/test_fsdp.py::FSDP2PluginIntegration::test_activation_checkpointing_requires_auto_wrap_policy
# 2 passed

ruff check src/accelerate/utils/dataclasses.py tests/fsdp/test_fsdp.py
ruff format --check src/accelerate/utils/dataclasses.py tests/fsdp/test_fsdp.py

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.

TypeError: NoneType object is not callable occurs when using fsdp_auto_wrap_policy=NO_WRAP with fsdp_activation_checkpointing=true

1 participant