diff --git a/kt-kernel/CMakeLists.txt b/kt-kernel/CMakeLists.txt index 1e7558b8f..f87c03df7 100644 --- a/kt-kernel/CMakeLists.txt +++ b/kt-kernel/CMakeLists.txt @@ -395,6 +395,7 @@ if(HOST_IS_X86) # 获取不带扩展名的文件名作为 target 名 get_filename_component(test_name ${test_src} NAME_WE) add_executable(${test_name} ${test_src} ${CMAKE_CURRENT_SOURCE_DIR}/cpu_backend/shared_mem_buffer.cpp) + target_compile_options(${test_name} PRIVATE ${ARCH_FLAGS}) target_link_libraries(${test_name} llama OpenMP::OpenMP_CXX numa) endforeach() endif() diff --git a/kt-kernel/ext_bindings.cpp b/kt-kernel/ext_bindings.cpp index 87074e417..9ce751752 100644 --- a/kt-kernel/ext_bindings.cpp +++ b/kt-kernel/ext_bindings.cpp @@ -325,22 +325,29 @@ class MOESFTBindings { intptr_t grad_down_lora_a; intptr_t grad_down_lora_b; intptr_t grad_weights; + intptr_t grad_gate_proj; + intptr_t grad_up_proj; + intptr_t grad_down_proj; }; static void inner(void* args) { Args* args_ = (Args*)args; args_->cpuinfer->enqueue(&TP_MOE_SFT::backward_binding, args_->moe, args_->grad_output, args_->grad_input, args_->grad_gate_lora_a, args_->grad_gate_lora_b, args_->grad_up_lora_a, args_->grad_up_lora_b, args_->grad_down_lora_a, args_->grad_down_lora_b, - args_->grad_weights); + args_->grad_weights, args_->grad_gate_proj, args_->grad_up_proj, + args_->grad_down_proj); } static std::pair cpuinfer_interface(std::shared_ptr> moe, intptr_t grad_output, intptr_t grad_input, intptr_t grad_gate_lora_a, intptr_t grad_gate_lora_b, intptr_t grad_up_lora_a, intptr_t grad_up_lora_b, intptr_t grad_down_lora_a, - intptr_t grad_down_lora_b, intptr_t grad_weights) { + intptr_t grad_down_lora_b, intptr_t grad_weights, + intptr_t grad_gate_proj, intptr_t grad_up_proj, + intptr_t grad_down_proj) { Args* args = new Args{nullptr, moe.get(), grad_output, grad_input, grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b, - grad_down_lora_a, grad_down_lora_b, grad_weights}; + grad_down_lora_a, grad_down_lora_b, grad_weights, + grad_gate_proj, grad_up_proj, grad_down_proj}; return std::make_pair((intptr_t)&inner, (intptr_t)args); } }; @@ -403,7 +410,15 @@ void bind_moe_sft_module(py::module_& moe_module, const char* name) { self.prepare_and_save_bwd((void*)gate, (void*)up, (void*)down, path); }) .def("submit_backward_repack", &MoeClass::submit_backward_repack) - .def("wait_backward_repack", &MoeClass::wait_backward_repack); + .def("wait_backward_repack", &MoeClass::wait_backward_repack) + .def("get_profile_stats", &MoeClass::get_profile_stats, py::arg("reset") = false) + .def("reset_profile_stats", &MoeClass::reset_profile_stats) + // Update base weight BF16 pointers for reload_base_weights (full mode training) + // After calling this, call load_weights_task() to re-quantize BF16->AMX + .def("set_base_weight_pointers", + [](MoeClass& self, intptr_t gate, intptr_t up, intptr_t down) { + self.set_base_weight_pointers((void*)gate, (void*)up, (void*)down); + }); } #endif // defined(__x86_64__) && defined(USE_AMX_AVX_KERNEL) @@ -779,7 +794,11 @@ PYBIND11_MODULE(kt_kernel_ext, m) { .DEF_PTR_PROPERTY(MOESFTConfig, up_lora_a) .DEF_PTR_PROPERTY(MOESFTConfig, up_lora_b) .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_a) - .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_b); + .DEF_PTR_PROPERTY(MOESFTConfig, down_lora_b) + .def_readwrite("full_weight_grad", &MOESFTConfig::full_weight_grad) + .DEF_PTR_PROPERTY(MOESFTConfig, grad_gate_proj) + .DEF_PTR_PROPERTY(MOESFTConfig, grad_up_proj) + .DEF_PTR_PROPERTY(MOESFTConfig, grad_down_proj); py::class_>(moe_module, "MoE_Interface"); @@ -800,7 +819,7 @@ PYBIND11_MODULE(kt_kernel_ext, m) { #endif #if defined(__AVX512BF16__) // SFT MoE with LoRA support (BF16, INT8, INT4, AWQ, K2) - bind_moe_sft_module>(moe_module, "AMXBF16_SFT_MOE"); + bind_moe_sft_module>(moe_module, "AMXBF16_SFT_MOE"); bind_moe_sft_module>(moe_module, "AMXInt8_SFT_MOE"); bind_moe_sft_module>(moe_module, "AMXInt4_SFT_MOE"); // bind_moe_sft_module>(moe_module, "AMXInt4_1_SFT_MOE"); @@ -809,7 +828,8 @@ PYBIND11_MODULE(kt_kernel_ext, m) { // bind_moe_sft_module>(moe_module, // "AMXInt4_KGroup_SFT_MOE"); // SFT MoE with SkipLoRA=true (skip all LoRA computation in backward, only compute base weight grad_input) - bind_moe_sft_module>(moe_module, "AMXBF16_SFT_MOE_SkipLoRA"); + bind_moe_sft_module>(moe_module, + "AMXBF16_SFT_MOE_SkipLoRA"); bind_moe_sft_module>(moe_module, "AMXInt8_SFT_MOE_SkipLoRA"); bind_moe_sft_module>(moe_module, "AMXInt4_SFT_MOE_SkipLoRA"); // bind_moe_sft_module>(moe_module, diff --git a/kt-kernel/operators/amx/bf16-moe.hpp b/kt-kernel/operators/amx/bf16-moe.hpp index 389446ed6..375aaef9d 100644 --- a/kt-kernel/operators/amx/bf16-moe.hpp +++ b/kt-kernel/operators/amx/bf16-moe.hpp @@ -30,6 +30,8 @@ template class AMX_BF16_MOE_TP : public AMX_MOE_BASE> { using Base = AMX_MOE_BASE>; + + protected: using Base::config_; using Base::down_ba_; using Base::down_bb_; diff --git a/kt-kernel/operators/amx/la/amx_raw_buffers.hpp b/kt-kernel/operators/amx/la/amx_raw_buffers.hpp index 86fee938a..fc3b67194 100644 --- a/kt-kernel/operators/amx/la/amx_raw_buffers.hpp +++ b/kt-kernel/operators/amx/la/amx_raw_buffers.hpp @@ -133,16 +133,13 @@ struct BufferBBF16Impl { } void set_data(void* new_ptr) { b = reinterpret_cast(new_ptr); } - void from_mat(ggml_bf16_t* src, int ith, int nth) { - auto [n_start, n_end] = K::split_range_n(n, ith, nth); - int n_block_begin = n_start; - int n_block_size = n_end - n_block_begin; + void pack_block(ggml_bf16_t* src, int src_stride, int n_block_begin, int n_block_size) { for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { int k_block_size = std::min(K_BLOCK, k - k_block_begin); for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { for (int i = 0; i < N_STEP; i++) { - __m512i* s = (__m512i*)(src + (n_block_begin + n_begin + i) * k + k_block_begin + k_begin); + __m512i* s = (__m512i*)(src + (n_begin + i) * src_stride + k_block_begin + k_begin); __m512i* d = (__m512i*)(b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP + i * K_STEP); avx512_copy_32xbf16(s, d); @@ -155,6 +152,116 @@ struct BufferBBF16Impl { } } } + + void from_mat(ggml_bf16_t* src, int ith, int nth) { + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + pack_block(src + n_block_begin * k, k, n_block_begin, n_block_size); + } + + void from_mat_strided(ggml_bf16_t* src, int src_stride, int ith, int nth) { + assert(src_stride >= k); + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + if (n_block_size <= 0) return; + pack_block(src + (size_t)n_block_begin * src_stride, src_stride, n_block_begin, n_block_size); + } + + void from_mat_transposed(ggml_bf16_t* src, int src_n, int src_k, int ith, int nth) { + assert(n == src_k && k == src_n); + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + if (n_block_size <= 0) return; + + thread_local std::vector strip; + strip.resize((size_t)n_block_size * k); + constexpr int TILE = 32; + for (int c_tile = 0; c_tile < k; c_tile += TILE) { + int c_end = std::min(c_tile + TILE, k); + for (int r_tile = 0; r_tile < n_block_size; r_tile += TILE) { + int r_end = std::min(r_tile + TILE, n_block_size); + for (int c = c_tile; c < c_end; c++) { + for (int r = r_tile; r < r_end; r++) { + strip[(size_t)r * k + c] = src[(size_t)c * src_k + n_block_begin + r]; + } + } + } + } + pack_block(strip.data(), k, n_block_begin, n_block_size); + } + + void to_mat(ggml_bf16_t* dst, int ith, int nth) const { + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + int n_block_begin = n_start; + int n_block_size = n_end - n_block_begin; + if (n_block_size <= 0) return; + + alignas(64) ggml_bf16_t tile_copy[N_STEP * K_STEP]; + for (int n_begin = 0; n_begin < n_block_size; n_begin += N_STEP) { + for (int k_block_begin = 0; k_block_begin < k; k_block_begin += K_BLOCK) { + int k_block_size = std::min(K_BLOCK, k - k_block_begin); + for (int k_begin = 0; k_begin < k_block_size; k_begin += K_STEP) { + const ggml_bf16_t* tile_src = + b + n_block_begin * k + k_block_begin * n_block_size + n_begin * k_block_size + k_begin * N_STEP; + memcpy(tile_copy, tile_src, sizeof(tile_copy)); + transpose_16x16_32bit((__m512i*)tile_copy); + transpose_16x16_32bit((__m512i*)(tile_copy + TILE_N * K_STEP)); + for (int i = 0; i < N_STEP; i++) { + __m512i* s = (__m512i*)(tile_copy + i * K_STEP); + __m512i* d = (__m512i*)(dst + (size_t)(n_block_begin + n_begin + i) * k + k_block_begin + k_begin); + avx512_copy_32xbf16(s, d); + } + } + } + } + } + + void from_bb_transposed(const BufferBBF16Impl& src, int ith, int nth) { + assert(n == src.k && k == src.n); + auto [n_start, n_end] = K::split_range_n(n, ith, nth); + int dst_n_block_begin = n_start; + int dst_n_block_size = n_end - dst_n_block_begin; + if (dst_n_block_size <= 0) return; + + auto tile_ptr = [](ggml_bf16_t* base, int total_n, int total_k, int abs_n, int abs_k) { + int n_block_begin = abs_n / N_BLOCK * N_BLOCK; + int n_within = abs_n - n_block_begin; + int n_block_size = std::min(N_BLOCK, total_n - n_block_begin); + int k_block_begin = abs_k / K_BLOCK * K_BLOCK; + int k_within = abs_k - k_block_begin; + return base + (size_t)n_block_begin * total_k + (size_t)k_block_begin * n_block_size + + (size_t)n_within * std::min(K_BLOCK, total_k - k_block_begin) + (size_t)k_within * N_STEP; + }; + + alignas(64) ggml_bf16_t src_tile[N_STEP * K_STEP]; + alignas(64) ggml_bf16_t dst_tile[N_STEP * K_STEP]; + for (int dst_n = 0; dst_n < dst_n_block_size; dst_n += N_STEP) { + for (int dst_k_block = 0; dst_k_block < k; dst_k_block += K_BLOCK) { + int dst_k_block_size = std::min(K_BLOCK, k - dst_k_block); + for (int dst_k = 0; dst_k < dst_k_block_size; dst_k += K_STEP) { + int abs_dst_n = dst_n_block_begin + dst_n; + int abs_dst_k = dst_k_block + dst_k; + ggml_bf16_t* src_ptr = tile_ptr(src.b, src.n, src.k, abs_dst_k, abs_dst_n); + memcpy(src_tile, src_ptr, sizeof(src_tile)); + transpose_16x16_32bit((__m512i*)src_tile); + transpose_16x16_32bit((__m512i*)(src_tile + TILE_N * K_STEP)); + + for (int i = 0; i < N_STEP; i++) { + for (int j = 0; j < K_STEP; j++) { + dst_tile[j * K_STEP + i] = src_tile[i * K_STEP + j]; + } + } + transpose_16x16_32bit((__m512i*)dst_tile); + transpose_16x16_32bit((__m512i*)(dst_tile + TILE_N * K_STEP)); + ggml_bf16_t* dst_ptr = tile_ptr(b, n, k, abs_dst_n, abs_dst_k); + memcpy(dst_ptr, dst_tile, sizeof(dst_tile)); + } + } + } + } ggml_bf16_t* get_submat(int n, int k, int n_begin, int k_begin) { int n_block_begin = n_begin / N_BLOCK * N_BLOCK; n_begin -= n_block_begin; diff --git a/kt-kernel/operators/amx/la/bf16_dweight.hpp b/kt-kernel/operators/amx/la/bf16_dweight.hpp new file mode 100644 index 000000000..190917c59 --- /dev/null +++ b/kt-kernel/operators/amx/la/bf16_dweight.hpp @@ -0,0 +1,223 @@ +#ifndef AMX_BF16_DWEIGHT_HPP +#define AMX_BF16_DWEIGHT_HPP + +#include +#include +#include +#include +#include + +#include "amx_kernels.hpp" +#include "amx_raw_kernels.hpp" + +namespace amx { + +struct BF16DWeightTimings { + uint64_t pack_a_ns = 0; + uint64_t pack_a_calls = 0; + uint64_t pack_b_ns = 0; + uint64_t pack_b_calls = 0; + uint64_t panel_input_ns = 0; + uint64_t panel_input_calls = 0; + uint64_t panel_grad_output_ns = 0; + uint64_t panel_grad_output_calls = 0; + uint64_t kernel_gate_up_ns = 0; + uint64_t kernel_gate_up_calls = 0; + uint64_t kernel_down_ns = 0; + uint64_t kernel_down_calls = 0; + uint64_t store_ns = 0; + uint64_t store_calls = 0; + + void reset() { *this = {}; } +}; + +inline BF16DWeightTimings& bf16_dweight_timings() { + static thread_local BF16DWeightTimings timings; + return timings; +} + +class BF16DWeightScratch { + public: + using Kernel = GemmKernel224BF16; + static constexpr int M_STEP = Kernel::M_STEP; + static constexpr int N_STEP = Kernel::N_STEP; + + BF16DWeightScratch() = default; + BF16DWeightScratch(const BF16DWeightScratch&) = delete; + BF16DWeightScratch& operator=(const BF16DWeightScratch&) = delete; + + ~BF16DWeightScratch() { + std::free(a0_); + std::free(a1_); + std::free(b_); + } + + void ensure(int padded_k) { + if (padded_k <= capacity_k_) return; + const size_t a_elements = static_cast(M_STEP) * padded_k; + const size_t b_elements = static_cast(N_STEP) * padded_k; + resize(a0_, a_elements); + resize(a1_, a_elements); + resize(b_, b_elements); + capacity_k_ = padded_k; + } + + ggml_bf16_t* a0() { return a0_; } + ggml_bf16_t* a1() { return a1_; } + ggml_bf16_t* b() { return b_; } + float* c0() { return c0_; } + float* c1() { return c1_; } + + private: + static void resize(ggml_bf16_t*& buffer, size_t elements) { + void* replacement = nullptr; + if (posix_memalign(&replacement, 64, elements * sizeof(ggml_bf16_t)) != 0 || replacement == nullptr) { + throw std::runtime_error("failed to allocate BF16 dWeight scratch"); + } + std::free(buffer); + buffer = static_cast(replacement); + } + + int capacity_k_ = 0; + ggml_bf16_t* a0_ = nullptr; + ggml_bf16_t* a1_ = nullptr; + ggml_bf16_t* b_ = nullptr; + alignas(64) float c0_[M_STEP * N_STEP]; + alignas(64) float c1_[M_STEP * N_STEP]; +}; + +inline BF16DWeightScratch& bf16_dweight_scratch() { + static thread_local BF16DWeightScratch scratch; + return scratch; +} + +class BF16DWeightKernel { + public: + using Kernel = GemmKernel224BF16; + using BufferA = Kernel::BufferA; + using BufferB = Kernel::BufferB; + static constexpr int M_STEP = Kernel::M_STEP; + static constexpr int N_STEP = Kernel::N_STEP; + static constexpr int K_STEP = Kernel::K_STEP; + + static int padded_k(int routes) { return std::max(K_STEP, (routes + K_STEP - 1) / K_STEP * K_STEP); } + + static void configure_worker() { Kernel::config(); } + + static void pack_a_transposed(BufferA& destination, const ggml_bf16_t* source, int source_stride, int source_column, + int row_count, int routes, int destination_row = 0) { + assert(destination_row >= 0 && destination_row + row_count <= destination.max_m); + assert(destination_row % M_STEP == 0 && row_count <= M_STEP); + const int k = destination.k; + for (int k_begin = 0; k_begin < k; k_begin += K_STEP) { + ggml_bf16_t* tile = destination.get_submat(destination.max_m, k, destination_row, k_begin); + std::memset(tile, 0, M_STEP * K_STEP * sizeof(ggml_bf16_t)); + const int valid_k = std::min(K_STEP, routes - k_begin); + if (valid_k <= 0) continue; + for (int row = 0; row < row_count; ++row) { + for (int kk = 0; kk < valid_k; ++kk) { + tile[row * K_STEP + kk] = source[static_cast(k_begin + kk) * source_stride + source_column + row]; + } + } + } + } + + static void pack_b_transposed(BufferB& destination, const ggml_bf16_t* source, int source_stride, int source_column, + int row_count, int routes, int destination_row = 0) { + assert(destination_row >= 0 && destination_row + row_count <= destination.n); + assert(destination_row % N_STEP == 0 && row_count <= N_STEP); + const int k = destination.k; + for (int k_begin = 0; k_begin < k; k_begin += K_STEP) { + ggml_bf16_t* tile = destination.get_submat(destination.n, k, destination_row, k_begin); + std::memset(tile, 0, N_STEP * K_STEP * sizeof(ggml_bf16_t)); + const int valid_k = std::min(K_STEP, routes - k_begin); + if (valid_k > 0) { + for (int row = 0; row < row_count; ++row) { + for (int kk = 0; kk < valid_k; ++kk) { + tile[row * K_STEP + kk] = source[static_cast(k_begin + kk) * source_stride + source_column + row]; + } + } + } + transpose_16x16_32bit(reinterpret_cast<__m512i*>(tile)); + transpose_16x16_32bit(reinterpret_cast<__m512i*>(tile + Kernel::TILE_N * K_STEP)); + } + } + + private: + template + static inline __attribute__((always_inline)) void multiply_avx_rows( + int padded_k, float* destination, BufferA& a, BufferB& b, int m_begin, int n_begin, int row_offset) { + static_assert(ROWS > 0 && ROWS <= M_STEP); + __m512 accum_lo[ROWS]; + __m512 accum_hi[ROWS]; + +#pragma GCC unroll 12 + for (int row = 0; row < ROWS; ++row) { + accum_lo[row] = _mm512_setzero_ps(); + accum_hi[row] = _mm512_setzero_ps(); + } + + for (int k_begin = 0; k_begin < padded_k; k_begin += K_STEP) { + const auto* a_pairs = reinterpret_cast(a.get_submat(a.max_m, padded_k, m_begin, k_begin)); + const auto* b_vectors = reinterpret_cast(b.get_submat(b.n, padded_k, n_begin, k_begin)); + + for (int k_pair = 0; k_pair < K_STEP / 2; ++k_pair) { + const __m512bh b_lo = b_vectors[k_pair]; + const __m512bh b_hi = b_vectors[Kernel::TILE_N + k_pair]; +#pragma GCC unroll 12 + for (int row = 0; row < ROWS; ++row) { + const __m512bh a_pair = + reinterpret_cast<__m512bh>(_mm512_set1_epi32(a_pairs[(row_offset + row) * (K_STEP / 2) + k_pair])); + accum_lo[row] = _mm512_dpbf16_ps(accum_lo[row], a_pair, b_lo); + accum_hi[row] = _mm512_dpbf16_ps(accum_hi[row], a_pair, b_hi); + } + } + } + +#pragma GCC unroll 12 + for (int row = 0; row < ROWS; ++row) { + _mm512_store_ps(destination + (row_offset + row) * N_STEP, accum_lo[row]); + _mm512_store_ps(destination + (row_offset + row) * N_STEP + Kernel::TILE_N, accum_hi[row]); + } + } + + public: + + static void multiply(int padded_k, float* destination, BufferA& a, BufferB& b, int m_begin = 0, int n_begin = 0) { + assert(m_begin >= 0 && m_begin + M_STEP <= a.max_m); + assert(n_begin >= 0 && n_begin + N_STEP <= b.n); + assert(m_begin % M_STEP == 0 && n_begin % N_STEP == 0); + if constexpr (AMX_AVAILABLE) { + for (int k_block_begin = 0; k_block_begin < padded_k; k_block_begin += Kernel::K_BLOCK) { + Kernel::amx_kernel(a.max_m, b.n, padded_k, m_begin, n_begin, k_block_begin, destination, &a, &b); + } + } else { + // Keep the complete 32x32 FP32 output tile in registers. The generic forward + // kernel spans all 32 rows at once and spills its 64 accumulators on AVX512. + multiply_avx_rows<12>(padded_k, destination, a, b, m_begin, n_begin, 0); + multiply_avx_rows<12>(padded_k, destination, a, b, m_begin, n_begin, 12); + multiply_avx_rows<8>(padded_k, destination, a, b, m_begin, n_begin, 24); + } + } + + static void store_bf16(const float* source, ggml_bf16_t* destination, int destination_stride, int row_count, + int column_count) { + for (int row = 0; row < row_count; ++row) { + const float* src_row = source + row * N_STEP; + ggml_bf16_t* dst_row = destination + static_cast(row) * destination_stride; + if (column_count == N_STEP) { + __m512 lo = _mm512_loadu_ps(src_row); + __m512 hi = _mm512_loadu_ps(src_row + 16); + avx512_32xfp32_to_32xbf16(&lo, &hi, reinterpret_cast<__m512i*>(dst_row)); + } else { + for (int column = 0; column < column_count; ++column) { + dst_row[column] = GGML_FP32_TO_BF16(src_row[column]); + } + } + } + } +}; + +} // namespace amx + +#endif // AMX_BF16_DWEIGHT_HPP diff --git a/kt-kernel/operators/amx/moe_base.hpp b/kt-kernel/operators/amx/moe_base.hpp index 48c4da9ee..7329ff01d 100644 --- a/kt-kernel/operators/amx/moe_base.hpp +++ b/kt-kernel/operators/amx/moe_base.hpp @@ -673,21 +673,27 @@ class AMX_MOE_BASE { } void apply_activation(int activated_expert, int nth, int qlen) { + apply_activation_to(activated_expert, nth, qlen, m_local_gate_output_ptr_); + } + + void apply_activation_to(int activated_expert, int nth, int qlen, + const std::vector& destination_ptrs) { auto pool = config_.pool->get_subpool(tp_part_idx); - auto fn = [this, nth](int task_id) { + auto fn = [this, nth, &destination_ptrs](int task_id) { int expert_idx = m_expert_id_map_[task_id / nth]; int ith = task_id % nth; auto [n_start, n_end] = T::split_range_n(config_.intermediate_size, ith, nth); for (int i = 0; i < m_local_num_[expert_idx]; i++) { ggml_bf16_t* gate_output_ptr = &m_local_gate_output_ptr_[expert_idx][i * config_.intermediate_size]; ggml_bf16_t* up_output_ptr = &m_local_up_output_ptr_[expert_idx][i * config_.intermediate_size]; + ggml_bf16_t* destination_ptr = &destination_ptrs[expert_idx][i * config_.intermediate_size]; for (int j = n_start; j < n_end; j += 32) { __m512 gate_val0, gate_val1, up_val0, up_val1; avx512_32xbf16_to_32xfp32((__m512i*)(gate_output_ptr + j), &gate_val0, &gate_val1); avx512_32xbf16_to_32xfp32((__m512i*)(up_output_ptr + j), &up_val0, &up_val1); __m512 result0 = amx::act_fn(gate_val0, up_val0, config_.swiglu_limit, config_.swiglu_alpha); __m512 result1 = amx::act_fn(gate_val1, up_val1, config_.swiglu_limit, config_.swiglu_alpha); - avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i*)(gate_output_ptr + j)); + avx512_32xfp32_to_32xbf16(&result0, &result1, (__m512i*)(destination_ptr + j)); } } }; diff --git a/kt-kernel/operators/amx/sft_moe.hpp b/kt-kernel/operators/amx/sft_moe.hpp index 357ad9384..c0aef51b0 100644 --- a/kt-kernel/operators/amx/sft_moe.hpp +++ b/kt-kernel/operators/amx/sft_moe.hpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -26,8 +27,11 @@ #include #include "../../cpu_backend/worker_pool.h" +#include "../sft_profile.hpp" #include "ggml.h" +#include "la/bf16_dweight.hpp" #include "la/amx_kernels.hpp" +#include "la/amx_raw_kernels.hpp" #include "la/avx_kernels.hpp" #include "moe.hpp" @@ -218,6 +222,8 @@ struct supports_standard_mat_mul : std::false_type {}; template <> struct supports_standard_mat_mul : std::true_type {}; template <> +struct supports_standard_mat_mul : std::true_type {}; +template <> struct supports_standard_mat_mul : std::true_type {}; template <> struct supports_standard_mat_mul : std::true_type {}; @@ -236,6 +242,8 @@ struct has_bb_transposed_repack : std::false_type {}; template <> struct has_bb_transposed_repack : std::true_type {}; template <> +struct has_bb_transposed_repack : std::true_type {}; +template <> struct has_bb_transposed_repack : std::true_type {}; template inline constexpr bool has_bb_transposed_repack_v = has_bb_transposed_repack::value; @@ -355,6 +363,7 @@ template class BaseMOE = AMX_MOE_TP, bool SkipLoRA = class AMX_SFT_MOE_TP : public BaseMOE { public: static constexpr bool kSkipLoRA = SkipLoRA; + static constexpr bool kSupportsDirectBf16Reload = std::is_same_v; protected: using Base = BaseMOE; @@ -434,6 +443,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { // SFT configuration MOESFTConfig sft_config_; + SFTProfiler profiler_; // LoRA configuration (from MOESFTConfig) int lora_rank_; @@ -464,10 +474,6 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Last backward expert token distribution (for load balancing analysis) std::vector last_backward_expert_tokens_; - // Experts that had non-zero contributions in last backward (for selective zeroing) - std::vector last_backward_active_experts_; - bool grad_outputs_initialized_ = false; - // Cache buffer pools void* cache_input_pool_ = nullptr; void* cache_gate_output_pool_ = nullptr; @@ -502,6 +508,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Precomputed offsets for cache operations (avoid repeated heap allocation) std::vector cache_offsets_; + std::vector direct_cache_intermediate_ptrs_; // ===================================================== // AMX-optimized LoRA GEMM buffers (performance optimization) @@ -576,10 +583,16 @@ class AMX_SFT_MOE_TP : public BaseMOE { // BF16 buffer for scattered grad_output (before quantization to BufferA) std::vector grad_output_bf16_ptr_; // [expert_num] + // Immutable copy of the route-weighted grad_output used by full-weight down_proj gradients. + // grad_output_bf16_ptr_ is reused as gate/up grad_input scratch later in backward. + std::vector base_grad_output_bf16_ptr_; // [expert_num] + // Backward buffer pools void* backward_ba_pool_ = nullptr; void* backward_bc_pool_ = nullptr; void* grad_output_bf16_pool_ = nullptr; + void* base_grad_output_bf16_pool_ = nullptr; + void* dweight_shared_panel_pool_ = nullptr; void* backward_pool_ = nullptr; size_t backward_pool_bytes_ = 0; @@ -587,6 +600,8 @@ class AMX_SFT_MOE_TP : public BaseMOE { size_t backward_ba_pool_bytes_ = 0; size_t backward_bc_pool_bytes_ = 0; size_t grad_output_bf16_pool_bytes_ = 0; + size_t base_grad_output_bf16_pool_bytes_ = 0; + size_t dweight_shared_panel_pool_bytes_ = 0; // LoRA gradient computation pools (FP32, used in bwd_down_lora_precompute and grad computation) float* lora_grad_out_pool_ = nullptr; // [max_len * num_experts_per_tok * hidden_size] @@ -696,6 +711,13 @@ class AMX_SFT_MOE_TP : public BaseMOE { free_transposed_lora_weights(); } + void append_profile_stats(std::map& out, const std::string& prefix, + bool reset_after = false) { + profiler_.append(out, prefix, reset_after); + } + + void reset_profile_stats() { profiler_.reset(); } + /** * @brief Allocate forward-phase buffers. * Called at the start of forward_sft. @@ -712,6 +734,13 @@ class AMX_SFT_MOE_TP : public BaseMOE { work_required += round_up(lora_bc_out_pool_bytes_, kAmxAlignment); work_required += round_up(lora_intermediate_bf16_pool_bytes_, kAmxAlignment); + // Base-weight gradients reuse the forward pool after LoRA backward is complete. + if (sft_config_.full_weight_grad) { + const size_t base_grad_accumulator_bytes = + 3 * (size_t)config_.intermediate_size * config_.hidden_size * sizeof(float); + work_required = std::max(work_required, base_grad_accumulator_bytes); + } + alloc_or_resize_forward_pool(work_required); SFT_POOL_LOG("fwd_work", config_.layer_idx, tp_part_idx, 0, cache_stack_top_, forward_pool_bytes_, @@ -814,6 +843,8 @@ class AMX_SFT_MOE_TP : public BaseMOE { required += round_up(backward_ba_pool_bytes_, kAmxAlignment); required += round_up(backward_bc_pool_bytes_, kAmxAlignment); required += round_up(grad_output_bf16_pool_bytes_, kAmxAlignment); + required += round_up(base_grad_output_bf16_pool_bytes_, kAmxAlignment); + required += round_up(dweight_shared_panel_pool_bytes_, kAmxAlignment); required += round_up(lora_grad_out_pool_bytes_, kAmxAlignment); required += round_up(lora_inter_proj_pool_bytes_, kAmxAlignment); required += round_up(lora_grad_times_b_pool_bytes_, kAmxAlignment); @@ -846,6 +877,8 @@ class AMX_SFT_MOE_TP : public BaseMOE { assign(&backward_ba_pool_, backward_ba_pool_bytes_); assign(&backward_bc_pool_, backward_bc_pool_bytes_); assign(&grad_output_bf16_pool_, grad_output_bf16_pool_bytes_); + assign(&base_grad_output_bf16_pool_, base_grad_output_bf16_pool_bytes_); + assign(&dweight_shared_panel_pool_, dweight_shared_panel_pool_bytes_); assign((void**)&lora_grad_out_pool_, lora_grad_out_pool_bytes_); assign((void**)&lora_inter_proj_pool_, lora_inter_proj_pool_bytes_); @@ -897,7 +930,23 @@ class AMX_SFT_MOE_TP : public BaseMOE { */ void set_lora_params(int rank, float alpha) { lora_rank_ = rank; - lora_scaling_ = alpha / rank; + lora_scaling_ = rank > 0 ? alpha / rank : 0.0f; + } + + void set_full_weight_grad(bool enabled) { + sft_config_.full_weight_grad = enabled; + if constexpr (std::is_same_v) { + if (enabled) { + const size_t max_routes = static_cast(config_.max_len) * config_.num_experts_per_tok; + const size_t max_padded_routes = + ((max_routes + static_cast(config_.expert_num) * (T::K_STEP - 1) + T::K_STEP - 1) / T::K_STEP) * + T::K_STEP; + dweight_shared_panel_pool_bytes_ = + 2 * max_padded_routes * config_.hidden_size * sizeof(ggml_bf16_t); + } else { + dweight_shared_panel_pool_bytes_ = 0; + } + } } /** @@ -916,7 +965,11 @@ class AMX_SFT_MOE_TP : public BaseMOE { */ void forward_sft(int qlen, int k, const int64_t* expert_ids, const float* weights, const void* input, void* output, bool save_for_backward) { - uint64_t _fwd_start_cycles = __rdtsc(); + SFTProfileScope total_scope(profiler_, SFTProfileStage::FwdTotal); + SFTProfileScope checkpoint_scope( + profiler_, save_for_backward && config_.share_cache_pool ? SFTProfileStage::FwdRecomputeTotal + : SFTProfileStage::FwdInitialTotal); + auto stage_start = profiler_.start(); SFT_POOL_LOG("fwd_enter", config_.layer_idx, tp_part_idx, qlen, cache_stack_top_, forward_pool_bytes_, cache_pool_bytes_, backward_pool_bytes_, 0, "save_bwd=%d", (int)save_for_backward); @@ -948,8 +1001,10 @@ class AMX_SFT_MOE_TP : public BaseMOE { transpose_lora_b_weights(); lora_b_transposed_ = true; } + profiler_.record(SFTProfileStage::FwdSetup, stage_start); // Step 1: Expert routing (reuse base class logic) + stage_start = profiler_.start(); int activated_expert = 0; std::fill(m_local_num_.begin(), m_local_num_.end(), 0); for (int i = 0; i < qlen; i++) { @@ -967,8 +1022,12 @@ class AMX_SFT_MOE_TP : public BaseMOE { activated_expert++; } } + profiler_.record(SFTProfileStage::FwdRoute, stage_start); + profiler_.record_workload(static_cast(qlen), static_cast(qlen) * k, + static_cast(activated_expert)); // Step 2: Buffer pool allocation (reuse base class logic) + stage_start = profiler_.start(); size_t offset = 0; void* gate_up_ba_pool_ptr = Base::gate_up_ba_pool_; void* gate_bc_pool_ptr = Base::gate_bc_pool_; @@ -1069,6 +1128,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { printf("[OVERFLOW DEBUG L%d] Total tokens processed: %zu (offset after loop)\n", config_.layer_idx, offset); } } + profiler_.record(SFTProfileStage::FwdBufferSetup, stage_start); // Step 3: Copy input to expert buffers auto direct_or_pool = [&](int count, auto&& fn) { @@ -1080,7 +1140,12 @@ class AMX_SFT_MOE_TP : public BaseMOE { pool->do_work_stealing_job(count, nullptr, fn, nullptr); } }; + constexpr bool kSupportsDirectFullFTCache = std::is_same_v; + const bool use_direct_fullft_cache = + kSupportsDirectFullFTCache && save_for_backward && sft_config_.full_weight_grad && lora_rank_ == 0; + ForwardCache* direct_cache = nullptr; + stage_start = profiler_.start(); direct_or_pool(qlen, [&](int i) { for (int j = 0; j < k; j++) { if (expert_ids[i * k + j] < config_.num_gpu_experts || expert_ids[i * k + j] >= config_.expert_num) { @@ -1090,6 +1155,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { (ggml_bf16_t*)input + i * config_.hidden_size, sizeof(ggml_bf16_t) * config_.hidden_size); } }); + profiler_.record(SFTProfileStage::FwdInputScatter, stage_start); // NaN Check: Step 3 - Packed input if (is_nan_check_enabled()) { @@ -1106,12 +1172,33 @@ class AMX_SFT_MOE_TP : public BaseMOE { } // Step 4: Quantize input + stage_start = profiler_.start(); direct_or_pool(activated_expert, [this](int task_id) { int expert_idx = m_expert_id_map_[task_id]; gate_up_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_input_ptr_[expert_idx], 0, 1); }); + profiler_.record(SFTProfileStage::FwdInputPack, stage_start); + + if constexpr (kSupportsDirectFullFTCache) { + if (use_direct_fullft_cache) { + stage_start = profiler_.start(); + ForwardCache& cache = (cache_stack_top_ > 0) ? cache_stack_[cache_stack_top_ - 1] : push_cache(); + prepare_cache_metadata(cache, qlen, k, expert_ids, weights, activated_expert); + bind_direct_cache_outputs(cache, activated_expert); + cache.valid = true; + direct_cache = &cache; + profiler_.record(SFTProfileStage::FwdCacheMetadata, stage_start); + + stage_start = profiler_.start(); + copy_input_to_cache(cache, input, qlen); + profiler_.record_bytes(SFTProfileStage::FwdCacheInput, + static_cast(qlen) * config_.hidden_size * sizeof(ggml_bf16_t)); + profiler_.record(SFTProfileStage::FwdCacheInput, stage_start); + } + } // Step 5: Gate + Up GEMM (base projection) + stage_start = profiler_.start(); int nth = T::recommended_nth(config_.intermediate_size); pool->do_work_stealing_job( nth * activated_expert * 2, [](int _) { T::config(); }, @@ -1128,6 +1215,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { } }, nullptr); + profiler_.record(SFTProfileStage::FwdGateUpBase, stage_start); // NaN Check: Step 5 - Gate/Up GEMM output (before LoRA) if (is_nan_check_enabled()) { @@ -1148,9 +1236,11 @@ class AMX_SFT_MOE_TP : public BaseMOE { } // Step 5.5: Gate + Up LoRA (AVX512 BF16 - no BufferB conversion needed) + stage_start = profiler_.start(); if (!SkipLoRA) { compute_lora_gate_up(qlen, activated_expert); } + profiler_.record(SFTProfileStage::FwdGateUpLora, stage_start); // NaN Check: Step 5.5 - Gate/Up output (after LoRA) if (is_nan_check_enabled()) { @@ -1171,13 +1261,16 @@ class AMX_SFT_MOE_TP : public BaseMOE { } // Save gate/up outputs before activation (for backward) + stage_start = profiler_.start(); if (save_for_backward) { // If a cache entry already exists (checkpoint recompute scenario), // overwrite it instead of pushing a new one. This keeps the cache // consistent with the current forward's buffer state (max_m, routing) // and avoids cache stack overflow from duplicate pushes. - ForwardCache& cache = (cache_stack_top_ > 0) ? cache_stack_[cache_stack_top_ - 1] : push_cache(); - save_to_cache(cache, qlen, k, expert_ids, weights, activated_expert, input); + ForwardCache& cache = direct_cache != nullptr + ? *direct_cache + : ((cache_stack_top_ > 0) ? cache_stack_[cache_stack_top_ - 1] : push_cache()); + if (direct_cache == nullptr) save_to_cache(cache, qlen, k, expert_ids, weights, activated_expert, input); // NaN Check: Forward Cache - input, gate_output, up_output if (is_nan_check_enabled()) { @@ -1215,9 +1308,20 @@ class AMX_SFT_MOE_TP : public BaseMOE { check_cache_bf16("up_output_cache", cache.up_output_cache, total_tokens * config_.intermediate_size); } } + profiler_.record(SFTProfileStage::FwdCacheGateUp, stage_start); // Step 6: Activation (silu(gate) * up) - { Base::apply_activation(activated_expert, nth, qlen); } + stage_start = profiler_.start(); + if (direct_cache != nullptr) { + Base::apply_activation_to(activated_expert, nth, qlen, direct_cache_intermediate_ptrs_); + for (int i = 0; i < activated_expert; ++i) { + const int expert_idx = m_expert_id_map_[i]; + m_local_gate_output_ptr_[expert_idx] = direct_cache_intermediate_ptrs_[expert_idx]; + } + } else { + Base::apply_activation(activated_expert, nth, qlen); + } + profiler_.record(SFTProfileStage::FwdActivation, stage_start); // NaN Check: Step 6 - Activation output (silu(gate) * up) if (is_nan_check_enabled()) { @@ -1234,9 +1338,10 @@ class AMX_SFT_MOE_TP : public BaseMOE { } // Save intermediate AFTER activation for backward_down (Bug #17c fix) + stage_start = profiler_.start(); if (save_for_backward) { ForwardCache& cache = cache_stack_[cache_stack_top_ - 1]; // Get the cache we just pushed - save_intermediate_to_cache(cache, activated_expert); + if (direct_cache == nullptr) save_intermediate_to_cache(cache, activated_expert); // NaN Check: Forward Cache - intermediate_cache if (is_nan_check_enabled()) { @@ -1270,8 +1375,10 @@ class AMX_SFT_MOE_TP : public BaseMOE { } } } + profiler_.record(SFTProfileStage::FwdCacheIntermediate, stage_start); // Step 7: Quantize intermediate for down projection + stage_start = profiler_.start(); pool->do_work_stealing_job( activated_expert, nullptr, [this](int task_id) { @@ -1279,8 +1386,10 @@ class AMX_SFT_MOE_TP : public BaseMOE { down_ba_[expert_idx]->from_mat(m_local_num_[expert_idx], m_local_gate_output_ptr_[expert_idx], 0, 1); }, nullptr); + profiler_.record(SFTProfileStage::FwdDownPack, stage_start); // Step 8: Down GEMM + stage_start = profiler_.start(); nth = T::recommended_nth(config_.hidden_size); pool->do_work_stealing_job( nth * activated_expert, [](int _) { T::config(); }, @@ -1291,6 +1400,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { down_bc_[expert_idx]->to_mat(m_local_num_[expert_idx], m_local_down_output_ptr_[expert_idx], ith, nth); }, nullptr); + profiler_.record(SFTProfileStage::FwdDownBase, stage_start); // NaN Check: Step 8 - Down GEMM output (before LoRA) if (is_nan_check_enabled()) { @@ -1307,10 +1417,12 @@ class AMX_SFT_MOE_TP : public BaseMOE { } // Step 8.5: Down LoRA (AVX512 BF16 - no BufferB conversion needed) + stage_start = profiler_.start(); if (down_lora_a_ != nullptr && down_lora_b_ != nullptr) { ForwardCache* cache_ptr = save_for_backward ? &cache_stack_[cache_stack_top_ - 1] : nullptr; compute_lora_down(qlen, activated_expert, cache_ptr); } + profiler_.record(SFTProfileStage::FwdDownLora, stage_start); // NaN Check: Step 8.5 - Down output (after LoRA) if (is_nan_check_enabled()) { @@ -1327,12 +1439,15 @@ class AMX_SFT_MOE_TP : public BaseMOE { } // Save down_output for grad_weights computation + stage_start = profiler_.start(); if (save_for_backward) { ForwardCache& cache = cache_stack_[cache_stack_top_ - 1]; // Get the cache we just pushed - save_down_output_to_cache(cache, activated_expert); + if (direct_cache == nullptr) save_down_output_to_cache(cache, activated_expert); } + profiler_.record(SFTProfileStage::FwdCacheDown, stage_start); // Step 9: Weighted merge + stage_start = profiler_.start(); pool->do_work_stealing_job( qlen, nullptr, [this, output, k, expert_ids, weights](int i) { @@ -1357,6 +1472,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { } }, nullptr); + profiler_.record(SFTProfileStage::FwdWeightedMerge, stage_start); // NaN Check: Step 9 - Final output (after weighted merge) if (is_nan_check_enabled()) { @@ -1394,7 +1510,10 @@ class AMX_SFT_MOE_TP : public BaseMOE { void backward(const void* grad_output, void* grad_input, void* grad_gate_lora_a, void* grad_gate_lora_b, void* grad_up_lora_a, void* grad_up_lora_b, void* grad_down_lora_a, void* grad_down_lora_b, void* grad_weights, int full_intermediate_size = 0, float* fp32_grad_down_lora_b = nullptr, - float* fp32_grad_gate_lora_a = nullptr, float* fp32_grad_up_lora_a = nullptr) { + float* fp32_grad_gate_lora_a = nullptr, float* fp32_grad_up_lora_a = nullptr, + void* grad_gate_proj = nullptr, void* grad_up_proj = nullptr, void* grad_down_proj = nullptr) { + SFTProfileScope total_scope(profiler_, SFTProfileStage::BwdTotal); + auto stage_start = profiler_.start(); // If full_intermediate_size not provided, use local (non-TP mode) if (full_intermediate_size == 0) full_intermediate_size = config_.intermediate_size; SFT_POOL_LOG("bwd_enter", config_.layer_idx, tp_part_idx, 0, cache_stack_top_, forward_pool_bytes_, @@ -1520,10 +1639,12 @@ class AMX_SFT_MOE_TP : public BaseMOE { m_local_down_output_ptr_[i] = m_local_down_output_ + offset * config_.hidden_size; offset += m_local_num_[i]; } + profiler_.record(SFTProfileStage::BwdSetup, stage_start); // Restore input data from cache into m_local_input_ (shared_mem_buffer may have been // overwritten by subsequent layers' forward passes). This is needed for gate/up LoRA // gradient computation which reads from m_local_input_ptr_. + stage_start = profiler_.start(); auto pool_local = config_.pool->get_subpool(tp_part_idx); auto restore_input = [&](int i) { for (int j = 0; j < k; j++) { @@ -1545,14 +1666,17 @@ class AMX_SFT_MOE_TP : public BaseMOE { } else { pool_local->do_work_stealing_job(qlen, nullptr, restore_input, nullptr); } + profiler_.record(SFTProfileStage::BwdCacheRestore, stage_start); // Step 1: Down projection backward + stage_start = profiler_.start(); if constexpr (supports_standard_mat_mul_v) { backward_down_amx(cache, grad_output, grad_down_lora_a, grad_down_lora_b, full_intermediate_size, fp32_grad_down_lora_b); } else { // backward_down(cache, grad_output, grad_down_lora_a, grad_down_lora_b); } + profiler_.record(SFTProfileStage::BwdDownTotal, stage_start); // // Compute total tokens for debug // size_t total_tokens = 0; @@ -1628,7 +1752,9 @@ class AMX_SFT_MOE_TP : public BaseMOE { // } // } + stage_start = profiler_.start(); backward_activation(cache); + profiler_.record(SFTProfileStage::BwdActivation, stage_start); // NaN Check: Step 2 - After backward_activation if (is_nan_check_enabled()) { @@ -1668,12 +1794,14 @@ class AMX_SFT_MOE_TP : public BaseMOE { // } // } + stage_start = profiler_.start(); if constexpr (supports_standard_mat_mul_v) { backward_gate_up_amx(cache, grad_input, grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b, full_intermediate_size, fp32_grad_gate_lora_a, fp32_grad_up_lora_a); } else { // backward_gate_up(cache, grad_input, grad_gate_lora_a, grad_gate_lora_b, grad_up_lora_a, grad_up_lora_b); } + profiler_.record(SFTProfileStage::BwdGateUpTotal, stage_start); // NaN Check: Step 3 - After backward_gate_up if (is_nan_check_enabled()) { @@ -1710,6 +1838,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Step 4: Compute grad_weights (gradient for routing weights) // grad_weights[token_idx, expert_pos] = dot(grad_output[token_idx], down_output[token, expert]) + stage_start = profiler_.start(); if (grad_weights != nullptr) { auto pool = config_.pool->get_subpool(tp_part_idx); float* grad_w = (float*)grad_weights; @@ -1761,6 +1890,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { pool->do_work_stealing_job(qlen, nullptr, compute_grad_weight, nullptr); } } + profiler_.record(SFTProfileStage::BwdRouterGrad, stage_start); // NaN Check: Step 4 - After grad_weights computation if (is_nan_check_enabled() && grad_weights != nullptr) { @@ -1878,7 +2008,16 @@ class AMX_SFT_MOE_TP : public BaseMOE { print_grad_stats_fp32("grad_weights", (const float*)grad_weights, qlen * k); } - // ★ Cache pool is NOT freed here — kept for reuse across steps. + // ===================================================================== + // Step 5: Base weight gradient accumulation (full weight grad mode) + // ===================================================================== + stage_start = profiler_.start(); + if (sft_config_.full_weight_grad && grad_gate_proj && grad_up_proj && grad_down_proj) { + backward_base_weight_grad(cache, full_intermediate_size, grad_gate_proj, grad_up_proj, grad_down_proj); + } + profiler_.record(SFTProfileStage::BwdBaseWeightGrad, stage_start); + + // \u2605 Cache pool is NOT freed here \u2014 kept for reuse across steps. // alloc_or_resize_cache_pool() is grow-only, so same-seqlen steps // reuse the existing allocation without malloc/free overhead. // Previously: free_seqlen_buffers() was called here, costing ~3.6ms per TP. @@ -1887,6 +2026,481 @@ class AMX_SFT_MOE_TP : public BaseMOE { cache.valid = false; } + /** + * @brief Compute base weight gradients via outer-product accumulation. + * + * For each activated expert e, computes: + * grad_gate_proj[e] = grad_gate_out[e]^T @ input[e] -> [I, H] + * grad_up_proj[e] = grad_up_out[e]^T @ input[e] -> [I, H] + * grad_down_proj[e] = grad_output[e]^T @ intermediate[e] -> [H, I] + * + * Uses FP32 accumulator for precision, writes BF16 output. + */ + void backward_base_weight_grad(const ForwardCache& cache, int full_intermediate_size, void* grad_gate_proj, + void* grad_up_proj, void* grad_down_proj) { + auto stage_start = profiler_.start(); + const int H = config_.hidden_size; + const int I = config_.intermediate_size; + const int F = full_intermediate_size; + const int activated_expert = cache.activated_expert_cache; + + if (F < I) { + throw std::runtime_error("full_intermediate_size must be at least the TP-local intermediate_size"); + } + + auto* ggp = static_cast(grad_gate_proj); // TP slice of [E, F, H] + auto* gup_ptr = static_cast(grad_up_proj); // TP slice of [E, F, H] + auto* gdp = static_cast(grad_down_proj); // TP slice of [E, H, F] + + std::vector expert_offsets(activated_expert); + size_t token_offset = 0; + int max_routes = 0; + for (int task_id = 0; task_id < activated_expert; task_id++) { + expert_offsets[task_id] = token_offset; + int expert_idx = cache.m_expert_id_map_cache[task_id]; + const int routes = cache.m_local_num_cache[expert_idx]; + token_offset += routes; + max_routes = std::max(max_routes, routes); + } + profiler_.record(SFTProfileStage::BwdBaseWeightGradOffsets, stage_start); + + if constexpr (std::is_same_v) { + using DWeightKernel = amx::BF16DWeightKernel; + constexpr int TILE_M = DWeightKernel::M_STEP; + constexpr int TILE_N = DWeightKernel::N_STEP; + const int i_tiles = (I + TILE_M - 1) / TILE_M; + const int h_tiles = (H + TILE_N - 1) / TILE_N; + const int tasks_per_expert = i_tiles * 2; + const int total_tasks = activated_expert * tasks_per_expert; + const int max_padded_k = DWeightKernel::padded_k(max_routes); + auto pool = config_.pool->get_subpool(tp_part_idx); + const bool profile_inner = profiler_.enabled(); + + std::vector expert_padded_k(activated_expert); + std::vector panel_offsets(activated_expert); + size_t panel_elements = 0; + for (int expert_task = 0; expert_task < activated_expert; ++expert_task) { + const int expert_idx = cache.m_expert_id_map_cache[expert_task]; + panel_offsets[expert_task] = panel_elements; + expert_padded_k[expert_task] = DWeightKernel::padded_k(cache.m_local_num_cache[expert_idx]); + panel_elements += static_cast(H) * expert_padded_k[expert_task]; + } + const size_t required_panel_bytes = 2 * panel_elements * sizeof(ggml_bf16_t); + if (required_panel_bytes > dweight_shared_panel_pool_bytes_) { + throw std::runtime_error("BF16 dWeight shared panel pool is too small"); + } + auto* panel_base = static_cast(dweight_shared_panel_pool_); + auto* input_panels = panel_base; + auto* grad_output_panels = panel_base + panel_elements; + + // X and route-weighted dY are reused by every I tile. Pack each H tile once per expert, + // then keep the existing expert x I-tile compute schedule for load balancing. + stage_start = profiler_.start(); + const int panel_tasks_per_expert = h_tiles * 2; + const int panel_tasks = activated_expert * panel_tasks_per_expert; + if (panel_tasks > 0) { + pool->do_work_stealing_job( + panel_tasks, + [](int _) { amx::bf16_dweight_timings().reset(); }, + [&, h_tiles, panel_tasks_per_expert](int task_id) { + const int expert_task = task_id / panel_tasks_per_expert; + const int local_task = task_id % panel_tasks_per_expert; + const bool pack_grad_output = local_task >= h_tiles; + const int h_tile = local_task % h_tiles; + const int h_start = h_tile * TILE_N; + const int h_count = std::min(TILE_N, H - h_start); + const int expert_idx = cache.m_expert_id_map_cache[expert_task]; + const int routes = cache.m_local_num_cache[expert_idx]; + if (routes == 0) return; + + const int padded_k = expert_padded_k[expert_task]; + const size_t panel_offset = panel_offsets[expert_task]; + auto& timings = amx::bf16_dweight_timings(); + const auto profile_operation = [profile_inner](uint64_t& elapsed_ns, uint64_t& calls, + auto&& operation) { + if (!profile_inner) { + operation(); + return; + } + const auto begin = SFTProfiler::Clock::now(); + operation(); + elapsed_ns += static_cast( + std::chrono::duration_cast(SFTProfiler::Clock::now() - begin).count()); + calls++; + }; + + if (pack_grad_output) { + typename DWeightKernel::BufferA panel(H, padded_k, grad_output_panels + panel_offset); + profile_operation(timings.panel_grad_output_ns, timings.panel_grad_output_calls, [&] { + DWeightKernel::pack_a_transposed(panel, base_grad_output_bf16_ptr_[expert_idx], H, h_start, h_count, + routes, h_start); + }); + } else { + typename DWeightKernel::BufferB panel(H, padded_k, input_panels + panel_offset); + profile_operation(timings.panel_input_ns, timings.panel_input_calls, [&] { + DWeightKernel::pack_b_transposed(panel, m_local_input_ptr_[expert_idx], H, h_start, h_count, routes, + h_start); + }); + } + }, + [this](int _) { + auto& timings = amx::bf16_dweight_timings(); + profiler_.record_ns(SFTProfileStage::BwdBaseWeightGradPanelInput, timings.panel_input_ns, + timings.panel_input_calls); + profiler_.record_ns(SFTProfileStage::BwdBaseWeightGradPanelGradOutput, timings.panel_grad_output_ns, + timings.panel_grad_output_calls); + }); + } + profiler_.record(SFTProfileStage::BwdBaseWeightGradPanelPack, stage_start); + profiler_.record_bytes(SFTProfileStage::BwdBaseWeightGradPanelInput, + panel_elements * sizeof(ggml_bf16_t)); + profiler_.record_bytes(SFTProfileStage::BwdBaseWeightGradPanelGradOutput, + panel_elements * sizeof(ggml_bf16_t)); + + stage_start = profiler_.start(); + if (total_tasks > 0) { + pool->do_work_stealing_job( + total_tasks, + [max_padded_k](int _) { + DWeightKernel::configure_worker(); + amx::bf16_dweight_scratch().ensure(max_padded_k); + amx::bf16_dweight_timings().reset(); + }, + [&, i_tiles, h_tiles, tasks_per_expert](int task_id) { + const int expert_task = task_id / tasks_per_expert; + const int local_task = task_id % tasks_per_expert; + const bool do_down = local_task >= i_tiles; + const int i_tile = local_task % i_tiles; + const int expert_idx = cache.m_expert_id_map_cache[expert_task]; + const int routes = cache.m_local_num_cache[expert_idx]; + if (routes == 0) return; + + const size_t pos_start = expert_offsets[expert_task]; + const int padded_k = DWeightKernel::padded_k(routes); + const int i_start = i_tile * TILE_M; + const int i_count = std::min(TILE_M, I - i_start); + const size_t panel_offset = panel_offsets[expert_task]; + auto& scratch = amx::bf16_dweight_scratch(); + auto& timings = amx::bf16_dweight_timings(); + typename DWeightKernel::BufferA a0(TILE_M, padded_k, scratch.a0()); + typename DWeightKernel::BufferA a1(TILE_M, padded_k, scratch.a1()); + typename DWeightKernel::BufferB b(TILE_N, padded_k, scratch.b()); + typename DWeightKernel::BufferA shared_grad_output(H, padded_k, grad_output_panels + panel_offset); + typename DWeightKernel::BufferB shared_input(H, padded_k, input_panels + panel_offset); + + auto profile_operation = [profile_inner](uint64_t& elapsed_ns, uint64_t& calls, auto&& operation) { + if (!profile_inner) { + operation(); + return; + } + const auto begin = SFTProfiler::Clock::now(); + operation(); + elapsed_ns += static_cast( + std::chrono::duration_cast(SFTProfiler::Clock::now() - begin).count()); + calls++; + }; + + if (do_down) { + const ggml_bf16_t* intermediate = cache.intermediate_cache + pos_start * I; + profile_operation(timings.pack_b_ns, timings.pack_b_calls, [&] { + DWeightKernel::pack_b_transposed(b, intermediate, I, i_start, i_count, routes); + }); + + ggml_bf16_t* down_dst = gdp + static_cast(expert_idx) * H * F; + for (int h_tile = 0; h_tile < h_tiles; ++h_tile) { + const int h_start = h_tile * TILE_M; + const int h_count = std::min(TILE_M, H - h_start); + profile_operation(timings.kernel_down_ns, timings.kernel_down_calls, + [&] { + DWeightKernel::multiply(padded_k, scratch.c0(), shared_grad_output, b, h_start, + 0); + }); + profile_operation(timings.store_ns, timings.store_calls, [&] { + DWeightKernel::store_bf16(scratch.c0(), down_dst + static_cast(h_start) * F + i_start, F, + h_count, i_count); + }); + } + return; + } + + const ggml_bf16_t* gate_grad = grad_gate_output_ + pos_start * I; + const ggml_bf16_t* up_grad = grad_up_output_ + pos_start * I; + profile_operation(timings.pack_a_ns, timings.pack_a_calls, [&] { + DWeightKernel::pack_a_transposed(a0, gate_grad, I, i_start, i_count, routes); + DWeightKernel::pack_a_transposed(a1, up_grad, I, i_start, i_count, routes); + }); + + ggml_bf16_t* gate_dst = ggp + static_cast(expert_idx) * F * H; + ggml_bf16_t* up_dst = gup_ptr + static_cast(expert_idx) * F * H; + for (int h_tile = 0; h_tile < h_tiles; ++h_tile) { + const int h_start = h_tile * TILE_N; + const int h_count = std::min(TILE_N, H - h_start); + profile_operation(timings.kernel_gate_up_ns, timings.kernel_gate_up_calls, [&] { + DWeightKernel::multiply(padded_k, scratch.c0(), a0, shared_input, 0, h_start); + DWeightKernel::multiply(padded_k, scratch.c1(), a1, shared_input, 0, h_start); + }); + profile_operation(timings.store_ns, timings.store_calls, [&] { + DWeightKernel::store_bf16(scratch.c0(), gate_dst + static_cast(i_start) * H + h_start, H, + i_count, h_count); + DWeightKernel::store_bf16(scratch.c1(), up_dst + static_cast(i_start) * H + h_start, H, + i_count, h_count); + }); + } + }, + [this](int _) { + auto& timings = amx::bf16_dweight_timings(); + profiler_.record_ns(SFTProfileStage::BwdBaseWeightGradPackA, timings.pack_a_ns, timings.pack_a_calls); + profiler_.record_ns(SFTProfileStage::BwdBaseWeightGradPackB, timings.pack_b_ns, timings.pack_b_calls); + profiler_.record_ns(SFTProfileStage::BwdBaseWeightGradKernelGateUp, timings.kernel_gate_up_ns, + timings.kernel_gate_up_calls); + profiler_.record_ns(SFTProfileStage::BwdBaseWeightGradKernelDown, timings.kernel_down_ns, + timings.kernel_down_calls); + profiler_.record_ns(SFTProfileStage::BwdBaseWeightGradStore, timings.store_ns, timings.store_calls); + }); + } + profiler_.record(SFTProfileStage::BwdBaseWeightGradMatMat, stage_start); + return; + } + + if constexpr (std::is_same_v && amx::AMX_AVAILABLE) { + constexpr int TILE_M = T::M_STEP; + constexpr int TILE_N = T::N_STEP; + constexpr int TILE_K = T::K_STEP; + static_assert(TILE_M == 32 && TILE_N == 32 && TILE_K == 32, + "base-weight gradient tile packing assumes 32x32x32 AMX BF16 tiles"); + + const int i_tiles = (I + TILE_M - 1) / TILE_M; + const int h_tiles = (H + TILE_N - 1) / TILE_N; + // Keep enough fixed-i strips for load balancing while amortizing task pickup and panel packing over H. + const int tasks_per_expert = i_tiles * 2; // fixed-i strips for fused gate/up plus down + const int total_tasks = activated_expert * tasks_per_expert; + auto pool = config_.pool->get_subpool(tp_part_idx); + + stage_start = profiler_.start(); + pool->do_work_stealing_job( + total_tasks, [](int _) { T::config(); }, + [&, i_tiles, h_tiles, tasks_per_expert](int task_id) { + const int expert_task = task_id / tasks_per_expert; + const int local_task = task_id % tasks_per_expert; + const bool do_down = local_task >= i_tiles; + const int i_tile = local_task % i_tiles; + const int expert_idx = cache.m_expert_id_map_cache[expert_task]; + const int m = cache.m_local_num_cache[expert_idx]; + if (m == 0) return; + + const size_t pos_start = expert_offsets[expert_task]; + const int k_tiles = (m + TILE_K - 1) / TILE_K; + constexpr size_t A_TILE_ELEMENTS = TILE_M * TILE_K; + constexpr size_t B_TILE_ELEMENTS = TILE_N * TILE_K; + constexpr size_t TILE_ALIGNMENT = 64; + constexpr size_t ALIGNMENT_PADDING = TILE_ALIGNMENT / sizeof(ggml_bf16_t); + const size_t packed_a_elements = (size_t)k_tiles * A_TILE_ELEMENTS; + const size_t packed_b_elements = (size_t)k_tiles * B_TILE_ELEMENTS; + + thread_local std::vector packed_a0_storage; + thread_local std::vector packed_a1_storage; + thread_local std::vector packed_b_storage; + auto resize_aligned = [](std::vector& storage, size_t elements) { + const size_t required = elements + ALIGNMENT_PADDING; + if (storage.size() < required) storage.resize(required); + const auto raw = reinterpret_cast(storage.data()); + const auto aligned = (raw + TILE_ALIGNMENT - 1) & ~(std::uintptr_t)(TILE_ALIGNMENT - 1); + return reinterpret_cast(aligned); + }; + + ggml_bf16_t* packed_a0 = resize_aligned(packed_a0_storage, packed_a_elements); + ggml_bf16_t* packed_b = resize_aligned(packed_b_storage, packed_b_elements); + alignas(64) float c0[TILE_M * TILE_N]; + + const int i_start = i_tile * TILE_N; + const int i_count = std::min(TILE_N, I - i_start); + + if (do_down) { + const ggml_bf16_t* grad_output = base_grad_output_bf16_ptr_[expert_idx]; + const ggml_bf16_t* intermediate = cache.intermediate_cache + pos_start * I; + std::memset(packed_b, 0, packed_b_elements * sizeof(ggml_bf16_t)); + for (int kt = 0; kt < k_tiles; kt++) { + const int k_start = kt * TILE_K; + const int k_count = std::min(TILE_K, m - k_start); + ggml_bf16_t* b_tile = packed_b + (size_t)kt * B_TILE_ELEMENTS; + for (int col = 0; col < i_count; col++) { + for (int kk = 0; kk < k_count; kk++) { + b_tile[col * TILE_K + kk] = intermediate[(size_t)(k_start + kk) * I + i_start + col]; + } + } + amx::transpose_16x16_32bit(reinterpret_cast<__m512i*>(b_tile)); + amx::transpose_16x16_32bit(reinterpret_cast<__m512i*>(b_tile + T::TILE_N * TILE_K)); + } + + ggml_bf16_t* down_dst = gdp + (size_t)expert_idx * H * F; + for (int h_tile = 0; h_tile < h_tiles; h_tile++) { + const int h_start = h_tile * TILE_M; + const int h_count = std::min(TILE_M, H - h_start); + std::memset(packed_a0, 0, packed_a_elements * sizeof(ggml_bf16_t)); + for (int kt = 0; kt < k_tiles; kt++) { + const int k_start = kt * TILE_K; + const int k_count = std::min(TILE_K, m - k_start); + ggml_bf16_t* a_tile = packed_a0 + (size_t)kt * A_TILE_ELEMENTS; + for (int row = 0; row < h_count; row++) { + for (int kk = 0; kk < k_count; kk++) { + a_tile[row * TILE_K + kk] = grad_output[(size_t)(k_start + kk) * H + h_start + row]; + } + } + } + + // Keep the full 32x32 FP32 C tile resident for the complete K reduction. + T::clean_c(); + for (int kt = 0; kt < k_tiles; kt++) { + T::load_b(packed_b + (size_t)kt * B_TILE_ELEMENTS, TILE_K * sizeof(ggml_bf16_t)); + T::load_a(packed_a0 + (size_t)kt * A_TILE_ELEMENTS, TILE_K * sizeof(ggml_bf16_t)); + T::run_tile(); + } + T::store_c(c0, TILE_N * sizeof(float)); + + for (int row = 0; row < h_count; row++) { + for (int col = 0; col < i_count; col++) { + down_dst[(size_t)(h_start + row) * F + i_start + col] = GGML_FP32_TO_BF16(c0[row * TILE_N + col]); + } + } + } + return; + } + + ggml_bf16_t* packed_a1 = resize_aligned(packed_a1_storage, packed_a_elements); + std::memset(packed_a0, 0, packed_a_elements * sizeof(ggml_bf16_t)); + std::memset(packed_a1, 0, packed_a_elements * sizeof(ggml_bf16_t)); + for (int kt = 0; kt < k_tiles; kt++) { + const int k_start = kt * TILE_K; + const int k_count = std::min(TILE_K, m - k_start); + ggml_bf16_t* gate_a_tile = packed_a0 + (size_t)kt * A_TILE_ELEMENTS; + ggml_bf16_t* up_a_tile = packed_a1 + (size_t)kt * A_TILE_ELEMENTS; + for (int row = 0; row < i_count; row++) { + for (int kk = 0; kk < k_count; kk++) { + gate_a_tile[row * TILE_K + kk] = grad_gate_output_[(pos_start + k_start + kk) * I + i_start + row]; + up_a_tile[row * TILE_K + kk] = grad_up_output_[(pos_start + k_start + kk) * I + i_start + row]; + } + } + } + + const ggml_bf16_t* input = m_local_input_ptr_[expert_idx]; + ggml_bf16_t* gate_dst = ggp + (size_t)expert_idx * F * H; + ggml_bf16_t* up_dst = gup_ptr + (size_t)expert_idx * F * H; + alignas(64) float c1[TILE_M * TILE_N]; + for (int h_tile = 0; h_tile < h_tiles; h_tile++) { + const int h_start = h_tile * TILE_N; + const int h_count = std::min(TILE_N, H - h_start); + std::memset(packed_b, 0, packed_b_elements * sizeof(ggml_bf16_t)); + for (int kt = 0; kt < k_tiles; kt++) { + const int k_start = kt * TILE_K; + const int k_count = std::min(TILE_K, m - k_start); + ggml_bf16_t* b_tile = packed_b + (size_t)kt * B_TILE_ELEMENTS; + for (int col = 0; col < h_count; col++) { + for (int kk = 0; kk < k_count; kk++) { + b_tile[col * TILE_K + kk] = input[(size_t)(k_start + kk) * H + h_start + col]; + } + } + amx::transpose_16x16_32bit(reinterpret_cast<__m512i*>(b_tile)); + amx::transpose_16x16_32bit(reinterpret_cast<__m512i*>(b_tile + T::TILE_N * TILE_K)); + } + + // Gate and up each consume all four C tiles, so retain C across K in two separate passes. + T::clean_c(); + for (int kt = 0; kt < k_tiles; kt++) { + T::load_b(packed_b + (size_t)kt * B_TILE_ELEMENTS, TILE_K * sizeof(ggml_bf16_t)); + T::load_a(packed_a0 + (size_t)kt * A_TILE_ELEMENTS, TILE_K * sizeof(ggml_bf16_t)); + T::run_tile(); + } + T::store_c(c0, TILE_N * sizeof(float)); + + T::clean_c(); + for (int kt = 0; kt < k_tiles; kt++) { + T::load_b(packed_b + (size_t)kt * B_TILE_ELEMENTS, TILE_K * sizeof(ggml_bf16_t)); + T::load_a(packed_a1 + (size_t)kt * A_TILE_ELEMENTS, TILE_K * sizeof(ggml_bf16_t)); + T::run_tile(); + } + T::store_c(c1, TILE_N * sizeof(float)); + + for (int row = 0; row < i_count; row++) { + for (int col = 0; col < h_count; col++) { + gate_dst[(size_t)(i_start + row) * H + h_start + col] = GGML_FP32_TO_BF16(c0[row * TILE_N + col]); + up_dst[(size_t)(i_start + row) * H + h_start + col] = GGML_FP32_TO_BF16(c1[row * TILE_N + col]); + } + } + } + }, + nullptr); + profiler_.record(SFTProfileStage::BwdBaseWeightGradAmx, stage_start); + return; + } + + for (int task_id = 0; task_id < activated_expert; task_id++) { + int expert_idx = cache.m_expert_id_map_cache[task_id]; + int m = cache.m_local_num_cache[expert_idx]; + if (m == 0) continue; + + const size_t pos_start = expert_offsets[task_id]; + + // Allocate FP32 accumulators from forward pool (safe during backward) + float* acc_gate = static_cast(forward_pool_); // [I, H] + float* acc_up = acc_gate + (size_t)I * H; // [I, H] + float* acc_down = acc_up + (size_t)I * H; // [H, I] + + stage_start = profiler_.start(); + std::memset(acc_gate, 0, (size_t)I * H * sizeof(float)); + std::memset(acc_up, 0, (size_t)I * H * sizeof(float)); + std::memset(acc_down, 0, (size_t)H * I * sizeof(float)); + profiler_.record(SFTProfileStage::BwdBaseWeightGradZero, stage_start); + + stage_start = profiler_.start(); + for (int t = 0; t < m; t++) { + const ggml_bf16_t* input_row = m_local_input_ptr_[expert_idx] + (size_t)t * H; + const ggml_bf16_t* gate_grad_row = grad_gate_output_ + (size_t)(pos_start + t) * I; + const ggml_bf16_t* up_grad_row = grad_up_output_ + (size_t)(pos_start + t) * I; + + // gate_proj grad: [I, H] += grad_gate_out[t]^T @ input[t] + for (int i = 0; i < I; i++) { + float gg = GGML_BF16_TO_FP32(gate_grad_row[i]); + float gu = GGML_BF16_TO_FP32(up_grad_row[i]); + for (int h = 0; h < H; h++) { + float inp = GGML_BF16_TO_FP32(input_row[h]); + acc_gate[i * H + h] += gg * inp; + acc_up[i * H + h] += gu * inp; + } + } + } + profiler_.record(SFTProfileStage::BwdBaseWeightGradGateUp, stage_start); + + stage_start = profiler_.start(); + for (int t = 0; t < m; t++) { + const ggml_bf16_t* inter_row = cache.intermediate_cache + (size_t)(pos_start + t) * I; + const ggml_bf16_t* grad_out_row = base_grad_output_bf16_ptr_[expert_idx] + (size_t)t * H; + // down_proj grad: [H, I] += grad_output[t]^T @ intermediate[t] + for (int h = 0; h < H; h++) { + float go = GGML_BF16_TO_FP32(grad_out_row[h]); + for (int i = 0; i < I; i++) { + acc_down[h * I + i] += go * GGML_BF16_TO_FP32(inter_row[i]); + } + } + } + profiler_.record(SFTProfileStage::BwdBaseWeightGradDown, stage_start); + + // Convert FP32 accumulators to BF16 and store + stage_start = profiler_.start(); + for (int i = 0; i < I; i++) { + for (int h = 0; h < H; h++) { + ggp[(size_t)expert_idx * F * H + (size_t)i * H + h] = GGML_FP32_TO_BF16(acc_gate[i * H + h]); + gup_ptr[(size_t)expert_idx * F * H + (size_t)i * H + h] = GGML_FP32_TO_BF16(acc_up[i * H + h]); + } + } + for (int h = 0; h < H; h++) { + for (int i = 0; i < I; i++) { + gdp[(size_t)expert_idx * H * F + (size_t)h * F + i] = GGML_FP32_TO_BF16(acc_down[h * I + i]); + } + } + profiler_.record(SFTProfileStage::BwdBaseWeightGradStore, stage_start); + } + } + /** * @brief Get qlen from the top of the forward cache stack. * @@ -2313,6 +2927,54 @@ class AMX_SFT_MOE_TP : public BaseMOE { backward_weights_prepared_ = true; } + void load_forward_weights_from_full_bf16(void* gate_proj, void* up_proj, void* down_proj, int full_intermediate_size, + int intermediate_offset) { + if constexpr (!kSupportsDirectBf16Reload) { + throw std::runtime_error("direct BF16 reload requires GemmKernel224BF16"); + } else { + const int H = config_.hidden_size; + const int I = config_.intermediate_size; + if (gate_proj == nullptr || up_proj == nullptr || down_proj == nullptr) { + throw std::runtime_error("direct BF16 reload requires all three base weights"); + } + if (intermediate_offset < 0 || full_intermediate_size < intermediate_offset + I) { + throw std::runtime_error("invalid TP slice for direct BF16 reload"); + } + + const auto* physical_to_logical_map = static_cast(config_.physical_to_logical_map); + auto* gate = static_cast(gate_proj); + auto* up = static_cast(up_proj); + auto* down = static_cast(down_proj); + auto pool = config_.pool->get_subpool(tp_part_idx); + + const int gate_up_nth = T::recommended_nth(I); + pool->do_work_stealing_job( + gate_up_nth * config_.expert_num, nullptr, + [&, gate_up_nth, physical_to_logical_map](int task_id) { + const int physical_expert = task_id / gate_up_nth; + const int ith = task_id % gate_up_nth; + const size_t logical_expert = expert_map(physical_to_logical_map, physical_expert); + const size_t source_offset = + logical_expert * full_intermediate_size * H + static_cast(intermediate_offset) * H; + gate_bb_[physical_expert]->from_mat_strided(gate + source_offset, H, ith, gate_up_nth); + up_bb_[physical_expert]->from_mat_strided(up + source_offset, H, ith, gate_up_nth); + }, + nullptr); + + const int down_nth = T::recommended_nth(H); + pool->do_work_stealing_job( + down_nth * config_.expert_num, nullptr, + [&, down_nth, physical_to_logical_map](int task_id) { + const int physical_expert = task_id / down_nth; + const int ith = task_id % down_nth; + const size_t logical_expert = expert_map(physical_to_logical_map, physical_expert); + const size_t source_offset = logical_expert * H * full_intermediate_size + intermediate_offset; + down_bb_[physical_expert]->from_mat_strided(down + source_offset, full_intermediate_size, ith, down_nth); + }, + nullptr); + } + } + /** * @brief Standalone method for async backward BB repack (Phase 2). * Called from TP_MOE_SFT::submit_backward_repack() on a separate thread. @@ -2320,6 +2982,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { * and sets the owner layer on the shared pool. */ void prepare_backward_bb_for_async() { + SFTProfileScope profile_scope(profiler_, SFTProfileStage::BackwardRepack); if constexpr (!supports_standard_mat_mul_v) return; if (backward_bb_pool_bytes_ == 0) return; @@ -2624,6 +3287,14 @@ class AMX_SFT_MOE_TP : public BaseMOE { cache_down_output_bytes_ = (size_t)max_cache_depth_ * ml * k_tok * H * sizeof(ggml_bf16_t); grad_buffer_bytes_ = ml * k_tok * I * sizeof(ggml_bf16_t); + if constexpr (std::is_same_v) { + if (sft_config_.full_weight_grad) { + const size_t max_padded_routes = + ((ml * k_tok + static_cast(config_.expert_num) * (K_STEP - 1) + K_STEP - 1) / K_STEP) * K_STEP; + dweight_shared_panel_pool_bytes_ = + 2 * max_padded_routes * H * sizeof(ggml_bf16_t); // X^T BufferB + route-weighted dY^T BufferA + } + } // ===================================================== // Calculate LoRA AMX buffer sizes @@ -2732,6 +3403,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { // BF16 buffer for scattered grad_output grad_output_bf16_pool_bytes_ = safe_alloc_tokens * config_.hidden_size * sizeof(ggml_bf16_t) + align_overhead; + base_grad_output_bf16_pool_bytes_ = grad_output_bf16_pool_bytes_; // LoRA gradient computation FP32 pools (used in bwd_down_lora_precompute and grad computation) // Total tokens across all activated experts = safe_alloc_tokens @@ -2764,6 +3436,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { backward_ba_pool_bytes_ = 0; backward_bc_pool_bytes_ = 0; grad_output_bf16_pool_bytes_ = 0; + base_grad_output_bf16_pool_bytes_ = 0; backward_bb_pool_bytes_ = 0; lora_grad_out_pool_bytes_ = 0; lora_inter_proj_pool_bytes_ = 0; @@ -2829,6 +3502,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { cache_stack_.resize(max_cache_depth_); // Preallocate cache offsets to avoid heap allocation in hot path cache_offsets_.resize(config_.expert_num + 1); + direct_cache_intermediate_ptrs_.resize(config_.expert_num, nullptr); for (int i = 0; i < max_cache_depth_; i++) { // Note: cache pointers (input_cache, gate_output_cache, etc.) are set in alloc_forward_buffers() cache_stack_[i].input_cache = nullptr; @@ -2904,6 +3578,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { grad_intermediate_bc_.resize(config_.expert_num); grad_gate_up_bc_.resize(config_.expert_num); grad_output_bf16_ptr_.resize(config_.expert_num); + base_grad_output_bf16_ptr_.resize(config_.expert_num); // Resize vectors - backward BufferB (transposed base weights) gate_backward_bb_.resize(config_.expert_num); @@ -2994,6 +3669,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { // BF16 pointer: will be assigned in backward grad_output_bf16_ptr_[i] = nullptr; + base_grad_output_bf16_ptr_[i] = nullptr; } // ===================================================== @@ -3854,10 +4530,8 @@ class AMX_SFT_MOE_TP : public BaseMOE { return cache_stack_[--cache_stack_top_]; } - void save_to_cache(ForwardCache& cache, int qlen, int k, const int64_t* expert_ids, const float* weights, - int activated_expert, const void* input) { - auto pool = config_.pool->get_subpool(tp_part_idx); - + void prepare_cache_metadata(ForwardCache& cache, int qlen, int k, const int64_t* expert_ids, const float* weights, + int activated_expert) { cache.qlen_cache = qlen; cache.k_cache = k; cache.activated_expert_cache = activated_expert; @@ -3883,6 +4557,39 @@ class AMX_SFT_MOE_TP : public BaseMOE { int expert_idx = m_expert_id_map_[i]; cache_offsets_[i + 1] = cache_offsets_[i] + m_local_num_[expert_idx]; } + } + + void copy_input_to_cache(ForwardCache& cache, const void* input, int qlen) { + auto pool = config_.pool->get_subpool(tp_part_idx); + const size_t total_bytes = static_cast(qlen) * config_.hidden_size * sizeof(ggml_bf16_t); + constexpr size_t kChunkBytes = 2 * 1024 * 1024; + const int chunks = static_cast((total_bytes + kChunkBytes - 1) / kChunkBytes); + pool->do_work_stealing_job( + chunks, nullptr, + [&, input, total_bytes](int chunk) { + const size_t offset = static_cast(chunk) * kChunkBytes; + const size_t bytes = std::min(kChunkBytes, total_bytes - offset); + std::memcpy(reinterpret_cast(cache.input_cache) + offset, + reinterpret_cast(input) + offset, bytes); + }, + nullptr); + } + + void bind_direct_cache_outputs(ForwardCache& cache, int activated_expert) { + for (int i = 0; i < activated_expert; ++i) { + const int expert_idx = m_expert_id_map_[i]; + const size_t offset = cache_offsets_[i]; + m_local_gate_output_ptr_[expert_idx] = cache.gate_output_cache + offset * config_.intermediate_size; + m_local_up_output_ptr_[expert_idx] = cache.up_output_cache + offset * config_.intermediate_size; + m_local_down_output_ptr_[expert_idx] = cache.down_output_cache + offset * config_.hidden_size; + direct_cache_intermediate_ptrs_[expert_idx] = cache.intermediate_cache + offset * config_.intermediate_size; + } + } + + void save_to_cache(ForwardCache& cache, int qlen, int k, const int64_t* expert_ids, const float* weights, + int activated_expert, const void* input) { + auto pool = config_.pool->get_subpool(tp_part_idx); + prepare_cache_metadata(cache, qlen, k, expert_ids, weights, activated_expert); // Parallel copy: input(1 task) + gate(N tasks) + up(N tasks) = 1 + 2N tasks // This parallelizes the ~1.8MB input copy that was previously serial @@ -4162,6 +4869,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { void backward_down_amx(const ForwardCache& cache, const void* grad_output, void* grad_down_lora_a, void* grad_down_lora_b, int full_intermediate_size = 0, float* fp32_grad_down_lora_b = nullptr) { + auto stage_start = profiler_.start(); if (full_intermediate_size == 0) full_intermediate_size = config_.intermediate_size; auto pool = config_.pool->get_subpool(tp_part_idx); int activated_expert = cache.activated_expert_cache; @@ -4194,6 +4902,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { char* backward_ba_ptr = (char*)backward_ba_pool_; char* backward_bc_ptr = (char*)backward_bc_pool_; char* grad_output_bf16_ptr = (char*)grad_output_bf16_pool_; + char* base_grad_output_bf16_ptr = (char*)base_grad_output_bf16_pool_; for (int task_id = 0; task_id < activated_expert; task_id++) { int expert_idx = m_expert_id_map_[task_id]; @@ -4215,9 +4924,16 @@ class AMX_SFT_MOE_TP : public BaseMOE { // Allocate BF16 buffer for scattered grad_output grad_output_bf16_ptr_[expert_idx] = (ggml_bf16_t*)grad_output_bf16_ptr; grad_output_bf16_ptr += align64(local_max_m * config_.hidden_size * sizeof(ggml_bf16_t)); + + // Preserve the route-weighted upstream gradient for down_proj full-weight gradients. + // The working grad_output buffer is overwritten by backward_gate_up_amx(). + base_grad_output_bf16_ptr_[expert_idx] = (ggml_bf16_t*)base_grad_output_bf16_ptr; + base_grad_output_bf16_ptr += align64(local_max_m * config_.hidden_size * sizeof(ggml_bf16_t)); } // NOTE: no full-buffer memset here; grad_intermediate_ is overwritten by to_mat() for active tokens. + profiler_.record(SFTProfileStage::BwdDownSetup, stage_start); + stage_start = profiler_.start(); // ===================================================== // Step 1: Zero per-expert grad_output buffers @@ -4274,6 +4990,19 @@ class AMX_SFT_MOE_TP : public BaseMOE { }); } + // Keep an immutable copy before grad_output_bf16_ptr_ is reused as gate/up grad_input scratch. + if (sft_config_.full_weight_grad) { + direct_or_pool(activated_expert, [this](int task_id) { + int expert_idx = m_expert_id_map_[task_id]; + int num_tokens = m_local_num_[expert_idx]; + if (num_tokens == 0) return; + std::memcpy(base_grad_output_bf16_ptr_[expert_idx], grad_output_bf16_ptr_[expert_idx], + static_cast(num_tokens) * config_.hidden_size * sizeof(ggml_bf16_t)); + }); + } + profiler_.record(SFTProfileStage::BwdDownScatter, stage_start); + stage_start = profiler_.start(); + // ===================================================== // Step 3: Quantize scattered grad_output to BufferA // ===================================================== @@ -4328,6 +5057,8 @@ class AMX_SFT_MOE_TP : public BaseMOE { bc->to_mat(m, grad_intermediate_ + expert_offsets[task_idx], ith, nth); }, nullptr); + profiler_.record(SFTProfileStage::BwdDownBaseDx, stage_start); + stage_start = profiler_.start(); // ===================================================== // Step 3.5: Add LoRA contribution to grad_intermediate (AVX512) @@ -4769,6 +5500,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { } } } + profiler_.record(SFTProfileStage::BwdDownLora, stage_start); } void backward_activation(const ForwardCache& cache) { @@ -4905,6 +5637,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { void backward_gate_up_amx(const ForwardCache& cache, void* grad_input, void* grad_gate_lora_a, void* grad_gate_lora_b, void* grad_up_lora_a, void* grad_up_lora_b, int full_intermediate_size = 0, float* fp32_grad_gate_lora_a = nullptr, float* fp32_grad_up_lora_a = nullptr) { + auto stage_start = profiler_.start(); if (full_intermediate_size == 0) full_intermediate_size = config_.intermediate_size; auto pool = config_.pool->get_subpool(tp_part_idx); int activated_expert = cache.activated_expert_cache; @@ -4928,7 +5661,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { ggml_bf16_t* grad_up_b = (ggml_bf16_t*)grad_up_lora_b; assert(backward_weights_prepared_); - if (gate_lora_a_ != nullptr && gate_lora_b_ != nullptr) { + if (lora_rank_ > 0 && gate_lora_a_ != nullptr && gate_lora_b_ != nullptr) { prepare_lora_backward_weights(); } @@ -5086,8 +5819,11 @@ class AMX_SFT_MOE_TP : public BaseMOE { scatter_to_grad_input(1.0f); }; + profiler_.record(SFTProfileStage::BwdGateUpSetup, stage_start); + stage_start = profiler_.start(); base_pass(false); // gate base_pass(true); // up + profiler_.record(SFTProfileStage::BwdGateUpBaseDx, stage_start); // // DEBUG: Check m_local_input_ptr_ AFTER base_pass (before LoRA) // { @@ -5113,9 +5849,10 @@ class AMX_SFT_MOE_TP : public BaseMOE { // } // Skip all LoRA computation when SkipLoRA is true - if (SkipLoRA || gate_lora_a_ == nullptr || gate_lora_b_ == nullptr) { + if (SkipLoRA || lora_rank_ <= 0 || gate_lora_a_ == nullptr || gate_lora_b_ == nullptr) { return; } + stage_start = profiler_.start(); const bool use_fp32_lora_a = (fp32_grad_gate_lora_a != nullptr); @@ -5483,6 +6220,7 @@ class AMX_SFT_MOE_TP : public BaseMOE { lora_pass_remainder(false); // gate: gb_gradin_fused, scatter, gradA lora_pass_remainder(true); // up: gb_gradin_fused, scatter, gradA + profiler_.record(SFTProfileStage::BwdGateUpLora, stage_start); } }; diff --git a/kt-kernel/operators/amx/test/test_bf16_dweight.cpp b/kt-kernel/operators/amx/test/test_bf16_dweight.cpp new file mode 100644 index 000000000..5afca4f87 --- /dev/null +++ b/kt-kernel/operators/amx/test/test_bf16_dweight.cpp @@ -0,0 +1,313 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../la/bf16_dweight.hpp" + +namespace { + +using DWeightKernel = amx::BF16DWeightKernel; +using Kernel = DWeightKernel::Kernel; + +void* alloc_buffer(size_t bytes) { + void* pointer = nullptr; + if (posix_memalign(&pointer, 64, (bytes + 63) / 64 * 64) != 0 || pointer == nullptr) std::abort(); + std::memset(pointer, 0, bytes); + return pointer; +} + +void fill_random(std::vector& values, unsigned seed) { + std::mt19937 generator(seed); + std::uniform_real_distribution distribution(-0.25f, 0.25f); + for (auto& value : values) value = GGML_FP32_TO_BF16(distribution(generator)); +} + +bool run_case(int routes, int rows, int columns) { + constexpr int source_column = 2; + constexpr int destination_column = 3; + const int lhs_stride = rows + source_column + 3; + const int rhs_stride = columns + source_column + 5; + const int destination_stride = columns + destination_column + 7; + const int padded_k = DWeightKernel::padded_k(routes); + + std::vector lhs(static_cast(routes) * lhs_stride); + std::vector rhs(static_cast(routes) * rhs_stride); + std::vector actual(static_cast(rows) * destination_stride); + std::vector expected(static_cast(rows) * columns); + fill_random(lhs, static_cast(routes * 17 + rows)); + fill_random(rhs, static_cast(routes * 31 + columns)); + + void* a_memory = alloc_buffer(Kernel::BufferA::required_size(Kernel::M_STEP, padded_k)); + void* b_memory = alloc_buffer(Kernel::BufferB::required_size(Kernel::N_STEP, padded_k)); + Kernel::BufferA a(Kernel::M_STEP, padded_k, a_memory); + Kernel::BufferB b(Kernel::N_STEP, padded_k, b_memory); + alignas(64) float accumulator[Kernel::M_STEP * Kernel::N_STEP]; + + DWeightKernel::pack_a_transposed(a, lhs.data(), lhs_stride, source_column, rows, routes); + DWeightKernel::pack_b_transposed(b, rhs.data(), rhs_stride, source_column, columns, routes); + DWeightKernel::multiply(padded_k, accumulator, a, b); + DWeightKernel::store_bf16(accumulator, actual.data() + destination_column, destination_stride, rows, columns); + + for (int row = 0; row < rows; ++row) { + for (int column = 0; column < columns; ++column) { + float sum = 0.0f; + for (int route = 0; route < routes; ++route) { + sum += GGML_BF16_TO_FP32(lhs[static_cast(route) * lhs_stride + source_column + row]) * + GGML_BF16_TO_FP32(rhs[static_cast(route) * rhs_stride + source_column + column]); + } + expected[static_cast(row) * columns + column] = GGML_FP32_TO_BF16(sum); + } + } + + double difference_sq = 0.0; + double expected_sq = 0.0; + double actual_sq = 0.0; + double dot = 0.0; + float max_abs = 0.0f; + for (int row = 0; row < rows; ++row) { + for (int column = 0; column < columns; ++column) { + const float expected_value = GGML_BF16_TO_FP32(expected[static_cast(row) * columns + column]); + const float actual_value = + GGML_BF16_TO_FP32(actual[static_cast(row) * destination_stride + destination_column + column]); + const double difference = static_cast(actual_value) - expected_value; + difference_sq += difference * difference; + expected_sq += static_cast(expected_value) * expected_value; + actual_sq += static_cast(actual_value) * actual_value; + dot += static_cast(expected_value) * actual_value; + max_abs = std::max(max_abs, std::fabs(actual_value - expected_value)); + } + } + + const double relative_l2 = std::sqrt(difference_sq / std::max(expected_sq, 1e-30)); + const double cosine = dot / std::sqrt(std::max(expected_sq * actual_sq, 1e-30)); + const bool passed = relative_l2 <= 0.01 && cosine >= 0.999; + std::printf("BF16 dWeight routes=%d shape=%dx%d: rel_l2=%.6e cosine=%.9f max_abs=%.6e %s\n", routes, rows, columns, + relative_l2, cosine, max_abs, passed ? "PASS" : "FAIL"); + + std::free(a_memory); + std::free(b_memory); + return passed; +} + +bool run_shared_panel_case(int routes, int rows, int columns) { + const int padded_k = DWeightKernel::padded_k(routes); + const int padded_rows = (rows + Kernel::M_STEP - 1) / Kernel::M_STEP * Kernel::M_STEP; + const int padded_columns = (columns + Kernel::N_STEP - 1) / Kernel::N_STEP * Kernel::N_STEP; + std::vector lhs(static_cast(routes) * rows); + std::vector rhs(static_cast(routes) * columns); + std::vector actual(static_cast(rows) * columns); + fill_random(lhs, static_cast(routes * 41 + rows)); + fill_random(rhs, static_cast(routes * 43 + columns)); + + void* a_memory = alloc_buffer(Kernel::BufferA::required_size(padded_rows, padded_k)); + void* b_memory = alloc_buffer(Kernel::BufferB::required_size(padded_columns, padded_k)); + Kernel::BufferA a(padded_rows, padded_k, a_memory); + Kernel::BufferB b(padded_columns, padded_k, b_memory); + for (int row = 0; row < rows; row += Kernel::M_STEP) { + DWeightKernel::pack_a_transposed(a, lhs.data(), rows, row, std::min(Kernel::M_STEP, rows - row), routes, row); + } + for (int column = 0; column < columns; column += Kernel::N_STEP) { + DWeightKernel::pack_b_transposed(b, rhs.data(), columns, column, + std::min(Kernel::N_STEP, columns - column), routes, column); + } + + alignas(64) float accumulator[Kernel::M_STEP * Kernel::N_STEP]; + for (int row = 0; row < rows; row += Kernel::M_STEP) { + const int row_count = std::min(Kernel::M_STEP, rows - row); + for (int column = 0; column < columns; column += Kernel::N_STEP) { + const int column_count = std::min(Kernel::N_STEP, columns - column); + DWeightKernel::multiply(padded_k, accumulator, a, b, row, column); + DWeightKernel::store_bf16(accumulator, actual.data() + static_cast(row) * columns + column, columns, + row_count, column_count); + } + } + + double difference_sq = 0.0; + double expected_sq = 0.0; + double actual_sq = 0.0; + double dot = 0.0; + for (int row = 0; row < rows; ++row) { + for (int column = 0; column < columns; ++column) { + float reference = 0.0f; + for (int route = 0; route < routes; ++route) { + reference += GGML_BF16_TO_FP32(lhs[static_cast(route) * rows + row]) * + GGML_BF16_TO_FP32(rhs[static_cast(route) * columns + column]); + } + const float expected = GGML_BF16_TO_FP32(GGML_FP32_TO_BF16(reference)); + const float value = GGML_BF16_TO_FP32(actual[static_cast(row) * columns + column]); + const double difference = static_cast(value) - expected; + difference_sq += difference * difference; + expected_sq += static_cast(expected) * expected; + actual_sq += static_cast(value) * value; + dot += static_cast(expected) * value; + } + } + const double relative_l2 = std::sqrt(difference_sq / std::max(expected_sq, 1e-30)); + const double cosine = dot / std::sqrt(std::max(expected_sq * actual_sq, 1e-30)); + const bool passed = relative_l2 <= 0.01 && cosine >= 0.999; + std::printf("BF16 dWeight shared panel routes=%d shape=%dx%d: rel_l2=%.6e cosine=%.9f %s\n", routes, rows, + columns, relative_l2, cosine, passed ? "PASS" : "FAIL"); + + std::free(a_memory); + std::free(b_memory); + return passed; +} + +bool run_amx_benchmark() { + if constexpr (!amx::AMX_AVAILABLE) { + std::printf("BF16 dWeight AMX benchmark: SKIP (AMX unavailable)\n"); + return true; + } + + constexpr int routes = 1024; + constexpr int iterations = 20000; + constexpr int rounds = 7; + const int padded_k = DWeightKernel::padded_k(routes); + std::vector lhs(static_cast(routes) * Kernel::M_STEP); + std::vector rhs(static_cast(routes) * Kernel::N_STEP); + fill_random(lhs, 20260717); + fill_random(rhs, 20260718); + + void* a_memory = alloc_buffer(Kernel::BufferA::required_size(Kernel::M_STEP, padded_k)); + void* b_memory = alloc_buffer(Kernel::BufferB::required_size(Kernel::N_STEP, padded_k)); + Kernel::BufferA a(Kernel::M_STEP, padded_k, a_memory); + Kernel::BufferB b(Kernel::N_STEP, padded_k, b_memory); + alignas(64) float accumulator[Kernel::M_STEP * Kernel::N_STEP]; + DWeightKernel::pack_a_transposed(a, lhs.data(), Kernel::M_STEP, 0, Kernel::M_STEP, routes); + DWeightKernel::pack_b_transposed(b, rhs.data(), Kernel::N_STEP, 0, Kernel::N_STEP, routes); + + auto legacy_tile_loop = [&] { + Kernel::clean_c(); + for (int k_begin = 0; k_begin < padded_k; k_begin += Kernel::K_STEP) { + Kernel::load_b(b.get_submat(Kernel::N_STEP, padded_k, 0, k_begin), + Kernel::K_STEP * sizeof(ggml_bf16_t)); + Kernel::load_a(a.get_submat(Kernel::M_STEP, padded_k, 0, k_begin), + Kernel::K_STEP * sizeof(ggml_bf16_t)); + Kernel::run_tile(); + } + Kernel::store_c(accumulator, Kernel::N_STEP * sizeof(float)); + }; + auto common_driver = [&] { DWeightKernel::multiply(padded_k, accumulator, a, b); }; + + for (int warmup = 0; warmup < 200; ++warmup) { + legacy_tile_loop(); + common_driver(); + } + + auto measure = [&](auto&& operation) { + const auto begin = std::chrono::steady_clock::now(); + for (int iteration = 0; iteration < iterations; ++iteration) operation(); + return std::chrono::duration(std::chrono::steady_clock::now() - begin).count() / iterations; + }; + std::vector legacy_ns; + std::vector common_ns; + for (int round = 0; round < rounds; ++round) { + if (round % 2 == 0) { + legacy_ns.push_back(measure(legacy_tile_loop)); + common_ns.push_back(measure(common_driver)); + } else { + common_ns.push_back(measure(common_driver)); + legacy_ns.push_back(measure(legacy_tile_loop)); + } + } + std::sort(legacy_ns.begin(), legacy_ns.end()); + std::sort(common_ns.begin(), common_ns.end()); + const double legacy_median = legacy_ns[rounds / 2]; + const double common_median = common_ns[rounds / 2]; + const double ratio = common_median / legacy_median; + const bool passed = ratio <= 1.05; + std::printf("BF16 dWeight AMX kernel routes=%d: legacy=%.1f ns common=%.1f ns ratio=%.4f %s\n", routes, + legacy_median, common_median, ratio, passed ? "PASS" : "FAIL"); + + std::free(a_memory); + std::free(b_memory); + return passed; +} + +bool run_avx_benchmark() { + if constexpr (amx::AMX_AVAILABLE) { + std::printf("BF16 dWeight AVX benchmark: SKIP (AMX enabled)\n"); + return true; + } + + constexpr int routes = 64; + constexpr int iterations = 50000; + constexpr int rounds = 7; + const int padded_k = DWeightKernel::padded_k(routes); + std::vector lhs(static_cast(routes) * Kernel::M_STEP); + std::vector rhs(static_cast(routes) * Kernel::N_STEP); + fill_random(lhs, 20260717); + fill_random(rhs, 20260718); + + void* a_memory = alloc_buffer(Kernel::BufferA::required_size(Kernel::M_STEP, padded_k)); + void* b_memory = alloc_buffer(Kernel::BufferB::required_size(Kernel::N_STEP, padded_k)); + Kernel::BufferA a(Kernel::M_STEP, padded_k, a_memory); + Kernel::BufferB b(Kernel::N_STEP, padded_k, b_memory); + alignas(64) float accumulator[Kernel::M_STEP * Kernel::N_STEP]; + DWeightKernel::pack_a_transposed(a, lhs.data(), Kernel::M_STEP, 0, Kernel::M_STEP, routes); + DWeightKernel::pack_b_transposed(b, rhs.data(), Kernel::N_STEP, 0, Kernel::N_STEP, routes); + + auto generic_driver = [&] { + for (int k_block_begin = 0; k_block_begin < padded_k; k_block_begin += Kernel::K_BLOCK) { + Kernel::avx_kernel_4(Kernel::M_STEP, Kernel::N_STEP, padded_k, 0, 0, k_block_begin, accumulator, &a, &b); + } + }; + auto register_blocked_driver = [&] { DWeightKernel::multiply(padded_k, accumulator, a, b); }; + + for (int warmup = 0; warmup < 200; ++warmup) { + generic_driver(); + register_blocked_driver(); + } + auto measure = [&](auto&& operation) { + const auto begin = std::chrono::steady_clock::now(); + for (int iteration = 0; iteration < iterations; ++iteration) operation(); + return std::chrono::duration(std::chrono::steady_clock::now() - begin).count() / iterations; + }; + + std::vector generic_ns; + std::vector register_blocked_ns; + for (int round = 0; round < rounds; ++round) { + if (round % 2 == 0) { + generic_ns.push_back(measure(generic_driver)); + register_blocked_ns.push_back(measure(register_blocked_driver)); + } else { + register_blocked_ns.push_back(measure(register_blocked_driver)); + generic_ns.push_back(measure(generic_driver)); + } + } + std::sort(generic_ns.begin(), generic_ns.end()); + std::sort(register_blocked_ns.begin(), register_blocked_ns.end()); + const double generic_median = generic_ns[rounds / 2]; + const double register_blocked_median = register_blocked_ns[rounds / 2]; + const double ratio = register_blocked_median / generic_median; + const bool passed = ratio <= 1.05; + std::printf("BF16 dWeight AVX kernel routes=%d: generic=%.1f ns register_blocked=%.1f ns ratio=%.4f %s\n", routes, + generic_median, register_blocked_median, ratio, passed ? "PASS" : "FAIL"); + + std::free(a_memory); + std::free(b_memory); + return passed; +} + +} // namespace + +int main(int argc, char** argv) { + DWeightKernel::configure_worker(); + if (argc == 2 && std::strcmp(argv[1], "--benchmark") == 0) { + return (run_amx_benchmark() && run_avx_benchmark()) ? 0 : 1; + } + bool passed = true; + for (int routes : {1, 31, 32, 33, 65, 1792, 1825}) { + passed = run_case(routes, 32, 32) && passed; + } + passed = run_case(33, 17, 29) && passed; + passed = run_case(65, 31, 7) && passed; + passed = run_shared_panel_case(33, 45, 77) && passed; + passed = run_shared_panel_case(65, 64, 96) && passed; + return passed ? 0 : 1; +} diff --git a/kt-kernel/operators/amx/test/test_raw_bf16_repack.cpp b/kt-kernel/operators/amx/test/test_raw_bf16_repack.cpp new file mode 100644 index 000000000..55a13097e --- /dev/null +++ b/kt-kernel/operators/amx/test/test_raw_bf16_repack.cpp @@ -0,0 +1,133 @@ +#include +#include +#include +#include +#include + +#include "../la/amx_kernels.hpp" +#include "../la/amx_raw_kernels.hpp" + +namespace { + +using Kernel = amx::GemmKernel224BF16; +using BufferB = Kernel::BufferB; + +void* alloc_buffer(size_t bytes) { + void* ptr = std::aligned_alloc(64, (bytes + 63) / 64 * 64); + if (ptr == nullptr) std::abort(); + std::memset(ptr, 0, bytes); + return ptr; +} + +void fill_random(std::vector& values, unsigned seed) { + std::mt19937 generator(seed); + std::uniform_real_distribution distribution(-1.0f, 1.0f); + for (auto& value : values) value = GGML_FP32_TO_BF16(distribution(generator)); +} + +void from_mat(BufferB& buffer, ggml_bf16_t* source) { + int nth = Kernel::recommended_nth(buffer.n); + for (int ith = 0; ith < nth; ith++) buffer.from_mat(source, ith, nth); +} + +void from_mat_strided(BufferB& buffer, ggml_bf16_t* source, int source_stride) { + int nth = Kernel::recommended_nth(buffer.n); + for (int ith = 0; ith < nth; ith++) buffer.from_mat_strided(source, source_stride, ith, nth); +} + +void to_mat(const BufferB& buffer, ggml_bf16_t* destination) { + int nth = Kernel::recommended_nth(buffer.n); + for (int ith = 0; ith < nth; ith++) buffer.to_mat(destination, ith, nth); +} + +void from_mat_transposed(BufferB& buffer, ggml_bf16_t* source, int source_n, int source_k) { + int nth = Kernel::recommended_nth(buffer.n); + for (int ith = 0; ith < nth; ith++) buffer.from_mat_transposed(source, source_n, source_k, ith, nth); +} + +void from_bb_transposed(BufferB& destination, const BufferB& source) { + int nth = Kernel::recommended_nth(destination.n); + for (int ith = 0; ith < nth; ith++) destination.from_bb_transposed(source, ith, nth); +} + +bool run_case(int n, int k) { + const size_t count = (size_t)n * k; + std::vector source(count); + std::vector roundtrip(count); + std::vector transposed(count); + std::vector direct_transposed(count); + fill_random(source, (unsigned)(n * 31 + k)); + + void* forward_memory = alloc_buffer(BufferB::required_size(n, k)); + void* expected_memory = alloc_buffer(BufferB::required_size(k, n)); + void* direct_memory = alloc_buffer(BufferB::required_size(k, n)); + BufferB forward(n, k, forward_memory); + BufferB expected(k, n, expected_memory); + BufferB direct(k, n, direct_memory); + + from_mat(forward, source.data()); + to_mat(forward, roundtrip.data()); + from_mat_transposed(expected, source.data(), n, k); + from_bb_transposed(direct, forward); + to_mat(expected, transposed.data()); + to_mat(direct, direct_transposed.data()); + + bool roundtrip_ok = std::memcmp(source.data(), roundtrip.data(), count * sizeof(ggml_bf16_t)) == 0; + bool direct_ok = std::memcmp(transposed.data(), direct_transposed.data(), count * sizeof(ggml_bf16_t)) == 0; + bool transpose_ok = true; + for (int row = 0; row < n && transpose_ok; row++) { + for (int column = 0; column < k; column++) { + if (source[(size_t)row * k + column].bits != transposed[(size_t)column * n + row].bits) { + transpose_ok = false; + break; + } + } + } + + std::free(forward_memory); + std::free(expected_memory); + std::free(direct_memory); + std::printf("raw BF16 repack %dx%d: roundtrip=%s transpose=%s direct=%s\n", n, k, roundtrip_ok ? "PASS" : "FAIL", + transpose_ok ? "PASS" : "FAIL", direct_ok ? "PASS" : "FAIL"); + return roundtrip_ok && transpose_ok && direct_ok; +} + +bool run_strided_case(int n, int k, int source_stride, size_t source_offset) { + const size_t source_count = source_offset + (size_t)(n - 1) * source_stride + k; + const size_t output_count = (size_t)n * k; + std::vector source(source_count); + std::vector output(output_count); + fill_random(source, (unsigned)(n * 17 + k * 13 + source_stride)); + + void* memory = alloc_buffer(BufferB::required_size(n, k)); + BufferB packed(n, k, memory); + from_mat_strided(packed, source.data() + source_offset, source_stride); + to_mat(packed, output.data()); + + bool passed = true; + for (int row = 0; row < n && passed; row++) { + for (int column = 0; column < k; column++) { + if (source[source_offset + (size_t)row * source_stride + column].bits != output[(size_t)row * k + column].bits) { + passed = false; + break; + } + } + } + + std::free(memory); + std::printf("raw BF16 strided repack %dx%d stride=%d offset=%zu: %s\n", n, k, source_stride, source_offset, + passed ? "PASS" : "FAIL"); + return passed; +} + +} // namespace + +int main() { + bool passed = true; + passed = run_case(64, 64) && passed; + passed = run_case(768, 2048) && passed; + passed = run_case(2048, 768) && passed; + passed = run_strided_case(64, 64, 64, 32 * 64) && passed; + passed = run_strided_case(64, 64, 96, 17) && passed; + return passed ? 0 : 1; +} diff --git a/kt-kernel/operators/common.hpp b/kt-kernel/operators/common.hpp index 86f3464c6..5c39c00d8 100644 --- a/kt-kernel/operators/common.hpp +++ b/kt-kernel/operators/common.hpp @@ -351,6 +351,15 @@ struct MOESFTConfig : public GeneralMOEConfig { void* down_lora_a = nullptr; // [expert_num, lora_rank, intermediate_size] void* down_lora_b = nullptr; // [expert_num, hidden_size, lora_rank] + // Full weight gradient configuration + bool full_weight_grad = false; + + // Base weight gradient buffer pointers (directly pointing to Python tensor memory, zero-copy) + // Only used when full_weight_grad == true + void* grad_gate_proj = nullptr; // [expert_num, intermediate_size, hidden_size] + void* grad_up_proj = nullptr; // [expert_num, intermediate_size, hidden_size] + void* grad_down_proj = nullptr; // [expert_num, hidden_size, intermediate_size] + MOESFTConfig() : GeneralMOEConfig() {} MOESFTConfig(int expert_num, int routed_expert_num, int hidden_size, int intermediate_size) diff --git a/kt-kernel/operators/moe-sft-tp.hpp b/kt-kernel/operators/moe-sft-tp.hpp index 4e8da6ad0..1468fe5b9 100644 --- a/kt-kernel/operators/moe-sft-tp.hpp +++ b/kt-kernel/operators/moe-sft-tp.hpp @@ -26,6 +26,7 @@ #include "amx/la/amx.hpp" #include "moe-tp.hpp" +#include "sft_profile.hpp" struct TPBf16Stats { double abs_mean = 0.0; @@ -201,6 +202,7 @@ class TP_MOE_SFT : public TP_MOE { // Async backward repack state (Phase 2: overlap repack with GPU attention backward) std::thread repack_thread_; std::atomic repack_in_flight_{false}; + SFTProfiler profiler_; // Per-instance references to shared per-TP backward temporary pools. std::vector backward_temp_pools_; @@ -215,6 +217,14 @@ class TP_MOE_SFT : public TP_MOE { std::vector part_grad_input_; std::vector part_grad_weights_; + // Full-FT dWeight outputs persist across calls. Active experts are overwritten in full; + // only experts that became inactive need to be cleared after the first invocation. + std::vector last_base_grad_active_experts_; + bool base_grad_outputs_initialized_ = false; + void* last_grad_gate_proj_ = nullptr; + void* last_grad_up_proj_ = nullptr; + void* last_grad_down_proj_ = nullptr; + public: TP_MOE_SFT(const MOESFTConfig& config) : Base(static_cast(config)), sft_config(config) { printf("Creating TP_MOE_SFT layer %d\n", config.layer_idx); @@ -229,6 +239,11 @@ class TP_MOE_SFT : public TP_MOE { part_grad_input_.assign(tp_count, nullptr); part_grad_weights_.assign(tp_count, nullptr); + // TP_MOE stores GeneralMOEConfig and slices off SFT-only fields. + for (int i = 0; i < tp_count; i++) { + tps[i]->set_full_weight_grad(config.full_weight_grad); + } + if constexpr (!kSkipLoRA) { // Bug #16 fix: TP_MOE base class uses GeneralMOEConfig (object slicing) which loses // LoRA pointers. We need to propagate LoRA pointers to all NUMA node instances. @@ -245,6 +260,22 @@ class TP_MOE_SFT : public TP_MOE { } } + std::map get_profile_stats(bool reset_after = false) { + std::map stats; + stats["layer_idx"] = static_cast(config.layer_idx); + stats["tp_count"] = static_cast(tp_count); + profiler_.append(stats, "wrapper.", reset_after); + for (int i = 0; i < tp_count; ++i) { + tps[i]->append_profile_stats(stats, "tp." + std::to_string(i) + ".", reset_after); + } + return stats; + } + + void reset_profile_stats() { + profiler_.reset(); + for (int i = 0; i < tp_count; ++i) tps[i]->reset_profile_stats(); + } + /** * @brief Load weights on all NUMA nodes with TP partitioning. * @@ -254,6 +285,7 @@ class TP_MOE_SFT : public TP_MOE { * resulting in 2x the expected output after merge. */ void load_weights() override { + SFTProfileScope profile_scope(profiler_, SFTProfileStage::BaseWeightReload); auto pool = config.pool; const uint64_t* physical_to_logical_map = (const uint64_t*)config.physical_to_logical_map; @@ -345,71 +377,111 @@ class TP_MOE_SFT : public TP_MOE { throw std::runtime_error("K2 pre-quantized mode does not support TP > 1 yet"); } } else if (config.gate_proj != nullptr) { - printf("TP_MOE_SFT: From BF16 with partitioning\n"); + if constexpr (T::kSupportsDirectBf16Reload) { + std::vector intermediate_offsets(tp_count); + int intermediate_offset = 0; + for (int i = 0; i < tp_count; ++i) { + const auto& tpc = tp_configs[i]; + if (tpc.hidden_size != config.hidden_size || tpc.expert_num != config.expert_num) { + throw std::runtime_error("incompatible TP config for direct BF16 reload"); + } + intermediate_offsets[i] = intermediate_offset; + intermediate_offset += tpc.intermediate_size; + tps[i]->set_physical_to_logical_map(config.physical_to_logical_map); + } + if (intermediate_offset != config.intermediate_size) { + throw std::runtime_error("TP intermediate slices do not cover the full BF16 weight"); + } - // Temporary storage for partitioned weights - std::vector temp_gate(tp_count); - std::vector temp_up(tp_count); - std::vector temp_down(tp_count); + auto reload_stage_start = profiler_.start(); + pool->dispense_backend()->do_numa_job([&, this](int numa_id) { + tps[numa_id]->load_forward_weights_from_full_bf16(config.gate_proj, config.up_proj, config.down_proj, + config.intermediate_size, intermediate_offsets[numa_id]); + }); + profiler_.record(SFTProfileStage::BaseWeightReloadDirectPack, reload_stage_start); - // Step 1: For each NUMA, allocate and copy partitioned weights - for (int i = 0; i < tp_count; i++) { - // Use tp_configs[i] instead of tps[i]->config_ (which is protected) - auto& tpc = tp_configs[i]; - size_t gate_up_elcount = (size_t)tpc.intermediate_size * tpc.hidden_size; - - // Allocate partitioned weight space - temp_gate[i] = new ggml_bf16_t[tpc.expert_num * gate_up_elcount]; - temp_up[i] = new ggml_bf16_t[tpc.expert_num * gate_up_elcount]; - temp_down[i] = new ggml_bf16_t[tpc.expert_num * gate_up_elcount]; - - // Copy partitioned weights - pool->get_subpool(i)->do_work_stealing_job( - tpc.expert_num, nullptr, - [&, i, gate_up_elcount](int expert_id_) { - size_t expert_id = expert_map(physical_to_logical_map, expert_id_); - - // gate_proj/up_proj: [intermediate_size, hidden_size] - contiguous block slice - memcpy(temp_gate[i] + expert_id * gate_up_elcount, - (ggml_bf16_t*)config.gate_proj + expert_id * config.intermediate_size * config.hidden_size + - i * gate_up_elcount, - sizeof(ggml_bf16_t) * gate_up_elcount); - - memcpy(temp_up[i] + expert_id * gate_up_elcount, - (ggml_bf16_t*)config.up_proj + expert_id * config.intermediate_size * config.hidden_size + - i * gate_up_elcount, - sizeof(ggml_bf16_t) * gate_up_elcount); - - // down_proj: [hidden_size, intermediate_size] - row-wise slice - for (size_t col = 0; col < config.hidden_size; col++) { - memcpy(temp_down[i] + expert_id * tpc.hidden_size * tpc.intermediate_size + col * tpc.intermediate_size, - (ggml_bf16_t*)config.down_proj + expert_id * config.intermediate_size * config.hidden_size + - col * config.intermediate_size + i * tpc.intermediate_size, - sizeof(ggml_bf16_t) * tpc.intermediate_size); - } - }, - nullptr); - } + if (!config.share_backward_bb) { + reload_stage_start = profiler_.start(); + pool->dispense_backend()->do_numa_job( + [this](int numa_id) { tps[numa_id]->prepare_backward_weights_from_forward(); }); + profiler_.record(SFTProfileStage::BaseWeightReloadBackwardPack, reload_stage_start); + } + } else { + // printf("TP_MOE_SFT: From BF16 with partitioning\n"); + auto reload_stage_start = profiler_.start(); - // Step 2: Set weight pointers BEFORE load_weights (Bug #24 fix) - for (int i = 0; i < tp_count; i++) { - tps[i]->set_physical_to_logical_map(config.physical_to_logical_map); - tps[i]->set_weight_pointers_for_forward(temp_gate[i], temp_up[i], temp_down[i]); - } + // Temporary storage for partitioned weights + std::vector temp_gate(tp_count); + std::vector temp_up(tp_count); + std::vector temp_down(tp_count); - pool->dispense_backend()->do_numa_job([this](int numa_id) { tps[numa_id]->load_weights(); }); + // Step 1: For each NUMA, allocate and copy partitioned weights + for (int i = 0; i < tp_count; i++) { + // Use tp_configs[i] instead of tps[i]->config_ (which is protected) + auto& tpc = tp_configs[i]; + size_t gate_up_elcount = (size_t)tpc.intermediate_size * tpc.hidden_size; - // Step 3: Prepare backward weights (this also clears weight pointers) - for (int i = 0; i < tp_count; i++) { - if (!config.share_backward_bb) { - tps[i]->prepare_bwd(temp_gate[i], temp_up[i], temp_down[i]); + // Allocate partitioned weight space + temp_gate[i] = new ggml_bf16_t[tpc.expert_num * gate_up_elcount]; + temp_up[i] = new ggml_bf16_t[tpc.expert_num * gate_up_elcount]; + temp_down[i] = new ggml_bf16_t[tpc.expert_num * gate_up_elcount]; + + // Copy partitioned weights + pool->get_subpool(i)->do_work_stealing_job( + tpc.expert_num, nullptr, + [&, i, gate_up_elcount](int expert_id_) { + size_t expert_id = expert_map(physical_to_logical_map, expert_id_); + + // gate_proj/up_proj: [intermediate_size, hidden_size] - contiguous block slice + memcpy(temp_gate[i] + expert_id * gate_up_elcount, + (ggml_bf16_t*)config.gate_proj + expert_id * config.intermediate_size * config.hidden_size + + i * gate_up_elcount, + sizeof(ggml_bf16_t) * gate_up_elcount); + + memcpy(temp_up[i] + expert_id * gate_up_elcount, + (ggml_bf16_t*)config.up_proj + expert_id * config.intermediate_size * config.hidden_size + + i * gate_up_elcount, + sizeof(ggml_bf16_t) * gate_up_elcount); + + // down_proj: [hidden_size, intermediate_size] - row-wise slice + for (size_t col = 0; col < config.hidden_size; col++) { + memcpy( + temp_down[i] + expert_id * tpc.hidden_size * tpc.intermediate_size + col * tpc.intermediate_size, + (ggml_bf16_t*)config.down_proj + expert_id * config.intermediate_size * config.hidden_size + + col * config.intermediate_size + i * tpc.intermediate_size, + sizeof(ggml_bf16_t) * tpc.intermediate_size); + } + }, + nullptr); } - } + profiler_.record(SFTProfileStage::BaseWeightReloadPartition, reload_stage_start); - for (int i = 0; i < tp_count; i++) { - delete[] (temp_gate[i]); - delete[] (temp_up[i]); - delete[] (temp_down[i]); + // Step 2: Set weight pointers BEFORE load_weights (Bug #24 fix) + reload_stage_start = profiler_.start(); + for (int i = 0; i < tp_count; i++) { + tps[i]->set_physical_to_logical_map(config.physical_to_logical_map); + tps[i]->set_weight_pointers_for_forward(temp_gate[i], temp_up[i], temp_down[i]); + } + + pool->dispense_backend()->do_numa_job([this](int numa_id) { tps[numa_id]->load_weights(); }); + profiler_.record(SFTProfileStage::BaseWeightReloadForwardPack, reload_stage_start); + + // Step 3: Prepare backward weights (this also clears weight pointers) + reload_stage_start = profiler_.start(); + for (int i = 0; i < tp_count; i++) { + if (!config.share_backward_bb) { + tps[i]->prepare_bwd(temp_gate[i], temp_up[i], temp_down[i]); + } + } + profiler_.record(SFTProfileStage::BaseWeightReloadBackwardPack, reload_stage_start); + + reload_stage_start = profiler_.start(); + for (int i = 0; i < tp_count; i++) { + delete[] (temp_gate[i]); + delete[] (temp_up[i]); + delete[] (temp_down[i]); + } + profiler_.record(SFTProfileStage::BaseWeightReloadCleanup, reload_stage_start); } } else { // Other loading methods (from loader or file) @@ -490,6 +562,7 @@ class TP_MOE_SFT : public TP_MOE { void forward_sft(int* qlen_ptr, int k, const int64_t* expert_ids, const float* weights, const void* input, void* output, bool save_for_backward) { + SFTProfileScope total_scope(profiler_, SFTProfileStage::TpFwdTotal); if (weights_loaded == false) [[unlikely]] { throw std::runtime_error("Weights not loaded"); } @@ -503,10 +576,12 @@ class TP_MOE_SFT : public TP_MOE { } // Run forward on each NUMA node + auto stage_start = profiler_.start(); pool->dispense_backend()->do_numa_job([this, qlen, k, expert_ids, input, weights, save_for_backward](int numa_id) { tps[numa_id]->forward_sft(qlen, k, expert_ids, weights, input, this->local_output_numa[numa_id], save_for_backward); }); + profiler_.record(SFTProfileStage::TpFwdNumaCompute, stage_start); // // Collect per-thread timing from all NUMA subpools // for (int i = 0; i < tp_count; i++) { @@ -515,7 +590,9 @@ class TP_MOE_SFT : public TP_MOE { // // Print per-thread forward timing // Merge results from all NUMA nodes + stage_start = profiler_.start(); this->merge_results(qlen, output); + profiler_.record(SFTProfileStage::TpFwdMerge, stage_start); pool->dispense_backend()->do_numa_job([&](int numa_id) {}); } @@ -549,7 +626,10 @@ class TP_MOE_SFT : public TP_MOE { */ void backward(const void* grad_output, void* grad_input, void* grad_gate_lora_a, void* grad_gate_lora_b, void* grad_up_lora_a, void* grad_up_lora_b, void* grad_down_lora_a, void* grad_down_lora_b, - void* grad_weights) { + void* grad_weights, void* grad_gate_proj = nullptr, void* grad_up_proj = nullptr, + void* grad_down_proj = nullptr) { + SFTProfileScope total_scope(profiler_, SFTProfileStage::TpBwdTotal); + auto stage_start = profiler_.start(); auto pool = config.pool; // Get full intermediate_size (before TP partitioning) @@ -561,6 +641,8 @@ class TP_MOE_SFT : public TP_MOE { int k = sft_config.num_experts_per_tok; const bool need_grad_weights = (grad_weights != nullptr); + const bool need_base_weight_grad = sft_config.full_weight_grad && grad_gate_proj != nullptr && + grad_up_proj != nullptr && grad_down_proj != nullptr; // SkipLoRA: zero out lora_rank to skip all LoRA buffer allocations if constexpr (kSkipLoRA) lora_rank = 0; @@ -571,6 +653,8 @@ class TP_MOE_SFT : public TP_MOE { if (active_count > 0) { std::memcpy(active_expert_map.data(), tps[0]->get_cache_expert_id_map(), active_count * sizeof(int)); } + profiler_.record_workload(static_cast(qlen), static_cast(qlen) * k, + static_cast(active_count)); // ===================================================================== // Allocate per-TP temporary buffers. @@ -627,8 +711,7 @@ class TP_MOE_SFT : public TP_MOE { clear_bytes[i] = offset; } - // Parallel memset: zero only per-TP sparse partials and per-TP grad_input/grad_weights partials. - // The caller is responsible for passing zero-initialized final grad tensors. + // Parallel memset for per-TP partials and final base-weight gradients. struct ClearSeg { uint8_t* ptr; size_t len; @@ -649,12 +732,54 @@ class TP_MOE_SFT : public TP_MOE { } } + size_t base_grad_clear_bytes = 0; + const bool can_selectively_clear_base_grad = + need_base_weight_grad && T::kSupportsDirectBf16Reload && sft_config.lora_rank == 0 && + config.num_gpu_experts == 0; + if (need_base_weight_grad) { + const size_t base_grad_bytes = (size_t)expert_num * full_intermediate_size * hidden_size * sizeof(ggml_bf16_t); + const size_t expert_grad_bytes = (size_t)full_intermediate_size * hidden_size * sizeof(ggml_bf16_t); + auto append_clear_segments = [&](void* ptr, size_t offset, size_t bytes) { + auto* base = static_cast(ptr) + offset; + for (size_t off = 0; off < bytes; off += kChunkBytes) { + size_t len = std::min(kChunkBytes, bytes - off); + clear_segs.push_back(ClearSeg{base + off, len}); + } + }; + + const bool pointers_unchanged = last_grad_gate_proj_ == grad_gate_proj && last_grad_up_proj_ == grad_up_proj && + last_grad_down_proj_ == grad_down_proj; + if (!can_selectively_clear_base_grad || !base_grad_outputs_initialized_ || !pointers_unchanged) { + append_clear_segments(grad_gate_proj, 0, base_grad_bytes); + append_clear_segments(grad_up_proj, 0, base_grad_bytes); + append_clear_segments(grad_down_proj, 0, base_grad_bytes); + base_grad_clear_bytes = 3 * base_grad_bytes; + } else { + std::vector active_mask(expert_num, 0); + for (int expert_idx : active_expert_map) { + if (expert_idx >= 0 && expert_idx < expert_num) active_mask[expert_idx] = 1; + } + for (int expert_idx : last_base_grad_active_experts_) { + if (expert_idx < 0 || expert_idx >= expert_num || active_mask[expert_idx]) continue; + const size_t expert_offset = static_cast(expert_idx) * expert_grad_bytes; + append_clear_segments(grad_gate_proj, expert_offset, expert_grad_bytes); + append_clear_segments(grad_up_proj, expert_offset, expert_grad_bytes); + append_clear_segments(grad_down_proj, expert_offset, expert_grad_bytes); + base_grad_clear_bytes += 3 * expert_grad_bytes; + } + } + } + pool->do_work_stealing_job((int)clear_segs.size(), nullptr, [&](int seg_idx) { const auto& seg = clear_segs[(size_t)seg_idx]; std::memset(seg.ptr, 0, seg.len); }, nullptr); + size_t partial_clear_bytes = 0; + for (size_t bytes : clear_bytes) partial_clear_bytes += bytes; + profiler_.record_bytes(SFTProfileStage::TpBwdBufferClear, partial_clear_bytes + base_grad_clear_bytes); + profiler_.record(SFTProfileStage::TpBwdBufferClear, stage_start); // Compute TP-slice pointers for copy-type direct writes // Each TP writes to its own I-slice of the final output tensor @@ -664,6 +789,9 @@ class TP_MOE_SFT : public TP_MOE { std::vector tp_fp32_down_b(tp_count); std::vector tp_fp32_gate_a(tp_count); std::vector tp_fp32_up_a(tp_count); + std::vector tp_grad_gate_proj(tp_count, nullptr); + std::vector tp_grad_up_proj(tp_count, nullptr); + std::vector tp_grad_down_proj(tp_count, nullptr); if constexpr (!kSkipLoRA) { int tp_offset = 0; @@ -682,7 +810,24 @@ class TP_MOE_SFT : public TP_MOE { } } + int tp_offset = 0; + for (int i = 0; i < tp_count; i++) { + if (need_base_weight_grad) { + tp_grad_gate_proj[i] = static_cast(grad_gate_proj) + (size_t)tp_offset * hidden_size; + tp_grad_up_proj[i] = static_cast(grad_up_proj) + (size_t)tp_offset * hidden_size; + tp_grad_down_proj[i] = static_cast(grad_down_proj) + tp_offset; + } + tp_offset += tp_configs[i].intermediate_size; + } + if (tp_offset != full_intermediate_size) { + throw std::runtime_error("TP intermediate_size slices do not cover the full intermediate_size"); + } + + // A failed dWeight computation must force a full clear on the next invocation. + if (can_selectively_clear_base_grad) base_grad_outputs_initialized_ = false; + // Run backward on each NUMA node + stage_start = profiler_.start(); pool->dispense_backend()->do_numa_job([&](int numa_id) { tps[numa_id]->backward(grad_output, part_grad_input_[numa_id], // reduce-type: BF16 pointer unused (FP32 sparse used instead) @@ -693,8 +838,17 @@ class TP_MOE_SFT : public TP_MOE { tp_down_a_ptr[numa_id], /* copy-type: direct write */ nullptr, /* grad_down_lora_b — unused, FP32 path below */ part_grad_weights_[numa_id], full_intermediate_size, tp_fp32_down_b[numa_id], - tp_fp32_gate_a[numa_id], tp_fp32_up_a[numa_id]); + tp_fp32_gate_a[numa_id], tp_fp32_up_a[numa_id], tp_grad_gate_proj[numa_id], + tp_grad_up_proj[numa_id], tp_grad_down_proj[numa_id]); }); + profiler_.record(SFTProfileStage::TpBwdNumaCompute, stage_start); + if (can_selectively_clear_base_grad) { + last_base_grad_active_experts_ = active_expert_map; + last_grad_gate_proj_ = grad_gate_proj; + last_grad_up_proj_ = grad_up_proj; + last_grad_down_proj_ = grad_down_proj; + base_grad_outputs_initialized_ = true; + } // // Collect per-thread timing from all NUMA subpools // for (int i = 0; i < tp_count; i++) { @@ -720,6 +874,7 @@ class TP_MOE_SFT : public TP_MOE { // } // Bug #22 fix: Merge grad_input from all NUMA nodes (sum them together) + stage_start = profiler_.start(); { auto* out = (ggml_bf16_t*)grad_input; pool->do_work_stealing_job( @@ -766,9 +921,11 @@ class TP_MOE_SFT : public TP_MOE { }, nullptr); } + profiler_.record(SFTProfileStage::TpBwdGradInputMerge, stage_start); // Merge reduce-type LoRA gradients: sparse FP32 sum across TPs → BF16 final output // Copy-type grads (gate/up_lora_b, down_lora_a) were written directly — no merge needed. + stage_start = profiler_.start(); if constexpr (!kSkipLoRA) { // Sparse merge for gate_lora_a, up_lora_a: [active_count, r, H] FP32 → [E, r, H] BF16 { @@ -842,9 +999,11 @@ class TP_MOE_SFT : public TP_MOE { nullptr); } } // if constexpr (!kSkipLoRA) + profiler_.record(SFTProfileStage::TpBwdLoraMerge, stage_start); // Merge grad_weights from all NUMA nodes (sum them together) // Each NUMA computes partial grad_weights based on its down_output partition + stage_start = profiler_.start(); if (grad_weights != nullptr) { float* out_grad_weights = (float*)grad_weights; const size_t total = (size_t)qlen * (size_t)k; @@ -880,6 +1039,7 @@ class TP_MOE_SFT : public TP_MOE { }, nullptr); } + profiler_.record(SFTProfileStage::TpBwdRouterGradMerge, stage_start); pool->dispense_backend()->do_numa_job([&](int numa_id) {}); } @@ -889,10 +1049,11 @@ class TP_MOE_SFT : public TP_MOE { */ void backward_binding(intptr_t grad_output, intptr_t grad_input, intptr_t grad_gate_lora_a, intptr_t grad_gate_lora_b, intptr_t grad_up_lora_a, intptr_t grad_up_lora_b, intptr_t grad_down_lora_a, - intptr_t grad_down_lora_b, intptr_t grad_weights) { + intptr_t grad_down_lora_b, intptr_t grad_weights, intptr_t grad_gate_proj, + intptr_t grad_up_proj, intptr_t grad_down_proj) { backward((const void*)grad_output, (void*)grad_input, (void*)grad_gate_lora_a, (void*)grad_gate_lora_b, (void*)grad_up_lora_a, (void*)grad_up_lora_b, (void*)grad_down_lora_a, (void*)grad_down_lora_b, - (void*)grad_weights); + (void*)grad_weights, (void*)grad_gate_proj, (void*)grad_up_proj, (void*)grad_down_proj); } /** @@ -1068,6 +1229,7 @@ class TP_MOE_SFT : public TP_MOE { repack_in_flight_.store(true, std::memory_order_release); repack_thread_ = std::thread([this]() { + SFTProfileScope profile_scope(profiler_, SFTProfileStage::BackwardRepack); config.pool->dispense_backend()->do_numa_job( [this](int numa_id) { tps[numa_id]->prepare_backward_bb_for_async(); }); repack_in_flight_.store(false, std::memory_order_release); @@ -1080,6 +1242,7 @@ class TP_MOE_SFT : public TP_MOE { */ void wait_backward_repack() { if (repack_thread_.joinable()) { + SFTProfileScope profile_scope(profiler_, SFTProfileStage::BackwardRepackWait); repack_thread_.join(); } } @@ -1099,6 +1262,22 @@ class TP_MOE_SFT : public TP_MOE { update_lora_weights((void*)gate_lora_a, (void*)gate_lora_b, (void*)up_lora_a, (void*)up_lora_b, (void*)down_lora_a, (void*)down_lora_b); } + + /** + * @brief Update base weight BF16 pointers for reload_base_weights (full mode training). + * + * After calling this, call load_weights_task() to re-quantize BF16->AMX + * and update the C++ kernel's internal quantized buffers. + * This avoids creating a new C++ MOE object (~0.6s/layer for quantization + * vs ~1.9s/layer for full object recreation). + */ + void set_base_weight_pointers(void* gate, void* up, void* down) { + config.gate_proj = gate; + config.up_proj = up; + config.down_proj = down; + // Mark that weights need re-loading (partitioning + quantization) + weights_loaded = false; + } }; #endif // CPUINFER_OPERATOR_MOE_SFT_TP_HPP diff --git a/kt-kernel/operators/sft_profile.hpp b/kt-kernel/operators/sft_profile.hpp new file mode 100644 index 000000000..22dd6b1cf --- /dev/null +++ b/kt-kernel/operators/sft_profile.hpp @@ -0,0 +1,259 @@ +// Lightweight staged profiling for KT SFT MoE operators. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef CPUINFER_OPERATOR_SFT_PROFILE_HPP +#define CPUINFER_OPERATOR_SFT_PROFILE_HPP + +#include +#include +#include +#include +#include +#include +#include +#include + +enum class SFTProfileStage : uint8_t { + // NUMA-local forward stages. + FwdTotal, + FwdInitialTotal, + FwdRecomputeTotal, + FwdSetup, + FwdRoute, + FwdBufferSetup, + FwdInputScatter, + FwdInputPack, + FwdCacheMetadata, + FwdCacheInput, + FwdGateUpBase, + FwdGateUpLora, + FwdCacheGateUp, + FwdActivation, + FwdCacheIntermediate, + FwdDownPack, + FwdDownBase, + FwdDownLora, + FwdCacheDown, + FwdWeightedMerge, + + // NUMA-local backward stages. + BwdTotal, + BwdSetup, + BwdCacheRestore, + BwdDownTotal, + BwdDownSetup, + BwdDownScatter, + BwdDownBaseDx, + BwdDownLora, + BwdActivation, + BwdGateUpTotal, + BwdGateUpSetup, + BwdGateUpBaseDx, + BwdGateUpLora, + BwdRouterGrad, + BwdBaseWeightGrad, + BwdBaseWeightGradOffsets, + BwdBaseWeightGradPanelPack, + BwdBaseWeightGradPanelInput, + BwdBaseWeightGradPanelGradOutput, + BwdBaseWeightGradMatMat, + BwdBaseWeightGradPackA, + BwdBaseWeightGradPackB, + BwdBaseWeightGradKernelGateUp, + BwdBaseWeightGradKernelDown, + BwdBaseWeightGradAmx, + BwdBaseWeightGradZero, + BwdBaseWeightGradGateUp, + BwdBaseWeightGradDown, + BwdBaseWeightGradStore, + + // TP wrapper and weight-layout stages. + TpFwdTotal, + TpFwdNumaCompute, + TpFwdMerge, + TpBwdTotal, + TpBwdBufferClear, + TpBwdNumaCompute, + TpBwdGradInputMerge, + TpBwdLoraMerge, + TpBwdRouterGradMerge, + BackwardRepack, + BackwardRepackWait, + BaseWeightReload, + BaseWeightReloadPartition, + BaseWeightReloadForwardPack, + BaseWeightReloadDirectPack, + BaseWeightReloadBackwardPack, + BaseWeightReloadCleanup, + + Count, +}; + +inline constexpr std::array(SFTProfileStage::Count)> kSFTProfileStageNames = { + "forward.total", + "forward.initial_total", + "forward.recompute_total", + "forward.setup", + "forward.route", + "forward.buffer_setup", + "forward.input_scatter", + "forward.input_pack", + "forward.cache_metadata", + "forward.cache_input", + "forward.gate_up_base", + "forward.gate_up_lora", + "forward.cache_gate_up", + "forward.activation", + "forward.cache_intermediate", + "forward.down_pack", + "forward.down_base", + "forward.down_lora", + "forward.cache_down", + "forward.weighted_merge", + "backward.total", + "backward.setup", + "backward.cache_restore", + "backward.down.total", + "backward.down.setup", + "backward.down.scatter", + "backward.down.base_dx", + "backward.down.lora", + "backward.activation", + "backward.gate_up.total", + "backward.gate_up.setup", + "backward.gate_up.base_dx", + "backward.gate_up.lora", + "backward.router_grad", + "backward.base_weight_grad", + "backward.base_weight_grad.offsets", + "backward.base_weight_grad.panel_pack", + "backward.base_weight_grad.worker_cpu.panel_input", + "backward.base_weight_grad.worker_cpu.panel_grad_output", + "backward.base_weight_grad.matmat", + "backward.base_weight_grad.worker_cpu.pack_a", + "backward.base_weight_grad.worker_cpu.pack_b", + "backward.base_weight_grad.worker_cpu.kernel_gate_up", + "backward.base_weight_grad.worker_cpu.kernel_down", + "backward.base_weight_grad.amx", + "backward.base_weight_grad.zero", + "backward.base_weight_grad.gate_up", + "backward.base_weight_grad.down", + "backward.base_weight_grad.worker_cpu.store", + "tp.forward.total", + "tp.forward.numa_compute", + "tp.forward.merge", + "tp.backward.total", + "tp.backward.buffer_clear", + "tp.backward.numa_compute", + "tp.backward.grad_input_merge", + "tp.backward.lora_merge", + "tp.backward.router_grad_merge", + "weights.backward_repack", + "weights.backward_repack_wait", + "weights.base_reload", + "weights.base_reload.partition", + "weights.base_reload.forward_pack", + "weights.base_reload.direct_pack", + "weights.base_reload.backward_pack", + "weights.base_reload.cleanup", +}; + +inline bool sft_profile_enabled_from_env() { + const char* value = std::getenv("KT_SFT_PROFILE"); + return value != nullptr && value[0] != '\0' && std::strcmp(value, "0") != 0 && std::strcmp(value, "false") != 0 && + std::strcmp(value, "False") != 0; +} + +class SFTProfiler { + public: + using Clock = std::chrono::steady_clock; + using TimePoint = Clock::time_point; + + explicit SFTProfiler(bool enabled = sft_profile_enabled_from_env()) : enabled_(enabled) { reset(); } + + bool enabled() const { return enabled_; } + + TimePoint start() const { return enabled_ ? Clock::now() : TimePoint{}; } + + void record(SFTProfileStage stage, TimePoint start) { + if (!enabled_) return; + const auto elapsed = std::chrono::duration_cast(Clock::now() - start).count(); + record_ns(stage, static_cast(elapsed)); + } + + void record_ns(SFTProfileStage stage, uint64_t elapsed_ns, uint64_t calls = 1) { + if (!enabled_) return; + const size_t idx = static_cast(stage); + total_ns_[idx].fetch_add(elapsed_ns, std::memory_order_relaxed); + calls_[idx].fetch_add(calls, std::memory_order_relaxed); + } + + void record_bytes(SFTProfileStage stage, uint64_t bytes) { + if (!enabled_) return; + bytes_[static_cast(stage)].fetch_add(bytes, std::memory_order_relaxed); + } + + void record_workload(uint64_t tokens, uint64_t routed_rows, uint64_t active_experts) { + if (!enabled_) return; + tokens_.fetch_add(tokens, std::memory_order_relaxed); + routed_rows_.fetch_add(routed_rows, std::memory_order_relaxed); + active_experts_.fetch_add(active_experts, std::memory_order_relaxed); + workloads_.fetch_add(1, std::memory_order_relaxed); + } + + void append(std::map& out, const std::string& prefix, bool reset_after = false) { + out[prefix + "enabled"] = enabled_ ? 1.0 : 0.0; + out[prefix + "workloads"] = static_cast(load_or_exchange(workloads_, reset_after)); + out[prefix + "tokens"] = static_cast(load_or_exchange(tokens_, reset_after)); + out[prefix + "routed_rows"] = static_cast(load_or_exchange(routed_rows_, reset_after)); + out[prefix + "active_experts"] = static_cast(load_or_exchange(active_experts_, reset_after)); + for (size_t i = 0; i < static_cast(SFTProfileStage::Count); ++i) { + const std::string stage_prefix = prefix + kSFTProfileStageNames[i] + "."; + out[stage_prefix + "total_ns"] = static_cast(load_or_exchange(total_ns_[i], reset_after)); + out[stage_prefix + "calls"] = static_cast(load_or_exchange(calls_[i], reset_after)); + out[stage_prefix + "bytes"] = static_cast(load_or_exchange(bytes_[i], reset_after)); + } + } + + void reset() { + for (auto& value : total_ns_) value.store(0, std::memory_order_relaxed); + for (auto& value : calls_) value.store(0, std::memory_order_relaxed); + for (auto& value : bytes_) value.store(0, std::memory_order_relaxed); + workloads_.store(0, std::memory_order_relaxed); + tokens_.store(0, std::memory_order_relaxed); + routed_rows_.store(0, std::memory_order_relaxed); + active_experts_.store(0, std::memory_order_relaxed); + } + + private: + static uint64_t load_or_exchange(std::atomic& value, bool reset_after) { + return reset_after ? value.exchange(0, std::memory_order_relaxed) : value.load(std::memory_order_relaxed); + } + + bool enabled_; + std::array, static_cast(SFTProfileStage::Count)> total_ns_{}; + std::array, static_cast(SFTProfileStage::Count)> calls_{}; + std::array, static_cast(SFTProfileStage::Count)> bytes_{}; + std::atomic workloads_{0}; + std::atomic tokens_{0}; + std::atomic routed_rows_{0}; + std::atomic active_experts_{0}; +}; + +class SFTProfileScope { + public: + SFTProfileScope(SFTProfiler& profiler, SFTProfileStage stage) + : profiler_(profiler), stage_(stage), start_(profiler.start()) {} + + ~SFTProfileScope() { profiler_.record(stage_, start_); } + + SFTProfileScope(const SFTProfileScope&) = delete; + SFTProfileScope& operator=(const SFTProfileScope&) = delete; + + private: + SFTProfiler& profiler_; + SFTProfileStage stage_; + SFTProfiler::TimePoint start_; +}; + +#endif // CPUINFER_OPERATOR_SFT_PROFILE_HPP diff --git a/kt-kernel/python/experts.py b/kt-kernel/python/experts.py index 318fe3170..d31e42801 100644 --- a/kt-kernel/python/experts.py +++ b/kt-kernel/python/experts.py @@ -144,6 +144,8 @@ def __new__( # Quantization config (for K-Group SFT methods) group_size: int = 128, zero_point: bool = True, + # Full weight gradient mode (for full fine-tuning without LoRA) + full_weight_grad: bool = False, # V4-Flash 2604B SwiGLU clamp limit. 0.0 = disabled (default for # every dtype except DSV4-2604B routed experts, which set this to # 10.0 to match trtllm gemm1_clamp_limit / deep_gemm @@ -252,6 +254,7 @@ def __new__( max_cache_depth=max_cache_depth, group_size=group_size, zero_point=zero_point, + full_weight_grad=full_weight_grad, ) # Forward static methods to the base class @@ -297,6 +300,7 @@ def clear_sft_buffer_cache(): to reset the buffer state or free memory during SFT. """ from .sft.base import KExpertsSFTBuffer + KExpertsSFTBuffer.clear_cache() @@ -400,6 +404,7 @@ def _create_sft_wrapper( max_cache_depth: int, group_size: int, zero_point: bool, + full_weight_grad: bool = False, ): """ Create an SFT wrapper based on the method. @@ -430,4 +435,5 @@ def _create_sft_wrapper( method=method, group_size=group_size, zero_point=zero_point, + full_weight_grad=full_weight_grad, ) diff --git a/kt-kernel/python/sft/__init__.py b/kt-kernel/python/sft/__init__.py index 7cab43bd2..b17ed6873 100644 --- a/kt-kernel/python/sft/__init__.py +++ b/kt-kernel/python/sft/__init__.py @@ -14,8 +14,15 @@ from .base import BaseSFTMoEWrapper, KExpertsSFTBuffer from .amx import AMXSFTMoEWrapper from .arch import ( - MOEArchConfig, get_moe_arch_config, get_moe_module, move_non_experts_to_gpu, get_expert_device, - KTAMXError, KTAMXNotAvailableError, KTAMXModelNotSupportedError, KTAMXConfigError, + MOEArchConfig, + get_moe_arch_config, + get_moe_module, + move_non_experts_to_gpu, + get_expert_device, + KTAMXError, + KTAMXNotAvailableError, + KTAMXModelNotSupportedError, + KTAMXConfigError, ) from .autograd import KTMoEFunction from .layer import KTMoELayerWrapper @@ -28,6 +35,7 @@ from .lora import ( kt_adapt_peft_lora, get_kt_lora_params, + get_kt_trainable_params, update_kt_lora_pointers, sync_kt_lora_gradients, save_lora_experts_to_adapter, @@ -44,6 +52,7 @@ get_kt_loading_kwargs, load_kt_model, ) +from .profiler import collect_kt_sft_profile, format_kt_sft_profile, reset_kt_sft_profile __all__ = [ "KTConfig", @@ -67,6 +76,7 @@ "INT8ExpertWeights", "kt_adapt_peft_lora", "get_kt_lora_params", + "get_kt_trainable_params", "update_kt_lora_pointers", "sync_kt_lora_gradients", "save_lora_experts_to_adapter", @@ -80,4 +90,7 @@ "build_kt_device_map_simplified", "get_kt_loading_kwargs", "load_kt_model", + "collect_kt_sft_profile", + "format_kt_sft_profile", + "reset_kt_sft_profile", ] diff --git a/kt-kernel/python/sft/amx.py b/kt-kernel/python/sft/amx.py index effa7dfd2..c59a8b455 100644 --- a/kt-kernel/python/sft/amx.py +++ b/kt-kernel/python/sft/amx.py @@ -9,6 +9,7 @@ from __future__ import annotations import ctypes +import logging import os import glob as _glob import torch @@ -16,6 +17,8 @@ from kt_kernel_ext.moe import MOESFTConfig +logger = logging.getLogger(__name__) + from ..utils.loader import BF16SafeTensorLoader, SafeTensorLoader try: @@ -81,6 +84,7 @@ def __init__( method: str = "AMXBF16_SFT", group_size: int = 128, zero_point: bool = True, + full_weight_grad: bool = False, ): if not _HAS_AMX_SFT_SUPPORT: raise RuntimeError( @@ -102,6 +106,7 @@ def __init__( lora_rank=lora_rank, lora_alpha=lora_alpha, max_cache_depth=max_cache_depth, + full_weight_grad=full_weight_grad, ) self.method = method @@ -140,19 +145,42 @@ def _make_backward_task(self, buffer: KExpertsSFTBuffer): return self.moe.backward_task( buffer.grad_output_cpu.data_ptr(), buffer.grad_input_cpu.data_ptr(), - 0, 0, 0, 0, 0, 0, + 0, + 0, + 0, + 0, + 0, + 0, buffer.grad_weights.data_ptr(), + 0, + 0, + 0, # grad_gate_proj, grad_up_proj, grad_down_proj ) + + # Base weight grad pointers (nullptr if not in full mode) + grad_gate_proj_ptr = ( + self.grad_gate_proj_buf.data_ptr() if self._full_weight_grad and self.grad_gate_proj_buf is not None else 0 + ) + grad_up_proj_ptr = ( + self.grad_up_proj_buf.data_ptr() if self._full_weight_grad and self.grad_up_proj_buf is not None else 0 + ) + grad_down_proj_ptr = ( + self.grad_down_proj_buf.data_ptr() if self._full_weight_grad and self.grad_down_proj_buf is not None else 0 + ) + return self.moe.backward_task( buffer.grad_output_cpu.data_ptr(), buffer.grad_input_cpu.data_ptr(), - self.grad_gate_lora_a.data_ptr(), - self.grad_gate_lora_b.data_ptr(), - self.grad_up_lora_a.data_ptr(), - self.grad_up_lora_b.data_ptr(), - self.grad_down_lora_a.data_ptr(), - self.grad_down_lora_b.data_ptr(), + self.grad_gate_lora_a.data_ptr() if self.lora_rank > 0 else 0, + self.grad_gate_lora_b.data_ptr() if self.lora_rank > 0 else 0, + self.grad_up_lora_a.data_ptr() if self.lora_rank > 0 else 0, + self.grad_up_lora_b.data_ptr() if self.lora_rank > 0 else 0, + self.grad_down_lora_a.data_ptr() if self.lora_rank > 0 else 0, + self.grad_down_lora_b.data_ptr() if self.lora_rank > 0 else 0, buffer.grad_weights.data_ptr(), + grad_gate_proj_ptr, + grad_up_proj_ptr, + grad_down_proj_ptr, ) # ========== Weight loading ========== @@ -187,6 +215,7 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None: config.layer_idx = self.layer_idx config.share_backward_bb = getattr(self, "share_backward_bb", False) config.share_cache_pool = getattr(self, "share_cache_pool", False) + config.full_weight_grad = self._full_weight_grad config.physical_to_logical_map = self._physical_to_logical_map_cpu.data_ptr() if getattr(self, "_use_kt_direct_load", False): @@ -229,6 +258,13 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None: config.quant_config.group_size = self.group_size config.quant_config.zero_point = self.zero_point + # Release old C++ MOE object before creating a new one to avoid memory leak + old_moe = getattr(self, "moe", None) + if old_moe is not None: + del old_moe + import gc + gc.collect() + self.moe = self._moe_class(config) self.cpu_infer.submit(self.moe.load_weights_task()) @@ -239,9 +275,11 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None: self.cpu_infer.sync() # Release Python-side weight tensors (C++ copied them) - self.gate_proj = None - self.up_proj = None - self.down_proj = None + # In full_weight_grad mode, keep them for nn.Parameter initialization + if not self._full_weight_grad: + self.gate_proj = None + self.up_proj = None + self.down_proj = None if getattr(self, "_bf16_gate_proj", None) is not None: self._bf16_gate_proj = None @@ -250,18 +288,34 @@ def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None: if getattr(self, "_use_projs_path", False): for attr in [ - "_gate_weights_per_numa", "_up_weights_per_numa", "_down_weights_per_numa", - "_gate_scales_per_numa", "_up_scales_per_numa", "_down_scales_per_numa", - "_gate_projs_ptrs", "_up_projs_ptrs", "_down_projs_ptrs", - "_gate_scale_ptrs", "_up_scale_ptrs", "_down_scale_ptrs", + "_gate_weights_per_numa", + "_up_weights_per_numa", + "_down_weights_per_numa", + "_gate_scales_per_numa", + "_up_scales_per_numa", + "_down_scales_per_numa", + "_gate_projs_ptrs", + "_up_projs_ptrs", + "_down_projs_ptrs", + "_gate_scale_ptrs", + "_up_scale_ptrs", + "_down_scale_ptrs", ]: setattr(self, attr, None) if getattr(self, "_has_bwd_projs", False): for attr in [ - "_gate_bwd_weights_per_numa", "_up_bwd_weights_per_numa", "_down_bwd_weights_per_numa", - "_gate_bwd_scales_per_numa", "_up_bwd_scales_per_numa", "_down_bwd_scales_per_numa", - "_gate_bwd_projs_ptrs", "_up_bwd_projs_ptrs", "_down_bwd_projs_ptrs", - "_gate_bwd_scale_ptrs", "_up_bwd_scale_ptrs", "_down_bwd_scale_ptrs", + "_gate_bwd_weights_per_numa", + "_up_bwd_weights_per_numa", + "_down_bwd_weights_per_numa", + "_gate_bwd_scales_per_numa", + "_up_bwd_scales_per_numa", + "_down_bwd_scales_per_numa", + "_gate_bwd_projs_ptrs", + "_up_bwd_projs_ptrs", + "_down_bwd_projs_ptrs", + "_gate_bwd_scale_ptrs", + "_up_bwd_scale_ptrs", + "_down_bwd_scale_ptrs", ]: setattr(self, attr, None) @@ -312,6 +366,7 @@ def _load_base_weights_from_file(self) -> None: self.up_proj = torch.stack(up_weights, dim=0).contiguous() self.down_proj = torch.stack(down_weights, dim=0).contiguous() else: + def _make_ptrs(arrays_per_numa): return [ [ @@ -416,12 +471,18 @@ def _validate_prepartitioned_weights(self) -> None: def init_lora_weights( self, - gate_lora_a: torch.Tensor, gate_lora_b: torch.Tensor, - up_lora_a: torch.Tensor, up_lora_b: torch.Tensor, - down_lora_a: torch.Tensor, down_lora_b: torch.Tensor, - grad_gate_lora_a: torch.Tensor, grad_gate_lora_b: torch.Tensor, - grad_up_lora_a: torch.Tensor, grad_up_lora_b: torch.Tensor, - grad_down_lora_a: torch.Tensor, grad_down_lora_b: torch.Tensor, + gate_lora_a: torch.Tensor, + gate_lora_b: torch.Tensor, + up_lora_a: torch.Tensor, + up_lora_b: torch.Tensor, + down_lora_a: torch.Tensor, + down_lora_b: torch.Tensor, + grad_gate_lora_a: torch.Tensor, + grad_gate_lora_b: torch.Tensor, + grad_up_lora_a: torch.Tensor, + grad_up_lora_b: torch.Tensor, + grad_down_lora_a: torch.Tensor, + grad_down_lora_b: torch.Tensor, ) -> None: expected_shapes = { "gate_lora_a": (self.num_experts, self.lora_rank, self.hidden_size), @@ -432,9 +493,12 @@ def init_lora_weights( "down_lora_b": (self.num_experts, self.hidden_size, self.lora_rank), } provided = { - "gate_lora_a": gate_lora_a, "gate_lora_b": gate_lora_b, - "up_lora_a": up_lora_a, "up_lora_b": up_lora_b, - "down_lora_a": down_lora_a, "down_lora_b": down_lora_b, + "gate_lora_a": gate_lora_a, + "gate_lora_b": gate_lora_b, + "up_lora_a": up_lora_a, + "up_lora_b": up_lora_b, + "down_lora_a": down_lora_a, + "down_lora_b": down_lora_b, } for name, tensor in provided.items(): expected = expected_shapes[name] @@ -470,6 +534,8 @@ def update_lora_weights(self) -> None: if self._is_skip_lora: return if not self._lora_initialized: + if self.lora_rank <= 0: + return # Full mode without LoRA — no LoRA weights to update raise RuntimeError("LoRA weights not initialized. Call init_lora_weights() first.") # Weight pointer updates are load-time synchronous work. Calling the @@ -484,6 +550,52 @@ def update_lora_weights(self) -> None: self.down_lora_b.data_ptr(), ) + def update_base_weights(self) -> None: + """Sync updated base weight parameters back to C++ kernel after optimizer step.""" + if not self._weights_loaded: + raise RuntimeError("Weights not loaded. Call load_weights() first.") + if not self._full_weight_grad: + return # No base weights to update in LoRA mode + if self.gate_proj_buf is None: + raise RuntimeError("Base weight buffers not initialized. Call init_full_weight_grad_buffers() first.") + + logger.info(f"Layer {self.layer_idx}: update_base_weights() - syncing updated weights to C++ kernel") + + # Preferred path: update config pointers on existing C++ object and re-quantize. + # This avoids full C++ MOE object recreation (~0.6s/layer vs ~1.9s/layer). + if hasattr(self.moe, "set_base_weight_pointers"): + self.moe.set_base_weight_pointers( + self.gate_proj_buf.data.data_ptr(), + self.up_proj_buf.data.data_ptr(), + self.down_proj_buf.data.data_ptr(), + ) + self.cpu_infer.submit(self.moe.load_weights_task()) + self.cpu_infer.sync() + logger.info(f"Layer {self.layer_idx}: update_base_weights() - re-quantized existing kernel") + return + + # Fallback: full reload path (creates new C++ MOE object) + # This is slower but works without C++ set_base_weight_pointers support. + logger.warning( + f"Layer {self.layer_idx}: set_base_weight_pointers not available, " + f"falling back to full C++ MOE object recreation" + ) + old_moe = getattr(self, "moe", None) + if old_moe is not None: + del old_moe + + self.gate_proj = self.gate_proj_buf.data + self.up_proj = self.up_proj_buf.data + self.down_proj = self.down_proj_buf.data + physical_to_logical_map = torch.arange(self.num_experts, dtype=torch.int64, device="cpu") + self._weights_loaded = False # Allow re-load + self.load_weights_from_tensors( + gate_proj=self.gate_proj, + up_proj=self.up_proj, + down_proj=self.down_proj, + physical_to_logical_map_cpu=physical_to_logical_map, + ) + def save_backward_weights_from_tensors( self, gate_proj: torch.Tensor, diff --git a/kt-kernel/python/sft/arch.py b/kt-kernel/python/sft/arch.py index 43b2e2cbc..234926d61 100644 --- a/kt-kernel/python/sft/arch.py +++ b/kt-kernel/python/sft/arch.py @@ -110,6 +110,18 @@ def get_moe_arch_config(config) -> MOEArchConfig: num_experts_per_tok=cfg.num_experts_per_tok, has_shared_experts=getattr(cfg, "shared_expert_intermediate_size", 0) > 0, ) + if "Glm4Moe" in arch: + return MOEArchConfig( + moe_layer_attr="mlp", + router_attr="gate", + experts_attr="experts", + weight_names=("gate_proj", "up_proj", "down_proj"), + expert_num=config.n_routed_experts, + intermediate_size=config.moe_intermediate_size, + num_experts_per_tok=config.num_experts_per_tok, + has_shared_experts=getattr(config, "n_shared_experts", 0) > 0, + router_type="glm4_moe_gate", + ) if "Mixtral" in arch: return MOEArchConfig( moe_layer_attr="block_sparse_moe", @@ -124,7 +136,7 @@ def get_moe_arch_config(config) -> MOEArchConfig: raise KTAMXModelNotSupportedError( f"Model architecture {arch} not supported for KT AMX. " - "Supported architectures: DeepseekV2, DeepseekV3, Qwen2Moe, Qwen3Moe, Qwen3_5Moe, Mixtral" + "Supported architectures: DeepseekV2, DeepseekV3, Qwen2Moe, Qwen3Moe, Qwen3_5Moe, Glm4Moe, Mixtral" ) diff --git a/kt-kernel/python/sft/autograd.py b/kt-kernel/python/sft/autograd.py index 0264e9de6..615429a7c 100644 --- a/kt-kernel/python/sft/autograd.py +++ b/kt-kernel/python/sft/autograd.py @@ -38,12 +38,19 @@ def forward( training: bool, train_lora: bool, all_qlens: list[int] | tuple[int, ...] | None, + cache_checkpoint_forward: bool = False, + reuse_cached_forward: bool = False, + gate_proj_param: torch.Tensor | None = None, + up_proj_param: torch.Tensor | None = None, + down_proj_param: torch.Tensor | None = None, ) -> torch.Tensor: if _KT_SFT_DEBUG: logging.debug( "KTMoEFunction.forward: layer=%d training=%s train_lora=%s", - layer_idx, training, train_lora, + layer_idx, + training, + train_lora, ) original_device = hidden_states.device @@ -52,6 +59,7 @@ def forward( qlen = batch_size * seq_len import torch.distributed as dist + dist_on = dist.is_initialized() and dist.get_world_size() > 1 rank = dist.get_rank() if dist.is_initialized() else 0 world_size = dist.get_world_size() if dist_on else 1 @@ -65,18 +73,24 @@ def forward( else: all_qlens_list = [int(q) for q in all_qlens] if len(all_qlens_list) != world_size: - raise RuntimeError( - f"all_qlens length mismatch: got {len(all_qlens_list)}, expected {world_size}" - ) + raise RuntimeError(f"all_qlens length mismatch: got {len(all_qlens_list)}, expected {world_size}") if int(all_qlens_list[rank]) != qlen: - raise RuntimeError( - f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens_list[rank]}" - ) + raise RuntimeError(f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens_list[rank]}") total_qlen = sum(all_qlens_list) # Rank 0: sync CPU result and split by real lengths if rank == 0: - cpu_output = wrapper.sync_forward(output_device=original_device) + if reuse_cached_forward: + with torch.profiler.record_function("kt.sft.checkpoint_cached_cpu_moe"): + cpu_output = wrapper.get_checkpoint_output(total_qlen, output_device=original_device) + elif cache_checkpoint_forward: + with torch.profiler.record_function("kt.sft.cpu_forward_sync"): + cached_output = wrapper.sync_forward(output_device=None) + wrapper.cache_checkpoint_output(cached_output, total_qlen) + cpu_output = cached_output.to(device=original_device, non_blocking=True) + else: + with torch.profiler.record_function("kt.sft.cpu_forward_sync"): + cpu_output = wrapper.sync_forward(output_device=original_device) cpu_output = cpu_output.to(dtype=original_dtype).view(total_qlen, hidden_size) offsets = _qlen_offsets(all_qlens_list) scatter_list = [cpu_output[offsets[i] : offsets[i + 1]].contiguous() for i in range(world_size)] @@ -96,13 +110,21 @@ def forward( del output_flat elif wrapper is not None: # Single-GPU: sync directly - cpu_output = wrapper.sync_forward(output_device=original_device) + if reuse_cached_forward: + with torch.profiler.record_function("kt.sft.checkpoint_cached_cpu_moe"): + cpu_output = wrapper.get_checkpoint_output(qlen, output_device=original_device) + elif cache_checkpoint_forward: + with torch.profiler.record_function("kt.sft.cpu_forward_sync"): + cached_output = wrapper.sync_forward(output_device=None) + wrapper.cache_checkpoint_output(cached_output, qlen) + cpu_output = cached_output.to(device=original_device, non_blocking=True) + else: + with torch.profiler.record_function("kt.sft.cpu_forward_sync"): + cpu_output = wrapper.sync_forward(output_device=original_device) output = cpu_output.view(batch_size, seq_len, hidden_size).to(dtype=original_dtype) else: # Broadcast-only rank (no wrapper) - output = torch.empty( - batch_size, seq_len, hidden_size, device=original_device, dtype=original_dtype - ) + output = torch.empty(batch_size, seq_len, hidden_size, device=original_device, dtype=original_dtype) ctx.wrapper = wrapper ctx.hidden_size = hidden_size @@ -120,6 +142,11 @@ def forward( ctx.num_experts_per_tok = num_experts_per_tok ctx.layer_idx = layer_idx + # Store base weight param references for gradient flow in full mode + ctx.full_weight_grad = ( + wrapper is not None and getattr(wrapper, "_full_weight_grad", False) and gate_proj_param is not None + ) + # Save a sentinel tensor so non-reentrant checkpoint's saved_tensors # hooks can intercept it. When backward accesses ctx.saved_tensors, # the checkpoint unpack hook triggers a full recompute of the decoder @@ -135,13 +162,15 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor): # Wait for any in-flight async repack before recompute forward uses the pool - if getattr(ctx.wrapper, 'share_backward_bb', False): - ctx.wrapper.wait_backward_repack() + if getattr(ctx.wrapper, "share_backward_bb", False): + with torch.profiler.record_function("kt.sft.wait_backward_repack"): + ctx.wrapper.wait_backward_repack() # Access saved_tensors FIRST — under non-reentrant checkpoint this # triggers the unpack hook which runs a full decoder-layer recompute, # populating the C++ cache before we call wrapper.backward(). - _ = ctx.saved_tensors + with torch.profiler.record_function("kt.sft.checkpoint_recompute"): + _ = ctx.saved_tensors qlen = ctx.qlen hidden_size = ctx.hidden_size @@ -152,12 +181,15 @@ def backward(ctx, grad_output: torch.Tensor): num_experts_per_tok = ctx.num_experts_per_tok import torch.distributed as dist + rank = dist.get_rank() if dist.is_initialized() else 0 if _KT_SFT_DEBUG: logging.debug( "KTMoEFunction.backward: layer=%d dist_on=%s qlen=%d", - getattr(ctx, "layer_idx", -1), dist_on, qlen, + getattr(ctx, "layer_idx", -1), + dist_on, + qlen, ) if dist_on: @@ -183,10 +215,11 @@ def backward(ctx, grad_output: torch.Tensor): all_go = torch.cat(gathered_go, dim=0) total_qlen = int(all_go.shape[0]) - backward_out = ctx.wrapper.backward( - all_go, - output_device=ctx.original_device, - ) + with torch.profiler.record_function("kt.sft.cpu_backward"): + backward_out = ctx.wrapper.backward( + all_go, + output_device=ctx.original_device, + ) if isinstance(backward_out, tuple) and len(backward_out) == 2: all_grad_input, all_grad_weights = backward_out elif isinstance(backward_out, tuple) and len(backward_out) == 3: @@ -228,11 +261,14 @@ def backward(ctx, grad_output: torch.Tensor): elif not ctx.use_broadcast: # ---- Single-GPU path ---- grad_output_flat = grad_output.view(qlen, hidden_size) - backward_out = ctx.wrapper.backward( - grad_output_flat, - output_device=ctx.original_device, - ) - ctx.wrapper._kt_has_cached_forward = False + try: + with torch.profiler.record_function("kt.sft.cpu_backward"): + backward_out = ctx.wrapper.backward( + grad_output_flat, + output_device=ctx.original_device, + ) + finally: + ctx.wrapper.clear_checkpoint_output() if isinstance(backward_out, tuple) and len(backward_out) == 2: grad_input, grad_weights = backward_out elif isinstance(backward_out, tuple) and len(backward_out) == 3: @@ -243,12 +279,42 @@ def backward(ctx, grad_output: torch.Tensor): grad_weights = grad_weights.to(dtype=torch.bfloat16) else: # No wrapper, no dist — shouldn't happen in normal flow - grad_input = torch.zeros(batch_size, seq_len, hidden_size, device=ctx.original_device, dtype=ctx.original_dtype) + grad_input = torch.zeros( + batch_size, seq_len, hidden_size, device=ctx.original_device, dtype=ctx.original_dtype + ) grad_weights = torch.zeros(ctx.weights_shape, device=ctx.weights_device, dtype=ctx.weights_dtype) # Trigger async repack for next MoE layer in backward order - next_bwd = getattr(ctx.wrapper, '_next_backward_wrapper', None) - if next_bwd is not None and getattr(next_bwd, 'share_backward_bb', False): - next_bwd.submit_backward_repack() - - return grad_input, None, grad_weights, None, None, None, None, None, None, None, None + next_bwd = getattr(ctx.wrapper, "_next_backward_wrapper", None) + if next_bwd is not None and getattr(next_bwd, "share_backward_bb", False): + with torch.profiler.record_function("kt.sft.submit_backward_repack"): + next_bwd.submit_backward_repack() + + # Base weight gradients: return C++-written grad buffers in full mode, None otherwise + if ctx.full_weight_grad and ctx.wrapper is not None: + grad_gate_proj = ctx.wrapper.grad_gate_proj_buf + grad_up_proj = ctx.wrapper.grad_up_proj_buf + grad_down_proj = ctx.wrapper.grad_down_proj_buf + else: + grad_gate_proj = None + grad_up_proj = None + grad_down_proj = None + + return ( + grad_input, + None, + grad_weights, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + grad_gate_proj, + grad_up_proj, + grad_down_proj, + ) diff --git a/kt-kernel/python/sft/base.py b/kt-kernel/python/sft/base.py index 57ba8e424..d402dc3fe 100644 --- a/kt-kernel/python/sft/base.py +++ b/kt-kernel/python/sft/base.py @@ -145,6 +145,7 @@ def __init__( lora_rank: int = 16, lora_alpha: float = 32.0, max_cache_depth: int = 1, + full_weight_grad: bool = False, ): self.cpu_infer = self._get_cpu_infer(cpuinfer_threads, threadpool_count) @@ -154,7 +155,7 @@ def __init__( moe_intermediate_size=moe_intermediate_size, num_experts_per_tok=num_experts_per_tok, ) - self._validate_sft_config(lora_rank, lora_alpha, max_cache_depth) + self._validate_sft_config(lora_rank, lora_alpha, max_cache_depth, full_weight_grad=full_weight_grad) self.layer_idx = layer_idx self.num_experts = num_experts @@ -168,9 +169,11 @@ def __init__( self.lora_rank = lora_rank self.lora_alpha = lora_alpha - self.lora_scaling = lora_alpha / lora_rank + self.lora_scaling = lora_alpha / lora_rank if lora_rank > 0 else 0.0 self.max_cache_depth = max_cache_depth + self._full_weight_grad = full_weight_grad + self.gate_lora_a: Optional[torch.Tensor] = None self.gate_lora_b: Optional[torch.Tensor] = None self.up_lora_a: Optional[torch.Tensor] = None @@ -178,22 +181,79 @@ def __init__( self.down_lora_a: Optional[torch.Tensor] = None self.down_lora_b: Optional[torch.Tensor] = None + # Base weight parameters for full fine-tuning + self.gate_proj_buf: Optional[torch.Tensor] = None + self.up_proj_buf: Optional[torch.Tensor] = None + self.down_proj_buf: Optional[torch.Tensor] = None + self.grad_gate_proj_buf: Optional[torch.Tensor] = None + self.grad_up_proj_buf: Optional[torch.Tensor] = None + self.grad_down_proj_buf: Optional[torch.Tensor] = None + self._weights_loaded: bool = False self._lora_initialized: bool = False self._cache_depth: int = 0 self._is_skip_lora: bool = False + self._base_weights_dirty: bool = False + self.reuse_checkpoint_forward: bool = False + self._kt_has_cached_forward: bool = False + self._checkpoint_output_cpu: Optional[torch.Tensor] = None + self._checkpoint_output_qlen: int = 0 self.moe = None @staticmethod - def _validate_sft_config(lora_rank: int, lora_alpha: float, max_cache_depth: int) -> None: - if lora_rank <= 0: - raise ValueError(f"lora_rank must be positive, got {lora_rank}") - if lora_alpha <= 0: + def _validate_sft_config( + lora_rank: int, lora_alpha: float, max_cache_depth: int, full_weight_grad: bool = False + ) -> None: + if not full_weight_grad and lora_rank <= 0: + raise ValueError( + f"lora_rank must be positive in LoRA mode, got {lora_rank}. " + "Set kt_train_mode='full' for full fine-tuning." + ) + if lora_rank > 0 and lora_alpha <= 0: raise ValueError(f"lora_alpha must be positive, got {lora_alpha}") if max_cache_depth <= 0: raise ValueError(f"max_cache_depth must be positive, got {max_cache_depth}") + # ========== Full weight grad methods ========== + + def init_full_weight_grad_buffers( + self, gate_proj: torch.Tensor, up_proj: torch.Tensor, down_proj: torch.Tensor + ) -> None: + """Initialize base weight nn.Parameter buffers and gradient buffers for full fine-tuning. + + Args: + gate_proj: [num_experts, intermediate_size, hidden_size] BF16 CPU tensor + up_proj: [num_experts, intermediate_size, hidden_size] BF16 CPU tensor + down_proj: [num_experts, hidden_size, intermediate_size] BF16 CPU tensor + """ + import torch.nn as nn + + dtype = torch.bfloat16 + E = self.num_experts + I = self.moe_intermediate_size + H = self.hidden_size + + # Create nn.Parameter buffers (optimizer-visible) + self.gate_proj_buf = nn.Parameter(gate_proj.to(dtype=dtype, device="cpu").contiguous(), requires_grad=True) + self.up_proj_buf = nn.Parameter(up_proj.to(dtype=dtype, device="cpu").contiguous(), requires_grad=True) + self.down_proj_buf = nn.Parameter(down_proj.to(dtype=dtype, device="cpu").contiguous(), requires_grad=True) + + # Create gradient buffers (C++ writes directly to these) + self.grad_gate_proj_buf = torch.zeros(E, I, H, dtype=dtype, device="cpu") + self.grad_up_proj_buf = torch.zeros(E, I, H, dtype=dtype, device="cpu") + self.grad_down_proj_buf = torch.zeros(E, H, I, dtype=dtype, device="cpu") + + # Note: .grad is NOT pre-assigned here. PyTorch autograd will set it + # when KTMoEFunction.backward() returns the gradient buffers. + # The C++ kernel writes directly to grad_gate_proj_buf etc., + # and backward returns them so PyTorch can propagate correctly. + + @abstractmethod + def update_base_weights(self) -> None: + """Sync updated base weight parameters back to C++ kernel after optimizer step.""" + ... + # ========== Abstract methods for subclasses ========== @abstractmethod @@ -207,24 +267,27 @@ def _make_backward_task(self, buffer: KExpertsSFTBuffer): ... @abstractmethod - def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None: - ... + def load_weights(self, physical_to_logical_map_cpu: torch.Tensor) -> None: ... @abstractmethod def init_lora_weights( self, - gate_lora_a: torch.Tensor, gate_lora_b: torch.Tensor, - up_lora_a: torch.Tensor, up_lora_b: torch.Tensor, - down_lora_a: torch.Tensor, down_lora_b: torch.Tensor, - grad_gate_lora_a: torch.Tensor, grad_gate_lora_b: torch.Tensor, - grad_up_lora_a: torch.Tensor, grad_up_lora_b: torch.Tensor, - grad_down_lora_a: torch.Tensor, grad_down_lora_b: torch.Tensor, - ) -> None: - ... + gate_lora_a: torch.Tensor, + gate_lora_b: torch.Tensor, + up_lora_a: torch.Tensor, + up_lora_b: torch.Tensor, + down_lora_a: torch.Tensor, + down_lora_b: torch.Tensor, + grad_gate_lora_a: torch.Tensor, + grad_gate_lora_b: torch.Tensor, + grad_up_lora_a: torch.Tensor, + grad_up_lora_b: torch.Tensor, + grad_down_lora_a: torch.Tensor, + grad_down_lora_b: torch.Tensor, + ) -> None: ... @abstractmethod - def update_lora_weights(self) -> None: - ... + def update_lora_weights(self) -> None: ... # ========== Buffer helpers ========== @@ -242,7 +305,7 @@ def _get_buffer(self, qlen: int) -> KExpertsSFTBuffer: def _validate_forward_inputs(self, hidden_states: torch.Tensor, expert_ids: torch.Tensor, weights: torch.Tensor): if not self._weights_loaded: raise RuntimeError("Weights not loaded. Call load_weights() or load_weights_from_tensors() first.") - if not self._lora_initialized and not self._is_skip_lora: + if not self._lora_initialized and not self._is_skip_lora and not self._full_weight_grad: raise RuntimeError("LoRA weights not initialized. Call init_lora_weights() first.") qlen = hidden_states.shape[0] if qlen > self.chunked_prefill_size: @@ -255,12 +318,16 @@ def _validate_forward_inputs(self, hidden_states: torch.Tensor, expert_ids: torc f"expert_ids shape {tuple(expert_ids.shape)} must be ({qlen}, {self.num_experts_per_tok})." ) if weights.shape[0] != qlen or weights.shape[1] != self.num_experts_per_tok: - raise ValueError( - f"weights shape {tuple(weights.shape)} must be ({qlen}, {self.num_experts_per_tok})." - ) + raise ValueError(f"weights shape {tuple(weights.shape)} must be ({qlen}, {self.num_experts_per_tok}).") - def _copy_inputs_to_buffer(self, buffer: KExpertsSFTBuffer, hidden_states: torch.Tensor, - expert_ids: torch.Tensor, weights: torch.Tensor, qlen: int) -> torch.device: + def _copy_inputs_to_buffer( + self, + buffer: KExpertsSFTBuffer, + hidden_states: torch.Tensor, + expert_ids: torch.Tensor, + weights: torch.Tensor, + qlen: int, + ) -> torch.device: """Copy inputs to CPU buffer, return input device.""" input_device = hidden_states.device buffer.input_cpu[:qlen].copy_(hidden_states.to(torch.bfloat16), non_blocking=True) @@ -284,6 +351,30 @@ def _return_output(self, buffer: KExpertsSFTBuffer, qlen: int, output_device: Op else: return buffer.output_cpu[:qlen].clone() + def cache_checkpoint_output(self, output_cpu: torch.Tensor, qlen: int) -> None: + if output_cpu.device.type != "cpu": + raise ValueError("checkpoint CPU expert output must reside on CPU") + if output_cpu.shape[0] < qlen: + raise ValueError(f"checkpoint output is shorter than qlen: {output_cpu.shape[0]} < {qlen}") + self._checkpoint_output_cpu = output_cpu[:qlen].contiguous() + self._checkpoint_output_qlen = qlen + self._kt_has_cached_forward = True + + def get_checkpoint_output(self, qlen: int, output_device: Optional[torch.device] = None) -> torch.Tensor: + if not self._kt_has_cached_forward or self._checkpoint_output_cpu is None: + raise RuntimeError("No cached checkpoint forward output is available.") + if qlen != self._checkpoint_output_qlen: + raise RuntimeError(f"Cached checkpoint qlen mismatch: cached={self._checkpoint_output_qlen}, requested={qlen}") + output = self._checkpoint_output_cpu + if output_device is not None: + return output.to(device=output_device, non_blocking=True) + return output + + def clear_checkpoint_output(self) -> None: + self._checkpoint_output_cpu = None + self._checkpoint_output_qlen = 0 + self._kt_has_cached_forward = False + def _return_grads(self, buffer: KExpertsSFTBuffer, qlen: int, output_device: Optional[torch.device]): if output_device is not None: grad_input = buffer.grad_input_cpu[:qlen].to(device=output_device, non_blocking=True) @@ -510,11 +601,11 @@ def sync_backward(self) -> Tuple[torch.Tensor, torch.Tensor]: def submit_backward_repack(self): if not self._weights_loaded or self.moe is None: return - if hasattr(self.moe, 'submit_backward_repack'): + if hasattr(self.moe, "submit_backward_repack"): self.moe.submit_backward_repack() def wait_backward_repack(self): if not self._weights_loaded or self.moe is None: return - if hasattr(self.moe, 'wait_backward_repack'): + if hasattr(self.moe, "wait_backward_repack"): self.moe.wait_backward_repack() diff --git a/kt-kernel/python/sft/config.py b/kt-kernel/python/sft/config.py index 35af4d3d6..a4eca76da 100644 --- a/kt-kernel/python/sft/config.py +++ b/kt-kernel/python/sft/config.py @@ -12,11 +12,18 @@ from __future__ import annotations import dataclasses +import logging import os from dataclasses import dataclass, field +from pathlib import Path from typing import Any, Callable +logger = logging.getLogger(__name__) + +_CPU_TOPOLOGY_ROOT = Path("/sys/devices/system/cpu") + + def _env_int(key: str, default: int | None) -> int | None: value = os.environ.get(key, None) if value is None or value == "": @@ -38,6 +45,83 @@ def _env_bool(key: str, default: bool) -> bool: return value.lower() in ("1", "true", "yes") +def _available_cpu_ids() -> set[int]: + """Return CPUs available to this process, respecting affinity/cpuset limits.""" + try: + return set(os.sched_getaffinity(0)) + except (AttributeError, OSError): + return set(range(os.cpu_count() or 1)) + + +def _read_cpu_topology(cpu_id: int) -> tuple[int, int] | None: + topology = _CPU_TOPOLOGY_ROOT / f"cpu{cpu_id}" / "topology" + try: + package_id = int((topology / "physical_package_id").read_text().strip()) + core_id = int((topology / "core_id").read_text().strip()) + except (OSError, ValueError): + return None + return package_id, core_id + + +def detect_physical_cpu_count() -> int: + """Count physical cores available to the current process. + + Linux exposes a stable ``(physical_package_id, core_id)`` pair for every + logical CPU. Counting those pairs avoids assigning one OpenMP worker to + each SMT sibling. If topology is unavailable, fall back to the number of + affinity-visible logical CPUs. + """ + cpu_ids = _available_cpu_ids() + physical_cores = { + topology + for cpu_id in cpu_ids + if (topology := _read_cpu_topology(cpu_id)) is not None + } + return max(1, len(physical_cores) if physical_cores else len(cpu_ids)) + + +def _set_torch_num_threads(num_threads: int) -> None: + try: + import torch + except ImportError: + return + torch.set_num_threads(num_threads) + + +def configure_omp_threads() -> int: + """Configure OpenMP for KT SFT CPU tensor work. + + ``accelerate launch`` defaults GPU jobs to ``OMP_NUM_THREADS=1`` when the + caller did not choose a value. That makes Full-FT CPU gradient accumulation, + AdamW, and zeroing effectively serial. Treat that value as the launcher + default and select the affinity-visible physical core count instead. + + ``ACCELERATE_KT_OMP_NUM_THREADS`` is the unambiguous KT-specific override, + including when an intentional single-thread run is required. An existing + generic ``OMP_NUM_THREADS`` value greater than one is also preserved. + """ + kt_override = _env_int("ACCELERATE_KT_OMP_NUM_THREADS", None) + current_omp = _env_int("OMP_NUM_THREADS", None) + + if kt_override is not None: + num_threads = kt_override + source = "ACCELERATE_KT_OMP_NUM_THREADS" + elif current_omp is not None and current_omp > 1: + num_threads = current_omp + source = "OMP_NUM_THREADS" + else: + num_threads = detect_physical_cpu_count() + source = "available physical cores" + + if num_threads < 1: + raise ValueError(f"OpenMP thread count must be positive, got {num_threads}") + + os.environ["OMP_NUM_THREADS"] = str(num_threads) + _set_torch_num_threads(num_threads) + logger.info("KT SFT configured OMP_NUM_THREADS=%d from %s", num_threads, source) + return num_threads + + @dataclass class KTConfig: """ @@ -75,6 +159,10 @@ class KTConfig: kt_lora_rank: int | None = None kt_lora_alpha: float | None = None + # Training mode + kt_train_mode: str | None = None # "lora" | "full" | "hybrid" + kt_full_weight_grad: bool | None = None # auto-set True when train_mode in (full, hybrid) + # LoRA Experts (GPU-side extra experts) kt_use_lora_experts: bool | None = None kt_lora_expert_num: int | None = None @@ -100,6 +188,7 @@ def from_object(cls, obj: Any) -> "KTConfig": return cls(**kwargs) def __post_init__(self): + configure_omp_threads() if self.kt_backend is None: self.kt_backend = os.environ.get("ACCELERATE_KT_BACKEND", "AMXBF16") if self.kt_num_threads is None: @@ -132,6 +221,10 @@ def __post_init__(self): self.kt_lora_alpha = _env_float("ACCELERATE_KT_LORA_ALPHA", None) if self.kt_lora_alpha is None and self.kt_lora_rank is not None: self.kt_lora_alpha = float(self.kt_lora_rank * 2) + if self.kt_train_mode is None: + self.kt_train_mode = os.environ.get("ACCELERATE_KT_TRAIN_MODE", "lora") + if self.kt_full_weight_grad is None: + self.kt_full_weight_grad = self.kt_train_mode in ("full", "hybrid") if self.kt_model_max_length is None: self.kt_model_max_length = _env_int("ACCELERATE_KT_MODEL_MAX_LENGTH", None) if self.kt_skip_expert_loading is None: diff --git a/kt-kernel/python/sft/layer.py b/kt-kernel/python/sft/layer.py index e4cb2b657..4f0acd902 100644 --- a/kt-kernel/python/sft/layer.py +++ b/kt-kernel/python/sft/layer.py @@ -13,6 +13,7 @@ import logging import os +from contextlib import nullcontext from typing import Any import torch @@ -62,7 +63,29 @@ def __init__( # 1. gate/router FIRST - keep original attribute name for PEFT compatibility router_attr = moe_config.router_attr # "gate" for Qwen3/DeepSeek - setattr(self, router_attr, getattr(original_moe, router_attr, None)) + original_router = getattr(original_moe, router_attr, None) + self._original_router = None # Set when router is not nn.Linear (e.g. TopKRouter) + + if original_router is not None and isinstance(original_router, nn.Linear): + # transformers <=4.x / some models: gate is nn.Linear - register directly. + setattr(self, router_attr, original_router) + elif original_router is not None and hasattr(original_router, "weight") and isinstance( + getattr(original_router, "weight"), nn.Parameter + ): + # transformers v5+: gate is a TopKRouter with nn.Parameter weight. + # Wrap it in nn.Linear so PEFT can discover and inject LoRA. + # The nn.Linear shares the same weight tensor - LoRA applied to it + # is equivalent to LoRA on the original gate. + router_weight = original_router.weight + router_linear = nn.Linear( + router_weight.shape[1], router_weight.shape[0], bias=False, + ) + router_linear.weight = router_weight # share the same parameter + setattr(self, router_attr, router_linear) + # Keep the original router for forward (top-k selection logic) + self._original_router = original_router + else: + setattr(self, router_attr, original_router) self._router_attr = router_attr # 2. experts SECOND (this is what PEFT targets for LoRA) @@ -84,10 +107,13 @@ def __init__( self._peft_lora_modules: dict[int, dict[str, tuple[nn.Module, nn.Module]]] | None = None self._lora_pointers_dirty = False + # Full weight grad mode (set during wrapping or kt_adapt_peft_lora) + self._full_weight_grad = getattr(wrapper, "_full_weight_grad", False) if wrapper is not None else False + def _apply(self, fn, recurse=True): # Protect experts from device transfer (PEFT LoRA should stay on CPU for KT) saved_experts = None - experts_attr = getattr(self, '_experts_attr', None) + experts_attr = getattr(self, "_experts_attr", None) if experts_attr is not None and getattr(self, experts_attr, None) is not None: saved_experts = getattr(self, experts_attr) @@ -103,36 +129,59 @@ def _apply(self, fn, recurse=True): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: import torch.distributed as dist + dist_on = dist.is_initialized() and dist.get_world_size() > 1 rank = dist.get_rank() if dist.is_initialized() else 0 # Check if we need to use distributed broadcast (only rank 0 has KT kernel) use_broadcast = dist_on and self.wrapper is None - topk_ids, topk_weights = self._compute_routing(hidden_states) + with torch.profiler.record_function("kt.sft.routing"): + topk_ids, topk_weights = self._compute_routing(hidden_states) train_lora = self._peft_lora_modules is not None and len(self._peft_lora_modules) > 0 + full_weight_grad = self._full_weight_grad save_for_backward = ( self.training and torch.is_grad_enabled() - and (hidden_states.requires_grad or topk_weights.requires_grad or train_lora) + and (hidden_states.requires_grad or topk_weights.requires_grad or train_lora or full_weight_grad) ) use_autograd_path = save_for_backward + checkpoint_mode = _checkpoint_hook_mode() + reuse_checkpoint_forward = ( + not dist_on + and self.wrapper is not None + and getattr(self.wrapper, "reuse_checkpoint_forward", False) + ) + reuse_cached_forward = ( + reuse_checkpoint_forward + and checkpoint_mode == "recompute" + and getattr(self.wrapper, "_kt_has_cached_forward", False) + ) + cache_checkpoint_forward = reuse_checkpoint_forward and checkpoint_mode == "first_forward" save_for_backward_submit = use_autograd_path - if _checkpoint_hook_mode() == "first_forward": + if checkpoint_mode == "first_forward" and not cache_checkpoint_forward: save_for_backward_submit = False if train_lora and self._lora_pointers_dirty: self.update_lora_pointers() self._lora_pointers_dirty = False - gpu_output, all_qlens = self._submit_and_compute_gpu( - hidden_states, - topk_ids, - topk_weights, - save_for_backward_submit, - ) + # In full_weight_grad mode, sync base weights after optimizer step + if full_weight_grad and getattr(self.wrapper, "_base_weights_dirty", False): + with torch.profiler.record_function("kt.sft.base_weight_reload"): + self.wrapper.update_base_weights() + self.wrapper._base_weights_dirty = False + + with torch.profiler.record_function("kt.sft.submit_and_gpu_experts"): + gpu_output, all_qlens = self._submit_and_compute_gpu( + hidden_states, + topk_ids, + topk_weights, + save_for_backward_submit, + reuse_cached_forward, + ) # Use KTMoEFunction whenever backward is needed so KT backward and LoRA # gradient paths remain connected. @@ -141,25 +190,36 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: if train_lora and self._peft_lora_modules: for expert_loras in self._peft_lora_modules.values(): for lora_A, lora_B in expert_loras.values(): - if hasattr(lora_A, 'weight') and lora_A.weight.requires_grad: + if hasattr(lora_A, "weight") and lora_A.weight.requires_grad: lora_ref = lora_A.weight break if lora_ref.numel() > 0: break - - moe_output = KTMoEFunction.apply( - hidden_states, - topk_ids, - topk_weights, - self.wrapper, - lora_ref, - self.hidden_size, - self.moe_config.num_experts_per_tok, - self.layer_idx, - save_for_backward, - train_lora, - all_qlens, - ) + elif full_weight_grad and self.wrapper is not None: + # In full mode, use base weight param as autograd sentinel + if self.wrapper.gate_proj_buf is not None: + lora_ref = self.wrapper.gate_proj_buf + + with torch.profiler.record_function("kt.sft.autograd_apply_and_cpu_sync"): + moe_output = KTMoEFunction.apply( + hidden_states, + topk_ids, + topk_weights, + self.wrapper, + lora_ref, + self.hidden_size, + self.moe_config.num_experts_per_tok, + self.layer_idx, + save_for_backward, + train_lora, + all_qlens, + cache_checkpoint_forward, + reuse_cached_forward, + # Base weight params for full mode gradient flow + self.wrapper.gate_proj_buf if full_weight_grad and self.wrapper is not None else None, + self.wrapper.up_proj_buf if full_weight_grad and self.wrapper is not None else None, + self.wrapper.down_proj_buf if full_weight_grad and self.wrapper is not None else None, + ) else: moe_output = self._sync_forward_output_no_autograd( hidden_states=hidden_states, @@ -194,13 +254,9 @@ def _sync_forward_output_no_autograd( else: all_qlens_list = [int(q) for q in all_qlens] if len(all_qlens_list) != world_size: - raise RuntimeError( - f"all_qlens length mismatch: got {len(all_qlens_list)}, expected {world_size}" - ) + raise RuntimeError(f"all_qlens length mismatch: got {len(all_qlens_list)}, expected {world_size}") if int(all_qlens_list[rank]) != qlen: - raise RuntimeError( - f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens_list[rank]}" - ) + raise RuntimeError(f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens_list[rank]}") total_qlen = sum(all_qlens_list) if rank == 0: @@ -234,12 +290,11 @@ def _sync_forward_output_no_autograd( return torch.empty(batch_size, seq_len, self.hidden_size, device=original_device, dtype=original_dtype) def _compute_routing(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - # Run routing under no_grad to avoid creating autograd nodes whose - # SavedVariables become orphan holders inside gradient checkpoint. - # The gate is frozen during LoRA fine-tuning and the main gradient - # flows through KTMoEFunction.backward()'s grad_input, so the - # routing gradient contribution to hidden_states can be safely dropped. - with torch.no_grad(): + # In full_weight_grad mode, Router gradients should flow (no torch.no_grad). + # In LoRA mode, Router is frozen — wrap in no_grad to avoid orphan autograd nodes. + no_grad_ctx = torch.no_grad() if not self._full_weight_grad else nullcontext() + + with no_grad_ctx: router = getattr(self, self._router_attr) if self.router_type == "deepseek_gate": # DeepSeek V3's MoEGate has `assert not self.training` in its noaux_tc @@ -259,6 +314,60 @@ def _compute_routing(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, t topk_weights = topk_weights.to(torch.bfloat16) return topk_ids, topk_weights + # When _original_router is set, self.gate is an nn.Linear wrapper + # around the TopKRouter's weight. Use it (with PEFT LoRA if + # applied) for the linear projection, then replicate top-k logic. + if self._original_router is not None: + orig_router = self._original_router + router_logits = router(hidden_states.view(-1, self.hidden_size)) + if self.router_type == "glm4_moe_gate": + router_probs = torch.sigmoid(router_logits.float()) + correction_bias = getattr(orig_router, "e_score_correction_bias", None) + if correction_bias is None: + router_logits_for_choice = router_probs + else: + router_logits_for_choice = router_probs + correction_bias.to( + device=router_probs.device, + dtype=router_probs.dtype, + ) + n_group = getattr(orig_router, "n_group", 1) + topk_group = getattr(orig_router, "topk_group", n_group) + expert_num = self.moe_config.expert_num + group_scores = ( + router_logits_for_choice.view(-1, n_group, expert_num // n_group) + .topk(2, dim=-1)[0] + .sum(dim=-1) + ) + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(-1, n_group, expert_num // n_group) + .reshape(-1, expert_num) + ) + scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0) + top_k = getattr(orig_router, "top_k", self.moe_config.num_experts_per_tok) + topk_ids = torch.topk(scores_for_choice, k=top_k, dim=-1, sorted=False)[1] + topk_weights = router_probs.gather(1, topk_ids) + if getattr(orig_router, "norm_topk_prob", True): + topk_weights = topk_weights / (topk_weights.sum(dim=-1, keepdim=True) + 1e-20) + topk_weights = topk_weights * getattr(orig_router, "routed_scaling_factor", 1.0) + if topk_weights.is_floating_point(): + topk_weights = topk_weights.to(torch.bfloat16) + return topk_ids, topk_weights + + router_probs = F.softmax(router_logits, dtype=torch.float, dim=-1) + top_k = getattr(orig_router, "top_k", self.moe_config.num_experts_per_tok) + norm_topk_prob = getattr(orig_router, "norm_topk_prob", True) + topk_weights, topk_ids = torch.topk(router_probs, top_k, dim=-1) + if norm_topk_prob: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + topk_weights = topk_weights.to(router_logits.dtype) + if topk_weights.is_floating_point(): + topk_weights = topk_weights.to(torch.bfloat16) + return topk_ids, topk_weights + router_output = router(hidden_states.view(-1, self.hidden_size)) # transformers v5 TopKRouter returns (router_logits, router_scores, router_indices) # directly — scores/indices are already topk-normalized. @@ -283,6 +392,7 @@ def _submit_and_compute_gpu( topk_ids: torch.Tensor, topk_weights: torch.Tensor, save_for_backward: bool, + reuse_cached_forward: bool = False, ) -> tuple[torch.Tensor | None, list[int] | None]: import torch.distributed as dist @@ -299,48 +409,40 @@ def _submit_and_compute_gpu( if dist_on: all_qlens = _all_gather_qlens(qlen, original_device, world_size) if int(all_qlens[rank]) != qlen: - raise RuntimeError( - f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens[rank]}" - ) + raise RuntimeError(f"Rank {rank} qlen mismatch: local={qlen}, all_qlens[{rank}]={all_qlens[rank]}") total_qlen = sum(all_qlens) - hs_flat = hidden_states.view(qlen, self.hidden_size).contiguous() - expert_ids = topk_ids.view(qlen, self.moe_config.num_experts_per_tok).contiguous() - weights = topk_weights.view(qlen, self.moe_config.num_experts_per_tok).contiguous() - - submit_hs = hs_flat.detach() - submit_ids = expert_ids.detach() - submit_wts = weights.detach() + if not reuse_cached_forward: + hs_flat = hidden_states.view(qlen, self.hidden_size).contiguous() + expert_ids = topk_ids.view(qlen, self.moe_config.num_experts_per_tok).contiguous() + weights = topk_weights.view(qlen, self.moe_config.num_experts_per_tok).contiguous() - gathered_hs = _dist_gather_varlen_to_rank0( - submit_hs, - all_qlens=all_qlens, - rank=rank, - world_size=world_size, - ) - gathered_ids = _dist_gather_varlen_to_rank0( - submit_ids, - all_qlens=all_qlens, - rank=rank, - world_size=world_size, - ) - gathered_wts = _dist_gather_varlen_to_rank0( - submit_wts, - all_qlens=all_qlens, - rank=rank, - world_size=world_size, - ) - - if rank == 0: - all_hs = torch.cat(gathered_hs, dim=0) - all_ids = torch.cat(gathered_ids, dim=0) - all_wts = torch.cat(gathered_wts, dim=0) - self.wrapper.submit_forward( - all_hs, - all_ids, - all_wts, - save_for_backward=save_for_backward, + gathered_hs = _dist_gather_varlen_to_rank0( + hs_flat.detach(), + all_qlens=all_qlens, + rank=rank, + world_size=world_size, + ) + gathered_ids = _dist_gather_varlen_to_rank0( + expert_ids.detach(), + all_qlens=all_qlens, + rank=rank, + world_size=world_size, ) + gathered_wts = _dist_gather_varlen_to_rank0( + weights.detach(), + all_qlens=all_qlens, + rank=rank, + world_size=world_size, + ) + + if rank == 0: + self.wrapper.submit_forward( + torch.cat(gathered_hs, dim=0), + torch.cat(gathered_ids, dim=0), + torch.cat(gathered_wts, dim=0), + save_for_backward=save_for_backward, + ) # Keep shared/lora experts local to avoid qlen_max-style amplification. gpu_output = None @@ -364,12 +466,13 @@ def _submit_and_compute_gpu( submit_hs = input_flat.detach() submit_ids = expert_ids.detach() submit_wts = weights.detach() - self.wrapper.submit_forward( - submit_hs, - submit_ids, - submit_wts, - save_for_backward=save_for_backward, - ) + if not reuse_cached_forward: + self.wrapper.submit_forward( + submit_hs, + submit_ids, + submit_wts, + save_for_backward=save_for_backward, + ) # GPU compute: shared_experts + lora_experts gpu_output = None diff --git a/kt-kernel/python/sft/lora.py b/kt-kernel/python/sft/lora.py index 5a594ec8e..e9f8d6202 100644 --- a/kt-kernel/python/sft/lora.py +++ b/kt-kernel/python/sft/lora.py @@ -98,15 +98,10 @@ def _find_kt_wrappers(model: nn.Module): return wrappers -def get_kt_lora_params(model: nn.Module) -> list[nn.Parameter]: - """Get all MoE LoRA parameters from KT model. - - Returns PEFT LoRA parameters from expert modules and lora_experts parameters. - """ +def _collect_kt_lora_params(wrappers) -> list[nn.Parameter]: + """Collect LoRA-only trainable parameters from KT wrappers.""" params: list[nn.Parameter] = [] - wrappers = _find_kt_wrappers(model) - if wrappers: for wrapper in wrappers: # PEFT LoRA parameters (from _peft_lora_modules) @@ -114,9 +109,9 @@ def get_kt_lora_params(model: nn.Module) -> list[nn.Parameter]: if peft_lora_modules is not None: for expert_loras in peft_lora_modules.values(): for lora_A, lora_B in expert_loras.values(): - if hasattr(lora_A, 'weight') and lora_A.weight.requires_grad: + if hasattr(lora_A, "weight") and lora_A.weight.requires_grad: params.append(lora_A.weight) - if hasattr(lora_B, 'weight') and lora_B.weight.requires_grad: + if hasattr(lora_B, "weight") and lora_B.weight.requires_grad: params.append(lora_B.weight) # Fused expert LoRA parameters (KT-managed, not PEFT) fused_params = getattr(wrapper, "_fused_expert_lora_params", None) @@ -129,6 +124,61 @@ def get_kt_lora_params(model: nn.Module) -> list[nn.Parameter]: return params +def _collect_kt_full_weight_params(wrappers) -> list[nn.Parameter]: + """Collect optimizer-visible base expert parameters for full/hybrid KT SFT.""" + params: list[nn.Parameter] = [] + + if wrappers: + for wrapper in wrappers: + if getattr(wrapper, "_full_weight_grad", False) and wrapper.wrapper is not None: + if wrapper.wrapper.gate_proj_buf is not None: + params.append(wrapper.wrapper.gate_proj_buf) + if wrapper.wrapper.up_proj_buf is not None: + params.append(wrapper.wrapper.up_proj_buf) + if wrapper.wrapper.down_proj_buf is not None: + params.append(wrapper.wrapper.down_proj_buf) + + return params + + +def get_kt_lora_params(model: nn.Module) -> list[nn.Parameter]: + """Get KT parameters for legacy Trainer optimizer injection. + + Historically the patched Trainer calls this function after optimizer + creation. In full_weight_grad mode, returning only LoRA params silently + drops expert base weights from the optimizer, so this compatibility entry + point delegates to the full trainable collector when needed. + """ + wrappers = _find_kt_wrappers(model) + if not wrappers: + return [] + + if any(getattr(w, "_full_weight_grad", False) for w in wrappers): + return _collect_kt_full_weight_params(wrappers) + _collect_kt_lora_params(wrappers) + + return _collect_kt_lora_params(wrappers) + + +def get_kt_trainable_params(model: nn.Module) -> list[nn.Parameter]: + """Get all trainable parameters from KT model based on training mode. + + In full mode: returns base weight nn.Parameter buffers from wrappers. + In LoRA mode: returns LoRA parameters (same as get_kt_lora_params). + """ + wrappers = _find_kt_wrappers(model) + if not wrappers: + return [] + + # Check if any wrapper is in full_weight_grad mode + has_full_weight_grad = any(getattr(w, "_full_weight_grad", False) for w in wrappers) + + if has_full_weight_grad: + return _collect_kt_full_weight_params(wrappers) + _collect_kt_lora_params(wrappers) + else: + # LoRA mode: return LoRA parameters + return _collect_kt_lora_params(wrappers) + + # ============================================================================= # PEFT LoRA Adaptation # ============================================================================= @@ -175,8 +225,24 @@ def kt_adapt_peft_lora(model: nn.Module) -> None: # wrap as nn.Parameter for optimizer, and pre-assign .grad for C++ backward. if getattr(wrapper, "_fused_experts", False): lora_rank = getattr(wrapper, "_lora_rank", 1) + + # In full mode (lora_rank=0), skip LoRA buffer creation entirely. + # C++ kernel will not compute LoRA contributions when lora_rank=0. + if lora_rank == 0: + wrapper._fused_expert_lora_params = [] + wrapper._peft_lora_modules = None + logger.info( + f"[kt_adapt_peft_lora] Layer {layer_idx}: fused expert, " + f"full mode (lora_rank=0, no LoRA buffers)" + ) + adapted_count += 1 + continue + lora_buffers, lora_grad_buffers, lora_params = _create_fused_expert_lora_buffers( - wrapper, moe_config, lora_rank, torch.bfloat16, + wrapper, + moe_config, + lora_rank, + torch.bfloat16, ) if is_rank_0 and wrapper.wrapper is not None: @@ -197,6 +263,18 @@ def kt_adapt_peft_lora(model: nn.Module) -> None: if len(experts) == 0: continue + # In full mode (lora_rank=0), PEFT does not inject LoRA on experts. + # Skip LoRA detection and initialization entirely. + if getattr(wrapper, "_lora_rank", 1) == 0: + wrapper._peft_lora_modules = None + wrapper._fused_expert_lora_params = [] + logger.info( + f"[kt_adapt_peft_lora] Layer {layer_idx}: non-fused expert, " + f"full mode (lora_rank=0, no LoRA)" + ) + adapted_count += 1 + continue + # Collect references to PEFT LoRA modules for each expert # Structure: {expert_idx: {proj_name: (lora_A_module, lora_B_module)}} peft_lora_modules = {} @@ -228,7 +306,16 @@ def kt_adapt_peft_lora(model: nn.Module) -> None: # Store PEFT LoRA references on wrapper wrapper._peft_lora_modules = peft_lora_modules + # In full_weight_grad mode, PEFT LoRA is not injected by LlamaFactory, + # so no PEFT LoRA found is expected — skip the error. if not peft_lora_modules: + if getattr(wrapper, "_full_weight_grad", False): + logger.info( + f"[kt_adapt_peft_lora] Layer {layer_idx}: No PEFT LoRA found " + f"(full_weight_grad mode — expected, skipping)" + ) + adapted_count += 1 + continue raise RuntimeError( f"[kt_adapt_peft_lora] Layer {layer_idx}: No PEFT LoRA found on any expert. " f"Check that PEFT lora_target includes expert modules." @@ -510,9 +597,16 @@ def _replace_peft_weights_with_views( "[_replace_peft_weights_with_views] first param: " "id %s->%s (same=%s) data_ptr %s->%s buf_ptr=%s (match=%s) " "has_grad=%s requires_grad=%s shape=%s", - _old_id_a, _new_id_a, _old_id_a == _new_id_a, - _old_ptr_a, _new_ptr_a, _buf_ptr_a, _new_ptr_a == _buf_ptr_a, - _has_grad, lora_A.weight.requires_grad, tuple(lora_A.weight.shape), + _old_id_a, + _new_id_a, + _old_id_a == _new_id_a, + _old_ptr_a, + _new_ptr_a, + _buf_ptr_a, + _new_ptr_a == _buf_ptr_a, + _has_grad, + lora_A.weight.requires_grad, + tuple(lora_A.weight.shape), ) _first_logged = True _replaced += 1 @@ -526,12 +620,15 @@ def _replace_peft_weights_with_views( def update_kt_lora_pointers(model: nn.Module): - """Mark KT wrapper LoRA pointers as dirty after optimizer.step().""" + """Mark KT wrapper LoRA pointers and base weight pointers as dirty after optimizer.step().""" wrappers = _find_kt_wrappers(model) if wrappers: for wrapper in wrappers: wrapper._lora_pointers_dirty = True + # In full mode, base weights also need re-sync after optimizer step + if getattr(wrapper, "_full_weight_grad", False) and wrapper.wrapper is not None: + wrapper.wrapper._base_weights_dirty = True # ============================================================================= @@ -541,12 +638,10 @@ def update_kt_lora_pointers(model: nn.Module): def sync_kt_lora_gradients(model: nn.Module) -> None: """ - Synchronize KT-managed LoRA gradients across ranks. + Synchronize KT-managed gradients across ranks. - KT computes expert LoRA gradients only on rank 0 (gather/scatter path). This function broadcasts the - per-layer contiguous grad buffers from rank 0 to all ranks so that: - - gradient clipping sees identical grads on every rank - - optimizer.step() applies identical updates + In LoRA mode: synchronizes LoRA gradients only. + In full mode: synchronizes both base weight and LoRA gradients. """ import torch.distributed as dist @@ -557,17 +652,36 @@ def sync_kt_lora_gradients(model: nn.Module) -> None: if world_size <= 1: return - params = get_kt_lora_params(model) + # Sync base weight gradients in full mode + wrappers = _find_kt_wrappers(model) + if wrappers: + for wrapper in wrappers: + if not getattr(wrapper, "_full_weight_grad", False): + continue + if wrapper.wrapper is None: + continue + for grad_buf in ( + wrapper.wrapper.grad_gate_proj_buf, + wrapper.wrapper.grad_up_proj_buf, + wrapper.wrapper.grad_down_proj_buf, + ): + if grad_buf is not None: + grad_gpu = grad_buf.cuda() + dist.all_reduce(grad_gpu, op=dist.ReduceOp.SUM) + grad_gpu.div_(world_size) + grad_buf.copy_(grad_gpu.cpu()) + + # Sync LoRA gradients. Use the LoRA-only helper here because base gradients + # were synchronized above and get_kt_lora_params() is full-aware for legacy + # optimizer injection compatibility. + params = _collect_kt_lora_params(wrappers) if not params: return for param in params: if param.grad is not None: - # Move grad to the same device as the parameter for all-reduce - # Then move back to CPU original_device = param.grad.device if original_device.type == "cpu": - # All-reduce on CPU might be slow; consider using a GPU buffer grad_gpu = param.grad.cuda() dist.all_reduce(grad_gpu, op=dist.ReduceOp.SUM) grad_gpu.div_(world_size) diff --git a/kt-kernel/python/sft/profiler.py b/kt-kernel/python/sft/profiler.py new file mode 100644 index 000000000..b3436fdbb --- /dev/null +++ b/kt-kernel/python/sft/profiler.py @@ -0,0 +1,151 @@ +# Staged profiling helpers for KT SFT MoE. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from collections import defaultdict +from typing import Any + + +def _find_kt_wrappers(model: Any): + wrappers = getattr(model, "_kt_wrappers", None) + if wrappers is None: + base_model = model + for attr in ("base_model", "model"): + if hasattr(base_model, attr): + base_model = getattr(base_model, attr) + wrappers = getattr(base_model, "_kt_wrappers", None) + if wrappers: + break + return wrappers + + +def _profile_moe(layer_wrapper: Any): + backend = getattr(layer_wrapper, "wrapper", layer_wrapper) + return getattr(backend, "moe", None) + + +def collect_kt_sft_profile(model: Any, reset: bool = False) -> dict[str, Any]: + """Collect JSON-compatible staged timing snapshots from all KT MoE layers.""" + layers: dict[int, dict[str, float]] = {} + enabled = False + for layer_wrapper in _find_kt_wrappers(model) or []: + moe = _profile_moe(layer_wrapper) + if moe is None or not hasattr(moe, "get_profile_stats"): + continue + raw = {str(key): float(value) for key, value in moe.get_profile_stats(reset).items()} + layer_idx = int(raw.get("layer_idx", getattr(layer_wrapper, "layer_idx", len(layers)))) + layers[layer_idx] = raw + enabled = enabled or bool(raw.get("wrapper.enabled", 0.0)) + return {"enabled": enabled, "layers": layers} + + +def reset_kt_sft_profile(model: Any) -> None: + """Reset staged timing counters on all KT MoE layers.""" + for layer_wrapper in _find_kt_wrappers(model) or []: + moe = _profile_moe(layer_wrapper) + if moe is not None and hasattr(moe, "reset_profile_stats"): + moe.reset_profile_stats() + + +def _split_timer_key(key: str) -> tuple[str, str] | None: + suffix = ".total_ns" + if not key.endswith(suffix): + return None + timer = key[: -len(suffix)] + if timer.startswith("wrapper."): + return "wrapper", timer[len("wrapper.") :] + if timer.startswith("tp."): + parts = timer.split(".", 2) + if len(parts) == 3: + return f"tp.{parts[1]}", parts[2] + return None + + +def _parent_stage(stage: str) -> str | None: + if stage.startswith("backward.base_weight_grad.worker_cpu."): + return None + if stage.startswith("backward.base_weight_grad."): + return "backward.base_weight_grad" + if stage.startswith("weights.base_reload."): + return "weights.base_reload" + if stage == "backward.down.total" or stage == "backward.gate_up.total": + return "backward.total" + if stage.startswith("backward.down."): + return "backward.down.total" + if stage.startswith("backward.gate_up."): + return "backward.gate_up.total" + if stage.endswith(".total"): + return None + if stage.startswith("forward."): + return "forward.total" + if stage.startswith("backward."): + return "backward.total" + if stage.startswith("tp.forward."): + return "tp.forward.total" + if stage.startswith("tp.backward."): + return "tp.backward.total" + return None + + +def _aggregate_rows(profile: dict[str, Any]) -> list[dict[str, float | str]]: + totals: dict[tuple[str, str], dict[str, float]] = defaultdict( + lambda: {"total_ns": 0.0, "calls": 0.0, "tokens": 0.0, "bytes": 0.0} + ) + for raw in profile.get("layers", {}).values(): + scope_tokens: dict[str, float] = {"wrapper": raw.get("wrapper.tokens", 0.0)} + tp_count = int(raw.get("tp_count", 0.0)) + for tp_idx in range(tp_count): + scope_tokens[f"tp.{tp_idx}"] = raw.get(f"tp.{tp_idx}.tokens", 0.0) + + for key, total_ns in raw.items(): + parsed = _split_timer_key(key) + if parsed is None or total_ns <= 0.0: + continue + scope, stage = parsed + calls = raw.get(key[: -len("total_ns")] + "calls", 0.0) + row = totals[(scope, stage)] + row["total_ns"] += total_ns + row["calls"] += calls + row["tokens"] += scope_tokens.get(scope, 0.0) + row["bytes"] += raw.get(key[: -len("total_ns")] + "bytes", 0.0) + + rows: list[dict[str, float | str]] = [] + for (scope, stage), values in totals.items(): + calls = values["calls"] + tokens = values["tokens"] + parent = _parent_stage(stage) + parent_ns = totals.get((scope, parent), {}).get("total_ns", 0.0) if parent else 0.0 + rows.append( + { + "scope": scope, + "stage": stage, + "calls": calls, + "total_ms": values["total_ns"] / 1e6, + "avg_ms": values["total_ns"] / calls / 1e6 if calls else 0.0, + "us_per_token": values["total_ns"] / tokens / 1e3 if tokens else 0.0, + "mib": values["bytes"] / (1024.0 * 1024.0), + "parent_pct": values["total_ns"] / parent_ns * 100.0 if parent_ns else 0.0, + } + ) + return sorted(rows, key=lambda row: (str(row["scope"]), -float(row["total_ms"]), str(row["stage"]))) + + +def format_kt_sft_profile(profile: dict[str, Any]) -> str: + """Format a compact cross-layer timing table while preserving TP scopes.""" + if not profile.get("enabled"): + return "KT SFT profiler disabled or no profiled KT layers found. Set KT_SFT_PROFILE=1 before model creation." + + rows = _aggregate_rows(profile) + header = ( + f"{'scope':<9} {'stage':<34} {'calls':>7} {'total_ms':>11} " + f"{'avg_ms':>10} {'us/token':>10} {'MiB':>10} {'parent%':>9}" + ) + lines = [header, "-" * len(header)] + for row in rows: + lines.append( + f"{row['scope']:<9} {row['stage']:<34} {row['calls']:>7.0f} " + f"{row['total_ms']:>11.3f} {row['avg_ms']:>10.3f} " + f"{row['us_per_token']:>10.3f} {row['mib']:>10.2f} {row['parent_pct']:>8.1f}%" + ) + return "\n".join(lines) diff --git a/kt-kernel/python/sft/weights.py b/kt-kernel/python/sft/weights.py index c15e22638..e535e1c48 100644 --- a/kt-kernel/python/sft/weights.py +++ b/kt-kernel/python/sft/weights.py @@ -108,9 +108,16 @@ def get_weight_tensor(mod): return gate_proj, up_proj, down_proj -def _clear_original_expert_weights(moe_module: nn.Module, moe_config: MOEArchConfig) -> None: +def _clear_original_expert_weights( + moe_module: nn.Module, moe_config: MOEArchConfig, full_weight_grad: bool = False +) -> None: """ Clear original expert weights to free memory after KT weights are loaded. + + In full_weight_grad mode, gate_proj_buf/up_proj_buf/down_proj_buf serve as + the authoritative copies for the optimizer. The original expert weights in + the model tree are redundant and cause double-counting in count_parameters(). + Clear them just like in LoRA mode. """ from .arch import detect_fused_experts @@ -127,10 +134,14 @@ def _clear_original_expert_weights(moe_module: nn.Module, moe_config: MOEArchCon original_dtype = param.dtype tiny_storage = torch.UntypedStorage(1, device="cpu") fake_tensor = torch.tensor([], dtype=original_dtype, device="cpu").set_( - tiny_storage, storage_offset=0, size=param.shape, + tiny_storage, + storage_offset=0, + size=param.shape, stride=[0] * len(param.shape), ) - experts._parameters[name] = nn.Parameter(fake_tensor, requires_grad=False) + placeholder = nn.Parameter(fake_tensor, requires_grad=False) + placeholder._kt_zero_storage = True # Mark for _setup_full_tuning / count_parameters to skip + experts._parameters[name] = placeholder return def _iter_weight_params(): @@ -141,7 +152,9 @@ def _iter_weight_params(): continue parametrizations = getattr(proj, "parametrizations", None) - parametrized_weight = getattr(parametrizations, "weight", None) if parametrizations is not None else None + parametrized_weight = ( + getattr(parametrizations, "weight", None) if parametrizations is not None else None + ) if parametrized_weight is not None: original = getattr(parametrized_weight, "original", None) if isinstance(original, torch.nn.Parameter): @@ -178,10 +191,13 @@ def _iter_weight_params(): # only used for shape/dtype discovery by PEFT. tiny_storage = torch.UntypedStorage(1, device="cpu") fake_tensor = torch.tensor([], dtype=original_dtype, device="cpu").set_( - tiny_storage, storage_offset=0, size=weight_param.shape, + tiny_storage, + storage_offset=0, + size=weight_param.shape, stride=[0] * len(weight_param.shape), ) new_param = nn.Parameter(fake_tensor, requires_grad=False) + new_param._kt_zero_storage = True # Mark for _setup_full_tuning / count_parameters to skip replaced_count += 1 # Avoid `KeyError: attribute 'weight' already exists` for parametrized modules @@ -201,9 +217,7 @@ def _iter_weight_params(): try: setattr(container, param_name, new_param) except Exception as exc: - logger.warning( - f"Failed to clear expert weight {type(proj).__name__}.{param_name}: {exc}" - ) + logger.warning(f"Failed to clear expert weight {type(proj).__name__}.{param_name}: {exc}") logger.info(f"Replaced {replaced_count} expert weight params") @@ -256,7 +270,9 @@ def _load_kt_weight_index(kt_weight_path: str) -> dict[str, str]: return index -def _dequant_fp8_experts(weights: list[torch.Tensor], scales: list[torch.Tensor | None], block_size: tuple[int, int]) -> torch.Tensor: +def _dequant_fp8_experts( + weights: list[torch.Tensor], scales: list[torch.Tensor | None], block_size: tuple[int, int] +) -> torch.Tensor: """Dequantize a list of FP8 expert weights and stack them (batched, vectorized). Args: @@ -468,9 +484,7 @@ def load_experts_from_kt_weight_path( f"Expected keys like 'blk.{layer_idx}.ffn_gate_exps.0.numa.0.weight'" ) - logger.info( - f"Loading INT8 weights for layer {layer_idx}: {num_experts} experts, {numa_count} NUMA partitions" - ) + logger.info(f"Loading INT8 weights for layer {layer_idx}: {num_experts} experts, {numa_count} NUMA partitions") gate_weights_list = [] gate_scales_list = [] diff --git a/kt-kernel/python/sft/wrapper.py b/kt-kernel/python/sft/wrapper.py index a53ea88af..4e3dd4614 100644 --- a/kt-kernel/python/sft/wrapper.py +++ b/kt-kernel/python/sft/wrapper.py @@ -95,9 +95,7 @@ def build_kt_device_map(config, kt_plugin, device: str = "cuda:0") -> dict[str, else: device_map[expert_key] = "cpu" - logger.info( - f"Built KT device_map: {num_gpu_experts} GPU experts, {num_experts - num_gpu_experts} CPU experts" - ) + logger.info(f"Built KT device_map: {num_gpu_experts} GPU experts, {num_experts - num_gpu_experts} CPU experts") return device_map @@ -163,8 +161,24 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT cfg = _get_kt_config(kt_plugin) # Read lora_rank/lora_alpha for C++ wrapper initialization (buffer allocation only) - lora_rank = getattr(cfg, "kt_lora_rank", 1) or 1 - lora_alpha = getattr(cfg, "kt_lora_alpha", 1.0) or 1.0 + # Use explicit None checks: lora_rank=0 is a valid value (full mode, no LoRA), + # but `or` pattern would treat 0 as falsy and replace it with 1. + _raw_rank = getattr(cfg, "kt_lora_rank", None) + lora_rank = _raw_rank if _raw_rank is not None else 1 + _raw_alpha = getattr(cfg, "kt_lora_alpha", None) + lora_alpha = _raw_alpha if _raw_alpha is not None else 1.0 + + # Read full_weight_grad mode + _raw_fwg = getattr(cfg, "kt_full_weight_grad", None) + full_weight_grad = _raw_fwg if _raw_fwg is not None else False + + # In full mode, lora_rank should be 0 (no LoRA, only base weight grad) + # If user explicitly set lora_rank > 0 in full mode (hybrid), keep it. + # Otherwise, auto-set lora_rank=0. + if full_weight_grad and lora_rank > 0: + _has_explicit_lora_rank = getattr(cfg, "kt_lora_rank", None) is not None + if not _has_explicit_lora_rank: + lora_rank = 0 # Read LoRA Experts configuration _raw_le = getattr(cfg, "kt_use_lora_experts", None) @@ -177,6 +191,8 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT f"LoRA Experts config: use_lora_experts={use_lora_experts}, " f"num={lora_expert_num}, intermediate_size={lora_expert_intermediate_size}" ) + if full_weight_grad: + logger.info(f"Full weight gradient mode enabled (lora_rank={lora_rank})") wrappers: list[KTMoELayerWrapper] = [] moe_layer_count = 0 @@ -225,7 +241,9 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT cfg.kt_sharded_metadata = sharded_metadata logger.info(f"Resolved {len(checkpoint_files)} checkpoint files from kt_expert_checkpoint_path") else: - logger.warning(f"Failed to resolve checkpoint files from kt_expert_checkpoint_path={kt_expert_checkpoint_path!r}") + logger.warning( + f"Failed to resolve checkpoint files from kt_expert_checkpoint_path={kt_expert_checkpoint_path!r}" + ) use_checkpoint_files = bool(checkpoint_files) and not use_kt_weight_path @@ -260,6 +278,7 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT ) import torch.distributed as _dist + _rank = _dist.get_rank() if _dist.is_initialized() else 0 model_container, layers = _get_model_container_and_layers(model, purpose="wrapping") @@ -329,11 +348,23 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT lora_rank=lora_rank, lora_alpha=lora_alpha, max_cache_depth=getattr(cfg, "kt_max_cache_depth", 2), + full_weight_grad=full_weight_grad, ) # Set share_backward_bb and share_cache_pool BEFORE load_weights (config is built during load) wrapper.share_backward_bb = cfg.kt_share_backward_bb - wrapper.share_cache_pool = cfg.kt_share_cache_pool + single_process = not dist.is_initialized() or dist.get_world_size() == 1 + reuse_checkpoint_forward = ( + single_process + and full_weight_grad + and lora_rank == 0 + and kt_method == "AMXBF16_SFT" + and os.environ.get("KT_REUSE_CHECKPOINT_FORWARD", "1") != "0" + ) + wrapper.reuse_checkpoint_forward = reuse_checkpoint_forward + # Reusing the first checkpoint forward requires each layer's C++ + # activations to remain valid until its backward invocation. + wrapper.share_cache_pool = False if reuse_checkpoint_forward else cfg.kt_share_cache_pool physical_to_logical_map = torch.arange(moe_config.expert_num, dtype=torch.int64, device="cpu") @@ -352,9 +383,18 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT physical_to_logical_map_cpu=physical_to_logical_map, ) - wrapper.gate_proj = None - wrapper.up_proj = None - wrapper.down_proj = None + # In full_weight_grad mode, keep weight references for nn.Parameter initialization + # and initialize the base weight buffers + if full_weight_grad: + wrapper.init_full_weight_grad_buffers( + gate_proj=wrapper.gate_proj if wrapper.gate_proj is not None else gate_proj, + up_proj=wrapper.up_proj if wrapper.up_proj is not None else up_proj, + down_proj=wrapper.down_proj if wrapper.down_proj is not None else down_proj, + ) + else: + wrapper.gate_proj = None + wrapper.up_proj = None + wrapper.down_proj = None # Create LoRA Experts if enabled lora_experts = None @@ -381,16 +421,18 @@ def wrap_moe_layers_with_kt_wrapper(model: nn.Module, kt_plugin: Any) -> list[KT setattr(layer, moe_config.moe_layer_attr, layer_wrapper) # Base weights have been copied into the C++ kernel's internal BufferB format. - # Do not hold a Python-side reference --- it wastes ~1 GB/layer. + # In full_weight_grad mode, the authoritative copies are gate_proj_buf etc. + # Always release local references to save ~1 GB/layer. del gate_proj, up_proj, down_proj wrappers.append(layer_wrapper) moe_layer_count += 1 - # Replace original expert weights with meta placeholders. + # Replace original expert weights with zero-storage placeholders. # Experts remain in the model tree (via wrapper.experts) so PEFT can discover them. # Rank 0 already copied weights to C++ kernel via load_weights_from_tensors. - _clear_original_expert_weights(moe_module, moe_config) + # gate_proj_buf serves as the authoritative copy in full_weight_grad mode. + _clear_original_expert_weights(moe_module, moe_config, full_weight_grad=full_weight_grad) logger.info(f"Wrapped {moe_layer_count} MoE layers with KTMoEWrapper") @@ -420,6 +462,17 @@ class should import it from the appropriate dataclasses module. from .config import KTConfig from accelerate.utils.dataclasses import KTransformersPlugin + # Map LlamaFactory finetuning_type to kt_train_mode + finetuning_type = getattr(finetuning_args, "finetuning_type", None) if finetuning_args else None + kt_train_mode_map = { + "full": "full", + "freeze": "hybrid", + "lora": "lora", + "galore": "full", + "badam": "full", + } + kt_train_mode = kt_train_mode_map.get(finetuning_type, None) if finetuning_type else None + kt_config = KTConfig( kt_backend=getattr(model_args, "kt_backend", None), kt_num_threads=getattr(model_args, "kt_num_threads", None), @@ -435,6 +488,7 @@ class should import it from the appropriate dataclasses module. kt_lora_rank=getattr(finetuning_args, "lora_rank", None) if finetuning_args else None, kt_lora_alpha=getattr(finetuning_args, "lora_alpha", None) if finetuning_args else None, kt_model_max_length=getattr(model_args, "model_max_length", None), + kt_train_mode=kt_train_mode, ) return KTransformersPlugin(enabled=True, kt_config=kt_config) @@ -509,7 +563,13 @@ def load_kt_model( **kwargs, ) -> nn.Module: """Load model with KTMoEWrapper backend.""" - from .arch import get_moe_arch_config, move_non_experts_to_gpu, get_expert_device, KTAMXNotAvailableError, KTAMXConfigError + from .arch import ( + get_moe_arch_config, + move_non_experts_to_gpu, + get_expert_device, + KTAMXNotAvailableError, + KTAMXConfigError, + ) if kt_plugin is None: if model_args is None: @@ -536,8 +596,11 @@ def load_kt_model( from transformers.integrations.kt import set_kt_config, unset_kt_config loading_kwargs = get_kt_loading_kwargs( - config, kt_plugin, torch_dtype=torch_dtype, - trust_remote_code=trust_remote_code, token=token, + config, + kt_plugin, + torch_dtype=torch_dtype, + trust_remote_code=trust_remote_code, + token=token, ) if model_args is not None: for key in ("cache_dir", "revision"): @@ -551,8 +614,10 @@ def load_kt_model( if getattr(cfg, "kt_skip_expert_loading", None) is None: checkpoint_files, sharded_metadata = _resolve_checkpoint_files( model_name_or_path=model_name_or_path, - cache_dir=cache_dir, revision=revision, - token=token, trust_remote_code=trust_remote_code, + cache_dir=cache_dir, + revision=revision, + token=token, + trust_remote_code=trust_remote_code, ) if checkpoint_files and all(f.endswith(".safetensors") for f in checkpoint_files): if getattr(cfg, "kt_weight_path", None) is None: diff --git a/kt-kernel/test/per_commit/test_sft_checkpoint_reuse.py b/kt-kernel/test/per_commit/test_sft_checkpoint_reuse.py new file mode 100644 index 000000000..b3d2e3037 --- /dev/null +++ b/kt-kernel/test/per_commit/test_sft_checkpoint_reuse.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch.utils.checkpoint import checkpoint + +from kt_kernel.sft.autograd import KTMoEFunction +from kt_kernel.sft.dist_utils import _checkpoint_hook_mode + + +class _FakeWrapper: + def __init__(self): + self._full_weight_grad = False + self.share_backward_bb = False + self._kt_has_cached_forward = False + self._checkpoint_output_cpu = None + self.submit_calls = 0 + self.sync_calls = 0 + self.cached_output_calls = 0 + self.backward_calls = 0 + + def submit_forward(self, hidden_states, _expert_ids, weights, save_for_backward=True): + self.submit_calls += 1 + self.input = hidden_states.detach().clone() + self.weights = weights.detach().clone() + self.output = self.input * self.weights + assert save_for_backward + + def sync_forward(self, output_device=None): + self.sync_calls += 1 + output = self.output.clone() + return output if output_device is None else output.to(output_device) + + def cache_checkpoint_output(self, output, _qlen): + self._checkpoint_output_cpu = output + self._kt_has_cached_forward = True + + def get_checkpoint_output(self, _qlen, output_device=None): + self.cached_output_calls += 1 + output = self._checkpoint_output_cpu + return output if output_device is None else output.to(output_device) + + def clear_checkpoint_output(self): + self._checkpoint_output_cpu = None + self._kt_has_cached_forward = False + + def backward(self, grad_output, output_device=None): + self.backward_calls += 1 + grad_input = grad_output * self.weights + grad_weights = (grad_output * self.input).sum(dim=-1, keepdim=True) + if output_device is not None: + grad_input = grad_input.to(output_device) + grad_weights = grad_weights.to(output_device) + return grad_input, grad_weights + + +class _CheckpointedExpert(torch.nn.Module): + def __init__(self): + super().__init__() + self.wrapper = _FakeWrapper() + self.route_weights = torch.nn.Parameter(torch.tensor([[0.25], [0.5], [0.75]], dtype=torch.bfloat16)) + + def forward(self, hidden_states): + batch, seq_len, hidden_size = hidden_states.shape + qlen = batch * seq_len + mode = _checkpoint_hook_mode() + cache_forward = mode == "first_forward" + reuse_forward = mode == "recompute" and self.wrapper._kt_has_cached_forward + expert_ids = torch.zeros((batch, seq_len, 1), dtype=torch.int64) + + if not reuse_forward: + self.wrapper.submit_forward( + hidden_states.view(qlen, hidden_size), + expert_ids.view(qlen, 1), + self.route_weights, + save_for_backward=True, + ) + + return KTMoEFunction.apply( + hidden_states, + expert_ids, + self.route_weights, + self.wrapper, + hidden_states.new_empty(()), + hidden_size, + 1, + 0, + True, + False, + None, + cache_forward, + reuse_forward, + None, + None, + None, + ) + + +def test_non_reentrant_checkpoint_reuses_cpu_expert_forward_and_preserves_gradients(): + module = _CheckpointedExpert() + hidden_states = torch.arange(12, dtype=torch.float32).view(1, 3, 4).requires_grad_(True) + + output = checkpoint(module, hidden_states, use_reentrant=False) + output.sum().backward() + + expected_input_grad = module.route_weights.detach().float().view(1, 3, 1).expand_as(hidden_states) + expected_weight_grad = hidden_states.detach().sum(dim=-1).view(3, 1).to(torch.bfloat16) + torch.testing.assert_close(hidden_states.grad, expected_input_grad) + torch.testing.assert_close(module.route_weights.grad, expected_weight_grad) + assert module.wrapper.submit_calls == 1 + assert module.wrapper.sync_calls == 1 + assert module.wrapper.cached_output_calls == 1 + assert module.wrapper.backward_calls == 1 + assert not module.wrapper._kt_has_cached_forward diff --git a/kt-kernel/test/per_commit/test_sft_omp_threads.py b/kt-kernel/test/per_commit/test_sft_omp_threads.py new file mode 100644 index 000000000..e180614dc --- /dev/null +++ b/kt-kernel/test/per_commit/test_sft_omp_threads.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import os +import sys +from pathlib import Path +from unittest.mock import patch + + +CONFIG_PATH = Path(__file__).resolve().parents[2] / "python" / "sft" / "config.py" +SPEC = importlib.util.spec_from_file_location("kt_sft_config_under_test", CONFIG_PATH) +assert SPEC is not None and SPEC.loader is not None +config = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = config +SPEC.loader.exec_module(config) + + +def test_detect_physical_cpu_count_deduplicates_smt_siblings(): + topology = { + 0: (0, 0), + 1: (0, 1), + 2: (0, 0), + 3: (0, 1), + 4: (1, 0), + 5: (1, 0), + } + with ( + patch.object(config, "_available_cpu_ids", return_value=set(topology)), + patch.object(config, "_read_cpu_topology", side_effect=topology.get), + ): + assert config.detect_physical_cpu_count() == 3 + + +def test_configure_omp_threads_replaces_accelerate_single_thread_default(): + with ( + patch.dict(os.environ, {"OMP_NUM_THREADS": "1"}, clear=False), + patch.object(config, "detect_physical_cpu_count", return_value=96), + patch.object(config, "_set_torch_num_threads") as set_torch_threads, + ): + os.environ.pop("ACCELERATE_KT_OMP_NUM_THREADS", None) + assert config.configure_omp_threads() == 96 + assert os.environ["OMP_NUM_THREADS"] == "96" + set_torch_threads.assert_called_once_with(96) + + +def test_configure_omp_threads_preserves_explicit_generic_value(): + with ( + patch.dict(os.environ, {"OMP_NUM_THREADS": "48"}, clear=False), + patch.object(config, "_set_torch_num_threads") as set_torch_threads, + ): + os.environ.pop("ACCELERATE_KT_OMP_NUM_THREADS", None) + assert config.configure_omp_threads() == 48 + set_torch_threads.assert_called_once_with(48) + + +def test_configure_omp_threads_supports_explicit_single_thread_override(): + with ( + patch.dict( + os.environ, + {"OMP_NUM_THREADS": "96", "ACCELERATE_KT_OMP_NUM_THREADS": "1"}, + clear=False, + ), + patch.object(config, "_set_torch_num_threads") as set_torch_threads, + ): + assert config.configure_omp_threads() == 1 + assert os.environ["OMP_NUM_THREADS"] == "1" + set_torch_threads.assert_called_once_with(1) diff --git a/kt-kernel/test/per_commit/test_sft_profiler.py b/kt-kernel/test/per_commit/test_sft_profiler.py new file mode 100644 index 000000000..fc9a0d8b3 --- /dev/null +++ b/kt-kernel/test/per_commit/test_sft_profiler.py @@ -0,0 +1,76 @@ +from types import SimpleNamespace + +from kt_kernel.sft.profiler import collect_kt_sft_profile, format_kt_sft_profile, reset_kt_sft_profile + + +class _FakeMoe: + def __init__(self): + self.reset_calls = 0 + self.get_reset_values = [] + + def get_profile_stats(self, reset=False): + self.get_reset_values.append(reset) + return { + "layer_idx": 3, + "tp_count": 1, + "wrapper.enabled": 1, + "wrapper.tokens": 8, + "wrapper.tp.forward.total.total_ns": 2_000_000, + "wrapper.tp.forward.total.calls": 2, + "wrapper.tp.forward.numa_compute.total_ns": 1_500_000, + "wrapper.tp.forward.numa_compute.calls": 2, + "wrapper.tp.forward.numa_compute.bytes": 2 * 1024 * 1024, + "tp.0.enabled": 1, + "tp.0.tokens": 8, + "tp.0.forward.total.total_ns": 1_400_000, + "tp.0.forward.total.calls": 2, + "tp.0.forward.route.total_ns": 140_000, + "tp.0.forward.route.calls": 2, + "tp.0.backward.total.total_ns": 5_000_000, + "tp.0.backward.total.calls": 1, + "tp.0.backward.down.total.total_ns": 2_000_000, + "tp.0.backward.down.total.calls": 1, + "tp.0.backward.base_weight_grad.total_ns": 1_000_000, + "tp.0.backward.base_weight_grad.calls": 1, + "tp.0.backward.base_weight_grad.worker_cpu.store.total_ns": 2_000_000, + "tp.0.backward.base_weight_grad.worker_cpu.store.calls": 4, + } + + def reset_profile_stats(self): + self.reset_calls += 1 + + +def _fake_model(moe): + backend = SimpleNamespace(moe=moe) + layer = SimpleNamespace(layer_idx=3, wrapper=backend) + return SimpleNamespace(_kt_wrappers=[layer]) + + +def test_collect_and_format_profile(): + moe = _FakeMoe() + profile = collect_kt_sft_profile(_fake_model(moe), reset=True) + + assert profile["enabled"] is True + assert list(profile["layers"]) == [3] + assert moe.get_reset_values == [True] + + output = format_kt_sft_profile(profile) + assert "tp.forward.numa_compute" in output + assert "forward.route" in output + assert "75.0%" in output + assert "2.00" in output + assert "10.0%" in output + assert "40.0%" in output + worker_store = next(line for line in output.splitlines() if "worker_cpu.store" in line) + assert worker_store.endswith("0.0%") + + +def test_reset_and_disabled_profile(): + moe = _FakeMoe() + model = _fake_model(moe) + reset_kt_sft_profile(model) + assert moe.reset_calls == 1 + + empty = collect_kt_sft_profile(SimpleNamespace(_kt_wrappers=[])) + assert empty == {"enabled": False, "layers": {}} + assert "disabled" in format_kt_sft_profile(empty)