add building workflow for TE/Jax

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.

[jeng1220]

This is the first commit for TE/Jax.

The build process for TE/PyTorch and TE/Jax components is different. In the building step, the TE/Jax only depends on CPP STL and CUDA, so we simply add needed building commands in CMakeList.txt. Unlike PyTorch, Jax has nothing like torch.utils.cpp_extension.BuildExtension and torch.utils.cpp_extension.CUDAExtension.

To ensure that users can build only TE/PyTorch, only TE/Jax, or both, I created FrameworkBuilder to manage framework-specific stuff. For example, if Jax users want to build TE/Jax, then PyTorch should not be required in their environment.

The user can set up an environment variable - FRAMEWORK to select framework, such as:

FRAMEWORK=jax pip install .
FRAMEWORK=pytorch pip install .
FRAMEWORK=all pip install .

The default value is all

The other Jax modules will be submitted after this PR is merged.

[jeng1220]

Addition, TE/Jax need pybind11 to build Python-CPP binding, so we need

apt install pybind11-dev
pip install pybind11

Both of above packages are needed. Otherwise, CMake will return:

CMake Error at CMakeLists.txt:32 (find_package):
Could not find a package configuration file provided by "pybind11" with any
of the following names:
pybind11Config.cmake
pybind11-config.cmake

[jeng1220]

@timmoon10 ,
The CI triggers build failed:
https://github.com/NVIDIA/TransformerEngine/actions/runs/3901104182/jobs/6662679011#step:4:85

As I mentioned, TE/Jax needs pybind11
We have to install pybind11-dev in the CI container.

$ apt install pybind11-dev

[timmoon10]

Sorry for the delay. I think this is mostly fine. I could imagine some future build errors if, for example, the CMake flags for the PyTorch and JAX extensions are incompatible. That seems unlikely though, and we can refactor setup.py as any issues come up.

The biggest issues are merge conflicts. In particular, #51 removes the scale_inv arg from many of the TE functions and makes it the responsibility of the frameworks. This will likely also affect #54.

[timmoon10]

Once we iron out the style and build issues, I think this is okay to merge.

I'm not thrilled with how we're refactoring layer norm and RMSNorm to use the same helper function. I think it's a premature optimization. In principle they are independent operations that just happen to look similar. But in practice I don't expect them to undergo significant API changes in the future, and this is an internal implementation detail anyways.

[timmoon10]

I see the linter found some style issues.

[timmoon10]

/te-ci

[timmoon10]

I have a feeling enums are the wrong data structure here. The number of enums grows exponentially with the number of sharding dimensions and I frequently see messy patterns like:

ifsharding_typein (ShardingType.TP_COL, ShardingType.DP_TP_COL):
tp_col_impl()
else:
default_impl()

It would feel better to keep track of separate bools for DP and TP. In fact, I wonder if it's better to split up TP and treat TP_COL and TP_ROW as completely orthogonal sharding axes. Treating the sharding axes as orthogonal would help with code encapsulation and make it easier to do things like add new sharding axes (e.g. over the sequence dim) without needing to touch everything.

Maybe not something to do right now if deadlines are urgent, but something to think about for a future refactor.

[jeng1220]

This is an action item in our schedule. But to refactor this part, probably will take more than 2 weeks. That will be too late to release TE/JAX.

Those sharding stuff is low-level implementation in TE/JAX. The TE/JAX users don't and shouldn't call them directly. In the next PR, we will provide high-level module to hide these stuff.

I fully agree this part should be better. We will continue refactoring after first TE/JAX release is done.

[timmoon10]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

[jeng1220]

The test failures are unrelated to this PR (see #67 (comment)), so I think this is ready after addressing @ksivaman's concerns. I don't have full confidence that I got everything and in the future I would try to keep each PR with a clearly defined scope. It's much faster and safer to review multiple simple PRs than a single massive one.

@timmoon10 ,
I understand it is difficult to review so much files at once.
The next PR will include less files. Hope it can be easier for code review.
I appreciate your help very much.

[timmoon10]

/te-ci

[jeng1220]

@ksivaman and @timmoon10 ,
I have merged the main

[timmoon10]

@ksivaman and I think this is ready to merge, but we'd like to discuss with @ptrendx tomorrow. Barring any issues, we'll merge it by California afternoon.