Skip to content

Add native FSDP2 module and import migration (Phase A PR-2)#1

Closed
3outeille wants to merge 1 commit into
split/a-pr-1-distributed-configfrom
split/a-pr-2-fsdp-module
Closed

Add native FSDP2 module and import migration (Phase A PR-2)#1
3outeille wants to merge 1 commit into
split/a-pr-1-distributed-configfrom
split/a-pr-2-fsdp-module

Conversation

@3outeille

Copy link
Copy Markdown
Owner

Summary

  • Add src/transformers/distributed/fsdp.py with FSDP2 wrapping (apply_fully_shard_data_parallel), plan expansion/verification, and FSDP1/2 detection
  • Keep integrations/fsdp.py as a backward-compatible re-export shim
  • Migrate core imports in accelerate.py, trainer.py, trainer_seq2seq.py, and generation/utils.py
  • Export is_fsdp_enabled, is_fsdp_managed_module, verify_fsdp_plan from transformers.distributed
  • Add unit tests in tests/test_distributed_fsdp.py

Stacks on: #??? (Phase A PR-1: split/a-pr-1-distributed-config)
Part of: splitting huggingface/transformers#46269

Not included (later PRs)

  • tests/test_fsdp_mixin.py integration tests (needs A-PR-4 orchestration + A-PR-5 model FSDP plans)
  • distributed/utils.py / from_pretrained(distributed_config=...) (A-PR-4)
  • Removal of integrations/__init__.py fsdp exports (kept for BC via re-export)

Test plan

  • pytest tests/test_distributed_fsdp.py tests/test_distributed_config.py -v
  • make style
  • Accelerate FSDP1 path unchanged (integrations.fsdp.is_fsdp_enabled still works)
  • Model files importing integrations.fsdp unchanged (re-export shim)

AI assistance was used for this PR.

Made with Cursor

Move FSDP2 wrapping and plan verification to distributed/fsdp.py, keep
integrations/fsdp.py as a backward-compatible re-export, and update core
call sites to import from transformers.distributed.fsdp.
@3outeille

Copy link
Copy Markdown
Owner Author

Opened on huggingface/transformers instead.

@3outeille 3outeille closed this Jun 17, 2026
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