Add plugins documentation #2207
Merged
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
github-actions
bot
added
the
documentation
Borda
approved these changes
Jun 10, 2025
k223kim
approved these changes
Jun 10, 2025
ethanwharris
approved these changes
Jun 10, 2025
t-vi
approved these changes
Jun 10, 2025
Collaborator
t-vi
left a comment
t-vi
left a comment
There was a problem hiding this comment.
Supergood to have those documented. Thank you!
crcrpar
reviewed
Jun 10, 2025
| This plugin applies the necessary transforms to bucket and synchronize gradients across | ||
| multiple processes, using a specified process group for communication. | ||
| See https://github.com/pytorch/pytorch/blob/v2.7.0/torch/nn/parallel/distributed.py#L326 for more details. |
Collaborator
There was a problem hiding this comment.
How about cross-referencing DDP instead of referencing pytorch doc of a certain version?
e.g. :class:~torch.nn.parallel.distributed.DistributedDataParallel would work
Comment on lines
+22
to
+27
| Args: | ||
| bucket_size_in_mb: float, default 25.0 | ||
| Size in megabytes of the gradient bucket in DDP. | ||
| broadcast_from: int | None, default None | ||
| Global rank ID to broadcast model parameters from at initialization. If None, no explicit broadcast is performed. | ||
| process_group: Optional[ProcessGroup], default is the current default process group |
Collaborator
There was a problem hiding this comment.
type annotation and default values can be obviated looking at
Comment on lines
+69
to
+83
| Args: | ||
| device: torch.device | None, default None | ||
| Device on which to place sharded modules. If None, modules remain on their existing devices. | ||
| broadcast_from: int | None, default None | ||
| Global rank ID to broadcast parameters from before sharding. If None, no broadcast is performed. | ||
| sharding_strategy: FSDPType, default FSDPType.ZERO2 | ||
| Strategy for parameter sharding (e.g., ZERO2 for sharding both parameters and optimizer state). | ||
| bucketing_strategy: FSDPBucketingStrategy, default FSDPBucketingStrategy.NONE | ||
| Bucketing strategy to use when saving or loading FSDP checkpoints. | ||
| move_state_dict_to_cpu: bool, default False | ||
| Whether to move the state dict parameters to CPU after serialization to reduce GPU memory usage. | ||
| ddp_bucket_size_in_mb: float, default 25.0 | ||
| Bucket size in megabytes for the DDP transform when used in a combined mesh with FSDP. | ||
| process_group: Optional[ProcessGroup or DeviceMesh], default is the current default process group | ||
| The process group or device mesh to use for distributed communication. If None, uses the default process group. |
Collaborator
There was a problem hiding this comment.
same. I think we can skip type annotations and default values
| """ | ||
| Plugin for enabling FP8 precision via NVIDIA Transformer Engine, enabling higher throughput of matrix operations in FP8. | ||
| See `lightning-thunder/thunder/executors/transformer_engineex.py` for implementation details. |
Collaborator
There was a problem hiding this comment.
Can we reference this file with a relative path?
| model weights, reducing memory footprint and improving | ||
| throughput for both training and inference. | ||
| See https://github.com/bitsandbytes-foundation/bitsandbytes/blob/main/bitsandbytes/functional.py#L889 for more details. |
Collaborator
There was a problem hiding this comment.
it'd be better if this link is a permalink
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 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.
Before submitting
What does this PR do?
This PR helps first-time users understanding plugins better by adding documentation for plugins, namely the DDP, FSDP, QuantizeInt4, FP8 and ReduceOverhead plugins.
PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.
Did you have fun?
Make sure you had fun coding 🙃
cc @Borda @lantiga