Adding torch accelerator and requirements file to FSDP2 example

[dggaytan]

Adding torch accelerator support to FSDP2 example and

Updates to FSDP2 example:

• Script Renaming and Documentation Updates:

  • Renamed train.py to example.py and updated references in README.md to reflect the new filename. Added instructions to install dependencies via requirements.txt before running the example.

• GPU Verification and Device Initialization:

  • Added a verify_min_gpu_count function to ensure at least two GPUs are available before running the example.
  • Updated device initialization in main() to dynamically detect and configure the device type using torch.accelerator. This improves compatibility with different hardware setups.

New supporting files:

• Dependency Management:

  • Added a requirements.txt file listing required dependencies (torch>=2.7 and numpy).

• Script for Running Examples:

  • Introduced run_example.sh to simplify launching FSDP2 example.

• Integration into Distributed Examples:

  • Added a new function distributed_FSDP2 in run_distributed_examples.sh to include the FSDP2 example in the distributed testing workflow.

  CC: @msaroufim @malfet @dvrogozh

[netlify]

✅ Deploy Preview for pytorch-examples-preview canceled.

Name	Link
🔨 Latest commit	5e960d8
🔍 Latest deploy log	https://app.netlify.com/projects/pytorch-examples-preview/deploys/68826ce9e58ebb000857417b

[dvrogozh]

Why do we need device: torch.device = instead of just device =?

[dggaytan]

It was just a flag for me, but I'll change it to use just torch.device

[dggaytan]

done :)

[dvrogozh]

I think these 2 lines should work for cpu as well. You can simplify the code:

if torch.accelerator.is_available():
     ...
else:
     device = torch.device("cpu")

backend = torch.distributed.get_default_backend_for_device(device)
torch.distributed.init_process_group(backend=backend, device_id=device)

[dggaytan]

done

[soumith]

thank you!