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 <[email protected]>
dggaytan
force-pushed
the
dggaytan/distributed_FSDP2
branch
from
1f0d7d3 to
5e960d8
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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