Adding torch accelerator and requirements file to FSDP2 example #1375
+46 −8
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
✅ Deploy Preview for pytorch-examples-preview canceled.
|
dvrogozh reviewed Jul 23, 2025
distributed/FSDP2/example.py Outdated
| torch.distributed.init_process_group(backend="nccl", device_id=device) | ||
| if torch.accelerator.is_available(): | ||
| device_type = torch.accelerator.current_accelerator() | ||
| device: torch.device = torch.device(f"{device_type}:{rank}") |
There was a problem hiding this comment.
Why do we need device: torch.device = instead of just device =?
There was a problem hiding this comment.
It was just a flag for me, but I'll change it to use just torch.device
dvrogozh reviewed Jul 23, 2025
distributed/FSDP2/example.py Outdated
Comment on lines 47 to 48
| backend = torch.distributed.get_default_backend_for_device(device) | ||
| torch.distributed.init_process_group(backend=backend, device_id=device) |
There was a problem hiding this comment.
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) 
Signed-off-by: dggaytan <diana.gaytan.munoz@intel.com>
dggaytan force-pushed the dggaytan/distributed_FSDP2 branch from 1f0d7d3 to 5e960d8 Compare 
soumith approved these changes Aug 6, 2025
| thank you! |
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit. This suggestion is invalid because no changes were made to the code. Suggestions cannot be applied while the pull request is closed. Suggestions cannot be applied while viewing a subset of changes. Only one suggestion per line can be applied in a batch. Add this suggestion to a batch that can be applied as a single commit. Applying suggestions on deleted lines is not supported. You must change the existing code in this line in order to create a valid suggestion. Outdated suggestions cannot be applied. This suggestion has been applied or marked resolved. Suggestions cannot be applied from pending reviews. Suggestions cannot be applied on multi-line comments. Suggestions cannot be applied while the pull request is queued to merge. Suggestion cannot be applied right now. Please check back later.
Adding torch accelerator support to FSDP2 example and
Updates to FSDP2 example:
Script Renaming and Documentation Updates:
train.pytoexample.pyand updated references inREADME.mdto reflect the new filename. Added instructions to install dependencies viarequirements.txtbefore running the example.GPU Verification and Device Initialization:
verify_min_gpu_countfunction to ensure at least two GPUs are available before running the example.main()to dynamically detect and configure the device type usingtorch.accelerator. This improves compatibility with different hardware setups.New supporting files:
Dependency Management:
requirements.txtfile listing required dependencies (torch>=2.7andnumpy).Script for Running Examples:
run_example.shto simplify launching FSDP2 example.Integration into Distributed Examples:
distributed_FSDP2inrun_distributed_examples.shto include the FSDP2 example in the distributed testing workflow.CC: @msaroufim @malfet @dvrogozh