[FSDP] Removed clamp to NO_SHARD for world size 1

[awgu]

Stack from ghstack (oldest at bottom):

• -> [FSDP] Removed clamp to NO_SHARD for world size 1 #120334
• [FSDP] Warned on wrapping ModuleList/ModuleDict #124764

This is a workaround for #119937. NO_SHARD always using unsharded views breaks with DTensor extensions since some logic assumes that sharded views are torch.Tensors, but NO_SHARD makes sharded views and unsharded views the same (as DTensors).

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225 @chauhang @d4l3k @rohan-varma

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/120334

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures, 9 Unrelated Failures

As of commit 0c04620 with merge base c82fcb7 ():

NEW FAILURES - The following jobs have failed:

• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (cpu_inductor_timm, 1, 2, linux.12xlarge) (gh)
  beit_base_patch16_224
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (cpu_inductor_timm, 2, 2, linux.12xlarge) (gh)
  mnasnet_100
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (dynamic_cpu_inductor_timm, 1, 2, linux.12xlarge) (gh)
  beit_base_patch16_224
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (dynamic_cpu_inductor_torchbench, 2, 2, linux.12xlarge) (gh)
  moondream
• pull / linux-focal-cuda11.8-py3.10-gcc9 / test (distributed, 3, 3, linux.8xlarge.nvidia.gpu) (gh)
  distributed/fsdp/test_wrap.py::TestAutoWrap::test_auto_wrap_preset_exclude_wrap

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• inductor / cuda12.1-py3.10-gcc9-sm86 / test (dynamic_inductor_timm, 2, 2, linux.g5.4xlarge.nvidia.gpu) (gh)
  sebotnet33ts_256
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (cpu_inductor_huggingface, 1, 1, linux.12xlarge) (gh)
  Process completed with exit code 1.
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (cpu_inductor_torchbench, 2, 2, linux.12xlarge) (gh)
  moondream
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (dynamic_cpu_inductor_huggingface, 1, 1, linux.12xlarge) (gh)
  Process completed with exit code 1.
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (dynamic_cpu_inductor_timm, 2, 2, linux.12xlarge) (gh)
  mixnet_l
• inductor / linux-jammy-cpu-py3.8-gcc11-inductor / test (inductor_torchbench_cpu_smoketest_perf, 1, 1, linux.12xlarge) (gh)
  Process completed with exit code 1.
• inductor / rocm6.0-py3.8-inductor / test (inductor, 1, 1, linux.rocm.gpu.2) (gh)
  test/distributed/test_c10d_functional_native.py::TestWithNCCL::test_threading
• pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
  test_foreach.py::TestForeachCUDA::test_binary_op_list_error_cases__foreach_clamp_max_cuda_float64

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• inductor / cuda12.1-py3.10-gcc9-sm80 / test (inductor_torchbench_smoketest_perf, 1, 1, linux.gcp.a100) (gh)
  hf_T5

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[mvpatel2000]

@awgu maybe it's worth emitting a warning still that NO_SHARD might be more performant in this case? Not strong opinion though

[awgu]

We have been seeing some issues with NO_SHARD internally, so we were planning to deprecate it. We do want to provide a good solution for a replicated data parallelism that can compose with FSDP (and we have started discussing this).

[mvpatel2000]

#119937 (comment)

Flagging that I still see SHARD_GRAD_OP issue with this :( (but FULL_SHARD works fine!)

[github-actions]

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.