From eb604b0e1edb68e27a7dcc96335c74ea78a1909a Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 17 Mar 2026 21:44:23 -0700 Subject: [PATCH 01/18] Update cudnnFE to v1.20.0 (#2774) Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Varun Thumbe --- 3rdparty/cudnn-frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 8d19d3182b..d33027a41a 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 8d19d3182bfbc304046a15e9236bec9ff31511fc +Subproject commit d33027a41a93af9c85f089c6364ab415fce98982 From 0608bdee9158dc3c5878fe442797154727853304 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 2 Mar 2026 03:10:24 +0000 Subject: [PATCH 02/18] fix merge conflicts, now things working Signed-off-by: Varun Thumbe --- 3rdparty/cudnn-frontend | 2 +- tests/cpp/operator/CMakeLists.txt | 1 + .../operator/test_multi_tensor_adam_mxfp8.cu | 266 ++++++++++++ tests/cpp/test_common.h | 10 + .../distributed/run_fsdp2_fused_adam.py | 8 +- tests/pytorch/distributed/test_torch_fsdp2.py | 5 - .../include/transformer_engine/multi_tensor.h | 35 ++ .../common/multi_tensor/adam.cu | 384 +++++++++++++++--- .../multi_tensor/multi_tensor_apply.cuh | 94 +++++ transformer_engine/pytorch/csrc/extensions.h | 6 + .../csrc/extensions/multi_tensor/adam.cpp | 20 + .../pytorch/csrc/extensions/pybind.cpp | 2 + .../pytorch/optimizers/fused_adam.py | 45 +- 13 files changed, 812 insertions(+), 66 deletions(-) create mode 100644 tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index d33027a41a..8d19d3182b 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit d33027a41a93af9c85f089c6364ab415fce98982 +Subproject commit 8d19d3182bfbc304046a15e9236bec9ff31511fc diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 5e73675f4f..4241ada3ba 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -27,6 +27,7 @@ add_executable(test_operator test_memset.cu test_splits_to_offsets.cu test_multi_cast_transpose.cu + test_multi_tensor_adam_mxfp8.cu test_multi_padding.cu test_multi_unpadding.cu test_causal_softmax.cu diff --git a/tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu b/tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu new file mode 100644 index 0000000000..470917580f --- /dev/null +++ b/tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu @@ -0,0 +1,266 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include + +#include "../test_common.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +uint8_t fp8_to_u8(fp8e4m3 v) { + uint8_t out = 0; + std::memcpy(&out, &v, sizeof(uint8_t)); + return out; +} + +uint8_t fp8_to_u8(fp8e5m2 v) { + uint8_t out = 0; + std::memcpy(&out, &v, sizeof(uint8_t)); + return out; +} + +void run_mxfp8_adam_test(DType fp8_dtype) { + const std::vector shape1{64, 128}; + const std::vector shape2{32, 64}; + const float lr = 1e-3f; + const float beta1 = 0.9f; + const float beta2 = 0.999f; + const float eps = 1e-8f; + const int step = 1; + const int mode = 1; + const int bias_correction = 1; + const float weight_decay = 0.0f; + + // Run with 25 tensors > 24[MXFP8_MAX_TENSORS] to check + // the chunking logic + const size_t tensor_count = 25; + std::vector> shapes; + shapes.reserve(tensor_count); + for (size_t i = 0; i < tensor_count; ++i) { + shapes.push_back((i % 2 == 0) ? shape1 : shape2); + } + + std::vector names; + names.reserve(tensor_count * 11); + std::vector g; + std::vector p; + std::vector m; + std::vector v; + std::vector p_ref_t; + std::vector m_ref_t; + std::vector v_ref_t; + std::vector q_ref; + std::vector dq; + std::vector dq_ref; + std::vector q; + g.reserve(tensor_count); + p.reserve(tensor_count); + m.reserve(tensor_count); + v.reserve(tensor_count); + p_ref_t.reserve(tensor_count); + m_ref_t.reserve(tensor_count); + v_ref_t.reserve(tensor_count); + q_ref.reserve(tensor_count); + dq.reserve(tensor_count); + dq_ref.reserve(tensor_count); + q.reserve(tensor_count); + + for (size_t i = 0; i < tensor_count; ++i) { + const std::vector &shape = shapes[i]; + names.push_back("g" + std::to_string(i)); + g.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + names.push_back("p" + std::to_string(i)); + p.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + names.push_back("m" + std::to_string(i)); + m.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + names.push_back("v" + std::to_string(i)); + v.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + + fillUniform(&g.back()); + fillUniform(&p.back()); + std::fill_n(m.back().rowwise_cpu_dptr(), product(m.back().rowwise_shape()), 0.0f); + std::fill_n(v.back().rowwise_cpu_dptr(), product(v.back().rowwise_shape()), 0.0f); + m.back().from_cpu(); + v.back().from_cpu(); + + names.push_back("p_ref_" + std::to_string(i)); + p_ref_t.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + names.push_back("m_ref_" + std::to_string(i)); + m_ref_t.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + names.push_back("v_ref_" + std::to_string(i)); + v_ref_t.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + const size_t n = shape[0] * shape[1]; + std::memcpy(p_ref_t.back().rowwise_cpu_dptr(), p.back().rowwise_cpu_dptr(), + n * sizeof(float)); + std::memcpy(m_ref_t.back().rowwise_cpu_dptr(), m.back().rowwise_cpu_dptr(), + n * sizeof(float)); + std::memcpy(v_ref_t.back().rowwise_cpu_dptr(), v.back().rowwise_cpu_dptr(), + n * sizeof(float)); + p_ref_t.back().from_cpu(); + m_ref_t.back().from_cpu(); + v_ref_t.back().from_cpu(); + + names.push_back("q_ref_" + std::to_string(i)); + q_ref.emplace_back(names.back().c_str(), shape, fp8_dtype, true, true, NVTE_MXFP8_1D_SCALING); + q_ref.back().set_with_gemm_swizzled_scales(false); + + names.push_back("dq" + std::to_string(i)); + dq.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + names.push_back("dq_ref_" + std::to_string(i)); + dq_ref.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); + + names.push_back("q" + std::to_string(i)); + q.emplace_back(names.back().c_str(), shape, fp8_dtype, true, true, NVTE_MXFP8_1D_SCALING); + q.back().set_with_gemm_swizzled_scales(false); + } + + Tensor noop("noop", std::vector{1}, DType::kInt32, true, false); + int zero = 0; + std::memcpy(noop.rowwise_cpu_dptr(), &zero, sizeof(int)); + noop.from_cpu(); + + std::vector> lists(8); + std::vector extra_wrappers; + extra_wrappers.reserve(tensor_count * 4); + + auto add_tensor = [&](Tensor &g, Tensor &p, Tensor &m, Tensor &v, Tensor &q) { + lists[0].push_back(g.data()); + lists[1].push_back(p.data()); + lists[2].push_back(m.data()); + lists[3].push_back(v.data()); + + extra_wrappers.emplace_back(q.rowwise_dptr(), q.rowwise_shape(), fp8_dtype); + lists[4].push_back(extra_wrappers.back().data()); + extra_wrappers.emplace_back(q.columnwise_dptr(), q.columnwise_shape(), fp8_dtype); + lists[5].push_back(extra_wrappers.back().data()); + extra_wrappers.emplace_back(q.rowwise_scale_inv_dptr(), q.rowwise_scale_inv_shape(), + DType::kByte); + lists[6].push_back(extra_wrappers.back().data()); + extra_wrappers.emplace_back(q.columnwise_scale_inv_dptr(), q.columnwise_scale_inv_shape(), + DType::kByte); + lists[7].push_back(extra_wrappers.back().data()); + }; + + for (size_t i = 0; i < tensor_count; ++i) { + add_tensor(g[i], p[i], m[i], v[i], q[i]); + } + + std::vector list_ptrs; + list_ptrs.reserve(lists.size()); + for (auto &l : lists) { + list_ptrs.push_back(l.data()); + } + + nvte_multi_tensor_adam_mxfp8_cuda(65536, noop.data(), list_ptrs.data(), lists.size(), + lists[0].size(), static_cast(fp8_dtype), lr, beta1, + beta2, eps, step, mode, bias_correction, weight_decay, 0); + + std::vector> ref_lists(4); + for (size_t i = 0; i < tensor_count; ++i) { + ref_lists[0].push_back(g[i].data()); + ref_lists[1].push_back(p_ref_t[i].data()); + ref_lists[2].push_back(m_ref_t[i].data()); + ref_lists[3].push_back(v_ref_t[i].data()); + } + std::vector ref_list_ptrs; + ref_list_ptrs.reserve(ref_lists.size()); + for (auto &l : ref_lists) { + ref_list_ptrs.push_back(l.data()); + } + + nvte_multi_tensor_adam_cuda(65536, noop.data(), ref_list_ptrs.data(), ref_lists.size(), + ref_lists[0].size(), lr, beta1, beta2, eps, step, mode, + bias_correction, weight_decay, 0); + + for (size_t i = 0; i < tensor_count; ++i) { + nvte_quantize(p_ref_t[i].data(), q_ref[i].data(), 0); + nvte_dequantize(q[i].data(), dq[i].data(), 0); + nvte_dequantize(q_ref[i].data(), dq_ref[i].data(), 0); + } + + cudaDeviceSynchronize(); + + for (size_t i = 0; i < tensor_count; ++i) { + q[i].to_cpu(); + p[i].to_cpu(); + m[i].to_cpu(); + v[i].to_cpu(); + q_ref[i].to_cpu(); + dq[i].to_cpu(); + dq_ref[i].to_cpu(); + p_ref_t[i].to_cpu(); + m_ref_t[i].to_cpu(); + v_ref_t[i].to_cpu(); + } + + for (size_t i = 0; i < lists[0].size(); ++i) { + const Tensor &g_i = g[i]; + const Tensor &p_i = p[i]; + const Tensor &m_i = m[i]; + const Tensor &v_i = v[i]; + Tensor &q_i = q[i]; + const Tensor &p_ref_t_i = p_ref_t[i]; + const Tensor &m_ref_t_i = m_ref_t[i]; + const Tensor &v_ref_t_i = v_ref_t[i]; + Tensor &q_ref_i = q_ref[i]; + + compareResults("p", p_i, p_ref_t_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, 0); + compareResults("m", m_i, m_ref_t_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, 0); + compareResults("v", v_i, v_ref_t_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, 0); + + const Tensor &dq_i = dq[i]; + const Tensor &dq_ref_i = dq_ref[i]; + compareResults("dequantized", dq_i, dq_ref_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, + 0); + + const size_t rs = q_i.rowwise_scale_inv_shape().data[1]; + const size_t cs = q_i.columnwise_scale_inv_shape().data[1]; + const size_t rowwise_scale_size = q_i.rowwise_scale_inv_shape().data[0] * rs; + const size_t colwise_scale_size = q_i.columnwise_scale_inv_shape().data[0] * cs; + compareResults("rowwise_scale", q_i.rowwise_cpu_scale_inv_ptr(), + q_ref_i.rowwise_cpu_scale_inv_ptr(), rowwise_scale_size, 0.0f); + compareResults("colwise_scale", q_i.columnwise_cpu_scale_inv_ptr(), + q_ref_i.columnwise_cpu_scale_inv_ptr(), colwise_scale_size, 0.0f); + + uint8_t *row_data = nullptr; + uint8_t *col_data = nullptr; + uint8_t *row_data_ref = nullptr; + uint8_t *col_data_ref = nullptr; + if (fp8_dtype == DType::kFloat8E4M3) { + row_data = reinterpret_cast(q_i.rowwise_cpu_dptr()); + col_data = reinterpret_cast(q_i.columnwise_cpu_dptr()); + row_data_ref = reinterpret_cast(q_ref_i.rowwise_cpu_dptr()); + col_data_ref = reinterpret_cast(q_ref_i.columnwise_cpu_dptr()); + } else { + row_data = reinterpret_cast(q_i.rowwise_cpu_dptr()); + col_data = reinterpret_cast(q_i.columnwise_cpu_dptr()); + row_data_ref = reinterpret_cast(q_ref_i.rowwise_cpu_dptr()); + col_data_ref = reinterpret_cast(q_ref_i.columnwise_cpu_dptr()); + } + const size_t data_size = q_i.rowwise_shape().data[0] * q_i.rowwise_shape().data[1]; + compareResults("rowwise_data", row_data, row_data_ref, data_size, 0.0f); + compareResults("colwise_data", col_data, col_data_ref, data_size, 0.0f); + } +} + +} // namespace + +TEST(MultiTensorAdamMXFP8, E4M3) { run_mxfp8_adam_test(DType::kFloat8E4M3); } + +TEST(MultiTensorAdamMXFP8, E5M2) { run_mxfp8_adam_test(DType::kFloat8E5M2); } diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 927407f478..eab181fa82 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -200,6 +200,16 @@ class Tensor { return tensor_.get_columnwise_data().data_ptr; } + void *rowwise_scale_inv_dptr() const { + NVTE_CHECK(rowwise_, "Tensor does not have rowwise data!"); + return tensor_.get_rowwise_scale_inv().data_ptr; + } + + void *columnwise_scale_inv_dptr() const { + NVTE_CHECK(columnwise_, "Tensor does not have columnwise data!"); + return tensor_.get_columnwise_scale_inv().data_ptr; + } + template T *rowwise_cpu_dptr() const { NVTE_CHECK(TypeInfo::dtype == tensor_.dtype(), "Invalid type!"); diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/run_fsdp2_fused_adam.py index c39957cf13..34764d4e0a 100644 --- a/tests/pytorch/distributed/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/run_fsdp2_fused_adam.py @@ -36,7 +36,11 @@ def get_recipe_from_string(recipe): SEQ_LEN = 32 BATCH_PER_RANK = 2 NUM_STEPS = 3 +LOCAL_RANK = None +def dist_print(msg): + if LOCAL_RANK == 0: + print(msg) def save_custom_attrs(module): custom_attrs = {} @@ -151,6 +155,8 @@ def test_fused_adam_fp8_master_weights(recipe=None): - Training loop completes without error - DTensor wrapping and QuantizedTensor local tensors are preserved """ + global LOCAL_RANK + LOCAL_RANK = int(os.environ["LOCAL_RANK"]) world_size, _, device = _setup() model = _build_model(fp8_init=True, recipe=recipe) @@ -183,7 +189,7 @@ def test_fused_adam_fp8_master_weights(recipe=None): loss = F.mse_loss(output, target) loss.backward() optimizer.step() - + dist_print(f"Step {step} completed with loss {loss.item()}") # Verify optimizer states for param in model.parameters(): state = optimizer.state[param] diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 02e45d99cb..6d7ae4d7bb 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -224,11 +224,6 @@ def test_fsdp2_dcp_output_parity_async(fp_recipe): @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") def test_fsdp2_safetensors_fp32_export(fp_recipe): """Export FP32 model from optimizer master weights to safetensors.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" - ) _run_fused_adam_test("safetensors_fp32_export", fp_recipe) diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index 09ab260f15..90c87b166e 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -149,6 +149,41 @@ void nvte_multi_tensor_adam_fp8_cuda(int chunk_size, NVTETensor noop_flag, const float weight_decay, const NVTEDType fp8_dtype, cudaStream_t stream); +/*! \brief Compute and apply gradient update to parameters for Adam optimizer + * when model parameters are in MXFP8 precision. + * + * The update is applied to FP32 master parameters, then the master + * parameters are quantized to MXFP8 rowwise and columnwise data + * (both are always required). + * + * \warning This API is **experimental** and subject to change. + * + * \param[in] chunk_size Number of tensor elements processed by a CUDA block. + * \param[in] noop_flag If this single element tensor has non-zero value, kernel will exit immediately. + * \param[in,out] tensor_lists 2D array of input tensors with 8 lists in order: + * (0) gradients, (1) FP32 master params, (2) first moment, + * (3) second moment, (4) rowwise MXFP8 data, + * (5) columnwise MXFP8 data, (6) rowwise scale-inv, + * (7) columnwise scale-inv. + * \param[in] num_tensor_lists Size (dim0) of tensor_lists. Must be 8. + * \param[in] num_tensors_per_list Size (dim1) of tensor_lists. + * \param[in] fp8_dtype MXFP8 element type for quantization (E4M3/E5M2). + * \param[in] lr Learning rate. + * \param[in] beta1 Coefficient for first moment of gradient. + * \param[in] beta2 Coefficient for second moment of gradient. + * \param[in] epsilon Term added to the denominator for numerical stability. + * \param[in] step Iteration counter. + * \param[in] mode Whether to use AdamW (L2 penalty applied to params). + * \param[in] bias_correction Whether to apply correction factor for moment estimates. + * \param[in] weight_decay L2 penalty for weight decay. + * \param[in] stream CUDA stream used for this operation. + */ +void nvte_multi_tensor_adam_mxfp8_cuda( + int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, + const size_t num_tensor_lists, const size_t num_tensors_per_list, const NVTEDType fp8_dtype, + const float lr, const float beta1, const float beta2, const float epsilon, const int step, + const int mode, const int bias_correction, const float weight_decay, cudaStream_t stream); + /*! \brief Compute and apply gradient update to parameters for Adam optimizer * with CUDA graph support and LR scheduling. * diff --git a/transformer_engine/common/multi_tensor/adam.cu b/transformer_engine/common/multi_tensor/adam.cu index 29a073be84..fa75c645f3 100644 --- a/transformer_engine/common/multi_tensor/adam.cu +++ b/transformer_engine/common/multi_tensor/adam.cu @@ -4,12 +4,16 @@ * See LICENSE for license information. ************************************************************************/ +#include #include #include #include #include +#include "../common.h" +#include "../util/math.h" #include "../utils.cuh" +#include "../util/ptx.cuh" #include "multi_tensor_apply.cuh" namespace transformer_engine { @@ -27,6 +31,7 @@ typedef enum { using MATH_T = float; using fp8e4m3 = __nv_fp8_e4m3; using fp8e5m2 = __nv_fp8_e5m2; +using e8m0_t = transformer_engine::e8m0_t; template struct is_fp8 : std::false_type {}; @@ -49,6 +54,31 @@ struct FP8Data { template <> struct FP8Data {}; +template +__device__ __forceinline__ void adam_update(T &r_g, T &r_p, T &r_m, T &r_v, const float beta1, + const float beta2, const float beta1_correction, + const float beta2_correction, const float epsilon, + const float lr, adamMode_t mode, const float decay) { + if (mode == ADAM_MODE_0) { // L2 + r_g = r_g + (decay * r_p); + r_m = beta1 * r_m + (1 - beta1) * r_g; + r_v = beta2 * r_v + (1 - beta2) * r_g * r_g; + T next_m_unbiased = r_m / beta1_correction; + T next_v_unbiased = r_v / beta2_correction; + T denom = sqrtf(next_v_unbiased) + epsilon; + T update = next_m_unbiased / denom; + r_p = r_p - (lr * update); + } else { // weight decay + r_m = beta1 * r_m + (1 - beta1) * r_g; + r_v = beta2 * r_v + (1 - beta2) * r_g * r_g; + T next_m_unbiased = r_m / beta1_correction; + T next_v_unbiased = r_v / beta2_correction; + T denom = sqrtf(next_v_unbiased) + epsilon; + T update = (next_m_unbiased / denom) + (decay * r_p); + r_p = r_p - (lr * update); + } +} + template struct AdamFunctorMaster { static constexpr bool is_fp8_type = is_fp8::value; @@ -122,24 +152,8 @@ struct AdamFunctorMaster { } #pragma unroll for (int ii = 0; ii < ILP; ii++) { - if (mode == ADAM_MODE_0) { // L2 - r_g[ii] = r_g[ii] + (decay * r_p[ii]); - r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; - r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; - MATH_T next_m_unbiased = r_m[ii] / beta1_correction; - MATH_T next_v_unbiased = r_v[ii] / beta2_correction; - MATH_T denom = sqrtf(next_v_unbiased) + epsilon; - MATH_T update = next_m_unbiased / denom; - r_p[ii] = r_p[ii] - (lr * update); - } else { // weight decay - r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; - r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; - MATH_T next_m_unbiased = r_m[ii] / beta1_correction; - MATH_T next_v_unbiased = r_v[ii] / beta2_correction; - MATH_T denom = sqrtf(next_v_unbiased) + epsilon; - MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]); - r_p[ii] = r_p[ii] - (lr * update); - } + adam_update(r_g[ii], r_p[ii], r_m[ii], r_v[ii], beta1, beta2, beta1_correction, + beta2_correction, epsilon, lr, mode, decay); } #pragma unroll @@ -572,6 +586,188 @@ struct AdamCapturableMasterFunctor { } }; +template +__device__ __forceinline__ FP8_T cast_to_fp8(float x) { + return static_cast(x); +} + +__device__ __forceinline__ float fp8_max_norm_rcp(uint8_t fp8_dtype) { + if (fp8_dtype == static_cast(transformer_engine::DType::kFloat8E4M3)) { + return transformer_engine::Quantized_Limits::max_norm_rcp; + } + return transformer_engine::Quantized_Limits::max_norm_rcp; +} + +template +__global__ void adam_mxfp8_fused_kernel( + int64_t chunk_size, volatile int *noop_gmem, MXFP8TensorListMetadata tl, float beta1, + float beta2, float beta1_correction, float beta2_correction, float epsilon, float lr, int mode, + float weight_decay) { + // Stage 0: optional early-exit if a noop flag is set. + if (noop_gmem != nullptr && *noop_gmem == 1) { + return; + } + (void)chunk_size; + + // Stage 1: map this block to a specific tensor tile. + const int block_idx = blockIdx.x; + const int tensor_idx = tl.block_to_tensor[block_idx]; + const int tile_idx = tl.block_to_tile[block_idx]; + const int64_t rows_val = tl.rows[tensor_idx]; + const int64_t cols_val = tl.cols[tensor_idx]; + if (rows_val == 0 || cols_val == 0) { + return; + } + + const int64_t tiles_per_row = (cols_val + MXFP8_TILE - 1) / MXFP8_TILE; + const int64_t tile_row = tile_idx / tiles_per_row; + const int64_t tile_col = tile_idx % tiles_per_row; + const int64_t row_base = tile_row * MXFP8_TILE; + const int64_t col_base = tile_col * MXFP8_TILE; + + // Stage 2: load pointers for grads/params/moments and MXFP8 outputs/scales. + GRAD_T *g = reinterpret_cast(tl.addresses[0][tensor_idx]); + PARAM_T *p = reinterpret_cast(tl.addresses[1][tensor_idx]); + MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_idx]); + MOMENT_T *v = reinterpret_cast(tl.addresses[3][tensor_idx]); + + auto *rowwise_data = reinterpret_cast(tl.addresses[4][tensor_idx]); + auto *colwise_data = reinterpret_cast(tl.addresses[5][tensor_idx]); + auto *rowwise_scale_inv = reinterpret_cast(tl.addresses[6][tensor_idx]); + auto *colwise_scale_inv = reinterpret_cast(tl.addresses[7][tensor_idx]); + + const int64_t unpadded_scales_X_rowwise = (cols_val + MXFP8_TILE - 1) / MXFP8_TILE; + constexpr int64_t kRowwiseScaleAlign = 4; + const int64_t row_stride = + DIVUP_TO_MULTIPLE(unpadded_scales_X_rowwise, kRowwiseScaleAlign); + constexpr int64_t kColwiseScaleAlign = 128; + const int64_t col_stride = DIVUP_TO_MULTIPLE(cols_val, kColwiseScaleAlign); + const uint8_t dtype = tl.fp8_dtype[tensor_idx]; + const auto adam_mode = static_cast(mode); + + // Stage 3: initialize shared amax accumulators per row/col within the tile. + __shared__ float row_max_vals[MXFP8_TILE]; + __shared__ float col_max_vals[MXFP8_TILE]; + if (threadIdx.x < MXFP8_TILE) { + row_max_vals[threadIdx.x] = 0.0f; + col_max_vals[threadIdx.x] = 0.0f; + } + __syncthreads(); + + for (int t = threadIdx.x; t < MXFP8_TILE_ELEMS; t += blockDim.x) { + const int local_r = t / MXFP8_TILE; + const int local_c = t % MXFP8_TILE; + const int64_t r = row_base + local_r; + const int64_t c = col_base + local_c; + if (r >= rows_val || c >= cols_val) { + continue; + } + const index_t idx = static_cast(r * cols_val + c); + + float r_g = static_cast(g[idx]); + float r_p = static_cast(p[idx]); + float r_m = static_cast(m[idx]); + float r_v = static_cast(v[idx]); + + // Stage 4: apply Adam update in FP32 and write back updated p/m/v. + transformer_engine::multi_tensor_adam::adam_update( + r_g, r_p, r_m, r_v, beta1, beta2, beta1_correction, beta2_correction, epsilon, lr, + adam_mode, weight_decay); + + p[idx] = static_cast(r_p); + m[idx] = static_cast(r_m); + v[idx] = static_cast(r_v); + + // Stage 5: accumulate per-row/col absmax for MXFP8 scaling. + const float abs_p = fabsf(r_p); + transformer_engine::atomicMaxFloat(&row_max_vals[local_r], abs_p); + transformer_engine::atomicMaxFloat(&col_max_vals[local_c], abs_p); + } + + __syncthreads(); + + // Stage 6: write rowwise/colwise scale-inverse exponents for the tile. + const float max_norm_rcp = fp8_max_norm_rcp(dtype); + + for (int r = threadIdx.x; r < MXFP8_TILE; r += blockDim.x) { + const int64_t row = row_base + r; + if (row >= rows_val) { + continue; + } + const float amax = row_max_vals[r]; + const ::transformer_engine::e8m0_t biased_exponent = + transformer_engine::ptx::float_to_e8m0(amax * max_norm_rcp); + const size_t scale_idx = static_cast(row * row_stride + tile_col); + rowwise_scale_inv[scale_idx] = reinterpret_cast(biased_exponent); + } + + for (int c = threadIdx.x; c < MXFP8_TILE; c += blockDim.x) { + const int64_t col = col_base + c; + if (col >= cols_val) { + continue; + } + const float amax = col_max_vals[c]; + const ::transformer_engine::e8m0_t biased_exponent = + transformer_engine::ptx::float_to_e8m0(amax * max_norm_rcp); + const size_t scale_idx = static_cast(tile_row * col_stride + col); + colwise_scale_inv[scale_idx] = reinterpret_cast(biased_exponent); + } + + __syncthreads(); + + // Stage 7: quantize updated params to MXFP8 using rowwise and colwise scales. + for (int t = threadIdx.x; t < MXFP8_TILE_ELEMS; t += blockDim.x) { + const int local_r = t / MXFP8_TILE; + const int local_c = t % MXFP8_TILE; + const int64_t r = row_base + local_r; + const int64_t c = col_base + local_c; + if (r >= rows_val || c >= cols_val) { + continue; + } + const index_t idx = static_cast(r * cols_val + c); + const float r_p = static_cast(p[idx]); + + const size_t row_scale_idx = static_cast(r * row_stride + tile_col); + const uint8_t row_raw = rowwise_scale_inv[row_scale_idx]; + const ::transformer_engine::e8m0_t row_biased = + reinterpret_cast(row_raw); + const float row_scale_inv = transformer_engine::ptx::exp2f_rcp(row_biased); + if (dtype == static_cast(transformer_engine::DType::kFloat8E4M3)) { + auto *out = reinterpret_cast(rowwise_data); + out[idx] = cast_to_fp8(r_p * row_scale_inv); + } else { + auto *out = reinterpret_cast(rowwise_data); + out[idx] = cast_to_fp8(r_p * row_scale_inv); + } + + const size_t col_scale_idx = static_cast(tile_row * col_stride + c); + const uint8_t col_raw = colwise_scale_inv[col_scale_idx]; + const ::transformer_engine::e8m0_t col_biased = + reinterpret_cast(col_raw); + const float col_scale_inv = transformer_engine::ptx::exp2f_rcp(col_biased); + if (dtype == static_cast(transformer_engine::DType::kFloat8E4M3)) { + auto *out = reinterpret_cast(colwise_data); + out[idx] = cast_to_fp8(r_p * col_scale_inv); + } else { + auto *out = reinterpret_cast(colwise_data); + out[idx] = cast_to_fp8(r_p * col_scale_inv); + } + } +} + +inline bool requires_64bit_indexing(const std::vector> &tensor_lists) { + const size_t num_tensor_lists = tensor_lists.size(); + const size_t num_tensors_per_list = tensor_lists[0].size(); + for (size_t i = 0; i < num_tensor_lists; ++i) { + for (size_t j = 0; j < num_tensors_per_list; ++j) { + if (tensor_lists[i][j]->numel() >= INT_MAX) { + return true; + } + } + } + return false; +} + void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, std::vector> tensor_lists, const float lr, const float beta1, const float beta2, const float epsilon, @@ -624,25 +820,13 @@ void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, } } - // Check if 64-bit indices are required - bool requires_64bit_indexing = false; - for (size_t i = 0; i < num_tensor_lists; i++) { - for (size_t j = 0; j < num_tensors_per_list; j++) { - if (tensor_lists[i][j]->numel() >= INT_MAX) { - requires_64bit_indexing = true; - break; - } - } - if (requires_64bit_indexing) { - break; - } - } + const bool use_64bit_indexing = requires_64bit_indexing(tensor_lists); // Get moment dtype (m and v have the same dtype, already validated above) const auto moment_type_te = tensor_lists[2][0]->dtype(); // Launch kernel - if (requires_64bit_indexing) { + if (use_64bit_indexing) { if (num_tensor_lists == 4) { // g, p, m, v TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( @@ -766,28 +950,41 @@ void multi_tensor_adam_param_remainder_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK_CUDA(cudaGetLastError()); } -void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, - std::vector> tensor_lists, const float lr, - const float beta1, const float beta2, const float epsilon, - const int step, const int mode, const int bias_correction, - const float weight_decay, const DType fp8_dtype, - cudaStream_t stream) { - // Handle bias correction mode - float bias_correction1 = 1.0f, bias_correction2 = 1.0f; +inline std::pair compute_bias_correction(int bias_correction, float beta1, + float beta2, int step) { + float bias_correction1 = 1.0f; + float bias_correction2 = 1.0f; if (bias_correction == 1) { bias_correction1 = 1 - std::pow(beta1, step); bias_correction2 = 1 - std::pow(beta2, step); } + return {bias_correction1, bias_correction2}; +} - // Check tensor list sizes - // 8 tensor lists: g, p_fp8, m, v, p_master, scale, amax, scale_inv +inline void check_tensor_list_sizes(const std::vector> &tensor_lists, + size_t expected_lists) { const size_t num_tensor_lists = tensor_lists.size(); - NVTE_CHECK(num_tensor_lists == 8, "Expected 8 tensor lists, but found ", num_tensor_lists); + NVTE_CHECK(num_tensor_lists == expected_lists, "Expected ", expected_lists, + " tensor lists, but found ", num_tensor_lists); const size_t num_tensors_per_list = tensor_lists[0].size(); - for (size_t i = 1; i < num_tensor_lists; i++) { + for (size_t i = 1; i < num_tensor_lists; ++i) { NVTE_CHECK(tensor_lists[i].size() == num_tensors_per_list, "Tensor list ", i, " has size=", tensor_lists[i].size(), ", but expected size=", num_tensors_per_list); } +} + + +void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, + std::vector> tensor_lists, const float lr, + const float beta1, const float beta2, const float epsilon, + const int step, const int mode, const int bias_correction, + const float weight_decay, const DType fp8_dtype, + cudaStream_t stream) { + auto [bias_correction1, bias_correction2] = + compute_bias_correction(bias_correction, beta1, beta2, step); + check_tensor_list_sizes(tensor_lists, 8); + const size_t num_tensor_lists = tensor_lists.size(); + const size_t num_tensors_per_list = tensor_lists[0].size(); // Check tensor dtypes const auto g_in_type_te = tensor_lists[0][0]->dtype(); @@ -819,22 +1016,10 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, ", but expected dtype=", to_string(DType::kFloat32)); } - // Check if 64-bit indices are required - bool requires_64bit_indexing = false; - for (size_t i = 0; i < num_tensor_lists; i++) { - for (size_t j = 0; j < num_tensors_per_list; j++) { - if (tensor_lists[i][j]->numel() >= INT_MAX) { - requires_64bit_indexing = true; - break; - } - } - if (requires_64bit_indexing) { - break; - } - } + const bool use_64bit_indexing = requires_64bit_indexing(tensor_lists); // Launch kernel - if (requires_64bit_indexing) { + if (use_64bit_indexing) { TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( fp8_dtype, FP8_T, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( @@ -856,6 +1041,76 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK_CUDA(cudaGetLastError()); } +void multi_tensor_adam_mxfp8_cuda(int chunk_size, Tensor noop_flag, + std::vector> tensor_lists, const float lr, + const float beta1, const float beta2, const float epsilon, + const int step, const int mode, const int bias_correction, + const float weight_decay, const DType fp8_dtype, + cudaStream_t stream) { + auto [bias_correction1, bias_correction2] = + compute_bias_correction(bias_correction, beta1, beta2, step); + check_tensor_list_sizes(tensor_lists, 8); + const size_t num_tensor_lists = tensor_lists.size(); + const size_t num_tensors_per_list = tensor_lists[0].size(); + + NVTE_CHECK(fp8_dtype == DType::kFloat8E4M3 || fp8_dtype == DType::kFloat8E5M2, + "fp8_dtype must be E4M3 or E5M2 for MXFP8 fused Adam."); + + // Check tensor dtypes + const auto g_in_type_te = tensor_lists[0][0]->dtype(); + const auto p_in_type_te = tensor_lists[1][0]->dtype(); + const auto moment_type_te = tensor_lists[2][0]->dtype(); + for (size_t j = 0; j < num_tensors_per_list; ++j) { + NVTE_CHECK(tensor_lists[0][j]->dtype() == g_in_type_te, "Grad tensor ", j, + " has dtype=", to_string(tensor_lists[0][j]->dtype()), + ", but expected dtype=", to_string(g_in_type_te)); + NVTE_CHECK(tensor_lists[1][j]->dtype() == p_in_type_te, "Param tensor ", j, + " has dtype=", to_string(tensor_lists[1][j]->dtype()), + ", but expected dtype=", to_string(p_in_type_te)); + { + const bool m_is_fp32 = tensor_lists[2][j]->dtype() == DType::kFloat32; + const bool m_is_bf16 = tensor_lists[2][j]->dtype() == DType::kBFloat16; + const bool v_is_fp32 = tensor_lists[3][j]->dtype() == DType::kFloat32; + const bool v_is_bf16 = tensor_lists[3][j]->dtype() == DType::kBFloat16; + NVTE_CHECK((m_is_fp32 && v_is_fp32) || (m_is_bf16 && v_is_bf16), + "First and second moment tensors must both be Float32 or both be BFloat16, but " + "tensor ", + j, " has first moment dtype=", to_string(tensor_lists[2][j]->dtype()), + " and second moment dtype=", to_string(tensor_lists[3][j]->dtype())); + } + } + + const bool use_64bit_indexing = requires_64bit_indexing(tensor_lists); + + if (use_64bit_indexing) { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + p_in_type_te, p_in_type, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + g_in_type_te, g_in_type, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply_mxfp8< + transformer_engine::multi_tensor_adam::adam_mxfp8_fused_kernel< + p_in_type, g_in_type, moment_type, int64_t>>( + chunk_size, noop_flag, tensor_lists, static_cast(fp8_dtype), stream, + beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, mode, + weight_decay);))); + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + p_in_type_te, p_in_type, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + g_in_type_te, g_in_type, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply_mxfp8< + transformer_engine::multi_tensor_adam::adam_mxfp8_fused_kernel< + p_in_type, g_in_type, moment_type, int32_t>>( + chunk_size, noop_flag, tensor_lists, static_cast(fp8_dtype), stream, + beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, mode, + weight_decay);))); + } +} + void multi_tensor_adam_capturable_cuda(int chunk_size, Tensor noop_flag, std::vector> tensor_lists, Tensor lr, const float beta1, const float beta2, const float epsilon, @@ -1018,6 +1273,19 @@ void nvte_multi_tensor_adam_fp8_cuda(int chunk_size, NVTETensor noop_flag, epsilon, step, mode, bias_correction, weight_decay, static_cast(fp8_dtype), stream); } +void nvte_multi_tensor_adam_mxfp8_cuda( + int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, + const size_t num_tensor_lists, const size_t num_tensors_per_list, const NVTEDType fp8_dtype, + const float lr, const float beta1, const float beta2, const float epsilon, const int step, + const int mode, const int bias_correction, const float weight_decay, cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_adam_mxfp8_cuda); + using namespace transformer_engine; + multi_tensor_adam::multi_tensor_adam_mxfp8_cuda( + chunk_size, *convertNVTETensorCheck(noop_flag), + convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list), lr, beta1, beta2, + epsilon, step, mode, bias_correction, weight_decay, static_cast(fp8_dtype), stream); +} + void nvte_multi_tensor_adam_capturable_cuda( int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, const size_t num_tensor_lists, const size_t num_tensors_per_list, NVTETensor lr, const float beta1, const float beta2, diff --git a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh index 3062ead551..c334f3908e 100644 --- a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh +++ b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh @@ -35,6 +35,23 @@ struct TensorListMetadata : public TensorListMetadataBase { void *fp8_meta_addresses[3][depth_to_max_tensors[n - 1]]; }; +constexpr int MXFP8_TILE = 32; +constexpr int MXFP8_TILE_ELEMS = MXFP8_TILE * MXFP8_TILE; +constexpr int MXFP8_BLOCK_THREADS = 256; +constexpr int MXFP8_MAX_TENSORS = 24; +constexpr int MXFP8_MAX_BLOCKS = 320; + +struct MXFP8TensorListMetadata { + void *addresses[8][MXFP8_MAX_TENSORS]; + int sizes[MXFP8_MAX_TENSORS]; + int rows[MXFP8_MAX_TENSORS]; + int cols[MXFP8_MAX_TENSORS]; + uint8_t fp8_dtype[MXFP8_MAX_TENSORS]; + unsigned char block_to_tensor[MXFP8_MAX_BLOCKS]; + int block_to_tile[MXFP8_MAX_BLOCKS]; + int start_tensor_this_launch; +}; + template __global__ void multi_tensor_apply_kernel(int64_t chunk_size, volatile int *noop_flag, T tl, U callable, ArgTypes... args) { @@ -113,3 +130,80 @@ void multi_tensor_apply(int64_t block_size, int64_t chunk_size, } } } + +template +void multi_tensor_apply_mxfp8(int64_t chunk_size, const transformer_engine::Tensor &noop_flag, + std::vector> tensor_lists, + uint8_t fp8_dtype, cudaStream_t stream, ArgTypes... args) { + constexpr size_t kNumTensorLists = 8; + NVTE_CHECK(tensor_lists.size() == kNumTensorLists, + "Expected 8 tensor lists for MXFP8, but found ", tensor_lists.size()); + + const size_t num_tensors_per_list = tensor_lists[0].size(); + if (num_tensors_per_list == 0) { + return; + } + for (size_t i = 1; i < tensor_lists.size(); ++i) { + NVTE_CHECK(tensor_lists[i].size() == num_tensors_per_list, "Tensor list ", i, + " has size=", tensor_lists[i].size(), ", but expected size=", num_tensors_per_list); + } + + MXFP8TensorListMetadata tl; + tl.start_tensor_this_launch = 0; + int loc_block_info = 0; + int loc_tensor_info = 0; + + for (size_t t = 0; t < num_tensors_per_list; ++t) { + + const auto &g = tensor_lists[0][t]; + const auto &rowwise_data = tensor_lists[4][t]; + const auto &colwise_data = tensor_lists[5][t]; + + const int rows_val = static_cast(rowwise_data->data.shape[0]); + const int cols_val = static_cast(rowwise_data->data.shape[1]); + + tl.sizes[loc_tensor_info] = g->numel(); + tl.rows[loc_tensor_info] = rows_val; + tl.cols[loc_tensor_info] = cols_val; + tl.fp8_dtype[loc_tensor_info] = fp8_dtype; + + for (int d = 0; d < kNumTensorLists; ++d) { + tl.addresses[d][loc_tensor_info] = tensor_lists[d][t]->data.dptr; + } + loc_tensor_info++; + + const int tiles_y = (rows_val + MXFP8_TILE - 1) / MXFP8_TILE; + const int tiles_x = (cols_val + MXFP8_TILE - 1) / MXFP8_TILE; + const int tiles_this_tensor = tiles_y * tiles_x; + + for (int tile = 0; tile < tiles_this_tensor; ++tile) { + tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; + tl.block_to_tile[loc_block_info] = tile; + loc_block_info++; + + const bool blocks_full = (loc_block_info == MXFP8_MAX_BLOCKS); + const bool tensors_full = + (loc_tensor_info == MXFP8_MAX_TENSORS && tile == tiles_this_tensor - 1); + const bool last_tile = (t == num_tensors_per_list - 1 && tile == tiles_this_tensor - 1); + if (blocks_full || tensors_full || last_tile) { + Kernel<<>>( + chunk_size, reinterpret_cast(noop_flag.data.dptr), tl, args...); + NVTE_CHECK_CUDA(cudaGetLastError()); + loc_block_info = 0; + if (tile == tiles_this_tensor - 1) { + loc_tensor_info = 0; + tl.start_tensor_this_launch = t + 1; + } else { + tl.rows[0] = tl.rows[loc_tensor_info - 1]; + tl.cols[0] = tl.cols[loc_tensor_info - 1]; + tl.fp8_dtype[0] = tl.fp8_dtype[loc_tensor_info - 1]; + for (int d = 0; d < kNumTensorLists; ++d) { + tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; + } + loc_tensor_info = 1; + tl.start_tensor_this_launch = t; + } + } + } + } +} diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1c5116a8da..65e2c54d67 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -517,6 +517,12 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, at::Tensor noop_flag, const int step, const int mode, const int bias_correction, const float weight_decay, DType fp8_dtype); +void multi_tensor_adam_mxfp8_cuda(int chunk_size, at::Tensor noop_flag, + std::vector> tensor_lists, const float lr, + const float beta1, const float beta2, const float epsilon, + const int step, const int mode, const int bias_correction, + const float weight_decay, DType fp8_dtype); + void multi_tensor_adam_capturable_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists, at::Tensor lr, const float beta1, const float beta2, diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp index 145e1d4b40..01a21d44bb 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp @@ -5,6 +5,7 @@ ************************************************************************/ #include "../../extensions.h" +#include "pybind.h" namespace transformer_engine::pytorch { @@ -51,6 +52,25 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, at::Tensor noop_flag, at::cuda::getCurrentCUDAStream()); } +void multi_tensor_adam_mxfp8_cuda(int chunk_size, at::Tensor noop_flag, + std::vector> tensor_lists, const float lr, + const float beta1, const float beta2, const float epsilon, + const int step, const int mode, const int bias_correction, + const float weight_decay, DType fp8_dtype) { + auto noop_flag_cu = makeTransformerEngineTensor(noop_flag); + auto [_, __, tensor_lists_ptr, num_lists, num_tensors] = + makeTransformerEngineTensorList(tensor_lists); + + NVTE_CHECK(num_lists == 8, + "Expected 8 tensor lists (g, p_master, m, v, rowwise_data, colwise_data, " + "rowwise_scale_inv, colwise_scale_inv), but found ", + num_lists); + nvte_multi_tensor_adam_mxfp8_cuda( + chunk_size, noop_flag_cu.data(), tensor_lists_ptr.data(), num_lists, num_tensors, + static_cast(fp8_dtype), lr, beta1, beta2, epsilon, step, mode, bias_correction, + weight_decay, at::cuda::getCurrentCUDAStream()); +} + void multi_tensor_adam_capturable_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists, at::Tensor lr, const float beta1, const float beta2, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c590a3c9e2..6def07b08e 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -525,6 +525,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("multi_tensor_adam_fp8", &transformer_engine::pytorch::multi_tensor_adam_fp8_cuda, "Compute and apply gradient update to parameters for Adam optimizer", py::call_guard()); + m.def("multi_tensor_adam_mxfp8", &transformer_engine::pytorch::multi_tensor_adam_mxfp8_cuda, + "Compute and apply gradient update to parameters for Adam optimizer"); m.def("multi_tensor_adam_capturable", &transformer_engine::pytorch::multi_tensor_adam_capturable_cuda, "Compute and apply gradient update to parameters for Adam optimizer with CUDA graph " diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index bcfd2bef19..f4ab2e7c37 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -14,6 +14,7 @@ from torch.distributed._tensor import DTensor import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from .multi_tensor_apply import multi_tensor_applier @@ -189,6 +190,7 @@ def __init__( self.multi_tensor_adam = tex.multi_tensor_adam self.multi_tensor_adam_param_remainder = tex.multi_tensor_adam_param_remainder self.multi_tensor_adam_fp8 = tex.multi_tensor_adam_fp8 + self.multi_tensor_adam_mxfp8 = tex.multi_tensor_adam_mxfp8 self.multi_tensor_adam_capturable = tex.multi_tensor_adam_capturable self.multi_tensor_adam_capturable_master = tex.multi_tensor_adam_capturable_master @@ -544,18 +546,27 @@ def step(self, closure=None, grad_scaler=None): # create lists for multi-tensor apply p_main_of_fp8_model = [] p_main_of_f16_model = [] + p_main_of_mxfp8_model = [] g_of_fp8_model = [] g_of_f16_model = [] g_of_f32_model = [] + g_of_mxfp8_model = [] m_of_fp8_model = [] m_of_f16_model = [] m_of_f32_model = [] + m_of_mxfp8_model = [] v_of_fp8_model = [] v_of_f16_model = [] v_of_f32_model = [] + v_of_mxfp8_model = [] p_fp8_model = [] p_f16_model = [] p_f32_model = [] + # mxfp8 meta + p_mxfp8_rowwise = [] + p_mxfp8_colwise = [] + p_mxfp8_rowwise_scale_inv = [] + p_mxfp8_colwise_scale_inv = [] # fp8 meta scales = [] amaxes = [] @@ -623,10 +634,30 @@ def step(self, closure=None, grad_scaler=None): g_of_fp8_model.append(p_grad.data) m_of_fp8_model.append(unscaled_state["exp_avg"]) v_of_fp8_model.append(unscaled_state["exp_avg_sq"]) + elif isinstance(p, MXFP8Tensor) or ( + isinstance(p, DTensor) and isinstance(p._local_tensor, MXFP8Tensor) + ): + p = p._local_tensor if isinstance(p, DTensor) else p + if p._rowwise_data is None or p._columnwise_data is None: + raise RuntimeError("MXFP8Tensor does not have one of rowwise/columnwise data.") + if self.capturable: + raise RuntimeError( + "FusedAdam does not support MXFP8 model weights with capturable=True." + ) + if self.master_weights: + p_main_of_mxfp8_model.append(unscaled_state["master_param"].data) + g_of_mxfp8_model.append(p_grad.data) + m_of_mxfp8_model.append(unscaled_state["exp_avg"]) + v_of_mxfp8_model.append(unscaled_state["exp_avg_sq"]) + p_mxfp8_rowwise.append(p._rowwise_data) + p_mxfp8_colwise.append(p._columnwise_data) + p_mxfp8_rowwise_scale_inv.append(p._rowwise_scale_inv) + p_mxfp8_colwise_scale_inv.append(p._columnwise_scale_inv) + out_dtype = p._fp8_dtype elif isinstance(p, QuantizedTensor) or ( isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) ): - # Block-scaling quantized params (MXFP8Tensor, Float8BlockwiseQTensor, + # Block-scaling quantized params (Float8BlockwiseQTensor, # NVFP4Tensor). Operate on FP32 master weights, requantize back after # Adam update. # Note: a fused Adam+requantize kernel (like multi_tensor_adam_fp8 @@ -797,6 +828,18 @@ def apply_multi_tensor_adam(adam_func, tensor_lists, inv_scale=None, out_dtype=N scale_invs, ] apply_multi_tensor_adam(self.multi_tensor_adam_fp8, tensor_lists, out_dtype) + if len(p_mxfp8_rowwise) > 0 and len(p_mxfp8_colwise) > 0: + tensor_lists = [ + g_of_mxfp8_model, + p_main_of_mxfp8_model, + m_of_mxfp8_model, + v_of_mxfp8_model, + p_mxfp8_rowwise, + p_mxfp8_colwise, + p_mxfp8_rowwise_scale_inv, + p_mxfp8_colwise_scale_inv, + ] + apply_multi_tensor_adam(self.multi_tensor_adam_mxfp8, tensor_lists, out_dtype) if len(p_f32_model) > 0: tensor_lists = [ g_of_f32_model, From 93e8b9aff82c8e8a473bf39698f7a8c9e24d681a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Wed, 18 Mar 2026 19:33:32 +0100 Subject: [PATCH 03/18] [PyTorch] torch.compile support for permutation functions (#2686) * init Signed-off-by: Pawel Gadzinski * work finished Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint fixes Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: root * removed warning.warn Signed-off-by: root * [PyTorch] Remove dead None-check for num_out_tokens in moe_permute_mask_map_forward num_out_tokens is typed as int in the custom_op signature and can never be None; the check was incorrectly carried over from the class-based upstream version during merge conflict resolution. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Signed-off-by: root Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 Signed-off-by: Varun Thumbe --- tests/pytorch/test_permutation.py | 148 +- transformer_engine/pytorch/permutation.py | 1630 ++++++++++------- .../pytorch/quantized_tensor.py | 12 + 3 files changed, 1137 insertions(+), 653 deletions(-) diff --git a/tests/pytorch/test_permutation.py b/tests/pytorch/test_permutation.py index be1ff30472..66c685e139 100644 --- a/tests/pytorch/test_permutation.py +++ b/tests/pytorch/test_permutation.py @@ -218,6 +218,17 @@ def backward_wrapper( return act.backward(backward_input, retain_graph=retain_graph) +def _maybe_compile(fn, use_torch_compile): + """Wrap fn with torch.compile(fullgraph=True) if requested.""" + if use_torch_compile: + torch._dynamo.reset() + import torch._functorch.config as functorch_config + + functorch_config.donated_buffer = False + return torch.compile(fn, fullgraph=True) + return fn + + def _test_permutation_index_map( te_dtype, num_tokens, @@ -227,6 +238,7 @@ def _test_permutation_index_map( num_out_tokens, with_probs, BENCHMARK=False, + use_torch_compile=False, ): if not with_probs and topK > 1: pytest.skip("Only permutations with topK=1 and without probabilities are supported.") @@ -298,9 +310,13 @@ def _test_permutation_index_map( te_permute_fwd_input.requires_grad_(True) te_permute_bwd_input = pytorch_permute_bwd_input.detach() - te_permute_output, row_id_map = te_permute( - te_permute_fwd_input, indices, num_out_tokens, map_type="index" + _permute = _maybe_compile( + lambda inp, idx, num_out, max_token: te_permute( + inp, idx, num_out, max_token, map_type="index" + ), + use_torch_compile, ) + te_permute_output, row_id_map = _permute(te_permute_fwd_input, indices, num_out_tokens, -1) te_permute_output.backward(te_permute_bwd_input, retain_graph=True) te_probs = None @@ -311,9 +327,11 @@ def _test_permutation_index_map( te_unpermute_fwd_input.requires_grad_(True) te_unpermute_bwd_input = pytorch_unpermute_bwd_input.detach() - te_unpermute_output = te_unpermute( - te_unpermute_fwd_input, row_id_map, te_probs, map_type="index" + _unpermute = _maybe_compile( + lambda inp, row_map, probs_val: te_unpermute(inp, row_map, probs_val, map_type="index"), + use_torch_compile, ) + te_unpermute_output = _unpermute(te_unpermute_fwd_input, row_id_map, te_probs) te_unpermute_output.backward(te_unpermute_bwd_input, retain_graph=True) ################################################################################################################################### @@ -444,6 +462,7 @@ def _test_permutation_mask_map( num_out_tokens, with_probs, BENCHMARK=False, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -514,9 +533,11 @@ def _test_permutation_mask_map( te_permute_fwd_input.requires_grad_(True) te_permute_bwd_input = pytorch_permute_bwd_input.detach() - te_permute_output, row_id_map = te_permute( - te_permute_fwd_input, routing_map, num_out_tokens=num_out_tokens, map_type="mask" + _permute = _maybe_compile( + lambda inp, rmap, n_out: te_permute(inp, rmap, num_out_tokens=n_out, map_type="mask"), + use_torch_compile, ) + te_permute_output, row_id_map = _permute(te_permute_fwd_input, routing_map, num_out_tokens) te_permute_output.backward(te_permute_bwd_input, retain_graph=True) te_probs = None @@ -527,9 +548,11 @@ def _test_permutation_mask_map( te_unpermute_fwd_input.requires_grad_(True) te_unpermute_bwd_input = pytorch_unpermute_bwd_input.detach() - te_unpermute_output = te_unpermute( - te_unpermute_fwd_input, row_id_map, te_probs, restore_shape, map_type="mask" + _unpermute = _maybe_compile( + lambda inp, row_map, p, rs: te_unpermute(inp, row_map, p, rs, map_type="mask"), + use_torch_compile, ) + te_unpermute_output = _unpermute(te_unpermute_fwd_input, row_id_map, te_probs, restore_shape) te_unpermute_output.backward(te_unpermute_bwd_input, retain_graph=True) ################################################################################################################################### @@ -666,6 +689,7 @@ def _test_permutation_and_padding_mask_map( with_merging_probs=False, align_size=16, BENCHMARK=False, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -957,6 +981,7 @@ def _test_permutation_and_padding_with_merging_probs( num_out_tokens, align_size=16, BENCHMARK=False, + use_torch_compile=False, ): """ Test the combination of merging_probs AND pad_offsets together in moe_unpermute. @@ -1180,6 +1205,7 @@ def _test_permutation_mask_map_fp8( topK, num_out_tokens, recipe, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -1255,9 +1281,11 @@ def _test_permutation_mask_map_fp8( ) # TE Permutation - permute_output, _ = te_permute( - permute_fwd_input_fp8, routing_map, num_out_tokens=num_out_tokens, map_type="mask" + _permute = _maybe_compile( + lambda inp, rmap, n_out: te_permute(inp, rmap, num_out_tokens=n_out, map_type="mask"), + use_torch_compile, ) + permute_output, _ = _permute(permute_fwd_input_fp8, routing_map, num_out_tokens) if recipe.float8_block_scaling(): te_permute_output = permute_output._rowwise_data te_permute_scale_output = permute_output._rowwise_scale_inv.T.contiguous() @@ -1291,6 +1319,7 @@ def _test_moe_chunk_sort( tp_size, hidden_size, BENCHMARK=False, + use_torch_compile=False, ): print( "chunk permute:" @@ -1340,7 +1369,11 @@ def _test_moe_chunk_sort( te_fwd_input.requires_grad_(True) te_bwd_input = pytorch_bwd_input.detach() - te_output = te_sort_chunks_by_index(te_fwd_input, split_sizes_cuda, sorted_idxs_cuda) + _sort = _maybe_compile( + lambda inp, ss, si: te_sort_chunks_by_index(inp, ss, si), + use_torch_compile, + ) + te_output = _sort(te_fwd_input, split_sizes_cuda, sorted_idxs_cuda) te_output.backward(te_bwd_input, retain_graph=True) ################################################################################################################################### @@ -1415,6 +1448,7 @@ def _test_permutation_mask_map_alongside_probs( num_out_tokens, tp_size, BENCHMARK=False, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -1510,30 +1544,27 @@ def _test_permutation_mask_map_alongside_probs( te_probs = probs.detach() te_probs.requires_grad_(True) - te_permute_output, te_permuted_probs, row_id_map = te_permute_with_probs( + def _alongside_probs_fn(fwd_inp, t_probs, rmap, ss1, si1, ss2, si2): + out, pprobs, rid = te_permute_with_probs( + fwd_inp, t_probs, rmap, num_out_tokens=num_out_tokens + ) + out, pprobs = te_sort_chunks_by_index_with_probs(out, pprobs, ss1, si1) + out_dtype = out.dtype + out = out * pprobs.unsqueeze(-1) + out = out.to(dtype=out_dtype) + out = te_sort_chunks_by_index(out, ss2, si2) + out = te_unpermute(out, rid, restore_shape=restore_shape, map_type="mask") + return out + + _fn = _maybe_compile(_alongside_probs_fn, use_torch_compile) + te_unpermute_output = _fn( te_permute_fwd_input, te_probs, routing_map, - num_out_tokens=num_out_tokens, - ) - - te_permute_output, te_permuted_probs = te_sort_chunks_by_index_with_probs( - te_permute_output, te_permuted_probs, split_sizes_cuda, sorted_idxs_cuda - ) - - te_permute_output_dtype = te_permute_output.dtype - te_permute_output = te_permute_output * te_permuted_probs.unsqueeze(-1) - te_permute_output = te_permute_output.to(dtype=te_permute_output_dtype) - - te_permute_output = te_sort_chunks_by_index( - te_permute_output, split_sizes_2_cuda, sorted_idxs_2_cuda - ) - - te_unpermute_output = te_unpermute( - te_permute_output, - row_id_map, - restore_shape=restore_shape, - map_type="mask", + split_sizes_cuda, + sorted_idxs_cuda, + split_sizes_2_cuda, + sorted_idxs_2_cuda, ) te_unpermute_output.backward(te_unpermute_bwd_input, retain_graph=True) @@ -1647,6 +1678,7 @@ def perf_test_cuda_kernel(cuda_kernel_fn): @pytest.mark.parametrize("hidden_size", [4096]) @pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_index_map( te_dtype, num_tokens, @@ -1654,7 +1686,10 @@ def test_permutation_index_map( hidden_size, topK, num_out_tokens, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or topK != 2): + pytest.skip("torch.compile tested with single config only") with_probs = True BENCHMARK = False @@ -1667,6 +1702,7 @@ def test_permutation_index_map( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1676,6 +1712,7 @@ def test_permutation_index_map( @pytest.mark.parametrize("hidden_size", [4096]) @pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_mask_map( te_dtype, num_tokens, @@ -1683,7 +1720,10 @@ def test_permutation_mask_map( hidden_size, topK, num_out_tokens, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or topK != 2): + pytest.skip("torch.compile tested with single config only") with_probs = True BENCHMARK = False @@ -1696,6 +1736,7 @@ def test_permutation_mask_map( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1711,6 +1752,7 @@ def test_permutation_mask_map( ], ) @pytest.mark.parametrize("with_merging_probs", [True, False]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_and_padding_mask_map( te_dtype, num_tokens, @@ -1719,7 +1761,10 @@ def test_permutation_and_padding_mask_map( topK, num_out_tokens, with_merging_probs, + use_torch_compile, ): + if use_torch_compile and (num_expert != 8 or topK != 2): + pytest.skip("torch.compile tested with single config only") BENCHMARK = False _test_permutation_and_padding_mask_map( @@ -1731,6 +1776,7 @@ def test_permutation_and_padding_mask_map( num_out_tokens=num_out_tokens, with_merging_probs=with_merging_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1745,6 +1791,7 @@ def test_permutation_and_padding_mask_map( (4096, 512, 9216, 8), ], ) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_and_padding_with_merging_probs( te_dtype, num_tokens, @@ -1752,8 +1799,11 @@ def test_permutation_and_padding_with_merging_probs( hidden_size, topK, num_out_tokens, + use_torch_compile, ): """Test moe_unpermute backward pass with BOTH merging_probs AND pad_offsets.""" + if use_torch_compile and (num_expert != 8 or topK != 2): + pytest.skip("torch.compile tested with single config only") BENCHMARK = False _test_permutation_and_padding_with_merging_probs( @@ -1764,11 +1814,13 @@ def test_permutation_and_padding_with_merging_probs( topK=topK, num_out_tokens=num_out_tokens, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @pytest.mark.parametrize("te_dtype", _te_dtypes) -def test_permutation_mask_map_empty_input(te_dtype): +@pytest.mark.parametrize("use_torch_compile", [False, True]) +def test_permutation_mask_map_empty_input(te_dtype, use_torch_compile): with_probs = True BENCHMARK = False @@ -1781,6 +1833,7 @@ def test_permutation_mask_map_empty_input(te_dtype): num_out_tokens=0, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1791,6 +1844,7 @@ def test_permutation_mask_map_empty_input(te_dtype): @pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) @pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_mask_map_alongside_probs( te_dtype, num_tokens, @@ -1799,7 +1853,10 @@ def test_permutation_mask_map_alongside_probs( topK, num_out_tokens, tp_size, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or topK != 2 or tp_size != 1): + pytest.skip("torch.compile tested with single config only") _test_permutation_mask_map_alongside_probs( te_dtype=te_dtype, num_tokens=num_tokens, @@ -1808,11 +1865,13 @@ def test_permutation_mask_map_alongside_probs( topK=topK, num_out_tokens=num_out_tokens, tp_size=tp_size, + use_torch_compile=use_torch_compile, ) @pytest.mark.parametrize("te_dtype", _te_dtypes) -def test_permutation_mask_map_alongside_probs_empty_input(te_dtype): +@pytest.mark.parametrize("use_torch_compile", [False, True]) +def test_permutation_mask_map_alongside_probs_empty_input(te_dtype, use_torch_compile): _test_permutation_mask_map_alongside_probs( te_dtype=te_dtype, num_tokens=0, @@ -1821,6 +1880,7 @@ def test_permutation_mask_map_alongside_probs_empty_input(te_dtype): topK=2, num_out_tokens=0, tp_size=2, + use_torch_compile=use_torch_compile, ) @@ -1868,6 +1928,7 @@ def test_permutation_mask_map_fp8( topK=topK, num_out_tokens=num_out_tokens, recipe=recipe, + use_torch_compile=False, # FP8 permutation is not yet supported under torch.compile ) @@ -1875,12 +1936,16 @@ def test_permutation_mask_map_fp8( @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_index_map_topk1_no_probs( te_dtype, num_tokens, num_expert, hidden_size, + use_torch_compile, ): + if use_torch_compile and num_expert != 7: + pytest.skip("torch.compile tested with single config only") topK = 1 num_out_tokens = None with_probs = False @@ -1895,6 +1960,7 @@ def test_permutation_index_map_topk1_no_probs( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1902,12 +1968,16 @@ def test_permutation_index_map_topk1_no_probs( @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_mask_map_topk1_no_probs( te_dtype, num_tokens, num_expert, hidden_size, + use_torch_compile, ): + if use_torch_compile and num_expert != 7: + pytest.skip("torch.compile tested with single config only") topK = 1 num_out_tokens = None with_probs = False @@ -1922,6 +1992,7 @@ def test_permutation_mask_map_topk1_no_probs( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1930,13 +2001,17 @@ def test_permutation_mask_map_topk1_no_probs( @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("tp_size", [2, 8]) @pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_chunk_permutation( te_dtype, num_tokens, num_expert, tp_size, hidden_size, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or tp_size != 2): + pytest.skip("torch.compile tested with single config only") BENCHMARK = False _test_moe_chunk_sort( @@ -1946,11 +2021,13 @@ def test_chunk_permutation( tp_size=tp_size, hidden_size=hidden_size, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @pytest.mark.parametrize("te_dtype", _te_dtypes) -def test_chunk_permutation_empty_input(te_dtype): +@pytest.mark.parametrize("use_torch_compile", [False, True]) +def test_chunk_permutation_empty_input(te_dtype, use_torch_compile): BENCHMARK = False _test_moe_chunk_sort( @@ -1960,6 +2037,7 @@ def test_chunk_permutation_empty_input(te_dtype): tp_size=2, hidden_size=4096, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index ca59a0ebf8..bc9a2660b7 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -6,11 +6,13 @@ import warnings from typing import Optional, Tuple import torch - import transformer_engine_torch as tex import transformer_engine.pytorch.triton.permutation as triton_permutation from transformer_engine.pytorch.constants import TE_DType -from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensor, + _quantized_tensor_passthrough_ops, +) from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor @@ -22,557 +24,829 @@ ] -class _moe_permute_index_map(torch.autograd.Function): - """functional Permute with index router map""" - - workspace = None - max_expanded_token_num = 0 - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - index: torch.Tensor, - num_out_tokens: int, - max_token_num: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # pylint: disable=missing-function-docstring - # Empty input check - if not inp.numel(): - return inp, torch.tensor([], device=inp.device) - - # Device check - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not index.is_cuda: - raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.") - # Shape check - if inp.size(0) != index.size(0): - raise ValueError( - f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " - f"index.size(0) ({index.size(0)})." - ) +# ===================== _moe_permute_index_map custom ops ===================== + +# Workspace state for moe_permute_index_map +_moe_permute_index_map_workspace = None +_moe_permute_index_map_max_expanded_token_num = 0 - # Data type check - dtype = TE_DType[inp.dtype] - if index.dtype != torch.int32: - warnings.warn( - f"The data type of the input `index` of Permute is {index.dtype}! " - "The recommended type is torch.int32." - ) - index = index.to(torch.int32) - topK = index.size(1) +@torch.library.custom_op("te_moe::permute_index_map", mutates_args=[]) +def moe_permute_index_map_forward( + inp: torch.Tensor, + index: torch.Tensor, + num_out_tokens: int, + max_token_num: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass for MoE permute with index router map.""" + global _moe_permute_index_map_workspace, _moe_permute_index_map_max_expanded_token_num - input_max_expanded_token_num = max(max_token_num, inp.size(0)) * topK - if _moe_permute_index_map.max_expanded_token_num < input_max_expanded_token_num: - _moe_permute_index_map.max_expanded_token_num = input_max_expanded_token_num - _moe_permute_index_map.workspace = [] + if not inp.numel(): + return inp.clone(), torch.tensor([], device=inp.device) - permuted_act, row_id_map, _moe_permute_index_map.workspace = tex.moe_permute_fwd( - inp, - dtype, - index, - num_out_tokens, - _moe_permute_index_map.workspace, - _moe_permute_index_map.max_expanded_token_num, + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not index.is_cuda: + raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.") + if inp.size(0) != index.size(0): + raise ValueError( + f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " + f"index.size(0) ({index.size(0)})." + ) + if index.dtype != torch.int32: + warnings.warn( + f"The data type of the input `index` of Permute is {index.dtype}! " + "The recommended type is torch.int32." ) + index = index.to(torch.int32) - ctx.row_id_map = row_id_map - ctx.num_tokens = index.size(0) - ctx.topK = index.size(1) - return permuted_act, row_id_map - - @staticmethod - def backward( - ctx, - permuted_act_grad: torch.Tensor, - _, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - # Empty input check - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, None - - if not permuted_act_grad.is_contiguous(): - permuted_act_grad = permuted_act_grad.contiguous() - - dtype = TE_DType[permuted_act_grad.dtype] - act_grad = None - if ctx.needs_input_grad[0]: - act_grad = tex.moe_permute_bwd( - permuted_act_grad, dtype, ctx.row_id_map, torch.empty(0), ctx.num_tokens, ctx.topK - ) + dtype = TE_DType[inp.dtype] - return act_grad, None, None, None + topK = index.size(1) + input_max_expanded_token_num = max(max_token_num, inp.size(0)) * topK + if _moe_permute_index_map_max_expanded_token_num < input_max_expanded_token_num: + _moe_permute_index_map_max_expanded_token_num = input_max_expanded_token_num + _moe_permute_index_map_workspace = [] -class _moe_unpermute_index_map(torch.autograd.Function): - """functional Unpermute with index router map""" + permuted_act, row_id_map, _moe_permute_index_map_workspace = tex.moe_permute_fwd( + inp, + dtype, + index, + num_out_tokens, + _moe_permute_index_map_workspace, + _moe_permute_index_map_max_expanded_token_num, + ) - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - row_id_map: torch.Tensor, - probs: torch.Tensor, - ) -> torch.Tensor: - # pylint: disable=missing-function-docstring - # Empty input check - if not inp.numel(): - ctx.probs = probs - return inp + return permuted_act, row_id_map - # None probs check - if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") - if probs.dtype != torch.float32: - warnings.warn( - f"The data type of the input `probs` of Unpermute is {probs.dtype}! " - "The recommended type is torch.float32." - ) - probs = probs.to(torch.float32) +@moe_permute_index_map_forward.register_fake +def _moe_permute_index_map_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + index: torch.Tensor, + num_out_tokens: int, + max_token_num: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake implementation for shape inference.""" + num_tokens = inp.shape[0] + topK = index.shape[1] - num_tokens = probs.size(0) - topK = probs.size(1) - else: - num_tokens = row_id_map.size(0) - topK = 1 - probs = torch.empty(0) + # Infer output shape + output_tokens = num_out_tokens if num_out_tokens > 0 else num_tokens * topK - # Device check - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not row_id_map.is_cuda: - raise ValueError( - f"row_id_map must be a CUDA tensor, but got tensor on {row_id_map.device}." - ) + # row_id_map is 1D with size = num_tokens * topK + fake_output = torch.empty((output_tokens, inp.shape[1]), dtype=inp.dtype, device=inp.device) + fake_row_id_map = torch.empty((num_tokens * topK,), dtype=torch.int32, device=inp.device) - # Data type check - dtype = TE_DType[inp.dtype] - if row_id_map.dtype != torch.int32: - warnings.warn( - f"The data type of the input `row_id_map` of Unpermute is {row_id_map.dtype}! " - "The recommended type is torch.int32." - ) - row_id_map = row_id_map.to(torch.int32) + return fake_output, fake_row_id_map + + +@torch.library.custom_op("te_moe::permute_index_map_bwd", mutates_args=[]) +def moe_permute_index_map_backward( + grad_permuted_act: torch.Tensor, + row_id_map: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Backward pass for MoE permute with index router map.""" + dtype = TE_DType[grad_permuted_act.dtype] + act_grad = tex.moe_permute_bwd( + grad_permuted_act, dtype, row_id_map, torch.empty(0), num_tokens, topK + ) + return act_grad + + +@moe_permute_index_map_backward.register_fake +def _moe_permute_index_map_backward_fake( # pylint: disable=unused-argument + grad_permuted_act: torch.Tensor, + row_id_map: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Fake implementation for shape inference of backward.""" + return torch.empty( + (num_tokens, grad_permuted_act.shape[1]), + dtype=grad_permuted_act.dtype, + device=grad_permuted_act.device, + ) + + +def _moe_permute_index_map_setup_context(ctx, inputs, output): + """Save context for backward pass.""" + inp, index, _num_out_tokens, _max_token_num = inputs + _permuted_act, row_id_map = output + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(row_id_map) + ctx.num_tokens = index.size(0) if not ctx.empty_input else 0 + ctx.topK = index.size(1) if not ctx.empty_input else 1 + + +def _moe_permute_index_map_backward_wrapper( + ctx, grad_permuted_act, grad_row_id_map +): # pylint: disable=unused-argument + """Backward pass wrapper that calls the custom backward op.""" + if ctx.empty_input: + return grad_permuted_act, None, None, None + + if not grad_permuted_act.is_contiguous(): + grad_permuted_act = grad_permuted_act.contiguous() + + (row_id_map,) = ctx.saved_tensors + act_grad = torch.ops.te_moe.permute_index_map_bwd( + grad_permuted_act, row_id_map, ctx.num_tokens, ctx.topK + ) - unpermuted_output = tex.moe_unpermute_fwd(inp, dtype, row_id_map, probs, num_tokens, topK) + return act_grad, None, None, None - ctx.save_for_backward(inp, row_id_map, probs) - return unpermuted_output - @staticmethod - def backward( - ctx, - unpermuted_act_grad: torch.Tensor, - ) -> Tuple[torch.Tensor, None, torch.Tensor]: - # pylint: disable=missing-function-docstring - # Empty input check - if not unpermuted_act_grad.numel(): - return unpermuted_act_grad, None, ctx.probs +moe_permute_index_map_forward.register_autograd( + _moe_permute_index_map_backward_wrapper, + setup_context=_moe_permute_index_map_setup_context, +) - if not unpermuted_act_grad.is_contiguous(): - unpermuted_act_grad = unpermuted_act_grad.contiguous() - dtype = TE_DType[unpermuted_act_grad.dtype] - inp, row_id_map, probs = ctx.saved_tensors +# ===================== _moe_unpermute_index_map custom ops ===================== - act_grad = None + +@torch.library.custom_op("te_moe::unpermute_index_map_fwd", mutates_args=[]) +def moe_unpermute_index_map_forward( + inp: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Forward pass for MoE unpermute with index router map.""" + if not inp.numel(): + return inp.clone() + dtype = TE_DType[inp.dtype] + return tex.moe_unpermute_fwd(inp, dtype, row_id_map, probs, num_tokens, topK) + + +@moe_unpermute_index_map_forward.register_fake +def _moe_unpermute_index_map_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Fake implementation for shape inference.""" + # Output shape: (num_tokens, hidden_size) + return torch.empty((num_tokens, inp.shape[1]), dtype=inp.dtype, device=inp.device) + + +@torch.library.custom_op("te_moe::unpermute_index_map_bwd", mutates_args=[]) +def moe_unpermute_index_map_backward( + unpermuted_act_grad: torch.Tensor, + fwd_input: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE unpermute with index router map.""" + dtype = TE_DType[unpermuted_act_grad.dtype] + act_grad, prob_grad = tex.moe_unpermute_bwd( + unpermuted_act_grad, fwd_input, dtype, row_id_map, probs + ) + return act_grad, prob_grad + + +@moe_unpermute_index_map_backward.register_fake +def _moe_unpermute_index_map_backward_fake( + unpermuted_act_grad: torch.Tensor, + fwd_input: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake implementation for shape inference of backward.""" + # act_grad shape: (fwd_input.size(0), hidden_size) + # prob_grad shape: (num_tokens, topK) + topK = probs.size(1) if probs.numel() > 0 else 1 + num_tokens = probs.size(0) if probs.numel() > 0 else row_id_map.size(0) + act_grad = torch.empty( + (fwd_input.size(0), unpermuted_act_grad.shape[1]), + dtype=unpermuted_act_grad.dtype, + device=unpermuted_act_grad.device, + ) + prob_grad = torch.empty( + (num_tokens, topK), dtype=torch.float32, device=unpermuted_act_grad.device + ) + return act_grad, prob_grad + + +def _moe_unpermute_index_map_setup_context(ctx, inputs, output): # pylint: disable=unused-argument + """Save context for backward pass.""" + inp, row_id_map, probs, _num_tokens, _topK = inputs + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(inp, row_id_map, probs) + ctx.needs_probs_grad = probs.requires_grad + + +def _moe_unpermute_index_map_backward_wrapper(ctx, unpermuted_act_grad): + """Backward pass wrapper that calls the custom backward op.""" + if ctx.empty_input: + prob_grad = torch.zeros_like(ctx.saved_tensors[2]) if ctx.needs_probs_grad else None + return unpermuted_act_grad, None, prob_grad, None, None + + if not unpermuted_act_grad.is_contiguous(): + unpermuted_act_grad = unpermuted_act_grad.contiguous() + + inp, row_id_map, probs = ctx.saved_tensors + + act_grad, prob_grad = torch.ops.te_moe.unpermute_index_map_bwd( + unpermuted_act_grad, inp, row_id_map, probs + ) + + if not ctx.needs_probs_grad: prob_grad = None - if ctx.needs_input_grad[0]: - act_grad, prob_grad = tex.moe_unpermute_bwd( - unpermuted_act_grad, inp, dtype, row_id_map, probs - ) - if not ctx.needs_input_grad[2]: - prob_grad = None - - return act_grad, None, prob_grad - - -class _moe_permute_mask_map(torch.autograd.Function): - """functional Permute with mask router map""" - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - routing_map: torch.Tensor, - num_out_tokens: int, - probs: torch.Tensor, - pad_offsets: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - # pylint: disable=missing-function-docstring - if not inp.numel(): - ctx.probs = probs - return inp, torch.tensor([], device=inp.device), torch.tensor([], device=inp.device) - - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not routing_map.is_cuda: + + return act_grad, None, prob_grad, None, None + + +moe_unpermute_index_map_forward.register_autograd( + _moe_unpermute_index_map_backward_wrapper, + setup_context=_moe_unpermute_index_map_setup_context, +) + + +# ===================== _moe_permute_mask_map custom ops ===================== + + +@torch.library.custom_op("te_moe::permute_mask_map_fwd", mutates_args=[]) +def moe_permute_mask_map_forward( + inp: torch.Tensor, + routing_map: torch.Tensor, + num_out_tokens: int, + probs: Optional[torch.Tensor], + pad_offsets: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass for MoE permute with mask router map.""" + if not inp.numel(): + return inp.clone(), torch.tensor([], device=inp.device), torch.tensor([], device=inp.device) + + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not routing_map.is_cuda: + raise ValueError( + f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}." + ) + if probs is not None: + if not probs.is_cuda: + raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") + if pad_offsets is not None: + if not pad_offsets.is_cuda: raise ValueError( - f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}." + f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." ) - if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") - if pad_offsets is not None: - if not pad_offsets.is_cuda: + if inp.size(0) != routing_map.size(0): + raise ValueError( + f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " + f"routing_map.size(0) ({routing_map.size(0)})." + ) + num_tokens, hidden_size = inp.size() + num_experts = routing_map.size(1) + + row_id_map = triton_permutation.make_row_id_map(routing_map, num_tokens, num_experts) + + # FP8 handling + fp8 = isinstance(inp, QuantizedTensor) + per_tensor_recipe = isinstance(inp, Float8Tensor) + blockwise_recipe = isinstance(inp, Float8BlockwiseQTensor) + mxfp8_recipe = isinstance(inp, MXFP8Tensor) + + if fp8: + fp8_dtype = inp._fp8_dtype + fake_dtype = inp.dtype + # blockwise scaling + if blockwise_recipe: + fp8_scale = inp._rowwise_scale_inv.T.contiguous() + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: raise ValueError( - f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Input shape: ({num_tokens}, {hidden_size}), " + f"scale shape: {tuple(fp8_scale.shape)}." ) - - if inp.size(0) != routing_map.size(0): - raise ValueError( - f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " - f"routing_map.size(0) ({routing_map.size(0)})." - ) - num_tokens, hidden_size = inp.size() - num_experts = routing_map.size(1) - if num_out_tokens is None: - raise ValueError("num_out_tokens must be provided to the fused permute function.") - - row_id_map = triton_permutation.make_row_id_map(routing_map, num_tokens, num_experts) - - fp8 = isinstance(inp, QuantizedTensor) - per_tensor_recipe = isinstance(inp, Float8Tensor) - blockwise_recipe = isinstance(inp, Float8BlockwiseQTensor) - mxfp8_recipe = isinstance(inp, MXFP8Tensor) - - if fp8: - fp8_dtype = inp._fp8_dtype - fake_dtype = inp.dtype - # blockwise scaling - if blockwise_recipe: - fp8_scale = inp._rowwise_scale_inv.T.contiguous() - scale_hidden_dim = fp8_scale.shape[1] - if num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Input shape: ({num_tokens}, {hidden_size}), " - f"scale shape: {tuple(fp8_scale.shape)}." - ) - inp = inp._rowwise_data - # mxfp8 scaling - elif mxfp8_recipe: - fp8_scale = inp._rowwise_scale_inv.contiguous() - scale_hidden_dim = fp8_scale.shape[1] - if num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Input shape: ({num_tokens}, {hidden_size}), " - f"scale shape: {tuple(fp8_scale.shape)}." - ) - inp = inp._rowwise_data - # per-tensor scaling - elif per_tensor_recipe: - # Kernel does not need scale in per-tensor scaling - fp8_scale = None - scale_hidden_dim = None - fp8_scale_inv = inp._scale_inv - inp = inp._data - else: - raise ValueError("Unsupported FP8 recipe") - else: + inp = inp._rowwise_data + # mxfp8 scaling + elif mxfp8_recipe: + fp8_scale = inp._rowwise_scale_inv.contiguous() + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: + raise ValueError( + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Input shape: ({num_tokens}, {hidden_size}), " + f"scale shape: {tuple(fp8_scale.shape)}." + ) + inp = inp._rowwise_data + # per-tensor scaling + elif per_tensor_recipe: + # Kernel does not need scale in per-tensor scaling fp8_scale = None - fp8_dtype = None scale_hidden_dim = None + fp8_scale_inv = inp._scale_inv + inp = inp._data + else: + raise ValueError("Unsupported FP8 recipe") + else: + fp8_scale = None + fp8_dtype = None + scale_hidden_dim = None + + output, permuted_scale, permuted_probs = triton_permutation.permute_with_mask_map( + inp, + row_id_map, + probs, + fp8_scale, + pad_offsets, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + scale_hidden_dim, + ) - output, permuted_scale, permuted_probs = triton_permutation.permute_with_mask_map( - inp, - row_id_map, - probs, - fp8_scale, - pad_offsets, - num_tokens, - num_experts, - num_out_tokens, - hidden_size, - scale_hidden_dim, - ) + if fp8: + if per_tensor_recipe: + output = Float8Tensor( + data=output, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=output.shape, + dtype=fake_dtype, + ) + elif blockwise_recipe: + output = Float8BlockwiseQTensor( + shape=output.shape, + dtype=fake_dtype, + rowwise_data=output, + rowwise_scale_inv=permuted_scale.T.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=fp8_dtype, + quantizer=None, + is_2D_scaled=False, + requires_grad=output.requires_grad, + ) + elif mxfp8_recipe: + output = MXFP8Tensor( + shape=output.shape, + dtype=fake_dtype, + fp8_dtype=fp8_dtype, + rowwise_data=output, + rowwise_scale_inv=permuted_scale.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=None, + requires_grad=output.requires_grad, + with_gemm_swizzled_scales=False, + ) - if fp8: - if per_tensor_recipe: - output = Float8Tensor( - data=output, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=output.shape, - dtype=fake_dtype, - ) - elif blockwise_recipe: - output = Float8BlockwiseQTensor( - shape=output.shape, - dtype=fake_dtype, - rowwise_data=output, - rowwise_scale_inv=permuted_scale.T.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - fp8_dtype=fp8_dtype, - quantizer=None, - is_2D_scaled=False, - requires_grad=output.requires_grad, - ) - elif mxfp8_recipe: - output = MXFP8Tensor( - shape=output.shape, - dtype=fake_dtype, - fp8_dtype=fp8_dtype, - rowwise_data=output, - rowwise_scale_inv=permuted_scale.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - quantizer=None, - requires_grad=output.requires_grad, - with_gemm_swizzled_scales=False, - ) + # If permuted_probs is None, return empty tensor (custom ops need concrete tensors) + if permuted_probs is None: + permuted_probs = torch.empty(0, device=inp.device) - ctx.save_for_backward(row_id_map, pad_offsets) - ctx.num_experts = num_experts - ctx.num_tokens = num_tokens - ctx.hidden_size = hidden_size - return output, row_id_map, permuted_probs - - @staticmethod - def backward( - ctx, - permuted_act_grad: torch.Tensor, - _, - permuted_probs_grad: torch.Tensor, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, ctx.probs, None - - act_grad = None - probs_grad = None - if ctx.needs_input_grad[0]: - row_id_map, pad_offsets = ctx.saved_tensors - if isinstance(permuted_act_grad, QuantizedTensor): - raise TypeError( - "The backward of moe_permute does not support FP8, but got " - f"QuantizedTensor of type {type(permuted_act_grad).__name__}." - ) - act_grad, probs_grad = triton_permutation.unpermute_with_mask_map( - permuted_act_grad, - row_id_map, - None, - permuted_probs_grad, - pad_offsets, - ctx.num_tokens, - ctx.num_experts, - ctx.hidden_size, + return output, row_id_map, permuted_probs + + +@moe_permute_mask_map_forward.register_fake +def _moe_permute_mask_map_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + routing_map: torch.Tensor, + num_out_tokens: int, + probs: Optional[torch.Tensor], + pad_offsets: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fake implementation for shape inference.""" + num_tokens = inp.shape[0] + hidden_size = inp.shape[1] + num_experts = routing_map.shape[1] + # row_id_map: (num_tokens, num_experts * 2 + 1) + fake_output = torch.empty((num_out_tokens, hidden_size), dtype=inp.dtype, device=inp.device) + fake_row_id_map = torch.empty( + (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=inp.device + ) + if probs is not None: + fake_permuted_probs = torch.empty((num_out_tokens,), dtype=probs.dtype, device=inp.device) + else: + fake_permuted_probs = torch.empty(0, device=inp.device) + return fake_output, fake_row_id_map, fake_permuted_probs + + +@torch.library.custom_op("te_moe::permute_mask_map_bwd", mutates_args=[]) +def moe_permute_mask_map_backward( + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE permute with mask router map.""" + act_grad, probs_grad = triton_permutation.unpermute_with_mask_map( + permuted_act_grad, + row_id_map, + None, + permuted_probs_grad, + pad_offsets, + num_tokens, + num_experts, + hidden_size, + ) + if probs_grad is None: + probs_grad = torch.empty(0, device=permuted_act_grad.device) + return act_grad, probs_grad + + +@moe_permute_mask_map_backward.register_fake +def _moe_permute_mask_map_backward_fake( # pylint: disable=unused-argument + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake for backward shape inference.""" + act_grad = torch.empty( + (num_tokens, hidden_size), dtype=permuted_act_grad.dtype, device=permuted_act_grad.device + ) + if permuted_probs_grad is not None: + probs_grad = torch.empty( + (num_tokens, num_experts), + dtype=permuted_probs_grad.dtype, + device=permuted_act_grad.device, + ) + else: + probs_grad = torch.empty(0, device=permuted_act_grad.device) + return act_grad, probs_grad + + +def _moe_permute_mask_map_setup_context(ctx, inputs, output): + """Save context for backward pass.""" + inp, routing_map, _num_out_tokens, probs, pad_offsets = inputs + _output_tensor, row_id_map, _permuted_probs = output + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(row_id_map, pad_offsets) + ctx.num_experts = routing_map.size(1) + ctx.num_tokens = inp.size(0) + ctx.hidden_size = inp.size(1) if not ctx.empty_input else 0 + ctx.needs_probs_grad = probs is not None and probs.requires_grad + + +def _moe_permute_mask_map_backward_wrapper( + ctx, grad_output, grad_row_id_map, grad_permuted_probs +): # pylint: disable=unused-argument + """Backward wrapper calling the custom backward op.""" + if ctx.empty_input: + if ctx.needs_probs_grad: + probs_grad = torch.zeros( + (ctx.num_tokens, ctx.num_experts), + dtype=grad_permuted_probs.dtype, + device=grad_permuted_probs.device, ) - if not ctx.needs_input_grad[3]: + else: probs_grad = None - return act_grad, None, None, probs_grad, None - - -class _moe_unpermute_mask_map(torch.autograd.Function): - """functional Unpermute with mask router map""" - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - row_id_map: torch.Tensor, - merging_probs: Optional[torch.Tensor], - restore_shape: Optional[torch.Size], - pad_offsets: Optional[torch.Tensor], - ) -> torch.Tensor: - # pylint: disable=missing-function-docstring - if not inp.numel(): - ctx.merging_probs = merging_probs - return inp + return grad_output, None, None, probs_grad, None - if restore_shape is None: - restore_shape = inp.shape - num_tokens, hidden_size = restore_shape - num_experts = (row_id_map.size(1) - 1) // 2 + assert not isinstance( + grad_output, QuantizedTensor + ), "The backward of moe_permute does not support FP8." + + row_id_map, pad_offsets = ctx.saved_tensors + + # Pass permuted_probs_grad only if it has content + probs_grad_input = grad_permuted_probs if grad_permuted_probs.numel() > 0 else None + + act_grad, probs_grad = torch.ops.te_moe.permute_mask_map_bwd( + grad_output, + probs_grad_input, + row_id_map, + pad_offsets, + ctx.num_tokens, + ctx.num_experts, + ctx.hidden_size, + ) + + if not ctx.needs_probs_grad or probs_grad.numel() == 0: + probs_grad = None + + return act_grad, None, None, probs_grad, None - with_probs = merging_probs is not None - if with_probs: - if not merging_probs.is_cuda: + +moe_permute_mask_map_forward.register_autograd( + _moe_permute_mask_map_backward_wrapper, + setup_context=_moe_permute_mask_map_setup_context, +) + + +# ===================== _moe_unpermute_mask_map custom ops ===================== + + +@torch.library.custom_op("te_moe::unpermute_mask_map_fwd", mutates_args=[]) +def moe_unpermute_mask_map_forward( + inp: torch.Tensor, + row_id_map: torch.Tensor, + merging_probs: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, + pad_offsets: Optional[torch.Tensor], +) -> torch.Tensor: + """Forward pass for MoE unpermute with mask router map.""" + if not inp.numel(): + return inp.clone() + assert not isinstance( + inp, QuantizedTensor + ), "The forward of moe_unpermute does not support FP8." + unpermuted_output, _ = triton_permutation.unpermute_with_mask_map( + inp, + row_id_map, + merging_probs, + None, + pad_offsets, + num_tokens, + num_experts, + hidden_size, + ) + return unpermuted_output + + +@moe_unpermute_mask_map_forward.register_fake +def _moe_unpermute_mask_map_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + row_id_map: torch.Tensor, + merging_probs: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, + pad_offsets: Optional[torch.Tensor], +) -> torch.Tensor: + """Fake implementation for shape inference.""" + return torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=inp.device) + + +@torch.library.custom_op("te_moe::unpermute_mask_map_bwd_with_probs", mutates_args=[]) +def moe_unpermute_mask_map_backward_with_probs( + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + fwd_input: torch.Tensor, + merging_probs: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE unpermute with merging probs.""" + act_grad, probs_grad = triton_permutation.unpermute_with_mask_map_bwd_with_merging_probs( + unpermuted_act_grad, + row_id_map, + fwd_input, + merging_probs, + pad_offsets, + num_tokens, + num_experts, + num_permuted_tokens, + hidden_size, + ) + return act_grad, probs_grad + + +@moe_unpermute_mask_map_backward_with_probs.register_fake +def _moe_unpermute_mask_map_bwd_with_probs_fake( # pylint: disable=unused-argument + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + fwd_input: torch.Tensor, + merging_probs: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake for backward shape inference with merging probs.""" + act_grad = torch.empty( + (num_permuted_tokens, hidden_size), + dtype=unpermuted_act_grad.dtype, + device=unpermuted_act_grad.device, + ) + probs_grad = torch.empty( + (num_tokens, num_experts), + dtype=merging_probs.dtype, + device=unpermuted_act_grad.device, + ) + return act_grad, probs_grad + + +@torch.library.custom_op("te_moe::unpermute_mask_map_bwd_no_probs", mutates_args=[]) +def moe_unpermute_mask_map_backward_no_probs( + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> torch.Tensor: + """Backward pass for MoE unpermute without merging probs (permute grad back).""" + # FP8 handling + fp8 = isinstance(unpermuted_act_grad, QuantizedTensor) + per_tensor_recipe = isinstance(unpermuted_act_grad, Float8Tensor) + blockwise_recipe = isinstance(unpermuted_act_grad, Float8BlockwiseQTensor) + mxfp8_recipe = isinstance(unpermuted_act_grad, MXFP8Tensor) + + if fp8: + fp8_dtype = unpermuted_act_grad._fp8_dtype + fake_dtype = unpermuted_act_grad.dtype + if per_tensor_recipe: + fp8_scale = None + scale_hidden_dim = None + fp8_scale_inv = unpermuted_act_grad._scale_inv + unpermuted_act_grad = unpermuted_act_grad._data + # blockwise scaling + elif blockwise_recipe: + fp8_scale = unpermuted_act_grad._rowwise_scale_inv.T.contiguous() + unpermuted_act_grad = unpermuted_act_grad._rowwise_data + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: raise ValueError( - "merging_probs must be a CUDA tensor, but got tensor on " - f"{merging_probs.device}." + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Scale shape: {tuple(fp8_scale.shape)}." ) - - # Device check - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not row_id_map.is_cuda: - raise ValueError( - f"row_id_map must be a CUDA tensor, but got tensor on {row_id_map.device}." - ) - if pad_offsets is not None: - if not pad_offsets.is_cuda: + # mxfp8 scaling + elif mxfp8_recipe: + fp8_scale = unpermuted_act_grad._rowwise_scale_inv.contiguous() + unpermuted_act_grad = unpermuted_act_grad._rowwise_data + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: raise ValueError( - f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Scale shape: {tuple(fp8_scale.shape)}." ) + else: + raise ValueError("Unsupported FP8 recipe") + else: + scale_hidden_dim = None + fp8_dtype = None + fp8_scale = None + + act_grad, permuted_scale, _ = triton_permutation.permute_with_mask_map( + unpermuted_act_grad, + row_id_map, + None, + fp8_scale, + pad_offsets, + num_tokens, + num_experts, + num_permuted_tokens, + hidden_size, + scale_hidden_dim, + ) - if isinstance(inp, QuantizedTensor): - raise TypeError( - "The forward of moe_unpermute does not support FP8, but got " - f"QuantizedTensor of type {type(inp).__name__}." + if fp8: + if per_tensor_recipe: + act_grad = Float8Tensor( + data=act_grad, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=act_grad.shape, + dtype=fake_dtype, ) - unpermuted_output, _ = triton_permutation.unpermute_with_mask_map( - inp, + elif blockwise_recipe: + act_grad = Float8BlockwiseQTensor( + shape=act_grad.shape, + dtype=fake_dtype, + rowwise_data=act_grad, + rowwise_scale_inv=permuted_scale.T.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=fp8_dtype, + quantizer=None, + is_2D_scaled=False, + requires_grad=act_grad.requires_grad, + ) + elif mxfp8_recipe: + act_grad = MXFP8Tensor( + shape=act_grad.shape, + dtype=fake_dtype, + fp8_dtype=fp8_dtype, + rowwise_data=act_grad, + rowwise_scale_inv=permuted_scale.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=None, + requires_grad=act_grad.requires_grad, + with_gemm_swizzled_scales=False, + ) + + return act_grad + + +@moe_unpermute_mask_map_backward_no_probs.register_fake +def _moe_unpermute_mask_map_bwd_no_probs_fake( # pylint: disable=unused-argument + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> torch.Tensor: + """Fake for backward shape inference without probs.""" + return torch.empty( + (num_permuted_tokens, hidden_size), + dtype=unpermuted_act_grad.dtype, + device=unpermuted_act_grad.device, + ) + + +def _moe_unpermute_mask_map_setup_context(ctx, inputs, output): # pylint: disable=unused-argument + """Save context for backward pass.""" + inp, row_id_map, merging_probs, num_tokens, num_experts, hidden_size, pad_offsets = inputs + ctx.empty_input = inp.size(0) == 0 + ctx.num_experts = num_experts + ctx.num_tokens = num_tokens + ctx.num_permuted_tokens = inp.size(0) + ctx.hidden_size = hidden_size + ctx.with_probs = merging_probs is not None + if ctx.with_probs: + ctx.save_for_backward(inp, row_id_map, merging_probs, pad_offsets) + ctx.needs_probs_grad = merging_probs.requires_grad + else: + ctx.save_for_backward(row_id_map, pad_offsets) + ctx.needs_probs_grad = False + + +def _moe_unpermute_mask_map_backward_wrapper(ctx, unpermuted_act_grad): + """Backward wrapper calling the appropriate custom backward op.""" + if ctx.empty_input: + if ctx.with_probs: + _, _, merging_probs, _ = ctx.saved_tensors + probs_grad = torch.zeros_like(merging_probs) if ctx.needs_probs_grad else None + return unpermuted_act_grad, None, probs_grad, None, None, None, None + return unpermuted_act_grad, None, None, None, None, None, None + + act_grad = None + probs_grad = None + + if ctx.with_probs: + fwd_input, row_id_map, merging_probs, pad_offsets = ctx.saved_tensors + assert not isinstance( + unpermuted_act_grad, QuantizedTensor + ), "The backward of moe_unpermute with merging probs does not support FP8." + act_grad, probs_grad = torch.ops.te_moe.unpermute_mask_map_bwd_with_probs( + unpermuted_act_grad, row_id_map, + fwd_input, merging_probs, - None, pad_offsets, - num_tokens, - num_experts, - hidden_size, + ctx.num_tokens, + ctx.num_experts, + ctx.num_permuted_tokens, + ctx.hidden_size, + ) + else: + row_id_map, pad_offsets = ctx.saved_tensors + act_grad = torch.ops.te_moe.unpermute_mask_map_bwd_no_probs( + unpermuted_act_grad, + row_id_map, + pad_offsets, + ctx.num_tokens, + ctx.num_experts, + ctx.num_permuted_tokens, + ctx.hidden_size, ) - if with_probs: - ctx.save_for_backward(inp, row_id_map, merging_probs, pad_offsets) - else: - ctx.save_for_backward(row_id_map, pad_offsets) - ctx.num_experts = num_experts - ctx.num_tokens = num_tokens - ctx.num_permuted_tokens = inp.size(0) - ctx.hidden_size = hidden_size - ctx.with_probs = with_probs - return unpermuted_output - - @staticmethod - def backward(ctx, unpermuted_act_grad): - # pylint: disable=missing-function-docstring - if not unpermuted_act_grad.numel(): - return unpermuted_act_grad, None, ctx.merging_probs, None, None - - act_grad = None + if not ctx.needs_probs_grad: probs_grad = None - if ctx.needs_input_grad[0]: - if ctx.with_probs: - fwd_input, row_id_map, merging_probs, pad_offsets = ctx.saved_tensors - else: - row_id_map, pad_offsets = ctx.saved_tensors - - fp8 = isinstance(unpermuted_act_grad, QuantizedTensor) - per_tensor_recipe = isinstance(unpermuted_act_grad, Float8Tensor) - blockwise_recipe = isinstance(unpermuted_act_grad, Float8BlockwiseQTensor) - mxfp8_recipe = isinstance(unpermuted_act_grad, MXFP8Tensor) - - if fp8: - fp8_dtype = unpermuted_act_grad._fp8_dtype - fake_dtype = unpermuted_act_grad.dtype - # per-tensor scaling - if per_tensor_recipe: - # Kernel does not need scale in per-tensor scaling - fp8_scale = None - scale_hidden_dim = None - fp8_scale_inv = unpermuted_act_grad._scale_inv - unpermuted_act_grad = unpermuted_act_grad._data - # blockwise scaling - elif blockwise_recipe: - fp8_scale = unpermuted_act_grad._rowwise_scale_inv.T.contiguous() - unpermuted_act_grad = unpermuted_act_grad._rowwise_data - scale_hidden_dim = fp8_scale.shape[1] - if ctx.num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({ctx.num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Scale shape: {tuple(fp8_scale.shape)}." - ) - # mxfp8 scaling - elif mxfp8_recipe: - fp8_scale = unpermuted_act_grad._rowwise_scale_inv.contiguous() - unpermuted_act_grad = unpermuted_act_grad._rowwise_data - scale_hidden_dim = fp8_scale.shape[1] - if ctx.num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({ctx.num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Scale shape: {tuple(fp8_scale.shape)}." - ) - else: - raise ValueError("Unsupported FP8 recipe") - else: - scale_hidden_dim = None - fp8_dtype = None - fp8_scale = None - - permuted_scale = None - if ctx.with_probs: - if fp8: - raise TypeError( - "The backward of moe_unpermute with merging probs does not support FP8, " - f"but got FP8 gradient with dtype {fp8_dtype}." - ) - act_grad, probs_grad = ( - triton_permutation.unpermute_with_mask_map_bwd_with_merging_probs( - unpermuted_act_grad, - row_id_map, - fwd_input, - merging_probs, - pad_offsets, - ctx.num_tokens, - ctx.num_experts, - ctx.num_permuted_tokens, - ctx.hidden_size, - ) - ) - else: - act_grad, permuted_scale, _ = triton_permutation.permute_with_mask_map( - unpermuted_act_grad, - row_id_map, - None, - fp8_scale, - pad_offsets, - ctx.num_tokens, - ctx.num_experts, - ctx.num_permuted_tokens, - ctx.hidden_size, - scale_hidden_dim, - ) - if fp8: - if per_tensor_recipe: - act_grad = Float8Tensor( - data=act_grad, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=act_grad.shape, - dtype=fake_dtype, - ) - elif blockwise_recipe: - act_grad = Float8BlockwiseQTensor( - shape=act_grad.shape, - dtype=fake_dtype, - rowwise_data=act_grad, - rowwise_scale_inv=permuted_scale.T.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - fp8_dtype=fp8_dtype, - quantizer=None, - is_2D_scaled=False, - requires_grad=act_grad.requires_grad, - ) - elif mxfp8_recipe: - act_grad = MXFP8Tensor( - shape=act_grad.shape, - dtype=fake_dtype, - fp8_dtype=fp8_dtype, - rowwise_data=act_grad, - rowwise_scale_inv=permuted_scale.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - quantizer=None, - requires_grad=act_grad.requires_grad, - with_gemm_swizzled_scales=False, - ) - - if not ctx.needs_input_grad[2]: - probs_grad = None - return act_grad, None, probs_grad, None, None + return act_grad, None, probs_grad, None, None, None, None + + +moe_unpermute_mask_map_forward.register_autograd( + _moe_unpermute_mask_map_backward_wrapper, + setup_context=_moe_unpermute_mask_map_setup_context, +) + +# Register all te_moe custom ops as passthrough in QuantizedTensor.__torch_dispatch__ +# so that FP8 tensors are not unwrapped before entering these ops. +_quantized_tensor_passthrough_ops.update( + { + torch.ops.te_moe.permute_mask_map_fwd.default, + torch.ops.te_moe.permute_mask_map_bwd.default, + torch.ops.te_moe.unpermute_mask_map_fwd.default, + torch.ops.te_moe.unpermute_mask_map_bwd_with_probs.default, + torch.ops.te_moe.unpermute_mask_map_bwd_no_probs.default, + } +) def moe_permute( @@ -609,10 +883,15 @@ def moe_permute( Options are: 'mask', 'index'. Refer to `routing_map` for more details. """ + if isinstance(inp, QuantizedTensor) and torch.compiler.is_compiling(): + raise RuntimeError( + "moe_permute with quantized (FP8) input is not supported under torch.compile. " + "Please move quantization outside the compiled region." + ) if map_type == "index": - return _moe_permute_index_map.apply(inp, routing_map, num_out_tokens, max_token_num) + return torch.ops.te_moe.permute_index_map(inp, routing_map, num_out_tokens, max_token_num) if map_type == "mask": - output, row_id_map, _ = _moe_permute_mask_map.apply( + output, row_id_map, _ = torch.ops.te_moe.permute_mask_map_fwd( inp, routing_map, num_out_tokens, None, None ) return output, row_id_map @@ -646,7 +925,12 @@ def moe_permute_with_probs( The effective output token count, representing the number of tokens not dropped. By default, set to '-1', meaning no tokens are dropped. """ - output, row_id_map, permuted_probs = _moe_permute_mask_map.apply( + if isinstance(inp, QuantizedTensor) and torch.compiler.is_compiling(): + raise RuntimeError( + "moe_permute_with_probs with quantized (FP8) input is not supported under " + "torch.compile. Please move quantization outside the compiled region." + ) + output, row_id_map, permuted_probs = torch.ops.te_moe.permute_mask_map_fwd( inp, routing_map, num_out_tokens, probs, None ) return output, permuted_probs, row_id_map @@ -681,6 +965,11 @@ def moe_permute_and_pad_with_probs( align_size : int the alignment size for the input tensor. """ + if isinstance(inp, QuantizedTensor) and torch.compiler.is_compiling(): + raise RuntimeError( + "moe_permute_and_pad_with_probs with quantized (FP8) input is not supported under " + "torch.compile. Please move quantization outside the compiled region." + ) if tokens_per_expert is None: raise ValueError( "tokens_per_expert must be provided to the fused permute padding function." @@ -704,7 +993,7 @@ def moe_permute_and_pad_with_probs( [torch.zeros(1, dtype=cum_pad.dtype, device=inp.device), cum_pad[:-1]] ) - output, row_id_map, permuted_probs = _moe_permute_mask_map.apply( + output, row_id_map, permuted_probs = torch.ops.te_moe.permute_mask_map_fwd( inp, routing_map, target_tokens_per_expert.sum().item(), probs, pad_offsets ) return output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert @@ -754,125 +1043,228 @@ def moe_unpermute( warnings.warn("probs kwarg is deprecated. Use merging_probs kwarg instead.") merging_probs = probs if map_type == "index": - return _moe_unpermute_index_map.apply(inp, row_id_map, merging_probs) + # Normalize probs + if merging_probs is not None: + if merging_probs.dtype != torch.float32: + warnings.warn( + f"The data type of the input `probs` of Unpermute is {merging_probs.dtype}! " + "The recommended type is torch.float32." + ) + merging_probs = merging_probs.to(torch.float32) + num_tokens = merging_probs.size(0) + topK = merging_probs.size(1) + else: + num_tokens = row_id_map.size(0) + topK = 1 + merging_probs = torch.empty(0, device=inp.device) + + return torch.ops.te_moe.unpermute_index_map_fwd( + inp, row_id_map, merging_probs, num_tokens, topK + ) if map_type == "mask": - return _moe_unpermute_mask_map.apply( - inp, row_id_map, merging_probs, restore_shape, pad_offsets + if restore_shape is None: + restore_shape = inp.shape + num_tokens, hidden_size = restore_shape + num_experts = (row_id_map.size(1) - 1) // 2 if row_id_map.dim() > 1 else 0 + + return torch.ops.te_moe.unpermute_mask_map_fwd( + inp, + row_id_map, + merging_probs, + num_tokens, + num_experts, + hidden_size, + pad_offsets, ) raise ValueError("map_type should be one of 'mask' or 'index'") -class _moe_chunk_sort(torch.autograd.Function): - """functional MoE chunk permute""" - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - split_sizes: torch.Tensor, - sorted_idxs: torch.Tensor, - probs: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # pylint: disable=missing-function-docstring - if not inp.numel(): - return inp, probs - - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not split_sizes.is_cuda: - raise ValueError( - f"split_sizes must be a CUDA tensor, but got tensor on {split_sizes.device}." - ) - if not sorted_idxs.is_cuda: - raise ValueError( - f"sorted_idxs must be a CUDA tensor, but got tensor on {sorted_idxs.device}." - ) - if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") +# ===================== _moe_chunk_sort custom ops ===================== - num_tokens, hidden_size = inp.shape - num_splits = split_sizes.size(0) - if num_splits != sorted_idxs.size(0): - raise ValueError( - f"split_sizes.size(0) ({num_splits}) must match " - f"sorted_idxs.size(0) ({sorted_idxs.size(0)})." - ) - fp8 = isinstance(inp, Float8Tensor) - if fp8: - fp8_dtype = inp._fp8_dtype - fp8_scale_inv = inp._scale_inv - fake_dtype = inp.dtype - inp = inp._data +@torch.library.custom_op("te_moe::chunk_sort_fwd", mutates_args=[]) +def moe_chunk_sort_forward( + inp: torch.Tensor, + split_sizes: torch.Tensor, + sorted_idxs: torch.Tensor, + probs: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass for MoE chunk sort. Returns (output, permuted_probs, row_id_map).""" + if not inp.numel(): + probs_out = probs.clone() if probs is not None else torch.empty(0, device=inp.device) + return inp.clone(), probs_out, torch.empty(0, device=inp.device, dtype=torch.int32) + + num_tokens, hidden_size = inp.shape + num_splits = split_sizes.size(0) + + fp8 = isinstance(inp, Float8Tensor) + if fp8: + fp8_dtype = inp._fp8_dtype + fp8_scale_inv = inp._scale_inv + fake_dtype = inp.dtype + inp = inp._data + + row_id_map = triton_permutation.make_chunk_sort_map( + split_sizes, + sorted_idxs, + num_tokens, + num_splits, + ) + output, permuted_probs = triton_permutation.sort_chunks_by_map( + inp, + row_id_map, + probs, + num_tokens, + hidden_size, + is_forward=True, + ) + if fp8: + output = Float8Tensor( + data=output, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=output.shape, + dtype=fake_dtype, + ) - row_id_map = triton_permutation.make_chunk_sort_map( - split_sizes, - sorted_idxs, - num_tokens, - num_splits, + if permuted_probs is None: + permuted_probs = torch.empty(0, device=output.device) + + return output, permuted_probs, row_id_map + + +@moe_chunk_sort_forward.register_fake +def _moe_chunk_sort_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + split_sizes: torch.Tensor, + sorted_idxs: torch.Tensor, + probs: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fake for shape inference.""" + num_tokens = inp.shape[0] + hidden_size = inp.shape[1] + fake_output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=inp.device) + if probs is not None: + fake_probs = torch.empty((num_tokens,), dtype=probs.dtype, device=inp.device) + else: + fake_probs = torch.empty(0, device=inp.device) + # row_id_map: 1D, size num_tokens + fake_row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device=inp.device) + return fake_output, fake_probs, fake_row_id_map + + +@torch.library.custom_op("te_moe::chunk_sort_bwd", mutates_args=[]) +def moe_chunk_sort_backward( + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + num_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE chunk sort.""" + fp8 = isinstance(permuted_act_grad, Float8Tensor) + if fp8: + fp8_dtype = permuted_act_grad._fp8_dtype + fp8_scale_inv = permuted_act_grad._scale_inv + fake_dtype = permuted_act_grad.dtype + permuted_act_grad = permuted_act_grad._data + + act_grad, probs_grad = triton_permutation.sort_chunks_by_map( + permuted_act_grad, + row_id_map, + permuted_probs_grad, + num_tokens, + hidden_size, + is_forward=False, + ) + + if fp8: + act_grad = Float8Tensor( + data=act_grad, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=act_grad.shape, + dtype=fake_dtype, ) - output, permuted_probs = triton_permutation.sort_chunks_by_map( - inp, - row_id_map, - probs, - num_tokens, - hidden_size, - is_forward=True, + + if probs_grad is None: + probs_grad = torch.empty(0, device=act_grad.device) + + return act_grad, probs_grad + + +@moe_chunk_sort_backward.register_fake +def _moe_chunk_sort_backward_fake( # pylint: disable=unused-argument + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + num_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake for backward shape inference.""" + fake_act_grad = torch.empty( + (num_tokens, hidden_size), + dtype=permuted_act_grad.dtype, + device=permuted_act_grad.device, + ) + if permuted_probs_grad is not None: + fake_probs_grad = torch.empty( + (num_tokens,), + dtype=permuted_probs_grad.dtype, + device=permuted_act_grad.device, ) - if fp8: - output = Float8Tensor( - data=output, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=output.shape, - dtype=fake_dtype, - ) + else: + fake_probs_grad = torch.empty(0, device=permuted_act_grad.device) + return fake_act_grad, fake_probs_grad + + +def _moe_chunk_sort_setup_context(ctx, inputs, output): + """Save context for backward pass.""" + inp, _split_sizes, _sorted_idxs, probs = inputs + _output_tensor, _permuted_probs, row_id_map = output + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(row_id_map) + ctx.num_tokens = inp.size(0) + ctx.hidden_size = inp.size(1) if not ctx.empty_input else 0 + ctx.needs_probs_grad = probs is not None and probs.requires_grad - ctx.save_for_backward(row_id_map) - ctx.num_tokens = num_tokens - ctx.hidden_size = hidden_size - return output, permuted_probs - - @staticmethod - def backward( - ctx, - permuted_act_grad: torch.Tensor, - permuted_probs_grad: torch.Tensor, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, permuted_probs_grad - - act_grad = None + +def _moe_chunk_sort_backward_wrapper(ctx, permuted_act_grad, permuted_probs_grad, _row_id_map_grad): + """Backward wrapper calling the custom backward op.""" + if ctx.empty_input: + probs_grad = permuted_probs_grad if ctx.needs_probs_grad else None + return permuted_act_grad, None, None, probs_grad + + (row_id_map,) = ctx.saved_tensors + + probs_grad_input = permuted_probs_grad if permuted_probs_grad.numel() > 0 else None + + act_grad, probs_grad = torch.ops.te_moe.chunk_sort_bwd( + permuted_act_grad, + probs_grad_input, + row_id_map, + ctx.num_tokens, + ctx.hidden_size, + ) + + if not ctx.needs_probs_grad or probs_grad.numel() == 0: probs_grad = None - if ctx.needs_input_grad[0]: - (row_id_map,) = ctx.saved_tensors - fp8 = isinstance(permuted_act_grad, Float8Tensor) - if fp8: - fp8_dtype = permuted_act_grad._fp8_dtype - fp8_scale_inv = permuted_act_grad._scale_inv - fake_dtype = permuted_act_grad.dtype - permuted_act_grad = permuted_act_grad._data - act_grad, probs_grad = triton_permutation.sort_chunks_by_map( - permuted_act_grad, - row_id_map, - permuted_probs_grad, - ctx.num_tokens, - ctx.hidden_size, - is_forward=False, - ) - if fp8: - act_grad = Float8Tensor( - data=act_grad, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=act_grad.shape, - dtype=fake_dtype, - ) - if not ctx.needs_input_grad[3]: - probs_grad = None - return act_grad, None, None, probs_grad + + return act_grad, None, None, probs_grad + + +moe_chunk_sort_forward.register_autograd( + _moe_chunk_sort_backward_wrapper, + setup_context=_moe_chunk_sort_setup_context, +) + +# Register chunk sort ops as passthrough in QuantizedTensor.__torch_dispatch__ +_quantized_tensor_passthrough_ops.update( + { + torch.ops.te_moe.chunk_sort_fwd.default, + torch.ops.te_moe.chunk_sort_bwd.default, + } +) def moe_sort_chunks_by_index( @@ -894,7 +1286,7 @@ def moe_sort_chunks_by_index( sorted_indices : torch.Tensor Chunk indices used to permute the chunks. """ - output, _ = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, None) + output, _, _ = torch.ops.te_moe.chunk_sort_fwd(inp, split_sizes, sorted_index, None) return output @@ -922,5 +1314,7 @@ def moe_sort_chunks_by_index_with_probs( sorted_indices : torch.Tensor Chunk indices used to permute the chunks. """ - output, permuted_probs = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, probs) + output, permuted_probs, _ = torch.ops.te_moe.chunk_sort_fwd( + inp, split_sizes, sorted_index, probs + ) return output, permuted_probs diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 807671e863..e40f42edd3 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -21,6 +21,12 @@ ) +# Custom ops that should pass through __torch_dispatch__ without unwrapping +# QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that +# handle quantized tensors internally. +_quantized_tensor_passthrough_ops: set = set() + + class QuantizedTensorStorage: r"""Base class for all TensorStorage classes. @@ -614,6 +620,12 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): return func(t) return False # Or error out? + # Pass through registered custom ops without unwrapping + if func in _quantized_tensor_passthrough_ops: + if kwargs is None: + kwargs = {} + return super().__torch_dispatch__(func, types, args, kwargs) + def maybe_unwrap(arg): if isinstance(arg, QuantizedTensor): return arg.dequantize() From 6da802e02833548a0e56d6ca298647227bb291ea Mon Sep 17 00:00:00 2001 From: Jacket <44538064+kainzhong@users.noreply.github.com> Date: Thu, 19 Mar 2026 00:11:39 -0700 Subject: [PATCH 04/18] [PyTorch] Add an API restore from function context to ensure tensors are detached (#2772) [PyTorch] Change the restore tensor API to ensure tensors are detached from ctx Signed-off-by: Kaining Zhong Co-authored-by: Kirthi Shankar Sivamani Signed-off-by: Varun Thumbe --- tests/pytorch/attention/test_attention.py | 7 ++---- transformer_engine/pytorch/__init__.py | 1 + .../dot_product_attention/backends.py | 4 ++-- .../dot_product_attention/context_parallel.py | 6 ++--- .../pytorch/module/grouped_linear.py | 4 ++-- .../pytorch/module/layernorm_linear.py | 9 ++------ .../pytorch/module/layernorm_mlp.py | 8 ++----- transformer_engine/pytorch/module/linear.py | 9 ++------ transformer_engine/pytorch/ops/fuser.py | 5 ++--- .../pytorch/quantized_tensor.py | 22 ++++++++++++++++++- transformer_engine/pytorch/tensor/__init__.py | 2 ++ 11 files changed, 41 insertions(+), 36 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 60ade522e3..2eb307aa48 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -49,7 +49,7 @@ from transformer_engine.pytorch.quantized_tensor import ( Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) _current_file = pathlib.Path(__file__).resolve() @@ -2701,10 +2701,7 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: with torch.cuda.nvtx.range("_DPA"): - saved_tensors = ctx.saved_tensors - (q, k, v, inp_fp8, qkv_weight_fp8, out) = restore_from_saved( - ctx.tensor_objects, saved_tensors - ) + (q, k, v, inp_fp8, qkv_weight_fp8, out) = restore_from_func_ctx(ctx) proj_dgrad = ctx.dO_quantizer(grad_output) fp8_dtype_backward = get_fp8_te_dtype(ctx.fp8_meta["recipe"], fprop_tensor=False) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index cd18ca75ad..bbc1d7fab6 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -68,6 +68,7 @@ from transformer_engine.pytorch.quantized_tensor import Quantizer from transformer_engine.pytorch.quantized_tensor import prepare_for_saving from transformer_engine.pytorch.quantized_tensor import restore_from_saved +from transformer_engine.pytorch.quantized_tensor import restore_from_func_ctx from transformer_engine.pytorch.tensor import Float8Quantizer from transformer_engine.pytorch.tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor import MXFP8Quantizer diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index a6a8b0b26a..442366035a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -32,7 +32,7 @@ from transformer_engine.pytorch.quantized_tensor import ( QuantizedTensorStorage, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.constants import ( @@ -1477,7 +1477,7 @@ def backward(ctx, d_out, *_args): cu_seqlens_q_padded, cu_seqlens_kv_padded, *other_tensors, - ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + ) = restore_from_func_ctx(ctx) aux_ctx_tensors = other_tensors diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 10ba99595b..7d9eb0cb05 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -38,7 +38,7 @@ from transformer_engine.pytorch.quantized_tensor import ( prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) # Import attention utils @@ -2085,7 +2085,7 @@ def backward(ctx, dout, *_args): cu_seqlens_q_padded, cu_seqlens_kv_padded, *other_tensors, - ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + ) = restore_from_func_ctx(ctx) cu_seqlens_q_per_step = other_tensors[:cp_size] cu_seqlens_kv_per_step = other_tensors[cp_size : cp_size * 2] rng_states = other_tensors[cp_size * 2 : cp_size * 3] @@ -3675,7 +3675,7 @@ def backward(ctx, dout, *_args): cu_seqlens_q_padded, cu_seqlens_kv_padded, *aux_ctx_tensors, - ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + ) = restore_from_func_ctx(ctx) qkv_format = ctx.qkv_format qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 30c1dbf408..0adda48e36 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -49,7 +49,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ...debug.pytorch.debug_quantization import DebugQuantizer from ...debug.pytorch.debug_state import TEDebugState @@ -316,7 +316,7 @@ def forward( def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): - saved_tensors = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + saved_tensors = restore_from_func_ctx(ctx) N = ctx.num_gemms inputmats = saved_tensors[:N] weights = saved_tensors[N : 2 * N] diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index d775dc3e8e..ed91bc1235 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -60,7 +60,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -546,7 +546,6 @@ def backward( nvtx_label = f"{nvtx_label}.{ctx.ub_name}" with get_nvtx_range_context("_LayerNormLinear_backward"): - saved_tensors = ctx.saved_tensors ( # pylint: disable=unbalanced-tuple-unpacking inputmat, weight, @@ -556,11 +555,7 @@ def backward( ln_out, mu, rsigma, - ) = restore_from_saved(ctx.tensor_objects, saved_tensors) - - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None + ) = restore_from_func_ctx(ctx) # Since main_grad can be modified inplace, it should not be a part of saved_tensors main_grad = ( diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 037fb6c858..cc3dcc4064 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -80,7 +80,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ..cpp_extensions import ( general_gemm, @@ -898,11 +898,7 @@ def forward( def _recompute(ctx): # pylint: disable=missing-function-docstring - saved_tensors = ctx.saved_tensors - tensors = restore_from_saved(ctx.tensor_objects, saved_tensors) - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None + tensors = restore_from_func_ctx(ctx) if ctx.checkpoint: # do recomputation from the original args diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 1e3eadc405..ea921341a4 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -61,7 +61,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -501,15 +501,10 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], nvtx_label = f"{nvtx_label}.{ctx.ub_name}" with get_nvtx_range_context("_Linear_backward"): - saved_tensors = ctx.saved_tensors inputmat, weight_fp8, weight, bias = ( # pylint: disable=unbalanced-tuple-unpacking - restore_from_saved(ctx.tensor_objects, saved_tensors) + restore_from_func_ctx(ctx) ) - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None - # Since main_grad can be modified inplace, it should not be a part of saved_tensors main_grad = ( ctx.main_grad_func() diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 80386db2d9..76606ec799 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -12,7 +12,7 @@ import torch from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling -from ..quantized_tensor import prepare_for_saving, restore_from_saved +from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx from .op import ( BasicOperation, FusibleOperation, @@ -212,8 +212,7 @@ def backward( basic_op_ctxs = func_ctx.basic_op_ctxs # Restore saved tensors - saved_tensors = restore_from_saved(func_ctx.tensor_objects, func_ctx.saved_tensors) - func_ctx.tensor_objects = None + saved_tensors = restore_from_func_ctx(func_ctx) # Unflatten list of saved tensors for ctx in basic_op_ctxs: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index e40f42edd3..a7722f777e 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -165,7 +165,9 @@ def restore_from_saved( list[Optional[torch.Tensor]], ] ): - """Recombine the tensor data and metadata during backward pass.""" + """Recombine the tensor data and metadata during backward pass. + Note: please use `restore_from_func_ctx` instead if you are restoring tensors from a function context to make sure tensor_objects is detached and its memory can be freed + """ tensor_objects = [] for tensor in tensors: if tensor is None or isinstance(tensor, torch.Tensor): @@ -180,6 +182,24 @@ def restore_from_saved( return tensor_objects +def restore_from_func_ctx(ctx: torch.autograd.function.FunctionCtx, return_saved_tensors=False) -> ( + list[Optional[torch.Tensor | QuantizedTensorStorage]] + | tuple[ + list[Optional[torch.Tensor | QuantizedTensorStorage]], + list[Optional[torch.Tensor]], + ] +): + """Recombine the tensor data and metadata during backward pass and delete tensor objects attached to function context.""" + if not hasattr(ctx, "tensor_objects") or ctx.tensor_objects is None: + raise AttributeError("ctx must have .tensor_objects to restore saved tensors") + out = restore_from_saved( + ctx.tensor_objects, ctx.saved_tensors, return_saved_tensors=return_saved_tensors + ) + # Delete the references to tensor objects once they've been consumed by the `restore_from_saved` method to construct back the actual tensors. + ctx.tensor_objects = None + return out + + class Quantizer(abc.ABC): """Builder class for quantized tensors. diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index 5668056700..426c656d47 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -12,6 +12,7 @@ Quantizer, prepare_for_saving, restore_from_saved, + restore_from_func_ctx, ) from .storage.float8_tensor_storage import Float8TensorStorage from .storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -46,6 +47,7 @@ "GroupedTensor", "prepare_for_saving", "restore_from_saved", + "restore_from_func_ctx", ] From 56366bb9e01f8d955582fb2cc3889fe810f8fa46 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 19 Mar 2026 10:17:23 -0700 Subject: [PATCH 05/18] [PyT] Install pytest in onnx L1 test as Pyt container no longer packages it (#2781) Install pytest in onnx L1 test as Pyt container no longer packages it Signed-off-by: Kshitij Janardan Lakhani Signed-off-by: Varun Thumbe --- qa/L1_pytorch_onnx_unittest/test.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/qa/L1_pytorch_onnx_unittest/test.sh b/qa/L1_pytorch_onnx_unittest/test.sh index 6f9ff54e48..0edf92c475 100644 --- a/qa/L1_pytorch_onnx_unittest/test.sh +++ b/qa/L1_pytorch_onnx_unittest/test.sh @@ -2,9 +2,15 @@ # # See LICENSE for license information. +function error_exit() { + echo "Error: $1" + exit 1 +} + : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # NVTE_UnfusedDPA_Emulate_FP8=1 enables FP8 attention emulation when no native backend is available NVTE_UnfusedDPA_Emulate_FP8=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_onnx_export.xml $TE_PATH/tests/pytorch/test_onnx_export.py From f943147517ba73a4b3162bd0cc36a80394d6b176 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 19 Mar 2026 16:16:26 -0700 Subject: [PATCH 06/18] [Core] Fix MXFP8 grouped quantize for zero-sized groups in update_tma_descriptors (#2782) * Fix zero-sized groups in update_tma_descriptors Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update test_cast_mxfp8_grouped.cu Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Varun Thumbe --- tests/cpp/operator/test_cast_mxfp8_grouped.cu | 1 + .../common/cast/mxfp8/group_quantize_mxfp8.cuh | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index e469ad0845..09bd21657a 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -649,6 +649,7 @@ std::vector> input_config = { {SAME_BOTH_DIMS, 2, 256,128}, {VARYING_FIRST_DIM, 2, 512,128, 128,384}, {VARYING_FIRST_DIM, 3, 1024,144, 128,384,512}, + {VARYING_FIRST_DIM, 4, 1024,144, 128,384,0,512}, {VARYING_FIRST_DIM, 4, 1536,160, 128,384,512,512}, {VARYING_FIRST_DIM, 5, 4096,512, 128,256,384,1024,2304}, {VARYING_LAST_DIM, 3, 256,896, 128,256,512}, diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 129d6724ac..d0d15d8d6c 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -189,6 +189,13 @@ __global__ void update_tma_descriptors( get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + // Zero-sized groups: skip TMA descriptor update. The main kernel already returns + // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension + // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. + if (rows == 0 || cols == 0) { + return; + } + const size_t offset_elts = offsets_ptr[tensor_id]; if (leading_thread && (tensor_id < num_tensors)) { From 4f0f7f92f8d1907177ced0687acef70f438cd139 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 22 Mar 2026 18:44:16 +0000 Subject: [PATCH 07/18] Revert "fix merge conflicts, now things working" This reverts commit e355c38c111175a7d46b7dd27c38523135e9429b. Signed-off-by: Varun Thumbe --- 3rdparty/cudnn-frontend | 2 +- tests/cpp/operator/CMakeLists.txt | 1 - .../operator/test_multi_tensor_adam_mxfp8.cu | 266 ------------ tests/cpp/test_common.h | 10 - .../distributed/run_fsdp2_fused_adam.py | 8 +- tests/pytorch/distributed/test_torch_fsdp2.py | 5 + .../include/transformer_engine/multi_tensor.h | 35 -- .../common/multi_tensor/adam.cu | 384 +++--------------- .../multi_tensor/multi_tensor_apply.cuh | 94 ----- transformer_engine/pytorch/csrc/extensions.h | 6 - .../csrc/extensions/multi_tensor/adam.cpp | 20 - .../pytorch/csrc/extensions/pybind.cpp | 2 - .../pytorch/optimizers/fused_adam.py | 45 +- 13 files changed, 66 insertions(+), 812 deletions(-) delete mode 100644 tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 8d19d3182b..d33027a41a 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 8d19d3182bfbc304046a15e9236bec9ff31511fc +Subproject commit d33027a41a93af9c85f089c6364ab415fce98982 diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 4241ada3ba..5e73675f4f 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -27,7 +27,6 @@ add_executable(test_operator test_memset.cu test_splits_to_offsets.cu test_multi_cast_transpose.cu - test_multi_tensor_adam_mxfp8.cu test_multi_padding.cu test_multi_unpadding.cu test_causal_softmax.cu diff --git a/tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu b/tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu deleted file mode 100644 index 470917580f..0000000000 --- a/tests/cpp/operator/test_multi_tensor_adam_mxfp8.cu +++ /dev/null @@ -1,266 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include -#include -#include - -#include -#include -#include -#include - -#include -#include - -#include "../test_common.h" - -using namespace transformer_engine; -using namespace test; - -namespace { - -uint8_t fp8_to_u8(fp8e4m3 v) { - uint8_t out = 0; - std::memcpy(&out, &v, sizeof(uint8_t)); - return out; -} - -uint8_t fp8_to_u8(fp8e5m2 v) { - uint8_t out = 0; - std::memcpy(&out, &v, sizeof(uint8_t)); - return out; -} - -void run_mxfp8_adam_test(DType fp8_dtype) { - const std::vector shape1{64, 128}; - const std::vector shape2{32, 64}; - const float lr = 1e-3f; - const float beta1 = 0.9f; - const float beta2 = 0.999f; - const float eps = 1e-8f; - const int step = 1; - const int mode = 1; - const int bias_correction = 1; - const float weight_decay = 0.0f; - - // Run with 25 tensors > 24[MXFP8_MAX_TENSORS] to check - // the chunking logic - const size_t tensor_count = 25; - std::vector> shapes; - shapes.reserve(tensor_count); - for (size_t i = 0; i < tensor_count; ++i) { - shapes.push_back((i % 2 == 0) ? shape1 : shape2); - } - - std::vector names; - names.reserve(tensor_count * 11); - std::vector g; - std::vector p; - std::vector m; - std::vector v; - std::vector p_ref_t; - std::vector m_ref_t; - std::vector v_ref_t; - std::vector q_ref; - std::vector dq; - std::vector dq_ref; - std::vector q; - g.reserve(tensor_count); - p.reserve(tensor_count); - m.reserve(tensor_count); - v.reserve(tensor_count); - p_ref_t.reserve(tensor_count); - m_ref_t.reserve(tensor_count); - v_ref_t.reserve(tensor_count); - q_ref.reserve(tensor_count); - dq.reserve(tensor_count); - dq_ref.reserve(tensor_count); - q.reserve(tensor_count); - - for (size_t i = 0; i < tensor_count; ++i) { - const std::vector &shape = shapes[i]; - names.push_back("g" + std::to_string(i)); - g.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - names.push_back("p" + std::to_string(i)); - p.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - names.push_back("m" + std::to_string(i)); - m.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - names.push_back("v" + std::to_string(i)); - v.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - - fillUniform(&g.back()); - fillUniform(&p.back()); - std::fill_n(m.back().rowwise_cpu_dptr(), product(m.back().rowwise_shape()), 0.0f); - std::fill_n(v.back().rowwise_cpu_dptr(), product(v.back().rowwise_shape()), 0.0f); - m.back().from_cpu(); - v.back().from_cpu(); - - names.push_back("p_ref_" + std::to_string(i)); - p_ref_t.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - names.push_back("m_ref_" + std::to_string(i)); - m_ref_t.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - names.push_back("v_ref_" + std::to_string(i)); - v_ref_t.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - const size_t n = shape[0] * shape[1]; - std::memcpy(p_ref_t.back().rowwise_cpu_dptr(), p.back().rowwise_cpu_dptr(), - n * sizeof(float)); - std::memcpy(m_ref_t.back().rowwise_cpu_dptr(), m.back().rowwise_cpu_dptr(), - n * sizeof(float)); - std::memcpy(v_ref_t.back().rowwise_cpu_dptr(), v.back().rowwise_cpu_dptr(), - n * sizeof(float)); - p_ref_t.back().from_cpu(); - m_ref_t.back().from_cpu(); - v_ref_t.back().from_cpu(); - - names.push_back("q_ref_" + std::to_string(i)); - q_ref.emplace_back(names.back().c_str(), shape, fp8_dtype, true, true, NVTE_MXFP8_1D_SCALING); - q_ref.back().set_with_gemm_swizzled_scales(false); - - names.push_back("dq" + std::to_string(i)); - dq.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - names.push_back("dq_ref_" + std::to_string(i)); - dq_ref.emplace_back(names.back().c_str(), shape, DType::kFloat32, true, false); - - names.push_back("q" + std::to_string(i)); - q.emplace_back(names.back().c_str(), shape, fp8_dtype, true, true, NVTE_MXFP8_1D_SCALING); - q.back().set_with_gemm_swizzled_scales(false); - } - - Tensor noop("noop", std::vector{1}, DType::kInt32, true, false); - int zero = 0; - std::memcpy(noop.rowwise_cpu_dptr(), &zero, sizeof(int)); - noop.from_cpu(); - - std::vector> lists(8); - std::vector extra_wrappers; - extra_wrappers.reserve(tensor_count * 4); - - auto add_tensor = [&](Tensor &g, Tensor &p, Tensor &m, Tensor &v, Tensor &q) { - lists[0].push_back(g.data()); - lists[1].push_back(p.data()); - lists[2].push_back(m.data()); - lists[3].push_back(v.data()); - - extra_wrappers.emplace_back(q.rowwise_dptr(), q.rowwise_shape(), fp8_dtype); - lists[4].push_back(extra_wrappers.back().data()); - extra_wrappers.emplace_back(q.columnwise_dptr(), q.columnwise_shape(), fp8_dtype); - lists[5].push_back(extra_wrappers.back().data()); - extra_wrappers.emplace_back(q.rowwise_scale_inv_dptr(), q.rowwise_scale_inv_shape(), - DType::kByte); - lists[6].push_back(extra_wrappers.back().data()); - extra_wrappers.emplace_back(q.columnwise_scale_inv_dptr(), q.columnwise_scale_inv_shape(), - DType::kByte); - lists[7].push_back(extra_wrappers.back().data()); - }; - - for (size_t i = 0; i < tensor_count; ++i) { - add_tensor(g[i], p[i], m[i], v[i], q[i]); - } - - std::vector list_ptrs; - list_ptrs.reserve(lists.size()); - for (auto &l : lists) { - list_ptrs.push_back(l.data()); - } - - nvte_multi_tensor_adam_mxfp8_cuda(65536, noop.data(), list_ptrs.data(), lists.size(), - lists[0].size(), static_cast(fp8_dtype), lr, beta1, - beta2, eps, step, mode, bias_correction, weight_decay, 0); - - std::vector> ref_lists(4); - for (size_t i = 0; i < tensor_count; ++i) { - ref_lists[0].push_back(g[i].data()); - ref_lists[1].push_back(p_ref_t[i].data()); - ref_lists[2].push_back(m_ref_t[i].data()); - ref_lists[3].push_back(v_ref_t[i].data()); - } - std::vector ref_list_ptrs; - ref_list_ptrs.reserve(ref_lists.size()); - for (auto &l : ref_lists) { - ref_list_ptrs.push_back(l.data()); - } - - nvte_multi_tensor_adam_cuda(65536, noop.data(), ref_list_ptrs.data(), ref_lists.size(), - ref_lists[0].size(), lr, beta1, beta2, eps, step, mode, - bias_correction, weight_decay, 0); - - for (size_t i = 0; i < tensor_count; ++i) { - nvte_quantize(p_ref_t[i].data(), q_ref[i].data(), 0); - nvte_dequantize(q[i].data(), dq[i].data(), 0); - nvte_dequantize(q_ref[i].data(), dq_ref[i].data(), 0); - } - - cudaDeviceSynchronize(); - - for (size_t i = 0; i < tensor_count; ++i) { - q[i].to_cpu(); - p[i].to_cpu(); - m[i].to_cpu(); - v[i].to_cpu(); - q_ref[i].to_cpu(); - dq[i].to_cpu(); - dq_ref[i].to_cpu(); - p_ref_t[i].to_cpu(); - m_ref_t[i].to_cpu(); - v_ref_t[i].to_cpu(); - } - - for (size_t i = 0; i < lists[0].size(); ++i) { - const Tensor &g_i = g[i]; - const Tensor &p_i = p[i]; - const Tensor &m_i = m[i]; - const Tensor &v_i = v[i]; - Tensor &q_i = q[i]; - const Tensor &p_ref_t_i = p_ref_t[i]; - const Tensor &m_ref_t_i = m_ref_t[i]; - const Tensor &v_ref_t_i = v_ref_t[i]; - Tensor &q_ref_i = q_ref[i]; - - compareResults("p", p_i, p_ref_t_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, 0); - compareResults("m", m_i, m_ref_t_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, 0); - compareResults("v", v_i, v_ref_t_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, 0); - - const Tensor &dq_i = dq[i]; - const Tensor &dq_ref_i = dq_ref[i]; - compareResults("dequantized", dq_i, dq_ref_i.rowwise_cpu_dptr(), true, 0.0, 0.0, true, - 0); - - const size_t rs = q_i.rowwise_scale_inv_shape().data[1]; - const size_t cs = q_i.columnwise_scale_inv_shape().data[1]; - const size_t rowwise_scale_size = q_i.rowwise_scale_inv_shape().data[0] * rs; - const size_t colwise_scale_size = q_i.columnwise_scale_inv_shape().data[0] * cs; - compareResults("rowwise_scale", q_i.rowwise_cpu_scale_inv_ptr(), - q_ref_i.rowwise_cpu_scale_inv_ptr(), rowwise_scale_size, 0.0f); - compareResults("colwise_scale", q_i.columnwise_cpu_scale_inv_ptr(), - q_ref_i.columnwise_cpu_scale_inv_ptr(), colwise_scale_size, 0.0f); - - uint8_t *row_data = nullptr; - uint8_t *col_data = nullptr; - uint8_t *row_data_ref = nullptr; - uint8_t *col_data_ref = nullptr; - if (fp8_dtype == DType::kFloat8E4M3) { - row_data = reinterpret_cast(q_i.rowwise_cpu_dptr()); - col_data = reinterpret_cast(q_i.columnwise_cpu_dptr()); - row_data_ref = reinterpret_cast(q_ref_i.rowwise_cpu_dptr()); - col_data_ref = reinterpret_cast(q_ref_i.columnwise_cpu_dptr()); - } else { - row_data = reinterpret_cast(q_i.rowwise_cpu_dptr()); - col_data = reinterpret_cast(q_i.columnwise_cpu_dptr()); - row_data_ref = reinterpret_cast(q_ref_i.rowwise_cpu_dptr()); - col_data_ref = reinterpret_cast(q_ref_i.columnwise_cpu_dptr()); - } - const size_t data_size = q_i.rowwise_shape().data[0] * q_i.rowwise_shape().data[1]; - compareResults("rowwise_data", row_data, row_data_ref, data_size, 0.0f); - compareResults("colwise_data", col_data, col_data_ref, data_size, 0.0f); - } -} - -} // namespace - -TEST(MultiTensorAdamMXFP8, E4M3) { run_mxfp8_adam_test(DType::kFloat8E4M3); } - -TEST(MultiTensorAdamMXFP8, E5M2) { run_mxfp8_adam_test(DType::kFloat8E5M2); } diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index eab181fa82..927407f478 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -200,16 +200,6 @@ class Tensor { return tensor_.get_columnwise_data().data_ptr; } - void *rowwise_scale_inv_dptr() const { - NVTE_CHECK(rowwise_, "Tensor does not have rowwise data!"); - return tensor_.get_rowwise_scale_inv().data_ptr; - } - - void *columnwise_scale_inv_dptr() const { - NVTE_CHECK(columnwise_, "Tensor does not have columnwise data!"); - return tensor_.get_columnwise_scale_inv().data_ptr; - } - template T *rowwise_cpu_dptr() const { NVTE_CHECK(TypeInfo::dtype == tensor_.dtype(), "Invalid type!"); diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/run_fsdp2_fused_adam.py index 34764d4e0a..c39957cf13 100644 --- a/tests/pytorch/distributed/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/run_fsdp2_fused_adam.py @@ -36,11 +36,7 @@ def get_recipe_from_string(recipe): SEQ_LEN = 32 BATCH_PER_RANK = 2 NUM_STEPS = 3 -LOCAL_RANK = None -def dist_print(msg): - if LOCAL_RANK == 0: - print(msg) def save_custom_attrs(module): custom_attrs = {} @@ -155,8 +151,6 @@ def test_fused_adam_fp8_master_weights(recipe=None): - Training loop completes without error - DTensor wrapping and QuantizedTensor local tensors are preserved """ - global LOCAL_RANK - LOCAL_RANK = int(os.environ["LOCAL_RANK"]) world_size, _, device = _setup() model = _build_model(fp8_init=True, recipe=recipe) @@ -189,7 +183,7 @@ def test_fused_adam_fp8_master_weights(recipe=None): loss = F.mse_loss(output, target) loss.backward() optimizer.step() - dist_print(f"Step {step} completed with loss {loss.item()}") + # Verify optimizer states for param in model.parameters(): state = optimizer.state[param] diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 6d7ae4d7bb..02e45d99cb 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -224,6 +224,11 @@ def test_fsdp2_dcp_output_parity_async(fp_recipe): @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") def test_fsdp2_safetensors_fp32_export(fp_recipe): """Export FP32 model from optimizer master weights to safetensors.""" + if fp_recipe == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access" + ) _run_fused_adam_test("safetensors_fp32_export", fp_recipe) diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index 90c87b166e..09ab260f15 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -149,41 +149,6 @@ void nvte_multi_tensor_adam_fp8_cuda(int chunk_size, NVTETensor noop_flag, const float weight_decay, const NVTEDType fp8_dtype, cudaStream_t stream); -/*! \brief Compute and apply gradient update to parameters for Adam optimizer - * when model parameters are in MXFP8 precision. - * - * The update is applied to FP32 master parameters, then the master - * parameters are quantized to MXFP8 rowwise and columnwise data - * (both are always required). - * - * \warning This API is **experimental** and subject to change. - * - * \param[in] chunk_size Number of tensor elements processed by a CUDA block. - * \param[in] noop_flag If this single element tensor has non-zero value, kernel will exit immediately. - * \param[in,out] tensor_lists 2D array of input tensors with 8 lists in order: - * (0) gradients, (1) FP32 master params, (2) first moment, - * (3) second moment, (4) rowwise MXFP8 data, - * (5) columnwise MXFP8 data, (6) rowwise scale-inv, - * (7) columnwise scale-inv. - * \param[in] num_tensor_lists Size (dim0) of tensor_lists. Must be 8. - * \param[in] num_tensors_per_list Size (dim1) of tensor_lists. - * \param[in] fp8_dtype MXFP8 element type for quantization (E4M3/E5M2). - * \param[in] lr Learning rate. - * \param[in] beta1 Coefficient for first moment of gradient. - * \param[in] beta2 Coefficient for second moment of gradient. - * \param[in] epsilon Term added to the denominator for numerical stability. - * \param[in] step Iteration counter. - * \param[in] mode Whether to use AdamW (L2 penalty applied to params). - * \param[in] bias_correction Whether to apply correction factor for moment estimates. - * \param[in] weight_decay L2 penalty for weight decay. - * \param[in] stream CUDA stream used for this operation. - */ -void nvte_multi_tensor_adam_mxfp8_cuda( - int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, - const size_t num_tensor_lists, const size_t num_tensors_per_list, const NVTEDType fp8_dtype, - const float lr, const float beta1, const float beta2, const float epsilon, const int step, - const int mode, const int bias_correction, const float weight_decay, cudaStream_t stream); - /*! \brief Compute and apply gradient update to parameters for Adam optimizer * with CUDA graph support and LR scheduling. * diff --git a/transformer_engine/common/multi_tensor/adam.cu b/transformer_engine/common/multi_tensor/adam.cu index fa75c645f3..29a073be84 100644 --- a/transformer_engine/common/multi_tensor/adam.cu +++ b/transformer_engine/common/multi_tensor/adam.cu @@ -4,16 +4,12 @@ * See LICENSE for license information. ************************************************************************/ -#include #include #include #include #include -#include "../common.h" -#include "../util/math.h" #include "../utils.cuh" -#include "../util/ptx.cuh" #include "multi_tensor_apply.cuh" namespace transformer_engine { @@ -31,7 +27,6 @@ typedef enum { using MATH_T = float; using fp8e4m3 = __nv_fp8_e4m3; using fp8e5m2 = __nv_fp8_e5m2; -using e8m0_t = transformer_engine::e8m0_t; template struct is_fp8 : std::false_type {}; @@ -54,31 +49,6 @@ struct FP8Data { template <> struct FP8Data {}; -template -__device__ __forceinline__ void adam_update(T &r_g, T &r_p, T &r_m, T &r_v, const float beta1, - const float beta2, const float beta1_correction, - const float beta2_correction, const float epsilon, - const float lr, adamMode_t mode, const float decay) { - if (mode == ADAM_MODE_0) { // L2 - r_g = r_g + (decay * r_p); - r_m = beta1 * r_m + (1 - beta1) * r_g; - r_v = beta2 * r_v + (1 - beta2) * r_g * r_g; - T next_m_unbiased = r_m / beta1_correction; - T next_v_unbiased = r_v / beta2_correction; - T denom = sqrtf(next_v_unbiased) + epsilon; - T update = next_m_unbiased / denom; - r_p = r_p - (lr * update); - } else { // weight decay - r_m = beta1 * r_m + (1 - beta1) * r_g; - r_v = beta2 * r_v + (1 - beta2) * r_g * r_g; - T next_m_unbiased = r_m / beta1_correction; - T next_v_unbiased = r_v / beta2_correction; - T denom = sqrtf(next_v_unbiased) + epsilon; - T update = (next_m_unbiased / denom) + (decay * r_p); - r_p = r_p - (lr * update); - } -} - template struct AdamFunctorMaster { static constexpr bool is_fp8_type = is_fp8::value; @@ -152,8 +122,24 @@ struct AdamFunctorMaster { } #pragma unroll for (int ii = 0; ii < ILP; ii++) { - adam_update(r_g[ii], r_p[ii], r_m[ii], r_v[ii], beta1, beta2, beta1_correction, - beta2_correction, epsilon, lr, mode, decay); + if (mode == ADAM_MODE_0) { // L2 + r_g[ii] = r_g[ii] + (decay * r_p[ii]); + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sqrtf(next_v_unbiased) + epsilon; + MATH_T update = next_m_unbiased / denom; + r_p[ii] = r_p[ii] - (lr * update); + } else { // weight decay + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sqrtf(next_v_unbiased) + epsilon; + MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]); + r_p[ii] = r_p[ii] - (lr * update); + } } #pragma unroll @@ -586,188 +572,6 @@ struct AdamCapturableMasterFunctor { } }; -template -__device__ __forceinline__ FP8_T cast_to_fp8(float x) { - return static_cast(x); -} - -__device__ __forceinline__ float fp8_max_norm_rcp(uint8_t fp8_dtype) { - if (fp8_dtype == static_cast(transformer_engine::DType::kFloat8E4M3)) { - return transformer_engine::Quantized_Limits::max_norm_rcp; - } - return transformer_engine::Quantized_Limits::max_norm_rcp; -} - -template -__global__ void adam_mxfp8_fused_kernel( - int64_t chunk_size, volatile int *noop_gmem, MXFP8TensorListMetadata tl, float beta1, - float beta2, float beta1_correction, float beta2_correction, float epsilon, float lr, int mode, - float weight_decay) { - // Stage 0: optional early-exit if a noop flag is set. - if (noop_gmem != nullptr && *noop_gmem == 1) { - return; - } - (void)chunk_size; - - // Stage 1: map this block to a specific tensor tile. - const int block_idx = blockIdx.x; - const int tensor_idx = tl.block_to_tensor[block_idx]; - const int tile_idx = tl.block_to_tile[block_idx]; - const int64_t rows_val = tl.rows[tensor_idx]; - const int64_t cols_val = tl.cols[tensor_idx]; - if (rows_val == 0 || cols_val == 0) { - return; - } - - const int64_t tiles_per_row = (cols_val + MXFP8_TILE - 1) / MXFP8_TILE; - const int64_t tile_row = tile_idx / tiles_per_row; - const int64_t tile_col = tile_idx % tiles_per_row; - const int64_t row_base = tile_row * MXFP8_TILE; - const int64_t col_base = tile_col * MXFP8_TILE; - - // Stage 2: load pointers for grads/params/moments and MXFP8 outputs/scales. - GRAD_T *g = reinterpret_cast(tl.addresses[0][tensor_idx]); - PARAM_T *p = reinterpret_cast(tl.addresses[1][tensor_idx]); - MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_idx]); - MOMENT_T *v = reinterpret_cast(tl.addresses[3][tensor_idx]); - - auto *rowwise_data = reinterpret_cast(tl.addresses[4][tensor_idx]); - auto *colwise_data = reinterpret_cast(tl.addresses[5][tensor_idx]); - auto *rowwise_scale_inv = reinterpret_cast(tl.addresses[6][tensor_idx]); - auto *colwise_scale_inv = reinterpret_cast(tl.addresses[7][tensor_idx]); - - const int64_t unpadded_scales_X_rowwise = (cols_val + MXFP8_TILE - 1) / MXFP8_TILE; - constexpr int64_t kRowwiseScaleAlign = 4; - const int64_t row_stride = - DIVUP_TO_MULTIPLE(unpadded_scales_X_rowwise, kRowwiseScaleAlign); - constexpr int64_t kColwiseScaleAlign = 128; - const int64_t col_stride = DIVUP_TO_MULTIPLE(cols_val, kColwiseScaleAlign); - const uint8_t dtype = tl.fp8_dtype[tensor_idx]; - const auto adam_mode = static_cast(mode); - - // Stage 3: initialize shared amax accumulators per row/col within the tile. - __shared__ float row_max_vals[MXFP8_TILE]; - __shared__ float col_max_vals[MXFP8_TILE]; - if (threadIdx.x < MXFP8_TILE) { - row_max_vals[threadIdx.x] = 0.0f; - col_max_vals[threadIdx.x] = 0.0f; - } - __syncthreads(); - - for (int t = threadIdx.x; t < MXFP8_TILE_ELEMS; t += blockDim.x) { - const int local_r = t / MXFP8_TILE; - const int local_c = t % MXFP8_TILE; - const int64_t r = row_base + local_r; - const int64_t c = col_base + local_c; - if (r >= rows_val || c >= cols_val) { - continue; - } - const index_t idx = static_cast(r * cols_val + c); - - float r_g = static_cast(g[idx]); - float r_p = static_cast(p[idx]); - float r_m = static_cast(m[idx]); - float r_v = static_cast(v[idx]); - - // Stage 4: apply Adam update in FP32 and write back updated p/m/v. - transformer_engine::multi_tensor_adam::adam_update( - r_g, r_p, r_m, r_v, beta1, beta2, beta1_correction, beta2_correction, epsilon, lr, - adam_mode, weight_decay); - - p[idx] = static_cast(r_p); - m[idx] = static_cast(r_m); - v[idx] = static_cast(r_v); - - // Stage 5: accumulate per-row/col absmax for MXFP8 scaling. - const float abs_p = fabsf(r_p); - transformer_engine::atomicMaxFloat(&row_max_vals[local_r], abs_p); - transformer_engine::atomicMaxFloat(&col_max_vals[local_c], abs_p); - } - - __syncthreads(); - - // Stage 6: write rowwise/colwise scale-inverse exponents for the tile. - const float max_norm_rcp = fp8_max_norm_rcp(dtype); - - for (int r = threadIdx.x; r < MXFP8_TILE; r += blockDim.x) { - const int64_t row = row_base + r; - if (row >= rows_val) { - continue; - } - const float amax = row_max_vals[r]; - const ::transformer_engine::e8m0_t biased_exponent = - transformer_engine::ptx::float_to_e8m0(amax * max_norm_rcp); - const size_t scale_idx = static_cast(row * row_stride + tile_col); - rowwise_scale_inv[scale_idx] = reinterpret_cast(biased_exponent); - } - - for (int c = threadIdx.x; c < MXFP8_TILE; c += blockDim.x) { - const int64_t col = col_base + c; - if (col >= cols_val) { - continue; - } - const float amax = col_max_vals[c]; - const ::transformer_engine::e8m0_t biased_exponent = - transformer_engine::ptx::float_to_e8m0(amax * max_norm_rcp); - const size_t scale_idx = static_cast(tile_row * col_stride + col); - colwise_scale_inv[scale_idx] = reinterpret_cast(biased_exponent); - } - - __syncthreads(); - - // Stage 7: quantize updated params to MXFP8 using rowwise and colwise scales. - for (int t = threadIdx.x; t < MXFP8_TILE_ELEMS; t += blockDim.x) { - const int local_r = t / MXFP8_TILE; - const int local_c = t % MXFP8_TILE; - const int64_t r = row_base + local_r; - const int64_t c = col_base + local_c; - if (r >= rows_val || c >= cols_val) { - continue; - } - const index_t idx = static_cast(r * cols_val + c); - const float r_p = static_cast(p[idx]); - - const size_t row_scale_idx = static_cast(r * row_stride + tile_col); - const uint8_t row_raw = rowwise_scale_inv[row_scale_idx]; - const ::transformer_engine::e8m0_t row_biased = - reinterpret_cast(row_raw); - const float row_scale_inv = transformer_engine::ptx::exp2f_rcp(row_biased); - if (dtype == static_cast(transformer_engine::DType::kFloat8E4M3)) { - auto *out = reinterpret_cast(rowwise_data); - out[idx] = cast_to_fp8(r_p * row_scale_inv); - } else { - auto *out = reinterpret_cast(rowwise_data); - out[idx] = cast_to_fp8(r_p * row_scale_inv); - } - - const size_t col_scale_idx = static_cast(tile_row * col_stride + c); - const uint8_t col_raw = colwise_scale_inv[col_scale_idx]; - const ::transformer_engine::e8m0_t col_biased = - reinterpret_cast(col_raw); - const float col_scale_inv = transformer_engine::ptx::exp2f_rcp(col_biased); - if (dtype == static_cast(transformer_engine::DType::kFloat8E4M3)) { - auto *out = reinterpret_cast(colwise_data); - out[idx] = cast_to_fp8(r_p * col_scale_inv); - } else { - auto *out = reinterpret_cast(colwise_data); - out[idx] = cast_to_fp8(r_p * col_scale_inv); - } - } -} - -inline bool requires_64bit_indexing(const std::vector> &tensor_lists) { - const size_t num_tensor_lists = tensor_lists.size(); - const size_t num_tensors_per_list = tensor_lists[0].size(); - for (size_t i = 0; i < num_tensor_lists; ++i) { - for (size_t j = 0; j < num_tensors_per_list; ++j) { - if (tensor_lists[i][j]->numel() >= INT_MAX) { - return true; - } - } - } - return false; -} - void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, std::vector> tensor_lists, const float lr, const float beta1, const float beta2, const float epsilon, @@ -820,13 +624,25 @@ void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, } } - const bool use_64bit_indexing = requires_64bit_indexing(tensor_lists); + // Check if 64-bit indices are required + bool requires_64bit_indexing = false; + for (size_t i = 0; i < num_tensor_lists; i++) { + for (size_t j = 0; j < num_tensors_per_list; j++) { + if (tensor_lists[i][j]->numel() >= INT_MAX) { + requires_64bit_indexing = true; + break; + } + } + if (requires_64bit_indexing) { + break; + } + } // Get moment dtype (m and v have the same dtype, already validated above) const auto moment_type_te = tensor_lists[2][0]->dtype(); // Launch kernel - if (use_64bit_indexing) { + if (requires_64bit_indexing) { if (num_tensor_lists == 4) { // g, p, m, v TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( @@ -950,41 +766,28 @@ void multi_tensor_adam_param_remainder_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK_CUDA(cudaGetLastError()); } -inline std::pair compute_bias_correction(int bias_correction, float beta1, - float beta2, int step) { - float bias_correction1 = 1.0f; - float bias_correction2 = 1.0f; +void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, + std::vector> tensor_lists, const float lr, + const float beta1, const float beta2, const float epsilon, + const int step, const int mode, const int bias_correction, + const float weight_decay, const DType fp8_dtype, + cudaStream_t stream) { + // Handle bias correction mode + float bias_correction1 = 1.0f, bias_correction2 = 1.0f; if (bias_correction == 1) { bias_correction1 = 1 - std::pow(beta1, step); bias_correction2 = 1 - std::pow(beta2, step); } - return {bias_correction1, bias_correction2}; -} -inline void check_tensor_list_sizes(const std::vector> &tensor_lists, - size_t expected_lists) { + // Check tensor list sizes + // 8 tensor lists: g, p_fp8, m, v, p_master, scale, amax, scale_inv const size_t num_tensor_lists = tensor_lists.size(); - NVTE_CHECK(num_tensor_lists == expected_lists, "Expected ", expected_lists, - " tensor lists, but found ", num_tensor_lists); + NVTE_CHECK(num_tensor_lists == 8, "Expected 8 tensor lists, but found ", num_tensor_lists); const size_t num_tensors_per_list = tensor_lists[0].size(); - for (size_t i = 1; i < num_tensor_lists; ++i) { + for (size_t i = 1; i < num_tensor_lists; i++) { NVTE_CHECK(tensor_lists[i].size() == num_tensors_per_list, "Tensor list ", i, " has size=", tensor_lists[i].size(), ", but expected size=", num_tensors_per_list); } -} - - -void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, - std::vector> tensor_lists, const float lr, - const float beta1, const float beta2, const float epsilon, - const int step, const int mode, const int bias_correction, - const float weight_decay, const DType fp8_dtype, - cudaStream_t stream) { - auto [bias_correction1, bias_correction2] = - compute_bias_correction(bias_correction, beta1, beta2, step); - check_tensor_list_sizes(tensor_lists, 8); - const size_t num_tensor_lists = tensor_lists.size(); - const size_t num_tensors_per_list = tensor_lists[0].size(); // Check tensor dtypes const auto g_in_type_te = tensor_lists[0][0]->dtype(); @@ -1016,10 +819,22 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, ", but expected dtype=", to_string(DType::kFloat32)); } - const bool use_64bit_indexing = requires_64bit_indexing(tensor_lists); + // Check if 64-bit indices are required + bool requires_64bit_indexing = false; + for (size_t i = 0; i < num_tensor_lists; i++) { + for (size_t j = 0; j < num_tensors_per_list; j++) { + if (tensor_lists[i][j]->numel() >= INT_MAX) { + requires_64bit_indexing = true; + break; + } + } + if (requires_64bit_indexing) { + break; + } + } // Launch kernel - if (use_64bit_indexing) { + if (requires_64bit_indexing) { TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( fp8_dtype, FP8_T, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( @@ -1041,76 +856,6 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK_CUDA(cudaGetLastError()); } -void multi_tensor_adam_mxfp8_cuda(int chunk_size, Tensor noop_flag, - std::vector> tensor_lists, const float lr, - const float beta1, const float beta2, const float epsilon, - const int step, const int mode, const int bias_correction, - const float weight_decay, const DType fp8_dtype, - cudaStream_t stream) { - auto [bias_correction1, bias_correction2] = - compute_bias_correction(bias_correction, beta1, beta2, step); - check_tensor_list_sizes(tensor_lists, 8); - const size_t num_tensor_lists = tensor_lists.size(); - const size_t num_tensors_per_list = tensor_lists[0].size(); - - NVTE_CHECK(fp8_dtype == DType::kFloat8E4M3 || fp8_dtype == DType::kFloat8E5M2, - "fp8_dtype must be E4M3 or E5M2 for MXFP8 fused Adam."); - - // Check tensor dtypes - const auto g_in_type_te = tensor_lists[0][0]->dtype(); - const auto p_in_type_te = tensor_lists[1][0]->dtype(); - const auto moment_type_te = tensor_lists[2][0]->dtype(); - for (size_t j = 0; j < num_tensors_per_list; ++j) { - NVTE_CHECK(tensor_lists[0][j]->dtype() == g_in_type_te, "Grad tensor ", j, - " has dtype=", to_string(tensor_lists[0][j]->dtype()), - ", but expected dtype=", to_string(g_in_type_te)); - NVTE_CHECK(tensor_lists[1][j]->dtype() == p_in_type_te, "Param tensor ", j, - " has dtype=", to_string(tensor_lists[1][j]->dtype()), - ", but expected dtype=", to_string(p_in_type_te)); - { - const bool m_is_fp32 = tensor_lists[2][j]->dtype() == DType::kFloat32; - const bool m_is_bf16 = tensor_lists[2][j]->dtype() == DType::kBFloat16; - const bool v_is_fp32 = tensor_lists[3][j]->dtype() == DType::kFloat32; - const bool v_is_bf16 = tensor_lists[3][j]->dtype() == DType::kBFloat16; - NVTE_CHECK((m_is_fp32 && v_is_fp32) || (m_is_bf16 && v_is_bf16), - "First and second moment tensors must both be Float32 or both be BFloat16, but " - "tensor ", - j, " has first moment dtype=", to_string(tensor_lists[2][j]->dtype()), - " and second moment dtype=", to_string(tensor_lists[3][j]->dtype())); - } - } - - const bool use_64bit_indexing = requires_64bit_indexing(tensor_lists); - - if (use_64bit_indexing) { - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - p_in_type_te, p_in_type, - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - g_in_type_te, g_in_type, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( - moment_type_te, moment_type, - multi_tensor_apply_mxfp8< - transformer_engine::multi_tensor_adam::adam_mxfp8_fused_kernel< - p_in_type, g_in_type, moment_type, int64_t>>( - chunk_size, noop_flag, tensor_lists, static_cast(fp8_dtype), stream, - beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, mode, - weight_decay);))); - } else { - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - p_in_type_te, p_in_type, - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - g_in_type_te, g_in_type, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( - moment_type_te, moment_type, - multi_tensor_apply_mxfp8< - transformer_engine::multi_tensor_adam::adam_mxfp8_fused_kernel< - p_in_type, g_in_type, moment_type, int32_t>>( - chunk_size, noop_flag, tensor_lists, static_cast(fp8_dtype), stream, - beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, mode, - weight_decay);))); - } -} - void multi_tensor_adam_capturable_cuda(int chunk_size, Tensor noop_flag, std::vector> tensor_lists, Tensor lr, const float beta1, const float beta2, const float epsilon, @@ -1273,19 +1018,6 @@ void nvte_multi_tensor_adam_fp8_cuda(int chunk_size, NVTETensor noop_flag, epsilon, step, mode, bias_correction, weight_decay, static_cast(fp8_dtype), stream); } -void nvte_multi_tensor_adam_mxfp8_cuda( - int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, - const size_t num_tensor_lists, const size_t num_tensors_per_list, const NVTEDType fp8_dtype, - const float lr, const float beta1, const float beta2, const float epsilon, const int step, - const int mode, const int bias_correction, const float weight_decay, cudaStream_t stream) { - NVTE_API_CALL(nvte_multi_tensor_adam_mxfp8_cuda); - using namespace transformer_engine; - multi_tensor_adam::multi_tensor_adam_mxfp8_cuda( - chunk_size, *convertNVTETensorCheck(noop_flag), - convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list), lr, beta1, beta2, - epsilon, step, mode, bias_correction, weight_decay, static_cast(fp8_dtype), stream); -} - void nvte_multi_tensor_adam_capturable_cuda( int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, const size_t num_tensor_lists, const size_t num_tensors_per_list, NVTETensor lr, const float beta1, const float beta2, diff --git a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh index c334f3908e..3062ead551 100644 --- a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh +++ b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh @@ -35,23 +35,6 @@ struct TensorListMetadata : public TensorListMetadataBase { void *fp8_meta_addresses[3][depth_to_max_tensors[n - 1]]; }; -constexpr int MXFP8_TILE = 32; -constexpr int MXFP8_TILE_ELEMS = MXFP8_TILE * MXFP8_TILE; -constexpr int MXFP8_BLOCK_THREADS = 256; -constexpr int MXFP8_MAX_TENSORS = 24; -constexpr int MXFP8_MAX_BLOCKS = 320; - -struct MXFP8TensorListMetadata { - void *addresses[8][MXFP8_MAX_TENSORS]; - int sizes[MXFP8_MAX_TENSORS]; - int rows[MXFP8_MAX_TENSORS]; - int cols[MXFP8_MAX_TENSORS]; - uint8_t fp8_dtype[MXFP8_MAX_TENSORS]; - unsigned char block_to_tensor[MXFP8_MAX_BLOCKS]; - int block_to_tile[MXFP8_MAX_BLOCKS]; - int start_tensor_this_launch; -}; - template __global__ void multi_tensor_apply_kernel(int64_t chunk_size, volatile int *noop_flag, T tl, U callable, ArgTypes... args) { @@ -130,80 +113,3 @@ void multi_tensor_apply(int64_t block_size, int64_t chunk_size, } } } - -template -void multi_tensor_apply_mxfp8(int64_t chunk_size, const transformer_engine::Tensor &noop_flag, - std::vector> tensor_lists, - uint8_t fp8_dtype, cudaStream_t stream, ArgTypes... args) { - constexpr size_t kNumTensorLists = 8; - NVTE_CHECK(tensor_lists.size() == kNumTensorLists, - "Expected 8 tensor lists for MXFP8, but found ", tensor_lists.size()); - - const size_t num_tensors_per_list = tensor_lists[0].size(); - if (num_tensors_per_list == 0) { - return; - } - for (size_t i = 1; i < tensor_lists.size(); ++i) { - NVTE_CHECK(tensor_lists[i].size() == num_tensors_per_list, "Tensor list ", i, - " has size=", tensor_lists[i].size(), ", but expected size=", num_tensors_per_list); - } - - MXFP8TensorListMetadata tl; - tl.start_tensor_this_launch = 0; - int loc_block_info = 0; - int loc_tensor_info = 0; - - for (size_t t = 0; t < num_tensors_per_list; ++t) { - - const auto &g = tensor_lists[0][t]; - const auto &rowwise_data = tensor_lists[4][t]; - const auto &colwise_data = tensor_lists[5][t]; - - const int rows_val = static_cast(rowwise_data->data.shape[0]); - const int cols_val = static_cast(rowwise_data->data.shape[1]); - - tl.sizes[loc_tensor_info] = g->numel(); - tl.rows[loc_tensor_info] = rows_val; - tl.cols[loc_tensor_info] = cols_val; - tl.fp8_dtype[loc_tensor_info] = fp8_dtype; - - for (int d = 0; d < kNumTensorLists; ++d) { - tl.addresses[d][loc_tensor_info] = tensor_lists[d][t]->data.dptr; - } - loc_tensor_info++; - - const int tiles_y = (rows_val + MXFP8_TILE - 1) / MXFP8_TILE; - const int tiles_x = (cols_val + MXFP8_TILE - 1) / MXFP8_TILE; - const int tiles_this_tensor = tiles_y * tiles_x; - - for (int tile = 0; tile < tiles_this_tensor; ++tile) { - tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; - tl.block_to_tile[loc_block_info] = tile; - loc_block_info++; - - const bool blocks_full = (loc_block_info == MXFP8_MAX_BLOCKS); - const bool tensors_full = - (loc_tensor_info == MXFP8_MAX_TENSORS && tile == tiles_this_tensor - 1); - const bool last_tile = (t == num_tensors_per_list - 1 && tile == tiles_this_tensor - 1); - if (blocks_full || tensors_full || last_tile) { - Kernel<<>>( - chunk_size, reinterpret_cast(noop_flag.data.dptr), tl, args...); - NVTE_CHECK_CUDA(cudaGetLastError()); - loc_block_info = 0; - if (tile == tiles_this_tensor - 1) { - loc_tensor_info = 0; - tl.start_tensor_this_launch = t + 1; - } else { - tl.rows[0] = tl.rows[loc_tensor_info - 1]; - tl.cols[0] = tl.cols[loc_tensor_info - 1]; - tl.fp8_dtype[0] = tl.fp8_dtype[loc_tensor_info - 1]; - for (int d = 0; d < kNumTensorLists; ++d) { - tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; - } - loc_tensor_info = 1; - tl.start_tensor_this_launch = t; - } - } - } - } -} diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 65e2c54d67..1c5116a8da 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -517,12 +517,6 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, at::Tensor noop_flag, const int step, const int mode, const int bias_correction, const float weight_decay, DType fp8_dtype); -void multi_tensor_adam_mxfp8_cuda(int chunk_size, at::Tensor noop_flag, - std::vector> tensor_lists, const float lr, - const float beta1, const float beta2, const float epsilon, - const int step, const int mode, const int bias_correction, - const float weight_decay, DType fp8_dtype); - void multi_tensor_adam_capturable_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists, at::Tensor lr, const float beta1, const float beta2, diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp index 01a21d44bb..145e1d4b40 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp @@ -5,7 +5,6 @@ ************************************************************************/ #include "../../extensions.h" -#include "pybind.h" namespace transformer_engine::pytorch { @@ -52,25 +51,6 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, at::Tensor noop_flag, at::cuda::getCurrentCUDAStream()); } -void multi_tensor_adam_mxfp8_cuda(int chunk_size, at::Tensor noop_flag, - std::vector> tensor_lists, const float lr, - const float beta1, const float beta2, const float epsilon, - const int step, const int mode, const int bias_correction, - const float weight_decay, DType fp8_dtype) { - auto noop_flag_cu = makeTransformerEngineTensor(noop_flag); - auto [_, __, tensor_lists_ptr, num_lists, num_tensors] = - makeTransformerEngineTensorList(tensor_lists); - - NVTE_CHECK(num_lists == 8, - "Expected 8 tensor lists (g, p_master, m, v, rowwise_data, colwise_data, " - "rowwise_scale_inv, colwise_scale_inv), but found ", - num_lists); - nvte_multi_tensor_adam_mxfp8_cuda( - chunk_size, noop_flag_cu.data(), tensor_lists_ptr.data(), num_lists, num_tensors, - static_cast(fp8_dtype), lr, beta1, beta2, epsilon, step, mode, bias_correction, - weight_decay, at::cuda::getCurrentCUDAStream()); -} - void multi_tensor_adam_capturable_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists, at::Tensor lr, const float beta1, const float beta2, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 6def07b08e..c590a3c9e2 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -525,8 +525,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("multi_tensor_adam_fp8", &transformer_engine::pytorch::multi_tensor_adam_fp8_cuda, "Compute and apply gradient update to parameters for Adam optimizer", py::call_guard()); - m.def("multi_tensor_adam_mxfp8", &transformer_engine::pytorch::multi_tensor_adam_mxfp8_cuda, - "Compute and apply gradient update to parameters for Adam optimizer"); m.def("multi_tensor_adam_capturable", &transformer_engine::pytorch::multi_tensor_adam_capturable_cuda, "Compute and apply gradient update to parameters for Adam optimizer with CUDA graph " diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index f4ab2e7c37..bcfd2bef19 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -14,7 +14,6 @@ from torch.distributed._tensor import DTensor import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer -from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from .multi_tensor_apply import multi_tensor_applier @@ -190,7 +189,6 @@ def __init__( self.multi_tensor_adam = tex.multi_tensor_adam self.multi_tensor_adam_param_remainder = tex.multi_tensor_adam_param_remainder self.multi_tensor_adam_fp8 = tex.multi_tensor_adam_fp8 - self.multi_tensor_adam_mxfp8 = tex.multi_tensor_adam_mxfp8 self.multi_tensor_adam_capturable = tex.multi_tensor_adam_capturable self.multi_tensor_adam_capturable_master = tex.multi_tensor_adam_capturable_master @@ -546,27 +544,18 @@ def step(self, closure=None, grad_scaler=None): # create lists for multi-tensor apply p_main_of_fp8_model = [] p_main_of_f16_model = [] - p_main_of_mxfp8_model = [] g_of_fp8_model = [] g_of_f16_model = [] g_of_f32_model = [] - g_of_mxfp8_model = [] m_of_fp8_model = [] m_of_f16_model = [] m_of_f32_model = [] - m_of_mxfp8_model = [] v_of_fp8_model = [] v_of_f16_model = [] v_of_f32_model = [] - v_of_mxfp8_model = [] p_fp8_model = [] p_f16_model = [] p_f32_model = [] - # mxfp8 meta - p_mxfp8_rowwise = [] - p_mxfp8_colwise = [] - p_mxfp8_rowwise_scale_inv = [] - p_mxfp8_colwise_scale_inv = [] # fp8 meta scales = [] amaxes = [] @@ -634,30 +623,10 @@ def step(self, closure=None, grad_scaler=None): g_of_fp8_model.append(p_grad.data) m_of_fp8_model.append(unscaled_state["exp_avg"]) v_of_fp8_model.append(unscaled_state["exp_avg_sq"]) - elif isinstance(p, MXFP8Tensor) or ( - isinstance(p, DTensor) and isinstance(p._local_tensor, MXFP8Tensor) - ): - p = p._local_tensor if isinstance(p, DTensor) else p - if p._rowwise_data is None or p._columnwise_data is None: - raise RuntimeError("MXFP8Tensor does not have one of rowwise/columnwise data.") - if self.capturable: - raise RuntimeError( - "FusedAdam does not support MXFP8 model weights with capturable=True." - ) - if self.master_weights: - p_main_of_mxfp8_model.append(unscaled_state["master_param"].data) - g_of_mxfp8_model.append(p_grad.data) - m_of_mxfp8_model.append(unscaled_state["exp_avg"]) - v_of_mxfp8_model.append(unscaled_state["exp_avg_sq"]) - p_mxfp8_rowwise.append(p._rowwise_data) - p_mxfp8_colwise.append(p._columnwise_data) - p_mxfp8_rowwise_scale_inv.append(p._rowwise_scale_inv) - p_mxfp8_colwise_scale_inv.append(p._columnwise_scale_inv) - out_dtype = p._fp8_dtype elif isinstance(p, QuantizedTensor) or ( isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) ): - # Block-scaling quantized params (Float8BlockwiseQTensor, + # Block-scaling quantized params (MXFP8Tensor, Float8BlockwiseQTensor, # NVFP4Tensor). Operate on FP32 master weights, requantize back after # Adam update. # Note: a fused Adam+requantize kernel (like multi_tensor_adam_fp8 @@ -828,18 +797,6 @@ def apply_multi_tensor_adam(adam_func, tensor_lists, inv_scale=None, out_dtype=N scale_invs, ] apply_multi_tensor_adam(self.multi_tensor_adam_fp8, tensor_lists, out_dtype) - if len(p_mxfp8_rowwise) > 0 and len(p_mxfp8_colwise) > 0: - tensor_lists = [ - g_of_mxfp8_model, - p_main_of_mxfp8_model, - m_of_mxfp8_model, - v_of_mxfp8_model, - p_mxfp8_rowwise, - p_mxfp8_colwise, - p_mxfp8_rowwise_scale_inv, - p_mxfp8_colwise_scale_inv, - ] - apply_multi_tensor_adam(self.multi_tensor_adam_mxfp8, tensor_lists, out_dtype) if len(p_f32_model) > 0: tensor_lists = [ g_of_f32_model, From 6b95e608e5f67b65eb3053e37b0d628b00315bc7 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 22 Mar 2026 22:18:48 +0000 Subject: [PATCH 08/18] change distributed tests infra for fsdp2 Signed-off-by: Varun Thumbe --- .../distributed/fsdp2_tests/conftest.py | 107 +++++++ .../{ => fsdp2_tests}/run_fsdp2_fused_adam.py | 207 +++++++++----- .../{ => fsdp2_tests}/run_fsdp2_model.py | 172 ++++++----- tests/pytorch/distributed/test_torch_fsdp2.py | 267 +++--------------- .../pytorch/tensor/float8_blockwise_tensor.py | 115 +++----- 5 files changed, 421 insertions(+), 447 deletions(-) create mode 100644 tests/pytorch/distributed/fsdp2_tests/conftest.py rename tests/pytorch/distributed/{ => fsdp2_tests}/run_fsdp2_fused_adam.py (80%) rename tests/pytorch/distributed/{ => fsdp2_tests}/run_fsdp2_model.py (77%) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py new file mode 100644 index 0000000000..4d75f91317 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -0,0 +1,107 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared pytest fixtures and utilities for FSDP2 distributed tests. + +Fixtures defined here (dist_init, _cleanup, recipe_name) are auto-discovered +by pytest for every test module in this directory. Utility functions +(get_recipe_from_string, save_custom_attrs, restore_custom_attrs) can be +imported normally: ``from conftest import get_recipe_from_string``. +""" + +import gc +import os + +import pytest + +import torch +import torch.distributed as dist + +from transformer_engine.pytorch import fp8, QuantizedTensor +import transformer_engine.common.recipe + + +# ── FP8 recipe parametrization ────────────────────────────────────── +def _check_nvfp4_support(): + supported, reason = fp8.check_nvfp4_support() + if supported and torch.cuda.get_device_capability()[0] == 12: + return ( + False, + ( + "NVFP4BlockScaling is failing on SM120 with " + "hadamard_transform/hadamard_transform_cast_fusion.cu:672 in function " + "rht_gemm_ntt_w_sfc: CUDA Error: invalid argument" + ), + ) + return supported, reason + + +_FP8_RECIPE_CONFIGS = [ + ("DelayedScaling", fp8.check_fp8_support), + ("Float8CurrentScaling", fp8.check_fp8_support), + ("Float8BlockScaling", fp8.check_fp8_block_scaling_support), + ("MXFP8BlockScaling", fp8.check_mxfp8_support), + ("NVFP4BlockScaling", _check_nvfp4_support), +] + + +def _parametrize_recipes(): + params = [] + for name, check_fn in _FP8_RECIPE_CONFIGS: + supported, reason = check_fn() + params.append( + pytest.param(name, id=name, marks=pytest.mark.skipif(not supported, reason=reason)) + ) + return params + + +# ── Session / per-test fixtures ────────────────────────────────────── +@pytest.fixture(scope="session", autouse=True) +def dist_init(): + """Initialize the distributed process group once for the entire pytest session.""" + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + yield + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.fixture(autouse=True) +def _cleanup(): + """Release GPU memory and stale NCCL state between tests.""" + yield + gc.collect() + torch.cuda.empty_cache() + + +@pytest.fixture(params=_parametrize_recipes()) +def recipe_name(request): + return request.param + + +# ── Other Shared helpers ─────────────────────────────────────────────────── +def get_recipe_from_string(recipe): + return getattr(transformer_engine.common.recipe, recipe)() + + +def save_custom_attrs(module): + custom_attrs = {} + for name, param in module.named_parameters(): + if isinstance(param, QuantizedTensor): + ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] + else: + ignore_keys = [] + attrs = vars(param) + custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} + return custom_attrs + + +def restore_custom_attrs(module, custom_attrs): + for name, param in module.named_parameters(): + if name in custom_attrs: + for attr_name, attr_value in custom_attrs[name].items(): + setattr(param, attr_name, attr_value) diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py similarity index 80% rename from tests/pytorch/distributed/run_fsdp2_fused_adam.py rename to tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index c39957cf13..049b935878 100644 --- a/tests/pytorch/distributed/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -6,12 +6,29 @@ """FSDP2 + FusedAdam compatibility tests. -Launched via torchrun from test_fused_optimizer.py. +Run all tests (via torchrun + pytest): + torchrun -m pytest -v --tb=short + +Run a single test standalone (for debugging): + torchrun --test --recipe + +Available --test values: + fused_adam_fp8_master_weights, fused_adam_fp8_master_weights_no_meta, + fused_adam_bf16, fused_adam_fp8_no_master, fused_adam_bf16_store_param_remainders, + fuse_wgrad_accumulation, dcp_output_parity, dcp_output_parity_async, + safetensors_fp32_export + +Available --recipe values: + DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, NVFP4BlockScaling """ import argparse import functools import os +import pathlib +import sys +import pytest import torch import torch.distributed as dist @@ -23,10 +40,8 @@ import transformer_engine.pytorch as te from transformer_engine.pytorch import QuantizedTensor import transformer_engine.common.recipe - - -def get_recipe_from_string(recipe): - return getattr(transformer_engine.common.recipe, recipe)() +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from conftest import get_recipe_from_string, save_custom_attrs, restore_custom_attrs HIDDEN_SIZE = 256 @@ -38,38 +53,6 @@ def get_recipe_from_string(recipe): NUM_STEPS = 3 -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - if isinstance(param, QuantizedTensor): - ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] - else: - ignore_keys = [] - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) - - -def _setup(): - """Common distributed setup. Returns (world_size, local_rank, device).""" - world_size = int(os.environ["WORLD_SIZE"]) - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - # CPU backend required for async save - dist.init_process_group(backend="cpu:gloo,cuda:nccl") - device = torch.device(f"cuda:{local_rank}") - torch.manual_seed(42) - torch.cuda.manual_seed(42) - return world_size, local_rank, device - - def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_device=True): """Build a Sequential of TransformerLayers, optionally with FP8 init. @@ -142,8 +125,13 @@ def _shard_model(model, world_size): restore_custom_attrs(model, custom_attrs) return model +def _get_dist_info(): + """Get world_size and device from environment (PG already initialized by session fixture).""" + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + return world_size, device -def test_fused_adam_fp8_master_weights(recipe=None): +def test_fused_adam_fp8_master_weights(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). Verifies: @@ -151,7 +139,15 @@ def test_fused_adam_fp8_master_weights(recipe=None): - Training loop completes without error - DTensor wrapping and QuantizedTensor local tensors are preserved """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name == "NVFP4BlockScaling": + pytest.xfail( + f"{recipe_name}: quantized_model_init and FSDP2 is not currently supported, since the " + "block tensor is dequantized before we flatten it for FSDP2." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -206,10 +202,8 @@ def test_fused_adam_fp8_master_weights(recipe=None): ) assert qt_count > 0, "No QuantizedTensor local tensors after training" - dist.destroy_process_group() - -def test_fused_adam_fp8_master_weights_no_meta(recipe=None): +def test_fused_adam_fp8_master_weights_no_meta(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init WITHOUT meta device. This is the legacy path that creates quantized params directly on CUDA. @@ -219,7 +213,16 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe=None): For per-tensor FP8 (DelayedScaling, Float8CurrentScaling) this works because Float8Tensor's storage is accessible via data_ptr(). """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): + pytest.xfail( + f"{recipe_name}: FSDP2 without meta-device init crashes on block-scaling " + "QuantizedTensor wrapper subclasses (data_ptr() == 0). " + "Use device='meta' + reset_parameters() after sharding." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe, use_meta_device=False) model = _shard_model(model, world_size) @@ -231,7 +234,9 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe=None): master_weight_dtype=torch.float32, ) - x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + x = torch.randn( + SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device + ) target = torch.randn_like(x) for step in range(NUM_STEPS): @@ -242,15 +247,15 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe=None): loss.backward() optimizer.step() - dist.destroy_process_group() - -def test_fused_adam_bf16(recipe=None): +def test_fused_adam_bf16(recipe_name): """FusedAdam with master_weights + FSDP2 + bf16 params (no FP8). Verifies the non-FP8 DTensor param path in step() works correctly. """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=False) model = _shard_model(model, world_size) @@ -284,15 +289,21 @@ def test_fused_adam_bf16(recipe=None): # Verify loss decreased (basic sanity) assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" - dist.destroy_process_group() - -def test_fused_adam_fp8_no_master(recipe=None): +def test_fused_adam_fp8_no_master(recipe_name): """FusedAdam without master_weights + FSDP2 + FP8 params. Verifies FusedAdam works with FSDP2 even without master weights enabled. """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): + pytest.xfail( + f"{recipe_name}: FusedAdam without master_weights does not support " + "block-scaling quantized tensors. Use master_weights=True." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -318,10 +329,8 @@ def test_fused_adam_fp8_no_master(recipe=None): for name, param in model.named_parameters(): assert isinstance(param, DTensor), f"{name} lost DTensor wrapping" - dist.destroy_process_group() - -def test_fused_adam_bf16_store_param_remainders(recipe=None): +def test_fused_adam_bf16_store_param_remainders(recipe_name): """FusedAdam with master_weights + store_param_remainders + FSDP2 + bf16 params. store_param_remainders stores only the trailing 16 remainder bits (int16) @@ -335,7 +344,8 @@ def test_fused_adam_bf16_store_param_remainders(recipe=None): - exp_avg and exp_avg_sq are float32 - Loss decreases (basic sanity) """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() model = _build_model(fp8_init=False) model = _shard_model(model, world_size) @@ -385,10 +395,18 @@ def test_fused_adam_bf16_store_param_remainders(recipe=None): # Verify loss decreased (basic sanity) assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" - dist.destroy_process_group() - -def test_fuse_wgrad_accumulation(recipe=None): +@pytest.mark.xfail( + reason=( + "fuse_wgrad_accumulation is incompatible with vanilla FSDP2: " + "autograd Function.apply unwraps DTensors to local tensors, so " + "main_grad (set on the DTensor) is inaccessible during backward. " + "Additionally, the fused wgrad GEMM bypasses FSDP2's reduce-scatter." + ), + raises=AttributeError, + strict=True, +) +def test_fuse_wgrad_accumulation(recipe_name): """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail. With vanilla FSDP2, PyTorch's autograd Function.apply unwraps DTensor @@ -400,8 +418,8 @@ def test_fuse_wgrad_accumulation(recipe=None): writes the gradient directly into main_grad and returns None to autograd, bypassing FSDP2's reduce-scatter. """ - world_size, _, device = _setup() - + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, fuse_wgrad_accumulation=True, recipe=recipe) # Allocate main_grad buffers on the DTensor params @@ -433,10 +451,8 @@ def test_fuse_wgrad_accumulation(recipe=None): loss = F.mse_loss(output, target) loss.backward() # Expected to raise AttributeError - dist.destroy_process_group() - -def test_safetensors_fp32_export(recipe=None): +def test_safetensors_fp32_export(recipe_name): """Export full-precision (FP32) model to safetensors from optimizer master weights. Verifies: @@ -446,6 +462,13 @@ def test_safetensors_fp32_export(recipe=None): - All saved tensors are float32 - Saved tensor shapes match expected (unsharded) shapes """ + recipe = get_recipe_from_string(recipe_name) + if recipe_name == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access" + ) + from safetensors.torch import load_file, save_file from torch.distributed.checkpoint.state_dict import ( StateDictOptions, @@ -453,8 +476,7 @@ def test_safetensors_fp32_export(recipe=None): get_optimizer_state_dict, ) - world_size, _, device = _setup() - + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -511,10 +533,9 @@ def test_safetensors_fp32_export(recipe=None): # Clean up. os.remove(save_path) - dist.destroy_process_group() - -def test_dcp_output_parity(recipe=None, async_save=False): +@pytest.mark.parametrize("async_save", [False, True], ids=["sync", "async"]) +def test_dcp_output_parity(recipe_name, async_save): """DCP save/load round-trip produces bitwise-identical model outputs. 1. Builds and trains a model for NUM_STEPS @@ -525,9 +546,42 @@ def test_dcp_output_parity(recipe=None, async_save=False): 6. Runs the same forward pass and asserts outputs are identical 7. Runs one more training step on both models and asserts outputs still match """ + recipe = get_recipe_from_string(recipe_name) + + if recipe_name == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access: " + "/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh:92 in function " + "multi_tensor_apply: CUDA Error: an illegal memory access was encountered" + ) + + if recipe_name == "NVFP4BlockScaling": + pytest.xfail( + "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " + "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" + ) + + + if recipe_name == "Float8BlockScaling" and not async_save and torch.cuda.get_device_capability()[0] == 12: + pytest.xfail( + "Float8BlockScaling is failing on SM120 with RuntimeError: " + "transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu:534 " + "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " + "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " + "requires using power of two scaling factors." + ) + if recipe_name == "Float8BlockScaling" and async_save: + pytest.xfail( + "Float8BlockScaling: async DCP save/load round-trip produces different model " + "outputs — quantization metadata (scales) is not correctly persisted through " + "async distributed checkpointing. On SM120, additionally fails with pow2_scale " + "assertion in quantize_transpose_vector_blockwise." + ) + import torch.distributed.checkpoint as dcp - world_size, local_rank, device = _setup() + world_size, device = _get_dist_info() # ── Build and train the original model ─────────────────────────── model = _build_model(fp8_init=True, recipe=recipe) @@ -674,7 +728,6 @@ def test_dcp_output_parity(recipe=None, async_save=False): if int(os.environ.get("RANK", "0")) == 0: shutil.rmtree(checkpoint_dir, ignore_errors=True) - dist.destroy_process_group() TESTS = { @@ -707,5 +760,13 @@ def test_dcp_output_parity(recipe=None, async_save=False): ], ) args = parser.parse_args() - recipe = get_recipe_from_string(args.recipe) - TESTS[args.test](recipe) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + try: + TESTS[args.test](args.recipe) + finally: + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/tests/pytorch/distributed/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py similarity index 77% rename from tests/pytorch/distributed/run_fsdp2_model.py rename to tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 60d7cd2023..e894a9d0b0 100644 --- a/tests/pytorch/distributed/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -4,9 +4,37 @@ # # See LICENSE for license information. +"""FSDP2 model sharding tests. + +Run all tests (via torchrun + pytest): + torchrun -m pytest -v --tb=short + +Run standalone (for debugging): + torchrun --recipe [options] + +Available --recipe values: + DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, NVFP4BlockScaling + +Other options: + --fp8-init Initialize weights in FP8 + --layer-type TYPE Linear, LayerNormLinear, LayerNormMLP, + MultiheadAttention, TransformerLayer (default) + --sharding-dims N [M] FSDP dims, e.g. "2" or "2 2" for HSDP + --num-layers N Number of layers (default: 4) + --iter N Training iterations (default: 10) + --device cuda|meta Device for init (default: meta) +""" + +import gc import os +import pathlib import sys import argparse +from types import SimpleNamespace +from contextlib import nullcontext + +import pytest import transformer_engine.pytorch as te import transformer_engine.common.recipe @@ -19,14 +47,13 @@ from torch.distributed import DeviceMesh from torch.distributed._composable.fsdp import fully_shard from torch.distributed.device_mesh import init_device_mesh -from transformer_engine.pytorch import QuantizedTensor -from contextlib import nullcontext -LOCAL_RANK = None +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) +from conftest import get_recipe_from_string, save_custom_attrs, restore_custom_attrs def dist_print(msg): - if LOCAL_RANK == 0: + if int(os.getenv("LOCAL_RANK", "0")) == 0: print(msg) @@ -114,10 +141,6 @@ def get_te_layer_from_string(layer_name): return te_layer_map[layer_name.lower()] -def get_recipe_from_string(recipe): - return getattr(transformer_engine.common.recipe, recipe)() - - def init_te_model(config): hidden_size = config.num_heads * config.head_dim args = [hidden_size, hidden_size] @@ -188,31 +211,8 @@ def shard_model_with_fsdp2(model, mesh): return model -#### Methods to save the custom attributes of QuantizedTensors before sharding -#### them with FSDP2, and restore them after sharding. -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - if isinstance(param, QuantizedTensor): - # Ignore FP8 metadata attributes. Otherwise we will save duplicate copies - # for data/transpose FP8 tensors on top of FP8 tensors that FSDP2 will save. - ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] - else: - ignore_keys = [] - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) - - @torch.no_grad() -def test_fp8_fsdp2_allgather(model): +def _check_fp8_fsdp2_allgather(model): # Do manual allgather in fp32 and match against fp8 allgather done # with fsdp2 # FP32 manual weight allgather @@ -249,37 +249,15 @@ def test_fp8_fsdp2_allgather(model): module.reshard() -def _train(args): - global LOCAL_RANK - assert "TORCHELASTIC_RUN_ID" in os.environ - WORLD_RANK = int(os.getenv("RANK", "0")) - WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) - LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) - LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) - assert LOCAL_SIZE == WORLD_SIZE - - # Set device and initialize RNG states - torch.cuda.set_device(WORLD_RANK) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) +def _run_training(args): + """Core training logic. Assumes dist is already initialized.""" + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + world_size = int(os.getenv("WORLD_SIZE", "1")) - # Initialize torch.distributed global process group and get DP/TP groups - dist_init_kwargs = { - "backend": "nccl", - "rank": WORLD_RANK, - "world_size": WORLD_SIZE, - } - assert dist.is_nccl_available() - dist.init_process_group(**dist_init_kwargs) - nccl_world = dist.new_group(backend="nccl") - device = torch.device(f"cuda:{LOCAL_RANK}") - - # FP8 Configuration fp8_recipe = get_recipe_from_string(args.recipe) build_model_context_args = {} if not args.fp8_init: - # Build model context (FP8 init) build_model_context = nullcontext else: from transformer_engine.pytorch import fp8_model_init @@ -289,7 +267,6 @@ def _train(args): build_model_context_args["recipe"] = fp8_recipe dist_print(f"Memory before model init: {torch.cuda.memory_allocated(device) / 1e6} MB") - # Create the model on the meta/cuda device as per args with build_model_context(**build_model_context_args): model, inp_shape, out_shape = init_te_model(args) dist_print( @@ -297,37 +274,30 @@ def _train(args): f" {torch.cuda.memory_allocated(device) / 1e6} MB" ) - # Creating a DeviceMesh for fully_shard - world_size = int(WORLD_SIZE) - # Setup the sharding mesh for FSDP/HSDP mesh = get_device_mesh(world_size, args.sharding_dims) custom_attrs = save_custom_attrs(model) model = shard_model_with_fsdp2(model, mesh) restore_custom_attrs(model, custom_attrs) - # model now has DTensors as its parameters if args.device == "meta": - # After FSDP2 has been applied, materialize and initialize the sharded parameters - # TE base.py's reset_parameters() handles DTensors with FP8 initialization for module in model.modules(): if hasattr(module, "reset_parameters"): module.reset_parameters() dist_print(f" Sharded parameters materialized and initialized on cuda device.") dist_print( - f"FSDP2 model in cuda, memory allocated: {torch.cuda.memory_allocated(device) / 1e6} MB" + f"FSDP2 model in cuda, memory allocated:" + f" {torch.cuda.memory_allocated(device) / 1e6} MB" ) optimizer = optim.Adam(model.parameters(), lr=1e-3) for iteration in range(args.iter): - # Zero the parameter gradients optimizer.zero_grad() input_data = torch.randn(inp_shape, device=device) target = torch.randn(out_shape, device=device) - # NVFP4BlockScaling requires bfloat16 inputs in both the forward and backward passes. with ( torch.autocast(device_type="cuda", dtype=torch.bfloat16) if args.recipe == "NVFP4BlockScaling" @@ -341,14 +311,72 @@ def _train(args): optimizer.step() dist_print(f"Iteration {iteration} completed with loss {loss.item()}") - # Some of the FSDP states are lazy initialized during FSDP forward pass - # so testing fp8 allgather at the end of the training loop. if args.fp8_init: - test_fp8_fsdp2_allgather(model) + _check_fp8_fsdp2_allgather(model) + + +def _train(args): + """Standalone entry point with full dist lifecycle.""" + assert "TORCHELASTIC_RUN_ID" in os.environ + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + assert LOCAL_SIZE == WORLD_SIZE + + torch.cuda.set_device(LOCAL_RANK) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + assert dist.is_nccl_available() + dist.init_process_group( + backend="nccl", rank=WORLD_RANK, world_size=WORLD_SIZE, + ) + try: + _run_training(args) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + torch.cuda.empty_cache() + gc.collect() - dist.destroy_process_group() return 0 +# ── Pytest test function ───────────────────────────────────────────── + +NUM_PROCS = int(os.environ.get("WORLD_SIZE", "1")) + + +@pytest.mark.parametrize("sharding_dims", [[NUM_PROCS], [2, NUM_PROCS // 2]]) +@pytest.mark.parametrize("fp8_init", [False, True]) +@pytest.mark.parametrize("layer_type", ["LayerNormLinear", "TransformerLayer"]) +def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): + if recipe_name in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: + pytest.xfail( + f"{recipe_name} + fp8_init: test_fp8_fsdp2_allgather is currently failing." + ) + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + args = SimpleNamespace( + recipe=recipe_name, + fp8_init=fp8_init, + sharding_dims=list(sharding_dims), + layer_type=layer_type, + seed=42, + num_heads=8, + head_dim=64, + batch_size=16, + seq_length=128, + params_dtype="float32", + num_layers=4, + iter=10, + device="meta", + ) + _run_training(args) + + if __name__ == "__main__": sys.exit(_train(_parse_args())) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 02e45d99cb..d5918adb8d 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -10,242 +10,55 @@ import torch import transformer_engine.pytorch as te -from transformer_engine.pytorch import fp8 NUM_PROCS: int = torch.cuda.device_count() -def check_nvfp4_support(): - supported, reason = fp8.check_nvfp4_support() - if supported and torch.cuda.get_device_capability()[0] == 12: - return ( - False, - ( - "NVFP4BlockScaling is failing on SM120 with " - "hadamard_transform/hadamard_transform_cast_fusion.cu:672 in function " - "rht_gemm_ntt_w_sfc: CUDA Error: invalid argument" - ), - ) - - return supported, reason - - -# Each entry: (recipe_class_name, check_fn) -_FP8_RECIPE_CONFIGS = [ - ("DelayedScaling", fp8.check_fp8_support), - ("Float8CurrentScaling", fp8.check_fp8_support), - ("Float8BlockScaling", fp8.check_fp8_block_scaling_support), - ("MXFP8BlockScaling", fp8.check_mxfp8_support), - ("NVFP4BlockScaling", check_nvfp4_support), -] - - -def _parametrize_fp8_recipes(): - """Generate pytest.param objects with skip marks for unsupported FP8 recipes.""" - params = [] - for name, check_fn in _FP8_RECIPE_CONFIGS: - supported, reason = check_fn() - params.append( - pytest.param( - name, - id=name, - marks=pytest.mark.skipif(not supported, reason=reason), - ) - ) - return params - - -@pytest.fixture(params=_parametrize_fp8_recipes()) -def fp_recipe(request): - """Parametrized fixture providing FP8 recipe Hydra overrides for each supported TE recipe.""" - return request.param - - -def _run_test(fp_init, sharding_dims, recipe, layer_type): - test_path = Path(__file__).parent.resolve() / "run_fsdp2_model.py" - test_cmd = ["torchrun", f"--nproc_per_node={NUM_PROCS}", str(test_path)] - - if fp_init: - test_cmd += ["--fp8-init"] - - if len(sharding_dims) == 1: - test_cmd += ["--sharding-dims", str(sharding_dims[0])] - elif len(sharding_dims) == 2: - test_cmd += ["--sharding-dims", str(sharding_dims[0]), str(sharding_dims[1])] - else: - assert False - test_cmd += ["--recipe", recipe] - test_cmd += ["--layer-type", layer_type] - - subprocess.run(test_cmd, env=os.environ, check=True) - - @pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") -@pytest.mark.parametrize("sharding_dims", ([NUM_PROCS], [2, NUM_PROCS // 2])) -@pytest.mark.parametrize("fp8_init", (False, True)) -@pytest.mark.parametrize("layer_type", ("LayerNormLinear", "TransformerLayer")) -def test_distributed(fp8_init, sharding_dims, fp_recipe, layer_type): - - if fp_recipe in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: - pytest.xfail(f"{fp_recipe} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") - - _run_test(fp8_init, sharding_dims, fp_recipe, layer_type) - - -## ── FusedAdam + FSDP2 tests ───────────────────────────────────────── - - -def _run_fused_adam_test(test_name, recipe="delayed_scaling"): - """Launch an FSDP2 + FusedAdam test via torchrun.""" - test_path = Path(__file__).parent.resolve() / "run_fsdp2_fused_adam.py" - nproc = min(NUM_PROCS, 2) # These tests only need 2 GPUs - test_cmd = [ - "torchrun", - f"--nproc_per_node={nproc}", - str(test_path), - "--test", - test_name, - "--recipe", - recipe, - ] - - subprocess.run(test_cmd, env=os.environ, check=True) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_master_weights(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (meta device init).""" - if fp_recipe in ("NVFP4BlockScaling",): - pytest.xfail( - f"{fp_recipe}: quantized_model_init and FSDP2 is not currently supported, since the " - "block tensor is dequantized before we flatten it for FSDP2." - ) - _run_fused_adam_test("fused_adam_fp8_master_weights", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_master_weights_no_meta(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (CUDA init, no meta device). - - Block-scaling QuantizedTensors (MXFP8, Float8Blockwise, NVFP4) are wrapper - subclasses with data_ptr() == 0. Without meta-device init, FSDP2's - reset_sharded_param() crashes with 'invalid python storage'. - Per-tensor FP8 (DelayedScaling, Float8CurrentScaling) works because - Float8Tensor's storage is accessible. - """ - if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): - pytest.xfail( - f"{fp_recipe}: FSDP2 without meta-device init crashes on block-scaling " - "QuantizedTensor wrapper subclasses (data_ptr() == 0). " - "Use device='meta' + reset_parameters() after sharding." - ) - _run_fused_adam_test("fused_adam_fp8_master_weights_no_meta", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_bf16(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + bf16 params (no FP8).""" - _run_fused_adam_test("fused_adam_bf16", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_no_master(fp_recipe): - """FusedAdam(master_weights=False) + FSDP2 + FP8 params.""" - if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): - pytest.xfail( - f"{fp_recipe}: FusedAdam without master_weights does not support " - "block-scaling quantized tensors. Use master_weights=True." - ) - _run_fused_adam_test("fused_adam_fp8_no_master", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_bf16_store_param_remainders(fp_recipe): - """FusedAdam(master_weights=True, store_param_remainders=True) + FSDP2 + bf16.""" - _run_fused_adam_test("fused_adam_bf16_store_param_remainders", fp_recipe) +def test_fsdp2_model_tests(): + """All FSDP2 model tests (parametrized internally by recipe, fp8_init, sharding, layer).""" + test_path = Path(__file__).parent.resolve() / "fsdp2_tests" / "run_fsdp2_model.py" + result = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={NUM_PROCS}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "--tb=short", + ], + env=os.environ, + ) + assert result.returncode in (0, 5), ( + f"Inner pytest failed with exit code {result.returncode}" + ) @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_dcp_output_parity(fp_recipe): - """DCP save/load round-trip into a fresh model produces identical outputs.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" - ) - - if fp_recipe == "NVFP4BlockScaling": - pytest.xfail( - "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " - "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" - ) - - if fp_recipe == "Float8BlockScaling" and torch.cuda.get_device_capability()[0] == 12: - pytest.xfail( - "Float8BlockScaling is failing on SM120 with RuntimeError: " - "transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu:534 " - "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " - "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " - "requires using power of two scaling factors." - ) - - _run_fused_adam_test("dcp_output_parity", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_dcp_output_parity_async(fp_recipe): - """DCP save/load round-trip into a fresh model produces identical outputs.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access: " - "/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh:92 in function " - "multi_tensor_apply: CUDA Error: an illegal memory access was encountered" - ) - - if fp_recipe == "NVFP4BlockScaling": - pytest.xfail( - "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " - "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" - ) - - if fp_recipe == "Float8BlockScaling": - pytest.xfail( - "Float8BlockScaling: async DCP save/load round-trip produces different model " - "outputs — quantization metadata (scales) is not correctly persisted through " - "async distributed checkpointing. On SM120, additionally fails with pow2_scale " - "assertion in quantize_transpose_vector_blockwise." - ) - - _run_fused_adam_test("dcp_output_parity_async", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_safetensors_fp32_export(fp_recipe): - """Export FP32 model from optimizer master weights to safetensors.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" - ) - _run_fused_adam_test("safetensors_fp32_export", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -@pytest.mark.xfail( - reason=( - "fuse_wgrad_accumulation is incompatible with vanilla FSDP2: " - "autograd Function.apply unwraps DTensors to local tensors, so " - "main_grad (set on the DTensor) is inaccessible during backward. " - "Additionally, the fused wgrad GEMM bypasses FSDP2's reduce-scatter." - ), - raises=subprocess.CalledProcessError, - strict=True, -) -def test_fsdp2_fuse_wgrad_accumulation(fp_recipe): - """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail.""" - _run_fused_adam_test("fuse_wgrad_accumulation", fp_recipe) +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_fused_adam_tests(): + """All FSDP2 FusedAdam tests (parametrized internally by recipe, test variant).""" + test_path = Path(__file__).parent.resolve() / "fsdp2_tests" / "run_fsdp2_fused_adam.py" + nproc = min(NUM_PROCS, 2) + result = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "--tb=short", + ], + env=os.environ, + ) + assert result.returncode in (0, 5), ( + f"Inner pytest failed with exit code {result.returncode}" + ) def test_dummy() -> None: diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ab496d5a9e..47a55806ce 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -10,7 +10,7 @@ from typing import Any, Optional, Tuple, Union import torch - +from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType from transformer_engine.common.recipe import Float8BlockScaling, Recipe @@ -634,42 +634,31 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m "layout has M in dim1, which is incompatible with FSDP2 dim0 all-gather." ) - block_len = self._quantizer.block_len # 128 - - # Prepare rowwise tensors — for 2D scaling, M is in dim0 of both data and scale_inv, - # so they naturally align with FSDP2's dim0 all-gather. No unpadding needed. - rowwise_data = self._rowwise_data - rowwise_scale_inv = self._rowwise_scale_inv - - # Prepare columnwise tensors — columnwise data is transposed (K, M) and - # columnwise scale_inv is (ceil(K/128), round_up(ceil(M/128), 4)). - # M is in dim1 for both, so we must transpose to put M in dim0 for all-gather. - columnwise_data = self._columnwise_data - columnwise_scale_inv = self._columnwise_scale_inv - - if columnwise_data is not None: - # Transpose (K, shard_M) -> (shard_M, K) so M is in dim0 - columnwise_data = columnwise_data.t().contiguous() - - if columnwise_scale_inv is not None: - # Original shape: (ceil(K/128), round_up(ceil(shard_M/128), 4)) - # Strip padding from dim1 (the M-block dimension), transpose, then all-gather - shard_M = math.prod(self.shape[:-1]) - m_blocks = (shard_M + block_len - 1) // block_len # ceil(shard_M/128) - columnwise_scale_inv = columnwise_scale_inv[:, :m_blocks] # unpad dim1 - columnwise_scale_inv = columnwise_scale_inv.t().contiguous() # (m_blocks, k_blocks) - - # Always send both rowwise and columnwise data. - # Unlike MXFP8 (where both forms share the same shape), Float8Blockwise has - # differently-shaped rowwise (M, K) and columnwise (K, M) data. The GEMM kernel - # needs both forms available to perform forward and backward operations, so we - # cannot optimize by sending only one usage based on forward/backward pass. - rowwise_usage = True - sharded_tensors = (rowwise_data, rowwise_scale_inv) - columnwise_usage = self._quantizer.columnwise_usage - if columnwise_usage: - sharded_tensors += (columnwise_data, columnwise_scale_inv) + assert self._rowwise_data is not None and self._rowwise_scale_inv is not None, ( + "Rowwise data must be available for FSDP2 all-gather with 2D block scaling." + ) + fsdp_state = _get_module_fsdp_state(module) + reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward + + # If weights are resharded after forward pass, only the relevant usage + # is needed based on whether it's a forward or backward pass. + # If not resharded, the same all-gathered weights are reused in backward, + # so both usages may be needed. + if reshard_after_forward: + training_state = fsdp_state._fsdp_param_group._training_state + is_backward_pass = training_state == TrainingState.PRE_BACKWARD + rowwise_usage = not is_backward_pass + columnwise_usage = is_backward_pass + else: + rowwise_usage = True + columnwise_usage = self._quantizer.columnwise_usage + + # For 2D block scaling (128x128 blocks), columnwise data and scales are + # the transpose of rowwise data and scales. Only all-gather the rowwise + # tensors; columnwise will be derived locally via _create_columnwise() + # in post_all_gather, halving all-gather communication volume. + sharded_tensors = (self._rowwise_data, self._rowwise_scale_inv) metadata = (self._fp8_dtype, self._is_2D_scaled, rowwise_usage, columnwise_usage) return sharded_tensors, metadata @@ -694,59 +683,35 @@ def fsdp_post_all_gather( """ fp8_dtype, is_2D_scaled, rowwise_usage, columnwise_usage = metadata - # Extract rowwise tensors from all-gather outputs - rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] if rowwise_usage else (None, None) - - # Extract columnwise tensors — they were transposed in pre_all_gather, - # so we need to transpose them back. - columnwise_data, columnwise_scale_inv = ( - all_gather_outputs[-2:] if columnwise_usage else (None, None) - ) - - if columnwise_data is not None: - # All-gathered shape is (full_M, K), transpose back to (K, full_M) - columnwise_data = columnwise_data.t().contiguous() - - if columnwise_scale_inv is not None: - # All-gathered shape is (full_m_blocks, k_blocks), - # transpose back to (k_blocks, full_m_blocks) - columnwise_scale_inv = columnwise_scale_inv.t().contiguous() - # Repad dim1 (M-block dimension) to multiple of 4 for GEMM alignment - current_m_blocks = columnwise_scale_inv.shape[1] - pad_amount = (4 - current_m_blocks % 4) % 4 - if pad_amount > 0: - columnwise_scale_inv = torch.nn.functional.pad( - columnwise_scale_inv, (0, pad_amount) - ) - - # Determine the logical shape from the all-gathered data - if rowwise_data is not None: - data_shape = rowwise_data.shape - else: - # columnwise_data is (K, full_M), logical shape is (full_M, K) - data_shape = (columnwise_data.shape[1], columnwise_data.shape[0]) + # Only rowwise data+scales were all-gathered (columnwise is derived locally). + rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] + data_shape = rowwise_data.shape if out is not None: - # Update existing tensor in-place (subsequent iterations) out._rowwise_data = rowwise_data out._rowwise_scale_inv = rowwise_scale_inv - out._columnwise_data = columnwise_data - out._columnwise_scale_inv = columnwise_scale_inv else: - # Construct new tensor (first iteration). - # Float8BlockwiseQTensor constructor copies the quantizer, - # so the sharded tensor's quantizer remains independent. out = Float8BlockwiseQTensor( shape=data_shape, dtype=param_dtype, fp8_dtype=fp8_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, + columnwise_data=None, + columnwise_scale_inv=None, quantizer=self._quantizer, is_2D_scaled=is_2D_scaled, ) + + # For 2D block scaling, derive columnwise data and scales from rowwise + # via local fp8 transpose instead of all-gathering them separately. + if columnwise_usage: + out._create_columnwise() + # remove usages if not needed. + out.update_usage( + rowwise_usage=rowwise_usage, + columnwise_usage=columnwise_usage, + ) out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) return out, all_gather_outputs From a91d17bb633fa90f11587b132e1d305891842d3e Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 22 Mar 2026 22:50:46 +0000 Subject: [PATCH 09/18] verbose flag for reporting Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/test_torch_fsdp2.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index d5918adb8d..4c2619bf43 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -12,13 +12,14 @@ import transformer_engine.pytorch as te NUM_PROCS: int = torch.cuda.device_count() +_FSDP2_DIR = Path(__file__).parent.resolve() / "fsdp2_tests" @pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") def test_fsdp2_model_tests(): """All FSDP2 model tests (parametrized internally by recipe, fp8_init, sharding, layer).""" - test_path = Path(__file__).parent.resolve() / "fsdp2_tests" / "run_fsdp2_model.py" + test_path = _FSDP2_DIR / "run_fsdp2_model.py" result = subprocess.run( [ "torchrun", @@ -28,6 +29,7 @@ def test_fsdp2_model_tests(): "pytest", str(test_path), "-v", + "-s", "--tb=short", ], env=os.environ, @@ -41,7 +43,7 @@ def test_fsdp2_model_tests(): @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") def test_fsdp2_fused_adam_tests(): """All FSDP2 FusedAdam tests (parametrized internally by recipe, test variant).""" - test_path = Path(__file__).parent.resolve() / "fsdp2_tests" / "run_fsdp2_fused_adam.py" + test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" nproc = min(NUM_PROCS, 2) result = subprocess.run( [ @@ -52,6 +54,7 @@ def test_fsdp2_fused_adam_tests(): "pytest", str(test_path), "-v", + "-s", "--tb=short", ], env=os.environ, From 2e26b05c8c1994a7c7e333121f03fd23bae4af40 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 22 Mar 2026 23:11:51 +0000 Subject: [PATCH 10/18] add back coments Signed-off-by: Varun Thumbe --- .../distributed/fsdp2_tests/run_fsdp2_model.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index e894a9d0b0..fcba3aba58 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -254,6 +254,7 @@ def _run_training(args): device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") world_size = int(os.getenv("WORLD_SIZE", "1")) + # FP8 Configuration fp8_recipe = get_recipe_from_string(args.recipe) build_model_context_args = {} @@ -267,6 +268,7 @@ def _run_training(args): build_model_context_args["recipe"] = fp8_recipe dist_print(f"Memory before model init: {torch.cuda.memory_allocated(device) / 1e6} MB") + # Create the model on the meta/cuda device as per args with build_model_context(**build_model_context_args): model, inp_shape, out_shape = init_te_model(args) dist_print( @@ -274,12 +276,17 @@ def _run_training(args): f" {torch.cuda.memory_allocated(device) / 1e6} MB" ) + # Creating a DeviceMesh for fully_shard + # Setup the sharding mesh for FSDP/HSDP mesh = get_device_mesh(world_size, args.sharding_dims) custom_attrs = save_custom_attrs(model) model = shard_model_with_fsdp2(model, mesh) restore_custom_attrs(model, custom_attrs) + # model now has DTensors as its parameters if args.device == "meta": + # After FSDP2 has been applied, materialize and initialize the sharded parameters + # TE base.py's reset_parameters() handles DTensors with FP8 initialization for module in model.modules(): if hasattr(module, "reset_parameters"): module.reset_parameters() @@ -293,11 +300,13 @@ def _run_training(args): optimizer = optim.Adam(model.parameters(), lr=1e-3) for iteration in range(args.iter): + # Zero the parameter gradients optimizer.zero_grad() input_data = torch.randn(inp_shape, device=device) target = torch.randn(out_shape, device=device) + # NVFP4BlockScaling requires bfloat16 inputs in both the forward and backward passes. with ( torch.autocast(device_type="cuda", dtype=torch.bfloat16) if args.recipe == "NVFP4BlockScaling" @@ -311,6 +320,8 @@ def _run_training(args): optimizer.step() dist_print(f"Iteration {iteration} completed with loss {loss.item()}") + # Some of the FSDP states are lazy initialized during FSDP forward pass + # so testing fp8 allgather at the end of the training loop. if args.fp8_init: _check_fp8_fsdp2_allgather(model) From 5ca65c944d59eaa5ebc2e764d1f1bd5fee596aef Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 22 Mar 2026 23:06:56 +0000 Subject: [PATCH 11/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Varun Thumbe --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 15 +++++++++------ .../distributed/fsdp2_tests/run_fsdp2_model.py | 11 +++++------ tests/pytorch/distributed/test_torch_fsdp2.py | 8 ++------ .../pytorch/tensor/float8_blockwise_tensor.py | 6 +++--- 4 files changed, 19 insertions(+), 21 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 049b935878..24b31e1c69 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -40,6 +40,7 @@ import transformer_engine.pytorch as te from transformer_engine.pytorch import QuantizedTensor import transformer_engine.common.recipe + sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) from conftest import get_recipe_from_string, save_custom_attrs, restore_custom_attrs @@ -125,12 +126,14 @@ def _shard_model(model, world_size): restore_custom_attrs(model, custom_attrs) return model + def _get_dist_info(): """Get world_size and device from environment (PG already initialized by session fixture).""" world_size = int(os.environ["WORLD_SIZE"]) device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") return world_size, device + def test_fused_adam_fp8_master_weights(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). @@ -234,9 +237,7 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe_name): master_weight_dtype=torch.float32, ) - x = torch.randn( - SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device - ) + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) target = torch.randn_like(x) for step in range(NUM_STEPS): @@ -562,8 +563,11 @@ def test_dcp_output_parity(recipe_name, async_save): "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" ) - - if recipe_name == "Float8BlockScaling" and not async_save and torch.cuda.get_device_capability()[0] == 12: + if ( + recipe_name == "Float8BlockScaling" + and not async_save + and torch.cuda.get_device_capability()[0] == 12 + ): pytest.xfail( "Float8BlockScaling is failing on SM120 with RuntimeError: " "transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu:534 " @@ -729,7 +733,6 @@ def test_dcp_output_parity(recipe_name, async_save): shutil.rmtree(checkpoint_dir, ignore_errors=True) - TESTS = { "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, "fused_adam_fp8_master_weights_no_meta": test_fused_adam_fp8_master_weights_no_meta, diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index fcba3aba58..fda47bfcbf 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -293,8 +293,7 @@ def _run_training(args): dist_print(f" Sharded parameters materialized and initialized on cuda device.") dist_print( - f"FSDP2 model in cuda, memory allocated:" - f" {torch.cuda.memory_allocated(device) / 1e6} MB" + f"FSDP2 model in cuda, memory allocated: {torch.cuda.memory_allocated(device) / 1e6} MB" ) optimizer = optim.Adam(model.parameters(), lr=1e-3) @@ -341,7 +340,9 @@ def _train(args): assert dist.is_nccl_available() dist.init_process_group( - backend="nccl", rank=WORLD_RANK, world_size=WORLD_SIZE, + backend="nccl", + rank=WORLD_RANK, + world_size=WORLD_SIZE, ) try: _run_training(args) @@ -364,9 +365,7 @@ def _train(args): @pytest.mark.parametrize("layer_type", ["LayerNormLinear", "TransformerLayer"]) def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): if recipe_name in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: - pytest.xfail( - f"{recipe_name} + fp8_init: test_fp8_fsdp2_allgather is currently failing." - ) + pytest.xfail(f"{recipe_name} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") torch.manual_seed(42) torch.cuda.manual_seed(42) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 4c2619bf43..ad876c2805 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -34,9 +34,7 @@ def test_fsdp2_model_tests(): ], env=os.environ, ) - assert result.returncode in (0, 5), ( - f"Inner pytest failed with exit code {result.returncode}" - ) + assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") @@ -59,9 +57,7 @@ def test_fsdp2_fused_adam_tests(): ], env=os.environ, ) - assert result.returncode in (0, 5), ( - f"Inner pytest failed with exit code {result.returncode}" - ) + assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" def test_dummy() -> None: diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 47a55806ce..c4511b2b4d 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -634,9 +634,9 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m "layout has M in dim1, which is incompatible with FSDP2 dim0 all-gather." ) - assert self._rowwise_data is not None and self._rowwise_scale_inv is not None, ( - "Rowwise data must be available for FSDP2 all-gather with 2D block scaling." - ) + assert ( + self._rowwise_data is not None and self._rowwise_scale_inv is not None + ), "Rowwise data must be available for FSDP2 all-gather with 2D block scaling." fsdp_state = _get_module_fsdp_state(module) reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward From 9e91cafd1f6189247d51de894b4a3c5b75d0fa26 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Sun, 22 Mar 2026 23:14:25 +0000 Subject: [PATCH 12/18] another minor fix Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index fda47bfcbf..4b9a2dff1c 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -259,6 +259,7 @@ def _run_training(args): build_model_context_args = {} if not args.fp8_init: + # Build model context (FP8 init) build_model_context = nullcontext else: from transformer_engine.pytorch import fp8_model_init From 47f85133b408e9a971469e2b84e0702d62ac5dc6 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 23 Mar 2026 01:30:19 +0000 Subject: [PATCH 13/18] not needed for this PR Signed-off-by: Varun Thumbe --- .../pytorch/tensor/float8_blockwise_tensor.py | 117 ++++++++++++------ 1 file changed, 76 insertions(+), 41 deletions(-) diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index c4511b2b4d..ab496d5a9e 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -10,7 +10,7 @@ from typing import Any, Optional, Tuple, Union import torch -from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState + import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType from transformer_engine.common.recipe import Float8BlockScaling, Recipe @@ -634,31 +634,42 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m "layout has M in dim1, which is incompatible with FSDP2 dim0 all-gather." ) - assert ( - self._rowwise_data is not None and self._rowwise_scale_inv is not None - ), "Rowwise data must be available for FSDP2 all-gather with 2D block scaling." - - fsdp_state = _get_module_fsdp_state(module) - reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward - - # If weights are resharded after forward pass, only the relevant usage - # is needed based on whether it's a forward or backward pass. - # If not resharded, the same all-gathered weights are reused in backward, - # so both usages may be needed. - if reshard_after_forward: - training_state = fsdp_state._fsdp_param_group._training_state - is_backward_pass = training_state == TrainingState.PRE_BACKWARD - rowwise_usage = not is_backward_pass - columnwise_usage = is_backward_pass - else: - rowwise_usage = True - columnwise_usage = self._quantizer.columnwise_usage - - # For 2D block scaling (128x128 blocks), columnwise data and scales are - # the transpose of rowwise data and scales. Only all-gather the rowwise - # tensors; columnwise will be derived locally via _create_columnwise() - # in post_all_gather, halving all-gather communication volume. - sharded_tensors = (self._rowwise_data, self._rowwise_scale_inv) + block_len = self._quantizer.block_len # 128 + + # Prepare rowwise tensors — for 2D scaling, M is in dim0 of both data and scale_inv, + # so they naturally align with FSDP2's dim0 all-gather. No unpadding needed. + rowwise_data = self._rowwise_data + rowwise_scale_inv = self._rowwise_scale_inv + + # Prepare columnwise tensors — columnwise data is transposed (K, M) and + # columnwise scale_inv is (ceil(K/128), round_up(ceil(M/128), 4)). + # M is in dim1 for both, so we must transpose to put M in dim0 for all-gather. + columnwise_data = self._columnwise_data + columnwise_scale_inv = self._columnwise_scale_inv + + if columnwise_data is not None: + # Transpose (K, shard_M) -> (shard_M, K) so M is in dim0 + columnwise_data = columnwise_data.t().contiguous() + + if columnwise_scale_inv is not None: + # Original shape: (ceil(K/128), round_up(ceil(shard_M/128), 4)) + # Strip padding from dim1 (the M-block dimension), transpose, then all-gather + shard_M = math.prod(self.shape[:-1]) + m_blocks = (shard_M + block_len - 1) // block_len # ceil(shard_M/128) + columnwise_scale_inv = columnwise_scale_inv[:, :m_blocks] # unpad dim1 + columnwise_scale_inv = columnwise_scale_inv.t().contiguous() # (m_blocks, k_blocks) + + # Always send both rowwise and columnwise data. + # Unlike MXFP8 (where both forms share the same shape), Float8Blockwise has + # differently-shaped rowwise (M, K) and columnwise (K, M) data. The GEMM kernel + # needs both forms available to perform forward and backward operations, so we + # cannot optimize by sending only one usage based on forward/backward pass. + rowwise_usage = True + sharded_tensors = (rowwise_data, rowwise_scale_inv) + columnwise_usage = self._quantizer.columnwise_usage + if columnwise_usage: + sharded_tensors += (columnwise_data, columnwise_scale_inv) + metadata = (self._fp8_dtype, self._is_2D_scaled, rowwise_usage, columnwise_usage) return sharded_tensors, metadata @@ -683,35 +694,59 @@ def fsdp_post_all_gather( """ fp8_dtype, is_2D_scaled, rowwise_usage, columnwise_usage = metadata - # Only rowwise data+scales were all-gathered (columnwise is derived locally). - rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] - data_shape = rowwise_data.shape + # Extract rowwise tensors from all-gather outputs + rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] if rowwise_usage else (None, None) + + # Extract columnwise tensors — they were transposed in pre_all_gather, + # so we need to transpose them back. + columnwise_data, columnwise_scale_inv = ( + all_gather_outputs[-2:] if columnwise_usage else (None, None) + ) + + if columnwise_data is not None: + # All-gathered shape is (full_M, K), transpose back to (K, full_M) + columnwise_data = columnwise_data.t().contiguous() + + if columnwise_scale_inv is not None: + # All-gathered shape is (full_m_blocks, k_blocks), + # transpose back to (k_blocks, full_m_blocks) + columnwise_scale_inv = columnwise_scale_inv.t().contiguous() + # Repad dim1 (M-block dimension) to multiple of 4 for GEMM alignment + current_m_blocks = columnwise_scale_inv.shape[1] + pad_amount = (4 - current_m_blocks % 4) % 4 + if pad_amount > 0: + columnwise_scale_inv = torch.nn.functional.pad( + columnwise_scale_inv, (0, pad_amount) + ) + + # Determine the logical shape from the all-gathered data + if rowwise_data is not None: + data_shape = rowwise_data.shape + else: + # columnwise_data is (K, full_M), logical shape is (full_M, K) + data_shape = (columnwise_data.shape[1], columnwise_data.shape[0]) if out is not None: + # Update existing tensor in-place (subsequent iterations) out._rowwise_data = rowwise_data out._rowwise_scale_inv = rowwise_scale_inv + out._columnwise_data = columnwise_data + out._columnwise_scale_inv = columnwise_scale_inv else: + # Construct new tensor (first iteration). + # Float8BlockwiseQTensor constructor copies the quantizer, + # so the sharded tensor's quantizer remains independent. out = Float8BlockwiseQTensor( shape=data_shape, dtype=param_dtype, fp8_dtype=fp8_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, - columnwise_data=None, - columnwise_scale_inv=None, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, quantizer=self._quantizer, is_2D_scaled=is_2D_scaled, ) - - # For 2D block scaling, derive columnwise data and scales from rowwise - # via local fp8 transpose instead of all-gathering them separately. - if columnwise_usage: - out._create_columnwise() - # remove usages if not needed. - out.update_usage( - rowwise_usage=rowwise_usage, - columnwise_usage=columnwise_usage, - ) out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) return out, all_gather_outputs From 7d9785f542ee814cc5d9f0b2464762d82863670b Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 23 Mar 2026 05:26:26 +0000 Subject: [PATCH 14/18] address review comments Signed-off-by: Varun Thumbe --- .../distributed/fsdp2_tests/conftest.py | 38 +- .../distributed/fsdp2_tests/fsdp2_utils.py | 38 ++ .../fsdp2_tests/run_fsdp2_fused_adam.py | 331 +++++++++--------- .../fsdp2_tests/run_fsdp2_model.py | 4 +- 4 files changed, 217 insertions(+), 194 deletions(-) create mode 100644 tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 4d75f91317..055ca82402 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -2,24 +2,20 @@ # # See LICENSE for license information. -"""Shared pytest fixtures and utilities for FSDP2 distributed tests. +"""Shared pytest fixtures for FSDP2 distributed tests. Fixtures defined here (dist_init, _cleanup, recipe_name) are auto-discovered -by pytest for every test module in this directory. Utility functions -(get_recipe_from_string, save_custom_attrs, restore_custom_attrs) can be -imported normally: ``from conftest import get_recipe_from_string``. +by pytest for every test module in this directory. Utility functions live in +``fsdp2_utils.py`` so that test/runner scripts can import them without +triggering a duplicate import of this conftest module. """ import gc import os - import pytest - import torch import torch.distributed as dist - -from transformer_engine.pytorch import fp8, QuantizedTensor -import transformer_engine.common.recipe +from transformer_engine.pytorch import fp8 # ── FP8 recipe parametrization ────────────────────────────────────── @@ -74,6 +70,8 @@ def dist_init(): def _cleanup(): """Release GPU memory and stale NCCL state between tests.""" yield + if dist.is_initialized(): + dist.barrier() gc.collect() torch.cuda.empty_cache() @@ -83,25 +81,3 @@ def recipe_name(request): return request.param -# ── Other Shared helpers ─────────────────────────────────────────────────── -def get_recipe_from_string(recipe): - return getattr(transformer_engine.common.recipe, recipe)() - - -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - if isinstance(param, QuantizedTensor): - ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] - else: - ignore_keys = [] - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py new file mode 100644 index 0000000000..a031aac381 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -0,0 +1,38 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared utility functions for FSDP2 distributed tests. + +This module holds helpers that are used by both conftest.py (for fixture +parametrization) and the individual test/runner scripts. Keeping them +here avoids the double-import problem that arises when test modules +``sys.path.insert`` + ``from conftest import …`` alongside pytest's own +conftest auto-discovery. +""" + +import transformer_engine.common.recipe +from transformer_engine.pytorch import QuantizedTensor + + +def get_recipe_from_string(recipe): + return getattr(transformer_engine.common.recipe, recipe)() + + +def save_custom_attrs(module): + custom_attrs = {} + for name, param in module.named_parameters(): + if isinstance(param, QuantizedTensor): + ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] + else: + ignore_keys = [] + attrs = vars(param) + custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} + return custom_attrs + + +def restore_custom_attrs(module, custom_attrs): + for name, param in module.named_parameters(): + if name in custom_attrs: + for attr_name, attr_value in custom_attrs[name].items(): + setattr(param, attr_name, attr_value) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 24b31e1c69..5a6ab38582 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -26,8 +26,7 @@ import argparse import functools import os -import pathlib -import sys +import shutil import pytest import torch @@ -41,8 +40,7 @@ from transformer_engine.pytorch import QuantizedTensor import transformer_engine.common.recipe -sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) -from conftest import get_recipe_from_string, save_custom_attrs, restore_custom_attrs +from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs HIDDEN_SIZE = 256 @@ -506,33 +504,35 @@ def test_safetensors_fp32_export(recipe_name): full_opt_state = get_optimizer_state_dict(model, optimizer, options=full_opts) rank = int(os.environ.get("RANK", "0")) - save_path = "/tmp/te_test_fsdp2_model_fp32.safetensors" + save_path = f"/tmp/te_test_fsdp2_model_fp32_{recipe_name}.safetensors" if rank == 0: - # Build FP32 state dict from optimizer master weights. - fp32_state = {} - opt_param_states = full_opt_state.get("state", {}) + if os.path.exists(save_path): + os.remove(save_path) - for key, value in full_model_state.items(): - if key in opt_param_states and "master_param" in opt_param_states[key]: - fp32_state[key] = opt_param_states[key]["master_param"].float() - else: - fp32_state[key] = value.float() + try: + fp32_state = {} + opt_param_states = full_opt_state.get("state", {}) - assert len(fp32_state) > 0, "FP32 state dict is empty" + for key, value in full_model_state.items(): + if key in opt_param_states and "master_param" in opt_param_states[key]: + fp32_state[key] = opt_param_states[key]["master_param"].float() + else: + fp32_state[key] = value.float() - # Save and verify. - save_file(fp32_state, save_path) - loaded = load_file(save_path) + assert len(fp32_state) > 0, "FP32 state dict is empty" - assert len(loaded) == len( - fp32_state - ), f"Loaded {len(loaded)} tensors, expected {len(fp32_state)}" - for k, v in loaded.items(): - assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + save_file(fp32_state, save_path) + loaded = load_file(save_path) - # Clean up. - os.remove(save_path) + assert len(loaded) == len( + fp32_state + ), f"Loaded {len(loaded)} tensors, expected {len(fp32_state)}" + for k, v in loaded.items(): + assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + finally: + if os.path.exists(save_path): + os.remove(save_path) @pytest.mark.parametrize("async_save", [False, True], ids=["sync", "async"]) @@ -586,151 +586,162 @@ def test_dcp_output_parity(recipe_name, async_save): import torch.distributed.checkpoint as dcp world_size, device = _get_dist_info() + rank = int(os.environ.get("RANK", "0")) + save_mode = "async" if async_save else "sync" + checkpoint_dir = f"/tmp/te_test_fsdp2_dcp_parity_{recipe_name}_{save_mode}" - # ── Build and train the original model ─────────────────────────── - model = _build_model(fp8_init=True, recipe=recipe) - model = _shard_model(model, world_size) + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + dist.barrier() - optimizer = te.optimizers.FusedAdam( - model.parameters(), - lr=1e-3, - master_weights=True, - master_weight_dtype=torch.float32, - ) + try: + # ── Build and train the original model ─────────────────────────── + model = _build_model(fp8_init=True, recipe=recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) - x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) - target = torch.randn_like(x) + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Record reference output from the trained model. + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + ref_output = model(x).clone() + + # ── Save checkpoint ────────────────────────────────────────────── + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + # We need to remove the _extra_state keys from the model state dict for + # DelayedScaling, since otherwise we'll run into an error that the tensor + # sizes are different. The alternative is a LoadPlanner that dynamically + # re-sizes the input tensors, see NVIDIA/TransformerEngine#1860 for more + # details. + model_state = { + k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") + } + else: + model_state = model.state_dict() + + save_state = {"model": model_state, "optimizer": optimizer.state_dict()} + + if not async_save: + dcp.save(save_state, checkpoint_id=checkpoint_dir) + else: + future = dcp.async_save(save_state, checkpoint_id=checkpoint_dir) + future.result() + + # ── Build a fresh model and load the checkpoint ────────────────── + model2 = _build_model(fp8_init=True, recipe=recipe) + model2 = _shard_model(model2, world_size) + + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) - for _ in range(NUM_STEPS): + # Populate optimizer state so load_state_dict has matching structure. + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + model2_state = { + k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") + } + else: + model2_state = model2.state_dict() + + state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} + + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict( + state_to_load["model"], + strict=( + False + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) + else True + ), + ) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + # ── Verify identical forward-pass output ───────────────────────── + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + loaded_output = model2(x) + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + # DelayedScaling stores amax history and scaling factors in _extra_state, + # which cannot be saved via DCP due to non-deterministic pickle sizes + # across ranks. The fresh model therefore uses default scaling factors, + # producing small numerical differences from FP8 re-quantization. + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0.05, + atol=0.1, + msg=lambda x: ( + f"Fresh model loaded from DCP checkpoint produces different output: {x}" + ), + ) + else: + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0, + atol=0, + msg=lambda x: ( + f"Fresh model loaded from DCP checkpoint produces different output: {x}" + ), + ) + + # ── Verify one more training step produces identical results ───── optimizer.zero_grad(set_to_none=True) with te.autocast(enabled=True, recipe=recipe): - output = model(x) - loss = F.mse_loss(output, target) - loss.backward() + out1 = model(x) + loss1 = F.mse_loss(out1, target) + loss1.backward() optimizer.step() - # Record reference output from the trained model. - with torch.no_grad(): - with te.autocast(enabled=True, recipe=recipe): - ref_output = model(x).clone() - - # ── Save checkpoint ────────────────────────────────────────────── - checkpoint_dir = "/tmp/te_test_fsdp2_dcp_parity" - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - # We need to remove the _extra_state keys from the model state dict for DelayedScaling, - # since otherwise we'll run into an error that the tensor sizes are different. The - # alternative is a LoadPlanner that dynamically re-sizes the input tensors, see - # NVIDIA/TransformerEngine#1860 for more details. - model_state = { - k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") - } - else: - model_state = model.state_dict() - - save_state = {"model": model_state, "optimizer": optimizer.state_dict()} - - if not async_save: - dcp.save(save_state, checkpoint_id=checkpoint_dir) - else: - future = dcp.async_save(save_state, checkpoint_id=checkpoint_dir) - future.result() # Block on async save completion - - # ── Build a fresh model and load the checkpoint ────────────────── - model2 = _build_model(fp8_init=True, recipe=recipe) - model2 = _shard_model(model2, world_size) - - optimizer2 = te.optimizers.FusedAdam( - model2.parameters(), - lr=1e-3, - master_weights=True, - master_weight_dtype=torch.float32, - ) - - # Populate optimizer state so load_state_dict has matching structure. - optimizer2.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out_tmp = model2(x) - F.mse_loss(out_tmp, target).backward() - optimizer2.step() - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - model2_state = { - k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") - } - else: - model2_state = model2.state_dict() - - state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} - - dcp.load(state_to_load, checkpoint_id=checkpoint_dir) - model2.load_state_dict( - state_to_load["model"], - strict=( - False if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) else True - ), - ) - optimizer2.load_state_dict(state_to_load["optimizer"]) - - # ── Verify identical forward-pass output ───────────────────────── - with torch.no_grad(): + optimizer2.zero_grad(set_to_none=True) with te.autocast(enabled=True, recipe=recipe): - loaded_output = model2(x) - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - # DelayedScaling stores amax history and scaling factors in _extra_state, - # which cannot be saved via DCP due to non-deterministic pickle sizes - # across ranks. The fresh model therefore uses default scaling factors, - # producing small numerical differences from FP8 re-quantization. - torch.testing.assert_close( - loaded_output, - ref_output, - rtol=0.05, - atol=0.1, - msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", - ) - else: - torch.testing.assert_close( - loaded_output, - ref_output, - rtol=0, - atol=0, - msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", - ) - - # ── Verify one more training step produces identical results ───── - optimizer.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out1 = model(x) - loss1 = F.mse_loss(out1, target) - loss1.backward() - optimizer.step() - - optimizer2.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out2 = model2(x) - loss2 = F.mse_loss(out2, target) - loss2.backward() - optimizer2.step() - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - torch.testing.assert_close( - out2, - out1, - rtol=0.05, - atol=0.1, - msg="Training step after DCP load produces different output", - ) - else: - torch.testing.assert_close( - out2, out1, msg="Training step after DCP load produces different output" - ) - - # ── Cleanup ────────────────────────────────────────────────────── - import shutil - - if int(os.environ.get("RANK", "0")) == 0: - shutil.rmtree(checkpoint_dir, ignore_errors=True) + out2 = model2(x) + loss2 = F.mse_loss(out2, target) + loss2.backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + torch.testing.assert_close( + out2, + out1, + rtol=0.05, + atol=0.1, + msg="Training step after DCP load produces different output", + ) + else: + torch.testing.assert_close( + out2, out1, msg="Training step after DCP load produces different output" + ) + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) TESTS = { diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 4b9a2dff1c..fce565ed9a 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -28,7 +28,6 @@ import gc import os -import pathlib import sys import argparse from types import SimpleNamespace @@ -48,8 +47,7 @@ from torch.distributed._composable.fsdp import fully_shard from torch.distributed.device_mesh import init_device_mesh -sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent)) -from conftest import get_recipe_from_string, save_custom_attrs, restore_custom_attrs +from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs def dist_print(msg): From 3ac0ccd639ec341024e2d6c9f0f1a1b8daaf6bc1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 23 Mar 2026 05:31:11 +0000 Subject: [PATCH 15/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/fsdp2_tests/conftest.py | 2 -- .../distributed/fsdp2_tests/run_fsdp2_fused_adam.py | 8 ++------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 055ca82402..bf63618e20 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -79,5 +79,3 @@ def _cleanup(): @pytest.fixture(params=_parametrize_recipes()) def recipe_name(request): return request.param - - diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 5a6ab38582..877fa66795 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -696,9 +696,7 @@ def test_dcp_output_parity(recipe_name, async_save): ref_output, rtol=0.05, atol=0.1, - msg=lambda x: ( - f"Fresh model loaded from DCP checkpoint produces different output: {x}" - ), + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", ) else: torch.testing.assert_close( @@ -706,9 +704,7 @@ def test_dcp_output_parity(recipe_name, async_save): ref_output, rtol=0, atol=0, - msg=lambda x: ( - f"Fresh model loaded from DCP checkpoint produces different output: {x}" - ), + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", ) # ── Verify one more training step produces identical results ───── From 075b6aa7b7ee31df5c401827ee2a678a2c567ec7 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 23 Mar 2026 05:31:57 +0000 Subject: [PATCH 16/18] unecessary comments Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/fsdp2_tests/conftest.py | 4 +--- tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py | 6 ------ 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index 055ca82402..19321a1134 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -5,9 +5,7 @@ """Shared pytest fixtures for FSDP2 distributed tests. Fixtures defined here (dist_init, _cleanup, recipe_name) are auto-discovered -by pytest for every test module in this directory. Utility functions live in -``fsdp2_utils.py`` so that test/runner scripts can import them without -triggering a duplicate import of this conftest module. +by pytest for every test module in this directory. """ import gc diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index a031aac381..1b4d787b02 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -3,12 +3,6 @@ # See LICENSE for license information. """Shared utility functions for FSDP2 distributed tests. - -This module holds helpers that are used by both conftest.py (for fixture -parametrization) and the individual test/runner scripts. Keeping them -here avoids the double-import problem that arises when test modules -``sys.path.insert`` + ``from conftest import …`` alongside pytest's own -conftest auto-discovery. """ import transformer_engine.common.recipe From 14c1c486277638f2b47e4553f11e1d5c1948071e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 23 Mar 2026 05:33:48 +0000 Subject: [PATCH 17/18] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py index 1b4d787b02..178ce62375 100644 --- a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -2,8 +2,7 @@ # # See LICENSE for license information. -"""Shared utility functions for FSDP2 distributed tests. -""" +"""Shared utility functions for FSDP2 distributed tests.""" import transformer_engine.common.recipe from transformer_engine.pytorch import QuantizedTensor From 947bb395c5be5abe38255cc27156d8922931c958 Mon Sep 17 00:00:00 2001 From: Varun Thumbe Date: Mon, 23 Mar 2026 06:04:18 +0000 Subject: [PATCH 18/18] address revire comments Signed-off-by: Varun Thumbe --- tests/pytorch/distributed/fsdp2_tests/conftest.py | 6 ++++++ tests/pytorch/distributed/test_torch_fsdp2.py | 2 ++ 2 files changed, 8 insertions(+) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py index a93851e67e..bf9db094d2 100644 --- a/tests/pytorch/distributed/fsdp2_tests/conftest.py +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -15,6 +15,12 @@ import torch.distributed as dist from transformer_engine.pytorch import fp8 +# Ensure the correct CUDA device is active before _parametrize_recipes() +# runs at collection time, since the session-scoped dist_init fixture +# has not executed yet. +_local_rank = int(os.environ.get("LOCAL_RANK", "0")) +torch.cuda.set_device(_local_rank) + # ── FP8 recipe parametrization ────────────────────────────────────── def _check_nvfp4_support(): diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index ad876c2805..aca8d6d692 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -33,6 +33,7 @@ def test_fsdp2_model_tests(): "--tb=short", ], env=os.environ, + timeout=600, ) assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" @@ -56,6 +57,7 @@ def test_fsdp2_fused_adam_tests(): "--tb=short", ], env=os.environ, + timeout=600, ) assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}"