From 43257936e2ce2b61e56c8d5b832c25d5b19fdf54 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 14 Aug 2026 22:13:30 +0800 Subject: [PATCH 1/2] feat(hygon): add Qwen3-235B BF16/W8A8 inference support --- csrc/config/model_config.hpp | 20 + csrc/config/quant_config.cpp | 97 ++ csrc/config/quant_config.hpp | 7 + csrc/engine/compiler/paged_compiler.cpp | 56 +- csrc/engine/compiler/paged_compiler.hpp | 2 + csrc/engine/infer_engine.cpp | 44 +- csrc/engine/rank_worker.cpp | 123 +- csrc/engine/rank_worker.hpp | 15 + csrc/layers/attention/backends/flash_attn.cpp | 105 +- csrc/layers/attention/backends/flash_attn.hpp | 26 +- .../causal_lm_templates/text_causal_lm.hpp | 63 +- csrc/layers/linear/base_linear.cpp | 4 + csrc/layers/linear/base_linear.hpp | 1 + csrc/layers/moe/common/moe_types.hpp | 25 +- .../moe/dispatcher/standard_dispatcher.cpp | 7 +- csrc/layers/moe/experts/fused_moe_experts.cpp | 147 ++- csrc/layers/moe/experts/fused_moe_experts.hpp | 8 + csrc/layers/moe/fused_moe.cpp | 37 +- csrc/layers/moe/fused_moe.hpp | 2 +- .../moe/runner/cuda_fused_moe_runner.cpp | 218 +++- .../moe/runner/cuda_fused_moe_runner.hpp | 31 +- csrc/layers/quantization/awq_marlin.hpp | 1 + .../layers/quantization/base_quantization.hpp | 10 + .../quantization/compressed_tensors.cpp | 88 ++ .../quantization/compressed_tensors.hpp | 2 + csrc/layers/quantization/gptq_marlin.hpp | 1 + csrc/models/infinilm_model.cpp | 24 + csrc/models/infinilm_model.hpp | 6 + csrc/models/qwen3/qwen3_attention.cpp | 32 +- csrc/pybind11/engine/engine.hpp | 32 +- .../qwen3moe_w8a8_status_and_vllm_dispatch.md | 147 +++ python/infinilm/base_config.py | 15 +- python/infinilm/infer_engine.py | 132 ++ python/infinilm/modeling_utils.py | 23 +- test/engine/test_nll_validation.py | 145 +++ test/ppl/qwen3_235b/README.md | 200 +++ test/ppl/qwen3_235b/scripts/_gpu_guard.py | 100 ++ test/ppl/qwen3_235b/scripts/_ppl_common.py | 338 +++++ .../calculate_infinilm_precision_ppl.py | 84 ++ .../qwen3_235b/scripts/calculate_true_ppl.py | 305 +++++ .../infinilm/infinilm_ppl_Qwen3_235B.py | 334 +++++ .../scripts/prepare_ppl_corpus_Qwen3_235B.py | 328 +++++ .../scripts/transformers/_pytorch_runner.py | 1144 +++++++++++++++++ .../transformers/pytorch_ppl_Qwen3_235B.py | 491 +++++++ 44 files changed, 4763 insertions(+), 257 deletions(-) create mode 100644 docs/qwen3moe_w8a8_status_and_vllm_dispatch.md create mode 100644 test/engine/test_nll_validation.py create mode 100644 test/ppl/qwen3_235b/README.md create mode 100755 test/ppl/qwen3_235b/scripts/_gpu_guard.py create mode 100755 test/ppl/qwen3_235b/scripts/_ppl_common.py create mode 100644 test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py create mode 100755 test/ppl/qwen3_235b/scripts/calculate_true_ppl.py create mode 100755 test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py create mode 100755 test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py create mode 100755 test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py create mode 100755 test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py diff --git a/csrc/config/model_config.hpp b/csrc/config/model_config.hpp index dc0e8928..5a41a1ac 100644 --- a/csrc/config/model_config.hpp +++ b/csrc/config/model_config.hpp @@ -88,6 +88,26 @@ class ModelConfig { return quant_config.get_quantization_method(); } + std::string get_moe_weight_method() const { + return quant_config.get_moe_weight_method(); + } + + std::string get_moe_weight_method(const infinicore::Device &device) const { + return quant_config.get_moe_weight_method(device); + } + + bool is_moe_w16a16_marlin_enabled() const { + return quant_config.is_moe_w16a16_marlin_enabled(); + } + + bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const { + return quant_config.is_moe_w16a16_marlin_enabled(device); + } + + bool is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const { + return quant_config.is_moe_w8a8_marlin_enabled(device); + } + infinicore::DataType get_dtype() const; infinilm::quantization::QuantScheme get_quant_scheme() const; diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index da261ce0..14cce4af 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -1,6 +1,66 @@ #include "quant_config.hpp" +#include +#include + namespace infinilm::config { +namespace { + +std::string lower_string(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + return value; +} + +bool is_w16a16_marlin_method(const std::string &method) { + return method == "w16a16_marlin" || method == "hygon_w16a16_marlin"; +} + +bool is_w8a8_marlin_method(const std::string &method) { + return method == "slimquant_marlin" || method == "slimquant_compressed_tensors_marlin" || + method == "w8a8_marlin" || method == "hygon_w8a8_marlin"; +} + +bool is_unquantized_config(const nlohmann::json &quantization_config) { + if (quantization_config.is_null()) { + return true; + } + if (!quantization_config.is_object()) { + return false; + } + auto it = quantization_config.find("quant_method"); + if (it == quantization_config.end() || it->is_null()) { + return true; + } + if (!it->is_string()) { + return false; + } + auto method = lower_string(it->get()); + return method.empty() || method == "none" || method == "dense"; +} + +std::string explicit_moe_weight_method(const nlohmann::json &quantization_config) { + if (!quantization_config.is_object()) { + return {}; + } + for (const char *key : {"moe_weight_method", "weight_method", "moe_kernel_method"}) { + auto it = quantization_config.find(key); + if (it != quantization_config.end() && it->is_string()) { + return lower_string(it->get()); + } + } + auto it = quantization_config.find("quant_method"); + if (it != quantization_config.end() && it->is_string()) { + auto method = lower_string(it->get()); + if (is_w16a16_marlin_method(method) || is_w8a8_marlin_method(method)) { + return method; + } + } + return {}; +} + +} // namespace QuantConfig::QuantConfig(const nlohmann::json &json) : quantization_config(json) { this->quantization_method = get_quantization_method(); } @@ -20,6 +80,9 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); } else if (quant_method == "gptq") { return std::make_shared(quantization_config); + } else if (quantization_config["quant_method"] == "w16a16_marlin" || + quantization_config["quant_method"] == "hygon_w16a16_marlin") { + return std::make_shared(quantization_config); } else { return std::make_shared(quantization_config); } @@ -27,4 +90,38 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); // Default case if no matching scheme } + +std::string QuantConfig::get_moe_weight_method() const { + return get_moe_weight_method(infinicore::Device(infinicore::Device::Type::CPU, 0)); +} + +std::string QuantConfig::get_moe_weight_method(const infinicore::Device &device) const { + auto configured_method = explicit_moe_weight_method(quantization_config); + if (!configured_method.empty()) { + return configured_method; + } + if (quantization_method != nullptr) { + auto method = quantization_method->get_moe_weight_method(device); + if (method != "dense") { + return method; + } + } + if (device.getType() == infinicore::Device::Type::HYGON && is_unquantized_config(quantization_config)) { + return "hygon_w16a16_marlin"; + } + return "dense"; +} + +bool QuantConfig::is_moe_w16a16_marlin_enabled() const { + return is_w16a16_marlin_method(get_moe_weight_method()); +} + +bool QuantConfig::is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const { + return is_w16a16_marlin_method(get_moe_weight_method(device)); +} + +bool QuantConfig::is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const { + return is_w8a8_marlin_method(get_moe_weight_method(device)); +} + } // namespace infinilm::config diff --git a/csrc/config/quant_config.hpp b/csrc/config/quant_config.hpp index fb0b8abf..a562ea15 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -1,8 +1,10 @@ #pragma once #include "../utils.hpp" #include "../layers/quantization/quantization.hpp" +#include "infinicore/device.hpp" #include "nlohmann/json.hpp" #include +#include #include namespace infinilm::config { @@ -15,6 +17,11 @@ class QuantConfig { QuantConfig(const nlohmann::json &json); std::shared_ptr get_quantization_method() const; + std::string get_moe_weight_method() const; + std::string get_moe_weight_method(const infinicore::Device &device) const; + bool is_moe_w16a16_marlin_enabled() const; + bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const; + bool is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const; infinilm::quantization::QuantScheme get_quant_scheme() const { if (quantization_method != nullptr) { diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index dee3123c..a7f915d3 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -71,13 +71,17 @@ void PagedCompiler::compile() { } size_t max_batch_size = *std::max_element(decode_batch_sizes_.begin(), decode_batch_sizes_.end()); + decode_graph_needs_runtime_state_reset_ = model_->needs_runtime_state_reset(); compiled_map_decode_.clear(); + // b * ceil(nblocks / b) is at most nblocks + b - 1. All decode + // graphs share this holder and only the selected graph runs at once. block_tables_holder_ = infinicore::Tensor::empty( - {nblocks * max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); + {nblocks + max_batch_size}, infinicore::DataType::I32, infinicore::context::getDevice()); set_zeros(block_tables_holder_); auto make_decode_input = [&](size_t b) { InfinilmModel::Input input; + input.last_token_only = true; input.input_ids = infinicore::Tensor::empty({1, b}, infinicore::DataType::I64, infinicore::context::getDevice()); input.position_ids = infinicore::Tensor::empty( position_id_axes > 1 @@ -98,7 +102,9 @@ void PagedCompiler::compile() { infinicore::context::memcpyH2D(input.input_offsets.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); input.cu_seqlens = infinicore::Tensor::empty({b + 1}, infinicore::DataType::I32, infinicore::context::getDevice()); infinicore::context::memcpyH2D(input.cu_seqlens.value()->data(), input_offsets_vec.data(), (b + 1) * sizeof(int32_t), false); - const size_t block_per_req = nblocks; + // Give each request its fair share of the global cache capacity. + // Wider runtime tables safely fall back to eager in get_compiled(). + const size_t block_per_req = (nblocks + b - 1) / b; input.block_tables = block_tables_holder_->as_strided({b, block_per_req}, {(ptrdiff_t)block_per_req, 1}); input.slot_mapping = infinicore::Tensor::empty({b}, infinicore::DataType::I64, infinicore::context::getDevice()); set_zeros(input.slot_mapping.value()); @@ -150,8 +156,10 @@ void PagedCompiler::compile() { // Warmup runs the eager Marlin path and may leave per-layer lock // workspaces dirty. Reset before CUDA graph capture so capture // starts from the same all-zero lock state as normal execution. - model_->reset_runtime_state(); - infinicore::context::syncStream(); + if (decode_graph_needs_runtime_state_reset_) { + model_->reset_runtime_state(); + infinicore::context::syncStream(); + } } for (size_t b : decode_batch_sizes_) { @@ -164,8 +172,10 @@ void PagedCompiler::compile() { // warmup/capture attempts. This reset is intentionally outside // graph capture; the current implementation still pays a memset // before every graph replay in get_compiled(). - model_->reset_runtime_state(); - infinicore::context::syncStream(); + if (decode_graph_needs_runtime_state_reset_) { + model_->reset_runtime_state(); + infinicore::context::syncStream(); + } infinicore::context::startGraphRecording(); auto output = model_->forward(input); auto graph = infinicore::context::stopGraphRecording(); @@ -192,20 +202,36 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & if (result == compiled_map_decode_.end()) { return {nullptr, nullptr}; } - auto &graph_input = result->second.input; - graph_input.input_ids.value()->copy_from(input.input_ids.value()); - graph_input.position_ids.value()->copy_from(input.position_ids.value()); - graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value()); - graph_input.input_offsets.value()->copy_from(input.input_offsets.value()); - graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value()); + // Decode graphs are captured with one token per request, so their + // input offsets are the fixed sequence [0, 1, ..., batch_size]. + // Reuse the captured tensor only after validating that the runtime + // input has the same layout; otherwise fall back to eager mode. + const auto &runtime_input_offsets = input.input_offsets.value(); + if (!runtime_input_offsets->is_contiguous() || + runtime_input_offsets->size(0) != batch_size + 1) { + return {nullptr, nullptr}; + } + const auto *offsets = reinterpret_cast(runtime_input_offsets->data()); + for (size_t i = 0; i <= batch_size; ++i) { + if (offsets[i] != static_cast(i)) { + return {nullptr, nullptr}; + } + } + auto &graph_input = result->second.input; const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1); if (block_per_req > compiled_block_per_req) { - // Runtime width exceeds compiled graph slot; fall back to eager path. + // Runtime width exceeds compiled graph slot; fall back before + // enqueueing copies that the eager path cannot consume. return {nullptr, nullptr}; } + graph_input.input_ids.value()->copy_from(input.input_ids.value()); + graph_input.position_ids.value()->copy_from(input.position_ids.value()); + graph_input.total_sequence_lengths.value()->copy_from(input.total_sequence_lengths.value()); + graph_input.cu_seqlens.value()->copy_from(input.cu_seqlens.value()); + // Initialize only the active graph rows to -1, then overwrite the // runtime logical region. Avoid clearing the full preallocated // holder on every decode token. @@ -230,7 +256,9 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & // one on the same stream before launch. This is correct but costs // decode latency; the intended follow-up is a reusable global // zero workspace/lock buffer shared by all Marlin layers. - model_->reset_runtime_state(); + if (decode_graph_needs_runtime_state_reset_) { + model_->reset_runtime_state(); + } auto graph = std::get<0>(result->second.compiled); if (graph != nullptr) { diff --git a/csrc/engine/compiler/paged_compiler.hpp b/csrc/engine/compiler/paged_compiler.hpp index a1125864..807f99a5 100644 --- a/csrc/engine/compiler/paged_compiler.hpp +++ b/csrc/engine/compiler/paged_compiler.hpp @@ -18,6 +18,8 @@ class PagedCompiler : public GraphCompiler { infinicore::Tensor block_tables_holder_; + bool decode_graph_needs_runtime_state_reset_ = true; + struct CompiledResult { InfinilmModel::Input input; Compiled compiled; diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index ba01e035..1fb764f9 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -146,6 +146,44 @@ std::vector InferEngine::state_dict_keys() { //------------------------------------------------------ // forward //------------------------------------------------------ +void InferEngine::Input::validate() const { + if (!return_nll) { + if (labels.has_value()) { + throw std::invalid_argument("labels require return_nll=true"); + } + if (score_start != 0) { + throw std::invalid_argument("score_start requires return_nll=true"); + } + return; + } + + if (!input_ids.has_value() || !input_ids.value()) { + throw std::invalid_argument("NLL scoring requires input_ids"); + } + if (!labels.has_value() || !labels.value()) { + throw std::invalid_argument("NLL scoring requires labels"); + } + + const auto &ids = input_ids.value(); + const auto &target = labels.value(); + if (ids->dtype() != infinicore::DataType::I64 + || target->dtype() != infinicore::DataType::I64) { + throw std::invalid_argument("NLL input_ids and labels must use I64 dtype"); + } + if (ids->ndim() != 2 || target->ndim() != 2) { + throw std::invalid_argument("NLL input_ids and labels must be rank-2 tensors"); + } + if (ids->shape() != target->shape()) { + throw std::invalid_argument("NLL input_ids and labels must have identical shapes"); + } + if (ids->size(0) != 1) { + throw std::invalid_argument("NLL scoring currently requires batch_size=1"); + } + if (score_start >= ids->size(1)) { + throw std::invalid_argument("NLL score_start must select at least one token"); + } +} + infinilm::InfinilmModel::Input InferEngine::Input::to_model_input(infinicore::Device device) const { @@ -184,7 +222,7 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { image_req_ids, visual_token_ranges, to_device(target_hidden_states)}; - + input.last_token_only = !sample_all_positions && !return_nll; infinilm::global_state::get_forward_context().attn_metadata = { input.past_sequence_lengths, input.total_sequence_lengths, @@ -206,6 +244,10 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } InferEngine::Output InferEngine::forward(const InferEngine::Input &input) { + // Validate before dispatch so malformed NLL requests cannot fail only one + // rank and leave the remaining workers waiting at a collective. + input.validate(); + // Trigger each worker to run inference for (auto &worker : workers_) { worker->run(input); diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 871b48d7..276fc6d0 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -418,9 +418,13 @@ void RankWorker::thread_loop() { infinicore::Tensor logits; infinicore::Tensor hidden_states; - // All-position speculative/MTP runs need eager mode because - // hidden states are not part of compiled graph outputs. - if (!local_args.sample_all_positions && compiler_ != nullptr && rank_info_.pp_size == 1) { + // Full-position and NLL runs need eager mode because generation + // graphs return last-token logits and omit hidden states. PP graph + // compilation is not supported yet. + if (!local_args.sample_all_positions + && !local_args.return_nll + && compiler_ != nullptr + && rank_info_.pp_size == 1) { auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu())); if (graph != nullptr && output != nullptr) { graph->run(); @@ -435,6 +439,11 @@ void RankWorker::thread_loop() { hidden_states = model_output.hidden_states; } + if (local_args.return_nll && rank_info_.pp_size > 1) { + throw std::runtime_error( + "NLL scoring with pipeline parallelism is not supported"); + } + if (rank_info_.pp_size > 1 && rank_info_.pp_stage + 1 != rank_info_.pp_size) { infinicore::Tensor output_ids; if (rank_info_.pp_stage == 0 && rank_info_.tp_rank == 0) { @@ -464,52 +473,96 @@ void RankWorker::thread_loop() { continue; } - // Random sampling (rank 0 only) + // Sampling and scoring both consume replicated full-vocabulary + // logits, so only rank 0 needs to materialize the result. if (rank_info_.tp_rank == 0) { - auto temperature{local_args.temperature}; - auto top_p{local_args.top_p}; - auto top_k{local_args.top_k}; - const auto &logits_shape{logits->shape()}; + if (logits_shape.size() != 3) { + throw std::runtime_error("InferEngine expected rank-3 logits"); + } const auto &vocab_size{logits_shape[2]}; const auto &total_len{logits_shape[1]}; const auto &batch_size{logits_shape[0]}; - auto n_req = local_args.input_offsets.value()->size(0) - 1; - int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); + if (local_args.return_nll) { + auto labels = local_args.labels.value()->to(rank_info_.device); + if (labels->dtype() != infinicore::DataType::I64 + || labels->ndim() != 2 + || labels->size(0) != batch_size + || labels->size(1) != total_len) { + throw std::runtime_error( + "NLL labels must be I64 with shape [batch, sequence]"); + } - const bool sample_all_positions = local_args.sample_all_positions; - const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; - auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; + const auto score_len = total_len - local_args.score_start; + auto score_logits = logits->narrow( + {{1, local_args.score_start, score_len}}); + auto score_labels = labels->narrow( + {{1, local_args.score_start, score_len}}); + + auto token_nll = infinicore::Tensor::empty( + score_labels->shape(), + infinicore::DataType::F32, + rank_info_.device); + infinicore::op::cross_entropy_( + token_nll, score_logits, score_labels); + token_nll = token_nll->to(infinicore::Device::cpu()); + infinicore::context::syncStream(); + output_ = Output{ + infinicore::Tensor{}, + infinicore::Tensor{}, + infinicore::Tensor{}, + token_nll, + score_len, + }; + } else { + auto temperature{local_args.temperature}; + auto top_p{local_args.top_p}; + auto top_k{local_args.top_k}; + auto n_req = local_args.input_offsets.value()->size(0) - 1; + int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); + + const bool sample_all_positions = local_args.sample_all_positions; + const size_t n_out = sample_all_positions + ? static_cast(input_offsets[n_req]) + : n_req; + auto output_ids{infinicore::Tensor::empty( + {n_out}, infinicore::DataType::I64, rank_info_.device)}; + + for (size_t i{0}; i < n_out; ++i) { + size_t score_idx = i; + if (!sample_all_positions) { + score_idx = total_len == n_req + ? i + : static_cast(input_offsets[i + 1] - 1); + } + auto score{logits->view({batch_size * total_len, vocab_size}) + ->narrow({{0, score_idx, 1}}) + ->view({vocab_size})}; + auto out{output_ids->narrow({{0, i, 1}})->view({})}; + float random_val = std::uniform_real_distribution(0, 1)(rng_); + infinicore::op::random_sample_( + out, score, random_val, top_p, top_k, temperature); + } - for (size_t i{0}; i < n_out; ++i) { - size_t score_idx = i; - if (!sample_all_positions) { - score_idx = static_cast(input_offsets[i + 1] - 1); + if (rank_info_.pp_size > 1) { + infinicore::op::distributed::send( + output_ids, + 0, + rank_info_.world_comm); } - auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; - auto out{output_ids->narrow({{0, i, 1}})->view({})}; - float random_val = std::uniform_real_distribution(0, 1)(rng_); - infinicore::op::random_sample_( - out, score, random_val, top_p, top_k, temperature); - } - if (rank_info_.pp_size > 1) { - infinicore::op::distributed::send( + // Tensor::to(CPU) uses the synchronous D2H contract. + output_ids = output_ids->to(infinicore::Device::cpu()); + output_ = Output{ output_ids, + logits, + hidden_states, + infinicore::Tensor{}, 0, - rank_info_.world_comm); + }; } - - output_ids = output_ids->to(infinicore::Device::cpu()); - - infinicore::context::syncStream(); - - auto out{Output{output_ids, logits, hidden_states}}; - - output_ = std::move(out); } - job_done_ = true; } cv_.notify_all(); diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index d396ef6f..19f2f1f2 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -73,12 +73,25 @@ class RankWorker { /// Sample logits at every packed input position instead of one token per request. bool sample_all_positions{false}; + /// Shifted causal-LM labels. Present only for explicit NLL scoring. + std::optional labels; + + /// First logits/label position included in NLL scoring. + size_t score_start{0}; + + /// Compute token NLL instead of sampling output IDs. + bool return_nll{false}; + float temperature{1}; int top_k{50}; float top_p{1}; + /// Validate invariants shared by Python and native callers before a + /// request is dispatched to any rank worker. + void validate() const; + infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const; }; @@ -86,6 +99,8 @@ class RankWorker { infinicore::Tensor output_ids; infinicore::Tensor logits; infinicore::Tensor hidden_states; + infinicore::Tensor nll; + size_t scored_tokens{0}; }; RankWorker(std::shared_ptr infinilm_config, diff --git a/csrc/layers/attention/backends/flash_attn.cpp b/csrc/layers/attention/backends/flash_attn.cpp index ec7e3772..eb2e3797 100644 --- a/csrc/layers/attention/backends/flash_attn.cpp +++ b/csrc/layers/attention/backends/flash_attn.cpp @@ -2,8 +2,6 @@ #include "../../../utils.hpp" #include "infinicore/ops.hpp" -#include "infinicore/ops/mha_kvcache.hpp" -#include "infinicore/ops/mha_varlen.hpp" namespace infinilm::layers::attention::backends { @@ -13,90 +11,39 @@ FlashAttentionImpl::FlashAttentionImpl(size_t num_heads, size_t num_kv_heads, size_t layer_idx) : num_heads_(num_heads), - head_size_(head_size), scale_(scale), num_kv_heads_(num_kv_heads), - layer_idx_(layer_idx), head_dim_(head_size) { - - const infinilm::global_state::InfinilmConfig &infinilm_config = infinilm::global_state::get_infinilm_config(); - if (!infinilm_config.model_config) { - throw std::runtime_error("infinilm::layers::attention::backends::FlashAttentionImpl: model_config is null"); - } - max_position_embeddings_ = infinilm_config.model_config->get("max_position_embeddings"); -} - -infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, - const infinicore::Tensor &query, - const infinicore::Tensor &key, - const infinicore::Tensor &value, - infinicore::Tensor &kv_cache, - const infinilm::global_state::AttentionMetadata &attn_metadata) const { - auto total_sequence_lengths = attn_metadata.total_sequence_lengths; - auto input_offsets = attn_metadata.input_offsets; - auto block_tables = attn_metadata.block_tables; - auto slot_mapping = attn_metadata.slot_mapping; - auto cu_seqlens = attn_metadata.cu_seqlens; - - ASSERT(block_tables.has_value()); - ASSERT(slot_mapping.has_value()); - - // 1. update paged kv cache - auto [k_total, v_total] = do_kv_cache_update(layer, key, value, kv_cache, slot_mapping.value()); - - size_t seq_len = query->shape()[0]; - bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]); - - // 2. Compute attention - infinicore::Tensor attn_output = infinicore::Tensor::empty({seq_len, num_heads_, head_dim_}, query->dtype(), query->device()); - if (is_prefill) { - infinicore::op::mha_varlen_( - attn_output, - query, - k_total, - v_total, - input_offsets.value(), - cu_seqlens.value(), - block_tables.value(), - max_position_embeddings_, - max_position_embeddings_, - std::nullopt, - scale_); - } else { - // FA2 decode path: flash::mha_fwd_kvcache - // In paged-attn mode, seq_len = actual batch_size (one query token per sequence). - // q_reshaped: [seq_len, num_heads, head_dim] → [seq_len, 1, num_heads, head_dim] - // k/v cache: [num_blocks, block_size, num_kv_heads, head_dim] - auto q_for_fa = query->view({seq_len, 1, num_heads_, head_dim_}); - auto attn_out_4d = infinicore::op::mha_kvcache( - q_for_fa, - k_total, // [num_blocks, block_size, num_kv_heads, head_dim] - v_total, - total_sequence_lengths.value(), // [seq_len] int32 (one entry per sequence) - block_tables.value(), // [seq_len, max_num_blocks_per_seq] int32 - std::nullopt, - scale_); - attn_output = attn_out_4d->view({seq_len, num_heads_, head_dim_}); - } - attn_output = attn_output->view({1, seq_len, num_heads_ * head_dim_}); - return attn_output; + (void)layer_idx; } -std::tuple FlashAttentionImpl::do_kv_cache_update(const AttentionLayer &layer, - const infinicore::Tensor key, - const infinicore::Tensor value, - infinicore::Tensor &kv_cache, - const infinicore::Tensor slot_mapping) const { - auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); - auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); - infinicore::op::paged_caching_( - k_cache_layer->permute({0, 2, 1, 3}), // permute to BHSD for paged_caching_ - v_cache_layer->permute({0, 2, 1, 3}), +infinicore::Tensor FlashAttentionImpl::forward( + const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const { + (void)layer; + + ASSERT(attn_metadata.total_sequence_lengths.has_value()); + ASSERT(attn_metadata.block_tables.has_value()); + ASSERT(attn_metadata.slot_mapping.has_value()); + + return infinicore::op::paged_flash_attention( + query, key, value, - slot_mapping); - - return {k_cache_layer, v_cache_layer}; + kv_cache, + attn_metadata.total_sequence_lengths.value(), + attn_metadata.input_offsets, + attn_metadata.cu_seqlens, + attn_metadata.block_tables.value(), + attn_metadata.slot_mapping.value(), + num_heads_, + num_kv_heads_, + head_dim_, + scale_); } } // namespace infinilm::layers::attention::backends diff --git a/csrc/layers/attention/backends/flash_attn.hpp b/csrc/layers/attention/backends/flash_attn.hpp index 93f61e8b..a592dddb 100644 --- a/csrc/layers/attention/backends/flash_attn.hpp +++ b/csrc/layers/attention/backends/flash_attn.hpp @@ -2,7 +2,6 @@ #include "../../../global_state/global_state.hpp" #include "infinicore/tensor.hpp" -#include namespace infinilm::layers::attention { class AttentionLayer; @@ -29,26 +28,19 @@ class FlashAttentionImpl { * @param attn_metadata: Attention metadata. * @return Attention output, shape `[1, num_tokens, num_heads * head_dim]`. */ - infinicore::Tensor forward(const AttentionLayer &layer, - const infinicore::Tensor &query, - const infinicore::Tensor &key, - const infinicore::Tensor &value, - infinicore::Tensor &kv_cache, - const infinilm::global_state::AttentionMetadata &attn_metadata) const; - - std::tuple do_kv_cache_update(const AttentionLayer &layer, - const infinicore::Tensor key, - const infinicore::Tensor value, - infinicore::Tensor &kv_cache, - const infinicore::Tensor slot_mapping) const; + infinicore::Tensor forward( + const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const; private: size_t num_heads_; - size_t head_size_; float scale_; size_t num_kv_heads_; - size_t layer_idx_; - size_t head_dim_; // Note: head_dim equals to head_size - size_t max_position_embeddings_; + size_t head_dim_; }; + } // namespace infinilm::layers::attention::backends diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp index d359a7f8..135a165a 100644 --- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp +++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp @@ -4,8 +4,12 @@ #include "../../models/infinilm_model.hpp" #include "../linear/linear.hpp" #include "infinicore/device.hpp" +#include "infinicore/ops/distributed/allgather.hpp" +#include "infinicore/ops/select_last_token_hidden_states.hpp" + #include + namespace infinilm::layers::causal_lm_templates { /** @@ -39,10 +43,23 @@ class TextCausalLM : public InfinilmModel { const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); pp_size_ = static_cast(rank_info.pp_size); pp_stage_ = static_cast(rank_info.pp_stage); + tp_size_ = static_cast(rank_info.tp_size); + tp_rank_ = static_cast(rank_info.tp_rank); + vocab_parallel_ = device.getType() == infinicore::Device::Type::HYGON + && tp_size_ > 1 + && vocab_size % tp_size_ == 0; model_ = this->register_module("model", model_config, device); if (is_last_pp_stage()) { - lm_head_ = this->register_module("lm_head", hidden_size, vocab_size, false, dtype, device); + lm_head_ = this->register_module( + "lm_head", + hidden_size, + vocab_size, + false, + dtype, + device, + vocab_parallel_ ? tp_rank_ : 0, + vocab_parallel_ ? tp_size_ : 1); } } @@ -54,7 +71,16 @@ class TextCausalLM : public InfinilmModel { if (!is_last_pp_stage()) { return {infinicore::Tensor(), hidden_states}; } - auto logits = lm_head_->forward(hidden_states); + + if (input.last_token_only) { + if (!input.input_offsets.has_value()) { + throw std::runtime_error("TextCausalLM: last_token_only requires input_offsets"); + } + hidden_states = infinicore::op::select_last_token_hidden_states( + hidden_states, input.input_offsets.value()); + } + + auto logits = gather_logits(lm_head_->forward(hidden_states)); return {logits, hidden_states}; } @@ -62,20 +88,49 @@ class TextCausalLM : public InfinilmModel { if (!lm_head_) { throw std::runtime_error("TextCausalLM::logits_from_hidden called on a non-last pipeline stage"); } - return lm_head_->forward(const_cast(hidden_states)); + return gather_logits( + lm_head_->forward(const_cast(hidden_states))); } Model &model() { return *model_; } protected: INFINICORE_NN_MODULE(Model, model); - INFINICORE_NN_MODULE(infinilm::layers::linear::ReplicatedLinear, lm_head); + INFINICORE_NN_MODULE(infinilm::layers::linear::ColumnParallelLinear, lm_head); private: bool is_last_pp_stage() const { return pp_stage_ + 1 == pp_size_; } + infinicore::Tensor gather_logits(const infinicore::Tensor &local_logits) const { + if (!vocab_parallel_) { + return local_logits; + } + + const auto &local_shape = local_logits->shape(); + if (local_shape.empty() || local_shape.back() == 0) { + throw std::runtime_error("TextCausalLM: invalid local logits shape"); + } + const size_t local_vocab_size = local_shape.back(); + const size_t num_rows = local_logits->numel() / local_vocab_size; + auto local_flat = local_logits->view({num_rows, local_vocab_size}); + const auto &rank_info = + infinilm::global_state::get_tensor_model_parallel_rank_info(); + auto gathered = infinicore::op::distributed::allgather( + local_flat, tp_size_, rank_info.comm); + + auto output_shape = local_shape; + output_shape.back() *= tp_size_; + return gathered->view({tp_size_, num_rows, local_vocab_size}) + ->permute({1, 0, 2}) + ->contiguous() + ->view(output_shape); + } + size_t pp_size_{1}; size_t pp_stage_{0}; + size_t tp_size_{1}; + size_t tp_rank_{0}; + bool vocab_parallel_{false}; }; } // namespace infinilm::layers::causal_lm_templates diff --git a/csrc/layers/linear/base_linear.cpp b/csrc/layers/linear/base_linear.cpp index dc4c77f6..79b10ca3 100644 --- a/csrc/layers/linear/base_linear.cpp +++ b/csrc/layers/linear/base_linear.cpp @@ -88,6 +88,10 @@ void BaseLinear::reset_runtime_state() const { quantization_->reset_runtime_state(); } +bool BaseLinear::needs_runtime_state_reset() const { + return quantization_->needs_runtime_state_reset(); +} + // Backward compatible accessors infinicore::Tensor BaseLinear::weight() const { diff --git a/csrc/layers/linear/base_linear.hpp b/csrc/layers/linear/base_linear.hpp index 8b452544..0566adb8 100644 --- a/csrc/layers/linear/base_linear.hpp +++ b/csrc/layers/linear/base_linear.hpp @@ -47,6 +47,7 @@ class BaseLinear : public infinicore::nn::Module { std::shared_ptr get_quantization() const { return quantization_; } void process_weights_after_loading() override; void reset_runtime_state() const override; + bool needs_runtime_state_reset() const; // Split fused linear parameters into named sub-parameters std::vector split_params( diff --git a/csrc/layers/moe/common/moe_types.hpp b/csrc/layers/moe/common/moe_types.hpp index 8b784cf2..84c9dcbd 100644 --- a/csrc/layers/moe/common/moe_types.hpp +++ b/csrc/layers/moe/common/moe_types.hpp @@ -2,6 +2,7 @@ #include "topk_output.hpp" +#include "infinicore/ops/hygon_moe_marlin.hpp" #include "infinicore/tensor.hpp" #include @@ -20,6 +21,12 @@ enum class CombineInputFormat { DeepEPLL, }; +enum class MoeWeightBackend { + Dense, + HygonW16A16Marlin, + HygonW8A8Marlin, +}; + struct DispatchOutput { DispatchOutputFormat format = DispatchOutputFormat::Standard; infinicore::Tensor hidden_states; @@ -53,6 +60,9 @@ struct CombineInput { struct MoeWeights { infinicore::Tensor packed_w13; infinicore::Tensor packed_w2; + infinicore::Tensor packed_w13_scale; + infinicore::Tensor packed_w2_scale; + MoeWeightBackend backend = MoeWeightBackend::Dense; bool empty() const { return !packed_w13 && !packed_w2; @@ -61,6 +71,19 @@ struct MoeWeights { bool has_packed_dense_weights() const { return packed_w13 && packed_w2; } + + bool has_packed_w8a8_marlin_weights() const { + return packed_w13 && packed_w2 + && packed_w13_scale && packed_w2_scale; + } + + bool is_hygon_w16a16_marlin() const { + return backend == MoeWeightBackend::HygonW16A16Marlin; + } + + bool is_hygon_w8a8_marlin() const { + return backend == MoeWeightBackend::HygonW8A8Marlin; + } }; struct MoeWorkspace { @@ -69,6 +92,7 @@ struct MoeWorkspace { infinicore::Tensor ep_gathered_topk_ids; infinicore::Tensor ep_reduced_hidden_states; infinicore::Tensor fused_moe_output; + infinicore::op::HygonMoeMarlinWorkspace hygon_marlin; infinicore::Tensor sorted_token_ids; infinicore::Tensor expert_ids; @@ -84,7 +108,6 @@ struct MoeWorkspace { size_t expert_ids_capacity = 0; size_t ep_gathered_tokens_capacity = 0; size_t ep_reduced_tokens_capacity = 0; - size_t fused_moe_output_tokens_capacity = 0; size_t blockscale_offsets_capacity = 0; size_t permutation_capacity = 0; size_t prepared_num_experts = 0; diff --git a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp index 27e5c0b2..96bc6205 100644 --- a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp +++ b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp @@ -30,11 +30,8 @@ infinicore::Tensor StandardDispatcher::combine(const CombineInput &combine_input MoeWorkspace &workspace) const { (void)workspace; if (tp_size_ > 1 && communicator_ != nullptr) { - infinicore::op::distributed::allreduce_( - combine_input.hidden_states, - combine_input.hidden_states, - INFINICCL_SUM, - communicator_); + return infinicore::op::distributed::allreduce( + combine_input.hidden_states, INFINICCL_SUM, communicator_); } return combine_input.hidden_states; } diff --git a/csrc/layers/moe/experts/fused_moe_experts.cpp b/csrc/layers/moe/experts/fused_moe_experts.cpp index b456395d..40fc8ec2 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.cpp +++ b/csrc/layers/moe/experts/fused_moe_experts.cpp @@ -3,16 +3,33 @@ #include "../../../global_state/global_state.hpp" #include "../ep/ep_config.hpp" +#include "infinicore/ops/moe_w16a16_marlin.hpp" +#include "infinicore/ops/moe_w8a8_marlin.hpp" + +#include + +#include #include namespace infinilm::layers::moe { FusedMoeExperts::FusedMoeExperts(std::shared_ptr model_config, const infinicore::Device &device) { + device_ = device; num_experts_ = model_config->get("num_experts"); hidden_size_ = model_config->get("hidden_size"); const size_t intermediate_size = model_config->get("moe_intermediate_size"); const auto dtype = model_config->get_dtype(); + const auto moe_weight_method = model_config->get_moe_weight_method(device); + enable_hygon_w16a16_marlin_ = model_config->is_moe_w16a16_marlin_enabled(device); + enable_hygon_w8a8_marlin_ = model_config->is_moe_w8a8_marlin_enabled(device); + if (enable_hygon_w16a16_marlin_ && enable_hygon_w8a8_marlin_) { + throw std::runtime_error("Only one Hygon MoE Marlin weight method can be enabled"); + } + if (moe_weight_method != "dense" && + !enable_hygon_w16a16_marlin_ && !enable_hygon_w8a8_marlin_) { + throw std::runtime_error("Unsupported MoE weight method: " + moe_weight_method); + } ASSERT(num_experts_ > 0); const auto ep_config = make_ep_config(); @@ -32,17 +49,31 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr const size_t expert_tp_rank = ep_enabled ? 0 : tp_rank; const size_t expert_tp_size = ep_enabled ? 1 : tp_size; + const auto expert_weight_dtype = enable_hygon_w8a8_marlin_ ? infinicore::DataType::I8 : dtype; w13_weight_ = infinicore::nn::Parameter( {num_local_experts, intermediate_size_per_partition_ * 2, hidden_size_}, - dtype, + expert_weight_dtype, device); w2_weight_ = infinicore::nn::Parameter( {num_local_experts, hidden_size_, intermediate_size_per_partition_}, - dtype, + expert_weight_dtype, device); this->register_parameter("w13_weight", w13_weight_); this->register_parameter("w2_weight", w2_weight_); + if (enable_hygon_w8a8_marlin_) { + w13_weight_scale_ = infinicore::nn::Parameter( + {num_local_experts, intermediate_size_per_partition_ * 2, 1}, + infinicore::DataType::F32, + device); + w2_weight_scale_ = infinicore::nn::Parameter( + {num_local_experts, hidden_size_, 1}, + infinicore::DataType::F32, + device); + this->register_parameter("w13_weight_scale", w13_weight_scale_); + this->register_parameter("w2_weight_scale", w2_weight_scale_); + } + for (size_t local_expert = 0; local_expert < num_local_experts; ++local_expert) { const size_t global_expert = expert_placement.local_expert_start + local_expert; auto gate_weight = w13_weight_ @@ -65,10 +96,122 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr this->register_parameter( prefix + "down_proj.weight", infinicore::nn::Parameter(down_weight, 1, expert_tp_rank, expert_tp_size)); + + if (enable_hygon_w8a8_marlin_) { + auto gate_scale = w13_weight_scale_ + ->narrow({{0, local_expert, 1}, {1, 0, intermediate_size_per_partition_}}) + ->squeeze(0); + auto up_scale = w13_weight_scale_ + ->narrow({{0, local_expert, 1}, {1, intermediate_size_per_partition_, intermediate_size_per_partition_}}) + ->squeeze(0); + auto down_scale = w2_weight_scale_ + ->narrow({{0, local_expert, 1}}) + ->squeeze(0); + this->register_parameter( + prefix + "gate_proj.weight_scale", + infinicore::nn::Parameter(gate_scale, 0, expert_tp_rank, expert_tp_size)); + this->register_parameter( + prefix + "up_proj.weight_scale", + infinicore::nn::Parameter(up_scale, 0, expert_tp_rank, expert_tp_size)); + this->register_parameter( + prefix + "down_proj.weight_scale", + infinicore::nn::Parameter(down_scale)); + } } moe_weights_.packed_w13 = w13_weight_; moe_weights_.packed_w2 = w2_weight_; + moe_weights_.backend = MoeWeightBackend::Dense; +} + +void FusedMoeExperts::process_weights_after_loading() { + if (enable_hygon_w8a8_marlin_ && !w8a8_marlin_packed_) { + if (device_.getType() != infinicore::Device::Type::HYGON) { + throw std::runtime_error("slimquant_marlin MoE weight method is only supported on HYGON"); + } + const auto ep_config = make_ep_config(); + if (ep_config.backend != EPBackend::Disabled) { + throw std::runtime_error("slimquant_marlin MoE weight method currently supports TP-split experts only; disable MoE EP"); + } + if (!w13_weight_ || !w2_weight_ || !w13_weight_scale_ || !w2_weight_scale_) { + throw std::runtime_error("slimquant_marlin MoE weight method requires loaded int8 w13/w2 weights and scales"); + } + if (w13_weight_->dtype() != infinicore::DataType::I8 || + w2_weight_->dtype() != infinicore::DataType::I8 || + w13_weight_scale_->dtype() != infinicore::DataType::F32 || + w2_weight_scale_->dtype() != infinicore::DataType::F32) { + throw std::runtime_error("slimquant_marlin MoE weight method requires int8 weights and fp32 weight scales"); + } + if (hidden_size_ % 64 != 0 || intermediate_size_per_partition_ % 64 != 0) { + throw std::runtime_error("slimquant_marlin MoE weight method requires hidden/intermediate sizes divisible by 64"); + } + + spdlog::debug( + "Packing MoE weights with Hygon W8A8 slimquant Marlin layout: experts={}, hidden={}, intermediate_per_partition={}", + w13_weight_->size(0), hidden_size_, intermediate_size_per_partition_); + + auto packed_w13 = infinicore::op::moe_w8a8_marlin_pack(w13_weight_); + auto packed_w2 = infinicore::op::moe_w8a8_marlin_pack(w2_weight_); + + parameters_.clear(); + w13_weight_ = infinicore::nn::Parameter(packed_w13); + w2_weight_ = infinicore::nn::Parameter(packed_w2); + w13_weight_scale_ = infinicore::nn::Parameter(w13_weight_scale_); + w2_weight_scale_ = infinicore::nn::Parameter(w2_weight_scale_); + this->register_parameter("w13_weight", w13_weight_); + this->register_parameter("w2_weight", w2_weight_); + this->register_parameter("w13_weight_scale", w13_weight_scale_); + this->register_parameter("w2_weight_scale", w2_weight_scale_); + + moe_weights_.packed_w13 = w13_weight_; + moe_weights_.packed_w2 = w2_weight_; + moe_weights_.packed_w13_scale = w13_weight_scale_; + moe_weights_.packed_w2_scale = w2_weight_scale_; + moe_weights_.backend = MoeWeightBackend::HygonW8A8Marlin; + w8a8_marlin_packed_ = true; + return; + } + + if (!enable_hygon_w16a16_marlin_ || w16a16_marlin_packed_) { + return; + } + if (device_.getType() != infinicore::Device::Type::HYGON) { + throw std::runtime_error("w16a16_marlin MoE weight method is only supported on HYGON"); + } + + const auto ep_config = make_ep_config(); + if (ep_config.backend != EPBackend::Disabled) { + throw std::runtime_error("w16a16_marlin MoE weight method currently supports TP-split experts only; disable MoE EP"); + } + if (!w13_weight_ || !w2_weight_) { + throw std::runtime_error("w16a16_marlin MoE weight method requires loaded dense w13/w2 weights"); + } + if (w13_weight_->dtype() != infinicore::DataType::F16 && + w13_weight_->dtype() != infinicore::DataType::BF16) { + throw std::runtime_error("w16a16_marlin MoE weight method requires FP16 or BF16 weights"); + } + if (hidden_size_ % 32 != 0 || intermediate_size_per_partition_ % 16 != 0 || + (intermediate_size_per_partition_ * 2) % 32 != 0) { + throw std::runtime_error("w16a16_marlin MoE weight method requires aligned hidden/intermediate sizes"); + } + + spdlog::debug( + "Packing MoE weights with Hygon W16A16 Marlin layout: experts={}, hidden={}, intermediate_per_partition={}", + w13_weight_->size(0), hidden_size_, intermediate_size_per_partition_); + + auto packed_w13 = infinicore::op::moe_w16a16_marlin_pack(w13_weight_); + auto packed_w2 = infinicore::op::moe_w16a16_marlin_pack(w2_weight_); + + parameters_.clear(); + w13_weight_ = infinicore::nn::Parameter(packed_w13); + w2_weight_ = infinicore::nn::Parameter(packed_w2); + this->register_parameter("w13_weight", w13_weight_); + this->register_parameter("w2_weight", w2_weight_); + + moe_weights_.packed_w13 = w13_weight_; + moe_weights_.packed_w2 = w2_weight_; + moe_weights_.backend = MoeWeightBackend::HygonW16A16Marlin; + w16a16_marlin_packed_ = true; } const MoeWeights &FusedMoeExperts::moe_weights() const { diff --git a/csrc/layers/moe/experts/fused_moe_experts.hpp b/csrc/layers/moe/experts/fused_moe_experts.hpp index 3f231ec6..a8f2a53b 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.hpp +++ b/csrc/layers/moe/experts/fused_moe_experts.hpp @@ -17,13 +17,21 @@ class FusedMoeExperts : public infinicore::nn::Module { const MoeWeights &moe_weights() const; + void process_weights_after_loading() override; + protected: INFINICORE_NN_PARAMETER(w13_weight); INFINICORE_NN_PARAMETER(w2_weight); + INFINICORE_NN_PARAMETER(w13_weight_scale); + INFINICORE_NN_PARAMETER(w2_weight_scale); size_t num_experts_{0}; size_t hidden_size_{0}; size_t intermediate_size_per_partition_{0}; + bool enable_hygon_w16a16_marlin_{false}; + bool enable_hygon_w8a8_marlin_{false}; + bool w16a16_marlin_packed_{false}; + bool w8a8_marlin_packed_{false}; MoeWeights moe_weights_; }; diff --git a/csrc/layers/moe/fused_moe.cpp b/csrc/layers/moe/fused_moe.cpp index fe1301a8..1e607207 100644 --- a/csrc/layers/moe/fused_moe.cpp +++ b/csrc/layers/moe/fused_moe.cpp @@ -10,6 +10,34 @@ namespace infinilm::layers::moe { +namespace { + +std::shared_ptr make_workspace( + const EPConfig &ep_config, + const std::shared_ptr &model_config, + const infinicore::Device &device) { + const bool use_hygon_marlin = + device.getType() == infinicore::Device::Type::HYGON && + ep_config.backend == EPBackend::Disabled && + (model_config->is_moe_w8a8_marlin_enabled(device) || + model_config->is_moe_w16a16_marlin_enabled(device)); + if (!use_hygon_marlin) { + return std::make_shared(); + } + + // Decoder layers execute sequentially on each rank, including graph replay. + // Reuse their large Marlin scratch buffers instead of retaining one copy per layer. + static thread_local std::weak_ptr shared_workspace; + auto workspace = shared_workspace.lock(); + if (!workspace) { + workspace = std::make_shared(); + shared_workspace = workspace; + } + return workspace; +} + +} // namespace + FusedMoE::FusedMoE(std::shared_ptr model_config, const infinicore::Device &device, size_t layer_id) { @@ -29,21 +57,22 @@ FusedMoE::FusedMoE(std::shared_ptr model_config, intermediate_size_per_partition = intermediate_size / tp_size; } + workspace_ = make_workspace(ep_config, model_config, device); dispatcher_ = make_dispatcher(ep_config, num_experts); runner_ = std::make_shared( expert_placement.local_num_experts, hidden_size, intermediate_size_per_partition, model_config->get_or("moe_align_block_size", 16)); - dispatcher_->initialize(device, workspace_); + dispatcher_->initialize(device, *workspace_); } infinicore::Tensor FusedMoE::forward(const infinicore::Tensor &hidden_states, const TopKOutput &topk_output, const MoeWeights &weights) const { - auto dispatch_output = dispatcher_->dispatch(hidden_states, topk_output, workspace_); - auto combine_input = runner_->run(dispatch_output, weights, workspace_); - return dispatcher_->combine(combine_input, workspace_); + auto dispatch_output = dispatcher_->dispatch(hidden_states, topk_output, *workspace_); + auto combine_input = runner_->run(dispatch_output, weights, *workspace_); + return dispatcher_->combine(combine_input, *workspace_); } } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/fused_moe.hpp b/csrc/layers/moe/fused_moe.hpp index 81b3cad7..7301a5d6 100644 --- a/csrc/layers/moe/fused_moe.hpp +++ b/csrc/layers/moe/fused_moe.hpp @@ -24,7 +24,7 @@ class FusedMoE final : public infinicore::nn::Module { private: std::shared_ptr dispatcher_; std::shared_ptr runner_; - mutable MoeWorkspace workspace_; + std::shared_ptr workspace_; }; } // namespace infinilm::layers::moe diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index c37abcc8..629f9345 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -1,38 +1,48 @@ #include "cuda_fused_moe_runner.hpp" #include "infinicore/context/context.hpp" +#include "infinicore/ops/hygon_moe_marlin.hpp" #include "infinicore/ops/moe_align.hpp" #include "infinicore/ops/moe_fused_dense.hpp" #include #include #include -#include namespace infinilm::layers::moe { -CudaFusedMoeRunner::CudaFusedMoeRunner(size_t num_local_experts, - size_t hidden_size, - size_t intermediate_size_per_partition, - size_t align_block_size) +CudaFusedMoeRunner::CudaFusedMoeRunner( + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size_per_partition, + size_t align_block_size) : num_local_experts_(num_local_experts), hidden_size_(hidden_size), - intermediate_size_per_partition_(intermediate_size_per_partition), + intermediate_size_per_partition_( + intermediate_size_per_partition), align_block_size_(align_block_size) {} namespace { -bool same_device(const infinicore::Tensor &tensor, const infinicore::Device &device) { - return tensor && tensor->device().getType() == device.getType() && tensor->device().getIndex() == device.getIndex(); +bool same_device( + const infinicore::Tensor &tensor, + const infinicore::Device &device) { + return tensor + && tensor->device().getType() == device.getType() + && tensor->device().getIndex() == device.getIndex(); } -void ensure_tensor(infinicore::Tensor &tensor, - const infinicore::Shape &shape, - infinicore::DataType dtype, - const infinicore::Device &device) { - if (!same_device(tensor, device) || tensor->dtype() != dtype || tensor->shape() != shape) { +void ensure_tensor( + infinicore::Tensor &tensor, + const infinicore::Shape &shape, + infinicore::DataType dtype, + const infinicore::Device &device) { + if (!same_device(tensor, device) + || tensor->dtype() != dtype + || tensor->shape() != shape) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE runner workspace tensor was not initialized before graph capture"); + throw std::runtime_error( + "MoE runner workspace tensor was not initialized before graph capture"); } tensor = infinicore::Tensor::empty(shape, dtype, device); } @@ -51,37 +61,86 @@ std::string shape_to_string(const infinicore::Shape &shape) { return oss.str(); } -void check_packed_weight_tensor(const infinicore::Tensor &tensor, - const std::string &name, - const infinicore::Device &device, - const infinicore::DataType dtype, - const infinicore::Shape &shape) { +void check_packed_weight_tensor( + const infinicore::Tensor &tensor, + const std::string &name, + const infinicore::Device &device, + infinicore::DataType dtype, + const infinicore::Shape &shape) { if (!tensor) { - throw std::runtime_error("MoE fused dense core requires " + name); + throw std::runtime_error( + "MoE fused dense core requires " + name); } - if (tensor->device().getType() != device.getType() || tensor->device().getIndex() != device.getIndex()) { - throw std::runtime_error("MoE fused dense core requires packed weights on the hidden_states device"); + if (tensor->device().getType() != device.getType() + || tensor->device().getIndex() != device.getIndex()) { + throw std::runtime_error( + "MoE fused dense core requires packed weights on the hidden_states device"); } if (tensor->dtype() != dtype) { - throw std::runtime_error("MoE fused dense core requires packed weights to have the same dtype as hidden_states"); + throw std::runtime_error( + "MoE fused dense core packed tensor dtype mismatch for " + + name); } if (tensor->shape() != shape) { throw std::runtime_error( - "MoE fused dense core packed weight shape mismatch for " + name + ": expected " + shape_to_string(shape) + ", got " + shape_to_string(tensor->shape())); + "MoE fused dense core packed weight shape mismatch for " + + name + ": expected " + shape_to_string(shape) + + ", got " + shape_to_string(tensor->shape())); } } } // namespace -CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, - const MoeWeights &weights, - MoeWorkspace &workspace) const { - auto runner_input = prepare_runner_input( - dispatch_output, - workspace); +CombineInput CudaFusedMoeRunner::run( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const { + if (weights.is_hygon_w16a16_marlin() + || weights.is_hygon_w8a8_marlin()) { + const auto format = weights.is_hygon_w16a16_marlin() + ? infinicore::op::HygonMoeMarlinWeightFormat::W16A16 + : infinicore::op::HygonMoeMarlinWeightFormat::W8A8; + const infinicore::op::HygonMoeMarlinWeights marlin_weights{ + weights.packed_w13, + weights.packed_w2, + weights.packed_w13_scale, + weights.packed_w2_scale, + format, + }; + const auto output = infinicore::op::hygon_moe_marlin_fused( + dispatch_output.hidden_states, + dispatch_output.topk_output.topk_weights, + dispatch_output.topk_output.topk_ids, + dispatch_output.expert_map, + marlin_weights, + workspace.hygon_marlin, + num_local_experts_, + hidden_size_, + intermediate_size_per_partition_, + align_block_size_); - auto runner_output = run_fused_core(runner_input, weights, workspace); + MoeRoutingMetadata routing_metadata; + if (output.has_routing_metadata) { + routing_metadata.sorted_token_ids = + output.sorted_token_ids; + routing_metadata.expert_ids = output.expert_ids; + routing_metadata.num_tokens_post_padded = + output.num_tokens_post_padded; + } + return CombineInput{ + CombineInputFormat::Standard, + output.hidden_states, + dispatch_output.topk_output, + routing_metadata, + }; + } + auto runner_input = prepare_runner_input( + dispatch_output, + workspace, + align_block_size_); + auto runner_output = + run_fused_core(runner_input, weights, workspace); return CombineInput{ CombineInputFormat::Standard, runner_output.hidden_states, @@ -90,51 +149,77 @@ CombineInput CudaFusedMoeRunner::run(const DispatchOutput &dispatch_output, }; } -CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchOutput &dispatch_output, - MoeWorkspace &workspace) const { - const auto &topk_ids = dispatch_output.topk_output.topk_ids; +CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input( + const DispatchOutput &dispatch_output, + MoeWorkspace &workspace, + size_t block_size) const { + const auto &topk_ids = + dispatch_output.topk_output.topk_ids; const auto &topk_shape = topk_ids->shape(); if (topk_shape.size() != 2) { - throw std::runtime_error("MoE runner requires topk_ids to be a 2D tensor"); + throw std::runtime_error( + "MoE runner requires topk_ids to be a 2D tensor"); } const size_t num_pairs = topk_shape[0] * topk_shape[1]; - const size_t block_size = align_block_size_; const size_t align_num_experts = num_local_experts_ + 1; - const size_t max_num_tokens_padded = num_pairs < align_num_experts - ? num_pairs * block_size - : num_pairs + align_num_experts * (block_size - 1); - const size_t sorted_token_ids_capacity = ((max_num_tokens_padded + 3) / 4) * 4; - const size_t max_num_blocks = (max_num_tokens_padded + block_size - 1) / block_size; + const size_t max_num_tokens_padded = + num_pairs < align_num_experts + ? num_pairs * block_size + : num_pairs + + align_num_experts * (block_size - 1); + const size_t sorted_token_ids_capacity = + ((max_num_tokens_padded + 3) / 4) * 4; + const size_t max_num_blocks = + (max_num_tokens_padded + block_size - 1) / block_size; const auto device = topk_ids->device(); - if (!workspace.sorted_token_ids || workspace.sorted_token_ids_capacity < sorted_token_ids_capacity) { + if (!workspace.sorted_token_ids + || workspace.sorted_token_ids_capacity + < sorted_token_ids_capacity) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE sorted_token_ids workspace was not initialized before graph capture"); + throw std::runtime_error( + "MoE sorted_token_ids workspace was not initialized before graph capture"); } workspace.sorted_token_ids = infinicore::Tensor::empty( - {sorted_token_ids_capacity}, infinicore::DataType::I32, device); - workspace.sorted_token_ids_capacity = sorted_token_ids_capacity; + {sorted_token_ids_capacity}, + infinicore::DataType::I32, + device); + workspace.sorted_token_ids_capacity = + sorted_token_ids_capacity; } - if (!workspace.expert_ids || workspace.expert_ids_capacity < max_num_blocks) { + if (!workspace.expert_ids + || workspace.expert_ids_capacity < max_num_blocks) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE expert_ids workspace was not initialized before graph capture"); + throw std::runtime_error( + "MoE expert_ids workspace was not initialized before graph capture"); } workspace.expert_ids = infinicore::Tensor::empty( - {max_num_blocks}, infinicore::DataType::I32, device); + {max_num_blocks}, + infinicore::DataType::I32, + device); workspace.expert_ids_capacity = max_num_blocks; } if (!workspace.num_tokens_post_padded) { if (infinicore::context::isGraphRecording()) { - throw std::runtime_error("MoE num_tokens_post_padded workspace was not initialized before graph capture"); + throw std::runtime_error( + "MoE num_tokens_post_padded workspace was not initialized before graph capture"); } - workspace.num_tokens_post_padded = infinicore::Tensor::empty( - {1}, infinicore::DataType::I32, device); + workspace.num_tokens_post_padded = + infinicore::Tensor::empty( + {1}, + infinicore::DataType::I32, + device); } + auto sorted_token_ids = + workspace.sorted_token_ids->narrow( + {{0, 0, sorted_token_ids_capacity}}); + auto expert_ids = workspace.expert_ids->narrow( + {{0, 0, max_num_blocks}}); if (dispatch_output.expert_map) { infinicore::op::moe_align_with_expert_map_( - workspace.sorted_token_ids, - workspace.expert_ids, + sorted_token_ids, + expert_ids, workspace.num_tokens_post_padded, topk_ids, dispatch_output.expert_map, @@ -143,8 +228,8 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO true); } else { infinicore::op::moe_align_( - workspace.sorted_token_ids, - workspace.expert_ids, + sorted_token_ids, + expert_ids, workspace.num_tokens_post_padded, topk_ids, num_local_experts_, @@ -155,37 +240,42 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input(const DispatchO dispatch_output.hidden_states, dispatch_output.topk_output, MoeRoutingMetadata{ - workspace.sorted_token_ids, - workspace.expert_ids, + sorted_token_ids, + expert_ids, workspace.num_tokens_post_padded, }, }; } -CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_fused_core(const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, - MoeWorkspace &workspace) const { +CudaFusedMoeRunnerOutput CudaFusedMoeRunner::run_fused_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace) const { if (!weights.has_packed_dense_weights()) { - throw std::runtime_error("MoE fused dense runner requires load-time packed w13/w2 weights"); + throw std::runtime_error( + "MoE fused dense runner requires load-time packed w13/w2 weights"); } check_packed_weight_tensor( weights.packed_w13, "w13", runner_input.hidden_states->device(), runner_input.hidden_states->dtype(), - {num_local_experts_, intermediate_size_per_partition_ * 2, hidden_size_}); + {num_local_experts_, + intermediate_size_per_partition_ * 2, + hidden_size_}); check_packed_weight_tensor( weights.packed_w2, "w2", runner_input.hidden_states->device(), runner_input.hidden_states->dtype(), - {num_local_experts_, hidden_size_, intermediate_size_per_partition_}); + {num_local_experts_, + hidden_size_, + intermediate_size_per_partition_}); ensure_tensor( workspace.fused_moe_output, runner_input.hidden_states->shape(), runner_input.hidden_states->dtype(), runner_input.hidden_states->device()); - workspace.fused_moe_output_tokens_capacity = runner_input.hidden_states->shape()[0]; infinicore::op::moe_fused_dense_( workspace.fused_moe_output, runner_input.hidden_states, diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp index 2a1a1f94..be557fa0 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.hpp @@ -16,22 +16,27 @@ struct CudaFusedMoeRunnerOutput { class CudaFusedMoeRunner final : public MoeRunnerCore { public: - CudaFusedMoeRunner(size_t num_local_experts, - size_t hidden_size, - size_t intermediate_size_per_partition, - size_t align_block_size); + CudaFusedMoeRunner( + size_t num_local_experts, + size_t hidden_size, + size_t intermediate_size_per_partition, + size_t align_block_size); - CombineInput run(const DispatchOutput &dispatch_output, - const MoeWeights &weights, - MoeWorkspace &workspace) const override; + CombineInput run( + const DispatchOutput &dispatch_output, + const MoeWeights &weights, + MoeWorkspace &workspace) const override; private: - CudaFusedMoeRunnerInput prepare_runner_input(const DispatchOutput &dispatch_output, - MoeWorkspace &workspace) const; - - CudaFusedMoeRunnerOutput run_fused_core(const CudaFusedMoeRunnerInput &runner_input, - const MoeWeights &weights, - MoeWorkspace &workspace) const; + CudaFusedMoeRunnerInput prepare_runner_input( + const DispatchOutput &dispatch_output, + MoeWorkspace &workspace, + size_t block_size) const; + + CudaFusedMoeRunnerOutput run_fused_core( + const CudaFusedMoeRunnerInput &runner_input, + const MoeWeights &weights, + MoeWorkspace &workspace) const; size_t num_local_experts_ = 0; size_t hidden_size_ = 0; diff --git a/csrc/layers/quantization/awq_marlin.hpp b/csrc/layers/quantization/awq_marlin.hpp index a22fb43b..5b6bc942 100644 --- a/csrc/layers/quantization/awq_marlin.hpp +++ b/csrc/layers/quantization/awq_marlin.hpp @@ -29,6 +29,7 @@ class AWQMarlin : public BaseQuantization { int tp_rank, int tp_size, int tp_num_heads) const override; void reset_runtime_state() const override; + bool needs_runtime_state_reset() const override { return true; } private: infinicore::Tensor get_workspace( diff --git a/csrc/layers/quantization/base_quantization.hpp b/csrc/layers/quantization/base_quantization.hpp index 4b17cc94..6142a4a3 100644 --- a/csrc/layers/quantization/base_quantization.hpp +++ b/csrc/layers/quantization/base_quantization.hpp @@ -79,6 +79,14 @@ class BaseQuantization : public std::enable_shared_from_this { // Default: raw size is already logical size. virtual size_t get_logical_dim_size(size_t raw_size) const { return raw_size; } + // Optional MoE weight backend selected from the same quantization config as Linear. + // Backends can specialize by device while keeping model code independent from + // vendor-specific kernel names. + virtual std::string get_moe_weight_method(const infinicore::Device &device) const { + (void)device; + return "dense"; + } + // Split fused linear parameters into named sub-parameters (for QKV/GateUp) // params: the fused linear's registered parameters (by name) // splits: description of each shard @@ -111,6 +119,8 @@ class BaseQuantization : public std::enable_shared_from_this { // runtime state can keep the default no-op implementation. virtual void reset_runtime_state() const {} + virtual bool needs_runtime_state_reset() const { return false; } + template T get(const std::string &key) const { if (!quant_config_.contains(key)) { diff --git a/csrc/layers/quantization/compressed_tensors.cpp b/csrc/layers/quantization/compressed_tensors.cpp index ff5617a1..72659554 100644 --- a/csrc/layers/quantization/compressed_tensors.cpp +++ b/csrc/layers/quantization/compressed_tensors.cpp @@ -2,10 +2,78 @@ #include "infinicore/ops/linear_w8a8i8.hpp" #include "infinicore/ops/mul_scalar.hpp" +#include +#include #include #include +#include namespace infinilm::quantization { +namespace { + +std::string lower_string(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char ch) { + return static_cast(std::tolower(ch)); + }); + return value; +} + +bool string_field_equals(const nlohmann::json &json, const char *key, const char *expected) { + auto it = json.find(key); + return it != json.end() && it->is_string() && lower_string(it->get()) == expected; +} + +bool bool_field_equals(const nlohmann::json &json, const char *key, bool expected) { + auto it = json.find(key); + return it != json.end() && it->is_boolean() && it->get() == expected; +} + +bool integer_field_equals(const nlohmann::json &json, const char *key, int expected) { + auto it = json.find(key); + return it != json.end() && it->is_number_integer() && it->get() == expected; +} + +bool has_linear_or_moe_target(const nlohmann::json &group) { + auto targets = group.find("targets"); + if (targets == group.end() || !targets->is_array()) { + return false; + } + for (const auto &target : *targets) { + if (!target.is_string()) { + continue; + } + const auto value = lower_string(target.get()); + if (value == "linear" || value == "fusedmoe") { + return true; + } + } + return false; +} + +bool is_dynamic_token_w8a8_group(const nlohmann::json &group) { + auto weights_it = group.find("weights"); + auto input_it = group.find("input_activations"); + if (weights_it == group.end() || input_it == group.end() || + !weights_it->is_object() || !input_it->is_object()) { + return false; + } + const auto &weights = *weights_it; + const auto &input = *input_it; + const bool weight_ok = + string_field_equals(weights, "type", "int") && + string_field_equals(weights, "strategy", "channel") && + integer_field_equals(weights, "num_bits", 8) && + bool_field_equals(weights, "symmetric", true); + const bool input_ok = + string_field_equals(input, "type", "int") && + string_field_equals(input, "strategy", "token") && + integer_field_equals(input, "num_bits", 8) && + bool_field_equals(input, "dynamic", true) && + bool_field_equals(input, "symmetric", true); + return weight_ok && input_ok; +} + +} // namespace std::vector CompressedTensors::get_param_layout( size_t in_features, size_t out_features, @@ -28,6 +96,26 @@ std::vector CompressedTensors::get_param_layout( return descs; } +std::string CompressedTensors::get_moe_weight_method(const infinicore::Device &device) const { + if (device.getType() != infinicore::Device::Type::HYGON || !quant_config_.is_object()) { + return "dense"; + } + if (!string_field_equals(quant_config_, "quant_method", "compressed-tensors")) { + return "dense"; + } + auto groups = quant_config_.find("config_groups"); + if (groups == quant_config_.end() || !groups->is_object()) { + return "dense"; + } + for (const auto &item : groups->items()) { + const auto &group = item.value(); + if (group.is_object() && has_linear_or_moe_target(group) && is_dynamic_token_w8a8_group(group)) { + return "slimquant_marlin"; + } + } + return "dense"; +} + infinicore::Tensor CompressedTensors::forward( const ParamsMap ¶ms, const infinicore::Tensor &input, diff --git a/csrc/layers/quantization/compressed_tensors.hpp b/csrc/layers/quantization/compressed_tensors.hpp index dcf65c2e..2a088728 100644 --- a/csrc/layers/quantization/compressed_tensors.hpp +++ b/csrc/layers/quantization/compressed_tensors.hpp @@ -25,6 +25,8 @@ class CompressedTensors : public BaseQuantization { bool has_bias, float alpha = 1.0f) const override; + std::string get_moe_weight_method(const infinicore::Device &device) const override; + std::vector split_params( const std::unordered_map ¶ms, const std::vector &splits, diff --git a/csrc/layers/quantization/gptq_marlin.hpp b/csrc/layers/quantization/gptq_marlin.hpp index c16c7943..1fc345b9 100644 --- a/csrc/layers/quantization/gptq_marlin.hpp +++ b/csrc/layers/quantization/gptq_marlin.hpp @@ -31,6 +31,7 @@ class GPTQMarlin : public BaseQuantization { int tp_rank, int tp_size, int tp_num_heads) const override; void reset_runtime_state() const override; + bool needs_runtime_state_reset() const override { return true; } private: infinicore::Tensor get_workspace( diff --git a/csrc/models/infinilm_model.cpp b/csrc/models/infinilm_model.cpp index 5d284a31..2c6a652d 100644 --- a/csrc/models/infinilm_model.cpp +++ b/csrc/models/infinilm_model.cpp @@ -1,6 +1,7 @@ #include "infinilm_model.hpp" #include "../cache/kv_cache.hpp" #include "../global_state/global_state.hpp" +#include "../layers/linear/base_linear.hpp" #include "../utils.hpp" #include @@ -104,6 +105,15 @@ void InfinilmModel::reset_runtime_state() const { } } +bool InfinilmModel::needs_runtime_state_reset() const { + for (const auto &[_, sub] : children()) { + if (needs_runtime_state_reset_recursive_(sub.get())) { + return true; + } + } + return false; +} + void InfinilmModel::process_weights_recursive_(infinicore::nn::Module *module) { for (const auto &[_, sub] : module->children()) { process_weights_recursive_(sub.get()); @@ -118,4 +128,18 @@ void InfinilmModel::reset_runtime_state_recursive_(const infinicore::nn::Module module->reset_runtime_state(); } +bool InfinilmModel::needs_runtime_state_reset_recursive_(const infinicore::nn::Module *module) { + if (const auto *linear = dynamic_cast(module)) { + if (linear->needs_runtime_state_reset()) { + return true; + } + } + for (const auto &[_, sub] : module->children()) { + if (needs_runtime_state_reset_recursive_(sub.get())) { + return true; + } + } + return false; +} + } // namespace infinilm diff --git a/csrc/models/infinilm_model.hpp b/csrc/models/infinilm_model.hpp index 27275ee2..6b7d6f34 100644 --- a/csrc/models/infinilm_model.hpp +++ b/csrc/models/infinilm_model.hpp @@ -55,6 +55,8 @@ class InfinilmModel : public infinicore::nn::Module { std::optional> visual_token_ranges; /// Target model hidden states consumed by draft/MTP models. std::optional target_hidden_states; + /// Return one logit row per request instead of logits for every input token. + bool last_token_only{false}; }; struct Output { @@ -77,6 +79,8 @@ class InfinilmModel : public infinicore::nn::Module { void process_weights_after_loading(); void reset_runtime_state() const; + bool needs_runtime_state_reset() const; + protected: std::vector default_allocate_kv_cache_tensors( const cache::CacheConfig *cache_config, @@ -89,5 +93,7 @@ class InfinilmModel : public infinicore::nn::Module { private: static void process_weights_recursive_(infinicore::nn::Module *module); static void reset_runtime_state_recursive_(const infinicore::nn::Module *module); + + static bool needs_runtime_state_reset_recursive_(const infinicore::nn::Module *module); }; } // namespace infinilm diff --git a/csrc/models/qwen3/qwen3_attention.cpp b/csrc/models/qwen3/qwen3_attention.cpp index 7d9beb04..57b84056 100644 --- a/csrc/models/qwen3/qwen3_attention.cpp +++ b/csrc/models/qwen3/qwen3_attention.cpp @@ -2,6 +2,7 @@ #include "../../global_state/global_state.hpp" #include "../../layers/attention/attention.hpp" #include "../../utils.hpp" +#include "infinicore/ops/rms_rotary_embedding.hpp" namespace infinilm::models::qwen3 { @@ -121,8 +122,6 @@ infinicore::Tensor Qwen3Attention::forward_paged_(const infinicore::Tensor &posi auto q_reshaped = q->view({seq_len, num_attention_heads_, head_dim_}); auto k_reshaped = k->view({seq_len, num_key_value_heads_, head_dim_}); auto v_reshaped = v->view({seq_len, num_key_value_heads_, head_dim_}); - q_reshaped = q_norm_->forward(q_reshaped); - k_reshaped = k_norm_->forward(k_reshaped); // 3. Prepare position_ids for RoPE auto pos_shape = position_ids->shape(); @@ -136,9 +135,32 @@ infinicore::Tensor Qwen3Attention::forward_paged_(const infinicore::Tensor &posi throw std::runtime_error("Unexpected position_ids shape"); } - // 4. Apply RoPE to QK - rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); - rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + // 4. Apply Q/K RMSNorm and RoPE. + const bool can_use_hygon_fused_rms_rope = + qkv_proj_->get_quantization()->get_quant_scheme() == infinilm::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8 + && q_reshaped->device().getType() == infinicore::Device::Type::HYGON + && rotary_emb_->rotary_dim() == head_dim_ + && !rotary_emb_->mrope_section().has_value() + && infinicore::op::rms_rotary_embedding_fuse_available(q_reshaped->device()); + if (can_use_hygon_fused_rms_rope) { + q_reshaped = q_reshaped->contiguous(); + k_reshaped = k_reshaped->contiguous(); + auto pos_ids_fused = pos_ids_for_rope->is_contiguous() ? pos_ids_for_rope : pos_ids_for_rope->contiguous(); + infinicore::op::rms_rotary_embedding_fuse_(q_reshaped, + k_reshaped, + pos_ids_fused, + static_cast(head_dim_), + rotary_emb_->cos_sin_cache(), + rotary_emb_->algo() == infinicore::nn::RoPE::Algo::GPT_NEOX, + q_norm_->weight(), + k_norm_->weight(), + static_cast(q_norm_->eps())); + } else { + q_reshaped = q_norm_->forward(q_reshaped); + k_reshaped = k_norm_->forward(k_reshaped); + rotary_emb_->forward(q_reshaped, pos_ids_for_rope, true); + rotary_emb_->forward(k_reshaped, pos_ids_for_rope, true); + } // 5. Attn Backend calculate auto attn_output = attn_->forward(q_reshaped, k_reshaped, v_reshaped); diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index c5e85577..cc9ecaee 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -161,6 +161,7 @@ inline void bind_infer_engine(py::module &m) { std::optional> visual_token_ranges, std::optional target_hidden_states, bool sample_all_positions, + std::optional labels, py::kwargs kwargs) { InferEngine::Input input{ std::move(input_ids), @@ -181,6 +182,7 @@ inline void bind_infer_engine(py::module &m) { std::move(visual_token_ranges), std::move(target_hidden_states), sample_all_positions, + std::move(labels), }; // Explicit defaults @@ -193,6 +195,8 @@ inline void bind_infer_engine(py::module &m) { "temperature", "top_p", "top_k", + "score_start", + "return_nll", }; for (auto &item : kwargs) { @@ -209,9 +213,24 @@ inline void bind_infer_engine(py::module &m) { input.top_p = py::cast(item.second); } else if (key == "top_k") { input.top_k = py::cast(item.second); + } else if (key == "score_start") { + if (py::isinstance(item.second)) { + throw py::type_error("score_start must be an integer, not bool"); + } + const auto score_start = py::cast(item.second); + if (score_start < 0) { + throw py::value_error("score_start must be non-negative"); + } + input.score_start = static_cast(score_start); + } else if (key == "return_nll") { + if (!py::isinstance(item.second)) { + throw py::type_error("return_nll must be a bool"); + } + input.return_nll = py::cast(item.second); } } + input.validate(); return input; }), py::arg("input_ids") = std::nullopt, @@ -231,7 +250,8 @@ inline void bind_infer_engine(py::module &m) { py::arg("image_req_ids") = std::nullopt, py::arg("visual_token_ranges") = std::nullopt, py::arg("target_hidden_states") = std::nullopt, - py::arg("sample_all_positions") = false) + py::arg("sample_all_positions") = false, + py::arg("labels") = std::nullopt) .def_readwrite("input_ids", &InferEngine::Input::input_ids) .def_readwrite("position_ids", &InferEngine::Input::position_ids) .def_readwrite("past_sequence_lengths", &InferEngine::Input::past_sequence_lengths) @@ -250,6 +270,9 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) + .def_readwrite("labels", &InferEngine::Input::labels) + .def_readwrite("score_start", &InferEngine::Input::score_start) + .def_readwrite("return_nll", &InferEngine::Input::return_nll) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p); @@ -257,7 +280,12 @@ inline void bind_infer_engine(py::module &m) { py::class_(infer_engine, "Output") .def_readwrite("output_ids", &InferEngine::Output::output_ids, "Sampled token IDs") .def_readwrite("logits", &InferEngine::Output::logits, "Raw logits tensor") - .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor"); + .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor") + .def_readwrite("nll", &InferEngine::Output::nll, "Per-token NLL tensor") + .def_readwrite( + "scored_tokens", + &InferEngine::Output::scored_tokens, + "Number of scored tokens"); } } // namespace infinilm::engine diff --git a/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md b/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md new file mode 100644 index 00000000..e8256838 --- /dev/null +++ b/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md @@ -0,0 +1,147 @@ +# Qwen3-MoE W8A8 InfiniLM/vLLM Status + +Date: 2026-07-10 + +## Remote Environment + +- Host: `qinyiqun@10.211.3.28` +- SSH key: `C:\Users\qinyi\.ssh\bw1000` +- Container: `qinyiqun` +- InfiniCore: `/home/qinyiqun/InfiniCore` +- InfiniLM: `/home/qinyiqun/InfiniLM` +- FP model: `/home_aclsylqidf/shared/Qwen3-30B-A3B` +- W8A8 model: `/home_aclsylqidf/shared/Qwen3-30B-A3B-Channel-INT8-w8a8` + +Runtime setup inside the container: + +```bash +unset http_proxy https_proxy all_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY +export PATH=/root/.local/bin:/opt/dtk/cuda/cuda/bin:$PATH +export XMAKE_ROOT=y +export LD_LIBRARY_PATH=/usr/local/lib/python3.10/dist-packages/torch/lib:/root/.infini/lib:${LD_LIBRARY_PATH:-} +export PYTHONPATH=/usr/local/:${PYTHONPATH:-} +``` + +InfiniCore configure must include `--graph=y`: + +```bash +xmake f --hygon-dcu=true --aten=true --flash-attn=/usr/local/lib/python3.10/dist-packages/ --cuda=/opt/dtk/cuda/cuda --ccl=true --graph=y -cv -y +xmake build +xmake install +xmake build _infinicore +xmake install _infinicore +pip install -e . +``` + +## Current Benchmark Contract + +- Model family: `Qwen3-30B-A3B` +- Target path: W8A8 quantized model +- Parallelism: `TP=2`, `DP=1`, `EP=1` +- MoE communication: TP only, no DeepEP/allgather EP path +- Benchmark length going forward: `input_len=4096`, `output_len=1280` +- Important guardrail: pass only one `input_len` value. A comma-separated input length list can hang the current benchmark path. +- Device/profiling tools: `hy-smi` and Hygon trace. + +## Current InfiniLM Findings + +The long-run stall was isolated to the W8A8 MoE path with long prefill. FP graph runs and short W8A8 decode runs can complete, so the issue is not simply long output length. + +Observed behavior before long-prefill slicing: + +- W8A8 `4096/128` graph timed out. +- W8A8 `4096/128` no-graph segfaulted. +- Backtraces showed one rank waiting in `RankWorker::wait`, while the other rank was inside the W8A8 Marlin MoE path and teardown/exit handling. + +The Hygon W8A8 Marlin MoE path now chunks long prefill internally: + +- Files: `csrc/layers/moe/runner/cuda_fused_moe_runner.cpp`, `.hpp` +- Fixed chunk size: `16384` tokens, matching vLLM's production chunk size +- The sliced path is selected before full-input routing metadata is prepared, so each token is aligned only once +- No W8A8 slice or debug environment switches are required + +This keeps long-prefill workspace bounded while decode continues to use the graph-captured Marlin path directly. + +## vLLM W8A8 MoE Path + +vLLM package path: + +- `/usr/local/lib/python3.10/dist-packages/vllm` +- Runtime version in logs: `v0.15.1` + +Important vLLM env: + +- `VLLM_FUSED_MOE_CHUNK_SIZE=16384` +- `VLLM_W8A8_BACKEND=3` + +Main call chain: + +1. `CompressedTensorsW8A8Int8MoEMethod.apply()` +2. `fused_experts(...)` +3. `lmslim.layers.fused_moe.fuse_moe_int8.fused_experts_impl_int8` + +vLLM does not repack Qwen3 MoE weights into the InfiniLM Marlin layout. It keeps ordinary channel-wise int8 tensors: + +- `w1`: `[E, 768, 2048]` +- `w2`: `[E, 2048, 384]` +- `w1_scale`: `[E, 768, 1]` +- `w2_scale`: `[E, 2048, 1]` +- `E=128`, `top_k=8` + +Operator sequence per chunk: + +1. Per-token quantize hidden states. +2. Align/count/sort tokens by expert. +3. GEMM1: `lightop.moe_gemm_w8a8(...)` +4. Activation and quantize: `fuse_silu_mul_quant(...)` +5. GEMM2: `lightop.moe_gemm_w8a8(...)` +6. Reduce top-k outputs: `moe_sum` / `moe_reduce_dispatch` + +## Representative vLLM Size Dispatch + +These are the useful anchor cases for InfiniLM implementation. We do not need to reproduce every tiny graph-capture size immediately. + +| Effective M | GEMM1 shape | GEMM1 config/kernel | GEMM2 shape | GEMM2 config/kernel | Notes | +| --- | --- | --- | --- | --- | --- | +| `1..32` | `N=768,K=2048` | small-M `lightop.moe_gemm_w8a8`, often `BLOCK_M=16` | `N=2048,K=384` | small-M `lightop.moe_gemm_w8a8` | decode/graph capture sizes | +| `896` | `N=768,K=2048` | `BLOCK_M=64, MODE=517, DELTA=1`, HIP NT prefill up | `N=2048,K=384` | `BLOCK_M=32, MODE=568, DELTA=2`, HIP NT prefill down | tail chunk | +| `4096` | `N=768,K=2048` | `BLOCK_M=128, MODE=1000, DELTA=1`, `MOE_W8A8_I8_PERCHANNEL_ASM_TN_MT128x256x128_WGM1_UP` | `N=2048,K=384` | `BLOCK_M=64, MODE=517, DELTA=2`, HIP NT prefill down | target single request prefill | +| `10240` | `N=768,K=2048` | `BLOCK_M=128, MODE=1000, DELTA=1`, same ASM UP kernel | `N=2048,K=384` | `BLOCK_M=64, MODE=523, DELTA=2`, HIP NT prefill down | vLLM chunked prefill example | + +vLLM with 16 concurrent 8K prompts enabled chunked prefill with `max_num_batched_tokens=10240`. The observed MoE effective sizes were `10240`, `8256`, `896`, plus small graph-capture sizes. This means scheduler chunking, not only `VLLM_FUSED_MOE_CHUNK_SIZE`, controls the actual large-M MoE calls. + +## vLLM W8A8 Dense Linear Path + +Main call chain: + +1. `CompressedTensorsW8A8Int8.apply_weights()` +2. `apply_int8_linear(..., w8a8_strategy=3)` +3. `per_token_quant_int8(...)` +4. `ops.blaslt_scaled_mm(...)` +5. backend 3: `hipblaslt_w8a8_channelwise_gemm` + +Representative kernels: + +- `M=1,N=4096,K=2048`: small `Cijk_Alik_Bljk_I8BS_MT64x16x256...` +- `M=4096,N=4096,K=2048`: large `Cijk_Alik_Bljk_I8BS_MT256x256x128...` + +## Implementation Direction + +The next code change should move InfiniLM W8A8 MoE toward vLLM's ordinary channel-wise path: + +1. Add a new W8A8 channel MoE backend in InfiniLM, keeping `[E,N,K]` weights and `[E,N,1]` scales instead of calling `moe_w8a8_marlin_pack`. +2. Add/route an InfiniCore wrapper around ordinary `lightop.moe_gemm_w8a8`, not the current `moe_gemm_marlin_w8a8` adaptor. +3. Reuse the existing MoE workspace pattern where possible: int8 hidden cache, int8 intermediate cache, per-token scales, BF16 intermediate/output buffers. +4. Select configs by effective `M` and GEMM shape, matching the vLLM anchors above first: small decode, `896`, `4096`, `10240`. +5. Default the MoE chunk cap to `16384` for the ordinary channel path, matching vLLM's fused MoE chunk cap. Scheduler-level chunking is still needed later for 16 concurrency and 8K-10K contexts. + +## Useful Remote Artifacts + +- vLLM probe log: `/tmp/vllm_w8a8_kernel_probe_i8192_c16_o16_20260710_110631.server.log` +- MoE micro traces: + - `/tmp/hygon_trace_lmslim_w8a8_moe_m10240_20260710_111507` + - `/tmp/hygon_trace_lmslim_w8a8_moe_m896_20260710_111603` + - `/tmp/hygon_trace_lmslim_w8a8_moe_m4096_20260710_111646` +- Dense linear micro traces: + - `/tmp/hygon_trace_vllm_w8a8_linear_m1_n4096_k2048_20260710_113331` + - `/tmp/hygon_trace_vllm_w8a8_linear_m4096_n4096_k2048_20260710_113413` diff --git a/python/infinilm/base_config.py b/python/infinilm/base_config.py index 4ae7665c..bbcd0233 100644 --- a/python/infinilm/base_config.py +++ b/python/infinilm/base_config.py @@ -79,7 +79,15 @@ def __init__(self): self.use_mla = self.args.use_mla self.pre_transpose = self.args.pre_transpose self.num_blocks = self.args.num_blocks - self.block_size = self.args.block_size + if self.args.block_size is None: + platform = ( + self.detect_device() + if self.device.lower() == "auto" + else self.device.lower() + ) + self.block_size = 64 if platform == "hygon" else 256 + else: + self.block_size = self.args.block_size self.max_cache_len = self.args.max_cache_len self.kv_cache_dtype = self.args.kv_cache_dtype self.skip_load = self.args.skip_load @@ -278,7 +286,10 @@ def _add_common_args(self): "--num-blocks", type=int, default=512, help="number of KV cache blocks" ) self.parser.add_argument( - "--block-size", type=int, default=256, help="size of each KV cache block" + "--block-size", + type=int, + default=None, + help="size of each KV cache block (default: 64 on Hygon, 256 otherwise)", ) self.parser.add_argument( "--max-cache-len", type=int, default=4096, help="maximum cache length" diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index 11bcdaa0..fb36dc31 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -1,4 +1,5 @@ import json +import operator import os import time from dataclasses import dataclass @@ -18,6 +19,45 @@ } +def _validate_nll_score_inputs(input_ids, labels, score_start): + """Validate an explicit batch-1, shifted-token NLL request.""" + for name, tensor in (("input_ids", input_ids), ("labels", labels)): + if tensor is None: + raise TypeError(f"{name} must be an infinicore tensor") + missing = [ + attr + for attr in ("ndim", "shape", "dtype", "_underlying") + if not hasattr(tensor, attr) + ] + if missing: + raise TypeError( + f"{name} must be an infinicore tensor; missing {', '.join(missing)}" + ) + if tensor.dtype != infinicore.int64: + raise ValueError(f"{name} must use infinicore.int64 dtype") + if tensor.ndim != 2: + raise ValueError(f"{name} must be a rank-2 tensor") + + input_shape = tuple(input_ids.shape) + label_shape = tuple(labels.shape) + if input_shape != label_shape: + raise ValueError("input_ids and labels must have identical shapes") + if input_shape[0] != 1: + raise ValueError("score_nll currently requires batch_size=1") + + if isinstance(score_start, bool): + raise TypeError("score_start must be an integer, not bool") + try: + score_start = operator.index(score_start) + except TypeError as error: + raise TypeError("score_start must be an integer") from error + + seq_len = input_shape[1] + if score_start < 0 or score_start >= seq_len: + raise ValueError("score_start must select at least one token") + return seq_len, score_start + + def _apply_torch_dtype_defaults(config: dict) -> dict: if config.get("torch_dtype") is None: config["torch_dtype"] = config.get("dtype") or _MODEL_DEFAULTS.get( @@ -684,6 +724,98 @@ def generate( return output_ids + def score_nll(self, input_ids, labels, *, score_start=0): + """Return summed shifted-token NLL and token count for a batch-1 window. + + This explicit evaluation path bypasses graph replay and never changes the + behavior of ``forward``/``generate``. ``input_ids`` and ``labels`` must + have the same ``[1, sequence]`` shape; callers perform the causal shift. + """ + try: + seq_len, score_start = _validate_nll_score_inputs( + input_ids, labels, score_start + ) + + block_tables = None + slot_mapping = None + if self.enable_paged_attn: + cache_config = self.get_cache_config() + if cache_config is None: + raise RuntimeError("paged attention requires a cache configuration") + paged_block_size = cache_config.block_size() + max_blocks_per_batch = ( + seq_len + paged_block_size - 1 + ) // paged_block_size + if max_blocks_per_batch > cache_config.num_blocks(): + raise ValueError( + "NLL sequence requires more paged KV-cache blocks than " + "the current cache configuration provides" + ) + block_tables = infinicore.from_list( + [list(range(max_blocks_per_batch))], + dtype=infinicore.int32, + ) + slot_mapping = infinicore.from_list( + list(range(seq_len)), dtype=infinicore.int64 + ) + position_ids = infinicore.from_list( + list(range(seq_len)), dtype=infinicore.int64 + ) + else: + position_ids = infinicore.from_list( + [list(range(seq_len))], dtype=infinicore.int64 + ) + past_kv_lengths = infinicore.from_list([0], dtype=infinicore.int32) + total_kv_lengths = infinicore.from_list( + [seq_len], dtype=infinicore.int32 + ) + cu_seqlens = infinicore.from_list( + [0, seq_len], dtype=infinicore.int32 + ) + input_offsets = infinicore.from_list( + [0, seq_len], dtype=infinicore.int32 + ) + + output = super().forward( + super().Input( + input_ids._underlying, + position_ids=position_ids._underlying, + past_sequence_lengths=past_kv_lengths._underlying, + total_sequence_lengths=total_kv_lengths._underlying, + input_offsets=input_offsets._underlying, + cu_seqlens=cu_seqlens._underlying, + block_tables=( + block_tables._underlying + if block_tables is not None + else None + ), + slot_mapping=( + slot_mapping._underlying + if slot_mapping is not None + else None + ), + labels=labels._underlying, + score_start=score_start, + return_nll=True, + ) + ) + token_nll = infinicore.Tensor(output.nll).to_numpy() + scored_tokens = int(output.scored_tokens) + expected_scored_tokens = seq_len - score_start + if scored_tokens != expected_scored_tokens: + raise RuntimeError( + "score_nll returned an invalid scored-token count: " + f"expected {expected_scored_tokens}, got {scored_tokens}" + ) + if token_nll.size != expected_scored_tokens: + raise RuntimeError( + "score_nll returned a token-loss vector with an invalid size" + ) + return float(token_nll.astype("float64").sum()), scored_tokens + except BaseException as e: + handle_oom_and_exit(e) + raise + def reset_cache(self, cache_config): infinicore.sync_device() self.enable_paged_attn = isinstance(cache_config, PagedKVCacheConfig) diff --git a/python/infinilm/modeling_utils.py b/python/infinilm/modeling_utils.py index 486df03b..70a9b21f 100644 --- a/python/infinilm/modeling_utils.py +++ b/python/infinilm/modeling_utils.py @@ -58,13 +58,17 @@ def parse_dtype(dtype_str: str): def _is_internal_moe_packed_weight(key: str) -> bool: # InfiniLM registers packed MoE parameters internally. HF checkpoints - # provide per-expert gate/up/down weights instead, so these packed tensors - # are expected missing keys during non-strict checkpoint loading. - return ( - key.endswith(".mlp.experts.w13_weight") - or key.endswith(".mlp.experts.w2_weight") - or key.endswith(".mlp.experts.w1") - or key.endswith(".mlp.experts.w2") + # provide per-expert gate/up/down weights and scales, so these internal + # packed tensors are expected missing keys during non-strict loading. + return key.endswith( + ( + ".mlp.experts.w13_weight", + ".mlp.experts.w2_weight", + ".mlp.experts.w1", + ".mlp.experts.w2", + ".mlp.experts.w13_weight_scale", + ".mlp.experts.w2_weight_scale", + ) ) @@ -268,7 +272,10 @@ def load_model_state_dict_by_file( # --------------------------------------------------------- # model_param_infini = {} for key in model_param.keys(): - model_param_infini[key] = infinicore.from_torch(model_param[key]) + tensor = model_param[key] + if key.endswith(".weight_scale") and tensor.dtype != torch.float32: + tensor = tensor.to(torch.float32) + model_param_infini[key] = infinicore.from_torch(tensor) model.load_state_dict(model_param_infini, strict=False) infinicore.sync_device() del model_param_infini diff --git a/test/engine/test_nll_validation.py b/test/engine/test_nll_validation.py new file mode 100644 index 00000000..cc2f64ab --- /dev/null +++ b/test/engine/test_nll_validation.py @@ -0,0 +1,145 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + + +class FakeTensor: + def __init__(self, shape, dtype): + self.shape = tuple(shape) + self.ndim = len(self.shape) + self.dtype = dtype + self._underlying = object() + + +@pytest.fixture +def validator(monkeypatch): + """Load the pure validator without importing the hardware runtime.""" + package_root = Path(__file__).resolve().parents[2] / "python" / "infinilm" + int64_dtype = object() + + infinilm_package = types.ModuleType("infinilm") + infinilm_package.__path__ = [str(package_root)] + monkeypatch.setitem(sys.modules, "infinilm", infinilm_package) + + fake_infinicore = types.ModuleType("infinicore") + fake_infinicore.int64 = int64_dtype + fake_infinicore.Tensor = type("Tensor", (), {}) + monkeypatch.setitem(sys.modules, "infinicore", fake_infinicore) + + cache_module = types.ModuleType("infinilm.cache") + cache_module.PagedKVCacheConfig = type("PagedKVCacheConfig", (), {}) + monkeypatch.setitem(sys.modules, "infinilm.cache", cache_module) + + distributed_module = types.ModuleType("infinilm.distributed") + distributed_module.DistConfig = type("DistConfig", (), {}) + monkeypatch.setitem(sys.modules, "infinilm.distributed", distributed_module) + + native_engine = type("InferEngine", (), {}) + lib_module = types.ModuleType("infinilm.lib") + lib_module._infinilm = types.SimpleNamespace(InferEngine=native_engine) + monkeypatch.setitem(sys.modules, "infinilm.lib", lib_module) + + exception_module = types.ModuleType("infinilm.exception_utils") + exception_module.handle_oom_and_exit = lambda error: None + monkeypatch.setitem(sys.modules, "infinilm.exception_utils", exception_module) + + modeling_module = types.ModuleType("infinilm.modeling_utils") + modeling_module.parse_dtype = lambda dtype: dtype + monkeypatch.setitem(sys.modules, "infinilm.modeling_utils", modeling_module) + + module_name = "infinilm.infer_engine" + module_path = package_root / "infer_engine.py" + spec = importlib.util.spec_from_file_location(module_name, module_path) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + assert spec.loader is not None + spec.loader.exec_module(module) + + return module._validate_nll_score_inputs, int64_dtype + + +def make_tensor(shape, int64_dtype, dtype=None): + return FakeTensor(shape, int64_dtype if dtype is None else dtype) + + +def test_validate_nll_score_inputs_accepts_valid_window(validator): + validate, int64_dtype = validator + input_ids = make_tensor((1, 8), int64_dtype) + labels = make_tensor((1, 8), int64_dtype) + + assert validate(input_ids, labels, 3) == (8, 3) + + +@pytest.mark.parametrize("name", ["input_ids", "labels"]) +def test_validate_nll_score_inputs_requires_tensor_protocol(validator, name): + validate, int64_dtype = validator + tensors = { + "input_ids": make_tensor((1, 8), int64_dtype), + "labels": make_tensor((1, 8), int64_dtype), + } + tensors[name] = object() + + with pytest.raises(TypeError, match=name): + validate(tensors["input_ids"], tensors["labels"], 0) + + +@pytest.mark.parametrize("name", ["input_ids", "labels"]) +def test_validate_nll_score_inputs_requires_int64(validator, name): + validate, int64_dtype = validator + tensors = { + "input_ids": make_tensor((1, 8), int64_dtype), + "labels": make_tensor((1, 8), int64_dtype), + } + tensors[name] = make_tensor((1, 8), int64_dtype, dtype=object()) + + with pytest.raises(ValueError, match=f"{name} must use infinicore.int64"): + validate(tensors["input_ids"], tensors["labels"], 0) + + +@pytest.mark.parametrize( + ("input_shape", "label_shape", "message"), + [ + ((8,), (8,), "rank-2"), + ((1, 8), (1, 7), "identical shapes"), + ((2, 8), (2, 8), "batch_size=1"), + ], +) +def test_validate_nll_score_inputs_rejects_invalid_shapes( + validator, input_shape, label_shape, message +): + validate, int64_dtype = validator + with pytest.raises(ValueError, match=message): + validate( + make_tensor(input_shape, int64_dtype), + make_tensor(label_shape, int64_dtype), + 0, + ) + + +@pytest.mark.parametrize("score_start", [-1, 8]) +def test_validate_nll_score_inputs_rejects_empty_score_range( + validator, score_start +): + validate, int64_dtype = validator + with pytest.raises(ValueError, match="select at least one token"): + validate( + make_tensor((1, 8), int64_dtype), + make_tensor((1, 8), int64_dtype), + score_start, + ) + + +@pytest.mark.parametrize("score_start", [True, 1.5, "1"]) +def test_validate_nll_score_inputs_requires_integer_score_start( + validator, score_start +): + validate, int64_dtype = validator + with pytest.raises(TypeError, match="score_start must be an integer"): + validate( + make_tensor((1, 8), int64_dtype), + make_tensor((1, 8), int64_dtype), + score_start, + ) diff --git a/test/ppl/qwen3_235b/README.md b/test/ppl/qwen3_235b/README.md new file mode 100644 index 00000000..559fd7c1 --- /dev/null +++ b/test/ppl/qwen3_235b/README.md @@ -0,0 +1,200 @@ +# Qwen3-235B true PPL CLI + +This directory contains reproducible token-level perplexity tools for: + +- Transformers BF16 on TP8 +- InfiniLM BF16 on TP8 +- InfiniLM W8A8 on TP8 + +The runners consume the same frozen token manifest and calculate causal, +shifted-token cross entropy: + +```text +mean_nll = sum(-log p(x_t | x_&1 | tee "$LOG_DIR/infinilm_w8a8_smoke.log" +rc=${PIPESTATUS[0]} +echo "INFINILM_W8A8_SMOKE_EXIT_CODE=$rc" +hy-smi --showpids +``` + +## Full WikiText-2 runs + +`--max-scored-tokens 0` scores every target token in the manifest. Use the same +`window`, `stride` and `max-scored-tokens` values for every backend. + +Transformers BF16: + +```bash +set -o pipefail +timeout --signal=TERM --kill-after=60s 21600s \ + python -u "$PPL_ROOT/scripts/transformers/pytorch_ppl_Qwen3_235B.py" \ + --model "$MODEL_BF16" \ + --token-manifest "$TOKEN_MANIFEST" \ + --window 256 \ + --stride 128 \ + --max-scored-tokens 0 \ + --tp-size 8 \ + --attention eager \ + --json-output "$LOG_DIR/transformers_bf16_full.json" \ + 2>&1 | tee "$LOG_DIR/transformers_bf16_full.log" +``` + +The Transformers entry point launches `torchrun` itself. Do not wrap it in a +second `torchrun` command. Eager attention is the validated Hygon path. + +InfiniLM BF16: + +```bash +set -o pipefail +timeout --signal=TERM --kill-after=60s 21600s \ + python -u "$PPL_ROOT/scripts/infinilm/infinilm_ppl_Qwen3_235B.py" \ + --model "$MODEL_BF16" \ + --token-manifest "$TOKEN_MANIFEST" \ + --window 256 \ + --stride 128 \ + --max-scored-tokens 0 \ + --tp-size 8 \ + --attention flash-attn \ + --json-output "$LOG_DIR/infinilm_bf16_full.json" \ + 2>&1 | tee "$LOG_DIR/infinilm_bf16_full.log" +``` + +InfiniLM W8A8: + +```bash +set -o pipefail +timeout --signal=TERM --kill-after=60s 21600s \ + python -u "$PPL_ROOT/scripts/infinilm/infinilm_ppl_Qwen3_235B.py" \ + --model "$MODEL_W8A8" \ + --token-manifest "$TOKEN_MANIFEST" \ + --window 256 \ + --stride 128 \ + --max-scored-tokens 0 \ + --tp-size 8 \ + --attention flash-attn \ + --json-output "$LOG_DIR/infinilm_w8a8_full.json" \ + 2>&1 | tee "$LOG_DIR/infinilm_w8a8_full.log" +``` + +For a bounded formal run, replace `0` with the same positive token count in all +three commands, for example `10240`. + +## Compare results + +Transformers BF16 versus InfiniLM W8A8: + +```bash +python -u "$PPL_ROOT/scripts/calculate_true_ppl.py" \ + --inputs \ + "$LOG_DIR/transformers_bf16_full.json" \ + "$LOG_DIR/infinilm_w8a8_full.json" \ + --max-ppl-increase-percent 20 \ + --json-out "$LOG_DIR/ppl_transformers_vs_w8a8.json" +``` + +InfiniLM BF16 versus InfiniLM W8A8: + +```bash +python -u "$PPL_ROOT/scripts/calculate_infinilm_precision_ppl.py" \ + --inputs \ + "$LOG_DIR/infinilm_bf16_full.json" \ + "$LOG_DIR/infinilm_w8a8_full.json" \ + --max-ppl-increase-percent 20 \ + --json-out "$LOG_DIR/ppl_bf16_vs_w8a8.json" +``` + +Exit code `0` means the configured PPL increase threshold passed, `1` means it +failed, and `2` means the input files are invalid or describe different +workloads. + +## Scope + +PPL is a quality test. InfiniLM intentionally disables graph only for the +explicit `score_nll` path because it must retain full token logits/losses. +Normal generation and formal performance tests keep their existing graph path. +Do not report PPL scoring throughput as inference performance. diff --git a/test/ppl/qwen3_235b/scripts/_gpu_guard.py b/test/ppl/qwen3_235b/scripts/_gpu_guard.py new file mode 100755 index 00000000..1ceaeddd --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/_gpu_guard.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +"""Fail closed unless all eight Hygon devices are idle.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + + +EXPECTED_DEVICES = set(range(8)) +SMI_TIMEOUT_SECONDS = 60 + + +def _local_gpu_processes() -> list[str]: + users: list[str] = [] + own_pid = os.getpid() + for process_dir in Path("/proc").glob("[0-9]*"): + try: + pid = int(process_dir.name) + except ValueError: + continue + if pid == own_pid: + continue + try: + targets = [entry.resolve() for entry in (process_dir / "fd").iterdir()] + except OSError: + continue + if not any( + str(target) == "/dev/kfd" or str(target).startswith("/dev/dri/renderD") + for target in targets + ): + continue + try: + command = (process_dir / "cmdline").read_bytes().replace(b"\0", b" ").decode( + "utf-8", errors="replace" + ).strip() + except OSError: + command = "" + users.append(f"pid={pid} command={command or '[unknown]'}") + return sorted(users) + + +def require_idle_gpu() -> None: + try: + result = subprocess.run( + ["hy-smi"], + check=False, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + timeout=SMI_TIMEOUT_SECONDS, + ) + except (OSError, subprocess.TimeoutExpired) as error: + print(f"拒绝启动：hy-smi 空闲检查失败：{error}", file=sys.stderr) + raise SystemExit(90) from error + if result.returncode != 0: + print( + f"拒绝启动：hy-smi 退出码为 {result.returncode}。\n{result.stdout}", + file=sys.stderr, + ) + raise SystemExit(90) + + utilization: dict[int, tuple[float, float]] = {} + for line in result.stdout.splitlines(): + fields = line.split() + if ( + len(fields) >= 7 + and fields[0].isdigit() + and fields[5].endswith("%") + and fields[6].endswith("%") + ): + try: + utilization[int(fields[0])] = ( + float(fields[5][:-1]), + float(fields[6][:-1]), + ) + except ValueError: + continue + if set(utilization) != EXPECTED_DEVICES: + print( + "拒绝启动：hy-smi 未完整报告 0-7 号设备。\n" + result.stdout, + file=sys.stderr, + ) + raise SystemExit(90) + + busy_devices = { + device: values + for device, values in utilization.items() + if values[0] > 0.0 or values[1] > 0.0 + } + local_users = _local_gpu_processes() + if busy_devices or local_users: + print( + "拒绝启动：GPU 未完全空闲；" + f"设备占用={busy_devices}，容器内进程={local_users}。\n{result.stdout}", + file=sys.stderr, + ) + raise SystemExit(90) diff --git a/test/ppl/qwen3_235b/scripts/_ppl_common.py b/test/ppl/qwen3_235b/scripts/_ppl_common.py new file mode 100755 index 00000000..bc720547 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/_ppl_common.py @@ -0,0 +1,338 @@ +#!/usr/bin/env python3 +"""Shared, deterministic corpus and sliding-window helpers for true PPL tests.""" + +from __future__ import annotations + +import ast +import array +import hashlib +import json +import operator +import re +import struct +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Iterator, Sequence + + +CORPUS_SCHEMA = "qw235_ppl_token_ids_v1" +SCORING_METHOD = ( + "sliding_window_shifted_cross_entropy_fp32_compute_fp64_accumulation" +) +SHA256_RE = re.compile(r"^[0-9a-f]{64}$") + + +def canonical_json_bytes(value: object) -> bytes: + return json.dumps( + value, ensure_ascii=True, sort_keys=True, separators=(",", ":") + ).encode("ascii") + + +def _canonical_int_sequence_sha256(values: Iterable[int], label: str) -> str: + """Hash an integer sequence exactly like compact JSON ``[1,2,3]``.""" + digest = hashlib.sha256() + digest.update(b"[") + for index, value in enumerate(values): + if isinstance(value, bool): + raise ValueError(f"{label}[{index}] 不是非负整数") + try: + parsed = operator.index(value) + except TypeError as error: + raise ValueError(f"{label}[{index}] 不是非负整数") from error + if parsed < 0: + raise ValueError(f"{label}[{index}] 不是非负整数：{value!r}") + if index: + digest.update(b",") + digest.update(str(parsed).encode("ascii")) + digest.update(b"]") + return digest.hexdigest() + + +def canonical_token_ids_sha256(token_ids: Iterable[int]) -> str: + return _canonical_int_sequence_sha256(token_ids, "token_ids") + + +def canonical_indices_sha256(indices: Iterable[int]) -> str: + return _canonical_int_sequence_sha256(indices, "indices") + + +@dataclass(frozen=True) +class PplCorpusManifest: + path: Path + payload: dict[str, Any] + token_ids: tuple[int, ...] + manifest_sha256: str + token_ids_sha256: str + + @property + def token_count(self) -> int: + return len(self.token_ids) + + +@dataclass(frozen=True) +class SlidingWindow: + """One causal-LM window using half-open global token index ranges. + + ``token_start:token_end`` is model input. Targets in + ``score_start:score_end`` are scored. ``prediction_*`` select the matching + logits before the causal shift, while ``target_*`` select labels locally. + """ + + index: int + token_start: int + token_end: int + score_start: int + score_end: int + token_ids: tuple[int, ...] + + @property + def scored_token_count(self) -> int: + return self.score_end - self.score_start + + @property + def prediction_start(self) -> int: + return self.score_start - self.token_start - 1 + + @property + def prediction_end(self) -> int: + return self.score_end - self.token_start - 1 + + @property + def target_start(self) -> int: + return self.score_start - self.token_start + + @property + def target_end(self) -> int: + return self.score_end - self.token_start + + +def _required_positive_int(payload: dict[str, Any], key: str, path: Path) -> int: + value = payload.get(key) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path} 的 {key} 必须是正整数") + parsed = value + if parsed <= 0: + raise ValueError(f"{path} 的 {key} 必须是正整数") + return parsed + + +def _required_sha(payload: dict[str, Any], key: str, path: Path) -> str: + value = str(payload.get(key, "")).lower() + if not SHA256_RE.fullmatch(value): + raise ValueError(f"{path} 的 {key} 不是有效 SHA256") + return value + + +def _load_npy(manifest_path: Path, relative_name: object) -> list[int]: + relative = Path(str(relative_name)) + if relative.is_absolute() or ".." in relative.parts: + raise ValueError(f"{manifest_path} 的 token_ids_file 必须是安全相对路径") + base = manifest_path.parent.resolve() + token_path = (base / relative).resolve() + try: + token_path.relative_to(base) + except ValueError as error: + raise ValueError(f"token_ids_file 越出 manifest 目录：{relative}") from error + try: + import numpy as np + except ImportError: + return _load_int64_npy_without_numpy(token_path) + try: + array = np.load(token_path, allow_pickle=False) + except FileNotFoundError: + raise ValueError(f"token_ids_file 不存在：{token_path}") from None + if array.ndim != 1 or array.dtype.kind not in "iu": + raise ValueError(f"{token_path} 必须是一维整数 .npy 数组") + return [int(value) for value in array.tolist()] + + +def _load_int64_npy_without_numpy(path: Path) -> list[int]: + try: + with path.open("rb") as handle: + if handle.read(6) != b"\x93NUMPY": + raise ValueError(f"{path} 不是有效 .npy 文件") + version = handle.read(2) + if version == b"\x01\x00": + header_length = struct.unpack(" None: + """Write a portable NumPy v1.0, one-dimensional little-endian int64 file.""" + output = Path(path) + values = list(token_ids) + # Validate before creating a partial file. + canonical_token_ids_sha256(values) + header_text = repr( + {"descr": " 65535: + raise ValueError(".npy header 超过 v1.0 长度限制") + output.parent.mkdir(parents=True, exist_ok=True) + packed = array.array("q", (int(value) for value in values)) + if packed.itemsize != 8: + raise RuntimeError("当前 Python 平台的 signed long long 不是 64 bit") + if sys.byteorder != "little": + packed.byteswap() + with output.open("wb") as handle: + handle.write(b"\x93NUMPY") + handle.write(b"\x01\x00") + handle.write(struct.pack(" PplCorpusManifest: + """Load and fully verify an inline or relative-``.npy`` corpus manifest.""" + manifest_path = Path(path) + try: + payload = json.loads(manifest_path.read_text(encoding="utf-8")) + except FileNotFoundError: + raise ValueError(f"PPL manifest 不存在：{manifest_path}") from None + except json.JSONDecodeError as error: + raise ValueError(f"PPL manifest JSON 无效：{manifest_path}: {error.msg}") from error + if not isinstance(payload, dict): + raise ValueError(f"PPL manifest 必须是 JSON 对象：{manifest_path}") + if payload.get("schema") != CORPUS_SCHEMA: + raise ValueError( + f"{manifest_path} schema 必须为 {CORPUS_SCHEMA!r}，" + f"实际为 {payload.get('schema')!r}" + ) + manifest_hash = _required_sha(payload, "manifest_sha256", manifest_path) + semantic_payload = dict(payload) + semantic_payload.pop("manifest_sha256", None) + calculated_manifest_hash = hashlib.sha256( + canonical_json_bytes(semantic_payload) + ).hexdigest() + if manifest_hash != calculated_manifest_hash: + raise ValueError(f"{manifest_path} 的 manifest_sha256 校验失败") + for key in ("source_sha256", "tokenizer_sha256"): + _required_sha(payload, key, manifest_path) + + has_inline = "token_ids" in payload + has_file = "token_ids_file" in payload + if has_inline == has_file: + raise ValueError( + f"{manifest_path} 必须且只能包含 token_ids 或 token_ids_file 之一" + ) + if has_inline: + raw_ids = payload["token_ids"] + if not isinstance(raw_ids, list): + raise ValueError(f"{manifest_path} 的 token_ids 必须是数组") + token_ids = list(raw_ids) + else: + token_ids = _load_npy(manifest_path, payload["token_ids_file"]) + + # The canonical hash validates type, integrality, sign, order and contents. + calculated_token_hash = canonical_token_ids_sha256(token_ids) + expected_token_hash = _required_sha( + payload, "token_ids_sha256", manifest_path + ) + if calculated_token_hash != expected_token_hash: + raise ValueError(f"{manifest_path} 的 token_ids_sha256 校验失败") + token_count = _required_positive_int(payload, "token_count", manifest_path) + if token_count != len(token_ids) or token_count < 2: + raise ValueError( + f"{manifest_path} token_count={token_count}，实际 token 数={len(token_ids)}" + ) + return PplCorpusManifest( + path=manifest_path, + payload=payload, + token_ids=tuple(int(value) for value in token_ids), + manifest_sha256=manifest_hash, + token_ids_sha256=expected_token_hash, + ) + + +def iter_sliding_windows( + token_ids: Sequence[int], + window_size: int, + stride: int, + max_scored_tokens: int | None = None, +) -> Iterator[SlidingWindow]: + """Yield windows that score global token indices ``1..N-1`` exactly once. + + The first window scores ``1:end``. Every later window scores only + ``previous_end:end``; overlapped prefix tokens provide context but are not + counted again. ``stride`` must be smaller than ``window_size`` so the first + new target in every later window retains its immediately preceding token. + """ + if isinstance(window_size, bool) or not isinstance(window_size, int): + raise ValueError("window_size 必须是整数") + if isinstance(stride, bool) or not isinstance(stride, int): + raise ValueError("stride 必须是整数") + if window_size < 2: + raise ValueError("window_size 必须至少为 2") + if stride < 1 or stride >= window_size: + raise ValueError("stride 必须满足 1 <= stride < window_size") + if len(token_ids) < 2: + raise ValueError("至少需要 2 个 token 才能计算 PPL") + if max_scored_tokens is not None: + if ( + isinstance(max_scored_tokens, bool) + or not isinstance(max_scored_tokens, int) + or max_scored_tokens < 1 + ): + raise ValueError("max_scored_tokens 必须是正整数") + + score_limit = len(token_ids) + if max_scored_tokens is not None: + score_limit = min(score_limit, 1 + max_scored_tokens) + previous_end = 1 + index = 0 + while previous_end < score_limit: + if index == 0: + token_end = min(score_limit, window_size) + token_start = 0 + score_start = 1 + else: + token_end = min(score_limit, previous_end + stride) + token_start = max(0, token_end - window_size) + score_start = previous_end + window = SlidingWindow( + index=index, + token_start=token_start, + token_end=token_end, + score_start=score_start, + score_end=token_end, + token_ids=tuple(int(value) for value in token_ids[token_start:token_end]), + ) + if window.prediction_start < 0 or window.prediction_end > len(window.token_ids): + raise AssertionError("内部错误：滑窗缺少 causal predecessor") + yield window + previous_end = token_end + index += 1 diff --git a/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py b/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py new file mode 100644 index 00000000..5bd49e86 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Compare InfiniLM BF16 and W8A8 true-PPL result JSON files.""" + +from __future__ import annotations + +import argparse +import json +import os +from dataclasses import asdict +from pathlib import Path +from typing import Any, Sequence + +from calculate_true_ppl import _validate_same_workload, load_result + + +SCHEMA = "qwen3_235b_infinilm_precision_ppl_comparison/v1" + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--inputs", nargs=2, type=Path, required=True) + parser.add_argument("--max-ppl-increase-percent", type=float, default=20.0) + parser.add_argument("--json-out", type=Path, required=True) + args = parser.parse_args(argv) + if args.max_ppl_increase_percent < 0: + parser.error("--max-ppl-increase-percent must be non-negative") + return args + + +def _atomic_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") + temporary.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + results = [load_result(path) for path in args.inputs] + if any(result.backend != "infinilm" for result in results): + raise ValueError("both inputs must be InfiniLM results") + precisions: list[str] = [] + for path in args.inputs: + payload = json.loads(path.read_text(encoding="utf-8")) + precision = str(payload.get("precision", "")).strip().upper() + if precision not in {"BF16", "W8A8"}: + raise ValueError(f"{path} has invalid precision: {precision!r}") + precisions.append(precision) + by_precision = dict(zip(precisions, results, strict=True)) + if set(by_precision) != {"BF16", "W8A8"}: + raise ValueError("inputs must contain one BF16 and one W8A8 result") + baseline = by_precision["BF16"] + candidate = by_precision["W8A8"] + _validate_same_workload(baseline, candidate) + increase = (candidate.ppl / baseline.ppl - 1.0) * 100.0 + threshold = float(args.max_ppl_increase_percent) + passed = increase <= threshold + payload = { + "schema": SCHEMA, + "status": "PASS" if passed else "FAIL", + "baseline": asdict(baseline), + "candidate": asdict(candidate), + "ppl_increase_percent": increase, + "max_ppl_increase_percent": threshold, + } + _atomic_json(args.json_out, payload) + except (OSError, ValueError, KeyError, TypeError, json.JSONDecodeError) as error: + print(f"PPL comparison error: {error}") + return 2 + + print(f"InfiniLM BF16 PPL: {baseline.ppl:.6f}") + print(f"InfiniLM W8A8 PPL: {candidate.ppl:.6f}") + print(f"W8A8 PPL increase: {increase:.2f}%") + print(f"Quality threshold: <= {threshold:.2f}%") + print(f"Result: {'PASS' if passed else 'FAIL'}") + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py b/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py new file mode 100755 index 00000000..cfefc232 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""严格比较 Transformers 与 InfiniLM 的真实 token-level PPL 结果。""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import sys +import tempfile +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, Sequence + +from _ppl_common import SCORING_METHOD, canonical_indices_sha256 + +RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" +COMPARISON_SCHEMA = "qwen3_235b_true_ppl_comparison/v1" +SHA256_RE = re.compile(r"^[0-9a-f]{64}$") + + +@dataclass(frozen=True) +class PplResult: + path: str + backend: str + model: str + corpus_manifest_sha256: str + corpus_token_ids_sha256: str + window_size: int + stride: int + scoring_method: str + first_scored_token_index: int + last_scored_token_index_exclusive: int + scored_token_count: int + scored_token_indices_sha256: str + total_nll: float + mean_nll: float + ppl: float + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("paths", nargs="*", type=Path) + parser.add_argument( + "--inputs", + nargs="+", + type=Path, + default=[], + help="Transformers 与 InfiniLM 结果 JSON，顺序可以互换", + ) + parser.add_argument( + "--max-ppl-increase-percent", + type=float, + default=20.0, + help="InfiniLM 相对 Transformers 的最大 PPL 增幅，默认 20%%", + ) + parser.add_argument("--json-out", type=Path) + parser.add_argument("--verbose", action="store_true", help="打印完整 JSON") + args = parser.parse_args(argv) + args.inputs = [*args.paths, *args.inputs] + if len(args.inputs) != 2: + parser.error("必须提供两个结果 JSON：Transformers 与 InfiniLM") + if ( + not math.isfinite(args.max_ppl_increase_percent) + or args.max_ppl_increase_percent < 0 + ): + parser.error("--max-ppl-increase-percent 必须是有限非负数") + return args + + +def _required(payload: dict[str, Any], key: str, path: Path) -> Any: + if key not in payload: + raise ValueError(f"{path} 缺少字段 {key}") + return payload[key] + + +def _sha256(payload: dict[str, Any], key: str, path: Path) -> str: + value = str(_required(payload, key, path)).lower() + if not SHA256_RE.fullmatch(value): + raise ValueError(f"{path} 的 {key} 不是有效 SHA256") + return value + + +def _positive_int(payload: dict[str, Any], key: str, path: Path) -> int: + value = _required(payload, key, path) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path} 的 {key} 必须是正整数") + parsed = value + if parsed <= 0: + raise ValueError(f"{path} 的 {key} 必须是正整数") + return parsed + + +def _nonnegative_int(payload: dict[str, Any], key: str, path: Path) -> int: + value = _required(payload, key, path) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path} 的 {key} 必须是非负整数") + parsed = value + if parsed < 0: + raise ValueError(f"{path} 的 {key} 必须是非负整数") + return parsed + + +def _finite(payload: dict[str, Any], key: str, path: Path) -> float: + try: + value = float(_required(payload, key, path)) + except (TypeError, ValueError) as error: + raise ValueError(f"{path} 的 {key} 必须是有限数") from error + if not math.isfinite(value): + raise ValueError(f"{path} 的 {key} 必须是有限数") + return value + + +def load_result(path: Path) -> PplResult: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + raise ValueError(f"结果文件不存在：{path}") from None + except json.JSONDecodeError as error: + raise ValueError(f"结果 JSON 无效：{path}: {error.msg}") from error + if not isinstance(payload, dict): + raise ValueError(f"结果 JSON 必须是对象：{path}") + if payload.get("status") != "PASS": + raise ValueError( + f"{path} 不是成功的 PPL 结果：status={payload.get('status')!r}" + ) + if payload.get("schema") != RESULT_SCHEMA: + raise ValueError( + f"{path} schema 必须为 {RESULT_SCHEMA!r}，实际为 {payload.get('schema')!r}" + ) + + backend = str(_required(payload, "backend", path)).strip().lower() + if backend not in {"transformers", "infinilm"}: + raise ValueError(f"{path} backend 必须是 transformers 或 infinilm") + model = str(_required(payload, "model", path)).strip() + if not model: + raise ValueError(f"{path} model 不能为空") + window_size = _positive_int(payload, "window_size", path) + stride = _positive_int(payload, "stride", path) + if stride >= window_size: + raise ValueError(f"{path} stride 必须小于 window_size") + scoring_method = str(_required(payload, "scoring_method", path)).strip() + if scoring_method != SCORING_METHOD: + raise ValueError( + f"{path} scoring_method 必须为 {SCORING_METHOD!r}" + ) + first_index = _nonnegative_int(payload, "first_scored_token_index", path) + last_index = _positive_int( + payload, "last_scored_token_index_exclusive", path + ) + scored_count = _positive_int(payload, "scored_token_count", path) + if last_index <= first_index or last_index - first_index != scored_count: + raise ValueError( + f"{path} 的计分范围 [{first_index}, {last_index}) 与 " + f"scored_token_count={scored_count} 不一致" + ) + if first_index != 1: + raise ValueError(f"{path} first_scored_token_index 必须为 1") + scored_indices_hash = _sha256( + payload, "scored_token_indices_sha256", path + ) + expected_indices_hash = canonical_indices_sha256( + range(first_index, last_index) + ) + if scored_indices_hash != expected_indices_hash: + raise ValueError(f"{path} 的 scored_token_indices_sha256 校验失败") + + total_nll = _finite(payload, "total_nll", path) + reported_mean = _finite(payload, "mean_nll", path) + reported_ppl = _finite(payload, "ppl", path) + if total_nll < 0 or reported_mean < 0 or reported_ppl < 1: + raise ValueError(f"{path} 的 NLL/PPL 超出有效范围") + calculated_mean = total_nll / scored_count + if calculated_mean > math.log(sys.float_info.max): + raise ValueError(f"{path} 的 mean NLL 过大，PPL 溢出") + calculated_ppl = math.exp(calculated_mean) + if not math.isclose(reported_mean, calculated_mean, rel_tol=1e-6, abs_tol=1e-8): + raise ValueError( + f"{path} 的 mean_nll 与 total_nll/scored_token_count 不一致" + ) + if not math.isclose(reported_ppl, calculated_ppl, rel_tol=1e-6, abs_tol=1e-8): + raise ValueError(f"{path} 的 ppl 与 exp(mean_nll) 不一致") + + return PplResult( + path=str(path), + backend=backend, + model=model, + corpus_manifest_sha256=_sha256( + payload, "corpus_manifest_sha256", path + ), + corpus_token_ids_sha256=_sha256( + payload, "corpus_token_ids_sha256", path + ), + window_size=window_size, + stride=stride, + scoring_method=scoring_method, + first_scored_token_index=first_index, + last_scored_token_index_exclusive=last_index, + scored_token_count=scored_count, + scored_token_indices_sha256=scored_indices_hash, + total_nll=total_nll, + mean_nll=calculated_mean, + ppl=calculated_ppl, + ) + + +def _ordered(results: Sequence[PplResult]) -> tuple[PplResult, PplResult]: + by_backend = {result.backend: result for result in results} + if len(by_backend) != 2 or set(by_backend) != {"transformers", "infinilm"}: + raise ValueError("必须且只能包含一份 Transformers 和一份 InfiniLM 结果") + return by_backend["transformers"], by_backend["infinilm"] + + +def _validate_same_workload(baseline: PplResult, candidate: PplResult) -> None: + fields = ( + "corpus_manifest_sha256", + "corpus_token_ids_sha256", + "window_size", + "stride", + "scoring_method", + "first_scored_token_index", + "last_scored_token_index_exclusive", + "scored_token_count", + "scored_token_indices_sha256", + ) + mismatches = [ + f"{field}: {getattr(baseline, field)!r} != {getattr(candidate, field)!r}" + for field in fields + if getattr(baseline, field) != getattr(candidate, field) + ] + if mismatches: + raise ValueError("两侧 PPL 工作负载不一致：" + "; ".join(mismatches)) + + +def compare( + baseline: PplResult, candidate: PplResult, threshold_percent: float +) -> dict[str, object]: + _validate_same_workload(baseline, candidate) + increase_percent = (candidate.ppl / baseline.ppl - 1.0) * 100.0 + passed = increase_percent <= threshold_percent + return { + "schema": COMPARISON_SCHEMA, + "status": "PASS" if passed else "FAIL", + "baseline": asdict(baseline), + "candidate": asdict(candidate), + "ppl_increase_percent": increase_percent, + "max_ppl_increase_percent": threshold_percent, + "pass": passed, + } + + +def _atomic_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, ensure_ascii=False, sort_keys=True, indent=2) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + except BaseException: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + baseline, candidate = _ordered([load_result(path) for path in args.inputs]) + report = compare( + baseline, candidate, float(args.max_ppl_increase_percent) + ) + if args.json_out is not None: + _atomic_json(args.json_out, report) + except (OSError, RuntimeError, ValueError) as error: + print(f"错误：{error}", file=sys.stderr) + return 2 + + print("真实 PPL 对比") + print( + f"Transformers：PPL={baseline.ppl:.6f} " + f"NLL={baseline.total_nll:.6f} Token={baseline.scored_token_count}" + ) + print( + f"InfiniLM： PPL={candidate.ppl:.6f} " + f"NLL={candidate.total_nll:.6f} Token={candidate.scored_token_count}" + ) + print(f"PPL 增幅：{report['ppl_increase_percent']:.2f}%") + print( + f"验收要求：增幅 <= {args.max_ppl_increase_percent:.2f}% " + f"结果={report['status']}" + ) + if args.verbose: + print(json.dumps(report, ensure_ascii=False, sort_keys=True, indent=2)) + return 0 if report["pass"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py new file mode 100755 index 00000000..9a2f3562 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +"""Calculate true token-level PPL with the current InfiniLM C++ TP engine.""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import os +import sys +import time +from pathlib import Path +from typing import Any, Sequence + + +SCRIPT_DIR = Path(__file__).resolve().parent +SCRIPTS_DIR = SCRIPT_DIR.parent +if str(SCRIPTS_DIR) not in sys.path: + sys.path.insert(0, str(SCRIPTS_DIR)) + +from _gpu_guard import require_idle_gpu +from _ppl_common import ( + SCORING_METHOD, + canonical_indices_sha256, + iter_sliding_windows, + load_manifest, +) + + +RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" +DEFAULT_MODEL = "/data1/Qwen3_235B" +EXPECTED_MODEL_TYPE = "qwen3_moe" +EXPECTED_VOCAB_SIZE = 151936 +PAGED_KV_BLOCK_SIZE = 256 + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="True shifted-token PPL for Qwen3_235B with InfiniLM TP8" + ) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--token-manifest", required=True) + parser.add_argument("--window", type=int, default=256) + parser.add_argument("--stride", type=int, default=128) + parser.add_argument( + "--max-scored-tokens", + type=int, + default=10240, + help="maximum target tokens to score; 0 scores the full manifest", + ) + parser.add_argument("--tp-size", type=int, default=8) + parser.add_argument("--attention", default="flash-attn") + parser.add_argument("--json-output") + args = parser.parse_args(argv) + + if not Path(args.model).is_dir(): + parser.error(f"model directory does not exist: {args.model}") + if not Path(args.token_manifest).is_file(): + parser.error(f"token manifest does not exist: {args.token_manifest}") + if args.window < 2: + parser.error("--window must be at least 2") + if args.stride < 1 or args.stride >= args.window: + parser.error("--stride must satisfy 1 <= stride < window") + if args.max_scored_tokens < 0: + parser.error("--max-scored-tokens must be non-negative") + if args.tp_size < 1: + parser.error("--tp-size must be positive") + + args.model = str(Path(args.model).resolve()) + args.token_manifest = str(Path(args.token_manifest).resolve()) + if args.json_output: + args.json_output = str(Path(args.json_output).resolve()) + return args + + +def _atomic_json(path_value: str, payload: dict[str, Any]) -> None: + path = Path(path_value) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") + temporary.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def _read_model_config(model_path: str) -> dict[str, Any]: + config_path = Path(model_path) / "config.json" + try: + config = json.loads(config_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError(f"cannot read model config {config_path}: {error}") from error + if not isinstance(config, dict): + raise RuntimeError(f"model config must be an object: {config_path}") + if config.get("model_type") != EXPECTED_MODEL_TYPE: + raise RuntimeError( + f"expected model_type={EXPECTED_MODEL_TYPE!r}, " + f"got {config.get('model_type')!r}" + ) + if int(config.get("vocab_size", 0)) != EXPECTED_VOCAB_SIZE: + raise RuntimeError( + f"expected vocab_size={EXPECTED_VOCAB_SIZE}, " + f"got {config.get('vocab_size')!r}" + ) + return config + + +def _is_quantized(config: dict[str, Any]) -> bool: + quantization = config.get("quantization_config") + return isinstance(quantization, dict) and bool(quantization) + + +def _run(args: argparse.Namespace) -> dict[str, Any]: + import infinicore + from infinilm.cache import PagedKVCacheConfig + from infinilm.distributed import DistConfig + from infinilm.infer_engine import InferEngine + from infinilm.modeling_utils import load_model_state_dict_by_file + + corpus = load_manifest(args.token_manifest) + model_config = _read_model_config(args.model) + if any(token >= EXPECTED_VOCAB_SIZE for token in corpus.token_ids): + raise RuntimeError("token manifest contains an ID outside the model vocabulary") + + available_targets = corpus.token_count - 1 + scored_token_count = ( + available_targets + if args.max_scored_tokens == 0 + else min(args.max_scored_tokens, available_targets) + ) + max_targets = None if args.max_scored_tokens == 0 else scored_token_count + windows = list( + iter_sliding_windows( + corpus.token_ids, + args.window, + args.stride, + max_targets, + ) + ) + if sum(window.scored_token_count for window in windows) != scored_token_count: + raise RuntimeError("sliding-window plan does not match scored token count") + + first_scored_token_index = 1 + last_scored_token_index_exclusive = 1 + scored_token_count + indices_sha256 = canonical_indices_sha256( + range(first_scored_token_index, last_scored_token_index_exclusive) + ) + precision = "W8A8" if _is_quantized(model_config) else "BF16" + config_payload = { + "backend": "infinilm", + "model": args.model, + "precision": precision, + "tp_size": args.tp_size, + "attention": args.attention, + "graph_enabled": False, + "window_size": args.window, + "stride": args.stride, + "scored_token_count": scored_token_count, + "scoring_method": SCORING_METHOD, + "corpus_manifest_sha256": corpus.manifest_sha256, + "corpus_token_ids_sha256": corpus.token_ids_sha256, + } + print( + "INFINILM_QWEN3_235B_PPL_CONFIG " + + json.dumps(config_payload, ensure_ascii=False, sort_keys=True), + flush=True, + ) + + device = infinicore.device("cuda", 0) + load_start = time.perf_counter() + model = InferEngine( + args.model, + device=device, + distributed_config=DistConfig(args.tp_size), + # This InfiniLM branch's flash-attn backend consumes the paged KV-cache + # layout while retaining flash-attn as the attention implementation. + cache_config=PagedKVCacheConfig( + num_blocks=(args.window + PAGED_KV_BLOCK_SIZE - 1) + // PAGED_KV_BLOCK_SIZE, + block_size=PAGED_KV_BLOCK_SIZE, + ), + enable_graph_compiling=False, + attention_backend=args.attention, + ) + if not hasattr(model, "score_nll"): + raise RuntimeError( + "installed InfiniLM lacks InferEngine.score_nll; rebuild the PPL scoring patch" + ) + load_model_state_dict_by_file(model, args.model, dtype=model.dtype) + model_load_seconds = time.perf_counter() - load_start + + window_nll_values: list[float] = [] + window_results: list[dict[str, Any]] = [] + infinicore.sync_device() + scoring_start = time.perf_counter() + for window in windows: + input_tokens = list(window.token_ids[:-1]) + label_tokens = list(window.token_ids[1:]) + if not input_tokens or len(input_tokens) != len(label_tokens): + raise RuntimeError(f"invalid causal shift in window {window.index}") + input_ids = infinicore.from_list( + [input_tokens], dtype=infinicore.int64 + ) + labels = infinicore.from_list( + [label_tokens], dtype=infinicore.int64 + ) + nll, returned_tokens = model.score_nll( + input_ids, + labels, + score_start=window.prediction_start, + ) + if returned_tokens != window.scored_token_count: + raise RuntimeError( + f"window {window.index} scored {returned_tokens} tokens, " + f"expected {window.scored_token_count}" + ) + if not math.isfinite(nll) or nll < 0: + raise RuntimeError(f"window {window.index} returned invalid NLL {nll}") + window_nll_values.append(nll) + window_results.append( + { + "index": window.index, + "context_start": window.token_start, + "target_start": window.score_start, + "target_end": window.score_end, + "input_token_count": len(window.token_ids), + "scored_token_count": returned_tokens, + "nll": nll, + } + ) + print( + f"PPL window {window.index + 1}/{len(windows)} " + f"tokens={returned_tokens} nll={nll:.6f}", + flush=True, + ) + + infinicore.sync_device() + scoring_seconds = time.perf_counter() - scoring_start + total_nll = math.fsum(window_nll_values) + mean_nll = total_nll / scored_token_count + try: + ppl = math.exp(mean_nll) + except OverflowError as error: + raise RuntimeError(f"PPL overflow at mean NLL={mean_nll}") from error + if not math.isfinite(ppl): + raise RuntimeError(f"PPL is not finite: {ppl}") + + result = { + "schema": RESULT_SCHEMA, + "status": "PASS", + "backend": "infinilm", + "model": args.model, + "precision": precision, + "tp_size": args.tp_size, + "attention": args.attention, + "graph_enabled": False, + "corpus_manifest": args.token_manifest, + "corpus_manifest_sha256": corpus.manifest_sha256, + "corpus_token_ids_sha256": corpus.token_ids_sha256, + "corpus_token_count": corpus.token_count, + "window_size": args.window, + "stride": args.stride, + "scoring_method": SCORING_METHOD, + "first_scored_token_index": first_scored_token_index, + "last_scored_token_index_exclusive": last_scored_token_index_exclusive, + "scored_token_indices_sha256": indices_sha256, + "scored_token_count": scored_token_count, + "total_nll": total_nll, + "mean_nll": mean_nll, + "ppl": ppl, + "windows": window_results, + "window_count": len(window_results), + "scoring_seconds": scoring_seconds, + "scored_tokens_per_second": scored_token_count / scoring_seconds, + "model_load_seconds": model_load_seconds, + "vocab_size": EXPECTED_VOCAB_SIZE, + } + if args.json_output: + _atomic_json(args.json_output, result) + print( + "INFINILM_QWEN3_235B_PPL_RESULT " + + json.dumps(result, ensure_ascii=False, sort_keys=True), + flush=True, + ) + print( + f"InfiniLM {precision} true PPL: {ppl:.6f} " + f"(mean NLL={mean_nll:.6f}, tokens={scored_token_count})", + flush=True, + ) + + del model + gc.collect() + infinicore.sync_device() + return result + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + # Validate the workload before checking or reserving GPUs. + load_manifest(args.token_manifest) + require_idle_gpu() + _run(args) + except BaseException as error: + completion = { + "schema": RESULT_SCHEMA, + "status": "ERROR", + "exit_code": 1, + "error": { + "type": type(error).__name__, + "message": str(error), + }, + } + print( + "INFINILM_QWEN3_235B_PPL_COMPLETE " + + json.dumps(completion, ensure_ascii=False, sort_keys=True), + flush=True, + ) + raise + print( + "INFINILM_QWEN3_235B_PPL_COMPLETE " + + json.dumps( + {"schema": RESULT_SCHEMA, "status": "PASS", "exit_code": 0}, + ensure_ascii=False, + sort_keys=True, + ), + flush=True, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py new file mode 100755 index 00000000..b501e453 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py @@ -0,0 +1,328 @@ +#!/usr/bin/env python3 +"""将本地纯文本固化为 Qwen3_235B PPL 测试使用的 token manifest。""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import operator +import os +import sys +import tempfile +from pathlib import Path +from typing import Any, Sequence + +from _ppl_common import ( + CORPUS_SCHEMA, + canonical_json_bytes, + canonical_token_ids_sha256, + write_token_ids_npy, +) + + +def _jsonable(value: Any) -> Any: + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + return str(value) + + +def _tokenizer_fingerprint(tokenizer: Any) -> tuple[str, str]: + backend = getattr(tokenizer, "backend_tokenizer", None) + if backend is not None and hasattr(backend, "to_str"): + try: + backend_payload: object = json.loads(backend.to_str()) + except (TypeError, ValueError, json.JSONDecodeError): + backend_payload = backend.to_str() + method = "backend_tokenizer+special_tokens/v1" + semantics = { + "backend_tokenizer": backend_payload, + "special_tokens_map": _jsonable( + getattr(tokenizer, "special_tokens_map", {}) + ), + } + else: + if not hasattr(tokenizer, "get_vocab"): + raise RuntimeError("tokenizer 既没有 backend_tokenizer，也没有 get_vocab()") + method = "vocab+init_kwargs+special_tokens/v1" + semantics = { + "vocab": tokenizer.get_vocab(), + "init_kwargs": _jsonable(getattr(tokenizer, "init_kwargs", {})), + "special_tokens_map": _jsonable( + getattr(tokenizer, "special_tokens_map", {}) + ), + } + return hashlib.sha256(canonical_json_bytes(semantics)).hexdigest(), method + + +def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--input", + nargs="+", + type=Path, + required=True, + help="本地 UTF-8 WikiText/raw text 文件，可按顺序提供多个文件", + ) + parser.add_argument( + "--tokenizer", + required=True, + help="本地 Qwen3_235B 模型或 tokenizer 目录", + ) + parser.add_argument("--output", required=True, type=Path, help="输出 JSON manifest") + parser.add_argument( + "--max-tokens", + type=int, + help="仅保留开头 N 个 token；省略时保留全部 token", + ) + parser.add_argument( + "--document-separator", + default="\n\n", + help=r"多个输入文件之间的分隔符，默认 '\n\n'", + ) + parser.add_argument( + "--storage", + choices=("inline", "npy"), + default="inline", + help="token IDs 内联到 JSON（默认），或写入相对路径 .npy 文件", + ) + parser.add_argument( + "--token-ids-file", + type=Path, + help="--storage=npy 时的相对路径；默认 .tokens.npy", + ) + parser.add_argument( + "--allow-download", + action="store_true", + help="允许 Transformers 访问网络；默认只读取本地文件/缓存", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="传给 AutoTokenizer；Qwen3 官方 tokenizer 通常不需要", + ) + parser.add_argument("--overwrite", action="store_true", help="覆盖已有输出") + args = parser.parse_args(argv) + + if args.max_tokens is not None and args.max_tokens < 2: + parser.error("--max-tokens 必须至少为 2") + if args.token_ids_file is not None and args.storage != "npy": + parser.error("--token-ids-file 只能与 --storage=npy 一起使用") + if len({path.name for path in args.input}) != len(args.input): + parser.error("输入文件名不能重复，否则 manifest 无法稳定区分来源") + return args + + +def _read_sources( + paths: Sequence[Path], separator: str +) -> tuple[str, list[dict[str, object]]]: + documents: list[str] = [] + files: list[dict[str, object]] = [] + for path in paths: + if not path.is_file(): + raise FileNotFoundError(f"输入文件不存在：{path}") + raw = path.read_bytes() + try: + text = raw.decode("utf-8") + except UnicodeDecodeError as error: + raise ValueError(f"输入文件不是有效 UTF-8：{path}: {error}") from error + documents.append(text) + files.append( + { + "name": path.name, + "byte_count": len(raw), + "sha256": hashlib.sha256(raw).hexdigest(), + } + ) + return separator.join(documents), files + + +def _load_tokenizer(identifier: str, allow_download: bool, trust_remote_code: bool) -> Any: + try: + from transformers import AutoTokenizer + except ImportError as error: + raise RuntimeError("缺少 transformers，无法加载 tokenizer") from error + return AutoTokenizer.from_pretrained( + identifier, + local_files_only=not allow_download, + trust_remote_code=trust_remote_code, + ) + + +def _encode(tokenizer: Any, text: str) -> list[int]: + encoded = tokenizer( + text, + add_special_tokens=False, + truncation=False, + return_attention_mask=False, + return_token_type_ids=False, + ) + raw_ids = encoded["input_ids"] + if not isinstance(raw_ids, (list, tuple)) or ( + raw_ids and isinstance(raw_ids[0], (list, tuple)) + ): + raise RuntimeError("tokenizer 必须为单条文本返回一维 input_ids") + token_ids: list[int] = [] + for index, value in enumerate(raw_ids): + if isinstance(value, bool): + raise RuntimeError(f"input_ids[{index}] 不是有效整数") + try: + token = operator.index(value) + except TypeError as error: + raise RuntimeError(f"input_ids[{index}] 不是有效整数") from error + if token < 0: + raise RuntimeError(f"input_ids[{index}] 不是非负整数：{value!r}") + token_ids.append(token) + if len(token_ids) < 2: + raise RuntimeError("语料 token 数不足 2，无法计算 causal LM PPL") + return token_ids + + +def _atomic_json(path: Path, payload: object) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(payload, handle, ensure_ascii=False, sort_keys=True, indent=2) + handle.write("\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + except BaseException: + try: + os.unlink(temporary) + except FileNotFoundError: + pass + raise + + +def _npy_relative_path(output: Path, requested: Path | None) -> Path: + relative = requested or Path(f"{output.stem}.tokens.npy") + if relative.is_absolute() or ".." in relative.parts or relative.name in {"", "."}: + raise ValueError("--token-ids-file 必须是 manifest 目录内的安全相对路径") + if relative.suffix != ".npy": + raise ValueError("--token-ids-file 必须以 .npy 结尾") + return relative + + +def _write_npy(path: Path, token_ids: Sequence[int], overwrite: bool) -> None: + if path.exists() and not overwrite: + raise FileExistsError(f"token 文件已存在（可加 --overwrite）：{path}") + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp.npy") + try: + write_token_ids_npy(temporary, token_ids) + os.replace(temporary, path) + finally: + try: + temporary.unlink() + except FileNotFoundError: + pass + + +def build_manifest( + *, + source_text: str, + source_files: list[dict[str, object]], + separator: str, + tokenizer: Any, + tokenizer_label: str, + token_ids: list[int], + original_token_count: int, + max_tokens: int | None, + storage: str, + token_ids_file: str | None, +) -> dict[str, object]: + tokenizer_hash, fingerprint_method = _tokenizer_fingerprint(tokenizer) + token_hash = canonical_token_ids_sha256(token_ids) + payload: dict[str, object] = { + "schema": CORPUS_SCHEMA, + "source_sha256": hashlib.sha256(source_text.encode("utf-8")).hexdigest(), + "tokenizer_sha256": tokenizer_hash, + "token_count": len(token_ids), + "token_ids_sha256": token_hash, + "source": { + "encoding": "utf-8", + "document_separator": separator, + "file_count": len(source_files), + "files": source_files, + }, + "tokenizer": { + "name": Path(tokenizer_label.rstrip("/")).name or tokenizer_label, + "class": tokenizer.__class__.__name__, + "vocab_size": int(getattr(tokenizer, "vocab_size", 0)), + "fingerprint_method": fingerprint_method, + }, + "tokenization": { + "add_special_tokens": False, + "original_token_count": original_token_count, + "max_tokens": max_tokens, + "truncated": len(token_ids) != original_token_count, + }, + } + if storage == "inline": + payload["token_ids"] = token_ids + else: + if token_ids_file is None: + raise ValueError("npy storage 缺少 token_ids_file") + payload["token_ids_file"] = token_ids_file + payload["token_ids_dtype"] = "int64" + payload["manifest_sha256"] = hashlib.sha256( + canonical_json_bytes(payload) + ).hexdigest() + return payload + + +def main(argv: Sequence[str] | None = None) -> int: + args = _parse_args(argv) + try: + if args.output.exists() and not args.overwrite: + raise FileExistsError(f"输出已存在（可加 --overwrite）：{args.output}") + source_text, source_files = _read_sources(args.input, args.document_separator) + tokenizer = _load_tokenizer( + args.tokenizer, args.allow_download, args.trust_remote_code + ) + all_token_ids = _encode(tokenizer, source_text) + original_token_count = len(all_token_ids) + token_ids = ( + all_token_ids[: args.max_tokens] + if args.max_tokens is not None + else all_token_ids + ) + + relative_npy: Path | None = None + if args.storage == "npy": + relative_npy = _npy_relative_path(args.output, args.token_ids_file) + + manifest = build_manifest( + source_text=source_text, + source_files=source_files, + separator=args.document_separator, + tokenizer=tokenizer, + tokenizer_label=args.tokenizer, + token_ids=token_ids, + original_token_count=original_token_count, + max_tokens=args.max_tokens, + storage=args.storage, + token_ids_file=str(relative_npy) if relative_npy is not None else None, + ) + if relative_npy is not None: + _write_npy(args.output.parent / relative_npy, token_ids, args.overwrite) + _atomic_json(args.output, manifest) + except (FileNotFoundError, FileExistsError, RuntimeError, ValueError) as error: + print(f"错误：{error}", file=sys.stderr) + return 2 + + print(f"PPL 语料已固化：{args.output}") + print(f"Token 数：{manifest['token_count']}") + print(f"Token SHA256：{manifest['token_ids_sha256']}") + print(f"Manifest SHA256：{manifest['manifest_sha256']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py b/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py new file mode 100755 index 00000000..027b5c87 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py @@ -0,0 +1,1144 @@ +#!/usr/bin/env python3 +"""Minimal Transformers TP benchmark runner for Qwen3_235B-A22B. + +The scenario wrappers are intentionally directly executable. When started as +``python wrapper.py`` this module replaces the process with a torchrun agent; +the original timeout therefore remains responsible for the agent, which in +turn terminates every worker on SIGTERM. +""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import importlib.metadata +import json +import os +import statistics +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + + +SCRIPT_ROOT = str(Path(__file__).resolve().parents[1]) +if SCRIPT_ROOT not in sys.path: + sys.path.insert(0, SCRIPT_ROOT) + +from _gpu_guard import require_idle_gpu as _require_idle_gpu + + +MODEL_NAME = "Qwen3_235B" +DEFAULT_MODEL = "/data1/Qwen3_235B" +DEFAULT_PROMPT_FILE = "examples/bench_prompt.md" +FALLBACK_PROMPT = """High-performance language-model inference processes a prompt in a prefill phase +and then produces one token per request during decode. Tensor-parallel ranks must +exchange identical partial results, while the key/value cache keeps each request +lane isolated. The benchmark uses deterministic prompt tokens and greedy decoding +so that every reported run has an exact, auditable token count.""" +MEASURED_INPUT_LENGTHS = 1 +REPEATS_PER_INPUT_LENGTH = 3 +MEASURED_ITERATIONS = REPEATS_PER_INPUT_LENGTH +MEASUREMENT_SEMANTICS = "one_fixed_shape_x_three_measurements" +SMOKE_OUTPUT_TOKENS = 64 +HYGON_TP_PLAN = { + "lm_head": "colwise_gather_output", + "model.layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "model.layers.*.mlp.experts.down_proj": "rowwise", + "model.layers.*.mlp.experts": "moe_tp_experts", +} +EXPECTED_QWEN3_235B_ARCHITECTURE = { + "hidden_size": 4096, + "intermediate_size": 12288, + "head_dim": 128, + "num_attention_heads": 64, + "num_key_value_heads": 4, + "num_hidden_layers": 94, + "num_experts": 128, + "num_experts_per_tok": 8, + "moe_intermediate_size": 1536, + "vocab_size": 151936, +} + + +@dataclass(frozen=True) +class Scenario: + name: str + batch_size: int + input_tokens: int + output_tokens: int + + @property + def input_lengths(self) -> tuple[int]: + return (self.input_tokens,) + + @property + def total_context_tokens(self) -> int: + return self.input_tokens + self.output_tokens + + +def _parse_args(scenario: Scenario) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Transformers Qwen3_235B TP8 benchmark: " + f"batch={scenario.batch_size}, input={scenario.input_tokens}, " + f"output={scenario.output_tokens}, " + f"total={scenario.total_context_tokens} tokens" + ) + ) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--prompt-file", default=DEFAULT_PROMPT_FILE) + parser.add_argument("--output-tokens", type=int, default=scenario.output_tokens) + parser.add_argument("--tp-size", type=int, default=8) + parser.add_argument( + "--smoke", + action="store_true", + help=( + "load the full model but run only batch=1, input=16 and output=64 " + "to validate the TP/attention/cache path" + ), + ) + parser.add_argument( + "--attention", + choices=("eager",), + default="eager", + help="BW1100 correctness path; SDPA is unsupported for this model stack", + ) + args = parser.parse_args() + + if not Path(args.model).is_dir(): + parser.error(f"model directory does not exist: {args.model}") + if args.smoke: + args.output_tokens = SMOKE_OUTPUT_TOKENS + if args.output_tokens < 2: + parser.error("--output-tokens must be at least 2 to measure decode speed") + if args.tp_size < 1: + parser.error("--tp-size must be positive") + if len(set(scenario.input_lengths)) != MEASURED_INPUT_LENGTHS: + parser.error( + f"scenario must define exactly {MEASURED_INPUT_LENGTHS} lengths" + ) + return args + + +def _effective_scenario(scenario: Scenario, smoke: bool) -> Scenario: + if not smoke: + return scenario + return Scenario(f"{scenario.name}_smoke", 1, 16, SMOKE_OUTPUT_TOKENS) + + +def _validate_qwen3_235b_architecture(model_config: Any) -> dict[str, int]: + if getattr(model_config, "model_type", None) != "qwen3_moe": + raise RuntimeError( + "this benchmark requires model_type='qwen3_moe', got " + f"{getattr(model_config, 'model_type', None)!r}" + ) + actual: dict[str, int] = {} + mismatches: list[str] = [] + for field, expected in EXPECTED_QWEN3_235B_ARCHITECTURE.items(): + value = getattr(model_config, field, None) + try: + parsed = int(value) + except (TypeError, ValueError): + mismatches.append(f"{field}={value!r} (expected {expected})") + continue + actual[field] = parsed + if parsed != expected: + mismatches.append(f"{field}={parsed} (expected {expected})") + architectures = tuple(getattr(model_config, "architectures", None) or ()) + if "Qwen3MoeForCausalLM" not in architectures: + mismatches.append( + "architectures does not contain 'Qwen3MoeForCausalLM': " + f"{architectures!r}" + ) + if mismatches: + raise RuntimeError( + "checkpoint is not the expected Qwen3_235B-A22B architecture: " + + "; ".join(mismatches) + ) + return actual + + +def _build_qwen3_moe_tp_plan( + model_config: Any, + tp_size: int, + scenario: Scenario, + output_tokens: int, +) -> tuple[str | dict[str, str], dict[str, Any]]: + """Use the correctness-first TP8 layout validated on BW1100. + + Attention stays replicated. Qwen3_235B has four KV heads, so attempting to + tensor-parallelize attention over eight ranks produces an invalid local GQA + layout in this Transformers/DTK stack. Only MoE experts and the LM head are + sharded, matching the working Hygon container example. + """ + if getattr(model_config, "model_type", None) != "qwen3_moe": + raise RuntimeError( + "this benchmark only supports model_type='qwen3_moe', got " + f"{getattr(model_config, 'model_type', None)!r}" + ) + + global_query_heads = int(model_config.num_attention_heads) + global_kv_heads = int(model_config.num_key_value_heads) + head_dim = int(model_config.head_dim) + num_hidden_layers = int(model_config.num_hidden_layers) + if min( + global_query_heads, + global_kv_heads, + head_dim, + num_hidden_layers, + tp_size, + ) < 1: + raise RuntimeError("TP and attention dimensions must all be positive") + if global_query_heads % global_kv_heads: + raise RuntimeError( + f"global Q heads ({global_query_heads}) must be divisible by global " + f"KV heads ({global_kv_heads})" + ) + tp_plan = dict(HYGON_TP_PLAN) + maximum_sequence_tokens = scenario.input_tokens + output_tokens + dtype_bytes = 2 # BF16 K and V elements. + kv_cache_bytes_per_rank = ( + scenario.batch_size + * maximum_sequence_tokens + * num_hidden_layers + * 2 + * global_kv_heads + * head_dim + * dtype_bytes + ) + plan_payload = tp_plan if isinstance(tp_plan, dict) else {"mode": tp_plan} + metadata = { + "tp_plan_mode": f"qwen3_moe_tp{tp_size}_experts_lm_head_only", + "tp_plan_sha256": _stable_hash(plan_payload), + "attention_strategy": "replicated_eager", + "kv_projection_strategy": "replicated", + "kv_cache_replication_factor_across_tp_ranks": tp_size, + "global_query_heads": global_query_heads, + "global_kv_heads": global_kv_heads, + "head_dim": head_dim, + "num_hidden_layers": num_hidden_layers, + "local_query_heads": global_query_heads, + "local_kv_heads": global_kv_heads, + "local_gqa_groups": global_query_heads // global_kv_heads, + "maximum_sequence_tokens": maximum_sequence_tokens, + "estimated_dense_bf16_kv_cache_gib_per_rank": ( + kv_cache_bytes_per_rank / (1024**3) + ), + "kv_cache_estimate_excludes_allocator_and_cache_metadata": True, + } + return tp_plan, metadata + + +def _validate_and_set_local_gqa( + model: Any, tp_metadata: dict[str, Any] +) -> dict[str, Any]: + """Verify that attention stayed fully replicated on every TP rank.""" + base_model_prefix = getattr(model, "base_model_prefix", None) + base_model = getattr(model, base_model_prefix, None) + layers = getattr(base_model, "layers", None) + if layers is None: + raise RuntimeError( + f"cannot locate {base_model_prefix!r}.layers on loaded model" + ) + + head_dim = int(tp_metadata["head_dim"]) + expected_query_heads = int(tp_metadata["local_query_heads"]) + expected_kv_heads = int(tp_metadata["local_kv_heads"]) + expected_groups = int(tp_metadata["local_gqa_groups"]) + expected_query_width = expected_query_heads * head_dim + expected_kv_width = expected_kv_heads * head_dim + observed_groups: set[int] = set() + + for layer_index, layer in enumerate(layers): + attention = getattr(layer, "self_attn", None) + if attention is None: + raise RuntimeError(f"layer {layer_index} has no self_attn module") + query_width = int(attention.q_proj.out_features) + key_width = int(attention.k_proj.out_features) + value_width = int(attention.v_proj.out_features) + if query_width != expected_query_width: + raise RuntimeError( + f"layer {layer_index} local Q width={query_width}, expected " + f"{expected_query_width} ({expected_query_heads} heads)" + ) + if key_width != expected_kv_width or value_width != expected_kv_width: + raise RuntimeError( + f"layer {layer_index} local K/V widths={key_width}/{value_width}, " + f"expected {expected_kv_width} ({expected_kv_heads} heads)" + ) + observed_groups.add(int(attention.num_key_value_groups)) + + if observed_groups != {expected_groups}: + raise RuntimeError( + "attention GQA metadata changed despite replicated attention: " + f"observed={sorted(observed_groups)}, expected={expected_groups}" + ) + + expected_layers = int(tp_metadata["num_hidden_layers"]) + if len(layers) != expected_layers: + raise RuntimeError( + f"loaded model has {len(layers)} transformer layers, expected " + f"{expected_layers}" + ) + return { + "validated_attention_layers": len(layers), + "local_query_projection_width": expected_query_width, + "local_kv_projection_width": expected_kv_width, + "attention_replication_validated": True, + "local_gqa_groups": expected_groups, + } + + +def _launch_torchrun(args: argparse.Namespace) -> None: + env = os.environ.copy() + target_path = ( + "/root/.local/bin:/opt/dtk/cuda/cuda/bin:/opt/dtk/bin:/opt/dtk/hip/bin" + ) + target_library_path = ":".join( + ( + "/usr/local/lib/python3.10/dist-packages/torch/lib", + "/opt/dtk/dcc/gcvm/lib", + "/opt/dtk/hip/lib", + "/opt/dtk/llvm/lib", + "/opt/dtk/lib", + "/opt/dtk/lib64", + "/opt/hyhal/lib", + "/opt/hyhal/lib64", + "/opt/dtk/dushmem/lib", + "/opt/dtk/opencl/lib", + "/opt/ucx/lib", + "/opt/mpi/lib", + "/opt/hwloc/lib", + ) + ) + env["PATH"] = f"{target_path}:{env.get('PATH', '')}" + inherited_library_path = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ( + f"{target_library_path}:{inherited_library_path}" + if inherited_library_path + else target_library_path + ) + inherited_python_path = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = ( + f"/usr/local:{inherited_python_path}" + if inherited_python_path + else "/usr/local" + ) + visible_devices = ",".join(str(index) for index in range(args.tp_size)) + env.setdefault("HIP_VISIBLE_DEVICES", visible_devices) + env.setdefault("CUDA_VISIBLE_DEVICES", visible_devices) + env.setdefault("OMP_NUM_THREADS", "1") + env.setdefault("TOKENIZERS_PARALLELISM", "false") + env.setdefault("PYTHONUNBUFFERED", "1") + env.setdefault("HSA_FORCE_FINE_GRAIN_PCIE", "1") + env.setdefault("NCCL_DEBUG", "WARN") + + script = str(Path(sys.argv[0]).resolve()) + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + f"--nproc-per-node={args.tp_size}", + "--max-restarts=0", + "--monitor-interval=1", + script, + *sys.argv[1:], + ] + os.execvpe(sys.executable, command, env) + + +def _package_version(name: str) -> str | None: + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def _install_hygon_grouped_mm_guard(torch: Any) -> bool: + """Install the grouped-MM fallback validated in the BW1100 image.""" + if getattr(torch.version, "hip", None) is None: + return False + + from transformers.integrations import moe as transformers_moe + + if getattr(transformers_moe, "_hygon_grouped_mm_guard_installed", False): + return True + + def grouped_mm(input_tensor: Any, weight: Any, offs: Any) -> Any: + ends = [int(value) for value in offs.detach().cpu().tolist()] + output = input_tensor.new_empty( + (input_tensor.shape[0], weight.shape[-1]), dtype=weight.dtype + ) + start = 0 + for expert_index, end in enumerate(ends): + if end > start: + output[start:end] = input_tensor[start:end].to(weight.dtype).matmul( + weight[expert_index] + ) + start = end + return output + + transformers_moe._grouped_mm = grouped_mm + transformers_moe._hygon_grouped_mm_guard_installed = True + return True + + +def _stable_hash(value: Any) -> str: + payload = json.dumps(value, ensure_ascii=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _emit(rank: int, tag: str, payload: dict[str, Any]) -> None: + if rank != 0: + return + if tag == "PYTORCH_QWEN3_235B_CONFIG": + print( + "[Transformers] " + f"model={payload['model']} batch={payload['batch_size']} " + f"input={payload['input_lengths'][0]} " + f"output={payload['output_tokens_per_request']} " + f"tp={payload['tp_size']} attention={payload['attention_implementation']}", + flush=True, + ) + print( + f" load weights over! {payload['model_load_seconds'] * 1000.0:.2f} ms ", + flush=True, + ) + elif tag == "PYTORCH_QWEN3_235B_COMPLETE": + print(f"Transformers benchmark status: {payload['status']}", flush=True) + + +def _print_infinilm_style_metrics( + rank: int, measurement: dict[str, Any], decoded_output: str +) -> None: + if rank != 0: + return + print( + f"\n Generation completed in {measurement['generation_seconds'] * 1000.0:.2f} ms", + flush=True, + ) + print( + f" Batchsize={measurement['batch_size']} " + f"Per_Batch_Input_Len={measurement['input_tokens_per_request']} " + f"Per_Batch_New_Tokens={measurement['output_tokens_per_request']}", + flush=True, + ) + print( + f"\n Prefill TTFT: {measurement['ttft_seconds'] * 1000.0:.2f} ms " + f"Throughput: {measurement['prefill_tokens_per_second']:.2f} tok/s", + flush=True, + ) + print( + f"\n Decode Avg ITL: {measurement['inter_token_latency_ms']:.2f} ms " + f"Throughput: {measurement['decode_tokens_per_second']:.2f} tok/s\n", + flush=True, + ) + print(decoded_output or "（未生成可显示文本）", flush=True) + + +def _validate_decoded_output(decoded_output: str) -> dict[str, Any]: + text = decoded_output.strip() + if not text: + raise RuntimeError("generated output decoded to an empty string") + if "\ufffd" in text: + raise RuntimeError("generated output contains Unicode replacement characters") + printable_characters = sum( + character.isprintable() or character in "\n\t" for character in text + ) + printable_ratio = printable_characters / len(text) + cjk_characters = sum("\u3400" <= character <= "\u9fff" for character in text) + url_fragments = text.lower().count("http") + if printable_ratio < 0.95: + raise RuntimeError( + f"generated output printable ratio is too low: {printable_ratio:.3f}" + ) + if cjk_characters < 8: + raise RuntimeError( + f"generated output is not a substantive Chinese response: CJK={cjk_characters}" + ) + if url_fragments > 2: + raise RuntimeError( + f"generated output contains suspicious URL fragments: {url_fragments}" + ) + return { + "nonempty_decoded_text": True, + "no_replacement_characters": True, + "printable_ratio": printable_ratio, + "cjk_character_count": cjk_characters, + "url_fragment_count": url_fragments, + } + + +def _max_across_ranks(value: float, torch: Any, dist: Any, device: Any) -> float: + tensor = torch.tensor(value, dtype=torch.float64, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return float(tensor.item()) + + +def _make_prompt_base( + tokenizer: Any, prompt_file: str +) -> tuple[list[int], dict[str, Any]]: + prompt_path = Path(prompt_file).resolve() + prompt_file_exists = prompt_path.is_file() + prompt_text = ( + prompt_path.read_text(encoding="utf-8") + if prompt_file_exists + else FALLBACK_PROMPT + ).strip() + if not prompt_text: + raise RuntimeError(f"benchmark prompt file is empty: {prompt_path}") + if not getattr(tokenizer, "chat_template", None): + raise RuntimeError("model tokenizer does not define a chat template") + rendered_prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt_text}], + tokenize=False, + add_generation_prompt=True, + ) + # Match InfiniLM's benchmark path: use tokenizer.encode defaults after the + # model's chat template, then repeat this exact base sequence to each length. + token_ids = list(tokenizer.encode(rendered_prompt)) + if not token_ids: + raise RuntimeError("the fixed benchmark prompt tokenized to an empty list") + return token_ids, { + "prompt_source": ( + "file_chat_template" if prompt_file_exists else "embedded_fallback" + ), + "prompt_file": str(prompt_path), + "prompt_file_exists": prompt_file_exists, + "prompt_file_sha256": hashlib.sha256( + prompt_text.encode("utf-8") + ).hexdigest(), + "rendered_prompt_sha256": hashlib.sha256( + rendered_prompt.encode("utf-8") + ).hexdigest(), + } + + +def _repeat_prompt(token_ids: Sequence[int], target_length: int) -> list[int]: + if target_length < 1: + raise ValueError("target_length must be positive") + base = list(token_ids) + if target_length <= len(base): + # Preserve the assistant-generation suffix instead of cutting it off. + result = base[-target_length:] + else: + prefix_length = target_length - len(base) + repeats = (prefix_length + len(base) - 1) // len(base) + result = (base * repeats)[:prefix_length] + base + if len(result) != target_length: + raise RuntimeError(f"expected {target_length} prompt tokens, got {len(result)}") + return result + + +def _materialize_logits(logits: Any) -> Any: + # A replicated TP lm_head returns Tensor. Keep this guard for TP plans that + # leave the vocabulary output as a DTensor. + if type(logits).__name__ == "DTensor" and hasattr(logits, "full_tensor"): + return logits.full_tensor() + return logits + + +def _forward( + model: Any, + logits_limit_argument: str, + input_ids: Any, + past_key_values: Any | None = None, +) -> tuple[Any, Any]: + kwargs: dict[str, Any] = { + "input_ids": input_ids, + "past_key_values": past_key_values, + "use_cache": True, + "return_dict": True, + logits_limit_argument: 1, + } + outputs = model(**kwargs) + if outputs.past_key_values is None: + raise RuntimeError("model did not return past_key_values with use_cache=True") + logits = _materialize_logits(outputs.logits) + if logits.ndim != 3 or logits.shape[0] != input_ids.shape[0]: + raise RuntimeError(f"unexpected logits shape: {tuple(logits.shape)}") + if logits.shape[1] != 1: + raise RuntimeError( + f"expected one retained logits position, got shape {tuple(logits.shape)}" + ) + return logits[:, -1, :], outputs.past_key_values + + +def _validate_output( + generated: Any, + last_logits: Any, + batch_size: int, + output_tokens: int, + vocab_size: int, + torch: Any, + dist: Any, +) -> tuple[Any, dict[str, Any]]: + expected_shape = (batch_size, output_tokens) + if tuple(generated.shape) != expected_shape: + raise RuntimeError( + f"expected generated shape {expected_shape}, got {tuple(generated.shape)}" + ) + + rank_minimum = generated.clone() + rank_maximum = generated.clone() + dist.all_reduce(rank_minimum, op=dist.ReduceOp.MIN) + dist.all_reduce(rank_maximum, op=dist.ReduceOp.MAX) + rank_consensus = bool(torch.equal(rank_minimum, rank_maximum)) + + valid_ids = bool( + torch.logical_and(generated >= 0, generated < vocab_size).all().item() + ) + finite_logits = bool(torch.isfinite(last_logits).all().item()) + checks = torch.tensor( + [int(rank_consensus), int(valid_ids), int(finite_logits)], + dtype=torch.int32, + device=generated.device, + ) + dist.all_reduce(checks, op=dist.ReduceOp.MIN) + rank_consensus, valid_ids, finite_logits = [bool(value) for value in checks.tolist()] + if not (rank_consensus and valid_ids and finite_logits): + raise RuntimeError( + "correctness validation failed: " + f"rank_consensus={rank_consensus}, valid_ids={valid_ids}, " + f"finite_logits={finite_logits}" + ) + + generated_cpu = generated.cpu() + matrix = generated_cpu.tolist() + return generated_cpu, { + "exact_output_shape": True, + "rank_consensus": rank_consensus, + "valid_token_ids": valid_ids, + "finite_last_logits": finite_logits, + "output_token_ids_sha256": _stable_hash(matrix), + "first_request_first_16_tokens": matrix[0][:16], + } + + +def _run_iteration( + model: Any, + prompt_base: Sequence[int], + batch_size: int, + input_tokens: int, + output_tokens: int, + vocab_size: int, + logits_limit_argument: str, + torch: Any, + dist: Any, + device: Any, +) -> dict[str, Any]: + prompt = _repeat_prompt(prompt_base, input_tokens) + input_ids = ( + torch.tensor(prompt, dtype=torch.long, device=device) + .unsqueeze(0) + .expand(batch_size, -1) + .contiguous() + ) + + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + prefill_start = time.perf_counter() + logits, past_key_values = _forward( + model, logits_limit_argument, input_ids, past_key_values=None + ) + next_token = torch.argmax(logits, dim=-1) + torch.cuda.synchronize(device) + prefill_local_seconds = time.perf_counter() - prefill_start + prefill_seconds = _max_across_ranks( + prefill_local_seconds, torch, dist, device + ) + + generated_tokens = [next_token] + decode_start = time.perf_counter() + for _ in range(output_tokens - 1): + logits, past_key_values = _forward( + model, + logits_limit_argument, + next_token.unsqueeze(1), + past_key_values=past_key_values, + ) + next_token = torch.argmax(logits, dim=-1) + generated_tokens.append(next_token) + torch.cuda.synchronize(device) + decode_local_seconds = time.perf_counter() - decode_start + decode_seconds = _max_across_ranks(decode_local_seconds, torch, dist, device) + + peak_allocated_gib = _max_across_ranks( + torch.cuda.max_memory_allocated(device) / (1024**3), torch, dist, device + ) + peak_reserved_gib = _max_across_ranks( + torch.cuda.max_memory_reserved(device) / (1024**3), torch, dist, device + ) + generated = torch.stack(generated_tokens, dim=1) + generated_cpu, correctness = _validate_output( + generated, + logits, + batch_size, + output_tokens, + vocab_size, + torch, + dist, + ) + + total_prompt_tokens = batch_size * input_tokens + decode_token_count = batch_size * (output_tokens - 1) + generated_token_count = batch_size * output_tokens + total_seconds = prefill_seconds + decode_seconds + result = { + "batch_size": batch_size, + "input_tokens_per_request": input_tokens, + "prompt_token_ids_sha256": _stable_hash(prompt), + "output_tokens_per_request": output_tokens, + "total_context_tokens_per_request": input_tokens + output_tokens, + "total_prompt_tokens": total_prompt_tokens, + "total_generated_tokens": generated_token_count, + "ttft_seconds": prefill_seconds, + "prefill_tokens_per_second": total_prompt_tokens / prefill_seconds, + "decode_seconds": decode_seconds, + "decode_tokens_per_second": decode_token_count / decode_seconds, + "decode_tokens_per_second_per_request": ( + (output_tokens - 1) / decode_seconds + ), + "inter_token_latency_ms": decode_seconds * 1000.0 / (output_tokens - 1), + "generation_seconds": total_seconds, + "generated_tokens_per_second": generated_token_count / total_seconds, + "peak_memory_allocated_gib_max_rank": peak_allocated_gib, + "peak_memory_reserved_gib_max_rank": peak_reserved_gib, + "correctness": correctness, + "first_request_output_token_ids": generated_cpu[0].tolist(), + } + + del generated_cpu, generated, generated_tokens, logits, next_token + del past_key_values, input_ids + return result + + +def _median_summary(measurements: Sequence[dict[str, Any]]) -> dict[str, Any]: + fields = ( + "ttft_seconds", + "prefill_tokens_per_second", + "decode_seconds", + "decode_tokens_per_second", + "decode_tokens_per_second_per_request", + "inter_token_latency_ms", + "generation_seconds", + "generated_tokens_per_second", + "peak_memory_allocated_gib_max_rank", + "peak_memory_reserved_gib_max_rank", + ) + return { + f"median_{field}": statistics.median( + float(measurement[field]) for measurement in measurements + ) + for field in fields + } + + +def _per_length_medians( + measurements: Sequence[dict[str, Any]], input_lengths: Sequence[int] +) -> list[dict[str, Any]]: + summaries: list[dict[str, Any]] = [] + for input_tokens in input_lengths: + records = [ + measurement + for measurement in measurements + if int(measurement["input_tokens_per_request"]) == input_tokens + ] + if len(records) != REPEATS_PER_INPUT_LENGTH: + raise RuntimeError( + f"input length {input_tokens}: recorded {len(records)} repeats; " + f"expected {REPEATS_PER_INPUT_LENGTH}" + ) + summaries.append( + { + "input_tokens_per_request": input_tokens, + "measured_repeats": len(records), + **_median_summary(records), + } + ) + return summaries + + +def _overall_median_of_per_length_medians( + per_length_medians: Sequence[dict[str, Any]], +) -> dict[str, float]: + metric_names = [ + name for name in per_length_medians[0] if name.startswith("median_") + ] + return { + name: statistics.median( + float(length_summary[name]) for length_summary in per_length_medians + ) + for name in metric_names + } + + +def _run_worker_impl(args: argparse.Namespace, scenario: Scenario) -> int: + import inspect + + import torch + import torch.distributed as dist + import transformers + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + + grouped_mm_fallback = _install_hygon_grouped_mm_guard(torch) + + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", world_size)) + if rank != 0: + transformers.utils.logging.disable_progress_bar() + if world_size != args.tp_size: + raise RuntimeError( + f"expected WORLD_SIZE={args.tp_size}, got WORLD_SIZE={world_size}" + ) + if local_world_size != args.tp_size: + raise RuntimeError( + "this benchmark requires all TP ranks on one host: " + f"LOCAL_WORLD_SIZE={local_world_size}, TP={args.tp_size}" + ) + if not torch.cuda.is_available(): + raise RuntimeError("torch.cuda is unavailable") + if torch.cuda.device_count() < local_world_size: + raise RuntimeError( + f"need {local_world_size} visible GPUs, found {torch.cuda.device_count()}" + ) + + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + # Each torchrun worker owns one device. Avoid manual_seed_all(), which can + # make every worker initialize a context on all eight visible GPUs. + torch.random.default_generator.manual_seed(0) + torch.cuda.manual_seed(0) + + model_config = AutoConfig.from_pretrained( + args.model, + local_files_only=True, + trust_remote_code=False, + ) + architecture_signature = _validate_qwen3_235b_architecture(model_config) + quantization_config = getattr(model_config, "quantization_config", None) + if quantization_config: + raise RuntimeError( + "the Transformers benchmark is BF16-only; refusing quantized " + f"checkpoint {args.model!r} with quantization_config=" + f"{quantization_config!r}" + ) + tp_plan, tp_metadata = _build_qwen3_moe_tp_plan( + model_config, + args.tp_size, + scenario, + args.output_tokens, + ) + + load_start = time.perf_counter() + model = AutoModelForCausalLM.from_pretrained( + args.model, + config=model_config, + dtype=torch.bfloat16, + attn_implementation=args.attention, + tp_plan=tp_plan, + local_files_only=True, + low_cpu_mem_usage=True, + trust_remote_code=False, + ) + tp_validation = _validate_and_set_local_gqa(model, tp_metadata) + model.eval() + torch.cuda.synchronize(device) + if not dist.is_initialized(): + raise RuntimeError( + "Transformers tp_plan='auto' did not initialize torch.distributed" + ) + load_seconds = _max_across_ranks( + time.perf_counter() - load_start, torch, dist, device + ) + + tp_plan = getattr(model, "_tp_plan", None) + if not tp_plan: + raise RuntimeError("model loaded without a non-empty Transformers TP plan") + resolved_attention = getattr(model.config, "_attn_implementation", None) + if resolved_attention != args.attention: + raise RuntimeError( + f"requested attention={args.attention!r}, loaded model resolved " + f"attention={resolved_attention!r}" + ) + forward_parameters = inspect.signature(model.forward).parameters + if "logits_to_keep" in forward_parameters: + logits_limit_argument = "logits_to_keep" + elif "num_logits_to_keep" in forward_parameters: + logits_limit_argument = "num_logits_to_keep" + else: + raise RuntimeError( + "model.forward has no logits_to_keep argument; refusing to materialize " + "full [batch, context, vocab] logits for this benchmark" + ) + + tokenizer = AutoTokenizer.from_pretrained( + args.model, + local_files_only=True, + trust_remote_code=False, + use_fast=True, + ) + prompt_base, prompt_metadata = _make_prompt_base(tokenizer, args.prompt_file) + vocab_size = int(model.config.vocab_size) + if any(token < 0 or token >= vocab_size for token in prompt_base): + raise RuntimeError("fixed prompt contains a token outside model vocabulary") + + maximum_position_embeddings = int( + getattr(model.config, "max_position_embeddings", 0) or 0 + ) + maximum_requested = scenario.input_tokens + args.output_tokens + if maximum_position_embeddings and maximum_requested > maximum_position_embeddings: + raise RuntimeError( + f"requested sequence length {maximum_requested} exceeds " + f"max_position_embeddings={maximum_position_embeddings}" + ) + + config = { + "framework": "transformers", + "scenario": scenario.name, + "model_name": MODEL_NAME, + "model": str(Path(args.model).absolute()), + "model_realpath": str(Path(args.model).resolve()), + "model_class": type(model).__name__, + "dtype": "bfloat16", + "checkpoint_quantized": False, + "validated_qwen3_235b_architecture": architecture_signature, + "attention_implementation": resolved_attention, + "tp_plan": tp_metadata["tp_plan_mode"], + "tp_plan_rules": tp_plan, + "tp_plan_rule_count": len(tp_plan) if isinstance(tp_plan, dict) else None, + "tp_size": args.tp_size, + "smoke": args.smoke, + "batch_size": scenario.batch_size, + "input_lengths": list(scenario.input_lengths), + "output_tokens_per_request": args.output_tokens, + "total_context_tokens_per_request": maximum_requested, + "measured_iterations": MEASURED_ITERATIONS, + "measured_input_lengths": MEASURED_INPUT_LENGTHS, + "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, + "measurement_semantics": MEASUREMENT_SEMANTICS, + "model_load_seconds": load_seconds, + "fixed_prompt_base_tokens": len(prompt_base), + "fixed_prompt_base_sha256": _stable_hash(prompt_base), + **prompt_metadata, + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + "flash_attn_version": _package_version("flash-attn"), + "hygon_transformers_grouped_mm_fallback": grouped_mm_fallback, + "gpu_name": torch.cuda.get_device_name(device), + **tp_metadata, + **tp_validation, + } + _emit(rank, "PYTORCH_QWEN3_235B_CONFIG", config) + + measurements: list[dict[str, Any]] = [] + iteration = 0 + for length_index, input_tokens in enumerate(scenario.input_lengths, start=1): + with torch.inference_mode(): + shape_warmup = _run_iteration( + model, + prompt_base, + scenario.batch_size, + input_tokens, + args.output_tokens, + vocab_size, + logits_limit_argument, + torch, + dist, + device, + ) + shape_warmup_hash = shape_warmup["correctness"][ + "output_token_ids_sha256" + ] + shape_warmup_prompt_hash = shape_warmup["prompt_token_ids_sha256"] + _emit( + rank, + "PYTORCH_QWEN3_235B_SHAPE_WARMUP", + { + "scenario": scenario.name, + "length_index": length_index, + "batch_size": scenario.batch_size, + "input_tokens_per_request": input_tokens, + "prompt_token_ids_sha256": shape_warmup_prompt_hash, + "output_tokens_per_request": args.output_tokens, + "output_token_ids_sha256": shape_warmup_hash, + "correctness": shape_warmup["correctness"], + }, + ) + del shape_warmup + gc.collect() + + for repeat in range(1, REPEATS_PER_INPUT_LENGTH + 1): + iteration += 1 + with torch.inference_mode(): + measurement = _run_iteration( + model, + prompt_base, + scenario.batch_size, + input_tokens, + args.output_tokens, + vocab_size, + logits_limit_argument, + torch, + dist, + device, + ) + measured_hash = measurement["correctness"][ + "output_token_ids_sha256" + ] + measured_prompt_hash = measurement["prompt_token_ids_sha256"] + if measured_prompt_hash != shape_warmup_prompt_hash: + raise RuntimeError( + f"input length {input_tokens} repeat {repeat}: measured prompt " + f"hash {measured_prompt_hash} does not match exact-shape " + f"warmup prompt hash {shape_warmup_prompt_hash}" + ) + # Hygon BF16 kernels can make numerically valid MoE routing choices + # differ across independent runs. Keep the replay hash observable, + # while treating per-run shape/range/finite/rank checks as correctness. + measurement["correctness"]["output_matches_exact_shape_warmup"] = ( + measured_hash == shape_warmup_hash + ) + measurement["exact_shape_warmup_output_sha256"] = shape_warmup_hash + measurement = { + "scenario": scenario.name, + "iteration": iteration, + "length_index": length_index, + "repeat": repeat, + **measurement, + } + output_token_ids = measurement.pop("first_request_output_token_ids") + decoded_output = tokenizer.decode( + output_token_ids, skip_special_tokens=True + ).strip() + measurement["correctness"]["decoded_output"] = ( + _validate_decoded_output(decoded_output) + ) + measurements.append(measurement) + _print_infinilm_style_metrics(rank, measurement, decoded_output) + _emit(rank, "PYTORCH_QWEN3_235B_ITERATION", measurement) + gc.collect() + + if len(measurements) != MEASURED_ITERATIONS: + raise RuntimeError( + f"expected {MEASURED_ITERATIONS} measurements, got {len(measurements)}" + ) + per_length_medians = _per_length_medians( + measurements, scenario.input_lengths + ) + overall_medians = _overall_median_of_per_length_medians(per_length_medians) + summary = { + "scenario": scenario.name, + "status": "PASS", + "measured_iterations": len(measurements), + "measured_input_lengths": MEASURED_INPUT_LENGTHS, + "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, + "measurement_semantics": MEASUREMENT_SEMANTICS, + "input_lengths": list(scenario.input_lengths), + "batch_size": scenario.batch_size, + "output_tokens_per_request": args.output_tokens, + "total_context_tokens_per_request": maximum_requested, + "per_length_medians": per_length_medians, + "overall_aggregate": { + "aggregation_method": "median_of_three_fixed_shape_measurements", + "measurement_count": len(measurements), + "mixed_input_lengths": False, + **overall_medians, + }, + # Compatibility aliases for existing table consumers. Their scope is the + # explicitly labeled overall aggregate above, not a single input length. + **overall_medians, + "output_token_ids_sha256": [ + item["correctness"]["output_token_ids_sha256"] for item in measurements + ], + } + _emit(rank, "PYTORCH_QWEN3_235B_SUMMARY", summary) + return len(measurements) + + +def _run_worker(args: argparse.Namespace, scenario: Scenario) -> None: + import torch.distributed as dist + + rank = int(os.environ["RANK"]) + measured_iterations = 0 + caught: BaseException | None = None + caught_traceback: Any = None + teardown_errors: list[str] = [] + process_group_was_initialized = False + try: + measured_iterations = _run_worker_impl(args, scenario) + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + finally: + process_group_was_initialized = dist.is_initialized() + if process_group_was_initialized: + if caught is None: + try: + dist.barrier() + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"barrier: {type(error).__name__}: {error}" + ) + try: + dist.destroy_process_group() + except BaseException as error: + if caught is None: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"destroy_process_group: {type(error).__name__}: {error}" + ) + + teardown_complete = not dist.is_initialized() and not teardown_errors + status = ( + "PASS" + if caught is None + and measured_iterations == MEASURED_ITERATIONS + and teardown_complete + else "ERROR" + ) + completion: dict[str, Any] = { + "scenario": scenario.name, + "status": status, + "exit_code": 0 if status == "PASS" else 1, + "measured_iterations": measured_iterations, + "measured_input_lengths": MEASURED_INPUT_LENGTHS, + "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, + "measurement_semantics": MEASUREMENT_SEMANTICS, + "process_group_was_initialized": process_group_was_initialized, + "distributed_teardown_complete": teardown_complete, + } + if caught is not None: + completion["error"] = { + "type": type(caught).__name__, + "message": str(caught), + } + if teardown_errors: + completion["teardown_errors"] = teardown_errors + _emit(rank, "PYTORCH_QWEN3_235B_COMPLETE", completion) + + if caught is not None: + raise caught.with_traceback(caught_traceback) + + +def main(scenario: Scenario) -> None: + args = _parse_args(scenario) + scenario = _effective_scenario(scenario, args.smoke) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if "LOCAL_RANK" not in os.environ and world_size == 1: + _require_idle_gpu() + _launch_torchrun(args) + raise AssertionError("os.execvpe returned unexpectedly") + _run_worker(args, scenario) diff --git a/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py new file mode 100755 index 00000000..5474ca99 --- /dev/null +++ b/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py @@ -0,0 +1,491 @@ +#!/usr/bin/env python3 +"""Calculate true token-level PPL for Qwen3_235B with Transformers TP8. + +The input is a framework-neutral token manifest. Both the Transformers and +InfiniLM runners must consume the same manifest so their PPL values score the +same target tokens instead of independently tokenizing the source corpus. +""" + +from __future__ import annotations + +import argparse +import gc +import inspect +import json +import math +import os +import sys +import time +from pathlib import Path +from typing import Any + + +SCRIPT_DIR = Path(__file__).resolve().parent +SCRIPTS_DIR = SCRIPT_DIR.parent +for import_path in (SCRIPT_DIR, SCRIPTS_DIR): + if str(import_path) not in sys.path: + sys.path.insert(0, str(import_path)) + +import _pytorch_runner as benchmark_runner +from _ppl_common import ( + SCORING_METHOD, + canonical_indices_sha256, + iter_sliding_windows, + load_manifest, +) + + +RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" +DEFAULT_MODEL = "/data1/Qwen3_235B" +DEFAULT_WINDOW_SIZE = 256 +DEFAULT_STRIDE = 128 +DEFAULT_MAX_SCORED_TOKENS = 10240 +EXPECTED_VOCAB_SIZE = 151936 + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="True shifted-token PPL for Qwen3_235B BF16 on Hygon TP8" + ) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--token-manifest", required=True) + parser.add_argument("--window", type=int, default=DEFAULT_WINDOW_SIZE) + parser.add_argument("--stride", type=int, default=DEFAULT_STRIDE) + parser.add_argument( + "--max-scored-tokens", + type=int, + default=DEFAULT_MAX_SCORED_TOKENS, + help="maximum shifted target tokens to score; 0 scores the full manifest", + ) + parser.add_argument("--tp-size", type=int, default=8) + parser.add_argument( + "--attention", + choices=("eager",), + default="eager", + help="BW1100 correctness path; SDPA is unsupported for this model stack", + ) + parser.add_argument( + "--json-output", + help="optional rank-0 result path; the result is always printed as JSON", + ) + args = parser.parse_args() + + model_path = Path(args.model) + manifest_path = Path(args.token_manifest) + if not model_path.is_dir(): + parser.error(f"model directory does not exist: {model_path}") + if not manifest_path.is_file(): + parser.error(f"token manifest does not exist: {manifest_path}") + if args.window < 2: + parser.error("--window must be at least 2") + if args.stride < 1 or args.stride >= args.window: + parser.error("--stride must satisfy 1 <= stride < window") + if args.max_scored_tokens < 0: + parser.error("--max-scored-tokens must be non-negative") + if args.tp_size < 1: + parser.error("--tp-size must be positive") + + args.model = str(model_path.resolve()) + args.token_manifest = str(manifest_path.resolve()) + if args.json_output: + args.json_output = str(Path(args.json_output).resolve()) + return args + +def _write_json_atomic(path_value: str, payload: dict[str, Any]) -> None: + path = Path(path_value) + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") + temporary.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + temporary.replace(path) + + +def _run_worker_impl(args: argparse.Namespace) -> dict[str, Any]: + import torch + import torch.distributed as dist + import torch.nn.functional as functional + import transformers + from transformers import AutoConfig, AutoModelForCausalLM + + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", world_size)) + if world_size != args.tp_size or local_world_size != args.tp_size: + raise RuntimeError( + "PPL runner requires one-host TP with WORLD_SIZE=" + f"LOCAL_WORLD_SIZE={args.tp_size}; got {world_size}/{local_world_size}" + ) + if not torch.cuda.is_available() or torch.cuda.device_count() < local_world_size: + raise RuntimeError( + f"need {local_world_size} visible GPUs, found {torch.cuda.device_count()}" + ) + + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + # RCCL initializes lazily. Reserve communicator memory before the 235B + # checkpoint consumes nearly all device memory. + communicator_probe = torch.ones(1, dtype=torch.int32, device=device) + dist.all_reduce(communicator_probe) + if int(communicator_probe.item()) != args.tp_size: + raise RuntimeError("Transformers TP8 RCCL communicator probe failed") + dist.barrier() + torch.cuda.synchronize(device) + del communicator_probe + torch.random.default_generator.manual_seed(0) + torch.cuda.manual_seed(0) + if rank != 0: + transformers.utils.logging.disable_progress_bar() + + corpus_manifest = load_manifest(args.token_manifest) + token_ids = corpus_manifest.token_ids + corpus = { + "manifest_path": str(corpus_manifest.path.resolve()), + "manifest_sha256": corpus_manifest.manifest_sha256, + "token_ids_sha256": corpus_manifest.token_ids_sha256, + "token_count": corpus_manifest.token_count, + "source_sha256": corpus_manifest.payload["source_sha256"], + "tokenizer_sha256": corpus_manifest.payload["tokenizer_sha256"], + "source": corpus_manifest.payload.get("source"), + "tokenizer": corpus_manifest.payload.get("tokenizer"), + } + model_config = AutoConfig.from_pretrained( + args.model, local_files_only=True, trust_remote_code=False + ) + architecture = benchmark_runner._validate_qwen3_235b_architecture(model_config) + if getattr(model_config, "quantization_config", None): + raise RuntimeError("Transformers PPL baseline requires the BF16 checkpoint") + vocab_size = int(model_config.vocab_size) + if vocab_size != EXPECTED_VOCAB_SIZE: + raise RuntimeError( + f"expected complete Qwen3_235B vocabulary {EXPECTED_VOCAB_SIZE}, " + f"got {vocab_size}" + ) + if any(token >= vocab_size for token in token_ids): + raise RuntimeError("token manifest contains an ID outside the model vocabulary") + maximum_positions = int( + getattr(model_config, "max_position_embeddings", 0) or 0 + ) + if maximum_positions and args.window > maximum_positions: + raise RuntimeError( + f"window={args.window} exceeds max_position_embeddings={maximum_positions}" + ) + + scoring_scenario = benchmark_runner.Scenario("true_ppl", 1, args.window, 1) + tp_plan, tp_metadata = benchmark_runner._build_qwen3_moe_tp_plan( + model_config, args.tp_size, scoring_scenario, 1 + ) + grouped_mm_fallback = benchmark_runner._install_hygon_grouped_mm_guard(torch) + load_start = time.perf_counter() + model = AutoModelForCausalLM.from_pretrained( + args.model, + config=model_config, + dtype=torch.bfloat16, + attn_implementation=args.attention, + tp_plan=tp_plan, + local_files_only=True, + low_cpu_mem_usage=True, + trust_remote_code=False, + ) + tp_validation = benchmark_runner._validate_and_set_local_gqa(model, tp_metadata) + model.eval() + torch.cuda.synchronize(device) + load_seconds = benchmark_runner._max_across_ranks( + time.perf_counter() - load_start, torch, dist, device + ) + + loaded_tp_plan = getattr(model, "_tp_plan", None) + if not loaded_tp_plan: + raise RuntimeError("model loaded without a non-empty Transformers TP plan") + resolved_attention = getattr(model.config, "_attn_implementation", None) + if resolved_attention != args.attention: + raise RuntimeError( + f"requested attention={args.attention!r}, resolved={resolved_attention!r}" + ) + forward_parameters = inspect.signature(model.forward).parameters + if "logits_to_keep" in forward_parameters: + logits_limit_argument = "logits_to_keep" + elif "num_logits_to_keep" in forward_parameters: + logits_limit_argument = "num_logits_to_keep" + else: + raise RuntimeError( + "model.forward has no logits_to_keep argument; refusing full-context " + "vocabulary materialization" + ) + + available_targets = len(token_ids) - 1 + scored_token_count = ( + available_targets + if args.max_scored_tokens == 0 + else min(args.max_scored_tokens, available_targets) + ) + if scored_token_count < 1: + raise RuntimeError("the selected corpus range contains no shifted target token") + first_scored_token_index = 1 + last_scored_token_index_exclusive = 1 + scored_token_count + scored_token_indices_sha256 = canonical_indices_sha256( + range(first_scored_token_index, last_scored_token_index_exclusive) + ) + scoring_method = SCORING_METHOD + + config = { + "backend": "transformers", + "model": args.model, + "dtype": "bfloat16", + "tp_size": args.tp_size, + "attention": resolved_attention, + "window_size": args.window, + "stride": args.stride, + "requested_max_scored_tokens": args.max_scored_tokens, + "scored_token_count": scored_token_count, + "scoring_method": scoring_method, + "first_scored_token_index": first_scored_token_index, + "last_scored_token_index_exclusive": last_scored_token_index_exclusive, + "scored_token_indices_sha256": scored_token_indices_sha256, + "corpus_manifest": corpus, + "vocab_size": vocab_size, + "architecture": architecture, + "model_load_seconds": load_seconds, + "tp_plan": tp_metadata, + "tp_validation": tp_validation, + "hygon_transformers_grouped_mm_fallback": grouped_mm_fallback, + "torch_version": torch.__version__, + "transformers_version": transformers.__version__, + } + if rank == 0: + print( + "PYTORCH_QWEN3_235B_PPL_CONFIG " + + json.dumps(config, ensure_ascii=False, sort_keys=True), + flush=True, + ) + + total_nll = 0.0 + windows: list[dict[str, Any]] = [] + scored_by_windows = 0 + torch.cuda.synchronize(device) + scoring_start = time.perf_counter() + with torch.inference_mode(): + for window in iter_sliding_windows( + token_ids, + args.window, + args.stride, + None if args.max_scored_tokens == 0 else args.max_scored_tokens, + ): + display_index = window.index + 1 + target_count = window.scored_token_count + input_slice = window.token_ids + expected_prediction_start = len(input_slice) - target_count - 1 + if ( + window.prediction_start != expected_prediction_start + or window.prediction_end != len(input_slice) - 1 + ): + raise RuntimeError( + f"window {display_index} retained-logits alignment is invalid" + ) + input_ids = torch.tensor( + input_slice, dtype=torch.long, device=device + ).unsqueeze(0) + outputs = model( + input_ids=input_ids, + use_cache=False, + return_dict=True, + **{logits_limit_argument: target_count + 1}, + ) + logits = benchmark_runner._materialize_logits(outputs.logits) + expected_shape = (1, target_count + 1, vocab_size) + if tuple(logits.shape) != expected_shape: + raise RuntimeError( + "incomplete or unexpected logits: " + f"got {tuple(logits.shape)}, expected {expected_shape}" + ) + score_logits = logits[:, :-1, :] + labels = torch.tensor( + input_slice[window.target_start : window.target_end], + dtype=torch.long, + device=device, + ).unsqueeze(0) + finite = torch.isfinite(score_logits).all().to(dtype=torch.int32) + dist.all_reduce(finite, op=dist.ReduceOp.MIN) + if not bool(finite.item()): + raise RuntimeError( + f"window {display_index} contains non-finite logits" + ) + window_nll_tensor = functional.cross_entropy( + score_logits.float().reshape(-1, vocab_size), + labels.reshape(-1), + reduction="sum", + ) + window_nll = float(window_nll_tensor.double().item()) + nll_min = torch.tensor(window_nll, dtype=torch.float64, device=device) + nll_max = nll_min.clone() + dist.all_reduce(nll_min, op=dist.ReduceOp.MIN) + dist.all_reduce(nll_max, op=dist.ReduceOp.MAX) + rank_delta_per_token = float((nll_max - nll_min).item()) / target_count + if rank_delta_per_token > 1e-4: + raise RuntimeError( + f"window {display_index} rank NLL mismatch: " + f"delta/token={rank_delta_per_token:.6g}" + ) + total_nll += window_nll + scored_by_windows += target_count + windows.append( + { + "index": window.index, + "token_start": window.token_start, + "token_end": window.token_end, + "score_start": window.score_start, + "score_end": window.score_end, + "input_token_count": len(input_slice), + "scored_token_count": target_count, + "nll": window_nll, + } + ) + del outputs, logits, score_logits, labels, window_nll_tensor, input_ids + + torch.cuda.synchronize(device) + scoring_seconds = benchmark_runner._max_across_ranks( + time.perf_counter() - scoring_start, torch, dist, device + ) + if scored_by_windows != scored_token_count: + raise RuntimeError( + f"scored {scored_by_windows} tokens, expected {scored_token_count}" + ) + + mean_nll = total_nll / scored_token_count + if not math.isfinite(mean_nll): + raise RuntimeError(f"mean NLL is not finite: {mean_nll}") + try: + ppl = math.exp(mean_nll) + except OverflowError as error: + raise RuntimeError(f"PPL overflows float64 at mean NLL={mean_nll}") from error + if not math.isfinite(ppl): + raise RuntimeError(f"PPL is not finite: {ppl}") + + result: dict[str, Any] = {} + if rank == 0: + result = { + "schema": RESULT_SCHEMA, + "status": "PASS", + "backend": "transformers", + "model": args.model, + "dtype": "bfloat16", + "tp_size": args.tp_size, + "attention": resolved_attention, + "corpus_manifest": args.token_manifest, + "corpus_manifest_sha256": corpus["manifest_sha256"], + "corpus_token_ids_sha256": corpus["token_ids_sha256"], + "corpus_token_count": corpus["token_count"], + "window_size": args.window, + "stride": args.stride, + "scoring_method": scoring_method, + "first_scored_token_index": first_scored_token_index, + "last_scored_token_index_exclusive": ( + last_scored_token_index_exclusive + ), + "scored_token_indices_sha256": scored_token_indices_sha256, + "scored_token_count": scored_token_count, + "total_nll": total_nll, + "mean_nll": mean_nll, + "ppl": ppl, + "windows": windows, + "window_count": len(windows), + "scoring_seconds": scoring_seconds, + "scored_tokens_per_second": scored_token_count / scoring_seconds, + "model_load_seconds": load_seconds, + "vocab_size": vocab_size, + "full_vocab_logits_validated_every_window": True, + } + if args.json_output: + _write_json_atomic(args.json_output, result) + print( + "PYTORCH_QWEN3_235B_PPL_RESULT " + + json.dumps(result, ensure_ascii=False, sort_keys=True), + flush=True, + ) + print( + f"Transformers true PPL: {ppl:.6f} " + f"(mean NLL={mean_nll:.6f}, tokens={scored_token_count})", + flush=True, + ) + + del model + gc.collect() + torch.cuda.empty_cache() + return result + + +def _run_worker(args: argparse.Namespace) -> int: + import torch.distributed as dist + + rank = int(os.environ["RANK"]) + caught: BaseException | None = None + caught_traceback: Any = None + teardown_errors: list[str] = [] + try: + _run_worker_impl(args) + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + finally: + initialized = dist.is_initialized() + if initialized: + if caught is None: + try: + dist.barrier() + except BaseException as error: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"barrier: {type(error).__name__}: {error}" + ) + try: + dist.destroy_process_group() + except BaseException as error: + if caught is None: + caught = error + caught_traceback = error.__traceback__ + teardown_errors.append( + f"destroy_process_group: {type(error).__name__}: {error}" + ) + status = "PASS" if caught is None and not teardown_errors else "ERROR" + if rank == 0: + completion: dict[str, Any] = { + "schema": RESULT_SCHEMA, + "status": status, + "exit_code": 0 if status == "PASS" else 1, + "distributed_teardown_complete": not dist.is_initialized(), + } + if caught is not None: + completion["error"] = { + "type": type(caught).__name__, + "message": str(caught), + } + if teardown_errors: + completion["teardown_errors"] = teardown_errors + print( + "PYTORCH_QWEN3_235B_PPL_COMPLETE " + + json.dumps(completion, ensure_ascii=False, sort_keys=True), + flush=True, + ) + if caught is not None: + raise caught.with_traceback(caught_traceback) + return 0 + + +def main() -> int: + args = _parse_args() + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if "LOCAL_RANK" not in os.environ and world_size == 1: + # Fail on corpus/schema errors before reserving all eight devices. + load_manifest(args.token_manifest) + benchmark_runner._require_idle_gpu() + benchmark_runner._launch_torchrun(args) + raise AssertionError("os.execvpe returned unexpectedly") + return _run_worker(args) + + +if __name__ == "__main__": + raise SystemExit(main()) From cd1d7b723ac8c5c47d60e14ce5ee6055d68705a9 Mon Sep 17 00:00:00 2001 From: wooway777 Date: Fri, 14 Aug 2026 22:30:26 +0800 Subject: [PATCH 2/2] fix(hygon): refine Qwen3-235B inference integration --- csrc/config/model_config.hpp | 8 - csrc/config/quant_config.cpp | 14 +- csrc/config/quant_config.hpp | 6 +- csrc/engine/compiler/paged_compiler.cpp | 9 +- csrc/engine/infer_engine.cpp | 45 +- csrc/engine/rank_worker.cpp | 125 +- csrc/engine/rank_worker.hpp | 15 - csrc/layers/attention/backends/flash_attn.cpp | 122 +- csrc/layers/attention/backends/flash_attn.hpp | 26 +- .../causal_lm_templates/text_causal_lm.hpp | 8 +- .../moe/dispatcher/standard_dispatcher.cpp | 11 +- csrc/layers/moe/experts/fused_moe_experts.cpp | 14 +- csrc/layers/moe/fused_moe.cpp | 6 +- .../moe/runner/cuda_fused_moe_runner.cpp | 50 +- .../quantization/compressed_tensors.cpp | 16 +- .../quantization/compressed_tensors.hpp | 2 +- csrc/models/qwen3/qwen3_attention.cpp | 11 +- csrc/pybind11/engine/engine.hpp | 32 +- .../qwen3moe_w8a8_status_and_vllm_dispatch.md | 147 --- python/infinilm/infer_engine.py | 132 -- test/engine/test_nll_validation.py | 145 --- test/ppl/qwen3_235b/README.md | 200 --- test/ppl/qwen3_235b/scripts/_gpu_guard.py | 100 -- test/ppl/qwen3_235b/scripts/_ppl_common.py | 338 ----- .../calculate_infinilm_precision_ppl.py | 84 -- .../qwen3_235b/scripts/calculate_true_ppl.py | 305 ----- .../infinilm/infinilm_ppl_Qwen3_235B.py | 334 ----- .../scripts/prepare_ppl_corpus_Qwen3_235B.py | 328 ----- .../scripts/transformers/_pytorch_runner.py | 1144 ----------------- .../transformers/pytorch_ppl_Qwen3_235B.py | 491 ------- 30 files changed, 205 insertions(+), 4063 deletions(-) delete mode 100644 docs/qwen3moe_w8a8_status_and_vllm_dispatch.md delete mode 100644 test/engine/test_nll_validation.py delete mode 100644 test/ppl/qwen3_235b/README.md delete mode 100755 test/ppl/qwen3_235b/scripts/_gpu_guard.py delete mode 100755 test/ppl/qwen3_235b/scripts/_ppl_common.py delete mode 100644 test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py delete mode 100755 test/ppl/qwen3_235b/scripts/calculate_true_ppl.py delete mode 100755 test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py delete mode 100755 test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py delete mode 100755 test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py delete mode 100755 test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py diff --git a/csrc/config/model_config.hpp b/csrc/config/model_config.hpp index 5a41a1ac..72339160 100644 --- a/csrc/config/model_config.hpp +++ b/csrc/config/model_config.hpp @@ -88,18 +88,10 @@ class ModelConfig { return quant_config.get_quantization_method(); } - std::string get_moe_weight_method() const { - return quant_config.get_moe_weight_method(); - } - std::string get_moe_weight_method(const infinicore::Device &device) const { return quant_config.get_moe_weight_method(device); } - bool is_moe_w16a16_marlin_enabled() const { - return quant_config.is_moe_w16a16_marlin_enabled(); - } - bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const { return quant_config.is_moe_w16a16_marlin_enabled(device); } diff --git a/csrc/config/quant_config.cpp b/csrc/config/quant_config.cpp index 14cce4af..57ef20dd 100644 --- a/csrc/config/quant_config.cpp +++ b/csrc/config/quant_config.cpp @@ -18,8 +18,7 @@ bool is_w16a16_marlin_method(const std::string &method) { } bool is_w8a8_marlin_method(const std::string &method) { - return method == "slimquant_marlin" || method == "slimquant_compressed_tensors_marlin" || - method == "w8a8_marlin" || method == "hygon_w8a8_marlin"; + return method == "slimquant_marlin" || method == "slimquant_compressed_tensors_marlin" || method == "w8a8_marlin" || method == "hygon_w8a8_marlin"; } bool is_unquantized_config(const nlohmann::json &quantization_config) { @@ -80,8 +79,7 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); } else if (quant_method == "gptq") { return std::make_shared(quantization_config); - } else if (quantization_config["quant_method"] == "w16a16_marlin" || - quantization_config["quant_method"] == "hygon_w16a16_marlin") { + } else if (quantization_config["quant_method"] == "w16a16_marlin" || quantization_config["quant_method"] == "hygon_w16a16_marlin") { return std::make_shared(quantization_config); } else { return std::make_shared(quantization_config); @@ -91,10 +89,6 @@ QuantConfig::get_quantization_method() const { return std::make_shared(quantization_config); // Default case if no matching scheme } -std::string QuantConfig::get_moe_weight_method() const { - return get_moe_weight_method(infinicore::Device(infinicore::Device::Type::CPU, 0)); -} - std::string QuantConfig::get_moe_weight_method(const infinicore::Device &device) const { auto configured_method = explicit_moe_weight_method(quantization_config); if (!configured_method.empty()) { @@ -112,10 +106,6 @@ std::string QuantConfig::get_moe_weight_method(const infinicore::Device &device) return "dense"; } -bool QuantConfig::is_moe_w16a16_marlin_enabled() const { - return is_w16a16_marlin_method(get_moe_weight_method()); -} - bool QuantConfig::is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const { return is_w16a16_marlin_method(get_moe_weight_method(device)); } diff --git a/csrc/config/quant_config.hpp b/csrc/config/quant_config.hpp index a562ea15..2b1d7e5f 100644 --- a/csrc/config/quant_config.hpp +++ b/csrc/config/quant_config.hpp @@ -1,11 +1,11 @@ #pragma once -#include "../utils.hpp" #include "../layers/quantization/quantization.hpp" +#include "../utils.hpp" #include "infinicore/device.hpp" #include "nlohmann/json.hpp" #include -#include #include +#include namespace infinilm::config { @@ -17,9 +17,7 @@ class QuantConfig { QuantConfig(const nlohmann::json &json); std::shared_ptr get_quantization_method() const; - std::string get_moe_weight_method() const; std::string get_moe_weight_method(const infinicore::Device &device) const; - bool is_moe_w16a16_marlin_enabled() const; bool is_moe_w16a16_marlin_enabled(const infinicore::Device &device) const; bool is_moe_w8a8_marlin_enabled(const infinicore::Device &device) const; diff --git a/csrc/engine/compiler/paged_compiler.cpp b/csrc/engine/compiler/paged_compiler.cpp index a7f915d3..06e7d42a 100644 --- a/csrc/engine/compiler/paged_compiler.cpp +++ b/csrc/engine/compiler/paged_compiler.cpp @@ -208,16 +208,9 @@ PagedCompiler::Compiled PagedCompiler::get_compiled(const InfinilmModel::Input & // Reuse the captured tensor only after validating that the runtime // input has the same layout; otherwise fall back to eager mode. const auto &runtime_input_offsets = input.input_offsets.value(); - if (!runtime_input_offsets->is_contiguous() || - runtime_input_offsets->size(0) != batch_size + 1) { + if (!runtime_input_offsets->is_contiguous() || runtime_input_offsets->size(0) != batch_size + 1) { return {nullptr, nullptr}; } - const auto *offsets = reinterpret_cast(runtime_input_offsets->data()); - for (size_t i = 0; i <= batch_size; ++i) { - if (offsets[i] != static_cast(i)) { - return {nullptr, nullptr}; - } - } auto &graph_input = result->second.input; const size_t compiled_block_per_req = graph_input.block_tables.value()->size(1); diff --git a/csrc/engine/infer_engine.cpp b/csrc/engine/infer_engine.cpp index 1fb764f9..e43aa36b 100644 --- a/csrc/engine/infer_engine.cpp +++ b/csrc/engine/infer_engine.cpp @@ -146,44 +146,6 @@ std::vector InferEngine::state_dict_keys() { //------------------------------------------------------ // forward //------------------------------------------------------ -void InferEngine::Input::validate() const { - if (!return_nll) { - if (labels.has_value()) { - throw std::invalid_argument("labels require return_nll=true"); - } - if (score_start != 0) { - throw std::invalid_argument("score_start requires return_nll=true"); - } - return; - } - - if (!input_ids.has_value() || !input_ids.value()) { - throw std::invalid_argument("NLL scoring requires input_ids"); - } - if (!labels.has_value() || !labels.value()) { - throw std::invalid_argument("NLL scoring requires labels"); - } - - const auto &ids = input_ids.value(); - const auto &target = labels.value(); - if (ids->dtype() != infinicore::DataType::I64 - || target->dtype() != infinicore::DataType::I64) { - throw std::invalid_argument("NLL input_ids and labels must use I64 dtype"); - } - if (ids->ndim() != 2 || target->ndim() != 2) { - throw std::invalid_argument("NLL input_ids and labels must be rank-2 tensors"); - } - if (ids->shape() != target->shape()) { - throw std::invalid_argument("NLL input_ids and labels must have identical shapes"); - } - if (ids->size(0) != 1) { - throw std::invalid_argument("NLL scoring currently requires batch_size=1"); - } - if (score_start >= ids->size(1)) { - throw std::invalid_argument("NLL score_start must select at least one token"); - } -} - infinilm::InfinilmModel::Input InferEngine::Input::to_model_input(infinicore::Device device) const { @@ -222,7 +184,8 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { image_req_ids, visual_token_ranges, to_device(target_hidden_states)}; - input.last_token_only = !sample_all_positions && !return_nll; + input.last_token_only = !sample_all_positions; + infinilm::global_state::get_forward_context().attn_metadata = { input.past_sequence_lengths, input.total_sequence_lengths, @@ -244,10 +207,6 @@ InferEngine::Input::to_model_input(infinicore::Device device) const { } InferEngine::Output InferEngine::forward(const InferEngine::Input &input) { - // Validate before dispatch so malformed NLL requests cannot fail only one - // rank and leave the remaining workers waiting at a collective. - input.validate(); - // Trigger each worker to run inference for (auto &worker : workers_) { worker->run(input); diff --git a/csrc/engine/rank_worker.cpp b/csrc/engine/rank_worker.cpp index 276fc6d0..acb2d24c 100644 --- a/csrc/engine/rank_worker.cpp +++ b/csrc/engine/rank_worker.cpp @@ -418,13 +418,9 @@ void RankWorker::thread_loop() { infinicore::Tensor logits; infinicore::Tensor hidden_states; - // Full-position and NLL runs need eager mode because generation - // graphs return last-token logits and omit hidden states. PP graph - // compilation is not supported yet. - if (!local_args.sample_all_positions - && !local_args.return_nll - && compiler_ != nullptr - && rank_info_.pp_size == 1) { + // All-position speculative/MTP runs need eager mode because + // hidden states are not part of compiled graph outputs. + if (!local_args.sample_all_positions && compiler_ != nullptr && rank_info_.pp_size == 1) { auto [graph, output] = compiler_->get_compiled(local_args.to_model_input(infinicore::Device::cpu())); if (graph != nullptr && output != nullptr) { graph->run(); @@ -439,11 +435,6 @@ void RankWorker::thread_loop() { hidden_states = model_output.hidden_states; } - if (local_args.return_nll && rank_info_.pp_size > 1) { - throw std::runtime_error( - "NLL scoring with pipeline parallelism is not supported"); - } - if (rank_info_.pp_size > 1 && rank_info_.pp_stage + 1 != rank_info_.pp_size) { infinicore::Tensor output_ids; if (rank_info_.pp_stage == 0 && rank_info_.tp_rank == 0) { @@ -473,96 +464,54 @@ void RankWorker::thread_loop() { continue; } - // Sampling and scoring both consume replicated full-vocabulary - // logits, so only rank 0 needs to materialize the result. + // Random sampling (rank 0 only) if (rank_info_.tp_rank == 0) { + auto temperature{local_args.temperature}; + auto top_p{local_args.top_p}; + auto top_k{local_args.top_k}; + const auto &logits_shape{logits->shape()}; - if (logits_shape.size() != 3) { - throw std::runtime_error("InferEngine expected rank-3 logits"); - } const auto &vocab_size{logits_shape[2]}; const auto &total_len{logits_shape[1]}; const auto &batch_size{logits_shape[0]}; - if (local_args.return_nll) { - auto labels = local_args.labels.value()->to(rank_info_.device); - if (labels->dtype() != infinicore::DataType::I64 - || labels->ndim() != 2 - || labels->size(0) != batch_size - || labels->size(1) != total_len) { - throw std::runtime_error( - "NLL labels must be I64 with shape [batch, sequence]"); - } + auto n_req = local_args.input_offsets.value()->size(0) - 1; + int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); - const auto score_len = total_len - local_args.score_start; - auto score_logits = logits->narrow( - {{1, local_args.score_start, score_len}}); - auto score_labels = labels->narrow( - {{1, local_args.score_start, score_len}}); - - auto token_nll = infinicore::Tensor::empty( - score_labels->shape(), - infinicore::DataType::F32, - rank_info_.device); - infinicore::op::cross_entropy_( - token_nll, score_logits, score_labels); - token_nll = token_nll->to(infinicore::Device::cpu()); - infinicore::context::syncStream(); - output_ = Output{ - infinicore::Tensor{}, - infinicore::Tensor{}, - infinicore::Tensor{}, - token_nll, - score_len, - }; - } else { - auto temperature{local_args.temperature}; - auto top_p{local_args.top_p}; - auto top_k{local_args.top_k}; - auto n_req = local_args.input_offsets.value()->size(0) - 1; - int32_t *input_offsets = (int32_t *)local_args.input_offsets.value()->data(); - - const bool sample_all_positions = local_args.sample_all_positions; - const size_t n_out = sample_all_positions - ? static_cast(input_offsets[n_req]) - : n_req; - auto output_ids{infinicore::Tensor::empty( - {n_out}, infinicore::DataType::I64, rank_info_.device)}; - - for (size_t i{0}; i < n_out; ++i) { - size_t score_idx = i; - if (!sample_all_positions) { - score_idx = total_len == n_req - ? i - : static_cast(input_offsets[i + 1] - 1); - } - auto score{logits->view({batch_size * total_len, vocab_size}) - ->narrow({{0, score_idx, 1}}) - ->view({vocab_size})}; - auto out{output_ids->narrow({{0, i, 1}})->view({})}; - float random_val = std::uniform_real_distribution(0, 1)(rng_); - infinicore::op::random_sample_( - out, score, random_val, top_p, top_k, temperature); - } + const bool sample_all_positions = local_args.sample_all_positions; + const size_t n_out = sample_all_positions ? static_cast(input_offsets[n_req]) : n_req; + auto output_ids{infinicore::Tensor::empty({n_out}, infinicore::DataType::I64, rank_info_.device)}; - if (rank_info_.pp_size > 1) { - infinicore::op::distributed::send( - output_ids, - 0, - rank_info_.world_comm); + for (size_t i{0}; i < n_out; ++i) { + size_t score_idx = i; + if (!sample_all_positions) { + score_idx = total_len == n_req + ? i + : static_cast(input_offsets[i + 1] - 1); } + auto score{logits->view({batch_size * total_len, vocab_size})->narrow({{0, score_idx, 1}})->view({vocab_size})}; + auto out{output_ids->narrow({{0, i, 1}})->view({})}; + float random_val = std::uniform_real_distribution(0, 1)(rng_); + infinicore::op::random_sample_( + out, score, random_val, top_p, top_k, temperature); + } - // Tensor::to(CPU) uses the synchronous D2H contract. - output_ids = output_ids->to(infinicore::Device::cpu()); - output_ = Output{ + if (rank_info_.pp_size > 1) { + infinicore::op::distributed::send( output_ids, - logits, - hidden_states, - infinicore::Tensor{}, 0, - }; + rank_info_.world_comm); } + + output_ids = output_ids->to(infinicore::Device::cpu()); + + infinicore::context::syncStream(); + + auto out{Output{output_ids, logits, hidden_states}}; + + output_ = std::move(out); } + job_done_ = true; } cv_.notify_all(); diff --git a/csrc/engine/rank_worker.hpp b/csrc/engine/rank_worker.hpp index 19f2f1f2..d396ef6f 100644 --- a/csrc/engine/rank_worker.hpp +++ b/csrc/engine/rank_worker.hpp @@ -73,25 +73,12 @@ class RankWorker { /// Sample logits at every packed input position instead of one token per request. bool sample_all_positions{false}; - /// Shifted causal-LM labels. Present only for explicit NLL scoring. - std::optional labels; - - /// First logits/label position included in NLL scoring. - size_t score_start{0}; - - /// Compute token NLL instead of sampling output IDs. - bool return_nll{false}; - float temperature{1}; int top_k{50}; float top_p{1}; - /// Validate invariants shared by Python and native callers before a - /// request is dispatched to any rank worker. - void validate() const; - infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const; }; @@ -99,8 +86,6 @@ class RankWorker { infinicore::Tensor output_ids; infinicore::Tensor logits; infinicore::Tensor hidden_states; - infinicore::Tensor nll; - size_t scored_tokens{0}; }; RankWorker(std::shared_ptr infinilm_config, diff --git a/csrc/layers/attention/backends/flash_attn.cpp b/csrc/layers/attention/backends/flash_attn.cpp index eb2e3797..d593128e 100644 --- a/csrc/layers/attention/backends/flash_attn.cpp +++ b/csrc/layers/attention/backends/flash_attn.cpp @@ -2,6 +2,8 @@ #include "../../../utils.hpp" #include "infinicore/ops.hpp" +#include "infinicore/ops/mha_kvcache.hpp" +#include "infinicore/ops/mha_varlen.hpp" namespace infinilm::layers::attention::backends { @@ -11,39 +13,107 @@ FlashAttentionImpl::FlashAttentionImpl(size_t num_heads, size_t num_kv_heads, size_t layer_idx) : num_heads_(num_heads), + head_size_(head_size), scale_(scale), num_kv_heads_(num_kv_heads), + layer_idx_(layer_idx), head_dim_(head_size) { - (void)layer_idx; + + const infinilm::global_state::InfinilmConfig &infinilm_config = infinilm::global_state::get_infinilm_config(); + if (!infinilm_config.model_config) { + throw std::runtime_error("infinilm::layers::attention::backends::FlashAttentionImpl: model_config is null"); + } + max_position_embeddings_ = infinilm_config.model_config->get("max_position_embeddings"); } -infinicore::Tensor FlashAttentionImpl::forward( - const AttentionLayer &layer, - const infinicore::Tensor &query, - const infinicore::Tensor &key, - const infinicore::Tensor &value, - infinicore::Tensor &kv_cache, - const infinilm::global_state::AttentionMetadata &attn_metadata) const { - (void)layer; - - ASSERT(attn_metadata.total_sequence_lengths.has_value()); - ASSERT(attn_metadata.block_tables.has_value()); - ASSERT(attn_metadata.slot_mapping.has_value()); - - return infinicore::op::paged_flash_attention( - query, +infinicore::Tensor FlashAttentionImpl::forward(const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const { + auto total_sequence_lengths = attn_metadata.total_sequence_lengths; + auto input_offsets = attn_metadata.input_offsets; + auto block_tables = attn_metadata.block_tables; + auto slot_mapping = attn_metadata.slot_mapping; + auto cu_seqlens = attn_metadata.cu_seqlens; + + ASSERT(block_tables.has_value()); + ASSERT(slot_mapping.has_value()); + + if (query->device().getType() == infinicore::Device::Type::HYGON) { + return infinicore::op::paged_flash_attention( + query, + key, + value, + kv_cache, + total_sequence_lengths.value(), + input_offsets, + cu_seqlens, + block_tables.value(), + slot_mapping.value(), + num_heads_, + num_kv_heads_, + head_dim_, + scale_); + } + + // 1. update paged kv cache + auto [k_total, v_total] = do_kv_cache_update(layer, key, value, kv_cache, slot_mapping.value()); + + size_t seq_len = query->shape()[0]; + bool is_prefill = (seq_len != total_sequence_lengths.value()->shape()[0]); + + // 2. Compute attention + infinicore::Tensor attn_output = infinicore::Tensor::empty({seq_len, num_heads_, head_dim_}, query->dtype(), query->device()); + if (is_prefill) { + infinicore::op::mha_varlen_( + attn_output, + query, + k_total, + v_total, + input_offsets.value(), + cu_seqlens.value(), + block_tables.value(), + max_position_embeddings_, + max_position_embeddings_, + std::nullopt, + scale_); + } else { + // FA2 decode path: flash::mha_fwd_kvcache + // In paged-attn mode, seq_len = actual batch_size (one query token per sequence). + // q_reshaped: [seq_len, num_heads, head_dim] → [seq_len, 1, num_heads, head_dim] + // k/v cache: [num_blocks, block_size, num_kv_heads, head_dim] + auto q_for_fa = query->view({seq_len, 1, num_heads_, head_dim_}); + auto attn_out_4d = infinicore::op::mha_kvcache( + q_for_fa, + k_total, // [num_blocks, block_size, num_kv_heads, head_dim] + v_total, + total_sequence_lengths.value(), // [seq_len] int32 (one entry per sequence) + block_tables.value(), // [seq_len, max_num_blocks_per_seq] int32 + std::nullopt, + scale_); + attn_output = attn_out_4d->view({seq_len, num_heads_, head_dim_}); + } + attn_output = attn_output->view({1, seq_len, num_heads_ * head_dim_}); + return attn_output; +} + +std::tuple FlashAttentionImpl::do_kv_cache_update(const AttentionLayer &layer, + const infinicore::Tensor key, + const infinicore::Tensor value, + infinicore::Tensor &kv_cache, + const infinicore::Tensor slot_mapping) const { + auto k_cache_layer = kv_cache->narrow({{0, 0, 1}})->squeeze(0); + auto v_cache_layer = kv_cache->narrow({{0, 1, 1}})->squeeze(0); + infinicore::op::paged_caching_( + k_cache_layer->permute({0, 2, 1, 3}), // permute to BHSD for paged_caching_ + v_cache_layer->permute({0, 2, 1, 3}), key, value, - kv_cache, - attn_metadata.total_sequence_lengths.value(), - attn_metadata.input_offsets, - attn_metadata.cu_seqlens, - attn_metadata.block_tables.value(), - attn_metadata.slot_mapping.value(), - num_heads_, - num_kv_heads_, - head_dim_, - scale_); + slot_mapping); + + return {k_cache_layer, v_cache_layer}; } } // namespace infinilm::layers::attention::backends diff --git a/csrc/layers/attention/backends/flash_attn.hpp b/csrc/layers/attention/backends/flash_attn.hpp index a592dddb..93f61e8b 100644 --- a/csrc/layers/attention/backends/flash_attn.hpp +++ b/csrc/layers/attention/backends/flash_attn.hpp @@ -2,6 +2,7 @@ #include "../../../global_state/global_state.hpp" #include "infinicore/tensor.hpp" +#include namespace infinilm::layers::attention { class AttentionLayer; @@ -28,19 +29,26 @@ class FlashAttentionImpl { * @param attn_metadata: Attention metadata. * @return Attention output, shape `[1, num_tokens, num_heads * head_dim]`. */ - infinicore::Tensor forward( - const AttentionLayer &layer, - const infinicore::Tensor &query, - const infinicore::Tensor &key, - const infinicore::Tensor &value, - infinicore::Tensor &kv_cache, - const infinilm::global_state::AttentionMetadata &attn_metadata) const; + infinicore::Tensor forward(const AttentionLayer &layer, + const infinicore::Tensor &query, + const infinicore::Tensor &key, + const infinicore::Tensor &value, + infinicore::Tensor &kv_cache, + const infinilm::global_state::AttentionMetadata &attn_metadata) const; + + std::tuple do_kv_cache_update(const AttentionLayer &layer, + const infinicore::Tensor key, + const infinicore::Tensor value, + infinicore::Tensor &kv_cache, + const infinicore::Tensor slot_mapping) const; private: size_t num_heads_; + size_t head_size_; float scale_; size_t num_kv_heads_; - size_t head_dim_; + size_t layer_idx_; + size_t head_dim_; // Note: head_dim equals to head_size + size_t max_position_embeddings_; }; - } // namespace infinilm::layers::attention::backends diff --git a/csrc/layers/causal_lm_templates/text_causal_lm.hpp b/csrc/layers/causal_lm_templates/text_causal_lm.hpp index 135a165a..5dfbfa4a 100644 --- a/csrc/layers/causal_lm_templates/text_causal_lm.hpp +++ b/csrc/layers/causal_lm_templates/text_causal_lm.hpp @@ -9,7 +9,6 @@ #include - namespace infinilm::layers::causal_lm_templates { /** @@ -46,8 +45,8 @@ class TextCausalLM : public InfinilmModel { tp_size_ = static_cast(rank_info.tp_size); tp_rank_ = static_cast(rank_info.tp_rank); vocab_parallel_ = device.getType() == infinicore::Device::Type::HYGON - && tp_size_ > 1 - && vocab_size % tp_size_ == 0; + && tp_size_ > 1 + && vocab_size % tp_size_ == 0; model_ = this->register_module("model", model_config, device); if (is_last_pp_stage()) { @@ -113,8 +112,7 @@ class TextCausalLM : public InfinilmModel { const size_t local_vocab_size = local_shape.back(); const size_t num_rows = local_logits->numel() / local_vocab_size; auto local_flat = local_logits->view({num_rows, local_vocab_size}); - const auto &rank_info = - infinilm::global_state::get_tensor_model_parallel_rank_info(); + const auto &rank_info = infinilm::global_state::get_tensor_model_parallel_rank_info(); auto gathered = infinicore::op::distributed::allgather( local_flat, tp_size_, rank_info.comm); diff --git a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp index 96bc6205..b0a4d015 100644 --- a/csrc/layers/moe/dispatcher/standard_dispatcher.cpp +++ b/csrc/layers/moe/dispatcher/standard_dispatcher.cpp @@ -30,8 +30,15 @@ infinicore::Tensor StandardDispatcher::combine(const CombineInput &combine_input MoeWorkspace &workspace) const { (void)workspace; if (tp_size_ > 1 && communicator_ != nullptr) { - return infinicore::op::distributed::allreduce( - combine_input.hidden_states, INFINICCL_SUM, communicator_); + if (combine_input.hidden_states->device().getType() == infinicore::Device::Type::HYGON) { + return infinicore::op::distributed::allreduce( + combine_input.hidden_states, INFINICCL_SUM, communicator_); + } + infinicore::op::distributed::allreduce_( + combine_input.hidden_states, + combine_input.hidden_states, + INFINICCL_SUM, + communicator_); } return combine_input.hidden_states; } diff --git a/csrc/layers/moe/experts/fused_moe_experts.cpp b/csrc/layers/moe/experts/fused_moe_experts.cpp index 40fc8ec2..3a679386 100644 --- a/csrc/layers/moe/experts/fused_moe_experts.cpp +++ b/csrc/layers/moe/experts/fused_moe_experts.cpp @@ -26,8 +26,7 @@ FusedMoeExperts::FusedMoeExperts(std::shared_ptr if (enable_hygon_w16a16_marlin_ && enable_hygon_w8a8_marlin_) { throw std::runtime_error("Only one Hygon MoE Marlin weight method can be enabled"); } - if (moe_weight_method != "dense" && - !enable_hygon_w16a16_marlin_ && !enable_hygon_w8a8_marlin_) { + if (moe_weight_method != "dense" && !enable_hygon_w16a16_marlin_ && !enable_hygon_w8a8_marlin_) { throw std::runtime_error("Unsupported MoE weight method: " + moe_weight_method); } ASSERT(num_experts_ > 0); @@ -136,10 +135,7 @@ void FusedMoeExperts::process_weights_after_loading() { if (!w13_weight_ || !w2_weight_ || !w13_weight_scale_ || !w2_weight_scale_) { throw std::runtime_error("slimquant_marlin MoE weight method requires loaded int8 w13/w2 weights and scales"); } - if (w13_weight_->dtype() != infinicore::DataType::I8 || - w2_weight_->dtype() != infinicore::DataType::I8 || - w13_weight_scale_->dtype() != infinicore::DataType::F32 || - w2_weight_scale_->dtype() != infinicore::DataType::F32) { + if (w13_weight_->dtype() != infinicore::DataType::I8 || w2_weight_->dtype() != infinicore::DataType::I8 || w13_weight_scale_->dtype() != infinicore::DataType::F32 || w2_weight_scale_->dtype() != infinicore::DataType::F32) { throw std::runtime_error("slimquant_marlin MoE weight method requires int8 weights and fp32 weight scales"); } if (hidden_size_ % 64 != 0 || intermediate_size_per_partition_ % 64 != 0) { @@ -186,12 +182,10 @@ void FusedMoeExperts::process_weights_after_loading() { if (!w13_weight_ || !w2_weight_) { throw std::runtime_error("w16a16_marlin MoE weight method requires loaded dense w13/w2 weights"); } - if (w13_weight_->dtype() != infinicore::DataType::F16 && - w13_weight_->dtype() != infinicore::DataType::BF16) { + if (w13_weight_->dtype() != infinicore::DataType::F16 && w13_weight_->dtype() != infinicore::DataType::BF16) { throw std::runtime_error("w16a16_marlin MoE weight method requires FP16 or BF16 weights"); } - if (hidden_size_ % 32 != 0 || intermediate_size_per_partition_ % 16 != 0 || - (intermediate_size_per_partition_ * 2) % 32 != 0) { + if (hidden_size_ % 32 != 0 || intermediate_size_per_partition_ % 16 != 0 || (intermediate_size_per_partition_ * 2) % 32 != 0) { throw std::runtime_error("w16a16_marlin MoE weight method requires aligned hidden/intermediate sizes"); } diff --git a/csrc/layers/moe/fused_moe.cpp b/csrc/layers/moe/fused_moe.cpp index 1e607207..8defcbc8 100644 --- a/csrc/layers/moe/fused_moe.cpp +++ b/csrc/layers/moe/fused_moe.cpp @@ -16,11 +16,7 @@ std::shared_ptr make_workspace( const EPConfig &ep_config, const std::shared_ptr &model_config, const infinicore::Device &device) { - const bool use_hygon_marlin = - device.getType() == infinicore::Device::Type::HYGON && - ep_config.backend == EPBackend::Disabled && - (model_config->is_moe_w8a8_marlin_enabled(device) || - model_config->is_moe_w16a16_marlin_enabled(device)); + const bool use_hygon_marlin = device.getType() == infinicore::Device::Type::HYGON && ep_config.backend == EPBackend::Disabled && (model_config->is_moe_w8a8_marlin_enabled(device) || model_config->is_moe_w16a16_marlin_enabled(device)); if (!use_hygon_marlin) { return std::make_shared(); } diff --git a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp index 629f9345..e0561b11 100644 --- a/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp +++ b/csrc/layers/moe/runner/cuda_fused_moe_runner.cpp @@ -98,8 +98,8 @@ CombineInput CudaFusedMoeRunner::run( if (weights.is_hygon_w16a16_marlin() || weights.is_hygon_w8a8_marlin()) { const auto format = weights.is_hygon_w16a16_marlin() - ? infinicore::op::HygonMoeMarlinWeightFormat::W16A16 - : infinicore::op::HygonMoeMarlinWeightFormat::W8A8; + ? infinicore::op::HygonMoeMarlinWeightFormat::W16A16 + : infinicore::op::HygonMoeMarlinWeightFormat::W8A8; const infinicore::op::HygonMoeMarlinWeights marlin_weights{ weights.packed_w13, weights.packed_w2, @@ -121,11 +121,9 @@ CombineInput CudaFusedMoeRunner::run( MoeRoutingMetadata routing_metadata; if (output.has_routing_metadata) { - routing_metadata.sorted_token_ids = - output.sorted_token_ids; + routing_metadata.sorted_token_ids = output.sorted_token_ids; routing_metadata.expert_ids = output.expert_ids; - routing_metadata.num_tokens_post_padded = - output.num_tokens_post_padded; + routing_metadata.num_tokens_post_padded = output.num_tokens_post_padded; } return CombineInput{ CombineInputFormat::Standard, @@ -139,8 +137,7 @@ CombineInput CudaFusedMoeRunner::run( dispatch_output, workspace, align_block_size_); - auto runner_output = - run_fused_core(runner_input, weights, workspace); + auto runner_output = run_fused_core(runner_input, weights, workspace); return CombineInput{ CombineInputFormat::Standard, runner_output.hidden_states, @@ -153,8 +150,7 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input( const DispatchOutput &dispatch_output, MoeWorkspace &workspace, size_t block_size) const { - const auto &topk_ids = - dispatch_output.topk_output.topk_ids; + const auto &topk_ids = dispatch_output.topk_output.topk_ids; const auto &topk_shape = topk_ids->shape(); if (topk_shape.size() != 2) { throw std::runtime_error( @@ -162,20 +158,17 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input( } const size_t num_pairs = topk_shape[0] * topk_shape[1]; const size_t align_num_experts = num_local_experts_ + 1; - const size_t max_num_tokens_padded = - num_pairs < align_num_experts - ? num_pairs * block_size - : num_pairs - + align_num_experts * (block_size - 1); - const size_t sorted_token_ids_capacity = - ((max_num_tokens_padded + 3) / 4) * 4; - const size_t max_num_blocks = - (max_num_tokens_padded + block_size - 1) / block_size; + const size_t max_num_tokens_padded = num_pairs < align_num_experts + ? num_pairs * block_size + : num_pairs + + align_num_experts * (block_size - 1); + const size_t sorted_token_ids_capacity = ((max_num_tokens_padded + 3) / 4) * 4; + const size_t max_num_blocks = (max_num_tokens_padded + block_size - 1) / block_size; const auto device = topk_ids->device(); if (!workspace.sorted_token_ids || workspace.sorted_token_ids_capacity - < sorted_token_ids_capacity) { + < sorted_token_ids_capacity) { if (infinicore::context::isGraphRecording()) { throw std::runtime_error( "MoE sorted_token_ids workspace was not initialized before graph capture"); @@ -184,8 +177,7 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input( {sorted_token_ids_capacity}, infinicore::DataType::I32, device); - workspace.sorted_token_ids_capacity = - sorted_token_ids_capacity; + workspace.sorted_token_ids_capacity = sorted_token_ids_capacity; } if (!workspace.expert_ids || workspace.expert_ids_capacity < max_num_blocks) { @@ -204,16 +196,14 @@ CudaFusedMoeRunnerInput CudaFusedMoeRunner::prepare_runner_input( throw std::runtime_error( "MoE num_tokens_post_padded workspace was not initialized before graph capture"); } - workspace.num_tokens_post_padded = - infinicore::Tensor::empty( - {1}, - infinicore::DataType::I32, - device); + workspace.num_tokens_post_padded = infinicore::Tensor::empty( + {1}, + infinicore::DataType::I32, + device); } - auto sorted_token_ids = - workspace.sorted_token_ids->narrow( - {{0, 0, sorted_token_ids_capacity}}); + auto sorted_token_ids = workspace.sorted_token_ids->narrow( + {{0, 0, sorted_token_ids_capacity}}); auto expert_ids = workspace.expert_ids->narrow( {{0, 0, max_num_blocks}}); if (dispatch_output.expert_map) { diff --git a/csrc/layers/quantization/compressed_tensors.cpp b/csrc/layers/quantization/compressed_tensors.cpp index 72659554..9428875e 100644 --- a/csrc/layers/quantization/compressed_tensors.cpp +++ b/csrc/layers/quantization/compressed_tensors.cpp @@ -53,23 +53,13 @@ bool has_linear_or_moe_target(const nlohmann::json &group) { bool is_dynamic_token_w8a8_group(const nlohmann::json &group) { auto weights_it = group.find("weights"); auto input_it = group.find("input_activations"); - if (weights_it == group.end() || input_it == group.end() || - !weights_it->is_object() || !input_it->is_object()) { + if (weights_it == group.end() || input_it == group.end() || !weights_it->is_object() || !input_it->is_object()) { return false; } const auto &weights = *weights_it; const auto &input = *input_it; - const bool weight_ok = - string_field_equals(weights, "type", "int") && - string_field_equals(weights, "strategy", "channel") && - integer_field_equals(weights, "num_bits", 8) && - bool_field_equals(weights, "symmetric", true); - const bool input_ok = - string_field_equals(input, "type", "int") && - string_field_equals(input, "strategy", "token") && - integer_field_equals(input, "num_bits", 8) && - bool_field_equals(input, "dynamic", true) && - bool_field_equals(input, "symmetric", true); + const bool weight_ok = string_field_equals(weights, "type", "int") && string_field_equals(weights, "strategy", "channel") && integer_field_equals(weights, "num_bits", 8) && bool_field_equals(weights, "symmetric", true); + const bool input_ok = string_field_equals(input, "type", "int") && string_field_equals(input, "strategy", "token") && integer_field_equals(input, "num_bits", 8) && bool_field_equals(input, "dynamic", true) && bool_field_equals(input, "symmetric", true); return weight_ok && input_ok; } diff --git a/csrc/layers/quantization/compressed_tensors.hpp b/csrc/layers/quantization/compressed_tensors.hpp index 2a088728..15438da9 100644 --- a/csrc/layers/quantization/compressed_tensors.hpp +++ b/csrc/layers/quantization/compressed_tensors.hpp @@ -6,7 +6,7 @@ namespace infinilm::quantization { class CompressedTensors : public BaseQuantization { public: explicit CompressedTensors(const nlohmann::json &quant_config) - : BaseQuantization(quant_config) {}; + : BaseQuantization(quant_config){}; QuantScheme get_quant_scheme() const override { return QuantScheme::COMPRESSED_TENSOR_W8A8I8; diff --git a/csrc/models/qwen3/qwen3_attention.cpp b/csrc/models/qwen3/qwen3_attention.cpp index 57b84056..c7b794ae 100644 --- a/csrc/models/qwen3/qwen3_attention.cpp +++ b/csrc/models/qwen3/qwen3_attention.cpp @@ -136,12 +136,11 @@ infinicore::Tensor Qwen3Attention::forward_paged_(const infinicore::Tensor &posi } // 4. Apply Q/K RMSNorm and RoPE. - const bool can_use_hygon_fused_rms_rope = - qkv_proj_->get_quantization()->get_quant_scheme() == infinilm::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8 - && q_reshaped->device().getType() == infinicore::Device::Type::HYGON - && rotary_emb_->rotary_dim() == head_dim_ - && !rotary_emb_->mrope_section().has_value() - && infinicore::op::rms_rotary_embedding_fuse_available(q_reshaped->device()); + const bool can_use_hygon_fused_rms_rope = qkv_proj_->get_quantization()->get_quant_scheme() == infinilm::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8 + && q_reshaped->device().getType() == infinicore::Device::Type::HYGON + && rotary_emb_->rotary_dim() == head_dim_ + && !rotary_emb_->mrope_section().has_value() + && infinicore::op::rms_rotary_embedding_fuse_available(q_reshaped->device()); if (can_use_hygon_fused_rms_rope) { q_reshaped = q_reshaped->contiguous(); k_reshaped = k_reshaped->contiguous(); diff --git a/csrc/pybind11/engine/engine.hpp b/csrc/pybind11/engine/engine.hpp index cc9ecaee..c5e85577 100644 --- a/csrc/pybind11/engine/engine.hpp +++ b/csrc/pybind11/engine/engine.hpp @@ -161,7 +161,6 @@ inline void bind_infer_engine(py::module &m) { std::optional> visual_token_ranges, std::optional target_hidden_states, bool sample_all_positions, - std::optional labels, py::kwargs kwargs) { InferEngine::Input input{ std::move(input_ids), @@ -182,7 +181,6 @@ inline void bind_infer_engine(py::module &m) { std::move(visual_token_ranges), std::move(target_hidden_states), sample_all_positions, - std::move(labels), }; // Explicit defaults @@ -195,8 +193,6 @@ inline void bind_infer_engine(py::module &m) { "temperature", "top_p", "top_k", - "score_start", - "return_nll", }; for (auto &item : kwargs) { @@ -213,24 +209,9 @@ inline void bind_infer_engine(py::module &m) { input.top_p = py::cast(item.second); } else if (key == "top_k") { input.top_k = py::cast(item.second); - } else if (key == "score_start") { - if (py::isinstance(item.second)) { - throw py::type_error("score_start must be an integer, not bool"); - } - const auto score_start = py::cast(item.second); - if (score_start < 0) { - throw py::value_error("score_start must be non-negative"); - } - input.score_start = static_cast(score_start); - } else if (key == "return_nll") { - if (!py::isinstance(item.second)) { - throw py::type_error("return_nll must be a bool"); - } - input.return_nll = py::cast(item.second); } } - input.validate(); return input; }), py::arg("input_ids") = std::nullopt, @@ -250,8 +231,7 @@ inline void bind_infer_engine(py::module &m) { py::arg("image_req_ids") = std::nullopt, py::arg("visual_token_ranges") = std::nullopt, py::arg("target_hidden_states") = std::nullopt, - py::arg("sample_all_positions") = false, - py::arg("labels") = std::nullopt) + py::arg("sample_all_positions") = false) .def_readwrite("input_ids", &InferEngine::Input::input_ids) .def_readwrite("position_ids", &InferEngine::Input::position_ids) .def_readwrite("past_sequence_lengths", &InferEngine::Input::past_sequence_lengths) @@ -270,9 +250,6 @@ inline void bind_infer_engine(py::module &m) { .def_readwrite("visual_token_ranges", &InferEngine::Input::visual_token_ranges) .def_readwrite("target_hidden_states", &InferEngine::Input::target_hidden_states) .def_readwrite("sample_all_positions", &InferEngine::Input::sample_all_positions) - .def_readwrite("labels", &InferEngine::Input::labels) - .def_readwrite("score_start", &InferEngine::Input::score_start) - .def_readwrite("return_nll", &InferEngine::Input::return_nll) .def_readwrite("temperature", &InferEngine::Input::temperature) .def_readwrite("top_k", &InferEngine::Input::top_k) .def_readwrite("top_p", &InferEngine::Input::top_p); @@ -280,12 +257,7 @@ inline void bind_infer_engine(py::module &m) { py::class_(infer_engine, "Output") .def_readwrite("output_ids", &InferEngine::Output::output_ids, "Sampled token IDs") .def_readwrite("logits", &InferEngine::Output::logits, "Raw logits tensor") - .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor") - .def_readwrite("nll", &InferEngine::Output::nll, "Per-token NLL tensor") - .def_readwrite( - "scored_tokens", - &InferEngine::Output::scored_tokens, - "Number of scored tokens"); + .def_readwrite("hidden_states", &InferEngine::Output::hidden_states, "Raw hidden states tensor"); } } // namespace infinilm::engine diff --git a/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md b/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md deleted file mode 100644 index e8256838..00000000 --- a/docs/qwen3moe_w8a8_status_and_vllm_dispatch.md +++ /dev/null @@ -1,147 +0,0 @@ -# Qwen3-MoE W8A8 InfiniLM/vLLM Status - -Date: 2026-07-10 - -## Remote Environment - -- Host: `qinyiqun@10.211.3.28` -- SSH key: `C:\Users\qinyi\.ssh\bw1000` -- Container: `qinyiqun` -- InfiniCore: `/home/qinyiqun/InfiniCore` -- InfiniLM: `/home/qinyiqun/InfiniLM` -- FP model: `/home_aclsylqidf/shared/Qwen3-30B-A3B` -- W8A8 model: `/home_aclsylqidf/shared/Qwen3-30B-A3B-Channel-INT8-w8a8` - -Runtime setup inside the container: - -```bash -unset http_proxy https_proxy all_proxy HTTP_PROXY HTTPS_PROXY ALL_PROXY -export PATH=/root/.local/bin:/opt/dtk/cuda/cuda/bin:$PATH -export XMAKE_ROOT=y -export LD_LIBRARY_PATH=/usr/local/lib/python3.10/dist-packages/torch/lib:/root/.infini/lib:${LD_LIBRARY_PATH:-} -export PYTHONPATH=/usr/local/:${PYTHONPATH:-} -``` - -InfiniCore configure must include `--graph=y`: - -```bash -xmake f --hygon-dcu=true --aten=true --flash-attn=/usr/local/lib/python3.10/dist-packages/ --cuda=/opt/dtk/cuda/cuda --ccl=true --graph=y -cv -y -xmake build -xmake install -xmake build _infinicore -xmake install _infinicore -pip install -e . -``` - -## Current Benchmark Contract - -- Model family: `Qwen3-30B-A3B` -- Target path: W8A8 quantized model -- Parallelism: `TP=2`, `DP=1`, `EP=1` -- MoE communication: TP only, no DeepEP/allgather EP path -- Benchmark length going forward: `input_len=4096`, `output_len=1280` -- Important guardrail: pass only one `input_len` value. A comma-separated input length list can hang the current benchmark path. -- Device/profiling tools: `hy-smi` and Hygon trace. - -## Current InfiniLM Findings - -The long-run stall was isolated to the W8A8 MoE path with long prefill. FP graph runs and short W8A8 decode runs can complete, so the issue is not simply long output length. - -Observed behavior before long-prefill slicing: - -- W8A8 `4096/128` graph timed out. -- W8A8 `4096/128` no-graph segfaulted. -- Backtraces showed one rank waiting in `RankWorker::wait`, while the other rank was inside the W8A8 Marlin MoE path and teardown/exit handling. - -The Hygon W8A8 Marlin MoE path now chunks long prefill internally: - -- Files: `csrc/layers/moe/runner/cuda_fused_moe_runner.cpp`, `.hpp` -- Fixed chunk size: `16384` tokens, matching vLLM's production chunk size -- The sliced path is selected before full-input routing metadata is prepared, so each token is aligned only once -- No W8A8 slice or debug environment switches are required - -This keeps long-prefill workspace bounded while decode continues to use the graph-captured Marlin path directly. - -## vLLM W8A8 MoE Path - -vLLM package path: - -- `/usr/local/lib/python3.10/dist-packages/vllm` -- Runtime version in logs: `v0.15.1` - -Important vLLM env: - -- `VLLM_FUSED_MOE_CHUNK_SIZE=16384` -- `VLLM_W8A8_BACKEND=3` - -Main call chain: - -1. `CompressedTensorsW8A8Int8MoEMethod.apply()` -2. `fused_experts(...)` -3. `lmslim.layers.fused_moe.fuse_moe_int8.fused_experts_impl_int8` - -vLLM does not repack Qwen3 MoE weights into the InfiniLM Marlin layout. It keeps ordinary channel-wise int8 tensors: - -- `w1`: `[E, 768, 2048]` -- `w2`: `[E, 2048, 384]` -- `w1_scale`: `[E, 768, 1]` -- `w2_scale`: `[E, 2048, 1]` -- `E=128`, `top_k=8` - -Operator sequence per chunk: - -1. Per-token quantize hidden states. -2. Align/count/sort tokens by expert. -3. GEMM1: `lightop.moe_gemm_w8a8(...)` -4. Activation and quantize: `fuse_silu_mul_quant(...)` -5. GEMM2: `lightop.moe_gemm_w8a8(...)` -6. Reduce top-k outputs: `moe_sum` / `moe_reduce_dispatch` - -## Representative vLLM Size Dispatch - -These are the useful anchor cases for InfiniLM implementation. We do not need to reproduce every tiny graph-capture size immediately. - -| Effective M | GEMM1 shape | GEMM1 config/kernel | GEMM2 shape | GEMM2 config/kernel | Notes | -| --- | --- | --- | --- | --- | --- | -| `1..32` | `N=768,K=2048` | small-M `lightop.moe_gemm_w8a8`, often `BLOCK_M=16` | `N=2048,K=384` | small-M `lightop.moe_gemm_w8a8` | decode/graph capture sizes | -| `896` | `N=768,K=2048` | `BLOCK_M=64, MODE=517, DELTA=1`, HIP NT prefill up | `N=2048,K=384` | `BLOCK_M=32, MODE=568, DELTA=2`, HIP NT prefill down | tail chunk | -| `4096` | `N=768,K=2048` | `BLOCK_M=128, MODE=1000, DELTA=1`, `MOE_W8A8_I8_PERCHANNEL_ASM_TN_MT128x256x128_WGM1_UP` | `N=2048,K=384` | `BLOCK_M=64, MODE=517, DELTA=2`, HIP NT prefill down | target single request prefill | -| `10240` | `N=768,K=2048` | `BLOCK_M=128, MODE=1000, DELTA=1`, same ASM UP kernel | `N=2048,K=384` | `BLOCK_M=64, MODE=523, DELTA=2`, HIP NT prefill down | vLLM chunked prefill example | - -vLLM with 16 concurrent 8K prompts enabled chunked prefill with `max_num_batched_tokens=10240`. The observed MoE effective sizes were `10240`, `8256`, `896`, plus small graph-capture sizes. This means scheduler chunking, not only `VLLM_FUSED_MOE_CHUNK_SIZE`, controls the actual large-M MoE calls. - -## vLLM W8A8 Dense Linear Path - -Main call chain: - -1. `CompressedTensorsW8A8Int8.apply_weights()` -2. `apply_int8_linear(..., w8a8_strategy=3)` -3. `per_token_quant_int8(...)` -4. `ops.blaslt_scaled_mm(...)` -5. backend 3: `hipblaslt_w8a8_channelwise_gemm` - -Representative kernels: - -- `M=1,N=4096,K=2048`: small `Cijk_Alik_Bljk_I8BS_MT64x16x256...` -- `M=4096,N=4096,K=2048`: large `Cijk_Alik_Bljk_I8BS_MT256x256x128...` - -## Implementation Direction - -The next code change should move InfiniLM W8A8 MoE toward vLLM's ordinary channel-wise path: - -1. Add a new W8A8 channel MoE backend in InfiniLM, keeping `[E,N,K]` weights and `[E,N,1]` scales instead of calling `moe_w8a8_marlin_pack`. -2. Add/route an InfiniCore wrapper around ordinary `lightop.moe_gemm_w8a8`, not the current `moe_gemm_marlin_w8a8` adaptor. -3. Reuse the existing MoE workspace pattern where possible: int8 hidden cache, int8 intermediate cache, per-token scales, BF16 intermediate/output buffers. -4. Select configs by effective `M` and GEMM shape, matching the vLLM anchors above first: small decode, `896`, `4096`, `10240`. -5. Default the MoE chunk cap to `16384` for the ordinary channel path, matching vLLM's fused MoE chunk cap. Scheduler-level chunking is still needed later for 16 concurrency and 8K-10K contexts. - -## Useful Remote Artifacts - -- vLLM probe log: `/tmp/vllm_w8a8_kernel_probe_i8192_c16_o16_20260710_110631.server.log` -- MoE micro traces: - - `/tmp/hygon_trace_lmslim_w8a8_moe_m10240_20260710_111507` - - `/tmp/hygon_trace_lmslim_w8a8_moe_m896_20260710_111603` - - `/tmp/hygon_trace_lmslim_w8a8_moe_m4096_20260710_111646` -- Dense linear micro traces: - - `/tmp/hygon_trace_vllm_w8a8_linear_m1_n4096_k2048_20260710_113331` - - `/tmp/hygon_trace_vllm_w8a8_linear_m4096_n4096_k2048_20260710_113413` diff --git a/python/infinilm/infer_engine.py b/python/infinilm/infer_engine.py index fb36dc31..11bcdaa0 100644 --- a/python/infinilm/infer_engine.py +++ b/python/infinilm/infer_engine.py @@ -1,5 +1,4 @@ import json -import operator import os import time from dataclasses import dataclass @@ -19,45 +18,6 @@ } -def _validate_nll_score_inputs(input_ids, labels, score_start): - """Validate an explicit batch-1, shifted-token NLL request.""" - for name, tensor in (("input_ids", input_ids), ("labels", labels)): - if tensor is None: - raise TypeError(f"{name} must be an infinicore tensor") - missing = [ - attr - for attr in ("ndim", "shape", "dtype", "_underlying") - if not hasattr(tensor, attr) - ] - if missing: - raise TypeError( - f"{name} must be an infinicore tensor; missing {', '.join(missing)}" - ) - if tensor.dtype != infinicore.int64: - raise ValueError(f"{name} must use infinicore.int64 dtype") - if tensor.ndim != 2: - raise ValueError(f"{name} must be a rank-2 tensor") - - input_shape = tuple(input_ids.shape) - label_shape = tuple(labels.shape) - if input_shape != label_shape: - raise ValueError("input_ids and labels must have identical shapes") - if input_shape[0] != 1: - raise ValueError("score_nll currently requires batch_size=1") - - if isinstance(score_start, bool): - raise TypeError("score_start must be an integer, not bool") - try: - score_start = operator.index(score_start) - except TypeError as error: - raise TypeError("score_start must be an integer") from error - - seq_len = input_shape[1] - if score_start < 0 or score_start >= seq_len: - raise ValueError("score_start must select at least one token") - return seq_len, score_start - - def _apply_torch_dtype_defaults(config: dict) -> dict: if config.get("torch_dtype") is None: config["torch_dtype"] = config.get("dtype") or _MODEL_DEFAULTS.get( @@ -724,98 +684,6 @@ def generate( return output_ids - def score_nll(self, input_ids, labels, *, score_start=0): - """Return summed shifted-token NLL and token count for a batch-1 window. - - This explicit evaluation path bypasses graph replay and never changes the - behavior of ``forward``/``generate``. ``input_ids`` and ``labels`` must - have the same ``[1, sequence]`` shape; callers perform the causal shift. - """ - try: - seq_len, score_start = _validate_nll_score_inputs( - input_ids, labels, score_start - ) - - block_tables = None - slot_mapping = None - if self.enable_paged_attn: - cache_config = self.get_cache_config() - if cache_config is None: - raise RuntimeError("paged attention requires a cache configuration") - paged_block_size = cache_config.block_size() - max_blocks_per_batch = ( - seq_len + paged_block_size - 1 - ) // paged_block_size - if max_blocks_per_batch > cache_config.num_blocks(): - raise ValueError( - "NLL sequence requires more paged KV-cache blocks than " - "the current cache configuration provides" - ) - block_tables = infinicore.from_list( - [list(range(max_blocks_per_batch))], - dtype=infinicore.int32, - ) - slot_mapping = infinicore.from_list( - list(range(seq_len)), dtype=infinicore.int64 - ) - position_ids = infinicore.from_list( - list(range(seq_len)), dtype=infinicore.int64 - ) - else: - position_ids = infinicore.from_list( - [list(range(seq_len))], dtype=infinicore.int64 - ) - past_kv_lengths = infinicore.from_list([0], dtype=infinicore.int32) - total_kv_lengths = infinicore.from_list( - [seq_len], dtype=infinicore.int32 - ) - cu_seqlens = infinicore.from_list( - [0, seq_len], dtype=infinicore.int32 - ) - input_offsets = infinicore.from_list( - [0, seq_len], dtype=infinicore.int32 - ) - - output = super().forward( - super().Input( - input_ids._underlying, - position_ids=position_ids._underlying, - past_sequence_lengths=past_kv_lengths._underlying, - total_sequence_lengths=total_kv_lengths._underlying, - input_offsets=input_offsets._underlying, - cu_seqlens=cu_seqlens._underlying, - block_tables=( - block_tables._underlying - if block_tables is not None - else None - ), - slot_mapping=( - slot_mapping._underlying - if slot_mapping is not None - else None - ), - labels=labels._underlying, - score_start=score_start, - return_nll=True, - ) - ) - token_nll = infinicore.Tensor(output.nll).to_numpy() - scored_tokens = int(output.scored_tokens) - expected_scored_tokens = seq_len - score_start - if scored_tokens != expected_scored_tokens: - raise RuntimeError( - "score_nll returned an invalid scored-token count: " - f"expected {expected_scored_tokens}, got {scored_tokens}" - ) - if token_nll.size != expected_scored_tokens: - raise RuntimeError( - "score_nll returned a token-loss vector with an invalid size" - ) - return float(token_nll.astype("float64").sum()), scored_tokens - except BaseException as e: - handle_oom_and_exit(e) - raise - def reset_cache(self, cache_config): infinicore.sync_device() self.enable_paged_attn = isinstance(cache_config, PagedKVCacheConfig) diff --git a/test/engine/test_nll_validation.py b/test/engine/test_nll_validation.py deleted file mode 100644 index cc2f64ab..00000000 --- a/test/engine/test_nll_validation.py +++ /dev/null @@ -1,145 +0,0 @@ -import importlib.util -import sys -import types -from pathlib import Path - -import pytest - - -class FakeTensor: - def __init__(self, shape, dtype): - self.shape = tuple(shape) - self.ndim = len(self.shape) - self.dtype = dtype - self._underlying = object() - - -@pytest.fixture -def validator(monkeypatch): - """Load the pure validator without importing the hardware runtime.""" - package_root = Path(__file__).resolve().parents[2] / "python" / "infinilm" - int64_dtype = object() - - infinilm_package = types.ModuleType("infinilm") - infinilm_package.__path__ = [str(package_root)] - monkeypatch.setitem(sys.modules, "infinilm", infinilm_package) - - fake_infinicore = types.ModuleType("infinicore") - fake_infinicore.int64 = int64_dtype - fake_infinicore.Tensor = type("Tensor", (), {}) - monkeypatch.setitem(sys.modules, "infinicore", fake_infinicore) - - cache_module = types.ModuleType("infinilm.cache") - cache_module.PagedKVCacheConfig = type("PagedKVCacheConfig", (), {}) - monkeypatch.setitem(sys.modules, "infinilm.cache", cache_module) - - distributed_module = types.ModuleType("infinilm.distributed") - distributed_module.DistConfig = type("DistConfig", (), {}) - monkeypatch.setitem(sys.modules, "infinilm.distributed", distributed_module) - - native_engine = type("InferEngine", (), {}) - lib_module = types.ModuleType("infinilm.lib") - lib_module._infinilm = types.SimpleNamespace(InferEngine=native_engine) - monkeypatch.setitem(sys.modules, "infinilm.lib", lib_module) - - exception_module = types.ModuleType("infinilm.exception_utils") - exception_module.handle_oom_and_exit = lambda error: None - monkeypatch.setitem(sys.modules, "infinilm.exception_utils", exception_module) - - modeling_module = types.ModuleType("infinilm.modeling_utils") - modeling_module.parse_dtype = lambda dtype: dtype - monkeypatch.setitem(sys.modules, "infinilm.modeling_utils", modeling_module) - - module_name = "infinilm.infer_engine" - module_path = package_root / "infer_engine.py" - spec = importlib.util.spec_from_file_location(module_name, module_path) - module = importlib.util.module_from_spec(spec) - monkeypatch.setitem(sys.modules, module_name, module) - assert spec.loader is not None - spec.loader.exec_module(module) - - return module._validate_nll_score_inputs, int64_dtype - - -def make_tensor(shape, int64_dtype, dtype=None): - return FakeTensor(shape, int64_dtype if dtype is None else dtype) - - -def test_validate_nll_score_inputs_accepts_valid_window(validator): - validate, int64_dtype = validator - input_ids = make_tensor((1, 8), int64_dtype) - labels = make_tensor((1, 8), int64_dtype) - - assert validate(input_ids, labels, 3) == (8, 3) - - -@pytest.mark.parametrize("name", ["input_ids", "labels"]) -def test_validate_nll_score_inputs_requires_tensor_protocol(validator, name): - validate, int64_dtype = validator - tensors = { - "input_ids": make_tensor((1, 8), int64_dtype), - "labels": make_tensor((1, 8), int64_dtype), - } - tensors[name] = object() - - with pytest.raises(TypeError, match=name): - validate(tensors["input_ids"], tensors["labels"], 0) - - -@pytest.mark.parametrize("name", ["input_ids", "labels"]) -def test_validate_nll_score_inputs_requires_int64(validator, name): - validate, int64_dtype = validator - tensors = { - "input_ids": make_tensor((1, 8), int64_dtype), - "labels": make_tensor((1, 8), int64_dtype), - } - tensors[name] = make_tensor((1, 8), int64_dtype, dtype=object()) - - with pytest.raises(ValueError, match=f"{name} must use infinicore.int64"): - validate(tensors["input_ids"], tensors["labels"], 0) - - -@pytest.mark.parametrize( - ("input_shape", "label_shape", "message"), - [ - ((8,), (8,), "rank-2"), - ((1, 8), (1, 7), "identical shapes"), - ((2, 8), (2, 8), "batch_size=1"), - ], -) -def test_validate_nll_score_inputs_rejects_invalid_shapes( - validator, input_shape, label_shape, message -): - validate, int64_dtype = validator - with pytest.raises(ValueError, match=message): - validate( - make_tensor(input_shape, int64_dtype), - make_tensor(label_shape, int64_dtype), - 0, - ) - - -@pytest.mark.parametrize("score_start", [-1, 8]) -def test_validate_nll_score_inputs_rejects_empty_score_range( - validator, score_start -): - validate, int64_dtype = validator - with pytest.raises(ValueError, match="select at least one token"): - validate( - make_tensor((1, 8), int64_dtype), - make_tensor((1, 8), int64_dtype), - score_start, - ) - - -@pytest.mark.parametrize("score_start", [True, 1.5, "1"]) -def test_validate_nll_score_inputs_requires_integer_score_start( - validator, score_start -): - validate, int64_dtype = validator - with pytest.raises(TypeError, match="score_start must be an integer"): - validate( - make_tensor((1, 8), int64_dtype), - make_tensor((1, 8), int64_dtype), - score_start, - ) diff --git a/test/ppl/qwen3_235b/README.md b/test/ppl/qwen3_235b/README.md deleted file mode 100644 index 559fd7c1..00000000 --- a/test/ppl/qwen3_235b/README.md +++ /dev/null @@ -1,200 +0,0 @@ -# Qwen3-235B true PPL CLI - -This directory contains reproducible token-level perplexity tools for: - -- Transformers BF16 on TP8 -- InfiniLM BF16 on TP8 -- InfiniLM W8A8 on TP8 - -The runners consume the same frozen token manifest and calculate causal, -shifted-token cross entropy: - -```text -mean_nll = sum(-log p(x_t | x_&1 | tee "$LOG_DIR/infinilm_w8a8_smoke.log" -rc=${PIPESTATUS[0]} -echo "INFINILM_W8A8_SMOKE_EXIT_CODE=$rc" -hy-smi --showpids -``` - -## Full WikiText-2 runs - -`--max-scored-tokens 0` scores every target token in the manifest. Use the same -`window`, `stride` and `max-scored-tokens` values for every backend. - -Transformers BF16: - -```bash -set -o pipefail -timeout --signal=TERM --kill-after=60s 21600s \ - python -u "$PPL_ROOT/scripts/transformers/pytorch_ppl_Qwen3_235B.py" \ - --model "$MODEL_BF16" \ - --token-manifest "$TOKEN_MANIFEST" \ - --window 256 \ - --stride 128 \ - --max-scored-tokens 0 \ - --tp-size 8 \ - --attention eager \ - --json-output "$LOG_DIR/transformers_bf16_full.json" \ - 2>&1 | tee "$LOG_DIR/transformers_bf16_full.log" -``` - -The Transformers entry point launches `torchrun` itself. Do not wrap it in a -second `torchrun` command. Eager attention is the validated Hygon path. - -InfiniLM BF16: - -```bash -set -o pipefail -timeout --signal=TERM --kill-after=60s 21600s \ - python -u "$PPL_ROOT/scripts/infinilm/infinilm_ppl_Qwen3_235B.py" \ - --model "$MODEL_BF16" \ - --token-manifest "$TOKEN_MANIFEST" \ - --window 256 \ - --stride 128 \ - --max-scored-tokens 0 \ - --tp-size 8 \ - --attention flash-attn \ - --json-output "$LOG_DIR/infinilm_bf16_full.json" \ - 2>&1 | tee "$LOG_DIR/infinilm_bf16_full.log" -``` - -InfiniLM W8A8: - -```bash -set -o pipefail -timeout --signal=TERM --kill-after=60s 21600s \ - python -u "$PPL_ROOT/scripts/infinilm/infinilm_ppl_Qwen3_235B.py" \ - --model "$MODEL_W8A8" \ - --token-manifest "$TOKEN_MANIFEST" \ - --window 256 \ - --stride 128 \ - --max-scored-tokens 0 \ - --tp-size 8 \ - --attention flash-attn \ - --json-output "$LOG_DIR/infinilm_w8a8_full.json" \ - 2>&1 | tee "$LOG_DIR/infinilm_w8a8_full.log" -``` - -For a bounded formal run, replace `0` with the same positive token count in all -three commands, for example `10240`. - -## Compare results - -Transformers BF16 versus InfiniLM W8A8: - -```bash -python -u "$PPL_ROOT/scripts/calculate_true_ppl.py" \ - --inputs \ - "$LOG_DIR/transformers_bf16_full.json" \ - "$LOG_DIR/infinilm_w8a8_full.json" \ - --max-ppl-increase-percent 20 \ - --json-out "$LOG_DIR/ppl_transformers_vs_w8a8.json" -``` - -InfiniLM BF16 versus InfiniLM W8A8: - -```bash -python -u "$PPL_ROOT/scripts/calculate_infinilm_precision_ppl.py" \ - --inputs \ - "$LOG_DIR/infinilm_bf16_full.json" \ - "$LOG_DIR/infinilm_w8a8_full.json" \ - --max-ppl-increase-percent 20 \ - --json-out "$LOG_DIR/ppl_bf16_vs_w8a8.json" -``` - -Exit code `0` means the configured PPL increase threshold passed, `1` means it -failed, and `2` means the input files are invalid or describe different -workloads. - -## Scope - -PPL is a quality test. InfiniLM intentionally disables graph only for the -explicit `score_nll` path because it must retain full token logits/losses. -Normal generation and formal performance tests keep their existing graph path. -Do not report PPL scoring throughput as inference performance. diff --git a/test/ppl/qwen3_235b/scripts/_gpu_guard.py b/test/ppl/qwen3_235b/scripts/_gpu_guard.py deleted file mode 100755 index 1ceaeddd..00000000 --- a/test/ppl/qwen3_235b/scripts/_gpu_guard.py +++ /dev/null @@ -1,100 +0,0 @@ -#!/usr/bin/env python3 -"""Fail closed unless all eight Hygon devices are idle.""" - -from __future__ import annotations - -import os -import subprocess -import sys -from pathlib import Path - - -EXPECTED_DEVICES = set(range(8)) -SMI_TIMEOUT_SECONDS = 60 - - -def _local_gpu_processes() -> list[str]: - users: list[str] = [] - own_pid = os.getpid() - for process_dir in Path("/proc").glob("[0-9]*"): - try: - pid = int(process_dir.name) - except ValueError: - continue - if pid == own_pid: - continue - try: - targets = [entry.resolve() for entry in (process_dir / "fd").iterdir()] - except OSError: - continue - if not any( - str(target) == "/dev/kfd" or str(target).startswith("/dev/dri/renderD") - for target in targets - ): - continue - try: - command = (process_dir / "cmdline").read_bytes().replace(b"\0", b" ").decode( - "utf-8", errors="replace" - ).strip() - except OSError: - command = "" - users.append(f"pid={pid} command={command or '[unknown]'}") - return sorted(users) - - -def require_idle_gpu() -> None: - try: - result = subprocess.run( - ["hy-smi"], - check=False, - text=True, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - timeout=SMI_TIMEOUT_SECONDS, - ) - except (OSError, subprocess.TimeoutExpired) as error: - print(f"拒绝启动：hy-smi 空闲检查失败：{error}", file=sys.stderr) - raise SystemExit(90) from error - if result.returncode != 0: - print( - f"拒绝启动：hy-smi 退出码为 {result.returncode}。\n{result.stdout}", - file=sys.stderr, - ) - raise SystemExit(90) - - utilization: dict[int, tuple[float, float]] = {} - for line in result.stdout.splitlines(): - fields = line.split() - if ( - len(fields) >= 7 - and fields[0].isdigit() - and fields[5].endswith("%") - and fields[6].endswith("%") - ): - try: - utilization[int(fields[0])] = ( - float(fields[5][:-1]), - float(fields[6][:-1]), - ) - except ValueError: - continue - if set(utilization) != EXPECTED_DEVICES: - print( - "拒绝启动：hy-smi 未完整报告 0-7 号设备。\n" + result.stdout, - file=sys.stderr, - ) - raise SystemExit(90) - - busy_devices = { - device: values - for device, values in utilization.items() - if values[0] > 0.0 or values[1] > 0.0 - } - local_users = _local_gpu_processes() - if busy_devices or local_users: - print( - "拒绝启动：GPU 未完全空闲；" - f"设备占用={busy_devices}，容器内进程={local_users}。\n{result.stdout}", - file=sys.stderr, - ) - raise SystemExit(90) diff --git a/test/ppl/qwen3_235b/scripts/_ppl_common.py b/test/ppl/qwen3_235b/scripts/_ppl_common.py deleted file mode 100755 index bc720547..00000000 --- a/test/ppl/qwen3_235b/scripts/_ppl_common.py +++ /dev/null @@ -1,338 +0,0 @@ -#!/usr/bin/env python3 -"""Shared, deterministic corpus and sliding-window helpers for true PPL tests.""" - -from __future__ import annotations - -import ast -import array -import hashlib -import json -import operator -import re -import struct -import sys -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Iterable, Iterator, Sequence - - -CORPUS_SCHEMA = "qw235_ppl_token_ids_v1" -SCORING_METHOD = ( - "sliding_window_shifted_cross_entropy_fp32_compute_fp64_accumulation" -) -SHA256_RE = re.compile(r"^[0-9a-f]{64}$") - - -def canonical_json_bytes(value: object) -> bytes: - return json.dumps( - value, ensure_ascii=True, sort_keys=True, separators=(",", ":") - ).encode("ascii") - - -def _canonical_int_sequence_sha256(values: Iterable[int], label: str) -> str: - """Hash an integer sequence exactly like compact JSON ``[1,2,3]``.""" - digest = hashlib.sha256() - digest.update(b"[") - for index, value in enumerate(values): - if isinstance(value, bool): - raise ValueError(f"{label}[{index}] 不是非负整数") - try: - parsed = operator.index(value) - except TypeError as error: - raise ValueError(f"{label}[{index}] 不是非负整数") from error - if parsed < 0: - raise ValueError(f"{label}[{index}] 不是非负整数：{value!r}") - if index: - digest.update(b",") - digest.update(str(parsed).encode("ascii")) - digest.update(b"]") - return digest.hexdigest() - - -def canonical_token_ids_sha256(token_ids: Iterable[int]) -> str: - return _canonical_int_sequence_sha256(token_ids, "token_ids") - - -def canonical_indices_sha256(indices: Iterable[int]) -> str: - return _canonical_int_sequence_sha256(indices, "indices") - - -@dataclass(frozen=True) -class PplCorpusManifest: - path: Path - payload: dict[str, Any] - token_ids: tuple[int, ...] - manifest_sha256: str - token_ids_sha256: str - - @property - def token_count(self) -> int: - return len(self.token_ids) - - -@dataclass(frozen=True) -class SlidingWindow: - """One causal-LM window using half-open global token index ranges. - - ``token_start:token_end`` is model input. Targets in - ``score_start:score_end`` are scored. ``prediction_*`` select the matching - logits before the causal shift, while ``target_*`` select labels locally. - """ - - index: int - token_start: int - token_end: int - score_start: int - score_end: int - token_ids: tuple[int, ...] - - @property - def scored_token_count(self) -> int: - return self.score_end - self.score_start - - @property - def prediction_start(self) -> int: - return self.score_start - self.token_start - 1 - - @property - def prediction_end(self) -> int: - return self.score_end - self.token_start - 1 - - @property - def target_start(self) -> int: - return self.score_start - self.token_start - - @property - def target_end(self) -> int: - return self.score_end - self.token_start - - -def _required_positive_int(payload: dict[str, Any], key: str, path: Path) -> int: - value = payload.get(key) - if isinstance(value, bool) or not isinstance(value, int): - raise ValueError(f"{path} 的 {key} 必须是正整数") - parsed = value - if parsed <= 0: - raise ValueError(f"{path} 的 {key} 必须是正整数") - return parsed - - -def _required_sha(payload: dict[str, Any], key: str, path: Path) -> str: - value = str(payload.get(key, "")).lower() - if not SHA256_RE.fullmatch(value): - raise ValueError(f"{path} 的 {key} 不是有效 SHA256") - return value - - -def _load_npy(manifest_path: Path, relative_name: object) -> list[int]: - relative = Path(str(relative_name)) - if relative.is_absolute() or ".." in relative.parts: - raise ValueError(f"{manifest_path} 的 token_ids_file 必须是安全相对路径") - base = manifest_path.parent.resolve() - token_path = (base / relative).resolve() - try: - token_path.relative_to(base) - except ValueError as error: - raise ValueError(f"token_ids_file 越出 manifest 目录：{relative}") from error - try: - import numpy as np - except ImportError: - return _load_int64_npy_without_numpy(token_path) - try: - array = np.load(token_path, allow_pickle=False) - except FileNotFoundError: - raise ValueError(f"token_ids_file 不存在：{token_path}") from None - if array.ndim != 1 or array.dtype.kind not in "iu": - raise ValueError(f"{token_path} 必须是一维整数 .npy 数组") - return [int(value) for value in array.tolist()] - - -def _load_int64_npy_without_numpy(path: Path) -> list[int]: - try: - with path.open("rb") as handle: - if handle.read(6) != b"\x93NUMPY": - raise ValueError(f"{path} 不是有效 .npy 文件") - version = handle.read(2) - if version == b"\x01\x00": - header_length = struct.unpack(" None: - """Write a portable NumPy v1.0, one-dimensional little-endian int64 file.""" - output = Path(path) - values = list(token_ids) - # Validate before creating a partial file. - canonical_token_ids_sha256(values) - header_text = repr( - {"descr": " 65535: - raise ValueError(".npy header 超过 v1.0 长度限制") - output.parent.mkdir(parents=True, exist_ok=True) - packed = array.array("q", (int(value) for value in values)) - if packed.itemsize != 8: - raise RuntimeError("当前 Python 平台的 signed long long 不是 64 bit") - if sys.byteorder != "little": - packed.byteswap() - with output.open("wb") as handle: - handle.write(b"\x93NUMPY") - handle.write(b"\x01\x00") - handle.write(struct.pack(" PplCorpusManifest: - """Load and fully verify an inline or relative-``.npy`` corpus manifest.""" - manifest_path = Path(path) - try: - payload = json.loads(manifest_path.read_text(encoding="utf-8")) - except FileNotFoundError: - raise ValueError(f"PPL manifest 不存在：{manifest_path}") from None - except json.JSONDecodeError as error: - raise ValueError(f"PPL manifest JSON 无效：{manifest_path}: {error.msg}") from error - if not isinstance(payload, dict): - raise ValueError(f"PPL manifest 必须是 JSON 对象：{manifest_path}") - if payload.get("schema") != CORPUS_SCHEMA: - raise ValueError( - f"{manifest_path} schema 必须为 {CORPUS_SCHEMA!r}，" - f"实际为 {payload.get('schema')!r}" - ) - manifest_hash = _required_sha(payload, "manifest_sha256", manifest_path) - semantic_payload = dict(payload) - semantic_payload.pop("manifest_sha256", None) - calculated_manifest_hash = hashlib.sha256( - canonical_json_bytes(semantic_payload) - ).hexdigest() - if manifest_hash != calculated_manifest_hash: - raise ValueError(f"{manifest_path} 的 manifest_sha256 校验失败") - for key in ("source_sha256", "tokenizer_sha256"): - _required_sha(payload, key, manifest_path) - - has_inline = "token_ids" in payload - has_file = "token_ids_file" in payload - if has_inline == has_file: - raise ValueError( - f"{manifest_path} 必须且只能包含 token_ids 或 token_ids_file 之一" - ) - if has_inline: - raw_ids = payload["token_ids"] - if not isinstance(raw_ids, list): - raise ValueError(f"{manifest_path} 的 token_ids 必须是数组") - token_ids = list(raw_ids) - else: - token_ids = _load_npy(manifest_path, payload["token_ids_file"]) - - # The canonical hash validates type, integrality, sign, order and contents. - calculated_token_hash = canonical_token_ids_sha256(token_ids) - expected_token_hash = _required_sha( - payload, "token_ids_sha256", manifest_path - ) - if calculated_token_hash != expected_token_hash: - raise ValueError(f"{manifest_path} 的 token_ids_sha256 校验失败") - token_count = _required_positive_int(payload, "token_count", manifest_path) - if token_count != len(token_ids) or token_count < 2: - raise ValueError( - f"{manifest_path} token_count={token_count}，实际 token 数={len(token_ids)}" - ) - return PplCorpusManifest( - path=manifest_path, - payload=payload, - token_ids=tuple(int(value) for value in token_ids), - manifest_sha256=manifest_hash, - token_ids_sha256=expected_token_hash, - ) - - -def iter_sliding_windows( - token_ids: Sequence[int], - window_size: int, - stride: int, - max_scored_tokens: int | None = None, -) -> Iterator[SlidingWindow]: - """Yield windows that score global token indices ``1..N-1`` exactly once. - - The first window scores ``1:end``. Every later window scores only - ``previous_end:end``; overlapped prefix tokens provide context but are not - counted again. ``stride`` must be smaller than ``window_size`` so the first - new target in every later window retains its immediately preceding token. - """ - if isinstance(window_size, bool) or not isinstance(window_size, int): - raise ValueError("window_size 必须是整数") - if isinstance(stride, bool) or not isinstance(stride, int): - raise ValueError("stride 必须是整数") - if window_size < 2: - raise ValueError("window_size 必须至少为 2") - if stride < 1 or stride >= window_size: - raise ValueError("stride 必须满足 1 <= stride < window_size") - if len(token_ids) < 2: - raise ValueError("至少需要 2 个 token 才能计算 PPL") - if max_scored_tokens is not None: - if ( - isinstance(max_scored_tokens, bool) - or not isinstance(max_scored_tokens, int) - or max_scored_tokens < 1 - ): - raise ValueError("max_scored_tokens 必须是正整数") - - score_limit = len(token_ids) - if max_scored_tokens is not None: - score_limit = min(score_limit, 1 + max_scored_tokens) - previous_end = 1 - index = 0 - while previous_end < score_limit: - if index == 0: - token_end = min(score_limit, window_size) - token_start = 0 - score_start = 1 - else: - token_end = min(score_limit, previous_end + stride) - token_start = max(0, token_end - window_size) - score_start = previous_end - window = SlidingWindow( - index=index, - token_start=token_start, - token_end=token_end, - score_start=score_start, - score_end=token_end, - token_ids=tuple(int(value) for value in token_ids[token_start:token_end]), - ) - if window.prediction_start < 0 or window.prediction_end > len(window.token_ids): - raise AssertionError("内部错误：滑窗缺少 causal predecessor") - yield window - previous_end = token_end - index += 1 diff --git a/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py b/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py deleted file mode 100644 index 5bd49e86..00000000 --- a/test/ppl/qwen3_235b/scripts/calculate_infinilm_precision_ppl.py +++ /dev/null @@ -1,84 +0,0 @@ -#!/usr/bin/env python3 -"""Compare InfiniLM BF16 and W8A8 true-PPL result JSON files.""" - -from __future__ import annotations - -import argparse -import json -import os -from dataclasses import asdict -from pathlib import Path -from typing import Any, Sequence - -from calculate_true_ppl import _validate_same_workload, load_result - - -SCHEMA = "qwen3_235b_infinilm_precision_ppl_comparison/v1" - - -def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--inputs", nargs=2, type=Path, required=True) - parser.add_argument("--max-ppl-increase-percent", type=float, default=20.0) - parser.add_argument("--json-out", type=Path, required=True) - args = parser.parse_args(argv) - if args.max_ppl_increase_percent < 0: - parser.error("--max-ppl-increase-percent must be non-negative") - return args - - -def _atomic_json(path: Path, payload: dict[str, Any]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") - temporary.write_text( - json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - temporary.replace(path) - - -def main(argv: Sequence[str] | None = None) -> int: - args = _parse_args(argv) - try: - results = [load_result(path) for path in args.inputs] - if any(result.backend != "infinilm" for result in results): - raise ValueError("both inputs must be InfiniLM results") - precisions: list[str] = [] - for path in args.inputs: - payload = json.loads(path.read_text(encoding="utf-8")) - precision = str(payload.get("precision", "")).strip().upper() - if precision not in {"BF16", "W8A8"}: - raise ValueError(f"{path} has invalid precision: {precision!r}") - precisions.append(precision) - by_precision = dict(zip(precisions, results, strict=True)) - if set(by_precision) != {"BF16", "W8A8"}: - raise ValueError("inputs must contain one BF16 and one W8A8 result") - baseline = by_precision["BF16"] - candidate = by_precision["W8A8"] - _validate_same_workload(baseline, candidate) - increase = (candidate.ppl / baseline.ppl - 1.0) * 100.0 - threshold = float(args.max_ppl_increase_percent) - passed = increase <= threshold - payload = { - "schema": SCHEMA, - "status": "PASS" if passed else "FAIL", - "baseline": asdict(baseline), - "candidate": asdict(candidate), - "ppl_increase_percent": increase, - "max_ppl_increase_percent": threshold, - } - _atomic_json(args.json_out, payload) - except (OSError, ValueError, KeyError, TypeError, json.JSONDecodeError) as error: - print(f"PPL comparison error: {error}") - return 2 - - print(f"InfiniLM BF16 PPL: {baseline.ppl:.6f}") - print(f"InfiniLM W8A8 PPL: {candidate.ppl:.6f}") - print(f"W8A8 PPL increase: {increase:.2f}%") - print(f"Quality threshold: <= {threshold:.2f}%") - print(f"Result: {'PASS' if passed else 'FAIL'}") - return 0 if passed else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py b/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py deleted file mode 100755 index cfefc232..00000000 --- a/test/ppl/qwen3_235b/scripts/calculate_true_ppl.py +++ /dev/null @@ -1,305 +0,0 @@ -#!/usr/bin/env python3 -"""严格比较 Transformers 与 InfiniLM 的真实 token-level PPL 结果。""" - -from __future__ import annotations - -import argparse -import json -import math -import os -import re -import sys -import tempfile -from dataclasses import asdict, dataclass -from pathlib import Path -from typing import Any, Sequence - -from _ppl_common import SCORING_METHOD, canonical_indices_sha256 - -RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" -COMPARISON_SCHEMA = "qwen3_235b_true_ppl_comparison/v1" -SHA256_RE = re.compile(r"^[0-9a-f]{64}$") - - -@dataclass(frozen=True) -class PplResult: - path: str - backend: str - model: str - corpus_manifest_sha256: str - corpus_token_ids_sha256: str - window_size: int - stride: int - scoring_method: str - first_scored_token_index: int - last_scored_token_index_exclusive: int - scored_token_count: int - scored_token_indices_sha256: str - total_nll: float - mean_nll: float - ppl: float - - -def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("paths", nargs="*", type=Path) - parser.add_argument( - "--inputs", - nargs="+", - type=Path, - default=[], - help="Transformers 与 InfiniLM 结果 JSON，顺序可以互换", - ) - parser.add_argument( - "--max-ppl-increase-percent", - type=float, - default=20.0, - help="InfiniLM 相对 Transformers 的最大 PPL 增幅，默认 20%%", - ) - parser.add_argument("--json-out", type=Path) - parser.add_argument("--verbose", action="store_true", help="打印完整 JSON") - args = parser.parse_args(argv) - args.inputs = [*args.paths, *args.inputs] - if len(args.inputs) != 2: - parser.error("必须提供两个结果 JSON：Transformers 与 InfiniLM") - if ( - not math.isfinite(args.max_ppl_increase_percent) - or args.max_ppl_increase_percent < 0 - ): - parser.error("--max-ppl-increase-percent 必须是有限非负数") - return args - - -def _required(payload: dict[str, Any], key: str, path: Path) -> Any: - if key not in payload: - raise ValueError(f"{path} 缺少字段 {key}") - return payload[key] - - -def _sha256(payload: dict[str, Any], key: str, path: Path) -> str: - value = str(_required(payload, key, path)).lower() - if not SHA256_RE.fullmatch(value): - raise ValueError(f"{path} 的 {key} 不是有效 SHA256") - return value - - -def _positive_int(payload: dict[str, Any], key: str, path: Path) -> int: - value = _required(payload, key, path) - if isinstance(value, bool) or not isinstance(value, int): - raise ValueError(f"{path} 的 {key} 必须是正整数") - parsed = value - if parsed <= 0: - raise ValueError(f"{path} 的 {key} 必须是正整数") - return parsed - - -def _nonnegative_int(payload: dict[str, Any], key: str, path: Path) -> int: - value = _required(payload, key, path) - if isinstance(value, bool) or not isinstance(value, int): - raise ValueError(f"{path} 的 {key} 必须是非负整数") - parsed = value - if parsed < 0: - raise ValueError(f"{path} 的 {key} 必须是非负整数") - return parsed - - -def _finite(payload: dict[str, Any], key: str, path: Path) -> float: - try: - value = float(_required(payload, key, path)) - except (TypeError, ValueError) as error: - raise ValueError(f"{path} 的 {key} 必须是有限数") from error - if not math.isfinite(value): - raise ValueError(f"{path} 的 {key} 必须是有限数") - return value - - -def load_result(path: Path) -> PplResult: - try: - payload = json.loads(path.read_text(encoding="utf-8")) - except FileNotFoundError: - raise ValueError(f"结果文件不存在：{path}") from None - except json.JSONDecodeError as error: - raise ValueError(f"结果 JSON 无效：{path}: {error.msg}") from error - if not isinstance(payload, dict): - raise ValueError(f"结果 JSON 必须是对象：{path}") - if payload.get("status") != "PASS": - raise ValueError( - f"{path} 不是成功的 PPL 结果：status={payload.get('status')!r}" - ) - if payload.get("schema") != RESULT_SCHEMA: - raise ValueError( - f"{path} schema 必须为 {RESULT_SCHEMA!r}，实际为 {payload.get('schema')!r}" - ) - - backend = str(_required(payload, "backend", path)).strip().lower() - if backend not in {"transformers", "infinilm"}: - raise ValueError(f"{path} backend 必须是 transformers 或 infinilm") - model = str(_required(payload, "model", path)).strip() - if not model: - raise ValueError(f"{path} model 不能为空") - window_size = _positive_int(payload, "window_size", path) - stride = _positive_int(payload, "stride", path) - if stride >= window_size: - raise ValueError(f"{path} stride 必须小于 window_size") - scoring_method = str(_required(payload, "scoring_method", path)).strip() - if scoring_method != SCORING_METHOD: - raise ValueError( - f"{path} scoring_method 必须为 {SCORING_METHOD!r}" - ) - first_index = _nonnegative_int(payload, "first_scored_token_index", path) - last_index = _positive_int( - payload, "last_scored_token_index_exclusive", path - ) - scored_count = _positive_int(payload, "scored_token_count", path) - if last_index <= first_index or last_index - first_index != scored_count: - raise ValueError( - f"{path} 的计分范围 [{first_index}, {last_index}) 与 " - f"scored_token_count={scored_count} 不一致" - ) - if first_index != 1: - raise ValueError(f"{path} first_scored_token_index 必须为 1") - scored_indices_hash = _sha256( - payload, "scored_token_indices_sha256", path - ) - expected_indices_hash = canonical_indices_sha256( - range(first_index, last_index) - ) - if scored_indices_hash != expected_indices_hash: - raise ValueError(f"{path} 的 scored_token_indices_sha256 校验失败") - - total_nll = _finite(payload, "total_nll", path) - reported_mean = _finite(payload, "mean_nll", path) - reported_ppl = _finite(payload, "ppl", path) - if total_nll < 0 or reported_mean < 0 or reported_ppl < 1: - raise ValueError(f"{path} 的 NLL/PPL 超出有效范围") - calculated_mean = total_nll / scored_count - if calculated_mean > math.log(sys.float_info.max): - raise ValueError(f"{path} 的 mean NLL 过大，PPL 溢出") - calculated_ppl = math.exp(calculated_mean) - if not math.isclose(reported_mean, calculated_mean, rel_tol=1e-6, abs_tol=1e-8): - raise ValueError( - f"{path} 的 mean_nll 与 total_nll/scored_token_count 不一致" - ) - if not math.isclose(reported_ppl, calculated_ppl, rel_tol=1e-6, abs_tol=1e-8): - raise ValueError(f"{path} 的 ppl 与 exp(mean_nll) 不一致") - - return PplResult( - path=str(path), - backend=backend, - model=model, - corpus_manifest_sha256=_sha256( - payload, "corpus_manifest_sha256", path - ), - corpus_token_ids_sha256=_sha256( - payload, "corpus_token_ids_sha256", path - ), - window_size=window_size, - stride=stride, - scoring_method=scoring_method, - first_scored_token_index=first_index, - last_scored_token_index_exclusive=last_index, - scored_token_count=scored_count, - scored_token_indices_sha256=scored_indices_hash, - total_nll=total_nll, - mean_nll=calculated_mean, - ppl=calculated_ppl, - ) - - -def _ordered(results: Sequence[PplResult]) -> tuple[PplResult, PplResult]: - by_backend = {result.backend: result for result in results} - if len(by_backend) != 2 or set(by_backend) != {"transformers", "infinilm"}: - raise ValueError("必须且只能包含一份 Transformers 和一份 InfiniLM 结果") - return by_backend["transformers"], by_backend["infinilm"] - - -def _validate_same_workload(baseline: PplResult, candidate: PplResult) -> None: - fields = ( - "corpus_manifest_sha256", - "corpus_token_ids_sha256", - "window_size", - "stride", - "scoring_method", - "first_scored_token_index", - "last_scored_token_index_exclusive", - "scored_token_count", - "scored_token_indices_sha256", - ) - mismatches = [ - f"{field}: {getattr(baseline, field)!r} != {getattr(candidate, field)!r}" - for field in fields - if getattr(baseline, field) != getattr(candidate, field) - ] - if mismatches: - raise ValueError("两侧 PPL 工作负载不一致：" + "; ".join(mismatches)) - - -def compare( - baseline: PplResult, candidate: PplResult, threshold_percent: float -) -> dict[str, object]: - _validate_same_workload(baseline, candidate) - increase_percent = (candidate.ppl / baseline.ppl - 1.0) * 100.0 - passed = increase_percent <= threshold_percent - return { - "schema": COMPARISON_SCHEMA, - "status": "PASS" if passed else "FAIL", - "baseline": asdict(baseline), - "candidate": asdict(candidate), - "ppl_increase_percent": increase_percent, - "max_ppl_increase_percent": threshold_percent, - "pass": passed, - } - - -def _atomic_json(path: Path, payload: object) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) - try: - with os.fdopen(fd, "w", encoding="utf-8") as handle: - json.dump(payload, handle, ensure_ascii=False, sort_keys=True, indent=2) - handle.write("\n") - handle.flush() - os.fsync(handle.fileno()) - os.replace(temporary, path) - except BaseException: - try: - os.unlink(temporary) - except FileNotFoundError: - pass - raise - - -def main(argv: Sequence[str] | None = None) -> int: - args = _parse_args(argv) - try: - baseline, candidate = _ordered([load_result(path) for path in args.inputs]) - report = compare( - baseline, candidate, float(args.max_ppl_increase_percent) - ) - if args.json_out is not None: - _atomic_json(args.json_out, report) - except (OSError, RuntimeError, ValueError) as error: - print(f"错误：{error}", file=sys.stderr) - return 2 - - print("真实 PPL 对比") - print( - f"Transformers：PPL={baseline.ppl:.6f} " - f"NLL={baseline.total_nll:.6f} Token={baseline.scored_token_count}" - ) - print( - f"InfiniLM： PPL={candidate.ppl:.6f} " - f"NLL={candidate.total_nll:.6f} Token={candidate.scored_token_count}" - ) - print(f"PPL 增幅：{report['ppl_increase_percent']:.2f}%") - print( - f"验收要求：增幅 <= {args.max_ppl_increase_percent:.2f}% " - f"结果={report['status']}" - ) - if args.verbose: - print(json.dumps(report, ensure_ascii=False, sort_keys=True, indent=2)) - return 0 if report["pass"] else 1 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py deleted file mode 100755 index 9a2f3562..00000000 --- a/test/ppl/qwen3_235b/scripts/infinilm/infinilm_ppl_Qwen3_235B.py +++ /dev/null @@ -1,334 +0,0 @@ -#!/usr/bin/env python3 -"""Calculate true token-level PPL with the current InfiniLM C++ TP engine.""" - -from __future__ import annotations - -import argparse -import gc -import json -import math -import os -import sys -import time -from pathlib import Path -from typing import Any, Sequence - - -SCRIPT_DIR = Path(__file__).resolve().parent -SCRIPTS_DIR = SCRIPT_DIR.parent -if str(SCRIPTS_DIR) not in sys.path: - sys.path.insert(0, str(SCRIPTS_DIR)) - -from _gpu_guard import require_idle_gpu -from _ppl_common import ( - SCORING_METHOD, - canonical_indices_sha256, - iter_sliding_windows, - load_manifest, -) - - -RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" -DEFAULT_MODEL = "/data1/Qwen3_235B" -EXPECTED_MODEL_TYPE = "qwen3_moe" -EXPECTED_VOCAB_SIZE = 151936 -PAGED_KV_BLOCK_SIZE = 256 - - -def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="True shifted-token PPL for Qwen3_235B with InfiniLM TP8" - ) - parser.add_argument("--model", default=DEFAULT_MODEL) - parser.add_argument("--token-manifest", required=True) - parser.add_argument("--window", type=int, default=256) - parser.add_argument("--stride", type=int, default=128) - parser.add_argument( - "--max-scored-tokens", - type=int, - default=10240, - help="maximum target tokens to score; 0 scores the full manifest", - ) - parser.add_argument("--tp-size", type=int, default=8) - parser.add_argument("--attention", default="flash-attn") - parser.add_argument("--json-output") - args = parser.parse_args(argv) - - if not Path(args.model).is_dir(): - parser.error(f"model directory does not exist: {args.model}") - if not Path(args.token_manifest).is_file(): - parser.error(f"token manifest does not exist: {args.token_manifest}") - if args.window < 2: - parser.error("--window must be at least 2") - if args.stride < 1 or args.stride >= args.window: - parser.error("--stride must satisfy 1 <= stride < window") - if args.max_scored_tokens < 0: - parser.error("--max-scored-tokens must be non-negative") - if args.tp_size < 1: - parser.error("--tp-size must be positive") - - args.model = str(Path(args.model).resolve()) - args.token_manifest = str(Path(args.token_manifest).resolve()) - if args.json_output: - args.json_output = str(Path(args.json_output).resolve()) - return args - - -def _atomic_json(path_value: str, payload: dict[str, Any]) -> None: - path = Path(path_value) - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") - temporary.write_text( - json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - temporary.replace(path) - - -def _read_model_config(model_path: str) -> dict[str, Any]: - config_path = Path(model_path) / "config.json" - try: - config = json.loads(config_path.read_text(encoding="utf-8")) - except (OSError, json.JSONDecodeError) as error: - raise RuntimeError(f"cannot read model config {config_path}: {error}") from error - if not isinstance(config, dict): - raise RuntimeError(f"model config must be an object: {config_path}") - if config.get("model_type") != EXPECTED_MODEL_TYPE: - raise RuntimeError( - f"expected model_type={EXPECTED_MODEL_TYPE!r}, " - f"got {config.get('model_type')!r}" - ) - if int(config.get("vocab_size", 0)) != EXPECTED_VOCAB_SIZE: - raise RuntimeError( - f"expected vocab_size={EXPECTED_VOCAB_SIZE}, " - f"got {config.get('vocab_size')!r}" - ) - return config - - -def _is_quantized(config: dict[str, Any]) -> bool: - quantization = config.get("quantization_config") - return isinstance(quantization, dict) and bool(quantization) - - -def _run(args: argparse.Namespace) -> dict[str, Any]: - import infinicore - from infinilm.cache import PagedKVCacheConfig - from infinilm.distributed import DistConfig - from infinilm.infer_engine import InferEngine - from infinilm.modeling_utils import load_model_state_dict_by_file - - corpus = load_manifest(args.token_manifest) - model_config = _read_model_config(args.model) - if any(token >= EXPECTED_VOCAB_SIZE for token in corpus.token_ids): - raise RuntimeError("token manifest contains an ID outside the model vocabulary") - - available_targets = corpus.token_count - 1 - scored_token_count = ( - available_targets - if args.max_scored_tokens == 0 - else min(args.max_scored_tokens, available_targets) - ) - max_targets = None if args.max_scored_tokens == 0 else scored_token_count - windows = list( - iter_sliding_windows( - corpus.token_ids, - args.window, - args.stride, - max_targets, - ) - ) - if sum(window.scored_token_count for window in windows) != scored_token_count: - raise RuntimeError("sliding-window plan does not match scored token count") - - first_scored_token_index = 1 - last_scored_token_index_exclusive = 1 + scored_token_count - indices_sha256 = canonical_indices_sha256( - range(first_scored_token_index, last_scored_token_index_exclusive) - ) - precision = "W8A8" if _is_quantized(model_config) else "BF16" - config_payload = { - "backend": "infinilm", - "model": args.model, - "precision": precision, - "tp_size": args.tp_size, - "attention": args.attention, - "graph_enabled": False, - "window_size": args.window, - "stride": args.stride, - "scored_token_count": scored_token_count, - "scoring_method": SCORING_METHOD, - "corpus_manifest_sha256": corpus.manifest_sha256, - "corpus_token_ids_sha256": corpus.token_ids_sha256, - } - print( - "INFINILM_QWEN3_235B_PPL_CONFIG " - + json.dumps(config_payload, ensure_ascii=False, sort_keys=True), - flush=True, - ) - - device = infinicore.device("cuda", 0) - load_start = time.perf_counter() - model = InferEngine( - args.model, - device=device, - distributed_config=DistConfig(args.tp_size), - # This InfiniLM branch's flash-attn backend consumes the paged KV-cache - # layout while retaining flash-attn as the attention implementation. - cache_config=PagedKVCacheConfig( - num_blocks=(args.window + PAGED_KV_BLOCK_SIZE - 1) - // PAGED_KV_BLOCK_SIZE, - block_size=PAGED_KV_BLOCK_SIZE, - ), - enable_graph_compiling=False, - attention_backend=args.attention, - ) - if not hasattr(model, "score_nll"): - raise RuntimeError( - "installed InfiniLM lacks InferEngine.score_nll; rebuild the PPL scoring patch" - ) - load_model_state_dict_by_file(model, args.model, dtype=model.dtype) - model_load_seconds = time.perf_counter() - load_start - - window_nll_values: list[float] = [] - window_results: list[dict[str, Any]] = [] - infinicore.sync_device() - scoring_start = time.perf_counter() - for window in windows: - input_tokens = list(window.token_ids[:-1]) - label_tokens = list(window.token_ids[1:]) - if not input_tokens or len(input_tokens) != len(label_tokens): - raise RuntimeError(f"invalid causal shift in window {window.index}") - input_ids = infinicore.from_list( - [input_tokens], dtype=infinicore.int64 - ) - labels = infinicore.from_list( - [label_tokens], dtype=infinicore.int64 - ) - nll, returned_tokens = model.score_nll( - input_ids, - labels, - score_start=window.prediction_start, - ) - if returned_tokens != window.scored_token_count: - raise RuntimeError( - f"window {window.index} scored {returned_tokens} tokens, " - f"expected {window.scored_token_count}" - ) - if not math.isfinite(nll) or nll < 0: - raise RuntimeError(f"window {window.index} returned invalid NLL {nll}") - window_nll_values.append(nll) - window_results.append( - { - "index": window.index, - "context_start": window.token_start, - "target_start": window.score_start, - "target_end": window.score_end, - "input_token_count": len(window.token_ids), - "scored_token_count": returned_tokens, - "nll": nll, - } - ) - print( - f"PPL window {window.index + 1}/{len(windows)} " - f"tokens={returned_tokens} nll={nll:.6f}", - flush=True, - ) - - infinicore.sync_device() - scoring_seconds = time.perf_counter() - scoring_start - total_nll = math.fsum(window_nll_values) - mean_nll = total_nll / scored_token_count - try: - ppl = math.exp(mean_nll) - except OverflowError as error: - raise RuntimeError(f"PPL overflow at mean NLL={mean_nll}") from error - if not math.isfinite(ppl): - raise RuntimeError(f"PPL is not finite: {ppl}") - - result = { - "schema": RESULT_SCHEMA, - "status": "PASS", - "backend": "infinilm", - "model": args.model, - "precision": precision, - "tp_size": args.tp_size, - "attention": args.attention, - "graph_enabled": False, - "corpus_manifest": args.token_manifest, - "corpus_manifest_sha256": corpus.manifest_sha256, - "corpus_token_ids_sha256": corpus.token_ids_sha256, - "corpus_token_count": corpus.token_count, - "window_size": args.window, - "stride": args.stride, - "scoring_method": SCORING_METHOD, - "first_scored_token_index": first_scored_token_index, - "last_scored_token_index_exclusive": last_scored_token_index_exclusive, - "scored_token_indices_sha256": indices_sha256, - "scored_token_count": scored_token_count, - "total_nll": total_nll, - "mean_nll": mean_nll, - "ppl": ppl, - "windows": window_results, - "window_count": len(window_results), - "scoring_seconds": scoring_seconds, - "scored_tokens_per_second": scored_token_count / scoring_seconds, - "model_load_seconds": model_load_seconds, - "vocab_size": EXPECTED_VOCAB_SIZE, - } - if args.json_output: - _atomic_json(args.json_output, result) - print( - "INFINILM_QWEN3_235B_PPL_RESULT " - + json.dumps(result, ensure_ascii=False, sort_keys=True), - flush=True, - ) - print( - f"InfiniLM {precision} true PPL: {ppl:.6f} " - f"(mean NLL={mean_nll:.6f}, tokens={scored_token_count})", - flush=True, - ) - - del model - gc.collect() - infinicore.sync_device() - return result - - -def main(argv: Sequence[str] | None = None) -> int: - args = _parse_args(argv) - try: - # Validate the workload before checking or reserving GPUs. - load_manifest(args.token_manifest) - require_idle_gpu() - _run(args) - except BaseException as error: - completion = { - "schema": RESULT_SCHEMA, - "status": "ERROR", - "exit_code": 1, - "error": { - "type": type(error).__name__, - "message": str(error), - }, - } - print( - "INFINILM_QWEN3_235B_PPL_COMPLETE " - + json.dumps(completion, ensure_ascii=False, sort_keys=True), - flush=True, - ) - raise - print( - "INFINILM_QWEN3_235B_PPL_COMPLETE " - + json.dumps( - {"schema": RESULT_SCHEMA, "status": "PASS", "exit_code": 0}, - ensure_ascii=False, - sort_keys=True, - ), - flush=True, - ) - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py deleted file mode 100755 index b501e453..00000000 --- a/test/ppl/qwen3_235b/scripts/prepare_ppl_corpus_Qwen3_235B.py +++ /dev/null @@ -1,328 +0,0 @@ -#!/usr/bin/env python3 -"""将本地纯文本固化为 Qwen3_235B PPL 测试使用的 token manifest。""" - -from __future__ import annotations - -import argparse -import hashlib -import json -import operator -import os -import sys -import tempfile -from pathlib import Path -from typing import Any, Sequence - -from _ppl_common import ( - CORPUS_SCHEMA, - canonical_json_bytes, - canonical_token_ids_sha256, - write_token_ids_npy, -) - - -def _jsonable(value: Any) -> Any: - if value is None or isinstance(value, (bool, int, float, str)): - return value - if isinstance(value, dict): - return {str(key): _jsonable(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_jsonable(item) for item in value] - return str(value) - - -def _tokenizer_fingerprint(tokenizer: Any) -> tuple[str, str]: - backend = getattr(tokenizer, "backend_tokenizer", None) - if backend is not None and hasattr(backend, "to_str"): - try: - backend_payload: object = json.loads(backend.to_str()) - except (TypeError, ValueError, json.JSONDecodeError): - backend_payload = backend.to_str() - method = "backend_tokenizer+special_tokens/v1" - semantics = { - "backend_tokenizer": backend_payload, - "special_tokens_map": _jsonable( - getattr(tokenizer, "special_tokens_map", {}) - ), - } - else: - if not hasattr(tokenizer, "get_vocab"): - raise RuntimeError("tokenizer 既没有 backend_tokenizer，也没有 get_vocab()") - method = "vocab+init_kwargs+special_tokens/v1" - semantics = { - "vocab": tokenizer.get_vocab(), - "init_kwargs": _jsonable(getattr(tokenizer, "init_kwargs", {})), - "special_tokens_map": _jsonable( - getattr(tokenizer, "special_tokens_map", {}) - ), - } - return hashlib.sha256(canonical_json_bytes(semantics)).hexdigest(), method - - -def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--input", - nargs="+", - type=Path, - required=True, - help="本地 UTF-8 WikiText/raw text 文件，可按顺序提供多个文件", - ) - parser.add_argument( - "--tokenizer", - required=True, - help="本地 Qwen3_235B 模型或 tokenizer 目录", - ) - parser.add_argument("--output", required=True, type=Path, help="输出 JSON manifest") - parser.add_argument( - "--max-tokens", - type=int, - help="仅保留开头 N 个 token；省略时保留全部 token", - ) - parser.add_argument( - "--document-separator", - default="\n\n", - help=r"多个输入文件之间的分隔符，默认 '\n\n'", - ) - parser.add_argument( - "--storage", - choices=("inline", "npy"), - default="inline", - help="token IDs 内联到 JSON（默认），或写入相对路径 .npy 文件", - ) - parser.add_argument( - "--token-ids-file", - type=Path, - help="--storage=npy 时的相对路径；默认 .tokens.npy", - ) - parser.add_argument( - "--allow-download", - action="store_true", - help="允许 Transformers 访问网络；默认只读取本地文件/缓存", - ) - parser.add_argument( - "--trust-remote-code", - action="store_true", - help="传给 AutoTokenizer；Qwen3 官方 tokenizer 通常不需要", - ) - parser.add_argument("--overwrite", action="store_true", help="覆盖已有输出") - args = parser.parse_args(argv) - - if args.max_tokens is not None and args.max_tokens < 2: - parser.error("--max-tokens 必须至少为 2") - if args.token_ids_file is not None and args.storage != "npy": - parser.error("--token-ids-file 只能与 --storage=npy 一起使用") - if len({path.name for path in args.input}) != len(args.input): - parser.error("输入文件名不能重复，否则 manifest 无法稳定区分来源") - return args - - -def _read_sources( - paths: Sequence[Path], separator: str -) -> tuple[str, list[dict[str, object]]]: - documents: list[str] = [] - files: list[dict[str, object]] = [] - for path in paths: - if not path.is_file(): - raise FileNotFoundError(f"输入文件不存在：{path}") - raw = path.read_bytes() - try: - text = raw.decode("utf-8") - except UnicodeDecodeError as error: - raise ValueError(f"输入文件不是有效 UTF-8：{path}: {error}") from error - documents.append(text) - files.append( - { - "name": path.name, - "byte_count": len(raw), - "sha256": hashlib.sha256(raw).hexdigest(), - } - ) - return separator.join(documents), files - - -def _load_tokenizer(identifier: str, allow_download: bool, trust_remote_code: bool) -> Any: - try: - from transformers import AutoTokenizer - except ImportError as error: - raise RuntimeError("缺少 transformers，无法加载 tokenizer") from error - return AutoTokenizer.from_pretrained( - identifier, - local_files_only=not allow_download, - trust_remote_code=trust_remote_code, - ) - - -def _encode(tokenizer: Any, text: str) -> list[int]: - encoded = tokenizer( - text, - add_special_tokens=False, - truncation=False, - return_attention_mask=False, - return_token_type_ids=False, - ) - raw_ids = encoded["input_ids"] - if not isinstance(raw_ids, (list, tuple)) or ( - raw_ids and isinstance(raw_ids[0], (list, tuple)) - ): - raise RuntimeError("tokenizer 必须为单条文本返回一维 input_ids") - token_ids: list[int] = [] - for index, value in enumerate(raw_ids): - if isinstance(value, bool): - raise RuntimeError(f"input_ids[{index}] 不是有效整数") - try: - token = operator.index(value) - except TypeError as error: - raise RuntimeError(f"input_ids[{index}] 不是有效整数") from error - if token < 0: - raise RuntimeError(f"input_ids[{index}] 不是非负整数：{value!r}") - token_ids.append(token) - if len(token_ids) < 2: - raise RuntimeError("语料 token 数不足 2，无法计算 causal LM PPL") - return token_ids - - -def _atomic_json(path: Path, payload: object) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) - try: - with os.fdopen(fd, "w", encoding="utf-8") as handle: - json.dump(payload, handle, ensure_ascii=False, sort_keys=True, indent=2) - handle.write("\n") - handle.flush() - os.fsync(handle.fileno()) - os.replace(temporary, path) - except BaseException: - try: - os.unlink(temporary) - except FileNotFoundError: - pass - raise - - -def _npy_relative_path(output: Path, requested: Path | None) -> Path: - relative = requested or Path(f"{output.stem}.tokens.npy") - if relative.is_absolute() or ".." in relative.parts or relative.name in {"", "."}: - raise ValueError("--token-ids-file 必须是 manifest 目录内的安全相对路径") - if relative.suffix != ".npy": - raise ValueError("--token-ids-file 必须以 .npy 结尾") - return relative - - -def _write_npy(path: Path, token_ids: Sequence[int], overwrite: bool) -> None: - if path.exists() and not overwrite: - raise FileExistsError(f"token 文件已存在（可加 --overwrite）：{path}") - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp.npy") - try: - write_token_ids_npy(temporary, token_ids) - os.replace(temporary, path) - finally: - try: - temporary.unlink() - except FileNotFoundError: - pass - - -def build_manifest( - *, - source_text: str, - source_files: list[dict[str, object]], - separator: str, - tokenizer: Any, - tokenizer_label: str, - token_ids: list[int], - original_token_count: int, - max_tokens: int | None, - storage: str, - token_ids_file: str | None, -) -> dict[str, object]: - tokenizer_hash, fingerprint_method = _tokenizer_fingerprint(tokenizer) - token_hash = canonical_token_ids_sha256(token_ids) - payload: dict[str, object] = { - "schema": CORPUS_SCHEMA, - "source_sha256": hashlib.sha256(source_text.encode("utf-8")).hexdigest(), - "tokenizer_sha256": tokenizer_hash, - "token_count": len(token_ids), - "token_ids_sha256": token_hash, - "source": { - "encoding": "utf-8", - "document_separator": separator, - "file_count": len(source_files), - "files": source_files, - }, - "tokenizer": { - "name": Path(tokenizer_label.rstrip("/")).name or tokenizer_label, - "class": tokenizer.__class__.__name__, - "vocab_size": int(getattr(tokenizer, "vocab_size", 0)), - "fingerprint_method": fingerprint_method, - }, - "tokenization": { - "add_special_tokens": False, - "original_token_count": original_token_count, - "max_tokens": max_tokens, - "truncated": len(token_ids) != original_token_count, - }, - } - if storage == "inline": - payload["token_ids"] = token_ids - else: - if token_ids_file is None: - raise ValueError("npy storage 缺少 token_ids_file") - payload["token_ids_file"] = token_ids_file - payload["token_ids_dtype"] = "int64" - payload["manifest_sha256"] = hashlib.sha256( - canonical_json_bytes(payload) - ).hexdigest() - return payload - - -def main(argv: Sequence[str] | None = None) -> int: - args = _parse_args(argv) - try: - if args.output.exists() and not args.overwrite: - raise FileExistsError(f"输出已存在（可加 --overwrite）：{args.output}") - source_text, source_files = _read_sources(args.input, args.document_separator) - tokenizer = _load_tokenizer( - args.tokenizer, args.allow_download, args.trust_remote_code - ) - all_token_ids = _encode(tokenizer, source_text) - original_token_count = len(all_token_ids) - token_ids = ( - all_token_ids[: args.max_tokens] - if args.max_tokens is not None - else all_token_ids - ) - - relative_npy: Path | None = None - if args.storage == "npy": - relative_npy = _npy_relative_path(args.output, args.token_ids_file) - - manifest = build_manifest( - source_text=source_text, - source_files=source_files, - separator=args.document_separator, - tokenizer=tokenizer, - tokenizer_label=args.tokenizer, - token_ids=token_ids, - original_token_count=original_token_count, - max_tokens=args.max_tokens, - storage=args.storage, - token_ids_file=str(relative_npy) if relative_npy is not None else None, - ) - if relative_npy is not None: - _write_npy(args.output.parent / relative_npy, token_ids, args.overwrite) - _atomic_json(args.output, manifest) - except (FileNotFoundError, FileExistsError, RuntimeError, ValueError) as error: - print(f"错误：{error}", file=sys.stderr) - return 2 - - print(f"PPL 语料已固化：{args.output}") - print(f"Token 数：{manifest['token_count']}") - print(f"Token SHA256：{manifest['token_ids_sha256']}") - print(f"Manifest SHA256：{manifest['manifest_sha256']}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py b/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py deleted file mode 100755 index 027b5c87..00000000 --- a/test/ppl/qwen3_235b/scripts/transformers/_pytorch_runner.py +++ /dev/null @@ -1,1144 +0,0 @@ -#!/usr/bin/env python3 -"""Minimal Transformers TP benchmark runner for Qwen3_235B-A22B. - -The scenario wrappers are intentionally directly executable. When started as -``python wrapper.py`` this module replaces the process with a torchrun agent; -the original timeout therefore remains responsible for the agent, which in -turn terminates every worker on SIGTERM. -""" - -from __future__ import annotations - -import argparse -import gc -import hashlib -import importlib.metadata -import json -import os -import statistics -import subprocess -import sys -import time -from dataclasses import dataclass -from pathlib import Path -from typing import Any, Sequence - - -SCRIPT_ROOT = str(Path(__file__).resolve().parents[1]) -if SCRIPT_ROOT not in sys.path: - sys.path.insert(0, SCRIPT_ROOT) - -from _gpu_guard import require_idle_gpu as _require_idle_gpu - - -MODEL_NAME = "Qwen3_235B" -DEFAULT_MODEL = "/data1/Qwen3_235B" -DEFAULT_PROMPT_FILE = "examples/bench_prompt.md" -FALLBACK_PROMPT = """High-performance language-model inference processes a prompt in a prefill phase -and then produces one token per request during decode. Tensor-parallel ranks must -exchange identical partial results, while the key/value cache keeps each request -lane isolated. The benchmark uses deterministic prompt tokens and greedy decoding -so that every reported run has an exact, auditable token count.""" -MEASURED_INPUT_LENGTHS = 1 -REPEATS_PER_INPUT_LENGTH = 3 -MEASURED_ITERATIONS = REPEATS_PER_INPUT_LENGTH -MEASUREMENT_SEMANTICS = "one_fixed_shape_x_three_measurements" -SMOKE_OUTPUT_TOKENS = 64 -HYGON_TP_PLAN = { - "lm_head": "colwise_gather_output", - "model.layers.*.mlp.experts.gate_up_proj": "packed_colwise", - "model.layers.*.mlp.experts.down_proj": "rowwise", - "model.layers.*.mlp.experts": "moe_tp_experts", -} -EXPECTED_QWEN3_235B_ARCHITECTURE = { - "hidden_size": 4096, - "intermediate_size": 12288, - "head_dim": 128, - "num_attention_heads": 64, - "num_key_value_heads": 4, - "num_hidden_layers": 94, - "num_experts": 128, - "num_experts_per_tok": 8, - "moe_intermediate_size": 1536, - "vocab_size": 151936, -} - - -@dataclass(frozen=True) -class Scenario: - name: str - batch_size: int - input_tokens: int - output_tokens: int - - @property - def input_lengths(self) -> tuple[int]: - return (self.input_tokens,) - - @property - def total_context_tokens(self) -> int: - return self.input_tokens + self.output_tokens - - -def _parse_args(scenario: Scenario) -> argparse.Namespace: - parser = argparse.ArgumentParser( - description=( - "Transformers Qwen3_235B TP8 benchmark: " - f"batch={scenario.batch_size}, input={scenario.input_tokens}, " - f"output={scenario.output_tokens}, " - f"total={scenario.total_context_tokens} tokens" - ) - ) - parser.add_argument("--model", default=DEFAULT_MODEL) - parser.add_argument("--prompt-file", default=DEFAULT_PROMPT_FILE) - parser.add_argument("--output-tokens", type=int, default=scenario.output_tokens) - parser.add_argument("--tp-size", type=int, default=8) - parser.add_argument( - "--smoke", - action="store_true", - help=( - "load the full model but run only batch=1, input=16 and output=64 " - "to validate the TP/attention/cache path" - ), - ) - parser.add_argument( - "--attention", - choices=("eager",), - default="eager", - help="BW1100 correctness path; SDPA is unsupported for this model stack", - ) - args = parser.parse_args() - - if not Path(args.model).is_dir(): - parser.error(f"model directory does not exist: {args.model}") - if args.smoke: - args.output_tokens = SMOKE_OUTPUT_TOKENS - if args.output_tokens < 2: - parser.error("--output-tokens must be at least 2 to measure decode speed") - if args.tp_size < 1: - parser.error("--tp-size must be positive") - if len(set(scenario.input_lengths)) != MEASURED_INPUT_LENGTHS: - parser.error( - f"scenario must define exactly {MEASURED_INPUT_LENGTHS} lengths" - ) - return args - - -def _effective_scenario(scenario: Scenario, smoke: bool) -> Scenario: - if not smoke: - return scenario - return Scenario(f"{scenario.name}_smoke", 1, 16, SMOKE_OUTPUT_TOKENS) - - -def _validate_qwen3_235b_architecture(model_config: Any) -> dict[str, int]: - if getattr(model_config, "model_type", None) != "qwen3_moe": - raise RuntimeError( - "this benchmark requires model_type='qwen3_moe', got " - f"{getattr(model_config, 'model_type', None)!r}" - ) - actual: dict[str, int] = {} - mismatches: list[str] = [] - for field, expected in EXPECTED_QWEN3_235B_ARCHITECTURE.items(): - value = getattr(model_config, field, None) - try: - parsed = int(value) - except (TypeError, ValueError): - mismatches.append(f"{field}={value!r} (expected {expected})") - continue - actual[field] = parsed - if parsed != expected: - mismatches.append(f"{field}={parsed} (expected {expected})") - architectures = tuple(getattr(model_config, "architectures", None) or ()) - if "Qwen3MoeForCausalLM" not in architectures: - mismatches.append( - "architectures does not contain 'Qwen3MoeForCausalLM': " - f"{architectures!r}" - ) - if mismatches: - raise RuntimeError( - "checkpoint is not the expected Qwen3_235B-A22B architecture: " - + "; ".join(mismatches) - ) - return actual - - -def _build_qwen3_moe_tp_plan( - model_config: Any, - tp_size: int, - scenario: Scenario, - output_tokens: int, -) -> tuple[str | dict[str, str], dict[str, Any]]: - """Use the correctness-first TP8 layout validated on BW1100. - - Attention stays replicated. Qwen3_235B has four KV heads, so attempting to - tensor-parallelize attention over eight ranks produces an invalid local GQA - layout in this Transformers/DTK stack. Only MoE experts and the LM head are - sharded, matching the working Hygon container example. - """ - if getattr(model_config, "model_type", None) != "qwen3_moe": - raise RuntimeError( - "this benchmark only supports model_type='qwen3_moe', got " - f"{getattr(model_config, 'model_type', None)!r}" - ) - - global_query_heads = int(model_config.num_attention_heads) - global_kv_heads = int(model_config.num_key_value_heads) - head_dim = int(model_config.head_dim) - num_hidden_layers = int(model_config.num_hidden_layers) - if min( - global_query_heads, - global_kv_heads, - head_dim, - num_hidden_layers, - tp_size, - ) < 1: - raise RuntimeError("TP and attention dimensions must all be positive") - if global_query_heads % global_kv_heads: - raise RuntimeError( - f"global Q heads ({global_query_heads}) must be divisible by global " - f"KV heads ({global_kv_heads})" - ) - tp_plan = dict(HYGON_TP_PLAN) - maximum_sequence_tokens = scenario.input_tokens + output_tokens - dtype_bytes = 2 # BF16 K and V elements. - kv_cache_bytes_per_rank = ( - scenario.batch_size - * maximum_sequence_tokens - * num_hidden_layers - * 2 - * global_kv_heads - * head_dim - * dtype_bytes - ) - plan_payload = tp_plan if isinstance(tp_plan, dict) else {"mode": tp_plan} - metadata = { - "tp_plan_mode": f"qwen3_moe_tp{tp_size}_experts_lm_head_only", - "tp_plan_sha256": _stable_hash(plan_payload), - "attention_strategy": "replicated_eager", - "kv_projection_strategy": "replicated", - "kv_cache_replication_factor_across_tp_ranks": tp_size, - "global_query_heads": global_query_heads, - "global_kv_heads": global_kv_heads, - "head_dim": head_dim, - "num_hidden_layers": num_hidden_layers, - "local_query_heads": global_query_heads, - "local_kv_heads": global_kv_heads, - "local_gqa_groups": global_query_heads // global_kv_heads, - "maximum_sequence_tokens": maximum_sequence_tokens, - "estimated_dense_bf16_kv_cache_gib_per_rank": ( - kv_cache_bytes_per_rank / (1024**3) - ), - "kv_cache_estimate_excludes_allocator_and_cache_metadata": True, - } - return tp_plan, metadata - - -def _validate_and_set_local_gqa( - model: Any, tp_metadata: dict[str, Any] -) -> dict[str, Any]: - """Verify that attention stayed fully replicated on every TP rank.""" - base_model_prefix = getattr(model, "base_model_prefix", None) - base_model = getattr(model, base_model_prefix, None) - layers = getattr(base_model, "layers", None) - if layers is None: - raise RuntimeError( - f"cannot locate {base_model_prefix!r}.layers on loaded model" - ) - - head_dim = int(tp_metadata["head_dim"]) - expected_query_heads = int(tp_metadata["local_query_heads"]) - expected_kv_heads = int(tp_metadata["local_kv_heads"]) - expected_groups = int(tp_metadata["local_gqa_groups"]) - expected_query_width = expected_query_heads * head_dim - expected_kv_width = expected_kv_heads * head_dim - observed_groups: set[int] = set() - - for layer_index, layer in enumerate(layers): - attention = getattr(layer, "self_attn", None) - if attention is None: - raise RuntimeError(f"layer {layer_index} has no self_attn module") - query_width = int(attention.q_proj.out_features) - key_width = int(attention.k_proj.out_features) - value_width = int(attention.v_proj.out_features) - if query_width != expected_query_width: - raise RuntimeError( - f"layer {layer_index} local Q width={query_width}, expected " - f"{expected_query_width} ({expected_query_heads} heads)" - ) - if key_width != expected_kv_width or value_width != expected_kv_width: - raise RuntimeError( - f"layer {layer_index} local K/V widths={key_width}/{value_width}, " - f"expected {expected_kv_width} ({expected_kv_heads} heads)" - ) - observed_groups.add(int(attention.num_key_value_groups)) - - if observed_groups != {expected_groups}: - raise RuntimeError( - "attention GQA metadata changed despite replicated attention: " - f"observed={sorted(observed_groups)}, expected={expected_groups}" - ) - - expected_layers = int(tp_metadata["num_hidden_layers"]) - if len(layers) != expected_layers: - raise RuntimeError( - f"loaded model has {len(layers)} transformer layers, expected " - f"{expected_layers}" - ) - return { - "validated_attention_layers": len(layers), - "local_query_projection_width": expected_query_width, - "local_kv_projection_width": expected_kv_width, - "attention_replication_validated": True, - "local_gqa_groups": expected_groups, - } - - -def _launch_torchrun(args: argparse.Namespace) -> None: - env = os.environ.copy() - target_path = ( - "/root/.local/bin:/opt/dtk/cuda/cuda/bin:/opt/dtk/bin:/opt/dtk/hip/bin" - ) - target_library_path = ":".join( - ( - "/usr/local/lib/python3.10/dist-packages/torch/lib", - "/opt/dtk/dcc/gcvm/lib", - "/opt/dtk/hip/lib", - "/opt/dtk/llvm/lib", - "/opt/dtk/lib", - "/opt/dtk/lib64", - "/opt/hyhal/lib", - "/opt/hyhal/lib64", - "/opt/dtk/dushmem/lib", - "/opt/dtk/opencl/lib", - "/opt/ucx/lib", - "/opt/mpi/lib", - "/opt/hwloc/lib", - ) - ) - env["PATH"] = f"{target_path}:{env.get('PATH', '')}" - inherited_library_path = env.get("LD_LIBRARY_PATH", "") - env["LD_LIBRARY_PATH"] = ( - f"{target_library_path}:{inherited_library_path}" - if inherited_library_path - else target_library_path - ) - inherited_python_path = env.get("PYTHONPATH", "") - env["PYTHONPATH"] = ( - f"/usr/local:{inherited_python_path}" - if inherited_python_path - else "/usr/local" - ) - visible_devices = ",".join(str(index) for index in range(args.tp_size)) - env.setdefault("HIP_VISIBLE_DEVICES", visible_devices) - env.setdefault("CUDA_VISIBLE_DEVICES", visible_devices) - env.setdefault("OMP_NUM_THREADS", "1") - env.setdefault("TOKENIZERS_PARALLELISM", "false") - env.setdefault("PYTHONUNBUFFERED", "1") - env.setdefault("HSA_FORCE_FINE_GRAIN_PCIE", "1") - env.setdefault("NCCL_DEBUG", "WARN") - - script = str(Path(sys.argv[0]).resolve()) - command = [ - sys.executable, - "-m", - "torch.distributed.run", - "--standalone", - f"--nproc-per-node={args.tp_size}", - "--max-restarts=0", - "--monitor-interval=1", - script, - *sys.argv[1:], - ] - os.execvpe(sys.executable, command, env) - - -def _package_version(name: str) -> str | None: - try: - return importlib.metadata.version(name) - except importlib.metadata.PackageNotFoundError: - return None - - -def _install_hygon_grouped_mm_guard(torch: Any) -> bool: - """Install the grouped-MM fallback validated in the BW1100 image.""" - if getattr(torch.version, "hip", None) is None: - return False - - from transformers.integrations import moe as transformers_moe - - if getattr(transformers_moe, "_hygon_grouped_mm_guard_installed", False): - return True - - def grouped_mm(input_tensor: Any, weight: Any, offs: Any) -> Any: - ends = [int(value) for value in offs.detach().cpu().tolist()] - output = input_tensor.new_empty( - (input_tensor.shape[0], weight.shape[-1]), dtype=weight.dtype - ) - start = 0 - for expert_index, end in enumerate(ends): - if end > start: - output[start:end] = input_tensor[start:end].to(weight.dtype).matmul( - weight[expert_index] - ) - start = end - return output - - transformers_moe._grouped_mm = grouped_mm - transformers_moe._hygon_grouped_mm_guard_installed = True - return True - - -def _stable_hash(value: Any) -> str: - payload = json.dumps(value, ensure_ascii=True, separators=(",", ":")) - return hashlib.sha256(payload.encode("utf-8")).hexdigest() - - -def _emit(rank: int, tag: str, payload: dict[str, Any]) -> None: - if rank != 0: - return - if tag == "PYTORCH_QWEN3_235B_CONFIG": - print( - "[Transformers] " - f"model={payload['model']} batch={payload['batch_size']} " - f"input={payload['input_lengths'][0]} " - f"output={payload['output_tokens_per_request']} " - f"tp={payload['tp_size']} attention={payload['attention_implementation']}", - flush=True, - ) - print( - f" load weights over! {payload['model_load_seconds'] * 1000.0:.2f} ms ", - flush=True, - ) - elif tag == "PYTORCH_QWEN3_235B_COMPLETE": - print(f"Transformers benchmark status: {payload['status']}", flush=True) - - -def _print_infinilm_style_metrics( - rank: int, measurement: dict[str, Any], decoded_output: str -) -> None: - if rank != 0: - return - print( - f"\n Generation completed in {measurement['generation_seconds'] * 1000.0:.2f} ms", - flush=True, - ) - print( - f" Batchsize={measurement['batch_size']} " - f"Per_Batch_Input_Len={measurement['input_tokens_per_request']} " - f"Per_Batch_New_Tokens={measurement['output_tokens_per_request']}", - flush=True, - ) - print( - f"\n Prefill TTFT: {measurement['ttft_seconds'] * 1000.0:.2f} ms " - f"Throughput: {measurement['prefill_tokens_per_second']:.2f} tok/s", - flush=True, - ) - print( - f"\n Decode Avg ITL: {measurement['inter_token_latency_ms']:.2f} ms " - f"Throughput: {measurement['decode_tokens_per_second']:.2f} tok/s\n", - flush=True, - ) - print(decoded_output or "（未生成可显示文本）", flush=True) - - -def _validate_decoded_output(decoded_output: str) -> dict[str, Any]: - text = decoded_output.strip() - if not text: - raise RuntimeError("generated output decoded to an empty string") - if "\ufffd" in text: - raise RuntimeError("generated output contains Unicode replacement characters") - printable_characters = sum( - character.isprintable() or character in "\n\t" for character in text - ) - printable_ratio = printable_characters / len(text) - cjk_characters = sum("\u3400" <= character <= "\u9fff" for character in text) - url_fragments = text.lower().count("http") - if printable_ratio < 0.95: - raise RuntimeError( - f"generated output printable ratio is too low: {printable_ratio:.3f}" - ) - if cjk_characters < 8: - raise RuntimeError( - f"generated output is not a substantive Chinese response: CJK={cjk_characters}" - ) - if url_fragments > 2: - raise RuntimeError( - f"generated output contains suspicious URL fragments: {url_fragments}" - ) - return { - "nonempty_decoded_text": True, - "no_replacement_characters": True, - "printable_ratio": printable_ratio, - "cjk_character_count": cjk_characters, - "url_fragment_count": url_fragments, - } - - -def _max_across_ranks(value: float, torch: Any, dist: Any, device: Any) -> float: - tensor = torch.tensor(value, dtype=torch.float64, device=device) - dist.all_reduce(tensor, op=dist.ReduceOp.MAX) - return float(tensor.item()) - - -def _make_prompt_base( - tokenizer: Any, prompt_file: str -) -> tuple[list[int], dict[str, Any]]: - prompt_path = Path(prompt_file).resolve() - prompt_file_exists = prompt_path.is_file() - prompt_text = ( - prompt_path.read_text(encoding="utf-8") - if prompt_file_exists - else FALLBACK_PROMPT - ).strip() - if not prompt_text: - raise RuntimeError(f"benchmark prompt file is empty: {prompt_path}") - if not getattr(tokenizer, "chat_template", None): - raise RuntimeError("model tokenizer does not define a chat template") - rendered_prompt = tokenizer.apply_chat_template( - [{"role": "user", "content": prompt_text}], - tokenize=False, - add_generation_prompt=True, - ) - # Match InfiniLM's benchmark path: use tokenizer.encode defaults after the - # model's chat template, then repeat this exact base sequence to each length. - token_ids = list(tokenizer.encode(rendered_prompt)) - if not token_ids: - raise RuntimeError("the fixed benchmark prompt tokenized to an empty list") - return token_ids, { - "prompt_source": ( - "file_chat_template" if prompt_file_exists else "embedded_fallback" - ), - "prompt_file": str(prompt_path), - "prompt_file_exists": prompt_file_exists, - "prompt_file_sha256": hashlib.sha256( - prompt_text.encode("utf-8") - ).hexdigest(), - "rendered_prompt_sha256": hashlib.sha256( - rendered_prompt.encode("utf-8") - ).hexdigest(), - } - - -def _repeat_prompt(token_ids: Sequence[int], target_length: int) -> list[int]: - if target_length < 1: - raise ValueError("target_length must be positive") - base = list(token_ids) - if target_length <= len(base): - # Preserve the assistant-generation suffix instead of cutting it off. - result = base[-target_length:] - else: - prefix_length = target_length - len(base) - repeats = (prefix_length + len(base) - 1) // len(base) - result = (base * repeats)[:prefix_length] + base - if len(result) != target_length: - raise RuntimeError(f"expected {target_length} prompt tokens, got {len(result)}") - return result - - -def _materialize_logits(logits: Any) -> Any: - # A replicated TP lm_head returns Tensor. Keep this guard for TP plans that - # leave the vocabulary output as a DTensor. - if type(logits).__name__ == "DTensor" and hasattr(logits, "full_tensor"): - return logits.full_tensor() - return logits - - -def _forward( - model: Any, - logits_limit_argument: str, - input_ids: Any, - past_key_values: Any | None = None, -) -> tuple[Any, Any]: - kwargs: dict[str, Any] = { - "input_ids": input_ids, - "past_key_values": past_key_values, - "use_cache": True, - "return_dict": True, - logits_limit_argument: 1, - } - outputs = model(**kwargs) - if outputs.past_key_values is None: - raise RuntimeError("model did not return past_key_values with use_cache=True") - logits = _materialize_logits(outputs.logits) - if logits.ndim != 3 or logits.shape[0] != input_ids.shape[0]: - raise RuntimeError(f"unexpected logits shape: {tuple(logits.shape)}") - if logits.shape[1] != 1: - raise RuntimeError( - f"expected one retained logits position, got shape {tuple(logits.shape)}" - ) - return logits[:, -1, :], outputs.past_key_values - - -def _validate_output( - generated: Any, - last_logits: Any, - batch_size: int, - output_tokens: int, - vocab_size: int, - torch: Any, - dist: Any, -) -> tuple[Any, dict[str, Any]]: - expected_shape = (batch_size, output_tokens) - if tuple(generated.shape) != expected_shape: - raise RuntimeError( - f"expected generated shape {expected_shape}, got {tuple(generated.shape)}" - ) - - rank_minimum = generated.clone() - rank_maximum = generated.clone() - dist.all_reduce(rank_minimum, op=dist.ReduceOp.MIN) - dist.all_reduce(rank_maximum, op=dist.ReduceOp.MAX) - rank_consensus = bool(torch.equal(rank_minimum, rank_maximum)) - - valid_ids = bool( - torch.logical_and(generated >= 0, generated < vocab_size).all().item() - ) - finite_logits = bool(torch.isfinite(last_logits).all().item()) - checks = torch.tensor( - [int(rank_consensus), int(valid_ids), int(finite_logits)], - dtype=torch.int32, - device=generated.device, - ) - dist.all_reduce(checks, op=dist.ReduceOp.MIN) - rank_consensus, valid_ids, finite_logits = [bool(value) for value in checks.tolist()] - if not (rank_consensus and valid_ids and finite_logits): - raise RuntimeError( - "correctness validation failed: " - f"rank_consensus={rank_consensus}, valid_ids={valid_ids}, " - f"finite_logits={finite_logits}" - ) - - generated_cpu = generated.cpu() - matrix = generated_cpu.tolist() - return generated_cpu, { - "exact_output_shape": True, - "rank_consensus": rank_consensus, - "valid_token_ids": valid_ids, - "finite_last_logits": finite_logits, - "output_token_ids_sha256": _stable_hash(matrix), - "first_request_first_16_tokens": matrix[0][:16], - } - - -def _run_iteration( - model: Any, - prompt_base: Sequence[int], - batch_size: int, - input_tokens: int, - output_tokens: int, - vocab_size: int, - logits_limit_argument: str, - torch: Any, - dist: Any, - device: Any, -) -> dict[str, Any]: - prompt = _repeat_prompt(prompt_base, input_tokens) - input_ids = ( - torch.tensor(prompt, dtype=torch.long, device=device) - .unsqueeze(0) - .expand(batch_size, -1) - .contiguous() - ) - - torch.cuda.synchronize(device) - torch.cuda.reset_peak_memory_stats(device) - prefill_start = time.perf_counter() - logits, past_key_values = _forward( - model, logits_limit_argument, input_ids, past_key_values=None - ) - next_token = torch.argmax(logits, dim=-1) - torch.cuda.synchronize(device) - prefill_local_seconds = time.perf_counter() - prefill_start - prefill_seconds = _max_across_ranks( - prefill_local_seconds, torch, dist, device - ) - - generated_tokens = [next_token] - decode_start = time.perf_counter() - for _ in range(output_tokens - 1): - logits, past_key_values = _forward( - model, - logits_limit_argument, - next_token.unsqueeze(1), - past_key_values=past_key_values, - ) - next_token = torch.argmax(logits, dim=-1) - generated_tokens.append(next_token) - torch.cuda.synchronize(device) - decode_local_seconds = time.perf_counter() - decode_start - decode_seconds = _max_across_ranks(decode_local_seconds, torch, dist, device) - - peak_allocated_gib = _max_across_ranks( - torch.cuda.max_memory_allocated(device) / (1024**3), torch, dist, device - ) - peak_reserved_gib = _max_across_ranks( - torch.cuda.max_memory_reserved(device) / (1024**3), torch, dist, device - ) - generated = torch.stack(generated_tokens, dim=1) - generated_cpu, correctness = _validate_output( - generated, - logits, - batch_size, - output_tokens, - vocab_size, - torch, - dist, - ) - - total_prompt_tokens = batch_size * input_tokens - decode_token_count = batch_size * (output_tokens - 1) - generated_token_count = batch_size * output_tokens - total_seconds = prefill_seconds + decode_seconds - result = { - "batch_size": batch_size, - "input_tokens_per_request": input_tokens, - "prompt_token_ids_sha256": _stable_hash(prompt), - "output_tokens_per_request": output_tokens, - "total_context_tokens_per_request": input_tokens + output_tokens, - "total_prompt_tokens": total_prompt_tokens, - "total_generated_tokens": generated_token_count, - "ttft_seconds": prefill_seconds, - "prefill_tokens_per_second": total_prompt_tokens / prefill_seconds, - "decode_seconds": decode_seconds, - "decode_tokens_per_second": decode_token_count / decode_seconds, - "decode_tokens_per_second_per_request": ( - (output_tokens - 1) / decode_seconds - ), - "inter_token_latency_ms": decode_seconds * 1000.0 / (output_tokens - 1), - "generation_seconds": total_seconds, - "generated_tokens_per_second": generated_token_count / total_seconds, - "peak_memory_allocated_gib_max_rank": peak_allocated_gib, - "peak_memory_reserved_gib_max_rank": peak_reserved_gib, - "correctness": correctness, - "first_request_output_token_ids": generated_cpu[0].tolist(), - } - - del generated_cpu, generated, generated_tokens, logits, next_token - del past_key_values, input_ids - return result - - -def _median_summary(measurements: Sequence[dict[str, Any]]) -> dict[str, Any]: - fields = ( - "ttft_seconds", - "prefill_tokens_per_second", - "decode_seconds", - "decode_tokens_per_second", - "decode_tokens_per_second_per_request", - "inter_token_latency_ms", - "generation_seconds", - "generated_tokens_per_second", - "peak_memory_allocated_gib_max_rank", - "peak_memory_reserved_gib_max_rank", - ) - return { - f"median_{field}": statistics.median( - float(measurement[field]) for measurement in measurements - ) - for field in fields - } - - -def _per_length_medians( - measurements: Sequence[dict[str, Any]], input_lengths: Sequence[int] -) -> list[dict[str, Any]]: - summaries: list[dict[str, Any]] = [] - for input_tokens in input_lengths: - records = [ - measurement - for measurement in measurements - if int(measurement["input_tokens_per_request"]) == input_tokens - ] - if len(records) != REPEATS_PER_INPUT_LENGTH: - raise RuntimeError( - f"input length {input_tokens}: recorded {len(records)} repeats; " - f"expected {REPEATS_PER_INPUT_LENGTH}" - ) - summaries.append( - { - "input_tokens_per_request": input_tokens, - "measured_repeats": len(records), - **_median_summary(records), - } - ) - return summaries - - -def _overall_median_of_per_length_medians( - per_length_medians: Sequence[dict[str, Any]], -) -> dict[str, float]: - metric_names = [ - name for name in per_length_medians[0] if name.startswith("median_") - ] - return { - name: statistics.median( - float(length_summary[name]) for length_summary in per_length_medians - ) - for name in metric_names - } - - -def _run_worker_impl(args: argparse.Namespace, scenario: Scenario) -> int: - import inspect - - import torch - import torch.distributed as dist - import transformers - from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer - - grouped_mm_fallback = _install_hygon_grouped_mm_guard(torch) - - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", world_size)) - if rank != 0: - transformers.utils.logging.disable_progress_bar() - if world_size != args.tp_size: - raise RuntimeError( - f"expected WORLD_SIZE={args.tp_size}, got WORLD_SIZE={world_size}" - ) - if local_world_size != args.tp_size: - raise RuntimeError( - "this benchmark requires all TP ranks on one host: " - f"LOCAL_WORLD_SIZE={local_world_size}, TP={args.tp_size}" - ) - if not torch.cuda.is_available(): - raise RuntimeError("torch.cuda is unavailable") - if torch.cuda.device_count() < local_world_size: - raise RuntimeError( - f"need {local_world_size} visible GPUs, found {torch.cuda.device_count()}" - ) - - torch.cuda.set_device(local_rank) - device = torch.device("cuda", local_rank) - if not dist.is_initialized(): - dist.init_process_group(backend="nccl") - # Each torchrun worker owns one device. Avoid manual_seed_all(), which can - # make every worker initialize a context on all eight visible GPUs. - torch.random.default_generator.manual_seed(0) - torch.cuda.manual_seed(0) - - model_config = AutoConfig.from_pretrained( - args.model, - local_files_only=True, - trust_remote_code=False, - ) - architecture_signature = _validate_qwen3_235b_architecture(model_config) - quantization_config = getattr(model_config, "quantization_config", None) - if quantization_config: - raise RuntimeError( - "the Transformers benchmark is BF16-only; refusing quantized " - f"checkpoint {args.model!r} with quantization_config=" - f"{quantization_config!r}" - ) - tp_plan, tp_metadata = _build_qwen3_moe_tp_plan( - model_config, - args.tp_size, - scenario, - args.output_tokens, - ) - - load_start = time.perf_counter() - model = AutoModelForCausalLM.from_pretrained( - args.model, - config=model_config, - dtype=torch.bfloat16, - attn_implementation=args.attention, - tp_plan=tp_plan, - local_files_only=True, - low_cpu_mem_usage=True, - trust_remote_code=False, - ) - tp_validation = _validate_and_set_local_gqa(model, tp_metadata) - model.eval() - torch.cuda.synchronize(device) - if not dist.is_initialized(): - raise RuntimeError( - "Transformers tp_plan='auto' did not initialize torch.distributed" - ) - load_seconds = _max_across_ranks( - time.perf_counter() - load_start, torch, dist, device - ) - - tp_plan = getattr(model, "_tp_plan", None) - if not tp_plan: - raise RuntimeError("model loaded without a non-empty Transformers TP plan") - resolved_attention = getattr(model.config, "_attn_implementation", None) - if resolved_attention != args.attention: - raise RuntimeError( - f"requested attention={args.attention!r}, loaded model resolved " - f"attention={resolved_attention!r}" - ) - forward_parameters = inspect.signature(model.forward).parameters - if "logits_to_keep" in forward_parameters: - logits_limit_argument = "logits_to_keep" - elif "num_logits_to_keep" in forward_parameters: - logits_limit_argument = "num_logits_to_keep" - else: - raise RuntimeError( - "model.forward has no logits_to_keep argument; refusing to materialize " - "full [batch, context, vocab] logits for this benchmark" - ) - - tokenizer = AutoTokenizer.from_pretrained( - args.model, - local_files_only=True, - trust_remote_code=False, - use_fast=True, - ) - prompt_base, prompt_metadata = _make_prompt_base(tokenizer, args.prompt_file) - vocab_size = int(model.config.vocab_size) - if any(token < 0 or token >= vocab_size for token in prompt_base): - raise RuntimeError("fixed prompt contains a token outside model vocabulary") - - maximum_position_embeddings = int( - getattr(model.config, "max_position_embeddings", 0) or 0 - ) - maximum_requested = scenario.input_tokens + args.output_tokens - if maximum_position_embeddings and maximum_requested > maximum_position_embeddings: - raise RuntimeError( - f"requested sequence length {maximum_requested} exceeds " - f"max_position_embeddings={maximum_position_embeddings}" - ) - - config = { - "framework": "transformers", - "scenario": scenario.name, - "model_name": MODEL_NAME, - "model": str(Path(args.model).absolute()), - "model_realpath": str(Path(args.model).resolve()), - "model_class": type(model).__name__, - "dtype": "bfloat16", - "checkpoint_quantized": False, - "validated_qwen3_235b_architecture": architecture_signature, - "attention_implementation": resolved_attention, - "tp_plan": tp_metadata["tp_plan_mode"], - "tp_plan_rules": tp_plan, - "tp_plan_rule_count": len(tp_plan) if isinstance(tp_plan, dict) else None, - "tp_size": args.tp_size, - "smoke": args.smoke, - "batch_size": scenario.batch_size, - "input_lengths": list(scenario.input_lengths), - "output_tokens_per_request": args.output_tokens, - "total_context_tokens_per_request": maximum_requested, - "measured_iterations": MEASURED_ITERATIONS, - "measured_input_lengths": MEASURED_INPUT_LENGTHS, - "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, - "measurement_semantics": MEASUREMENT_SEMANTICS, - "model_load_seconds": load_seconds, - "fixed_prompt_base_tokens": len(prompt_base), - "fixed_prompt_base_sha256": _stable_hash(prompt_base), - **prompt_metadata, - "torch_version": torch.__version__, - "transformers_version": transformers.__version__, - "flash_attn_version": _package_version("flash-attn"), - "hygon_transformers_grouped_mm_fallback": grouped_mm_fallback, - "gpu_name": torch.cuda.get_device_name(device), - **tp_metadata, - **tp_validation, - } - _emit(rank, "PYTORCH_QWEN3_235B_CONFIG", config) - - measurements: list[dict[str, Any]] = [] - iteration = 0 - for length_index, input_tokens in enumerate(scenario.input_lengths, start=1): - with torch.inference_mode(): - shape_warmup = _run_iteration( - model, - prompt_base, - scenario.batch_size, - input_tokens, - args.output_tokens, - vocab_size, - logits_limit_argument, - torch, - dist, - device, - ) - shape_warmup_hash = shape_warmup["correctness"][ - "output_token_ids_sha256" - ] - shape_warmup_prompt_hash = shape_warmup["prompt_token_ids_sha256"] - _emit( - rank, - "PYTORCH_QWEN3_235B_SHAPE_WARMUP", - { - "scenario": scenario.name, - "length_index": length_index, - "batch_size": scenario.batch_size, - "input_tokens_per_request": input_tokens, - "prompt_token_ids_sha256": shape_warmup_prompt_hash, - "output_tokens_per_request": args.output_tokens, - "output_token_ids_sha256": shape_warmup_hash, - "correctness": shape_warmup["correctness"], - }, - ) - del shape_warmup - gc.collect() - - for repeat in range(1, REPEATS_PER_INPUT_LENGTH + 1): - iteration += 1 - with torch.inference_mode(): - measurement = _run_iteration( - model, - prompt_base, - scenario.batch_size, - input_tokens, - args.output_tokens, - vocab_size, - logits_limit_argument, - torch, - dist, - device, - ) - measured_hash = measurement["correctness"][ - "output_token_ids_sha256" - ] - measured_prompt_hash = measurement["prompt_token_ids_sha256"] - if measured_prompt_hash != shape_warmup_prompt_hash: - raise RuntimeError( - f"input length {input_tokens} repeat {repeat}: measured prompt " - f"hash {measured_prompt_hash} does not match exact-shape " - f"warmup prompt hash {shape_warmup_prompt_hash}" - ) - # Hygon BF16 kernels can make numerically valid MoE routing choices - # differ across independent runs. Keep the replay hash observable, - # while treating per-run shape/range/finite/rank checks as correctness. - measurement["correctness"]["output_matches_exact_shape_warmup"] = ( - measured_hash == shape_warmup_hash - ) - measurement["exact_shape_warmup_output_sha256"] = shape_warmup_hash - measurement = { - "scenario": scenario.name, - "iteration": iteration, - "length_index": length_index, - "repeat": repeat, - **measurement, - } - output_token_ids = measurement.pop("first_request_output_token_ids") - decoded_output = tokenizer.decode( - output_token_ids, skip_special_tokens=True - ).strip() - measurement["correctness"]["decoded_output"] = ( - _validate_decoded_output(decoded_output) - ) - measurements.append(measurement) - _print_infinilm_style_metrics(rank, measurement, decoded_output) - _emit(rank, "PYTORCH_QWEN3_235B_ITERATION", measurement) - gc.collect() - - if len(measurements) != MEASURED_ITERATIONS: - raise RuntimeError( - f"expected {MEASURED_ITERATIONS} measurements, got {len(measurements)}" - ) - per_length_medians = _per_length_medians( - measurements, scenario.input_lengths - ) - overall_medians = _overall_median_of_per_length_medians(per_length_medians) - summary = { - "scenario": scenario.name, - "status": "PASS", - "measured_iterations": len(measurements), - "measured_input_lengths": MEASURED_INPUT_LENGTHS, - "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, - "measurement_semantics": MEASUREMENT_SEMANTICS, - "input_lengths": list(scenario.input_lengths), - "batch_size": scenario.batch_size, - "output_tokens_per_request": args.output_tokens, - "total_context_tokens_per_request": maximum_requested, - "per_length_medians": per_length_medians, - "overall_aggregate": { - "aggregation_method": "median_of_three_fixed_shape_measurements", - "measurement_count": len(measurements), - "mixed_input_lengths": False, - **overall_medians, - }, - # Compatibility aliases for existing table consumers. Their scope is the - # explicitly labeled overall aggregate above, not a single input length. - **overall_medians, - "output_token_ids_sha256": [ - item["correctness"]["output_token_ids_sha256"] for item in measurements - ], - } - _emit(rank, "PYTORCH_QWEN3_235B_SUMMARY", summary) - return len(measurements) - - -def _run_worker(args: argparse.Namespace, scenario: Scenario) -> None: - import torch.distributed as dist - - rank = int(os.environ["RANK"]) - measured_iterations = 0 - caught: BaseException | None = None - caught_traceback: Any = None - teardown_errors: list[str] = [] - process_group_was_initialized = False - try: - measured_iterations = _run_worker_impl(args, scenario) - except BaseException as error: - caught = error - caught_traceback = error.__traceback__ - finally: - process_group_was_initialized = dist.is_initialized() - if process_group_was_initialized: - if caught is None: - try: - dist.barrier() - except BaseException as error: - caught = error - caught_traceback = error.__traceback__ - teardown_errors.append( - f"barrier: {type(error).__name__}: {error}" - ) - try: - dist.destroy_process_group() - except BaseException as error: - if caught is None: - caught = error - caught_traceback = error.__traceback__ - teardown_errors.append( - f"destroy_process_group: {type(error).__name__}: {error}" - ) - - teardown_complete = not dist.is_initialized() and not teardown_errors - status = ( - "PASS" - if caught is None - and measured_iterations == MEASURED_ITERATIONS - and teardown_complete - else "ERROR" - ) - completion: dict[str, Any] = { - "scenario": scenario.name, - "status": status, - "exit_code": 0 if status == "PASS" else 1, - "measured_iterations": measured_iterations, - "measured_input_lengths": MEASURED_INPUT_LENGTHS, - "repeats_per_input_length": REPEATS_PER_INPUT_LENGTH, - "measurement_semantics": MEASUREMENT_SEMANTICS, - "process_group_was_initialized": process_group_was_initialized, - "distributed_teardown_complete": teardown_complete, - } - if caught is not None: - completion["error"] = { - "type": type(caught).__name__, - "message": str(caught), - } - if teardown_errors: - completion["teardown_errors"] = teardown_errors - _emit(rank, "PYTORCH_QWEN3_235B_COMPLETE", completion) - - if caught is not None: - raise caught.with_traceback(caught_traceback) - - -def main(scenario: Scenario) -> None: - args = _parse_args(scenario) - scenario = _effective_scenario(scenario, args.smoke) - world_size = int(os.environ.get("WORLD_SIZE", "1")) - if "LOCAL_RANK" not in os.environ and world_size == 1: - _require_idle_gpu() - _launch_torchrun(args) - raise AssertionError("os.execvpe returned unexpectedly") - _run_worker(args, scenario) diff --git a/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py b/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py deleted file mode 100755 index 5474ca99..00000000 --- a/test/ppl/qwen3_235b/scripts/transformers/pytorch_ppl_Qwen3_235B.py +++ /dev/null @@ -1,491 +0,0 @@ -#!/usr/bin/env python3 -"""Calculate true token-level PPL for Qwen3_235B with Transformers TP8. - -The input is a framework-neutral token manifest. Both the Transformers and -InfiniLM runners must consume the same manifest so their PPL values score the -same target tokens instead of independently tokenizing the source corpus. -""" - -from __future__ import annotations - -import argparse -import gc -import inspect -import json -import math -import os -import sys -import time -from pathlib import Path -from typing import Any - - -SCRIPT_DIR = Path(__file__).resolve().parent -SCRIPTS_DIR = SCRIPT_DIR.parent -for import_path in (SCRIPT_DIR, SCRIPTS_DIR): - if str(import_path) not in sys.path: - sys.path.insert(0, str(import_path)) - -import _pytorch_runner as benchmark_runner -from _ppl_common import ( - SCORING_METHOD, - canonical_indices_sha256, - iter_sliding_windows, - load_manifest, -) - - -RESULT_SCHEMA = "qwen3_235b_true_ppl_result/v1" -DEFAULT_MODEL = "/data1/Qwen3_235B" -DEFAULT_WINDOW_SIZE = 256 -DEFAULT_STRIDE = 128 -DEFAULT_MAX_SCORED_TOKENS = 10240 -EXPECTED_VOCAB_SIZE = 151936 - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="True shifted-token PPL for Qwen3_235B BF16 on Hygon TP8" - ) - parser.add_argument("--model", default=DEFAULT_MODEL) - parser.add_argument("--token-manifest", required=True) - parser.add_argument("--window", type=int, default=DEFAULT_WINDOW_SIZE) - parser.add_argument("--stride", type=int, default=DEFAULT_STRIDE) - parser.add_argument( - "--max-scored-tokens", - type=int, - default=DEFAULT_MAX_SCORED_TOKENS, - help="maximum shifted target tokens to score; 0 scores the full manifest", - ) - parser.add_argument("--tp-size", type=int, default=8) - parser.add_argument( - "--attention", - choices=("eager",), - default="eager", - help="BW1100 correctness path; SDPA is unsupported for this model stack", - ) - parser.add_argument( - "--json-output", - help="optional rank-0 result path; the result is always printed as JSON", - ) - args = parser.parse_args() - - model_path = Path(args.model) - manifest_path = Path(args.token_manifest) - if not model_path.is_dir(): - parser.error(f"model directory does not exist: {model_path}") - if not manifest_path.is_file(): - parser.error(f"token manifest does not exist: {manifest_path}") - if args.window < 2: - parser.error("--window must be at least 2") - if args.stride < 1 or args.stride >= args.window: - parser.error("--stride must satisfy 1 <= stride < window") - if args.max_scored_tokens < 0: - parser.error("--max-scored-tokens must be non-negative") - if args.tp_size < 1: - parser.error("--tp-size must be positive") - - args.model = str(model_path.resolve()) - args.token_manifest = str(manifest_path.resolve()) - if args.json_output: - args.json_output = str(Path(args.json_output).resolve()) - return args - -def _write_json_atomic(path_value: str, payload: dict[str, Any]) -> None: - path = Path(path_value) - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(f".{path.name}.tmp.{os.getpid()}") - temporary.write_text( - json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - temporary.replace(path) - - -def _run_worker_impl(args: argparse.Namespace) -> dict[str, Any]: - import torch - import torch.distributed as dist - import torch.nn.functional as functional - import transformers - from transformers import AutoConfig, AutoModelForCausalLM - - rank = int(os.environ["RANK"]) - local_rank = int(os.environ["LOCAL_RANK"]) - world_size = int(os.environ["WORLD_SIZE"]) - local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", world_size)) - if world_size != args.tp_size or local_world_size != args.tp_size: - raise RuntimeError( - "PPL runner requires one-host TP with WORLD_SIZE=" - f"LOCAL_WORLD_SIZE={args.tp_size}; got {world_size}/{local_world_size}" - ) - if not torch.cuda.is_available() or torch.cuda.device_count() < local_world_size: - raise RuntimeError( - f"need {local_world_size} visible GPUs, found {torch.cuda.device_count()}" - ) - - torch.cuda.set_device(local_rank) - device = torch.device("cuda", local_rank) - if not dist.is_initialized(): - dist.init_process_group(backend="nccl") - # RCCL initializes lazily. Reserve communicator memory before the 235B - # checkpoint consumes nearly all device memory. - communicator_probe = torch.ones(1, dtype=torch.int32, device=device) - dist.all_reduce(communicator_probe) - if int(communicator_probe.item()) != args.tp_size: - raise RuntimeError("Transformers TP8 RCCL communicator probe failed") - dist.barrier() - torch.cuda.synchronize(device) - del communicator_probe - torch.random.default_generator.manual_seed(0) - torch.cuda.manual_seed(0) - if rank != 0: - transformers.utils.logging.disable_progress_bar() - - corpus_manifest = load_manifest(args.token_manifest) - token_ids = corpus_manifest.token_ids - corpus = { - "manifest_path": str(corpus_manifest.path.resolve()), - "manifest_sha256": corpus_manifest.manifest_sha256, - "token_ids_sha256": corpus_manifest.token_ids_sha256, - "token_count": corpus_manifest.token_count, - "source_sha256": corpus_manifest.payload["source_sha256"], - "tokenizer_sha256": corpus_manifest.payload["tokenizer_sha256"], - "source": corpus_manifest.payload.get("source"), - "tokenizer": corpus_manifest.payload.get("tokenizer"), - } - model_config = AutoConfig.from_pretrained( - args.model, local_files_only=True, trust_remote_code=False - ) - architecture = benchmark_runner._validate_qwen3_235b_architecture(model_config) - if getattr(model_config, "quantization_config", None): - raise RuntimeError("Transformers PPL baseline requires the BF16 checkpoint") - vocab_size = int(model_config.vocab_size) - if vocab_size != EXPECTED_VOCAB_SIZE: - raise RuntimeError( - f"expected complete Qwen3_235B vocabulary {EXPECTED_VOCAB_SIZE}, " - f"got {vocab_size}" - ) - if any(token >= vocab_size for token in token_ids): - raise RuntimeError("token manifest contains an ID outside the model vocabulary") - maximum_positions = int( - getattr(model_config, "max_position_embeddings", 0) or 0 - ) - if maximum_positions and args.window > maximum_positions: - raise RuntimeError( - f"window={args.window} exceeds max_position_embeddings={maximum_positions}" - ) - - scoring_scenario = benchmark_runner.Scenario("true_ppl", 1, args.window, 1) - tp_plan, tp_metadata = benchmark_runner._build_qwen3_moe_tp_plan( - model_config, args.tp_size, scoring_scenario, 1 - ) - grouped_mm_fallback = benchmark_runner._install_hygon_grouped_mm_guard(torch) - load_start = time.perf_counter() - model = AutoModelForCausalLM.from_pretrained( - args.model, - config=model_config, - dtype=torch.bfloat16, - attn_implementation=args.attention, - tp_plan=tp_plan, - local_files_only=True, - low_cpu_mem_usage=True, - trust_remote_code=False, - ) - tp_validation = benchmark_runner._validate_and_set_local_gqa(model, tp_metadata) - model.eval() - torch.cuda.synchronize(device) - load_seconds = benchmark_runner._max_across_ranks( - time.perf_counter() - load_start, torch, dist, device - ) - - loaded_tp_plan = getattr(model, "_tp_plan", None) - if not loaded_tp_plan: - raise RuntimeError("model loaded without a non-empty Transformers TP plan") - resolved_attention = getattr(model.config, "_attn_implementation", None) - if resolved_attention != args.attention: - raise RuntimeError( - f"requested attention={args.attention!r}, resolved={resolved_attention!r}" - ) - forward_parameters = inspect.signature(model.forward).parameters - if "logits_to_keep" in forward_parameters: - logits_limit_argument = "logits_to_keep" - elif "num_logits_to_keep" in forward_parameters: - logits_limit_argument = "num_logits_to_keep" - else: - raise RuntimeError( - "model.forward has no logits_to_keep argument; refusing full-context " - "vocabulary materialization" - ) - - available_targets = len(token_ids) - 1 - scored_token_count = ( - available_targets - if args.max_scored_tokens == 0 - else min(args.max_scored_tokens, available_targets) - ) - if scored_token_count < 1: - raise RuntimeError("the selected corpus range contains no shifted target token") - first_scored_token_index = 1 - last_scored_token_index_exclusive = 1 + scored_token_count - scored_token_indices_sha256 = canonical_indices_sha256( - range(first_scored_token_index, last_scored_token_index_exclusive) - ) - scoring_method = SCORING_METHOD - - config = { - "backend": "transformers", - "model": args.model, - "dtype": "bfloat16", - "tp_size": args.tp_size, - "attention": resolved_attention, - "window_size": args.window, - "stride": args.stride, - "requested_max_scored_tokens": args.max_scored_tokens, - "scored_token_count": scored_token_count, - "scoring_method": scoring_method, - "first_scored_token_index": first_scored_token_index, - "last_scored_token_index_exclusive": last_scored_token_index_exclusive, - "scored_token_indices_sha256": scored_token_indices_sha256, - "corpus_manifest": corpus, - "vocab_size": vocab_size, - "architecture": architecture, - "model_load_seconds": load_seconds, - "tp_plan": tp_metadata, - "tp_validation": tp_validation, - "hygon_transformers_grouped_mm_fallback": grouped_mm_fallback, - "torch_version": torch.__version__, - "transformers_version": transformers.__version__, - } - if rank == 0: - print( - "PYTORCH_QWEN3_235B_PPL_CONFIG " - + json.dumps(config, ensure_ascii=False, sort_keys=True), - flush=True, - ) - - total_nll = 0.0 - windows: list[dict[str, Any]] = [] - scored_by_windows = 0 - torch.cuda.synchronize(device) - scoring_start = time.perf_counter() - with torch.inference_mode(): - for window in iter_sliding_windows( - token_ids, - args.window, - args.stride, - None if args.max_scored_tokens == 0 else args.max_scored_tokens, - ): - display_index = window.index + 1 - target_count = window.scored_token_count - input_slice = window.token_ids - expected_prediction_start = len(input_slice) - target_count - 1 - if ( - window.prediction_start != expected_prediction_start - or window.prediction_end != len(input_slice) - 1 - ): - raise RuntimeError( - f"window {display_index} retained-logits alignment is invalid" - ) - input_ids = torch.tensor( - input_slice, dtype=torch.long, device=device - ).unsqueeze(0) - outputs = model( - input_ids=input_ids, - use_cache=False, - return_dict=True, - **{logits_limit_argument: target_count + 1}, - ) - logits = benchmark_runner._materialize_logits(outputs.logits) - expected_shape = (1, target_count + 1, vocab_size) - if tuple(logits.shape) != expected_shape: - raise RuntimeError( - "incomplete or unexpected logits: " - f"got {tuple(logits.shape)}, expected {expected_shape}" - ) - score_logits = logits[:, :-1, :] - labels = torch.tensor( - input_slice[window.target_start : window.target_end], - dtype=torch.long, - device=device, - ).unsqueeze(0) - finite = torch.isfinite(score_logits).all().to(dtype=torch.int32) - dist.all_reduce(finite, op=dist.ReduceOp.MIN) - if not bool(finite.item()): - raise RuntimeError( - f"window {display_index} contains non-finite logits" - ) - window_nll_tensor = functional.cross_entropy( - score_logits.float().reshape(-1, vocab_size), - labels.reshape(-1), - reduction="sum", - ) - window_nll = float(window_nll_tensor.double().item()) - nll_min = torch.tensor(window_nll, dtype=torch.float64, device=device) - nll_max = nll_min.clone() - dist.all_reduce(nll_min, op=dist.ReduceOp.MIN) - dist.all_reduce(nll_max, op=dist.ReduceOp.MAX) - rank_delta_per_token = float((nll_max - nll_min).item()) / target_count - if rank_delta_per_token > 1e-4: - raise RuntimeError( - f"window {display_index} rank NLL mismatch: " - f"delta/token={rank_delta_per_token:.6g}" - ) - total_nll += window_nll - scored_by_windows += target_count - windows.append( - { - "index": window.index, - "token_start": window.token_start, - "token_end": window.token_end, - "score_start": window.score_start, - "score_end": window.score_end, - "input_token_count": len(input_slice), - "scored_token_count": target_count, - "nll": window_nll, - } - ) - del outputs, logits, score_logits, labels, window_nll_tensor, input_ids - - torch.cuda.synchronize(device) - scoring_seconds = benchmark_runner._max_across_ranks( - time.perf_counter() - scoring_start, torch, dist, device - ) - if scored_by_windows != scored_token_count: - raise RuntimeError( - f"scored {scored_by_windows} tokens, expected {scored_token_count}" - ) - - mean_nll = total_nll / scored_token_count - if not math.isfinite(mean_nll): - raise RuntimeError(f"mean NLL is not finite: {mean_nll}") - try: - ppl = math.exp(mean_nll) - except OverflowError as error: - raise RuntimeError(f"PPL overflows float64 at mean NLL={mean_nll}") from error - if not math.isfinite(ppl): - raise RuntimeError(f"PPL is not finite: {ppl}") - - result: dict[str, Any] = {} - if rank == 0: - result = { - "schema": RESULT_SCHEMA, - "status": "PASS", - "backend": "transformers", - "model": args.model, - "dtype": "bfloat16", - "tp_size": args.tp_size, - "attention": resolved_attention, - "corpus_manifest": args.token_manifest, - "corpus_manifest_sha256": corpus["manifest_sha256"], - "corpus_token_ids_sha256": corpus["token_ids_sha256"], - "corpus_token_count": corpus["token_count"], - "window_size": args.window, - "stride": args.stride, - "scoring_method": scoring_method, - "first_scored_token_index": first_scored_token_index, - "last_scored_token_index_exclusive": ( - last_scored_token_index_exclusive - ), - "scored_token_indices_sha256": scored_token_indices_sha256, - "scored_token_count": scored_token_count, - "total_nll": total_nll, - "mean_nll": mean_nll, - "ppl": ppl, - "windows": windows, - "window_count": len(windows), - "scoring_seconds": scoring_seconds, - "scored_tokens_per_second": scored_token_count / scoring_seconds, - "model_load_seconds": load_seconds, - "vocab_size": vocab_size, - "full_vocab_logits_validated_every_window": True, - } - if args.json_output: - _write_json_atomic(args.json_output, result) - print( - "PYTORCH_QWEN3_235B_PPL_RESULT " - + json.dumps(result, ensure_ascii=False, sort_keys=True), - flush=True, - ) - print( - f"Transformers true PPL: {ppl:.6f} " - f"(mean NLL={mean_nll:.6f}, tokens={scored_token_count})", - flush=True, - ) - - del model - gc.collect() - torch.cuda.empty_cache() - return result - - -def _run_worker(args: argparse.Namespace) -> int: - import torch.distributed as dist - - rank = int(os.environ["RANK"]) - caught: BaseException | None = None - caught_traceback: Any = None - teardown_errors: list[str] = [] - try: - _run_worker_impl(args) - except BaseException as error: - caught = error - caught_traceback = error.__traceback__ - finally: - initialized = dist.is_initialized() - if initialized: - if caught is None: - try: - dist.barrier() - except BaseException as error: - caught = error - caught_traceback = error.__traceback__ - teardown_errors.append( - f"barrier: {type(error).__name__}: {error}" - ) - try: - dist.destroy_process_group() - except BaseException as error: - if caught is None: - caught = error - caught_traceback = error.__traceback__ - teardown_errors.append( - f"destroy_process_group: {type(error).__name__}: {error}" - ) - status = "PASS" if caught is None and not teardown_errors else "ERROR" - if rank == 0: - completion: dict[str, Any] = { - "schema": RESULT_SCHEMA, - "status": status, - "exit_code": 0 if status == "PASS" else 1, - "distributed_teardown_complete": not dist.is_initialized(), - } - if caught is not None: - completion["error"] = { - "type": type(caught).__name__, - "message": str(caught), - } - if teardown_errors: - completion["teardown_errors"] = teardown_errors - print( - "PYTORCH_QWEN3_235B_PPL_COMPLETE " - + json.dumps(completion, ensure_ascii=False, sort_keys=True), - flush=True, - ) - if caught is not None: - raise caught.with_traceback(caught_traceback) - return 0 - - -def main() -> int: - args = _parse_args() - world_size = int(os.environ.get("WORLD_SIZE", "1")) - if "LOCAL_RANK" not in os.environ and world_size == 1: - # Fail on corpus/schema errors before reserving all eight devices. - load_manifest(args.token_manifest) - benchmark_runner._require_idle_gpu() - benchmark_runner._launch_torchrun(args) - raise AssertionError("os.execvpe returned unexpectedly") - return _run_worker(args) - - -if __name__ == "__main__": - raise SystemExit(main())