Add construction-scoped MFSDP fully shard options

[wujingyue]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Add a construction-scoped fully_shard_options() context manager for experimental MFSDP v2.

The symmetric-memory setting previously lived on each fully_shard() call. This change moves it into a context-local, immutable construction option so callers can configure a group of module sharding operations together. fully_shard() resolves the active option during construction and passes only the boolean through FsdpModule to FsdpParameterGroup; no options object is retained and the lazily initialized runtime FsdpContext is unchanged.

The experimental per-call use_symm_mem argument is removed without deprecation. Existing symmetric-memory parity, communication-overlap, zero-CTA, and NCCL symmetric-kernel assertions now use fully_shard_options(use_symmetric_memory=...).

Issue tracking

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Validation

• python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py tests/unit_tests/distributed/mfsdp_v2/test_symmetric_memory.py
  • 21 passed, 5 skipped per rank; skips require a 4-rank 2-D mesh.
• black --check on all changed Python files.
• ruff check on all changed Python files.
• isort on all changed Python files.
• git diff --check.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.