Skip to content

Fix to-fsdp2: drop REMOVED / NOT_YET_IMPLEMENTED FSDP1 keys instead of leaking them#4065

Open
lollinng wants to merge 1 commit into
huggingface:mainfrom
lollinng:accelerate-to-fsdp2-drop-removed-keys
Open

Fix to-fsdp2: drop REMOVED / NOT_YET_IMPLEMENTED FSDP1 keys instead of leaking them#4065
lollinng wants to merge 1 commit into
huggingface:mainfrom
lollinng:accelerate-to-fsdp2-drop-removed-keys

Conversation

@lollinng
Copy link
Copy Markdown

@lollinng lollinng commented Jun 4, 2026

Problem

accelerate to-fsdp2 converts an FSDP1 config to FSDP2. In convert_config_to_fsdp2, the first branch of the per-key loop was:

conversion_status = ARGUMENT_KEY_MAPPING.get(key, None)
if isinstance(conversion_status, ConversionStatus) or conversion_status is None:
    conversion_status = key
    new_fsdp_config[conversion_status] = value
    continue

This carries the key over and continues for both "key not in the mapping" (None) and the ConversionStatus enum values (REMOVED / NOT_YET_IMPLEMENTED). As a result, the REMOVED / NOT_YET_IMPLEMENTED handling just below was dead code, and FSDP1-only keys — fsdp_backward_prefetch, fsdp_sync_module_states, fsdp_use_orig_params (REMOVED) and fsdp_forward_prefetch (NOT_YET_IMPLEMENTED) — leaked into the converted FSDP2 config instead of being dropped.

Fix

Only carry over keys that are genuinely absent from the mapping (conversion_status is None); let ConversionStatus values fall through to the REMOVED / NOT_YET_IMPLEMENTED branches that drop them.

Testing done (CPU)

Added tests/test_cli.py::ToFSDP2Tester::test_convert_config_drops_removed_and_unimplemented_keys:

1 passed

It asserts the four FSDP1-only keys are dropped, that a renamed+value-mapped key (fsdp_sharding_strategy: FULL_SHARDfsdp_reshard_after_forward: True) and pass-through / unknown keys are preserved, and that fsdp_version is set to 2. ruff clean.

convert_config_to_fsdp2's first branch caught both 'key not in mapping'
(conversion_status is None) AND the ConversionStatus enum values, carrying
the key over and continuing. That made the REMOVED / NOT_YET_IMPLEMENTED
handling below unreachable, so FSDP1-only keys (fsdp_backward_prefetch,
fsdp_sync_module_states, fsdp_use_orig_params, fsdp_forward_prefetch) leaked
into the converted FSDP2 config. Only carry over keys that are genuinely
absent from the mapping; let ConversionStatus values fall through to be
dropped. Adds a regression test.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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.

1 participant