Skip to content

Add optimized TBE training forward #1804

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
Closed

Conversation

sryap
Copy link
Contributor

@sryap sryap commented Jun 2, 2023

Summary:
This diff adds the frontend changes and tests for TBE v2 (D43634651)

The FBGEMM_EXPERIMENTAL_TBE environment variable flag is added for
enabling/disabling the new implementation at runtime. If
FBGEMM_EXPERIMENTAL_TBE is not set, TBE will use the orignal
implementation. If FBGEMM_EXPERIMENTAL_TBE=1, TBE will use the new
implementation. If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation. By
default, FBGEMM_EXPERIMENTAL_TBE is not set.

This can also be enabled by passing use_experimental_tbe=True when
instantiating the TBE operator.

emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)

Differential Revision: D44479772

@netlify
Copy link

netlify bot commented Jun 2, 2023

Deploy Preview for pytorch-fbgemm-docs canceled.

Name Link
🔨 Latest commit 59cc254
🔍 Latest deploy log https://app.netlify.com/sites/pytorch-fbgemm-docs/deploys/64879de609503b00086fb1bd

@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 5, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: e5b6ad27fc6a914a4def4d988946555daf0072a0
@sryap sryap force-pushed the export-D44479772 branch from 149705a to 194a524 Compare June 5, 2023 17:36
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 5, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: b961811488a25904a3f34660c553067b1ab93c95
@sryap sryap force-pushed the export-D44479772 branch from 194a524 to 46f2559 Compare June 5, 2023 17:41
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 5, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: 38944e9654ea06c3bec20b8c72d39a576f325367
@sryap sryap force-pushed the export-D44479772 branch from 46f2559 to e90b058 Compare June 5, 2023 20:32
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 5, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: 86328c01d296182f0b9180b05b32cdfc4c5ea750
@sryap sryap force-pushed the export-D44479772 branch from e90b058 to 6bc90ac Compare June 5, 2023 20:42
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 6, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: cf6b6460194b7c6387dcfe82f2b7c6f4e8e06627
@sryap sryap force-pushed the export-D44479772 branch from 6bc90ac to 36acf18 Compare June 6, 2023 08:38
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 6, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: 232517acf727e0373876d1b31b45bbf059e7a2c0
@sryap sryap force-pushed the export-D44479772 branch from 36acf18 to 4458f1f Compare June 6, 2023 08:44
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 7, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: c5281c474671ebe8efb0064db4a8fd2a97c3d917
@sryap sryap force-pushed the export-D44479772 branch from 4458f1f to d351db3 Compare June 7, 2023 05:09
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 7, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: c5130ffdbee8bb6d8e8112386635609738ea8a84
@sryap sryap force-pushed the export-D44479772 branch from d351db3 to f9bbf0c Compare June 7, 2023 09:55
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 7, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: f1823db21a155794e6c2e7602ef7b37a7bd0ad59
@sryap sryap force-pushed the export-D44479772 branch from f9bbf0c to f36271f Compare June 7, 2023 15:17
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 7, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: df896e3aeb129d32a09cea6ac9873db43a96cbb3
@sryap sryap force-pushed the export-D44479772 branch from f36271f to 75ffd26 Compare June 7, 2023 15:24
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 8, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: e0530c653023eb57f408cf9b8d3d89486cc503f1
@sryap sryap force-pushed the export-D44479772 branch from 75ffd26 to f41b769 Compare June 8, 2023 07:13
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 8, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: e441f0b4d4c76bf69aa5587e2e58766a13b59f4c
@sryap sryap force-pushed the export-D44479772 branch from f41b769 to 2d96854 Compare June 8, 2023 07:21
@facebook-github-bot
Copy link
Contributor

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

sryap added a commit to sryap/FBGEMM that referenced this pull request Jun 12, 2023
Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: f56b60a071f8c44b046de8a0c15d5d0f46dd1570
@sryap sryap force-pushed the export-D44479772 branch from 2d96854 to d435a2d Compare June 12, 2023 22:29
@facebook-github-bot
Copy link
Contributor

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

Summary:
Pull Request resolved: pytorch#1804

This diff adds the frontend changes and tests for TBE v2 (D43634651)

The `FBGEMM_EXPERIMENTAL_TBE` environment variable flag is added for
enabling/disabling the new implementation at runtime.  If
`FBGEMM_EXPERIMENTAL_TBE` is not set, TBE will use the orignal
implementation.  If `FBGEMM_EXPERIMENTAL_TBE=1`, TBE will use the new
implementation.  If the TBE usecases are not supported in the new
implementation, TBE will fall back to the original implementation.  By
default, `FBGEMM_EXPERIMENTAL_TBE` is not set.

This can also be enabled by passing `use_experimental_tbe=True` when
instantiating the TBE operator.

```
emb_op = SplitTableBatchedEmbeddingBagsCodegen(
    embedding_specs=...,
    ...,
    use_experimental_tbe=True,
)
```

Reviewed By: jianyuh

Differential Revision: D44479772

fbshipit-source-id: 295fe8e28ba241cb640db15706dc15866f3effc8
@sryap sryap force-pushed the export-D44479772 branch from d435a2d to 59cc254 Compare June 12, 2023 22:36
@facebook-github-bot
Copy link
Contributor

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

@facebook-github-bot
Copy link
Contributor

This pull request has been merged in 38efe0d.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants