Skip to content

GitLab

Commit 87e3e56e authored by yuguo's avatar yuguo
Browse files

Merge commit '734bcedd' of... 

Merge commit '734bcedd' of https://github.com/NVIDIA/TransformerEngine
parents 2f11bd2e 734bcedd
Expand all Hide whitespace changes
Inline Side-by-side
Showing with 905 additions and 606 deletions
qa/L1_jax_distributed_unittest/test.sh
View file @ 87e3e56e
...@@ -8,4 +8,4 @@ set -xe ...@@ -8,4 +8,4 @@ set -xe
: ${XML_LOG_DIR:=/logs} : ${XML_LOG_DIR:=/logs}
mkdir -p "$XML_LOG_DIR" mkdir -p "$XML_LOG_DIR"
python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_* NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
qa/L1_pytorch_distributed_unittest/test.sh
View file @ 87e3e56e
...@@ -21,14 +21,20 @@ FAILED_CASES="" ...@@ -21,14 +21,20 @@ FAILED_CASES=""
mkdir -p "$XML_LOG_DIR" mkdir -p "$XML_LOG_DIR"
# It is not installed as a requirement,
# because it is not available on PyPI.
pip uninstall -y nvdlfw-inspect
pip install git+https://github.com/NVIDIA/nvidia-dlfw-inspect.git
pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py || test_fail "test_torch_fsdp2.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py || test_fail "test_torch_fsdp2.py"
python3 -m pytest -v -s --log-cli-level=INFO --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py" python3 -m pytest -v -s --log-cli-level=INFO --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fused_attn_with_cp.xml $TE_PATH/tests/pytorch/fused_attn/test_fused_attn_with_cp.py || test_fail "test_fused_attn_with_cp.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py"
python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py"
......
qa/L2_jax_distributed_unittest/test.sh 0 → 100644
View file @ 87e3e56e
# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.
set -xe
: ${TE_PATH:=/opt/transformerengine}
: ${XML_LOG_DIR:=/logs}
mkdir -p "$XML_LOG_DIR"
NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_*
qa/L2_jax_unittest/test.sh
View file @ 87e3e56e
...@@ -36,7 +36,7 @@ export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" ...@@ -36,7 +36,7 @@ export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops"
NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py"
# Test without custom calls # Test without custom calls
export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops"
NVTE_JAX_CUSTOM_CALLS_RE="" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py" NVTE_JAX_CUSTOM_CALLS="false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py"
if [ $RET -ne 0 ]; then if [ $RET -ne 0 ]; then
echo "Error: some sub-tests failed: $FAILED_CASES" echo "Error: some sub-tests failed: $FAILED_CASES"
......
qa/L3_pytorch_FA_versions_test/test.sh
View file @ 87e3e56e
...@@ -41,6 +41,6 @@ do ...@@ -41,6 +41,6 @@ do
fi fi
# Run tests # Run tests
NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/pytorch/fused_attn/test_fused_attn.py NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/pytorch/attention/test_attention.py
done done
tests/cpp/operator/CMakeLists.txt
View file @ 87e3e56e
...@@ -28,6 +28,7 @@ list(APPEND test_cuda_sources ...@@ -28,6 +28,7 @@ list(APPEND test_cuda_sources
test_multi_unpadding.cu test_multi_unpadding.cu
test_causal_softmax.cu test_causal_softmax.cu
test_swizzle.cu test_swizzle.cu
test_swap_first_dims.cu
../test_common.cu) ../test_common.cu)
if(USE_ROCM) if(USE_ROCM)
list(APPEND test_cuda_sources list(APPEND test_cuda_sources
......
tests/cpp/operator/test_cast_mxfp8.cu
View file @ 87e3e56e
This diff is collapsed.
tests/cpp/operator/test_swap_first_dims.cu 0 → 100644
View file @ 87e3e56e
/*************************************************************************
* Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
*
* See LICENSE for license information.
************************************************************************/
#include <memory>
#include <string>
#include <tuple>
#include <vector>
#include <cuda_runtime.h>
#include <gtest/gtest.h>
#include <transformer_engine/transpose.h>
#include "../test_common.h"
using namespace transformer_engine;
namespace {
template <typename Type>
void compute_ref(const Type *input, Type *output,
const std::vector<size_t> &shape) {
const size_t dim0 = shape[0];
const size_t dim1 = shape[1];
size_t dim2 = 1;
for (size_t i = 2; i < shape.size(); ++i) {
dim2 *= shape[i];
}
for (size_t i = 0; i < dim0; ++i) {
for (size_t j = 0; j < dim1; ++j) {
for (size_t k = 0; k < dim2; ++k) {
const size_t in_offset = i * dim1 * dim2 + j * dim2 + k;
const size_t out_offset = j * dim0 * dim2 + i * dim2 + k;
output[out_offset] = input[in_offset];
}
}
}
}
template <typename Type>
void performTest(const std::vector<size_t> &in_shape) {
using namespace test;
DType dtype = TypeInfo<Type>::dtype;
// Tensor dimensions
std::vector<size_t> out_shape = in_shape;
out_shape[0] = in_shape[1];
out_shape[1] = in_shape[0];
size_t numel = 1;
for (const auto& dim : in_shape) {
numel *= dim;
}
// Transformer engine implementation
Tensor input("input", in_shape, dtype);
Tensor output("output", out_shape, dtype);
fillUniform(&input);
nvte_swap_first_dims(input.data(), output.data(), 0);
// Reference implementation
std::unique_ptr<Type[]> ref_output = std::make_unique<Type[]>(numel);
compute_ref<Type>(input.rowwise_cpu_dptr<Type>(), ref_output.get(), in_shape);
// Check for CUDA failure
cudaDeviceSynchronize();
auto err = cudaGetLastError();
ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err);
// Check for exact numerics
compareResults("output", output, ref_output.get(), true, 0, 0);
}
std::vector<std::vector<size_t>> test_cases = {{4, 64, 1280},
{48, 8, 128, 16},
{229, 173}, // Primes 50, 40
{113, 71, 1, 1, 1, 29, 1, 1}}; // Primes 30, 20, 10
} // namespace
class SwapFirstDimsTestSuite : public ::testing::TestWithParam<std::tuple<transformer_engine::DType,
std::vector<size_t>>> {};
TEST_P(SwapFirstDimsTestSuite, TestSwapFirstDims) {
using namespace transformer_engine;
using namespace test;
const DType type = std::get<0>(GetParam());
const auto shape = std::get<1>(GetParam());
TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(type, T,
performTest<T>(shape);
);
}
INSTANTIATE_TEST_SUITE_P(
OperatorTest,
SwapFirstDimsTestSuite,
::testing::Combine(
::testing::ValuesIn(test::all_fp_types),
::testing::ValuesIn(test_cases)),
[](const testing::TestParamInfo<SwapFirstDimsTestSuite::ParamType>& info) {
std::string name = test::typeName(std::get<0>(info.param));
for (const auto& dim : std::get<1>(info.param)) {
name += "X";
name += std::to_string(dim);
}
return name;
});
tests/cpp/test_common.cu
View file @ 87e3e56e
...@@ -523,10 +523,13 @@ std::vector<size_t> unravel(const size_t i, const NVTEShape &shape) { ...@@ -523,10 +523,13 @@ std::vector<size_t> unravel(const size_t i, const NVTEShape &shape) {
void compareResults_sequential(const std::string &name, const Tensor &test, void compareResults_sequential(const std::string &name, const Tensor &test,
const void *ref, const bool rowwise, const void *ref, const bool rowwise,
double atol, double rtol, bool if_on_gpus) { double atol, double rtol, bool if_on_gpus,
const size_t tolerable_mismatches_limit) {
if (if_on_gpus) test.to_cpu(); if (if_on_gpus) test.to_cpu();
const auto& shape = rowwise ? test.rowwise_shape() : test.columnwise_shape(); const auto& shape = rowwise ? test.rowwise_shape() : test.columnwise_shape();
const size_t N = product(shape); const size_t N = product(shape);
size_t mismatches_num = 0;
int first_mismatch_idx = -1;
TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(test.dtype(), T, TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(test.dtype(), T,
const T *test_data = rowwise ? test.rowwise_cpu_dptr<T>() : test.columnwise_cpu_dptr<T>(); const T *test_data = rowwise ? test.rowwise_cpu_dptr<T>() : test.columnwise_cpu_dptr<T>();
const T *ref_data = reinterpret_cast<const T*>(ref); const T *ref_data = reinterpret_cast<const T*>(ref);
...@@ -562,27 +565,39 @@ void compareResults_sequential(const std::string &name, const Tensor &test, ...@@ -562,27 +565,39 @@ void compareResults_sequential(const std::string &name, const Tensor &test,
#endif #endif
} }
std::string direction = rowwise ? "rowwise" : "columnwise"; std::string direction = rowwise ? "rowwise" : "columnwise";
ASSERT_FALSE(assertion) << "Error in tensor " << name << " in " if (assertion) {
<< direction << " direction." << std::endl mismatches_num++;
<< "Mismatch at place " << to_string(unravel(i, shape)) if (first_mismatch_idx == -1) {
<< " (" << std::to_string(i) << "): " << t << " vs " << r; first_mismatch_idx = i;
}
}
if (mismatches_num > tolerable_mismatches_limit) {
const double first_mismatch_t = static_cast<double>(test_data[first_mismatch_idx]);
const double first_mismatch_r = static_cast<double>(ref_data[first_mismatch_idx]);
GTEST_FAIL() << mismatches_num << " mismatche(s) which is more than tolerable mismatch limit of "
<< tolerable_mismatches_limit << "." << std::endl
<< "Error in tensor " << name << " in "
<< direction << " direction." << std::endl
<< "First mismatch at place " << to_string(unravel(first_mismatch_idx, shape))
<< " (" << std::to_string(first_mismatch_idx) << "): "
<< first_mismatch_t << " vs " << first_mismatch_r;
}
} }
); );
} }
template <typename T> template <typename T>
static size_t getFirstMismatchIdx(const DType data_type, const T* test_data, const T* ref_data, static size_t getFirstMismatchIdx(const DType data_type, const T* test_data, const T* ref_data,
const size_t N, const double atol, const double rtol) { const size_t N, const double atol, const double rtol,
size_t& mismatches) {
int first_mismatch_idx = N; int first_mismatch_idx = N;
bool is_mismatch_found = false; #pragma omp parallel reduction(min: first_mismatch_idx) reduction(+: mismatches) proc_bind(spread)
#pragma omp parallel for schedule(static) firstprivate(is_mismatch_found) \ {
reduction(min: first_mismatch_idx) proc_bind(spread) size_t thread_mismatches = 0;
for (size_t i = 0; i < N; ++i) { #pragma omp for schedule(static)
if (is_mismatch_found) { // early escape of the omp thread for (size_t i = 0; i < N; ++i) {
continue;
}
#ifndef __HIP_PLATFORM_AMD__ #ifndef __HIP_PLATFORM_AMD__
double t = static_cast<double>(test_data[i]); double t = static_cast<double>(test_data[i]);
double r = static_cast<double>(ref_data[i]); double r = static_cast<double>(ref_data[i]);
...@@ -590,17 +605,16 @@ static size_t getFirstMismatchIdx(const DType data_type, const T* test_data, con ...@@ -590,17 +605,16 @@ static size_t getFirstMismatchIdx(const DType data_type, const T* test_data, con
double t = static_cast<double>(static_cast<float>(test_data[i])); double t = static_cast<double>(static_cast<float>(test_data[i]));
double r = static_cast<double>(static_cast<float>(ref_data[i])); double r = static_cast<double>(static_cast<float>(ref_data[i]));
#endif #endif
bool mismatch = fabs(t - r) > atol && (r == 0 || fabs((t - r) / r) > rtol);
bool mismatch = fabs(t - r) > atol && (r == 0 || fabs((t - r) / r) > rtol); /* For Float32 the floating point comparison is enough to error out */
/* For Float32 the floating point comparison is enough to error out */ bool assertion = mismatch && (data_type == DType::kFloat32);
bool assertion = mismatch && (data_type == DType::kFloat32); if (mismatch && !assertion) {
if (mismatch && !assertion) { /* Check if it is just a failure of round to nearest choosing different
/* Check if it is just a failure of round to nearest choosing different side of the real value */
side of the real value */ const double mean = (t + r) / 2;
const double mean = (t + r) / 2; const double mean_p = mean >= 0 ? mean * (1 + 1e-6) : mean * (1 - 1e-6);
const double mean_p = mean >= 0 ? mean * (1 + 1e-6) : mean * (1 - 1e-6); const double mean_m = mean >= 0 ? mean * (1 - 1e-6) : mean * (1 + 1e-6);
const double mean_m = mean >= 0 ? mean * (1 - 1e-6) : mean * (1 + 1e-6);
#ifndef __HIP_PLATFORM_AMD__ #ifndef __HIP_PLATFORM_AMD__
const double cast_mean_p = static_cast<double>(static_cast<T>(mean_p)); const double cast_mean_p = static_cast<double>(static_cast<T>(mean_p));
const double cast_mean_m = static_cast<double>(static_cast<T>(mean_m)); const double cast_mean_m = static_cast<double>(static_cast<T>(mean_m));
...@@ -608,32 +622,37 @@ static size_t getFirstMismatchIdx(const DType data_type, const T* test_data, con ...@@ -608,32 +622,37 @@ static size_t getFirstMismatchIdx(const DType data_type, const T* test_data, con
const double cast_mean_p = static_cast<double>(static_cast<float>(static_cast<T>(static_cast<float>(mean_p)))); const double cast_mean_p = static_cast<double>(static_cast<float>(static_cast<T>(static_cast<float>(mean_p))));
const double cast_mean_m = static_cast<double>(static_cast<float>(static_cast<T>(static_cast<float>(mean_m)))); const double cast_mean_m = static_cast<double>(static_cast<float>(static_cast<T>(static_cast<float>(mean_m))));
#endif #endif
#ifdef __HIP_PLATFORM_AMD__ #ifdef __HIP_PLATFORM_AMD__
assertion = !(cast_mean_m == std::min<double>(t,r) && cast_mean_p == std::max<double>(t,r)); assertion = !(cast_mean_m == std::min<double>(t,r) && cast_mean_p == std::max<double>(t,r));
#else #else
assertion = !(cast_mean_m == std::min(t,r) && cast_mean_p == std::max(t,r)); assertion = !(cast_mean_m == std::min(t,r) && cast_mean_p == std::max(t,r));
#endif #endif
}
if (assertion) {
if (i < first_mismatch_idx) {
first_mismatch_idx = i;
}
thread_mismatches++;
}
} }
if (assertion && i < first_mismatch_idx) { mismatches += thread_mismatches;
first_mismatch_idx = i;
is_mismatch_found = true;
}
} }
return first_mismatch_idx; return first_mismatch_idx;
} }
void compareResults_parallel(const std::string &name, const Tensor &test, const void *ref, void compareResults_parallel(const std::string &name, const Tensor &test, const void *ref,
const bool rowwise, double atol, double rtol, bool if_on_gpus) { const bool rowwise, double atol, double rtol, bool if_on_gpus,
const size_t tolerable_mismatches_limit) {
if (if_on_gpus) test.to_cpu(); if (if_on_gpus) test.to_cpu();
const auto& shape = rowwise ? test.rowwise_shape() : test.columnwise_shape(); const auto& shape = rowwise ? test.rowwise_shape() : test.columnwise_shape();
const size_t N = product(shape); const size_t N = product(shape);
size_t mismatches = 0;
TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(test.dtype(), T, TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(test.dtype(), T,
const T *test_data = rowwise ? test.rowwise_cpu_dptr<T>() : test.columnwise_cpu_dptr<T>(); const T *test_data = rowwise ? test.rowwise_cpu_dptr<T>() : test.columnwise_cpu_dptr<T>();
const T *ref_data = reinterpret_cast<const T*>(ref); const T *ref_data = reinterpret_cast<const T*>(ref);
const size_t i = getFirstMismatchIdx<T>(test.dtype(), test_data, ref_data, N, atol, rtol); const size_t i = getFirstMismatchIdx<T>(test.dtype(), test_data, ref_data, N, atol, rtol, mismatches);
if (i != N) { if ((i != N) && (mismatches > tolerable_mismatches_limit)) {
#ifndef __HIP_PLATFORM_AMD__ #ifndef __HIP_PLATFORM_AMD__
const double t = static_cast<double>(test_data[i]); const double t = static_cast<double>(test_data[i]);
const double r = static_cast<double>(ref_data[i]); const double r = static_cast<double>(ref_data[i]);
...@@ -642,21 +661,25 @@ void compareResults_parallel(const std::string &name, const Tensor &test, const ...@@ -642,21 +661,25 @@ void compareResults_parallel(const std::string &name, const Tensor &test, const
const double r = static_cast<double>(static_cast<float>(ref_data[i])); const double r = static_cast<double>(static_cast<float>(ref_data[i]));
#endif #endif
std::string direction = rowwise ? "rowwise" : "columnwise"; std::string direction = rowwise ? "rowwise" : "columnwise";
ASSERT_FALSE(true) << "Error in tensor " << name << " in "
<< direction << " direction." << std::endl GTEST_FAIL() << mismatches << " mismatche(s) which is more than tolerable mismatch limit of "
<< "Mismatch at place " << to_string(unravel(i, shape)) << tolerable_mismatches_limit << "." << std::endl
<< " (" << std::to_string(i) << "): " << t << " vs " << r; << "Error in tensor " << name << " in "
<< direction << " direction." << std::endl
<< "Mismatch at place " << to_string(unravel(i, shape))
<< " (" << std::to_string(i) << "): " << t << " vs " << r;
} }
); );
} }
void compareResults(const std::string &name, const Tensor &test, const void *ref, void compareResults(const std::string &name, const Tensor &test, const void *ref,
const bool rowwise, double atol, double rtol, bool if_on_gpus) { const bool rowwise, double atol, double rtol, bool if_on_gpus,
const size_t tolerable_mismatches_limit) {
constexpr bool sequential = false; constexpr bool sequential = false;
if constexpr (sequential) { if constexpr (sequential) {
compareResults_sequential(name, test, ref, rowwise, atol, rtol, if_on_gpus); compareResults_sequential(name, test, ref, rowwise, atol, rtol, if_on_gpus, tolerable_mismatches_limit);
} else { } else {
compareResults_parallel(name, test, ref, rowwise, atol, rtol, if_on_gpus); compareResults_parallel(name, test, ref, rowwise, atol, rtol, if_on_gpus, tolerable_mismatches_limit);
} }
} }
...@@ -698,25 +721,39 @@ void compareResults(const std::string &name, const uint8_t *test, const uint8_t ...@@ -698,25 +721,39 @@ void compareResults(const std::string &name, const uint8_t *test, const uint8_t
} }
void compare_e8m0_scaling_factors(const std::string &name, const uint8_t *test, const uint8_t *ref, void compare_e8m0_scaling_factors(const std::string &name, const uint8_t *test, const uint8_t *ref,
const size_t row_blocks, const size_t col_blocks, const size_t stride) const size_t row_blocks, const size_t col_blocks, const size_t stride,
size_t& mismatches_num, const size_t atol,
const double abs_tolerable_mismatches_limit,
const double rel_tolerable_mismatches_limit)
{ {
const size_t N = row_blocks * col_blocks;
const size_t tolerable_mismatches_limit = std::min(abs_tolerable_mismatches_limit,
std::floor(N * rel_tolerable_mismatches_limit));
mismatches_num = 0;
std::vector<int> mismatch_indices;
for (int i = 0; i < row_blocks; ++i) { for (int i = 0; i < row_blocks; ++i) {
for (int j = 0; j < col_blocks; ++j) { for (int j = 0; j < col_blocks; ++j) {
const int idx = i * stride + j; const int idx = i * stride + j;
ASSERT_FALSE(test[idx] != ref[idx]) << "Error in " << name << std::endl const int test_val = static_cast<int>(test[idx]);
<< "Mismatch: " << static_cast<int>(test[idx]) << " vs " const int ref_val = static_cast<int>(ref[idx]);
<< static_cast<int>(ref[idx]) << " at index " << idx; const int abs_delta = std::abs(test_val - ref_val);
}
}
}
void compare_e8m0_scaling_factors(const std::string &name, const uint8_t *test, const uint8_t *ref, if (abs_delta > atol) {
const size_t N) mismatches_num++;
{ mismatch_indices.push_back(idx);
for (int i = 0; i < N; i++) { }
ASSERT_FALSE(test[i] != ref[i]) << "Error in " << name << std::endl if (mismatches_num > tolerable_mismatches_limit) {
<< "Mismatch: " << static_cast<int>(test[i]) << " vs " std::cout << "Error in " << name << std::endl;
<< static_cast<int>(ref[i]) << " at index " << i; for (const int index : mismatch_indices) {
std::cout << "Mismatch at (" << index << "):"
<< static_cast<int>(test[index]) << " vs "
<< static_cast<int>(ref[index]) << std::endl;
}
GTEST_FAIL() << mismatches_num << " mismatche(s) which is more than tolerable mismatch limit of "
<< tolerable_mismatches_limit << ".";
}
}
} }
} }
......
tests/cpp/test_common.h
View file @ 87e3e56e
...@@ -430,7 +430,12 @@ inline fp8e8m0 float_to_e8m0(float val) { ...@@ -430,7 +430,12 @@ inline fp8e8m0 float_to_e8m0(float val) {
} }
inline float exp2f_rcp(fp8e8m0 biased_exp) { inline float exp2f_rcp(fp8e8m0 biased_exp) {
return (biased_exp == 0) ? 1 : exp2f(FP32_EXPONENT_BIAS - static_cast<float>(biased_exp)); if (biased_exp == 0) {
return 1.0f;
}
int32_t int_val = (254 - biased_exp) << FP32_MANTISSA_BITS; // 127 - (biased_exp - 127)
float fp32_val = *reinterpret_cast<float*>(&int_val);
return fp32_val;
} }
inline float identity(const float x) { return x; } inline float identity(const float x) { return x; }
...@@ -462,15 +467,18 @@ size_t last_dimension(const std::vector<size_t> &shape); ...@@ -462,15 +467,18 @@ size_t last_dimension(const std::vector<size_t> &shape);
bool areShapesEqual(const NVTEShape &s1, const NVTEShape &s2); bool areShapesEqual(const NVTEShape &s1, const NVTEShape &s2);
void compareResults(const std::string &name, const Tensor &test, const void *ref, void compareResults(const std::string &name, const Tensor &test, const void *ref,
bool rowwise, double atol = 1e-5, double rtol = 1e-8, bool if_on_gpus = true); bool rowwise, double atol = 1e-5, double rtol = 1e-8, bool if_on_gpus = true,
const size_t tolerable_mismatches_limit = 0);
void compareResults(const std::string &name, const float test, const float ref, void compareResults(const std::string &name, const float test, const float ref,
double atol = 1e-5, double rtol = 1e-8); double atol = 1e-5, double rtol = 1e-8);
void compareResults(const std::string &name, const uint8_t *test, const uint8_t *ref, void compareResults(const std::string &name, const uint8_t *test, const uint8_t *ref,
size_t N, float mismatch_rate_tol = 0.); size_t N, float mismatch_rate_tol = 0.);
void compare_e8m0_scaling_factors(const std::string &name, const uint8_t *test, const uint8_t *ref, void compare_e8m0_scaling_factors(const std::string &name, const uint8_t *test, const uint8_t *ref,
const size_t row_blocks, const size_t col_blocks, const size_t stride); const size_t row_blocks, const size_t col_blocks, const size_t stride,
void compare_e8m0_scaling_factors(const std::string &name, const uint8_t *test, const uint8_t *ref, size_t& mismatches_num,
const size_t N); const size_t scale_diff_abs_tolerance = 0,
const double abs_tolerable_mismatches_limit = 0,
const double rel_tolerable_mismatches_limit = 0);
std::array<size_t, 4> get_scale_tensor_dims(const size_t rows, const size_t cols, std::array<size_t, 4> get_scale_tensor_dims(const size_t rows, const size_t cols,
const size_t block_size_rows, const size_t block_size_cols); const size_t block_size_rows, const size_t block_size_cols);
......
tests/jax/conftest.py
View file @ 87e3e56e
...@@ -5,6 +5,8 @@ ...@@ -5,6 +5,8 @@
import os import os
import jax import jax
import pytest import pytest
from collections import defaultdict
import time
import transformer_engine.jax import transformer_engine.jax
...@@ -32,3 +34,54 @@ def enable_fused_attn_after_hopper(): ...@@ -32,3 +34,54 @@ def enable_fused_attn_after_hopper():
yield yield
if "NVTE_FUSED_ATTN" in os.environ: if "NVTE_FUSED_ATTN" in os.environ:
del os.environ["NVTE_FUSED_ATTN"] del os.environ["NVTE_FUSED_ATTN"]
class TestTimingPlugin:
"""
Plugin to measure test execution time. Enable test timing by setting NVTE_JAX_TEST_TIMING=1
in the environment.
"""
def __init__(self):
self.test_timings = defaultdict(list)
@pytest.hookimpl(tryfirst=True)
def pytest_runtest_setup(self, item):
item._timing_start = time.time()
@pytest.hookimpl(trylast=True)
def pytest_runtest_teardown(self, item, nextitem):
if hasattr(item, "_timing_start"):
duration = time.time() - item._timing_start
# Extract base function name without parameters
test_name = item.name
if "[" in test_name:
base_name = test_name.split("[")[0]
else:
base_name = test_name
self.test_timings[base_name].append(duration)
def pytest_sessionfinish(self, session, exitstatus):
print("\n" + "=" * 80)
print("TEST RUNTIME SUMMARY (grouped by function)")
print("=" * 80)
total_overall = 0
for test_name, durations in sorted(self.test_timings.items()):
total_time = sum(durations)
count = len(durations)
avg_time = total_time / count if count > 0 else 0
total_overall += total_time
print(f"{test_name:<60} | {count:3}x | {total_time:7.2f}s | avg: {avg_time:6.2f}s")
print("=" * 80)
print(f"{'TOTAL RUNTIME':<60} | {'':>3} | {total_overall:7.2f}s |")
print("=" * 80)
def pytest_configure(config):
if os.getenv("NVTE_JAX_TEST_TIMING", "0") == "1":
config.pluginmanager.register(TestTimingPlugin(), "test_timing")
tests/jax/distributed_test_base.py
View file @ 87e3e56e
...@@ -39,8 +39,10 @@ def generate_configs(): ...@@ -39,8 +39,10 @@ def generate_configs():
return configs return configs
def generate_context_parallel_configs(): def generate_context_parallel_configs_for_attn():
configs = [] """Generate CP combinations along with TP+DP for TestDistributedContextParallelSelfAttn only"""
configsL1 = []
configsL2 = []
mr = MeshResource(dp_resource="dp", cp_resource="cp", tp_resource="tp") mr = MeshResource(dp_resource="dp", cp_resource="cp", tp_resource="tp")
axes = ("dp", "cp", "tp") axes = ("dp", "cp", "tp")
DP_sizes = (1, 2) DP_sizes = (1, 2)
...@@ -49,10 +51,16 @@ def generate_context_parallel_configs(): ...@@ -49,10 +51,16 @@ def generate_context_parallel_configs():
for dp, cp, tp in product(DP_sizes, CP_sizes, TP_sizes): for dp, cp, tp in product(DP_sizes, CP_sizes, TP_sizes):
ndev = cp * tp * dp ndev = cp * tp * dp
if is_devices_enough(ndev): if is_devices_enough(ndev):
configs.append( # Do not run cp1 case in L1 as that is already covered in TestDistributedSelfAttn and TestDistributedCrossAttn (as these do not have any cp combinations)
pytest.param(ndev, (dp, cp, tp), axes, mr, id=f"n{ndev}_dp{dp}_cp{cp}_tp{tp}") if cp != 1:
) configsL1.append(
pytest.param(ndev, (dp, cp, tp), axes, mr, id=f"n{ndev}_dp{dp}_cp{cp}_tp{tp}")
)
else:
configsL2.append(
pytest.param(ndev, (dp, cp, tp), axes, mr, id=f"n{ndev}_dp{dp}_cp{cp}_tp{tp}")
)
configs = {"L0": [], "L1": configsL1, "L2": configsL2}
return configs return configs
......
tests/jax/test_custom_call_compute.py
View file @ 87e3e56e
...@@ -78,8 +78,14 @@ def is_shape_supported_by_mxfp8(input_shape): ...@@ -78,8 +78,14 @@ def is_shape_supported_by_mxfp8(input_shape):
return False return False
def assert_bitwise_scaled_tensors(a: ScaledTensor, b: ScaledTensor): def assert_bitwise_scaled_tensors(
a: ScaledTensor, b: ScaledTensor, precise_comparison: bool = True
):
if isinstance(a, ScaledTensor1x) and isinstance(b, ScaledTensor1x): if isinstance(a, ScaledTensor1x) and isinstance(b, ScaledTensor1x):
if not precise_comparison:
assert_allclose(a.dequantize(), b.dequantize(), dtype=a.data.dtype)
return
assert a.scaling_mode == b.scaling_mode assert a.scaling_mode == b.scaling_mode
assert a.scale_inv.dtype == b.scale_inv.dtype assert a.scale_inv.dtype == b.scale_inv.dtype
if a.scaling_mode.is_tensor_scaling(): if a.scaling_mode.is_tensor_scaling():
...@@ -94,8 +100,12 @@ def assert_bitwise_scaled_tensors(a: ScaledTensor, b: ScaledTensor): ...@@ -94,8 +100,12 @@ def assert_bitwise_scaled_tensors(a: ScaledTensor, b: ScaledTensor):
assert_allclose(a.data, b.data) assert_allclose(a.data, b.data)
elif isinstance(a, ScaledTensor2x) and isinstance(b, ScaledTensor2x): elif isinstance(a, ScaledTensor2x) and isinstance(b, ScaledTensor2x):
assert_bitwise_scaled_tensors(a.rowwise_tensor, b.rowwise_tensor) assert_bitwise_scaled_tensors(
assert_bitwise_scaled_tensors(a.colwise_tensor, b.colwise_tensor) a.rowwise_tensor, b.rowwise_tensor, precise_comparison=precise_comparison
)
assert_bitwise_scaled_tensors(
a.colwise_tensor, b.colwise_tensor, precise_comparison=precise_comparison
)
else: else:
pytest.fail("Unsupported input types") pytest.fail("Unsupported input types")
...@@ -481,24 +491,7 @@ class TestNorm: ...@@ -481,24 +491,7 @@ class TestNorm:
# if the input dtype is not float32 # if the input dtype is not float32
precise_comparison = False precise_comparison = False
if precise_comparison: assert_bitwise_scaled_tensors(output, ref_out, precise_comparison=precise_comparison)
assert_bitwise_scaled_tensors(output, ref_out)
else:
if isinstance(ref_out, ScaledTensor1x):
assert_allclose(output.dequantize(), ref_out.dequantize(), dtype=out_dtype)
elif isinstance(ref_out, ScaledTensor2x):
assert_allclose(
output.rowwise_tensor.dequantize(),
ref_out.rowwise_tensor.dequantize(),
dtype=out_dtype,
)
assert_allclose(
output.colwise_tensor.dequantize(),
ref_out.colwise_tensor.dequantize(),
dtype=out_dtype,
)
else:
pytest.fail("Unsupported output type")
assert_allclose(rsigma, ref_rsigma, dtype=inp_dtype) assert_allclose(rsigma, ref_rsigma, dtype=inp_dtype)
if norm_type == "layernorm": if norm_type == "layernorm":
...@@ -680,10 +673,6 @@ class TestGroupedQuantize: ...@@ -680,10 +673,6 @@ class TestGroupedQuantize:
n_groups=n_groups, n_groups=n_groups,
) )
# grouped_quantize does not work with cudaGraph yet, so the jitting will breaks
# To test it locally, export XLA_FLAGS="--xla_gpu_enable_command_buffer= $XLA_FLAGS" to
# disable cudaGraph, then use the following jitted function
scaled_tensor = tex.grouped_quantize( scaled_tensor = tex.grouped_quantize(
x, group_sizes=group_sizes, flatten_axis=flatten_axis, quantizer=grouped_quantizer x, group_sizes=group_sizes, flatten_axis=flatten_axis, quantizer=grouped_quantizer
) )
...@@ -768,12 +757,24 @@ class TestFusedQuantize: ...@@ -768,12 +757,24 @@ class TestFusedQuantize:
)(dz, x) )(dz, x)
if is_casted_output: if is_casted_output:
assert_bitwise_scaled_tensors(te_output, jax_output) # TE kernels cast the intermediate results to the input dtype which reduces precision compared to the JAX implementation
precise_comparison = not (
in_dtype != jnp.float32 and scaling_mode.is_1d_block_scaling()
)
assert_bitwise_scaled_tensors(
te_output, jax_output, precise_comparison=precise_comparison
)
else: else:
assert_allclose(te_output, jax_output) assert_allclose(te_output, jax_output)
if is_dbias: if is_dbias:
assert_allclose(te_dbias, jax_dbias) # TE kernels cast the intermediate results to the input dtype which reduces precision compared to the JAX implementation, for dbias this typically only affects bfloat16.
precise_comparison = not (
in_dtype == jnp.bfloat16 and scaling_mode.is_1d_block_scaling()
)
assert_allclose(
te_dbias, jax_dbias, dtype=in_dtype if precise_comparison else out_dtype
)
@pytest_parametrize_wrapper("activation_type", ACTIVATION_TYPES) @pytest_parametrize_wrapper("activation_type", ACTIVATION_TYPES)
@pytest_parametrize_wrapper("input_shape", ALL_ACTIVATION_SHAPES) @pytest_parametrize_wrapper("input_shape", ALL_ACTIVATION_SHAPES)
...@@ -858,15 +859,6 @@ valid_fp8_gemm_operand_types = [ ...@@ -858,15 +859,6 @@ valid_fp8_gemm_operand_types = [
] ]
def _use_jax_fp8_gemm(enabled=False):
import os
if enabled:
os.environ["NVTE_JAX_CUSTOM_CALLS_RE"] = "^(?!GemmPrimitive$).+$"
elif "NVTE_JAX_CUSTOM_CALLS_RE" in os.environ:
os.environ.pop("NVTE_JAX_CUSTOM_CALLS_RE")
class TestDense: class TestDense:
def _ref_gemm_with_jnp_dot(self, a, b, data_layout): def _ref_gemm_with_jnp_dot(self, a, b, data_layout):
if data_layout[0] == "T": if data_layout[0] == "T":
...@@ -1316,16 +1308,14 @@ class TestGroupedDense: ...@@ -1316,16 +1308,14 @@ class TestGroupedDense:
) )
ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims) ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims)
# grouped_gemm does not work with cudaGraph yet, so the jitting will breaks
# To test it locally, export XLA_FLAGS="--xla_gpu_enable_command_buffer= $XLA_FLAGS" to
# disable cudaGraph, then use the following jitted function
# jitting grouped_gemm # jitting grouped_gemm
# prim_out = jax.jit(tex.grouped_gemm, static_argnames=("contracting_dims",))( prim_out = jax.jit(tex.grouped_gemm, static_argnames=("contracting_dims",))(
# lhs, rhs, group_sizes, contracting_dims, lhs,
# ) rhs,
group_sizes,
contracting_dims,
)
prim_out = tex.grouped_gemm(lhs, rhs, group_sizes, contracting_dims)
self._assert_grouped_gemm_output(prim_out, group_sizes, ref_out, dtype) self._assert_grouped_gemm_output(prim_out, group_sizes, ref_out, dtype)
@pytest.mark.skipif(not is_fp8_supported, reason=fp8_unsupported_reason) @pytest.mark.skipif(not is_fp8_supported, reason=fp8_unsupported_reason)
...@@ -1354,12 +1344,7 @@ class TestGroupedDense: ...@@ -1354,12 +1344,7 @@ class TestGroupedDense:
) )
ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims) ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims)
# jitting grouped_gemm prim_out = jax.jit(tex.grouped_gemm, static_argnames=("contracting_dims",))(
# prim_out = jax.jit(tex.grouped_gemm, static_argnames=('contracting_dims',))(
# lhs, rhs, group_sizes, contracting_dims, quantizer_set=quantizer_set
# )
prim_out = tex.grouped_gemm(
lhs, rhs, group_sizes, contracting_dims, quantizer_set=quantizer_set lhs, rhs, group_sizes, contracting_dims, quantizer_set=quantizer_set
) )
...@@ -1395,9 +1380,9 @@ class TestGroupedDense: ...@@ -1395,9 +1380,9 @@ class TestGroupedDense:
value_n_grad_ref_func = value_and_grad(self._ref_sum_grouped_dense, (0, 1, 2)) value_n_grad_ref_func = value_and_grad(self._ref_sum_grouped_dense, (0, 1, 2))
# jitting the grouped_dense # jitting the grouped_dense
# value_n_grad_prim_func = jit(value_and_grad(self._primitive_sum_grouped_dense, (0, 1, 2)), value_n_grad_prim_func = jit(
# static_argnums=(4,)) value_and_grad(self._primitive_sum_grouped_dense, (0, 1, 2)), static_argnums=(4,)
value_n_grad_prim_func = value_and_grad(self._primitive_sum_grouped_dense, (0, 1, 2)) )
ref_out_sum, (ref_dgrad, ref_wgrad, ref_dbias) = value_n_grad_ref_func( ref_out_sum, (ref_dgrad, ref_wgrad, ref_dbias) = value_n_grad_ref_func(
x, kernel, bias, group_sizes, contracting_dims x, kernel, bias, group_sizes, contracting_dims
...@@ -1436,9 +1421,9 @@ class TestGroupedDense: ...@@ -1436,9 +1421,9 @@ class TestGroupedDense:
value_n_grad_ref_func = value_and_grad(self._ref_sum_grouped_dense, (0, 1, 2)) value_n_grad_ref_func = value_and_grad(self._ref_sum_grouped_dense, (0, 1, 2))
# jitting the grouped_dense # jitting the grouped_dense
# value_n_grad_prim_func = jit(value_and_grad(self._primitive_sum_grouped_dense, (0, 1, 2)), value_n_grad_prim_func = jit(
# static_argnums=(4,)) value_and_grad(self._primitive_sum_grouped_dense, (0, 1, 2)), static_argnums=(4,)
value_n_grad_prim_func = value_and_grad(self._primitive_sum_grouped_dense, (0, 1, 2)) )
ref_out_sum, (ref_dgrad, ref_wgrad, ref_dbias) = value_n_grad_ref_func( ref_out_sum, (ref_dgrad, ref_wgrad, ref_dbias) = value_n_grad_ref_func(
x, x,
......
tests/jax/test_distributed_fused_attn.py
View file @ 87e3e56e
...@@ -9,10 +9,11 @@ import jax.numpy as jnp ...@@ -9,10 +9,11 @@ import jax.numpy as jnp
from jax import random from jax import random
from distributed_test_base import ( from distributed_test_base import (
generate_configs, generate_configs,
generate_context_parallel_configs, generate_context_parallel_configs_for_attn,
generate_collectives_count, generate_collectives_count,
) )
from test_fused_attn import FusedAttnRunner, BiasShape, SeqDescFormat from test_fused_attn import FusedAttnRunner, BiasShape, SeqDescFormat
from utils import pytest_parametrize_wrapper
from transformer_engine.jax.attention import ( from transformer_engine.jax.attention import (
is_fused_attn_kernel_available, is_fused_attn_kernel_available,
AttnBiasType, AttnBiasType,
...@@ -28,6 +29,12 @@ from transformer_engine.jax.attention import ( ...@@ -28,6 +29,12 @@ from transformer_engine.jax.attention import (
DTYPES = [jnp.bfloat16] DTYPES = [jnp.bfloat16]
DISTRIBUTED_SELF_ATTN_DATA_SHAPES = {
"L0": [()],
"L1": [(32, 1024, 16, 128)],
"L2": [(32, 512, 12, 64)],
}
class TestDistributedSelfAttn: class TestDistributedSelfAttn:
...@@ -64,7 +71,6 @@ class TestDistributedSelfAttn: ...@@ -64,7 +71,6 @@ class TestDistributedSelfAttn:
jax.config.update("jax_use_shardy_partitioner", use_shardy) jax.config.update("jax_use_shardy_partitioner", use_shardy)
dropout_prob = 0.0 dropout_prob = 0.0
is_training = True is_training = True
batch, seqlen, num_head, hidden = data_shape batch, seqlen, num_head, hidden = data_shape
if not is_fused_attn_kernel_available( if not is_fused_attn_kernel_available(
...@@ -119,13 +125,7 @@ class TestDistributedSelfAttn: ...@@ -119,13 +125,7 @@ class TestDistributedSelfAttn:
runner.test_backward() runner.test_backward()
@pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs())
@pytest.mark.parametrize( @pytest_parametrize_wrapper("data_shape", DISTRIBUTED_SELF_ATTN_DATA_SHAPES)
"data_shape",
[
pytest.param((32, 512, 12, 64), id="32-512-12-64"),
pytest.param((32, 1024, 16, 128), id="32-1024-16-128"),
],
)
@pytest.mark.parametrize( @pytest.mark.parametrize(
"attn_bias_type, bias_shape", "attn_bias_type, bias_shape",
[ [
...@@ -193,6 +193,13 @@ class TestDistributedSelfAttn: ...@@ -193,6 +193,13 @@ class TestDistributedSelfAttn:
) )
DISTRIBUTED_CROSS_ATTN_DATA_SHAPES = {
"L0": [()],
"L1": [[32, 512, 16, 64]],
"L2": [[32, 128, 12, 64]],
}
class TestDistributedCrossAttn: class TestDistributedCrossAttn:
def generate_collectives_count_ref(self): def generate_collectives_count_ref(self):
...@@ -201,7 +208,7 @@ class TestDistributedCrossAttn: ...@@ -201,7 +208,7 @@ class TestDistributedCrossAttn:
return generate_collectives_count(allreduce=all_reduce_loss_bytes, allgather=0, other=0) return generate_collectives_count(allreduce=all_reduce_loss_bytes, allgather=0, other=0)
@pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs())
@pytest.mark.parametrize("data_shape", [[32, 128, 12, 64], [32, 512, 16, 64]]) @pytest_parametrize_wrapper("data_shape", DISTRIBUTED_CROSS_ATTN_DATA_SHAPES)
@pytest.mark.parametrize( @pytest.mark.parametrize(
"attn_mask_type", [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK] "attn_mask_type", [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK]
) )
...@@ -390,8 +397,9 @@ class TestDistributedContextParallelSelfAttn: ...@@ -390,8 +397,9 @@ class TestDistributedContextParallelSelfAttn:
runner.test_backward() runner.test_backward()
del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"]
@pytest.mark.parametrize( @pytest_parametrize_wrapper(
"device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs() "device_count,mesh_shape,mesh_axes,mesh_resource",
generate_context_parallel_configs_for_attn(),
) )
@pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1])
@pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")])
...@@ -426,8 +434,9 @@ class TestDistributedContextParallelSelfAttn: ...@@ -426,8 +434,9 @@ class TestDistributedContextParallelSelfAttn:
use_shardy=True, use_shardy=True,
) )
@pytest.mark.parametrize( @pytest_parametrize_wrapper(
"device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs() "device_count,mesh_shape,mesh_axes,mesh_resource",
generate_context_parallel_configs_for_attn(),
) )
@pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES)
@pytest.mark.parametrize("kv_groups", [1, 8]) @pytest.mark.parametrize("kv_groups", [1, 8])
...@@ -468,8 +477,9 @@ class TestDistributedContextParallelSelfAttn: ...@@ -468,8 +477,9 @@ class TestDistributedContextParallelSelfAttn:
use_shardy=False, use_shardy=False,
) )
@pytest.mark.parametrize( @pytest_parametrize_wrapper(
"device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs() "device_count,mesh_shape,mesh_axes,mesh_resource",
generate_context_parallel_configs_for_attn(),
) )
@pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES)
@pytest.mark.parametrize("kv_groups", [1, 8]) @pytest.mark.parametrize("kv_groups", [1, 8])
...@@ -532,8 +542,9 @@ class TestDistributedContextParallelSelfAttn: ...@@ -532,8 +542,9 @@ class TestDistributedContextParallelSelfAttn:
window_size=window_size, window_size=window_size,
) )
@pytest.mark.parametrize( @pytest_parametrize_wrapper(
"device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs() "device_count,mesh_shape,mesh_axes,mesh_resource",
generate_context_parallel_configs_for_attn(),
) )
@pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1])
@pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")])
...@@ -570,16 +581,16 @@ class TestDistributedContextParallelSelfAttn: ...@@ -570,16 +581,16 @@ class TestDistributedContextParallelSelfAttn:
) )
REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES = {
"L0": [[]],
"L1": [[3, 32, 8, 64]],
"L2": [[4, 32, 12, 32], [1, 16, 1, 1]],
}
class TestReorderCausalLoadBalancing: class TestReorderCausalLoadBalancing:
@pytest.mark.parametrize("cp_size", [2, 4, 8]) @pytest.mark.parametrize("cp_size", [2, 4, 8])
@pytest.mark.parametrize( @pytest_parametrize_wrapper("shape", REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES)
"shape",
[
pytest.param([1, 16, 1, 1], id="1-16-1-1"),
pytest.param([4, 32, 12, 32], id="4-32-12-32"),
pytest.param([3, 32, 8, 64], id="3-32-8-64"),
],
)
@pytest.mark.parametrize("qkv_format", [QKVFormat.BSHD, QKVFormat.SBHD]) @pytest.mark.parametrize("qkv_format", [QKVFormat.BSHD, QKVFormat.SBHD])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"reorder_strategy", "reorder_strategy",
......
tests/jax/test_distributed_layernorm.py
View file @ 87e3e56e
...@@ -25,6 +25,7 @@ DTYPES = [jnp.bfloat16, jnp.float32] ...@@ -25,6 +25,7 @@ DTYPES = [jnp.bfloat16, jnp.float32]
NORM_INPUT_SHAPES = { NORM_INPUT_SHAPES = {
"L0": [[64, 64]], "L0": [[64, 64]],
"L1": [[64, 64]],
"L2": [[64, 64]], "L2": [[64, 64]],
} }
......
tests/jax/test_distributed_layernorm_mlp.py
View file @ 87e3e56e
...@@ -333,7 +333,6 @@ class TestDistributedLayernormMLP: ...@@ -333,7 +333,6 @@ class TestDistributedLayernormMLP:
with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe): with fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe):
ln_mlp_single = LayerNormMLP( ln_mlp_single = LayerNormMLP(
layernorm_type=layernorm_type, layernorm_type=layernorm_type,
transpose_batch_sequence=False, # input: [batch, seqlen, hidden]
intermediate_dim=INTERMEDIATE, intermediate_dim=INTERMEDIATE,
activations=activation_type, activations=activation_type,
use_bias=use_bias, use_bias=use_bias,
...@@ -352,7 +351,6 @@ class TestDistributedLayernormMLP: ...@@ -352,7 +351,6 @@ class TestDistributedLayernormMLP:
): ):
ln_mlp_sharded = LayerNormMLP( ln_mlp_sharded = LayerNormMLP(
layernorm_type=layernorm_type, layernorm_type=layernorm_type,
transpose_batch_sequence=False,
intermediate_dim=INTERMEDIATE, intermediate_dim=INTERMEDIATE,
activations=activation_type, activations=activation_type,
scale_axes=LN_SCALE_AXES, scale_axes=LN_SCALE_AXES,
......
tests/jax/test_distributed_softmax.py
View file @ 87e3e56e
...@@ -135,7 +135,7 @@ class TestDistributedSoftmax: ...@@ -135,7 +135,7 @@ class TestDistributedSoftmax:
) )
@pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs())
@pytest.mark.parametrize("data_shape", [[32, 12, 128, 128], [64, 16, 1024, 1024]]) @pytest.mark.parametrize("data_shape", [[32, 12, 128, 128], [8, 8, 1024, 1024]])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"softmax_type", "softmax_type",
[SoftmaxType.SCALED, SoftmaxType.SCALED_MASKED, SoftmaxType.SCALED_UPPER_TRIANG_MASKED], [SoftmaxType.SCALED, SoftmaxType.SCALED_MASKED, SoftmaxType.SCALED_UPPER_TRIANG_MASKED],
...@@ -168,14 +168,14 @@ class TestDistributedSoftmax: ...@@ -168,14 +168,14 @@ class TestDistributedSoftmax:
dtype, dtype,
bad_sharding, bad_sharding,
broadcast_batch_mask, broadcast_batch_mask,
use_shardy=False, use_shardy=True,
) )
@pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs())
@pytest.mark.parametrize("softmax_type", [SoftmaxType.SCALED, SoftmaxType.SCALED_MASKED]) @pytest.mark.parametrize("softmax_type", [SoftmaxType.SCALED, SoftmaxType.SCALED_MASKED])
@pytest.mark.parametrize("bad_sharding", [False, True]) @pytest.mark.parametrize("bad_sharding", [False, True])
@pytest.mark.parametrize("broadcast_batch_mask", [False, True]) @pytest.mark.parametrize("broadcast_batch_mask", [False, True])
def test_softmax_shardy( def test_softmax_gspmd(
self, self,
device_count, device_count,
mesh_shape, mesh_shape,
...@@ -196,5 +196,5 @@ class TestDistributedSoftmax: ...@@ -196,5 +196,5 @@ class TestDistributedSoftmax:
dtype=DTYPES[0], dtype=DTYPES[0],
bad_sharding=bad_sharding, bad_sharding=bad_sharding,
broadcast_batch_mask=broadcast_batch_mask, broadcast_batch_mask=broadcast_batch_mask,
use_shardy=True, use_shardy=False,
) )
tests/jax/test_fused_attn.py
View file @ 87e3e56e
...@@ -372,7 +372,7 @@ class FusedAttnRunner: ...@@ -372,7 +372,7 @@ class FusedAttnRunner:
self.head_dim_v, self.head_dim_v,
(-1, -1) if self.window_size is None else self.window_size, (-1, -1) if self.window_size is None else self.window_size,
).get_fused_attn_backend() ).get_fused_attn_backend()
if self.backend == NVTE_Fused_Attn_Backend.NVTE_No_Backend: if self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen:
pytest.skip("Unsupported inputs combination or device compute capability.") pytest.skip("Unsupported inputs combination or device compute capability.")
if ( if (
......
tests/jax/test_helper.py
View file @ 87e3e56e
...@@ -58,7 +58,6 @@ class TestFP8Functions(unittest.TestCase): ...@@ -58,7 +58,6 @@ class TestFP8Functions(unittest.TestCase):
self.assertTrue(ref.amax_compute_algo == test.amax_compute_algo) self.assertTrue(ref.amax_compute_algo == test.amax_compute_algo)
def _compare_current_scaling(self, test): def _compare_current_scaling(self, test):
self.assertEqual(QuantizeConfig.MARGIN, test.margin)
self.assertEqual(QuantizeConfig.FP8_FORMAT, test.fp8_format) self.assertEqual(QuantizeConfig.FP8_FORMAT, test.fp8_format)
self.assertEqual(QuantizeConfig.SCALING_MODE, ScalingMode.CURRENT_TENSOR_SCALING) self.assertEqual(QuantizeConfig.SCALING_MODE, ScalingMode.CURRENT_TENSOR_SCALING)
...@@ -91,7 +90,7 @@ class TestFP8Functions(unittest.TestCase): ...@@ -91,7 +90,7 @@ class TestFP8Functions(unittest.TestCase):
self._check_default_state() self._check_default_state()
@unittest.skipIf(not is_mxfp8_supported, reason=mxfp8_reason) @unittest.skipIf(not is_fp8_supported, reason=reason)
def test_fp8_autocast_current_scaling(self): def test_fp8_autocast_current_scaling(self):
QuantizeConfig.finalize() # Ensure the testing not affect by previous tests. QuantizeConfig.finalize() # Ensure the testing not affect by previous tests.
self._check_default_state() self._check_default_state()
...@@ -101,14 +100,14 @@ class TestFP8Functions(unittest.TestCase): ...@@ -101,14 +100,14 @@ class TestFP8Functions(unittest.TestCase):
self._check_default_state() self._check_default_state()
cs = Float8CurrentScaling(margin=5.0, fp8_format=FP8Format.E4M3) cs = Float8CurrentScaling(fp8_format=FP8Format.E4M3)
with fp8_autocast(enabled=True, fp8_recipe=cs): with fp8_autocast(enabled=True, fp8_recipe=cs):
self.assertTrue(QuantizeConfig.is_fp8_enabled()) self.assertTrue(QuantizeConfig.is_fp8_enabled())
self._compare_current_scaling(cs) self._compare_current_scaling(cs)
self._check_default_state() self._check_default_state()
cs = Float8CurrentScaling(margin=3.0, fp8_format=FP8Format.HYBRID) cs = Float8CurrentScaling(fp8_format=FP8Format.HYBRID)
with fp8_autocast(enabled=True, fp8_recipe=cs): with fp8_autocast(enabled=True, fp8_recipe=cs):
self.assertTrue(QuantizeConfig.is_fp8_enabled()) self.assertTrue(QuantizeConfig.is_fp8_enabled())
self._compare_current_scaling(cs) self._compare_current_scaling(cs)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment