Skip to content

Enable More General Comms Bandwidth Estimates #2905

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from

Conversation

gregmacnamara
Copy link
Contributor

Summary:
Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators. We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 22, 2025
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D73229318

gregmacnamara added a commit to gregmacnamara/torchrec that referenced this pull request Apr 22, 2025
Summary:

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D73229318

gregmacnamara added a commit to gregmacnamara/torchrec that referenced this pull request Apr 22, 2025
Summary:
Pull Request resolved: pytorch#2905

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
gregmacnamara added a commit to gregmacnamara/torchrec that referenced this pull request Apr 22, 2025
Summary:

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D73229318

gregmacnamara added a commit to gregmacnamara/torchrec that referenced this pull request Apr 22, 2025
Summary:
Pull Request resolved: pytorch#2905

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
gregmacnamara added a commit to gregmacnamara/torchrec that referenced this pull request Apr 22, 2025
Summary:

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D73229318

gregmacnamara added a commit to gregmacnamara/torchrec that referenced this pull request Apr 22, 2025
Summary:
Pull Request resolved: pytorch#2905

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
gregmacnamara added a commit to gregmacnamara/torchrec that referenced this pull request Apr 23, 2025
Summary:

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
Summary:
Pull Request resolved: pytorch#2905

Optimizing sharding plans relies on good perf estimates which in turn relies on good bandwidth estimates.

Bandwidth of a collective communication depends on many factors including the world size, type of collective (all to all, reduce scatter, ...), and hardware.

In this change, we introduce an abstract class GeneralizedCommsBandwidths that implements a get_bw function. Users can subclass this and define the get_bw appropriately for the hardware or network setup being used. If passed to the Topology, it will be called from perf estimators and will enable setting correct configurations of bandwidth depending on hardware, collective, world_size. We prefer to define this as a class and function as opposed to a dictionary or other similar data structure because we want to enable users to specify bandwidth as precisely as is useful to them. For example, specifying a different bandwidth for every different worldsize rather than only <local_world_size or >=local_world_size.

For Topology objects, the comms_bandwidths argument is optional so we do not change existing behavior and perf estimates do not change. If not passed, we create a GeneralizedCommsBandwidth object with the passed values for inter and intra node comms bw and that replicates the existing behavior of perf estimators.  We re-define the signature of perf estimators downstream to ensure consistency and clarity for users so that there is a consistent comms bandwidth.

Reviewed By: iamzainhuda

Differential Revision: D73229318
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D73229318

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. fb-exported
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants