From 81cc44574038a3b391005e4529fb588bbbecb6c7 Mon Sep 17 00:00:00 2001 From: bolunz Date: Thu, 5 Nov 2026 11:36:07 +0800 Subject: [PATCH] feat: introduce suite-based training regression framework --- .gitignore | 4 + CMakeLists.txt | 3 + docs/test_usage_guide.md | 33 +- example/common/regression_output.cc | 388 +++++ example/common/regression_output.h | 49 + example/gpt2/checkpoint_loader.cc | 1 + example/gpt2/main.cc | 28 +- example/llama3/checkpoint_loader.cc | 1 + example/llama3/main.cc | 28 +- example/mixtral/main.cc | 24 + example/qwen3/checkpoint_loader.cc | 1 + example/qwen3/main.cc | 28 +- scripts/compare_loss.py | 212 +-- scripts/compare_tps.py | 216 +-- scripts/compare_utils.py | 57 + scripts/run_all_models.sh | 5 + scripts/run_all_tests.sh | 4 + scripts/run_models_and_profile.bash | 648 +------- tests/regression/README.md | 241 +++ tests/regression/TOLERANCES.md | 78 + tests/regression/assets.py | 92 ++ tests/regression/checks.py | 1134 +++++++++++++ tests/regression/environments/cuda.json | 17 + tests/regression/execution.py | 458 ++++++ tests/regression/fixtures/README.md | 113 ++ tests/regression/fixtures/classic.json | 39 + tests/regression/fixtures/generate.py | 206 +++ tests/regression/fixtures/manifest.json | 55 + tests/regression/fixtures/reference.py | 461 ++++++ tests/regression/history.py | 479 ++++++ tests/regression/models/gpt2.json | 204 +++ tests/regression/models/llama3.json | 201 +++ tests/regression/models/mixtral.json | 113 ++ tests/regression/models/qwen3.json | 201 +++ tests/regression/planning.py | 1414 +++++++++++++++++ tests/regression/profiling.py | 81 + tests/regression/quarantines.json | 4 + tests/regression/reporting.py | 190 +++ tests/regression/reproduction.py | 96 ++ tests/regression/runner.py | 1201 ++++++++++++++ .../regression/scenarios/classic_8_proc.json | 45 + tests/regression/scenarios/classic_basic.json | 57 + .../scenarios/classic_checkpoint.json | 49 + tests/regression/scenarios/classic_lora.json | 73 + .../scenarios/classic_lr_scheduler.json | 65 + tests/regression/scenarios/classic_moe.json | 163 ++ tests/regression/scenarios/classic_qwen3.json | 141 ++ .../scenarios/classic_qwen3_8_proc.json | 141 ++ .../scenarios/classic_recompute.json | 250 +++ .../scenarios/classic_recompute_qwen3.json | 270 ++++ tests/regression/scenarios/classic_zero.json | 53 + tests/regression/scenarios/kernel_basic.json | 13 + tests/regression/scenarios/kernel_ddp.json | 14 + .../regression/scenarios/long_stability.json | 44 + .../scenarios/tiny_accumulation.json | 40 + .../regression/scenarios/tiny_checkpoint.json | 52 + tests/regression/scenarios/tiny_lora.json | 52 + tests/regression/scenarios/tiny_oracle.json | 39 + tests/regression/scenarios/tiny_parallel.json | 108 ++ .../scenarios/tiny_parallel_anchors.json | 52 + .../regression/scenarios/tiny_recompute.json | 54 + .../regression/scenarios/tiny_scheduler.json | 49 + tests/regression/scenarios/tiny_smoke.json | 36 + tests/regression/scenarios/tiny_zero.json | 48 + tests/regression/schema/README.md | 196 +++ tests/regression/suites/nightly.json | 8 + tests/regression/suites/perf.json | 9 + tests/regression/suites/pr.json | 8 + tests/regression/suites/release.json | 8 + tests/regression/suites/smoke.json | 8 + tests/regression/test_checks.py | 743 +++++++++ tests/regression/tests/__init__.py | 0 tests/regression/tests/test_compatibility.py | 110 ++ tests/regression/tests/test_execution.py | 280 ++++ tests/regression/tests/test_fixtures.py | 64 + tests/regression/tests/test_history.py | 488 ++++++ tests/regression/tests/test_log_mode.py | 122 ++ .../tests/test_oracle_integration.py | 109 ++ tests/regression/tests/test_planning.py | 919 +++++++++++ tests/regression/tests/test_profiling.py | 128 ++ tests/regression/tests/test_recompute.py | 83 + .../tests/test_reference_batching.py | 57 + tests/regression/tests/test_reproduction.py | 146 ++ tests/regression/tests/test_runner_checks.py | 219 +++ .../regression/tests/test_runner_contracts.py | 203 +++ tests/regression/tests/test_tolerances.py | 93 ++ tests/regression/tolerances.json | 46 + tests/regression/tolerances.py | 164 ++ 88 files changed, 14011 insertions(+), 916 deletions(-) create mode 100644 example/common/regression_output.cc create mode 100644 example/common/regression_output.h create mode 100755 scripts/run_all_models.sh create mode 100755 scripts/run_all_tests.sh create mode 100644 tests/regression/README.md create mode 100644 tests/regression/TOLERANCES.md create mode 100644 tests/regression/assets.py create mode 100644 tests/regression/checks.py create mode 100644 tests/regression/environments/cuda.json create mode 100644 tests/regression/execution.py create mode 100644 tests/regression/fixtures/README.md create mode 100644 tests/regression/fixtures/classic.json create mode 100644 tests/regression/fixtures/generate.py create mode 100644 tests/regression/fixtures/manifest.json create mode 100644 tests/regression/fixtures/reference.py create mode 100644 tests/regression/history.py create mode 100644 tests/regression/models/gpt2.json create mode 100644 tests/regression/models/llama3.json create mode 100644 tests/regression/models/mixtral.json create mode 100644 tests/regression/models/qwen3.json create mode 100644 tests/regression/planning.py create mode 100644 tests/regression/profiling.py create mode 100644 tests/regression/quarantines.json create mode 100644 tests/regression/reporting.py create mode 100644 tests/regression/reproduction.py create mode 100644 tests/regression/runner.py create mode 100644 tests/regression/scenarios/classic_8_proc.json create mode 100644 tests/regression/scenarios/classic_basic.json create mode 100644 tests/regression/scenarios/classic_checkpoint.json create mode 100644 tests/regression/scenarios/classic_lora.json create mode 100644 tests/regression/scenarios/classic_lr_scheduler.json create mode 100644 tests/regression/scenarios/classic_moe.json create mode 100644 tests/regression/scenarios/classic_qwen3.json create mode 100644 tests/regression/scenarios/classic_qwen3_8_proc.json create mode 100644 tests/regression/scenarios/classic_recompute.json create mode 100644 tests/regression/scenarios/classic_recompute_qwen3.json create mode 100644 tests/regression/scenarios/classic_zero.json create mode 100644 tests/regression/scenarios/kernel_basic.json create mode 100644 tests/regression/scenarios/kernel_ddp.json create mode 100644 tests/regression/scenarios/long_stability.json create mode 100644 tests/regression/scenarios/tiny_accumulation.json create mode 100644 tests/regression/scenarios/tiny_checkpoint.json create mode 100644 tests/regression/scenarios/tiny_lora.json create mode 100644 tests/regression/scenarios/tiny_oracle.json create mode 100644 tests/regression/scenarios/tiny_parallel.json create mode 100644 tests/regression/scenarios/tiny_parallel_anchors.json create mode 100644 tests/regression/scenarios/tiny_recompute.json create mode 100644 tests/regression/scenarios/tiny_scheduler.json create mode 100644 tests/regression/scenarios/tiny_smoke.json create mode 100644 tests/regression/scenarios/tiny_zero.json create mode 100644 tests/regression/schema/README.md create mode 100644 tests/regression/suites/nightly.json create mode 100644 tests/regression/suites/perf.json create mode 100644 tests/regression/suites/pr.json create mode 100644 tests/regression/suites/release.json create mode 100644 tests/regression/suites/smoke.json create mode 100644 tests/regression/test_checks.py create mode 100644 tests/regression/tests/__init__.py create mode 100644 tests/regression/tests/test_compatibility.py create mode 100644 tests/regression/tests/test_execution.py create mode 100644 tests/regression/tests/test_fixtures.py create mode 100644 tests/regression/tests/test_history.py create mode 100644 tests/regression/tests/test_log_mode.py create mode 100644 tests/regression/tests/test_oracle_integration.py create mode 100644 tests/regression/tests/test_planning.py create mode 100644 tests/regression/tests/test_profiling.py create mode 100644 tests/regression/tests/test_recompute.py create mode 100644 tests/regression/tests/test_reference_batching.py create mode 100644 tests/regression/tests/test_reproduction.py create mode 100644 tests/regression/tests/test_runner_checks.py create mode 100644 tests/regression/tests/test_runner_contracts.py create mode 100644 tests/regression/tests/test_tolerances.py create mode 100644 tests/regression/tolerances.json create mode 100644 tests/regression/tolerances.py diff --git a/.gitignore b/.gitignore index 4ad6f92ff..65a36aeff 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,7 @@ build/ __pycache__/ /data/ + +# Local training regression environment and generated run artifacts. +/tests.local.json +/artifacts/test-runs/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 88911e84a..fe89a7c27 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -274,6 +274,9 @@ add_executable(qwen3 example/common/tokenizer.cc ) link_infini_train_exe(qwen3) +foreach(model_target gpt2 llama3 qwen3 mixtral) + target_sources(${model_target} PRIVATE example/common/regression_output.cc) +endforeach() endif() # Tools diff --git a/docs/test_usage_guide.md b/docs/test_usage_guide.md index 4267a351f..71abb0535 100644 --- a/docs/test_usage_guide.md +++ b/docs/test_usage_guide.md @@ -1,6 +1,37 @@ # 测试使用指南 -本指南介绍如何在 InfiniTrain 中构建、运行和编写测试。底层设计详见 [test_infrastructure_design.md](test_infrastructure_design.md) 和 [test_lifecycle.md](test_lifecycle.md)。 +本指南介绍如何构建、运行和编写测试。底层单元测试设计见 [test_infrastructure_design.md](test_infrastructure_design.md)。 + +## 一键训练回归 + +默认与上一次成功运行的原始日志进行 loss 和吞吐比较，保留配置矩阵与五类套件： + +```bash +./scripts/run_all_tests.sh # PR：87 个 GPU 训练任务及单元测试 +./scripts/run_all_tests.sh --suite smoke # 四个小模型 +./scripts/run_all_models.sh # nightly：299 个任务，含 212 个经典任务 +./scripts/run_all_tests.sh --suite perf # 14项基础/DDP算子profiling，独立PROFILE_MODE=ON构建 +./scripts/run_all_tests.sh --suite release # 303项普通构建回归及长程训练 + +./scripts/run_all_tests.sh --list +./scripts/run_all_tests.sh --model llama3 --tag zero2 +./scripts/run_all_tests.sh --tag recompute # 三模型小规模重计算 +./scripts/run_all_tests.sh --suite nightly --tag classic --tag recompute --model qwen3 +./scripts/run_all_tests.sh --reference-run artifacts/test-runs/OLD_RUN_ID +./scripts/run_all_tests.sh --rerun-failed artifacts/test-runs/RUN_ID/summary.json +``` + +首次没有历史时标记 `RECORDED`，保存本次记录；之后按相同 case 的训练参数、资产和步数匹配。FP32/BF16 loss 默认绝对误差分别为 `1e-5`/`1e-2`；性能比较排除首步后的平均 tok/s，默认降速超过 20% 失败。失败运行不会成为下一次参考。`--record-history` 可显式重新记录历史，`--strict` 则要求完整套件且完成实际对齐。 + +默认不运行额外单卡参考、PyTorch oracle 或完整张量快照。checkpoint 所需的保存/加载依赖仍自动执行。需要专项深度检查时，显式传入 `--alignment deep`；该模式在历史日志比较之外，执行场景声明的数值检查。 + +每次运行生成 `run.json`、`summary.md`、逐用例 `train.log`、指标及覆盖报告。记录 UTC 时间、完整 commit、dirty 标记、case ID、实际命令和退出码。首轮记录与真正对齐通过会分开显示。 + +本机的容器、数据、构建、历史及 checkpoint 目录放在 `tests.local.json`。建议大型 checkpoint 使用 `--checkpoint-root /fast-scratch/infinitrain-regression`；成功后只清理本次运行所属 checkpoint，失败现场及日志保留。CPU 仅支持 FP32 合法子集，不能代替 GPU 分布式验证。 + +普通新增组合在 `tests/regression/scenarios/` 添加 JSON；同接口新模型登记到 `models/`。配置约束见 [配置规范](../tests/regression/schema/README.md)，完整使用说明见 [训练回归指南](../tests/regression/README.md)。 + +CI 接入将另行提供。本次变更提供本地测试入口；后续计划在 CI 中运行 Python 自测、纯 CPU 构建、CTest 和四模型 smoke，并通过自托管 runner 执行 GPU 套件。CI 的历史目录应保存在 checkout 之外，避免被清理。 --- diff --git a/example/common/regression_output.cc b/example/common/regression_output.cc new file mode 100644 index 000000000..1b6fa831e --- /dev/null +++ b/example/common/regression_output.cc @@ -0,0 +1,388 @@ +#include "example/common/regression_output.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/modules/module.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel.h" +#include "infini_train/include/nn/parallel/ddp/distributed_data_parallel_config.h" +#include "infini_train/include/nn/parallel/ddp/param_and_grad_buffer.h" +#include "infini_train/include/nn/parallel/pp/pipeline_parallel.h" +#include "infini_train/include/nn/parallel/process_group.h" +#include "infini_train/include/optimizer.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::regression { +namespace { +std::string JsonString(const std::string &value) { + std::ostringstream out; + out << '"'; + for (unsigned char c : value) { + if (c == '"' || c == '\\') { + out << '\\' << c; + } else if (c < 32) { + out << "\\u" << std::hex << std::setw(4) << std::setfill('0') + << static_cast(c) << std::dec; + } else { + out << c; + } + } + return out.str() + '"'; +} + +// Tensor::SaveAsNpy currently only accepts FP32. Diagnostic state must preserve +// integer counters (including values above 2^24) and double-precision state. +std::string SaveSnapshotNpy(const Tensor &cpu, std::ostream &out) { + static_assert(std::endian::native == std::endian::little); + const std::map> types = { + {DataType::kFLOAT32, {"(text.size() & 255)); + out.put(static_cast((text.size() >> 8) & 255)); + out.write(text.data(), text.size()); + out.write(static_cast(cpu.DataPtr()), cpu.SizeInBytes()); + CHECK(out) << "Failed writing tensor snapshot"; + return dtype; +} + +} // namespace +Output::Output(const std::string &metrics_file, const std::string &snapshot_dir, + const std::string &model, const std::string &dtype, + const nn::TransformerConfig &config, int rank, int dp, int tp, + int pp, int sp, int dp_rank, int tp_rank, int pp_rank, int vpp, + const std::string &optimizer_type) + : snapshot_dir_(snapshot_dir), + config_(config), + rank_(rank), + dp_(dp), + tp_(tp), + pp_(pp), + tp_rank_(tp_rank), + pp_rank_(pp_rank) { + std::ostringstream context; + context << "\"schema_version\":1,\"model\":" << JsonString(model) + << ",\"dtype\":" << JsonString(dtype) << ",\"rank\":" << rank + << ",\"dp\":" << dp << ",\"tp\":" << tp << ",\"pp\":" << pp + << ",\"sp\":" << sp << ",\"dp_rank\":" << dp_rank + << ",\"tp_rank\":" << tp_rank << ",\"pp_rank\":" << pp_rank + << ",\"vpp\":" << vpp << ",\"world_size\":" << dp * tp * pp + << ",\"memory_resolution_bytes\":1048576,\"optimizer_type\":" + << JsonString(optimizer_type) << ",\"optimizer_stateful\":" + << (optimizer_type == "sgd" ? "false" : "true"); + context_ = context.str(); + if (!metrics_file.empty()) { + const std::filesystem::path path(metrics_file + ".rank_" + + std::to_string(rank) + ".jsonl"); + if (!path.parent_path().empty()) { + std::filesystem::create_directories(path.parent_path()); + } + metrics_.open(path, std::ios::trunc); + CHECK(metrics_) << "Cannot open regression metrics: " << path; + metrics_ << std::setprecision(17); + } + const auto info = nn::parallel::PipelineParallel::GetStageInfo( + config.n_layer, pp, pp_rank, vpp); + for (auto [begin, end] : info.layer_ranges_per_chunk) { + for (int i = begin; i < end; ++i) { + global_layers_.push_back(i); + } + } +} + +bool Output::enabled() const { + return metrics_.is_open() || !snapshot_dir_.empty(); +} + +void Output::Step(int64_t step, double loss, double lr, int64_t global_tokens, + double step_ms, size_t used_mb, size_t reserved_mb, + size_t consumed_samples) { + if (!metrics_.is_open()) { + return; + } + // JSON has no NaN/Inf literals: invalid metrics fail at their source. + CHECK(std::isfinite(lr) && std::isfinite(step_ms) && step_ms > 0); + if (pp_rank_ == pp_ - 1) { + CHECK(std::isfinite(loss)) << "Non-finite training loss"; + } + metrics_ << '{' << context_ << ",\"type\":\"step\",\"step\":" << step + << ",\"loss\":"; + if (pp_rank_ == pp_ - 1) { + metrics_ << loss; + } else { + metrics_ << "null"; + } + metrics_ << ",\"lr\":" << lr << ",\"global_tokens\":" << global_tokens + << ",\"step_ms\":" << step_ms + << ",\"peak_memory_bytes\":" << used_mb * 1048576ULL + << ",\"peak_reserved_bytes\":" << reserved_mb * 1048576ULL + << ",\"consumed_samples\":" << consumed_samples << "}\n"; + metrics_.flush(); + CHECK(metrics_) << "Failed writing regression metrics"; +} + +void Output::Complete(int64_t step, size_t consumed_samples) { + if (!metrics_.is_open()) { + return; + } + metrics_ << '{' << context_ << ",\"type\":\"complete\",\"step\":" << step + << ",\"consumed_samples\":" << consumed_samples << "}\n"; + metrics_.flush(); + CHECK(metrics_); +} + +void Output::Snapshot(int64_t step, const std::shared_ptr &model, + const Optimizer *optimizer) { + if (snapshot_dir_.empty()) { + return; + } + // Both ZeRO-1 and ZeRO-2 reduce-scatter gradients. Reconstruct only for + // diagnostics; never change the + // optimizer's shard storage or the parameter's own gradient binding. + std::unordered_map> gradients; + std::vector> roots{model}; + if (auto *pipeline = + dynamic_cast(model.get())) { + roots = *pipeline->mutable_chunks(); + } + std::set visited; + if (step > 0 && optimizer) { + for (const auto &root : roots) { + for (const auto &[_, module] : root->NamedModules()) { + auto *ddp = + dynamic_cast(module.get()); + if (!ddp || !visited.insert(ddp).second || + ddp->ddp_config().zero_stage < 1) { + continue; + } + const auto *pg = ddp->param_grad_buffers().front()->ddp_pg(); + for (const auto &group : ddp->bucket_groups()) { + for (size_t i = 0; i < group->buckets().size(); ++i) { + const auto &bucket = group->buckets()[i]; + std::shared_ptr shard; + if (ddp->ddp_config().zero_stage >= 2) { + shard = group->GetLocalGradShardBuffer(i); + } else { + // ZeRO-1 retains a full allocation, but only the + // local reduce-scatter slice is globally reduced. + const auto &storage = bucket->grad_data(); + CHECK_EQ(storage->NumElements() % dp_, 0); + const auto elements = storage->NumElements() / dp_; + const auto local_rank = pg->GetGroupRank(rank_); + shard = std::make_shared( + *storage, + local_rank * elements * kDataTypeToSize.at(storage->Dtype()), + std::vector{static_cast(elements)}); + } + auto full = std::make_shared( + std::vector{ + static_cast(shard->NumElements() * dp_)}, + shard->Dtype(), shard->GetDevice()); + pg->AllGather(full, shard); + for (const auto ¶meter : bucket->params()) { + size_t begin = 0, end = 0; + CHECK(bucket->GetTensorLocInBucket(parameter, begin, end)); + gradients[parameter.get()] = std::make_shared( + *full, begin * kDataTypeToSize.at(full->Dtype()), + parameter->Dims()); + } + } + } + } + } + } + const auto directory = std::filesystem::path(snapshot_dir_) / + ("step_" + std::to_string(step)) / + ("rank_" + std::to_string(rank_)); + std::filesystem::create_directories(directory); + std::ofstream tensors(directory / "tensors.bin", std::ios::binary); + CHECK(tensors); + std::ofstream manifest(directory / "manifest.json"); + CHECK(manifest); + manifest << '{' << context_ << ",\"step\":" << step << ",\"tensors\":["; + size_t index = 0; + std::unordered_map aliases; + auto write = [&](const std::string &name, const std::string &kind, + const std::shared_ptr &tensor, bool trainable) { + if (!tensor) { + return; + } + if (index) { + manifest << ','; + } + ++index; + const auto offset = static_cast(tensors.tellp()); + // Cast kernels require a shared-owned source. A temporary Tensor + // returned by To(Device) cannot call shared_from_this(). + auto host = std::make_shared(tensor->To(Device())); + host->set_grad(nullptr); + core::GetDeviceGuardImpl(tensor->GetDevice().type()) + ->SynchronizeDevice(tensor->GetDevice()); + auto cpu = host->To(host->Dtype() == DataType::kBFLOAT16 || + host->Dtype() == DataType::kFLOAT16 + ? DataType::kFLOAT32 + : host->Dtype()); + const auto storage_dtype = SaveSnapshotNpy(cpu, tensors); + const auto length = static_cast(tensors.tellp()) - offset; + manifest << "{\"name\":" << JsonString(name) + << ",\"kind\":" << JsonString(kind) + << ",\"dtype\":" << JsonString(storage_dtype) + << ",\"path\":\"tensors.bin\",\"offset\":" << offset + << ",\"length\":" << length + << ",\"trainable\":" << (trainable ? "true" : "false") + << ",\"shape\":["; + for (size_t i = 0; i < cpu.Dims().size(); ++i) { + if (i) { + manifest << ','; + } + manifest << cpu.Dims()[i]; + } + manifest << ']'; + if (kind == "parameter") { + if (aliases.contains(tensor.get())) { + manifest << ",\"alias_of\":" << JsonString(aliases.at(tensor.get())); + } else { + aliases.emplace(tensor.get(), name); + } + } + if (kind != "optimizer") { + WriteShard(manifest, name, cpu.Dims()); + } + manifest << '}'; + }; + std::set parameters; + for (const auto &[_, parameter] : model->NamedParameters()) { + parameters.insert(parameter.get()); + } + // StateDict avoids internal PP aliases and has stable logical names. + std::map> state; + for (const auto &[name, tensor] : model->StateDict()) { + if (parameters.contains(tensor.get())) { + state.emplace(CanonicalName(name), tensor); + } + } + for (const auto &[name, tensor] : state) { + write(name, "parameter", tensor, tensor->requires_grad()); + if (step > 0 && optimizer) { + auto gradient = gradients.contains(tensor.get()) + ? gradients.at(tensor.get()) + : tensor->grad(); + write(name, "gradient", gradient, tensor->requires_grad()); + } + } + if (optimizer && step > 0) { + std::map> states; + for (const auto &[name, tensor] : optimizer->StateDict()) { + states.emplace(name, tensor); + } + for (const auto &[name, tensor] : states) { + write("rank_" + std::to_string(rank_) + "." + name, "optimizer", tensor, + false); + } + } + tensors.flush(); + CHECK(tensors); + manifest << "]}\n"; + manifest.flush(); + CHECK(manifest); +} + +std::string Output::CanonicalName(std::string name) const { + while (name.starts_with("module.")) { + name.erase(0, 7); + } + const std::string prefix = "transformer.h."; + if (name.starts_with(prefix)) { + const auto end = name.find('.', prefix.size()); + const size_t local = + std::stoul(name.substr(prefix.size(), end - prefix.size())); + CHECK_LT(local, global_layers_.size()); + name.replace(prefix.size(), end - prefix.size(), + std::to_string(global_layers_[local])); + } + return name; +} + +void Output::WriteShard(std::ostream &out, const std::string &name, + const std::vector &shape) const { + if (tp_ == 1) { + return; + } + int axis = -1; + std::vector parts; + const bool column = name.find(".c_attn.") != std::string::npos || + name.find(".c_fc.") != std::string::npos; + if (column && name.ends_with(".lora_A")) { + return; + } + if (name.find(".c_proj.") != std::string::npos && name.ends_with(".lora_B")) { + return; + } + if (name == "transformer.wte.weight" || name == "lm_head.weight") { + axis = 0; + } else if (name.find(".c_attn.") != std::string::npos) { + axis = 0; + const auto kv = config_.n_kv_head * config_.n_embd / config_.n_head; + parts = {config_.n_embd, kv, kv}; + } else if (name.find(".c_fc.") != std::string::npos) { + axis = 0; + if (config_.activation_type == nn::MLPType::kSwiGLU) { + parts = {shape[0] * tp_ / 2, shape[0] * tp_ / 2}; + } + } else if (name.ends_with(".c_proj.weight") || + name.ends_with(".c_proj.lora_A")) { + axis = 1; + } + if (axis < 0) { + return; + } + out << ",\"shard\":{\"axis\":" << axis << ",\"index\":" << tp_rank_ + << ",\"count\":" << tp_; + if (!parts.empty()) { + out << ",\"part_sizes\":["; + for (size_t i = 0; i < parts.size(); ++i) { + if (i) { + out << ','; + } + out << parts[i]; + } + out << ']'; + } + out << '}'; +} +} // namespace infini_train::regression diff --git a/example/common/regression_output.h b/example/common/regression_output.h new file mode 100644 index 000000000..d96afd4a8 --- /dev/null +++ b/example/common/regression_output.h @@ -0,0 +1,49 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/nn/modules/transformer/transformer_config.h" +namespace infini_train { +class Optimizer; +namespace nn { +class Module; +} +} // namespace infini_train +namespace infini_train::regression { +class Output { + public: + Output(const std::string &metrics_file, const std::string &snapshot_dir, + const std::string &model, const std::string &dtype, + const nn::TransformerConfig &config, int rank = 0, int dp = 1, + int tp = 1, int pp = 1, int sp = 1, int dp_rank = 0, int tp_rank = 0, + int pp_rank = 0, int vpp = 1, + const std::string &optimizer_type = "adam"); + + bool enabled() const; + + void Step(int64_t step, double loss, double lr, int64_t global_tokens, + double step_ms, size_t used_mb, size_t reserved_mb, + size_t consumed_samples); + + void Complete(int64_t step, size_t consumed_samples); + + void Snapshot(int64_t step, const std::shared_ptr &model, + const Optimizer *optimizer = nullptr); + + private: + std::string CanonicalName(std::string name) const; + + void WriteShard(std::ostream &out, const std::string &name, + const std::vector &shape) const; + + std::ofstream metrics_; + std::string snapshot_dir_, context_; + nn::TransformerConfig config_; + int rank_, dp_, tp_, pp_, tp_rank_, pp_rank_; + std::vector global_layers_; +}; +} // namespace infini_train::regression diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc index b140edbd4..d8098be95 100644 --- a/example/gpt2/checkpoint_loader.cc +++ b/example/gpt2/checkpoint_loader.cc @@ -96,6 +96,7 @@ std::shared_ptr LoadFromLLMC( gpt2_config.original_vocab_size = vocab_size; gpt2_config.n_layer = n_layer; gpt2_config.n_head = n_head; + gpt2_config.n_kv_head = n_head; gpt2_config.n_embd = n_embd; gpt2::SanitizeGPT2Config(gpt2_config); nn::ApplyActivationRecomputeOptions(&gpt2_config, recompute_options); diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 37ce5024b..e46d2cd3c 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -62,6 +62,7 @@ #include "infini_train/include/utils/precision_check_config.h" #include "infini_train/include/utils/precision_checker.h" +#include "example/common/regression_output.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/common/tokenizer.h" #include "example/gpt2/checkpoint_loader.h" @@ -70,6 +71,10 @@ // TODO(jym): Reorganize CLI flags into categories for better readability and // maintainability. I/O DEFINE_string(input_bin, "", "input .bin to train on"); +DEFINE_string(metrics_file, "", + "Optional regression JSONL prefix (one file per rank)"); +DEFINE_string(regression_snapshot_dir, "", + "Optional per-step regression tensor snapshot directory"); DEFINE_string(input_val_bin, "", "input .bin to eval validation loss on"); DEFINE_string(tokenizer_bin, "", "input .bin to tokenizer"); // model bin file is downloaded and processed using the script at @@ -328,7 +333,10 @@ void Train(const nn::parallel::Rank &rank) { FLAGS_recompute_method, static_cast(FLAGS_recompute_num_layers)); if (!FLAGS_llmc_filepath.empty()) { - model = gpt2::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); + auto loaded_model = + gpt2::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); + model_config = loaded_model->Config(); + model = loaded_model; } else if (kModelToConfigs.count(FLAGS_model)) { model_config = kModelToConfigs.at(FLAGS_model); gpt2::SanitizeGPT2Config(model_config); @@ -512,6 +520,12 @@ void Train(const nn::parallel::Rank &rank) { .lr_scheduler = scheduler}); start_step = resume_result.global_step; size_t consumed_train_samples = resume_result.consumed_train_samples; + regression::Output regression_output( + FLAGS_metrics_file, FLAGS_regression_snapshot_dir, "gpt2", FLAGS_dtype, + model_config, rank.GlobalRank(), ddp_world_size, tp_world_size, + pp_world_size, sp_world_size, ddp_rank, tp_rank, pp_rank, + FLAGS_virtual_pipeline_parallel, "sgd"); + regression_output.Snapshot(start_step, model); auto advance_train_iter = [&]() { ++train_iter; @@ -679,11 +693,21 @@ void Train(const nn::parallel::Rank &rank) { static_cast(lossf_tensor->To(Device()).DataPtr())[0]; } + if (regression_output.enabled()) { + impl->SynchronizeDevice(device); + } const auto iter_end = std::chrono::high_resolution_clock::now(); const double duration_us = std::chrono::duration(iter_end - iter_start) .count(); const double tps = FLAGS_total_batch_size / (duration_us / 1e6); + if (regression_output.enabled()) { + const auto [peak_used, peak_reserved] = impl->GetMemPoolPeakMB(device); + regression_output.Step(step + 1, lossf, current_lr, + FLAGS_total_batch_size, duration_us / 1e3, + peak_used, peak_reserved, consumed_train_samples); + regression_output.Snapshot(step + 1, model, optimizer.get()); + } if (rank.IsLastRank()) { size_t used_mb = 0, reserved_mb = 0; @@ -722,6 +746,8 @@ void Train(const nn::parallel::Rank &rank) { } } + regression_output.Complete(FLAGS_num_iteration, consumed_train_samples); + // Save LoRA weights if enabled and path specified if (lora_enabled && !FLAGS_lora_save_path.empty()) { LOG(INFO) << "Saving LoRA weights to: " << FLAGS_lora_save_path; diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc index bbd6a9960..50e0ea8fb 100644 --- a/example/llama3/checkpoint_loader.cc +++ b/example/llama3/checkpoint_loader.cc @@ -97,6 +97,7 @@ std::shared_ptr LoadFromLLMC( nn::TransformerConfig llama3_config = llama3::LLaMA3Config(); llama3_config.block_size = block_size; llama3_config.vocab_size = vocab_size; + llama3_config.original_vocab_size = vocab_size; llama3_config.n_layer = n_layer; llama3_config.n_head = n_head; llama3_config.n_kv_head = n_kv_head; diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 2143ad98d..3f3c1c94b 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -61,6 +61,7 @@ #include "infini_train/include/profiler.h" #endif +#include "example/common/regression_output.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/common/tokenizer.h" #include "example/llama3/checkpoint_loader.h" @@ -69,6 +70,10 @@ // TODO(jym): Reorganize CLI flags into categories for better readability and // maintainability. I/O DEFINE_string(input_bin, "", "input .bin to train on"); +DEFINE_string(metrics_file, "", + "Optional regression JSONL prefix (one file per rank)"); +DEFINE_string(regression_snapshot_dir, "", + "Optional per-step regression tensor snapshot directory"); DEFINE_string(input_val_bin, "", "input .bin to eval validation loss on"); DEFINE_string(tokenizer_bin, "", "input .bin to tokenizer"); // model bin file is downloaded and processed using the script at @@ -297,7 +302,10 @@ void Train(const nn::parallel::Rank &rank) { FLAGS_activation_recompute, FLAGS_recompute_granularity, FLAGS_recompute_method, static_cast(FLAGS_recompute_num_layers)); if (!FLAGS_llmc_filepath.empty()) { - model = llama3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); + auto loaded_model = + llama3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); + model_config = loaded_model->Config(); + model = loaded_model; } else { nn::ApplyActivationRecomputeOptions(&model_config, recompute_options); llama3::SanitizeLLaMA3Config(model_config); @@ -490,6 +498,12 @@ void Train(const nn::parallel::Rank &rank) { start_step = resume_result.global_step; size_t consumed_train_samples = resume_result.consumed_train_samples; + regression::Output regression_output( + FLAGS_metrics_file, FLAGS_regression_snapshot_dir, "llama3", FLAGS_dtype, + model_config, rank.GlobalRank(), ddp_world_size, tp_world_size, + pp_world_size, sp_world_size, ddp_rank, tp_rank, pp_rank, + FLAGS_virtual_pipeline_parallel); + regression_output.Snapshot(start_step, model); auto advance_train_iter = [&]() { ++train_iter; @@ -654,11 +668,21 @@ void Train(const nn::parallel::Rank &rank) { static_cast(lossf_tensor->To(Device()).DataPtr())[0]; } + if (regression_output.enabled()) { + impl->SynchronizeDevice(device); + } const auto iter_end = std::chrono::high_resolution_clock::now(); const double duration_us = std::chrono::duration(iter_end - iter_start) .count(); const double tps = FLAGS_total_batch_size / (duration_us / 1e6); + if (regression_output.enabled()) { + const auto [peak_used, peak_reserved] = impl->GetMemPoolPeakMB(device); + regression_output.Step(step + 1, lossf, current_lr, + FLAGS_total_batch_size, duration_us / 1e3, + peak_used, peak_reserved, consumed_train_samples); + regression_output.Snapshot(step + 1, model, optimizer.get()); + } if (rank.IsLastRank()) { size_t used_mb = 0, reserved_mb = 0; @@ -697,6 +721,8 @@ void Train(const nn::parallel::Rank &rank) { } } + regression_output.Complete(FLAGS_num_iteration, consumed_train_samples); + // Save LoRA weights if enabled and path specified if (lora_enabled && !FLAGS_lora_save_path.empty()) { LOG(INFO) << "Saving LoRA weights to: " << FLAGS_lora_save_path; diff --git a/example/mixtral/main.cc b/example/mixtral/main.cc index 82ee2b392..9656f350c 100644 --- a/example/mixtral/main.cc +++ b/example/mixtral/main.cc @@ -49,12 +49,17 @@ #include "infini_train/include/profiler.h" #endif +#include "example/common/regression_output.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/mixtral/checkpoint_loader.h" #include "example/mixtral/config.h" // I/O DEFINE_string(input_bin, "", "input .bin to train on"); +DEFINE_string(metrics_file, "", + "Optional regression JSONL prefix (one file per rank)"); +DEFINE_string(regression_snapshot_dir, "", + "Optional per-step regression tensor snapshot directory"); DEFINE_string(llmc_filepath, "", "llmc model file path to load from"); // token layout for each step of the optimization DEFINE_uint32(micro_batch_size, 4, "batch size, in units of #batch dimensions"); @@ -172,6 +177,7 @@ void Train(const nn::parallel::Rank &rank) { std::shared_ptr model; if (!FLAGS_llmc_filepath.empty()) { + model_config = mixtral::ConfigFromLLMC(FLAGS_llmc_filepath); model = mixtral::LoadFromLLMC(FLAGS_llmc_filepath, model_config); } else { model = std::make_shared(model_config); @@ -259,6 +265,11 @@ void Train(const nn::parallel::Rank &rank) { step_duration_ms.reserve(FLAGS_num_iteration); const double tokens_per_step = static_cast(FLAGS_global_batch_size) * FLAGS_sequence_length; + regression::Output regression_output( + FLAGS_metrics_file, FLAGS_regression_snapshot_dir, "mixtral", FLAGS_dtype, + model_config, rank.GlobalRank(), ddp_world_size, 1, pp_world_size, 1, + ddp_rank, 0, pp_rank); + regression_output.Snapshot(0, model); LOG(INFO) << "Rank " << rank.GlobalRank() << ": start tiny Mixtral training on " << device.ToString(); @@ -323,6 +334,16 @@ void Train(const nn::parallel::Rank &rank) { step_end_time - step_start_time) .count(); step_duration_ms.push_back(duration_ms); + if (regression_output.enabled()) { + const auto [peak_used, peak_reserved] = + device_impl->GetMemPoolPeakMB(device); + regression_output.Step( + step + 1, lossf, FLAGS_learning_rate, + static_cast(tokens_per_step), duration_ms, peak_used, + peak_reserved, + static_cast(step + 1) * FLAGS_global_batch_size); + regression_output.Snapshot(step + 1, model, optimizer.get()); + } if (rank.IsLastRank() && FLAGS_log_interval > 0 && ((step + 1) % FLAGS_log_interval == 0 || @@ -341,6 +362,9 @@ void Train(const nn::parallel::Rank &rank) { << std::endl; } } + regression_output.Complete( + FLAGS_num_iteration, + static_cast(FLAGS_num_iteration) * FLAGS_global_batch_size); if (rank.IsLastRank() && FLAGS_print_timing && !step_duration_ms.empty()) { double duration_sum_ms = 0.0; diff --git a/example/qwen3/checkpoint_loader.cc b/example/qwen3/checkpoint_loader.cc index 7270f7bbd..fb66c3ee2 100644 --- a/example/qwen3/checkpoint_loader.cc +++ b/example/qwen3/checkpoint_loader.cc @@ -115,6 +115,7 @@ std::shared_ptr LoadFromLLMC( nn::TransformerConfig qwen3_config = qwen3::Qwen3Config(); qwen3_config.block_size = block_size; qwen3_config.vocab_size = vocab_size; + qwen3_config.original_vocab_size = vocab_size; qwen3_config.n_layer = n_layer; qwen3_config.n_head = n_head; qwen3_config.n_kv_head = n_kv_head; diff --git a/example/qwen3/main.cc b/example/qwen3/main.cc index 8759713f3..44b0cbdb1 100644 --- a/example/qwen3/main.cc +++ b/example/qwen3/main.cc @@ -61,6 +61,7 @@ #include "infini_train/include/profiler.h" #endif +#include "example/common/regression_output.h" #include "example/common/tiny_shakespeare_dataset.h" #include "example/common/tokenizer.h" #include "example/qwen3/checkpoint_loader.h" @@ -69,6 +70,10 @@ // TODO(jym): Reorganize CLI flags into categories for better readability and // maintainability. I/O DEFINE_string(input_bin, "", "input .bin to train on"); +DEFINE_string(metrics_file, "", + "Optional regression JSONL prefix (one file per rank)"); +DEFINE_string(regression_snapshot_dir, "", + "Optional per-step regression tensor snapshot directory"); DEFINE_string(input_val_bin, "", "input .bin to eval validation loss on"); DEFINE_string(tokenizer_bin, "", "input .bin to tokenizer"); // model bin file is downloaded and processed using the script at @@ -297,7 +302,10 @@ void Train(const nn::parallel::Rank &rank) { FLAGS_activation_recompute, FLAGS_recompute_granularity, FLAGS_recompute_method, static_cast(FLAGS_recompute_num_layers)); if (!FLAGS_llmc_filepath.empty()) { - model = qwen3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); + auto loaded_model = + qwen3::LoadFromLLMC(FLAGS_llmc_filepath, recompute_options); + model_config = loaded_model->Config(); + model = loaded_model; } else { nn::ApplyActivationRecomputeOptions(&model_config, recompute_options); qwen3::SanitizeQwen3Config(model_config); @@ -489,6 +497,12 @@ void Train(const nn::parallel::Rank &rank) { start_step = resume_result.global_step; size_t consumed_train_samples = resume_result.consumed_train_samples; + regression::Output regression_output( + FLAGS_metrics_file, FLAGS_regression_snapshot_dir, "qwen3", FLAGS_dtype, + model_config, rank.GlobalRank(), ddp_world_size, tp_world_size, + pp_world_size, sp_world_size, ddp_rank, tp_rank, pp_rank, + FLAGS_virtual_pipeline_parallel); + regression_output.Snapshot(start_step, model); auto advance_train_iter = [&]() { ++train_iter; @@ -652,11 +666,21 @@ void Train(const nn::parallel::Rank &rank) { static_cast(lossf_tensor->To(Device()).DataPtr())[0]; } + if (regression_output.enabled()) { + impl->SynchronizeDevice(device); + } const auto iter_end = std::chrono::high_resolution_clock::now(); const double duration_us = std::chrono::duration(iter_end - iter_start) .count(); const double tps = FLAGS_total_batch_size / (duration_us / 1e6); + if (regression_output.enabled()) { + const auto [peak_used, peak_reserved] = impl->GetMemPoolPeakMB(device); + regression_output.Step(step + 1, lossf, current_lr, + FLAGS_total_batch_size, duration_us / 1e3, + peak_used, peak_reserved, consumed_train_samples); + regression_output.Snapshot(step + 1, model, optimizer.get()); + } if (rank.IsLastRank()) { size_t used_mb = 0, reserved_mb = 0; @@ -695,6 +719,8 @@ void Train(const nn::parallel::Rank &rank) { } } + regression_output.Complete(FLAGS_num_iteration, consumed_train_samples); + // Save LoRA weights if enabled and path specified if (lora_enabled && !FLAGS_lora_save_path.empty()) { LOG(INFO) << "Saving LoRA weights to: " << FLAGS_lora_save_path; diff --git a/scripts/compare_loss.py b/scripts/compare_loss.py index da9c81d90..845d92cb8 100755 --- a/scripts/compare_loss.py +++ b/scripts/compare_loss.py @@ -1,171 +1,85 @@ #!/usr/bin/env python3 -# Usage: -# python tools/compare_loss.py \ -# /path/to/baseline/logs \ -# /path/to/test/logs \ -# --threshold-fp32 1e-5 --threshold-bf16 1e-2 - +"""Compare all legacy training logs, failing on incomplete or invalid evidence.""" +import math import re import sys from pathlib import Path from argparse import ArgumentParser -from compare_utils import collect_log_files, exit_if_duplicate_logs +from compare_utils import ( + matching_logs, + load_legacy_metrics, + compare_metrics, + MetricError, +) def get_dtype_from_filename(filename): - """Determine dtype from filename. Returns 'bfloat16' or 'fp32'.""" - bf16_suffix = re.search(r"_(?:bfloat16|bf16)(?:_|\.|$)", filename) - return "bfloat16" if bf16_suffix else "fp32" + # Explicit --dtype is preferred. Keep old names, including Qwen's bf16 alias. + return ( + "bfloat16" + if re.search(r"(?:^|[_\-.])(?:bfloat16|bf16)(?:[_\-.]|$)", filename) + else "float32" + ) def parse_log(file_path): - """Extract step -> loss mapping from log file.""" - pattern = re.compile(r"step\s+(\d+)/\d+\s+\|\s+train loss\s+([\d.]+)") - losses = {} - with open(file_path) as f: - for line in f: - match = pattern.search(line) - if match: - losses[int(match.group(1))] = float(match.group(2)) - return losses - - -def compare_files(file1, file2, threshold): - """Compare loss values from two log files.""" - losses1 = parse_log(file1) - losses2 = parse_log(file2) + return {row["step"]: row["loss"] for row in load_legacy_metrics(file_path).records} - all_steps = sorted(set(losses1.keys()) | set(losses2.keys())) - mismatches = [] - for step in all_steps: - if step not in losses1: - mismatches.append(f" Step {step}: missing in {file1.name}") - elif step not in losses2: - mismatches.append(f" Step {step}: missing in {file2.name}") - else: - loss1, loss2 = losses1[step], losses2[step] - diff = abs(loss1 - loss2) - if diff > threshold: - rel = ( - diff / max(abs(loss1), abs(loss2)) * 100 - if max(abs(loss1), abs(loss2)) > 0 - else 0 - ) - mismatches.append( - f" Step {step}: {loss1:.6f} vs {loss2:.6f} ✗ (diff: {diff:.2e}, {rel:.4f}%)" - ) - - return len(all_steps), len(mismatches), mismatches +def compare_files(file1, file2, threshold, dtype=None): + try: + dtype = dtype or get_dtype_from_filename(Path(file1).name) + reference = load_legacy_metrics(file1, dtype=dtype) + candidate = load_legacy_metrics(file2, dtype=dtype) + checked = compare_metrics( + reference, candidate, dtype=dtype, atol=threshold, rtol=0 + ) + return len(reference.records), len(checked["reasons"]), checked["reasons"] + except (MetricError, ValueError, OSError) as error: + return 0, 1, [str(error)] def main(): - parser = ArgumentParser( - description="Compare training loss between two log directories" - ) - parser.add_argument("dir1", type=Path, help="Baseline log directory") - parser.add_argument("dir2", type=Path, help="Test log directory") - parser.add_argument( - "--threshold", - type=float, - help="Loss difference threshold (deprecated, use --threshold-fp32 and --threshold-bf16)", - ) - parser.add_argument( - "--threshold-fp32", - type=float, - default=1e-5, - help="Loss difference threshold for fp32 (default: 1e-5)", - ) - parser.add_argument( - "--threshold-bf16", - type=float, - default=1e-2, - help="Loss difference threshold for bfloat16 (default: 1e-2)", - ) - parser.add_argument( - "--verbose", - action="store_true", - help="Print detailed output for all files, including passed ones", - ) + parser = ArgumentParser(description=__doc__) + parser.add_argument("dir1", type=Path) + parser.add_argument("dir2", type=Path) + parser.add_argument("--threshold", type=float) + parser.add_argument("--threshold-fp32", type=float, default=1e-5) + parser.add_argument("--threshold-bf16", type=float, default=1e-2) + parser.add_argument("--dtype", choices=("float32", "fp32", "bfloat16", "bf16")) + parser.add_argument("--verbose", action="store_true") args = parser.parse_args() - - # Support legacy --threshold argument if args.threshold is not None: - args.threshold_fp32 = args.threshold - args.threshold_bf16 = args.threshold - - print(f"Baseline: {args.dir1.resolve()}") - print(f"Test: {args.dir2.resolve()}") - print() - - files1, duplicates1 = collect_log_files(args.dir1) - files2, duplicates2 = collect_log_files(args.dir2) - exit_if_duplicate_logs(args.dir1, duplicates1) - exit_if_duplicate_logs(args.dir2, duplicates2) - - only_in_1 = set(files1.keys()) - set(files2.keys()) - only_in_2 = set(files2.keys()) - set(files1.keys()) - common = set(files1.keys()) & set(files2.keys()) - - if only_in_1: - print(f"Files only in baseline: {', '.join(sorted(only_in_1))}") - if only_in_2: - print(f"Files only in test: {', '.join(sorted(only_in_2))}") - if only_in_1 or only_in_2: - print() - - total_mismatches = 0 - fp32_total = 0 - fp32_passed = 0 - bf16_total = 0 - bf16_passed = 0 - - for name in sorted(common): - dtype = get_dtype_from_filename(name) - threshold = args.threshold_bf16 if dtype == "bfloat16" else args.threshold_fp32 - - if dtype == "bfloat16": - bf16_total += 1 - else: - fp32_total += 1 - - total_steps, num_mismatches, mismatches = compare_files( - files1[name], files2[name], threshold + args.threshold_fp32 = args.threshold_bf16 = args.threshold + if any( + not math.isfinite(t) or t < 0 + for t in (args.threshold_fp32, args.threshold_bf16) + ): + parser.error("thresholds must be finite and nonnegative") + try: + baseline, candidate = matching_logs(args.dir1, args.dir2) + except ValueError as error: + print(f"[FAIL] {error}") + return 1 + passed = 0 + for name in sorted(baseline): + dtype = args.dtype or get_dtype_from_filename(name) + threshold = ( + args.threshold_bf16 + if dtype in ("bf16", "bfloat16") + else args.threshold_fp32 ) - - if mismatches: - print(f"Comparing {name} ({dtype}, threshold: {threshold:.0e}):") - for msg in mismatches: - print(msg) - total_mismatches += num_mismatches - else: - if dtype == "bfloat16": - bf16_passed += 1 - else: - fp32_passed += 1 - - # Only print details when there are mismatches or verbose mode - if mismatches or args.verbose: - if mismatches: - matched = total_steps - num_mismatches - print(f" Summary: {matched}/{total_steps} steps matched") - print() - - print("=" * 50) - print(f"Overall Summary:") - print( - f" fp32: {fp32_passed}/{fp32_total} test cases passed (threshold: {args.threshold_fp32:.0e})" - ) - print( - f" bfloat16: {bf16_passed}/{bf16_total} test cases passed (threshold: {args.threshold_bf16:.0e})" - ) - print( - f" Total: {fp32_passed + bf16_passed}/{fp32_total + bf16_total} test cases passed" - ) - print("=" * 50) - - sys.exit(1 if total_mismatches > 0 else 0) + steps, errors, reasons = compare_files( + baseline[name], candidate[name], threshold, dtype + ) + passed += not errors + if errors or args.verbose: + print(f"[{'FAIL' if errors else 'PASS'}] {name} ({dtype}, {steps} steps)") + for reason in reasons: + print(f" {reason}") + print(f"Summary: {passed}/{len(baseline)} test cases passed") + return int(passed != len(baseline)) if __name__ == "__main__": - main() + sys.exit(main()) diff --git a/scripts/compare_tps.py b/scripts/compare_tps.py index 2b22ee40e..0fed85239 100755 --- a/scripts/compare_tps.py +++ b/scripts/compare_tps.py @@ -1,182 +1,72 @@ #!/usr/bin/env python3 -# Usage: -# python tools/compare_tps.py \ -# /path/to/baseline/logs \ -# /path/to/test/logs \ -# --threshold 0.20 - -import re +"""Legacy throughput comparison using the runner's arithmetic-mean log policy.""" +import math +import statistics import sys from pathlib import Path from argparse import ArgumentParser -from compare_utils import collect_log_files, exit_if_duplicate_logs +from compare_utils import matching_logs, load_legacy_metrics, MetricError def parse_log(file_path): - """Extract step -> tok/s mapping from log file.""" - pattern = re.compile(r"step\s+(\d+)/\d+.*?\|\s+(\d+)\s+tok/s") - tps_values = {} - with open(file_path) as f: - for line in f: - match = pattern.search(line) - if match: - tps_values[int(match.group(1))] = float(match.group(2)) - return tps_values + return { + row["step"]: row["tokens_per_second"] + for row in load_legacy_metrics( + file_path, required=("loss", "tokens_per_second") + ).records + } def compare_files(file1, file2, threshold): - """Compare tok/s values from two log files, excluding first step.""" - tps1 = parse_log(file1) - tps2 = parse_log(file2) - - # Remove step 1 - tps1 = {k: v for k, v in tps1.items() if k > 1} - tps2 = {k: v for k, v in tps2.items() if k > 1} - - if not tps1 or not tps2: - return ( - 0, - True, - [" No valid steps found (after excluding step 1)"], - 0, - 0, - 0, - 0, - 0, + try: + reference, candidate = parse_log(file1), parse_log(file2) + if reference.keys() != candidate.keys(): + raise MetricError("Throughput step sets differ") + steps = list(reference)[1:] + if not steps: + raise MetricError("No measurement steps after warmup") + # Match history.py: exclude the first step, then average logged tok/s. + avg1 = statistics.mean(reference[step] for step in steps) + avg2 = statistics.mean(candidate[step] for step in steps) + change = avg2 / avg1 - 1 + failed = change < -threshold + messages = ( + [f"Throughput {avg1:.4g} -> {avg2:.4g} tok/s ({change:+.2%})"] + if failed or change > threshold + else [] ) - - # Calculate averages - avg1 = sum(tps1.values()) / len(tps1) - avg2 = sum(tps2.values()) / len(tps2) - - # Calculate signed relative change of test vs baseline: positive means test faster, negative means test slower - signed_change = (avg2 - avg1) / avg1 if avg1 > 0 else 0 - - messages = [] - failed = False - if abs(signed_change) > threshold: - sign = "+" if signed_change >= 0 else "" - if signed_change < 0: - # test slower than baseline -> failure - label = "✗ SLOWER" - failed = True - else: - # test faster than baseline -> pass but notify - label = "↑ FASTER" - messages.append( - f" Average tok/s: {avg1:.2f} (baseline) vs {avg2:.2f} (test) {label} ({sign}{signed_change*100:.1f}%, threshold: ±{threshold*100:.0f}%)" - ) - messages.append( - f" Steps compared: {len(tps1)} vs {len(tps2)} (excluding step 1)" - ) - - return 1, failed, messages, avg1, avg2, signed_change, len(tps1), len(tps2) + return 1, failed, messages, avg1, avg2, change, len(steps), len(steps) + except (MetricError, OSError, ValueError) as error: + return 0, True, [str(error)], 0, 0, 0, 0, 0 def main(): - parser = ArgumentParser(description="Compare tok/s between two log directories") - parser.add_argument("dir1", type=Path, help="Baseline log directory") - parser.add_argument("dir2", type=Path, help="Test log directory") - parser.add_argument( - "--threshold", - type=float, - default=0.20, - help="Relative error threshold (default: 0.20 = 20%%)", - ) - parser.add_argument( - "--verbose", - action="store_true", - help="Print detailed output for all files, including passed ones", - ) + parser = ArgumentParser(description=__doc__) + parser.add_argument("dir1", type=Path) + parser.add_argument("dir2", type=Path) + parser.add_argument("--threshold", type=float, default=0.20) + parser.add_argument("--verbose", action="store_true") args = parser.parse_args() - - print(f"Baseline: {args.dir1.resolve()}") - print(f"Test: {args.dir2.resolve()}") - print() - - files1, duplicates1 = collect_log_files(args.dir1) - files2, duplicates2 = collect_log_files(args.dir2) - exit_if_duplicate_logs(args.dir1, duplicates1) - exit_if_duplicate_logs(args.dir2, duplicates2) - - only_in_1 = set(files1.keys()) - set(files2.keys()) - only_in_2 = set(files2.keys()) - set(files1.keys()) - common = set(files1.keys()) & set(files2.keys()) - - if only_in_1: - print(f"Files only in baseline: {', '.join(sorted(only_in_1))}") - if only_in_2: - print(f"Files only in test: {', '.join(sorted(only_in_2))}") - if only_in_1 or only_in_2: - print() - - improvements = [] # (name, avg1, avg2, signed_change) - regressions = [] - normal = [] - - for name in sorted(common): - _, failed, messages, avg1, avg2, signed_change, steps1, steps2 = compare_files( - files1[name], files2[name], args.threshold + if not math.isfinite(args.threshold) or not 0 <= args.threshold < 1: + parser.error("threshold must be finite, nonnegative and below one") + try: + reference, candidate = matching_logs(args.dir1, args.dir2) + except ValueError as error: + print(f"[FAIL] {error}") + return 1 + passed = 0 + for name in sorted(reference): + _, failed, messages, *_ = compare_files( + reference[name], candidate[name], args.threshold ) - if failed: - regressions.append((name, avg1, avg2, signed_change, steps1, steps2)) - elif messages: - improvements.append((name, avg1, avg2, signed_change, steps1, steps2)) - else: - normal.append((name, avg1, avg2, signed_change, steps1, steps2)) - - pct = f"{args.threshold*100:.0f}%" - - if improvements: - print(f"{'=' * 50}") - print(f"Large improvements (>{pct})") - print(f"{'=' * 50}") - for name, avg1, avg2, signed_change, steps1, steps2 in improvements: - sign = "+" if signed_change >= 0 else "" - print(f"[PASS] {name}") - print( - f" average tok/s: {avg1:.2f} vs {avg2:.2f} ({sign}{signed_change*100:.1f}%)" - ) - print(f" effective steps: {steps1} vs {steps2} (step > 1)") - print() - - if regressions: - print(f"{'=' * 50}") - print(f"FAILED regressions (>{pct})") - print(f"{'=' * 50}") - for name, avg1, avg2, signed_change, steps1, steps2 in regressions: - sign = "+" if signed_change >= 0 else "" - print(f"[FAIL] {name}") - print( - f" average tok/s: {avg1:.2f} vs {avg2:.2f} ({sign}{signed_change*100:.1f}%)" - ) - print(f" effective steps: {steps1} vs {steps2} (step > 1)") - print() - - if args.verbose and normal: - print(f"{'=' * 50}") - print(f"Within threshold (<={pct})") - print(f"{'=' * 50}") - for name, avg1, avg2, signed_change, steps1, steps2 in normal: - sign = "+" if signed_change >= 0 else "" - print(f"[PASS] {name}") - print( - f" tok/s: {avg1:.2f} vs {avg2:.2f} ({sign}{signed_change*100:.1f}%)" - ) - print() - - total = len(improvements) + len(regressions) + len(normal) - passed = len(improvements) + len(normal) - print(f"{'=' * 50}") - print(f"Summary: {passed}/{total} test cases passed") - print(f"failed regressions : {len(regressions)}") - print(f"large improvements : {len(improvements)}") - print(f"within threshold : {len(normal)}") - print(f"total cases : {total}") - print(f"{'=' * 50}") - - sys.exit(1 if regressions else 0) + passed += not failed + if failed or messages or args.verbose: + print(f"[{'FAIL' if failed else 'PASS'}] {name}") + for message in messages: + print(f" {message}") + print(f"Summary: {passed}/{len(reference)} test cases passed") + return int(passed != len(reference)) if __name__ == "__main__": - main() + sys.exit(main()) diff --git a/scripts/compare_utils.py b/scripts/compare_utils.py index 1f5b22b48..398dc8337 100644 --- a/scripts/compare_utils.py +++ b/scripts/compare_utils.py @@ -35,3 +35,60 @@ def exit_if_duplicate_logs(base_dir: Path, duplicates): for name, paths in sorted(duplicates.items()): print(f" {name}: {', '.join(str(p.relative_to(base_dir)) for p in paths)}") sys.exit(1) + + +def matching_logs(dir1, dir2): + """Fail closed on missing directories, empty selection, duplicates or case sets.""" + for directory in (dir1, dir2): + if not directory.is_dir(): + raise ValueError(f"Not a log directory: {directory}") + left, duplicates_left = collect_log_files(dir1) + right, duplicates_right = collect_log_files(dir2) + exit_if_duplicate_logs(dir1, duplicates_left) + exit_if_duplicate_logs(dir2, duplicates_right) + if not left or not right: + raise ValueError("No comparable training logs") + if left.keys() != right.keys(): + raise ValueError( + f"Case sets differ: missing in candidate={sorted(left.keys() - right.keys())}; " + f"missing in baseline={sorted(right.keys() - left.keys())}" + ) + return left, right + + +# Standalone historical CLI scripts use the same validated parser as the runner. +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tests" / "regression")) +from checks import MetricError, load_metrics, compare_metrics # noqa: E402 + + +def load_legacy_metrics(path, *, dtype=None, required=("loss",)): + """Recover explicit resume bounds for retained historical test_config cases. + + Unknown files remain subject to the strict 1..declared-total contract. + Known checkpoint starts come from configuration, never a truncated log. + """ + import json + import re + + configured = Path(__file__).with_name("test_config.json") + expected = None + if configured.exists(): + config = json.loads(configured.read_text()) + for group in config.get("test_groups", []): + for case in group.get("tests", []): + if Path(path).stem in ( + case["id"], + *( + model + "_" + case["id"] + for model in ("gpt2", "llama3", "qwen3", "mixtral") + ), + ): + args = case.get("args", {}) + checkpoint = re.search( + r"(?:checkpoint_step_|iter_)(\d+)", str(args.get("load", "")) + ) + if checkpoint and "num_iteration" in args: + expected = range( + int(checkpoint[1]) + 1, int(args["num_iteration"]) + 1 + ) + return load_metrics(path, expected_steps=expected, dtype=dtype, required=required) diff --git a/scripts/run_all_models.sh b/scripts/run_all_models.sh new file mode 100755 index 000000000..0e080c025 --- /dev/null +++ b/scripts/run_all_models.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +# Familiar all-model entry; explicit suite selection remains available. +set -euo pipefail +repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +exec "$repo_root/scripts/run_all_tests.sh" --suite nightly "$@" diff --git a/scripts/run_all_tests.sh b/scripts/run_all_tests.sh new file mode 100755 index 000000000..cf892883e --- /dev/null +++ b/scripts/run_all_tests.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -euo pipefail +repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +exec python3 "$repo_root/tests/regression/runner.py" "$@" diff --git a/scripts/run_models_and_profile.bash b/scripts/run_models_and_profile.bash index 84c6351f6..e4aa8fd98 100755 --- a/scripts/run_models_and_profile.bash +++ b/scripts/run_models_and_profile.bash @@ -1,613 +1,61 @@ -#!/bin/bash - -set -e -set -o pipefail - +#!/usr/bin/env bash +# Compatibility entry for classic groups; all execution uses one runner. +set -euo pipefail +repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)" +config="$repo_root/scripts/test_config.json" +only_run="" +only_set=false +forward=() usage() { cat <<'EOF' -Usage: run_models_and_profile.bash [--test-config path] [--only-run tag1,tag2] +Usage: run_models_and_profile.bash [--test-config PATH] [--only-run TAGS] [runner options] -Options: - --test-config PATH Path to test config JSON. Default: test_config.json. - --only-run TAGS Only run the specified tag groups, separated by commas. - -h, --help Show this help message. +Runs migrated classic tests through scripts/run_all_tests.sh --suite nightly. +--only-run basic,zero,... selects the UNION of legacy groups. +--test-config accepts only the repository's original scripts/test_config.json. +New/custom configurations belong in tests/regression/scenarios/*.json. +All other options are passed to run_all_tests.sh (for example --list, --dry-run, +--model llama3, --data-root PATH, --build-dir PATH, --no-build). + +Builds and checkpoints are retained. Profiling is no longer run automatically. +See tests/regression/README.md for suites, fixture paths and log-history alignment. EOF } - -CONFIG_FILE="test_config.json" -ONLY_RUN_TAGS="" - while [[ $# -gt 0 ]]; do case "$1" in - --test-config) - [[ $# -lt 2 ]] && { echo "Error: --test-config requires a file path."; exit 1; } - CONFIG_FILE="$2" - shift 2 - ;; - --test-config=*) - CONFIG_FILE="${1#*=}" - shift - ;; - --only-run) - [[ $# -lt 2 ]] && { echo "Error: --only-run requires a comma-separated tag list."; exit 1; } - ONLY_RUN_TAGS="$2" - shift 2 - ;; - --only-run=*) - ONLY_RUN_TAGS="${1#*=}" - shift - ;; - -h|--help) - usage - exit 0 - ;; - -*) - echo "Error: Unknown option: $1" - usage - exit 1 - ;; - *) - echo "Error: Unknown positional argument: $1" - usage - exit 1 - ;; + --test-config|--only-run) + [[ $# -ge 2 && "$2" != --* && -n "$2" ]] || { echo "Error: $1 requires a value." >&2; exit 2; } + if [[ "$1" == --test-config ]]; then config="$2"; else only_run="$2"; only_set=true; fi + shift 2 ;; + --test-config=*) config="${1#*=}"; shift ;; + --only-run=*) only_run="${1#*=}"; only_set=true; shift ;; + -h|--help) usage; exit 0 ;; + *) forward+=("$1"); shift ;; esac done - -# Dependencies check -if ! command -v jq >/dev/null 2>&1; then - echo "Error: jq is required. Install with: sudo apt-get install -y jq" - exit 1 -fi - -# Read variables -read_var() { - local key="$1" - jq -r --arg k "$key" '.variables[$k] // empty' "$CONFIG_FILE" -} - -BUILD_DIR="$(read_var BUILD_DIR)"; : "${BUILD_DIR:=../build}" -LOG_DIR="$(read_var LOG_DIR)"; : "${LOG_DIR:=logs}" -PROFILE_LOG_DIR="$(read_var PROFILE_LOG_DIR)"; : "${PROFILE_LOG_DIR:=./profile_logs}" -COMPARE_LOG_DIR="$(read_var COMPARE_LOG_DIR)"; : "${COMPARE_LOG_DIR:=}" -RUN_CTEST="$(read_var RUN_CTEST)"; : "${RUN_CTEST:=true}" -RUN_PROFILE_TEST="$(read_var RUN_PROFILE_TEST)"; : "${RUN_PROFILE_TEST:=true}" -CKPT_ROOT_DIR="$(read_var CKPT_ROOT_DIR)"; : "${CKPT_ROOT_DIR:=/data1/ckpt}" -MIXTRAL_INPUT_BIN="$(read_var MIXTRAL_INPUT_BIN)"; : "${MIXTRAL_INPUT_BIN:=/data/shared/InfiniTrain-dev/data/llmc/llama3/tinyshakespeare/tiny_shakespeare_train.bin}" -MIXTRAL_LLMC_FILEPATH="$(read_var MIXTRAL_LLMC_FILEPATH)"; : "${MIXTRAL_LLMC_FILEPATH:=/data/shared/InfiniTrain-dev/data/llmc/mixtral/mixtral_megatron_export.bin}" -GPT2_TEST_GROUPS="$(read_var GPT2_TEST_GROUPS)"; : "${GPT2_TEST_GROUPS:=basic,zero,lora,checkpoint}" -LLAMA3_TEST_GROUPS="$(read_var LLAMA3_TEST_GROUPS)"; : "${LLAMA3_TEST_GROUPS:=basic,zero,lora,checkpoint}" -QWEN3_TEST_GROUPS="$(read_var QWEN3_TEST_GROUPS)"; : "${QWEN3_TEST_GROUPS:=qwen3}" -MIXTRAL_TEST_GROUPS="$(read_var MIXTRAL_TEST_GROUPS)"; : "${MIXTRAL_TEST_GROUPS:=moe}" -DEVICE_BACKEND="$(read_var DEVICE_BACKEND)"; : "${DEVICE_BACKEND:=cuda}" - -# export custom variables from config first. LOG_DIR/PROFILE_LOG_DIR are normalized below. -while IFS="=" read -r k v; do - [[ -z "$k" || "$k" == "null" ]] && continue - export "$k"="$v" -done < <(jq -r '.variables | to_entries[] | "\(.key)=\(.value)"' "$CONFIG_FILE") - -# Global variable to save the last cmake command -LAST_CMAKE_CMD="" -declare -A SELECTED_TAGS=() - -RUN_STARTED_AT="$(date '+%Y-%m-%d %H:%M:%S')" -RUN_ID="$(date '+%Y%m%d_%H%M%S')" -RUN_DATE="$(date '+%Y%m%d')" -GIT_BRANCH="$(git rev-parse --abbrev-ref HEAD 2>/dev/null || true)" -: "${GIT_BRANCH:=unknown}" -GIT_COMMIT_FULL="$(git rev-parse HEAD 2>/dev/null || true)" -: "${GIT_COMMIT_FULL:=unknown}" -GIT_COMMIT_SHORT="${GIT_COMMIT_FULL:0:7}" -SAFE_GIT_BRANCH="${GIT_BRANCH//\//_}" -SAFE_GIT_BRANCH="${SAFE_GIT_BRANCH//[[:space:]]/_}" -SAFE_GIT_BRANCH="$(printf '%s' "$SAFE_GIT_BRANCH" | tr -cd '[:alnum:]_.-')" -: "${SAFE_GIT_BRANCH:=unknown}" - -LOG_DIR_PARENT="$(dirname "$LOG_DIR")" -if [[ "$LOG_DIR_PARENT" == "." ]]; then - RUN_OUTPUT_DIR="${RUN_DATE}/${SAFE_GIT_BRANCH}_${GIT_COMMIT_SHORT}" -else - RUN_OUTPUT_DIR="${LOG_DIR_PARENT}/${RUN_DATE}/${SAFE_GIT_BRANCH}_${GIT_COMMIT_SHORT}" +# Historically the default was relative to scripts/, even when no file was read. +if [[ "$config" == test_config.json ]]; then config="$repo_root/scripts/test_config.json"; fi +if [[ ! -f "$config" ]] || [[ "$(realpath -- "$config")" != "$(realpath -- "$repo_root/scripts/test_config.json")" ]]; then + echo "Error: custom --test-config files are not interpreted by the new runner." >&2 + echo "Migrate the configuration to tests/regression/scenarios/*.json; see tests/regression/schema/README.md." >&2 + exit 2 fi -LOG_DIR="${RUN_OUTPUT_DIR}/logs" -PROFILE_LOG_DIR="${RUN_OUTPUT_DIR}/profile_logs" - -mkdir -p "$BUILD_DIR" "$LOG_DIR" "$PROFILE_LOG_DIR" -export BUILD_DIR LOG_DIR PROFILE_LOG_DIR - -RUN_METADATA_FILE="${LOG_DIR}/run_metadata.log" -: > "$RUN_METADATA_FILE" -RUN_METADATA_FILE="$(realpath "$RUN_METADATA_FILE")" -{ - echo "[RUN_STARTED_AT] $RUN_STARTED_AT" - echo "[RUN_ID] $RUN_ID" - echo "[GIT_BRANCH] $GIT_BRANCH" - echo "[GIT_COMMIT] $GIT_COMMIT_FULL" - echo "[GIT_COMMIT_SHORT] $GIT_COMMIT_SHORT" - echo "[CONFIG_FILE] $CONFIG_FILE" - echo "[LOG_DIR] $(realpath "$LOG_DIR")" - echo "[PROFILE_LOG_DIR] $(realpath "$PROFILE_LOG_DIR")" -} > "$RUN_METADATA_FILE" -echo -e "\033[1;33mRun metadata:\033[0m $RUN_METADATA_FILE" -echo -e "\033[1;33mRun log dir:\033[0m $(realpath "$LOG_DIR")" -echo -e "\033[1;33mRun profile log dir:\033[0m $(realpath "$PROFILE_LOG_DIR")" - -normalize_tag() { - local raw="$1" - raw="${raw#"${raw%%[![:space:]]*}"}" - raw="${raw%"${raw##*[![:space:]]}"}" - printf '%s' "$raw" -} - -cmake_bool_value() { - if [[ "$1" == "true" ]]; then - printf 'ON' - else - printf 'OFF' - fi -} - -set_cmake_option() { - local cmd="$1" - local option="$2" - local value="$3" - local flag="-D${option}=${value}" - - if [[ "$cmd" =~ (^|[[:space:]])-D${option}= ]]; then - cmd="$(printf '%s' "$cmd" | sed -E "s#(^|[[:space:]])-D${option}=[^[:space:]]+#\1${flag}#")" - elif [[ "$cmd" == *" .."* ]]; then - cmd="${cmd/ ../ ${flag} ..}" - else - cmd="${cmd} ${flag}" - fi - - printf '%s' "$cmd" -} - -if [[ -n "$ONLY_RUN_TAGS" ]]; then - IFS=',' read -r -a requested_tags <<< "$ONLY_RUN_TAGS" - for raw_tag in "${requested_tags[@]}"; do - tag="$(normalize_tag "$raw_tag")" - [[ -z "$tag" ]] && continue - SELECTED_TAGS["$tag"]=1 - done - - if [[ ${#SELECTED_TAGS[@]} -eq 0 ]]; then - echo "Error: --only-run did not contain any valid tags." - exit 1 - fi -fi - -# Clean the build directory -clean_build_dir() { - echo -e "\033[1;31m[CLEAN] Removing all contents in: ${BUILD_DIR}\033[0m" - mkdir -p "$BUILD_DIR" - rm -rf "${BUILD_DIR:?}/"* -} - -# Clean checkpoint directories (called once at start of script) -clean_checkpoints() { - echo -e "\033[1;31m[CLEAN] Removing checkpoint directories from previous run\033[0m" - if [[ -d "$CKPT_ROOT_DIR" ]]; then - echo -e "\033[1;31m[CLEAN] Removing: ${CKPT_ROOT_DIR}\033[0m" - rm -rf "${CKPT_ROOT_DIR:?}" - fi -} - -run_ctest() { - local gpu_list=() - local cuda_tests=() - - if [[ -n "${CTEST_CUDA_GPUS:-}" ]]; then - IFS=',' read -r -a gpu_list <<< "$CTEST_CUDA_GPUS" - elif command -v nvidia-smi >/dev/null 2>&1; then - mapfile -t gpu_list < <(nvidia-smi --query-gpu=index --format=csv,noheader 2>/dev/null || true) - fi - - if [[ ${#gpu_list[@]} -eq 0 ]]; then - gpu_list=(0) - fi - - local filtered_gpu_list=() - local gpu - for gpu in "${gpu_list[@]}"; do - gpu="${gpu//[[:space:]]/}" - [[ -z "$gpu" ]] && continue - filtered_gpu_list+=("$gpu") - done - - if [[ ${#filtered_gpu_list[@]} -eq 0 ]]; then - filtered_gpu_list=(0) - fi - - ctest --output-on-failure -LE cuda -j"$(nproc)" - - mapfile -t cuda_tests < <(ctest -N -L cuda | sed -n 's/^ *Test *#[0-9][0-9]*: //p') - if [[ ${#cuda_tests[@]} -eq 0 ]]; then - return 0 - fi - - local worker_count="${#filtered_gpu_list[@]}" - local pids=() - local worker_idx - for ((worker_idx = 0; worker_idx < worker_count; worker_idx++)); do - ( - local worker_failed=0 - local test_idx="$worker_idx" - local test_name - local assigned_gpu="${filtered_gpu_list[$worker_idx]}" - - while ((test_idx < ${#cuda_tests[@]})); do - test_name="${cuda_tests[$test_idx]}" - echo "[CUDA GPU ${assigned_gpu}] ${test_name}" - if ! CUDA_VISIBLE_DEVICES="$assigned_gpu" ctest --output-on-failure -R "^${test_name}$" -j1; then - worker_failed=1 - fi - test_idx=$((test_idx + worker_count)) - done - - exit "$worker_failed" - ) & - pids+=("$!") - done - - local failed=0 - local pid - for pid in "${pids[@]}"; do - if ! wait "$pid"; then - failed=1 - fi +selection=() +if [[ "$only_set" == true ]]; then + [[ -n "${only_run//[[:space:],]/}" ]] || { echo 'Error: --only-run must contain at least one group.' >&2; exit 2; } + IFS=',' read -r -a tags <<< "$only_run" + for tag in "${tags[@]}"; do + tag="${tag//[[:space:]]/}" + case "$tag" in + recompute) + selection+=(--case "classic_recompute*/*") ;; + basic|zero|lora|checkpoint|lr_scheduler|8_proc|qwen3|qwen3_8_proc|moe) + selection+=(--case "classic_${tag}/*") ;; + *) echo "Error: unknown legacy group '$tag'." >&2; exit 2 ;; + esac done - - return "$failed" -} - -# Run a command and log output -run_and_log() { - local cmd="$1" - local log_name="$2" - local is_profile="$3" - local tag="${4:-basic}" - local timestamp - timestamp=$(date '+%Y-%m-%d %H:%M:%S') - local tag_log_dir="${LOG_DIR}/${tag}" - mkdir -p "$tag_log_dir" - local log_path="$(realpath "${tag_log_dir}/${log_name}.log")" - - echo -e "\033[1;32m============================================================\033[0m" - echo -e "\033[1;36m[$timestamp] [Running] ${log_name}\033[0m" - - # Print the command being executed - echo -e "\033[1;33mCommand:\033[0m $cmd" - - # Print the most recent CMake command - if [[ -n "$LAST_CMAKE_CMD" ]]; then - echo -e "\033[1;34mLast CMake Command:\033[0m $LAST_CMAKE_CMD" - fi - - echo -e "\033[1;33mLog file:\033[0m $log_path" - - # Notify if profiling mode is enabled - if [[ "$is_profile" == "yes" ]]; then - echo -e "\033[1;35m[PROFILE MODE ON] Profiling logs will be saved to: ${PROFILE_LOG_DIR}\033[0m" - fi - - echo -e "\033[1;32m============================================================\033[0m" - - pushd "$BUILD_DIR" > /dev/null - - # Write the last cmake command into the log file if available - if [[ -n "$LAST_CMAKE_CMD" ]]; then - echo "[LAST_CMAKE] $LAST_CMAKE_CMD" > "$log_path" - else - # If no cmake command has been run yet, clear the log - > "$log_path" - fi - - # Write the current run command to the log - echo "[RUN_METADATA] $RUN_METADATA_FILE" >> "$log_path" - echo "[RUN_STARTED_AT] $RUN_STARTED_AT" >> "$log_path" - echo "[GIT_BRANCH] $GIT_BRANCH" >> "$log_path" - echo "[GIT_COMMIT] $GIT_COMMIT_FULL" >> "$log_path" - echo "[GIT_COMMIT_SHORT] $GIT_COMMIT_SHORT" >> "$log_path" - local expanded_cmd="${cmd//\$LORA_WEIGHTS_DIR/$LORA_WEIGHTS_DIR}" - echo "[COMMAND] $expanded_cmd" >> "$log_path" - - # Run the command and append both stdout and stderr to the log file - if ! eval "$cmd" >> "$log_path" 2>&1; then - echo -e "\033[1;31m============================================================\033[0m" - echo -e "\033[1;31m[ERROR] Command failed: ${cmd}\033[0m" - echo -e "\033[1;31m[ERROR] See log file for details: ${log_path}\033[0m" - echo -e "\033[1;31m============================================================\033[0m" - echo "" - echo "[ERROR] Last 20 lines of log:" - tail -20 "$log_path" - exit 1 - fi - - popd > /dev/null - - # If profiling is enabled, move profiling files to the target directory - if [[ "$is_profile" == "yes" ]]; then - move_profile_logs "$log_name" "$tag" - fi -} - - -# Move profiling output logs -move_profile_logs() { - local prefix="$1" - local tag="${2:-basic}" - local tag_profile_dir="${PROFILE_LOG_DIR}/${tag}" - mkdir -p "$tag_profile_dir" - - # Move *.report.rankN files - for report_file in "${BUILD_DIR}"/*.report.rank*; do - if [[ -f "$report_file" ]]; then - local base_name - base_name=$(basename "$report_file") - mv "$report_file" "${tag_profile_dir}/${prefix}_${base_name}" - echo "Moved $base_name to ${tag_profile_dir}/${prefix}_${base_name}" - fi - done - - # Move *.records.log.rankN files - for record_file in "${BUILD_DIR}"/*.records.log.rank*; do - if [[ -f "$record_file" ]]; then - local base_name - base_name=$(basename "$record_file") - mv "$record_file" "${tag_profile_dir}/${prefix}_${base_name}" - echo "Moved $base_name to ${tag_profile_dir}/${prefix}_${base_name}" - fi - done -} - -# Build "--key=value" arg string from tests[i].args. -# For checkpoint-related args, automatically isolate by model and run mode -# (resume/no_resume) to avoid cross-test overwrites in one-click runs. -args_string_for_test() { - local group_idx="$1" - local test_idx="$2" - local model_name="$3" - local test_id="$4" - - jq -r --argjson g "$group_idx" --argjson t "$test_idx" --arg model "$model_name" --arg test_id "$test_id" ' - def namespaced_path($p; $model; $mode): - if ($p | test("/(?:checkpoint_step_|iter_)[0-9]+($|/)")) then - ($p | capture("^(?.*)/(?:checkpoint_step_|iter_)(?[0-9]+)(?/.*)?$")) as $m - | ("0000000" + $m.iteration)[-7:] as $iteration - | ($m.prefix + "/" + $model + "/" + $mode + "/iter_" + $iteration + ($m.suffix // "")) - else - ($p + "/" + $model + "/" + $mode) - end; - - .test_groups[$g].tests[$t].args as $args - | (if ($args | has("load")) then "resume" else "no_resume" end) as $run_mode - | (if (($args.load // "") | test("no_resume")) then "no_resume" else "resume" end) as $resume_src_mode - | $args - | (if has("save") then .save = namespaced_path(.save; $model; $run_mode) else . end) - | (if has("load") then .load = namespaced_path(.load; $model; $resume_src_mode) else . end) - | to_entries[] - | "--\(.key)=\(.value|tostring)" - ' "$CONFIG_FILE" | paste -sd' ' - | \ - sed "s|@CKPT_ROOT_DIR@|${CKPT_ROOT_DIR}|g" -} - -tag_enabled_for_model() { - local tag="$1" - local enabled_tags="$2" - - if [[ "$enabled_tags" == "*" ]]; then - return 0 - fi - - IFS=',' read -r -a tags <<< "$enabled_tags" - for raw_tag in "${tags[@]}"; do - local enabled_tag - enabled_tag="$(normalize_tag "$raw_tag")" - if [[ "$enabled_tag" == "$tag" ]]; then - return 0 - fi - done - return 1 -} - -model_has_selected_group() { - local enabled_tags="$1" - - for ((gi=0; gi&2 - exit 1 - fi - if [[ ! -f "$MIXTRAL_LLMC_FILEPATH" ]]; then - echo "Error: missing MIXTRAL_LLMC_FILEPATH: $MIXTRAL_LLMC_FILEPATH" >&2 - exit 1 - fi - fi -} - -infini_run_cmd_for_test() { - local model_bin="$1" - local input_bin="$2" - local llmc_filepath="$3" - local arg_str="$4" - local nproc_per_node="$5" - - printf './infini_run --nproc_per_node=%s %s --input_bin %q --llmc_filepath %q --device %q %s' \ - "$nproc_per_node" "$model_bin" \ - "$input_bin" "$llmc_filepath" "$DEVICE_BACKEND" "$arg_str" -} - -# Run tests -num_basic_compile_commands=$(jq '.basic_compile_commands | length' "$CONFIG_FILE") -num_groups=$(jq '.test_groups | length' "$CONFIG_FILE") - -if [[ "$num_basic_compile_commands" -eq 0 ]]; then - echo "Error: No basic compile commands found in basic_compile_commands." - exit 1 -fi - -selected_group_count=0 -for ((gi=0; gi&2 +exec "$repo_root/scripts/run_all_tests.sh" --suite nightly "${selection[@]}" "${forward[@]}" diff --git a/tests/regression/README.md b/tests/regression/README.md new file mode 100644 index 000000000..e1c743533 --- /dev/null +++ b/tests/regression/README.md @@ -0,0 +1,241 @@ +# 训练回归测试使用指南 + +默认流程是：选择套件 → 增量构建 → CTest 和模型训练 → 与上次成功运行的原始日志比较 → 保存本次记录。普通使用不需要导出完整张量或安装 PyTorch；参考数据直接来自历史运行目录。 + +## 一键运行 + +在仓库根目录执行： + +```bash +./scripts/run_all_tests.sh # 默认 PR 套件 +./scripts/run_all_tests.sh --suite smoke # 四个小模型冒烟 +./scripts/run_all_models.sh # nightly：经典模型和组合覆盖 +./scripts/run_all_tests.sh --suite perf # 基础/DDP 算子 profiling，PROFILE_MODE=ON +./scripts/run_all_tests.sh --suite release # 普通构建：nightly + 长程训练 +``` + +| 套件 | 默认日志模式的内容 | CUDA 训练任务数 | +|---|---|---:| +| smoke | 四个小模型，4 步训练 | 4 | +| pr | 小模型的梯度累积、并行、ZeRO、LoRA、学习率、checkpoint 及重计算配置 | 87 | +| nightly | PR 配置和 212 个经典任务（含重计算） | 299 | +| perf | 基础/DDP × FP32/BF16 的 kernel profiling | 14 | +| release | nightly 及四模型 200 步训练 | 303 | + +数量包含 checkpoint 必需的保存运行，实际值以 `--list` 为准。PR 最多需要四张 GPU，nightly/release 最多八张，perf 最多两张。CPU 支持合法 FP32 子集，perf 的 CPU 子集是四个单设备测例；不支持的 BF16、分布式组合会明确列在排除项中。原有 CTest 继续运行，`perf` 不运行 CTest。 + +默认不会生成单卡等价性参考、运行 PyTorch oracle 或保存梯度/参数快照。历史场景名字中的 `tiny_oracle` 在默认模式下也是普通训练任务。 + +## 日志怎么对齐 + +第一次没有可用历史时，训练成功的用例标为 `RECORDED`，保存日志，普通运行返回 0；这表示“已记录”，不表示已经完成回归比较。 + +后续自动选择同 suite/backend 的最近一次完整成功运行。筛选运行可使用同筛选范围或全套历史；全套运行不会拿一次局部运行充当完整参考。失败或未结束的运行不会成为参考。选择结果在本次运行开始时固定，报告会注明具体参考目录。 + +按 case ID 匹配后检查训练参数、数据/权重哈希和步数。checkpoint 的来源及恢复步数属于配置语义；临时输出目录变化不影响匹配。 + +| 指标 | 默认判定规则 | +|---|---| +| FP32 loss | 逐步绝对误差不超过 `1e-5` | +| BF16 loss | 逐步绝对误差不超过 `1e-2` | +| 吞吐 | 每次运行去掉首步，再取日志中 `tok/s` 的算术平均；降速超过 20% 失败 | +| 完整性 | 缺步骤、空日志、NaN/Inf、异常退出均失败 | + +性能比较要求硬件、后端、构建及软件环境信息兼容。速度提升不算失败，loss 错误也不能被速度提升抵消。不同配置或不可比环境显示 `NOT_COMPARABLE`，而非 PASS。新增 case 显示 `RECORDED`。 + +```bash +# 明确选择某次运行作为参考（两个参数名等价） +./scripts/run_all_tests.sh --reference-run artifacts/test-runs/OLD_RUN_ID +./scripts/run_all_tests.sh --compare-to artifacts/test-runs/OLD_RUN_ID + +# 使用独立的历史目录 +./scripts/run_all_tests.sh --history-root /test-history + +# 调整原有比较阈值 +./scripts/run_all_tests.sh --loss-atol-fp32 1e-5 --loss-atol-bf16 1e-2 \ + --performance-threshold 0.20 + +# 有意改变配置或输入后，显式重新记录历史 +./scripts/run_all_tests.sh --record-history +``` + +`--strict` 要求完整套件、包含单元测试，且每个用例都实际完成对齐；首次记录和新增 case 不能满足该门禁。`--record-history` 不能与 `--strict` 或显式参考目录组合。 + +## 每次保存什么 + +默认目录是 `artifacts/test-runs/-<唯一标识>/`，可用 `--output` 指定。已存在的运行目录不会被覆盖。 + +```text +/ + run.json # 时间、commit、环境及逐用例命令索引 + metadata.json # 源码/构建来源、环境、参考运行 + resolved_plan.json # 本次完整配置和参数模板 + fixtures.json # 实际输入文件及 SHA256 + summary.md # 首先查看这里：结果和失败原因 + summary.json # 机器可读结果 + junit.xml + coverage.json + cases//run-1/ + train.log # 原始 stdout/stderr，历史比较的数据来源 + metrics.normalized.json # 从日志解析的指标，便于查看 + cases//result.json +``` + +每个 `train.log` 都记录： + +- UTC 开始/结束时间、退出码； +- run ID、case ID； +- 当前 checkout 的完整 commit、dirty 标记和源码差异哈希； +- 可读的实际命令、无损的 JSON 参数数组及工作目录。 + +若存在构建清单，汇总还保存二进制对应的构建来源；尤其使用 `--no-build` 时，应区分当前源码和已编译二进制。每个用例的 `reproduce` 是从仓库根目录执行的复现命令，会保留 backend、实际输入、参数覆盖和参考设置。 + +`RECORDED` 与 `PASS` 在结果及覆盖报告中分开统计。原始日志和清单必须一起保留；可以整体移动运行目录，比较器会在选定目录内重新定位日志。 + +## 筛选和重跑 + +```bash +./scripts/run_all_tests.sh --list +./scripts/run_all_tests.sh --dry-run --explain +./scripts/run_all_tests.sh --model llama3 --tag zero2 +./scripts/run_all_tests.sh --case 'tiny_checkpoint/gpt2/*' +./scripts/run_all_tests.sh --no-build +./scripts/run_all_tests.sh --backend cpu +./scripts/run_all_tests.sh --rerun-failed artifacts/test-runs/RUN_ID/summary.json +./scripts/run_all_tests.sh --changed origin/main --dry-run --explain +``` + +重复 `--model`/`--case` 按并集选择，重复 `--tag` 要求同时满足。checkpoint 的生产者自动加入依赖；取消不必要的数值参考不会破坏保存/加载顺序。 + +`--changed` 保留固定锚点及受影响用例，公共或未知路径回退完整 PR。筛选、参数覆盖和 `--skip-unit` 属于部分验证，不能与严格全套门禁混用。 + +`--set learning_rate=0.0001` 使用 JSON 语法覆盖普通训练参数。输入文件改用 `--fixture ID=PATH`，以便校验实际内容。路径、batch 和拓扑的结构性修改应通过场景配置完成。 + +## 本机配置 + +小模型资产由标准库确定性生成，并校验 SHA256。经典模型需要本地权重和数据；根目录布局见 [经典资产清单](fixtures/classic.json)。 + +仓库根目录可放置不提交 Git 的 `tests.local.json`： + +```json +{ + "backend": "cuda", + "docker_container": "infinitrain-dev", + "data_root": "/datasets/llmc", + "checkpoint_root": "/fast-scratch/infinitrain-regression", + "history_root": "/test-history", + "gpus": "0,1,2,3,4,5,6,7" +} +``` + +不用容器时省略 `docker_container`。还支持 `build_dir`、`fixtures` 和 `reference_run`。显式命令行路径优先；经典数据也可通过 `INFINITRAIN_TEST_DATA_ROOT` 指定。跨宿主机/容器的路径应在执行环境中可访问。 + +普通构建默认位于 `.cache/regression/build-`，perf 使用独立的 `build--profile`。本机配置中的 build_dir 也会为 perf 自动追加 `-profile`，显式 `--build-dir` 则使用指定目录。`--no-build` 不会更新旧二进制；perf 要求构建清单和缓存都证明 PROFILE_MODE=ON，且二进制哈希匹配。普通套件拒绝 PROFILE_MODE=ON 的构建。`--jobs` 控制并发，GPU 按用例整体预留，perf 在本 runner 内串行并独占预留设备。 + +## perf：kernel profiling + +perf 用于评估框架 dispatcher 中的算子调用性能。默认使用经典模型权重，运行一次、10 步，micro batch=1、sequence length=64、梯度累积=1。 + +| 模型 | 单卡基础版本 | DDP 版本 | 精度 | +|---|---|---|---| +| GPT-2、Llama3、Qwen3 | 各 2 项 | 各 2 项，DP=2、线程启动 | FP32 / BF16 | +| Mixtral | 2 项 | 本套件暂未纳入；分布式入口已由 master 支持 | FP32 / BF16 | + +DDP 保持每张卡的 micro batch 和算子 shape，与基础版本相同；global batch 随 DP 数增加。TP/PP/VPP 均为1，SP、ZeRO、LoRA、重计算关闭。Qwen3 不再用 TP8 替代基础版本，单卡和 DDP 都需要每卡容纳完整模型及训练状态，显存不足时不会偷偷换成 TP 或小模型。 + +```bash +./scripts/run_all_tests.sh --suite perf --list +./scripts/run_all_tests.sh --suite perf --model llama3 +./scripts/run_all_tests.sh --suite perf --tag basic +./scripts/run_all_tests.sh --suite perf --tag ddp +``` + +每个测例的 `cases//run-1/` 保存： + +```text +train.log +<模型>.report.rankN # 按 step/tag、算子汇总次数、Host/Device 耗时及占比 +<模型>.records.log.rankN # 每次算子调用记录 +profiles.json # 文件索引、SHA256、各设备/算子的次数和总/平均耗时 +``` + +原始 profiler 文件保留所有 tag 和步骤；`profiles.json` 的算子汇总排除首个训练步及初始化记录，其余9步按算子、设备分别汇总。缺少任意设备的 report/records、缺训练步或无有效调用会使测例失败，不能仅因训练退出码为0就通过。 + +现有 profiler 的 rank 表示设备索引；当前 DDP 测例使用同进程多线程，不把它冒充跨节点全局 rank。算子记录包含同步开销，Host 时间可能包含设备等待；一个 dispatcher 算子也可能包含多个底层 CUDA kernel。记录中的 Peak Mem 来自预留内存池峰值，不代表单个算子的独占显存。 + +perf 仍保留同配置历史日志的 loss/tok/s 比较，PROFILE_MODE 属于环境兼容检查；不能把普通构建吞吐与 profiling 吞吐直接比较。当前未设置单算子跨版本耗时门禁，算子报告用于分析和后续数据汇总。首次无历史仍为 RECORDED，算子产物完整性始终检查。 + +release 的303项使用普通构建；完整发布检查可分别运行 release 和 perf。未来接入 FlashAttention 后，在 kernel 场景中增加有代表性的开关配置即可，无需展开所有并行组合。后续同步导出仍是独立工作，本次没有实现飞书写入。 + +大型 checkpoint 应使用容量充足的本地临时盘。空间预检考虑 DP 副本、优化器状态和保留数量；每次运行使用带归属标记的独立子目录。成功训练的 checkpoint 默认清理，失败现场保留；`--keep-checkpoints` 可保留成功产物。日志和报告始终保留。 + +## 配置如何扩展 + +- `models/`：模型入口、默认参数、能力和资产； +- `scenarios/`：测试参数、有限矩阵、标签及所属套件； +- `suites/`：套件策略； +- `fixtures/`：数据与权重清单； +- `tests.local.json`：本机路径和设备。 + +普通新增组合只增加场景 JSON。同接口新模型只增加模型清单。合法两两组合、高风险锚点、batch 单位和完整示例见 [配置规范](schema/README.md)。 + +当前 ZeRO-1/2 要求 DP > 1，ZeRO-3 未实现；CPU 训练只支持 FP32。非法配置在启动前拒绝，不会隐式改成另一种测试。 + +## 重计算测例 + +重计算使用 `activation_recompute`、`recompute_granularity`、`recompute_method` 和 +`recompute_num_layers` 四个参数；当前覆盖 `full` 的 `uniform`/`block`，不支持 `selective`。 + +| 场景 | 覆盖 | 不开重计算的对应配置 | +|---|---|---| +| `classic_recompute` | master 的 GPT-2/Llama3 各 14 个任务 | `classic_basic` 的 2/3/6/7/8 及 `classic_zero` 的 3/8 distopt，同精度 | +| `classic_recompute_qwen3` | Qwen3 现有 7 个经典任务各增加 uniform2/block2，共 14 个 | `classic_qwen3` 中去掉 `_recompute_uniform2`/`_recompute_block2` 后缀的同名任务 | +| `tiny_recompute` | GPT-2/Llama3/Qwen3 × FP32/BF16 × uniform2/block2，共 12 个 | `tiny_accumulation` 的同模型、同精度任务 | + +Qwen3 的经典配置包含 TP8、TP8+SP、PP8、TP8+LoRA、DP2+TP2+PP2+VPP2+SP。 +重计算与对应关闭版本的数据、权重、batch、步数、精度、学习率及并行配置一致； +配置自测会比较全部有效训练参数，确保差异只限于四个重计算选项。 +新增经典测例采用显式模板保存与现有命令的一对一关系，未来基准配置变化会由自测提醒同步。 + +```bash +# 日常开发：12 个小模型重计算任务 +./scripts/run_all_tests.sh --suite pr --tag recompute +# 查看全部 42 个经典重计算任务 +./scripts/run_all_tests.sh --suite nightly --tag classic --tag recompute --list +# 只跑 Qwen3 的 14 个经典重计算任务 +./scripts/run_all_tests.sh --suite nightly --tag classic --tag recompute --model qwen3 +# 旧入口也支持新增组，包括 Qwen3 +./scripts/run_models_and_profile.bash --only-run recompute +# 小模型开启/关闭重计算的同次直接数值对齐（额外运行关闭版本） +./scripts/run_all_tests.sh --suite pr --tag recompute --alignment deep +``` + +默认模式仍把开启重计算的测例与上次相同配置的历史日志比较，不把关闭重计算的吞吐作为性能门槛。 +`deep` 的小模型参考只关闭重计算，保持 batch 和拓扑不变，比较 loss、梯度和参数。 +新的测例首次没有历史会显示 `RECORDED`；这不表示完成了历史回归比较。 + +## 保留的专项能力 + +张量等价性、独立 PyTorch oracle、学习率公式等检查可通过显式模式启用： + +```bash +./scripts/run_all_tests.sh --alignment deep --suite pr +``` + +该模式仍与历史运行日志比较 loss 和吞吐，同时执行场景声明的深度检查。它会增加本次运行中的参考训练和快照开销，包含 oracle 的环境需要 PyTorch。历史参考仍使用 `--reference-run` 或自动选择的运行目录。详细数值契约见 [测试资产说明](fixtures/README.md) 和 [容差策略](TOLERANCES.md)。 + +## 旧脚本和 CI + +旧入口 `scripts/run_models_and_profile.bash` 运行 212 个经典任务：原有 155 个、master 新增的 28 个 GPT-2/Llama3 重计算任务、14 个 Qwen3 重计算任务，以及新迁入的7个 Qwen3 多进程和8个 Mixtral EP/EDP 任务。`--only-run basic,zero` 仍是分组并集，`--only-run recompute` 包含三个模型的经典重计算任务。`scripts/test_config.json` 保留为 master 迁移对照，Qwen3 重计算等新增用例使用新场景目录。旧入口不再自动清空 build 或重复 profiling。 + +本次变更提供本地运行入口，CI 接入将另行提供。计划在 CI 中运行 Python 自测、纯 CPU C++ 构建、CTest 和四模型 smoke。没有稳定历史的托管 CPU 环境可使用 `--record-history`，只记录本次结果，不宣称完成精度或性能对齐。 + +后续 GPU CI 计划使用自托管 runner，并配置模型数据路径和独立于 checkout 的持久历史目录。接入时可通过 `--history-root` 指定历史目录，首次使用 `--record-history` 建立记录，后续使用 `--strict` 检查完整套件的历史对齐结果。 + +本地执行测试框架自测： + +```bash +python3 -m unittest discover -s tests/regression -p 'test_*.py' -v +``` + +这些自测验证配置、日志解析、历史选择、失败传播和执行流程；真实训练验收应查看对应运行的 `summary.json`。 diff --git a/tests/regression/TOLERANCES.md b/tests/regression/TOLERANCES.md new file mode 100644 index 000000000..faa033d4f --- /dev/null +++ b/tests/regression/TOLERANCES.md @@ -0,0 +1,78 @@ +# 数值容差策略 + +本文用于显式 `--alignment deep` 中的数值检查。历史日志比较沿用原脚本的绝对 loss 容差（FP32 `1e-5`、BF16 `1e-2`）及 20% 吞吐降速阈值，可通过命令行覆盖；日志比较不使用本策略注册表。 + +`tolerances.json` 集中定义测试使用的数值容差。新增普通策略或修改某个 backend 的特定指标只需要编辑 JSON，再在场景的检查项中选择策略。 + +## 默认行为 + +| 策略 | dtype | atol | rtol | +|---|---|---:|---:| +| `oracle`、`distributed_equivalence` | `float32` | 1e-5 | 1e-5 | +| `oracle`、`distributed_equivalence` | `bfloat16` | 1e-2 | 1e-2 | +| `lr_schedule` | 两种 dtype | 1e-9 | 1e-5 | + +判定公式为 `abs(actual-reference) <= atol + rtol*abs(reference)`。这些默认值保留原检查器的数值行为。 + +指标名称为 `loss`、`lr`、`consumed_samples`、`parameter`、`gradient`、`optimizer`。`consumed_samples` 和整数／布尔 tensor state 始终精确比较；optimizer 策略的浮点容差只作用于浮点状态。初始参数一致性和真正的 DP 副本一致性也仍要求精确比较。 + +## 配置结构与覆盖顺序 + +顶层必须包含 `schema_version: 1`、`backends` 和 `policies`。每个策略包含两种 dtype 的完整 `defaults`，可以增加 `metrics` 和 `backends`： + +1. 选择策略的 `defaults[dtype]`。 +2. 应用 `metrics[metric][dtype]`。 +3. 应用 `backends[backend][metric][dtype]`。 + +每个叶子都必须同时提供有限、非负的 `atol` 和 `rtol`。没有继承链或隐式策略合并。已登记 backend 缺少专项覆盖时使用通用配置；未知 backend、策略、指标、dtype、字段或重复 JSON 键会在规划阶段报错。默认 backend 为 `cpu`、`cuda` 和 `privateuseone`，`privateuse1` 是后者的别名。 + +例如，在现有 `policies` 中增加一个条目： + +```json +"attention_strict": { + "defaults": { + "float32": {"atol": 0.00001, "rtol": 0.00001}, + "bfloat16": {"atol": 0.01, "rtol": 0.01} + }, + "metrics": { + "loss": { + "float32": {"atol": 0.0000001, "rtol": 0.000001} + } + }, + "backends": { + "cuda": { + "gradient": { + "bfloat16": {"atol": 0.002, "rtol": 0.003} + } + } + } +} +``` + +这个条目让 FP32 loss 使用更严格的阈值，并单独设置 CUDA BF16 gradient；其他指标和 backend 沿用条目中的 defaults。具体阈值应由已知正确版本的重复测量及独立参考校准，经审查后写入配置。 + +## 在场景中选择 + +独立参考检查： + +```json +{"kind": "oracle", "tolerance": "attention_strict"} +``` + +场景生成的等价性参考： + +```json +"reference": { + "single_device": true, + "metrics": ["loss", "gradient", "parameter"], + "tolerance": "attention_strict" +} +``` + +`oracle` 和 `lr_schedule` 检查同样支持 `tolerance` 字段。不指定时分别使用其默认命名策略。性能的时间、显存和噪声阈值在场景的 `performance` 配置中设置。 + +用 `./scripts/run_all_tests.sh --alignment deep --suite pr --model gpt2 --tag accumulation --dry-run --explain` 查看解析后的 `tolerance_policies` 与 `resolved_tolerances`。 + +## 配置与历史匹配 + +规划器将实际生效的非默认策略值纳入 case 配置哈希。修改某个 backend 的指标阈值只会影响使用该值的 case。历史日志仍按配置匹配，不会把配置不同的运行当作已经完成对齐；具体参考运行与解析后的策略保存在本次运行产物中。 diff --git a/tests/regression/assets.py b/tests/regression/assets.py new file mode 100644 index 000000000..6e7bae707 --- /dev/null +++ b/tests/regression/assets.py @@ -0,0 +1,92 @@ +"""Resolve versioned fixtures and case-local artifacts without machine paths.""" + +import importlib.util +import json +from pathlib import Path + +from execution import file_sha256, safe_case_path + + +class Assets: + def __init__(self, root, cache_dir, data_root=None, overrides=None): + self.root = Path(root) + self.cache_dir = Path(cache_dir) + self.data_root = Path(data_root).expanduser().resolve() if data_root else None + self.overrides = overrides or {} + self.paths = {} + self.hashes = {} + self.manifest = json.loads( + (self.root / "fixtures" / "manifest.json").read_text() + ) + + def prepare(self, cases): + needed = sorted({fixture for case in cases for fixture in case["fixtures"]}) + if any(f.startswith("tiny.") for f in needed): + tiny_dir = ( + self.cache_dir + / "fixtures" + / ("v" + str(self.manifest["fixture_version"])) + ) + missing = any( + not (tiny_dir / v["path"]).exists() + for v in self.manifest["files"].values() + ) + if missing: + spec = importlib.util.spec_from_file_location( + "regression_generate_fixtures", + self.root / "fixtures" / "generate.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + module.generate(tiny_dir) + errors = [] + for fixture in needed: + try: + if fixture in self.overrides: + path = Path(self.overrides[fixture]).expanduser().resolve() + expected = ( + self.manifest.get("files", {}).get(fixture, {}).get("sha256") + ) + elif fixture.startswith("tiny."): + item = self.manifest["files"][fixture] + path, expected = tiny_dir / item["path"], item["sha256"] + elif fixture.startswith("classic."): + if self.data_root is None: + raise ValueError( + "set --data-root, INFINITRAIN_TEST_DATA_ROOT or fixture override" + ) + manifest_path = self.root / "fixtures" / "classic.json" + mapping = json.loads(manifest_path.read_text()) + item = mapping.get("files", mapping)[fixture] + path = self.data_root / item["path"] + expected = item.get("sha256") + else: + raise ValueError("unknown fixture ID") + if not path.is_file(): + raise ValueError("file not found: " + str(path)) + digest = file_sha256(path, self.cache_dir / "hashes") + if expected and expected != digest: + raise ValueError( + "checksum mismatch: expected {}, got {}".format( + expected, digest + ) + ) + self.paths[fixture], self.hashes[fixture] = str(path), digest + except (OSError, KeyError, ValueError) as error: + errors.append("{}: {}".format(fixture, error)) + return errors + + def resolve(self, value, artifact_root, checkpoint_root=None): + if not isinstance(value, str): + return value + if value.startswith("fixture:"): + return self.paths[value[len("fixture:") :]] + if value.startswith("artifact:"): + relative = value[len("artifact:") :] + root = ( + checkpoint_root + if checkpoint_root is not None and "/checkpoints" in relative + else artifact_root + ) + return str(safe_case_path(root, relative)) + return value diff --git a/tests/regression/checks.py b/tests/regression/checks.py new file mode 100644 index 000000000..d4a9ca7b1 --- /dev/null +++ b/tests/regression/checks.py @@ -0,0 +1,1134 @@ +"""Fail-closed numerical contracts for training regression artifacts. + +No third party packages are required. A PASS always means the requested data +was present, finite and compared; missing evidence never becomes an empty pass. +""" + +from __future__ import annotations + +import ast +import json +import io +import math +import re +import statistics +import struct +from dataclasses import dataclass, field +from pathlib import Path + + +class MetricError(ValueError): + """An artifact is incomplete, ambiguous, corrupt or incompatible.""" + + +DTYPES = { + "fp32": "float32", + "float32": "float32", + "bf16": "bfloat16", + "bfloat16": "bfloat16", +} +TOLERANCES = {"float32": (1e-5, 1e-5), "bfloat16": (1e-2, 1e-2)} +NUMBER = r"[+-]?(?:nan|inf(?:inity)?|(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)" +STEP = re.compile(r"\bstep\s+(\d+)/(\d+)\b", re.I) +FIELDS = { + "loss": rf"train loss\s+({NUMBER})", + "lr": rf"\blr\s+({NUMBER})", + "step_ms": rf"({NUMBER})\s+ms\b", + "tokens_per_second": rf"({NUMBER})\s+tok/s", +} + + +def result(status, reasons=(), **data): + return {"status": status, "reasons": list(reasons), "data": data} + + +def normalize_dtype(dtype): + if dtype not in DTYPES: + raise MetricError(f"Unknown dtype: {dtype!r}") + return DTYPES[dtype] + + +@dataclass +class Metrics: + records: list[dict] + dtype: str | None = None + metadata: dict = field(default_factory=dict) + + @classmethod + def from_dict(cls, data): + return cls(**data) + + def to_dict(self): + return {"records": self.records, "dtype": self.dtype, "metadata": self.metadata} + + +def _finite(value, label): + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + ): + raise MetricError(f"{label} must be a finite number: {value!r}") + + +def _validate(records, *, expected_steps=None, required=("loss",)): + if not records: + raise MetricError("No training metrics found") + found = [] + for record in records: + step = record.get("step") + if isinstance(step, bool) or not isinstance(step, int) or step < 0: + raise MetricError(f"Invalid global step: {step!r}") + if step in found: + raise MetricError( + f"Duplicate step {step}; rank streams must be reduced explicitly" + ) + found.append(step) + for key in ( + "rank", + "dp_rank", + "tp_rank", + "pp_rank", + "consumed_samples", + "global_tokens", + "dp", + "tp", + "pp", + "sp", + "vpp", + "world_size", + ): + minimum = ( + 1 + if key in ("global_tokens", "dp", "tp", "pp", "sp", "vpp", "world_size") + else 0 + ) + if key in record and ( + type(record[key]) is not int or record[key] < minimum + ): + raise MetricError(f"Step {step}: {key} must be an integer >= {minimum}") + for key in required: + if key not in record: + raise MetricError(f"Step {step}: missing required metric {key}") + if ( + "optimizer_stateful" in record + and type(record["optimizer_stateful"]) is not bool + ): + raise MetricError( + f"Step {step}: optimizer_stateful must be boolean metadata" + ) + for key, value in record.items(): + if key not in ( + "dtype", + "type", + "event", + "schema_version", + "optimizer_stateful", + ) and isinstance(value, (int, float)): + _finite(value, f"step {step}/{key}") + for key in required: + _finite(record[key], f"step {step}/{key}") + for key in ("step_ms", "tokens_per_second"): + if key in record and record[key] <= 0: + raise MetricError(f"Step {step}: {key} must be positive") + if found != sorted(found): + raise MetricError("Steps are out of order") + expected = ( + list(expected_steps) + if expected_steps is not None + else list(range(found[0], found[-1] + 1)) + ) + if found != expected: + raise MetricError( + f"Incomplete/unexpected steps: expected {expected}, found {found}" + ) + + +def load_metrics(path, *, expected_steps=None, dtype=None, required=("loss",)): + """Read canonical-rank JSONL or historical text logs. Raises MetricError. + + Text step denominators detect truncated logs even without expected_steps. + For resumed text streams, callers must supply explicit expected_steps. + """ + path = Path(path) + records, metadata, observed_dtypes, totals = [], {}, set(), set() + try: + lines = path.read_text().splitlines() + except (OSError, UnicodeError) as error: + raise MetricError(str(error)) from error + structured = path.suffix == ".jsonl" + for number, line in enumerate(lines, 1): + if not line.strip(): + continue + if structured or line.lstrip().startswith("{"): + try: + row = json.loads(line) + except (ValueError, TypeError) as error: + raise MetricError(f"{path}:{number}: invalid JSON") from error + if not isinstance(row, dict): + raise MetricError(f"{path}:{number}: metric must be an object") + if "schema_version" in row and ( + type(row["schema_version"]) is not int or row["schema_version"] != 1 + ): + raise MetricError(f"{path}:{number}: unsupported metric schema") + if row.get("type") not in (None, "step", "metadata", "config", "complete"): + raise MetricError(f"{path}:{number}: unsupported metric record type") + if row.get("type") in ("metadata", "config") or row.get("event") in ( + "metadata", + "config", + ): + metadata.update(row) + if "dtype" in row: + observed_dtypes.add(normalize_dtype(row["dtype"])) + continue + if row.get("type") == "complete": + if "complete" in metadata: + raise MetricError("Duplicate completion record") + metadata["complete"] = row + continue + if "complete" in metadata: + raise MetricError("Training records after completion") + if "step" not in row: + raise MetricError(f"{path}:{number}: missing step") + if "dtype" in row: + observed_dtypes.add(normalize_dtype(row["dtype"])) + records.append(row) + else: + match = STEP.search(line) + if not match: + continue + row = {"step": int(match[1])} + totals.add(int(match[2])) + for key, pattern in FIELDS.items(): + match_field = re.search(pattern, line, re.I) + if match_field: + row[key] = float(match_field[1]) + memory = re.search(r"peak used:\s*(" + NUMBER + r")\s*MB", line, re.I) + if memory: + row["peak_memory_bytes"] = float(memory[1]) * 1048576 + records.append(row) + if dtype is not None: + observed_dtypes.add(normalize_dtype(dtype)) + if len(observed_dtypes) > 1: + raise MetricError(f"Conflicting dtype metadata: {sorted(observed_dtypes)}") + if len(totals) > 1: + raise MetricError(f"Inconsistent declared training lengths: {sorted(totals)}") + if expected_steps is None and totals: + total = next(iter(totals)) + expected_steps = range(1, total + 1) + _validate(records, expected_steps=expected_steps, required=required) + return Metrics(records, next(iter(observed_dtypes), None), metadata) + + +def check_complete(metrics, expected_steps, required=("loss",)): + try: + _validate(metrics.records, expected_steps=expected_steps, required=required) + return result("PASS", steps=len(metrics.records)) + except MetricError as error: + return result("FAIL", [str(error)]) + + +def _tolerances(dtype, atol, rtol): + defaults = TOLERANCES[normalize_dtype(dtype)] + values = ( + defaults[0] if atol is None else atol, + defaults[1] if rtol is None else rtol, + ) + if any(not math.isfinite(v) or v < 0 for v in values): + raise MetricError("Tolerances must be finite and nonnegative") + return values + + +def compare_metrics( + reference, candidate, *, dtype, metrics=("loss",), atol=None, rtol=None, steps=None +): + try: + if not metrics: + raise MetricError("No requested comparison metrics") + dtype = normalize_dtype(dtype) + atol, rtol = _tolerances(dtype, atol, rtol) + for label, run in (("reference", reference), ("candidate", candidate)): + if run.dtype is not None and normalize_dtype(run.dtype) != dtype: + raise MetricError( + f"{label} dtype {run.dtype} differs from configured {dtype}" + ) + _validate(run.records, required=metrics) + ref = {r["step"]: r for r in reference.records} + cand = {r["step"]: r for r in candidate.records} + selected = list(steps) if steps is not None else list(ref) + if not selected: + raise MetricError("Empty comparison step selection") + if steps is None and ref.keys() != cand.keys(): + raise MetricError( + f"Step sets differ: reference {list(ref)}, candidate {list(cand)}" + ) + reasons, worst = [], 0.0 + for step in selected: + if step not in ref or step not in cand: + raise MetricError(f"Missing comparison step {step}") + for metric in metrics: + expected, actual = ref[step][metric], cand[step][metric] + difference = abs(actual - expected) + worst = max(worst, difference) + exact = metric in ( + "step", + "global_step", + "global_tokens", + "consumed_samples", + "data_position", + ) + limit = 0 if exact else atol + rtol * abs(expected) + if difference > limit: + reasons.append( + f"Step {step}/{metric}: {actual:.12g} vs {expected:.12g}; diff {difference:.4g} > {limit:.4g}" + ) + return result( + "FAIL" if reasons else "PASS", + reasons[:50], + compared_steps=len(selected), + max_absolute_error=worst, + mismatch_count=len(reasons), + ) + except (MetricError, KeyError, TypeError) as error: + return result("FAIL", [str(error)]) + + +def compare_resume( + reference, resumed, *, dtype, resume_step, snapshot_paths=None, atol=None, rtol=None +): + selected = [r["step"] for r in reference.records if r["step"] > resume_step] + if [r["step"] for r in resumed.records] != selected: + return result( + "FAIL", + [ + "Resumed stream must contain exactly the continuous run steps after the checkpoint" + ], + ) + checked = compare_metrics( + reference, + resumed, + dtype=dtype, + metrics=("loss", "lr", "consumed_samples"), + steps=selected, + atol=atol, + rtol=rtol, + ) + if snapshot_paths is None: + return result( + "FAIL", + checked["reasons"] + + ["Resume requires optimizer/parameter snapshot evidence"], + ) + state = compare_snapshots( + *snapshot_paths, + dtype=dtype, + required_kinds=("parameter", "optimizer"), + atol=atol, + rtol=rtol, + ) + return result( + "PASS" if checked["status"] == state["status"] == "PASS" else "FAIL", + checked["reasons"] + state["reasons"], + metrics=checked, + snapshots=state, + ) + + +def _read_npy(path, data=None, *, with_dtype=False): + """Decode one complete NPY record, preserving integer state without rounding. + + Bundled callers supply a bounded bytes/memoryview slice; old callers keep + using a path. No numpy dependency is required for checking artifacts. + """ + raw = Path(path).read_bytes() if data is None else data + stream = io.BytesIO(raw) + if stream.read(6) != b"\x93NUMPY": + raise MetricError(f"Not a NPY record: {path}") + version = stream.read(2) + if len(version) != 2 or version[0] not in (1, 2, 3): + raise MetricError(f"Invalid NPY version: {path}") + length_bytes = 2 if version[0] == 1 else 4 + raw_length = stream.read(length_bytes) + if len(raw_length) != length_bytes: + raise MetricError(f"Truncated NPY header: {path}") + length = int.from_bytes(raw_length, "little") + if length > 1_000_000: + raise MetricError("Unreasonable NPY header size") + raw_header = stream.read(length) + if len(raw_header) != length: + raise MetricError(f"Truncated NPY header: {path}") + header = ast.literal_eval( + raw_header.decode("utf8" if version[0] == 3 else "latin1") + ) + if not isinstance(header, dict) or header.get("fortran_order") is not False: + raise MetricError("NPY header must describe a contiguous C-order tensor") + shape = tuple(header["shape"]) + if any(type(size) is not int or size < 0 for size in shape): + raise MetricError("NPY dimensions must be nonnegative integers") + descriptor = header["descr"] + suffixes = { + "f2": ("e", "float16"), + "f4": ("f", "float32"), + "f8": ("d", "float64"), + "i1": ("b", "int8"), + "i2": ("h", "int16"), + "i4": ("i", "int32"), + "i8": ("q", "int64"), + "b1": ("?", "bool"), + "u1": ("B", "uint8"), + "u2": ("H", "uint16"), + "u4": ("I", "uint32"), + "u8": ("Q", "uint64"), + } + if ( + not isinstance(descriptor, str) + or len(descriptor) != 3 + or descriptor[0] not in "<>|=" + or descriptor[1:] not in suffixes + ): + raise MetricError(f"Unsupported NPY dtype {descriptor}") + code, storage_dtype = suffixes[descriptor[1:]] + if descriptor[0] == "|" and descriptor[2] != "1": + raise MetricError( + "Byte-order-independent NPY dtype is only valid for one-byte values" + ) + byte_order = descriptor[0] if descriptor[0] in "<>=" else "<" + fmt = byte_order + code + payload = stream.read() + if len(payload) != math.prod(shape) * struct.calcsize(fmt): + raise MetricError(f"Truncated/extra NPY payload: {path}") + values = [value[0] for value in struct.iter_unpack(fmt, payload)] + if storage_dtype.startswith("float"): + for value in values: + _finite(value, str(path)) + return (shape, values, storage_dtype) if with_dtype else (shape, values) + + +def _join_shards(parts, axis, part_sizes=None): + shapes = [shape for shape, _ in parts] + if not shapes or axis < 0 or axis >= len(shapes[0]): + raise MetricError("Invalid tensor shard axis") + base = list(shapes[0]) + for shape in shapes: + if len(shape) != len(base) or any( + shape[i] != base[i] for i in range(len(base)) if i != axis + ): + raise MetricError("Incompatible shard shapes") + count = len(parts) + base[axis] = sum(shape[axis] for shape in shapes) + inner, outer = math.prod(base[axis + 1 :]), math.prod(base[:axis]) + if part_sizes is None: + part_sizes = [base[axis]] + if sum(part_sizes) != base[axis] or any(size % count for size in part_sizes): + raise MetricError( + "Packed shard part_sizes must sum to global axis and divide by TP" + ) + local_sizes = [size // count for size in part_sizes] + if any(shape[axis] != sum(local_sizes) for shape in shapes): + raise MetricError("Packed local shard axis size mismatch") + merged = [] + for outer_index in range(outer): + offset = 0 + for local_size in local_sizes: + for shape, values in parts: + begin = (outer_index * shape[axis] + offset) * inner + merged.extend(values[begin : begin + local_size * inner]) + offset += local_size + return tuple(base), merged + + +def _load_snapshot( + path, expected_world_size=None, replica_tolerance=(0, 0), allowed_kinds=None +): + path = Path(path) + manifests = sorted(path.rglob("manifest.json")) if path.is_dir() else [path] + if not manifests: + raise MetricError(f"No snapshot manifests: {path}") + grouped, steps, documents, ranks, tensor_dtypes = {}, set(), [], set(), {} + bundle_cache = {} + for manifest in manifests: + manifest_parent = manifest.parent.resolve() + resolved_paths = {} + document = json.loads(manifest.read_text()) + if document.get("schema_version") != 1 or not isinstance( + document.get("tensors"), list + ): + raise MetricError("Unsupported snapshot manifest") + steps.add(document.get("step")) + rank = document.get("rank", 0) + if rank in ranks: + raise MetricError("Duplicate rank manifest; select one snapshot step") + ranks.add(rank) + documents.append(document) + local_keys = set() + local_tensors = { + (tensor["kind"], tensor["name"]): tensor + for tensor in document["tensors"] + if allowed_kinds is None or tensor["kind"] in allowed_kinds + } + bundle_ranges = {} + + def alias_root(key): + seen = set() + while local_tensors[key].get("alias_of"): + if key in seen: + raise MetricError("Cyclic snapshot alias metadata") + seen.add(key) + key = (key[0], local_tensors[key]["alias_of"]) + if key not in local_tensors: + raise MetricError("Snapshot alias points to a missing tensor") + return key + + for tensor in document["tensors"]: + if allowed_kinds is not None and tensor["kind"] not in allowed_kinds: + continue + key = (tensor["kind"], tensor["name"]) + if key in local_keys: + raise MetricError(f"Duplicate logical tensor {key} within rank {rank}") + local_keys.add(key) + raw_path = tensor["path"] + if raw_path not in resolved_paths: + tensor_path = (manifest_parent / raw_path).resolve() + if manifest_parent not in tensor_path.parents: + raise MetricError("Snapshot path escapes artifact directory") + resolved_paths[raw_path] = tensor_path + tensor_path = resolved_paths[raw_path] + bundled = "offset" in tensor or "length" in tensor + data = None + if bundled: + offset, length = tensor.get("offset"), tensor.get("length") + if ( + type(offset) is not int + or type(length) is not int + or offset < 0 + or length <= 0 + ): + raise MetricError(f"Invalid bundle offset/length: {key}") + if tensor_path not in bundle_cache: + bundle_cache[tensor_path] = tensor_path.read_bytes() + bundle = bundle_cache[tensor_path] + if offset + length > len(bundle): + raise MetricError(f"Bundle record exceeds file bounds: {key}") + data = memoryview(bundle)[offset : offset + length] + bundle_ranges.setdefault(tensor_path, []).append( + (offset, offset + length, key) + ) + shape, values, storage_dtype = _read_npy(tensor_path, data, with_dtype=True) + if "dtype" in tensor and tensor["dtype"] != storage_dtype: + raise MetricError(f"Manifest/NPY dtype mismatch: {key}") + if key in tensor_dtypes and tensor_dtypes[key] != storage_dtype: + raise MetricError(f"Tensor shard/replica dtypes differ: {key}") + tensor_dtypes[key] = storage_dtype + if list(shape) != tensor["shape"]: + raise MetricError(f"Manifest/NPY shape mismatch: {key}") + replica = (document.get("tp_rank", 0), document.get("pp_rank", 0)) + grouped.setdefault(key, []).append( + (tensor.get("shard"), shape, values, replica) + ) + for tensor_path, ranges in bundle_ranges.items(): + ordered = sorted(ranges) + for previous, current in zip(ordered, ordered[1:]): + if current[0] < previous[1]: + if current[:2] != previous[:2] or alias_root( + current[2] + ) != alias_root(previous[2]): + raise MetricError( + f"Overlapping snapshot bundle records: {tensor_path}" + ) + if len(steps) != 1 or not isinstance(next(iter(steps)), int): + raise MetricError("Snapshots must have one explicit global step") + declared_world_sizes = { + document["world_size"] for document in documents if "world_size" in document + } + if expected_world_size is not None: + declared_world_sizes.add(expected_world_size) + if len(declared_world_sizes) > 1: + raise MetricError("Snapshot world sizes disagree") + expected_ranks = next(iter(declared_world_sizes), max(ranks) + 1) + if ranks != set(range(expected_ranks)): + raise MetricError("Missing rank snapshot manifest") + if not grouped: + raise MetricError("Empty tensor snapshot") + tensors, replica_variants = {}, {} + for key, entries in grouped.items(): + shard_counts = { + entry[0].get("count", 1) if entry[0] else 1 for entry in entries + } + if len(shard_counts) != 1: + raise MetricError(f"Inconsistent shard count: {key}") + count = next(iter(shard_counts)) + if count < 1: + raise MetricError("Invalid shard count") + parts, descriptor, dp_copies, representatives = {}, None, {}, {} + for shard, shape, values, replica in entries: + index = shard.get("index", 0) if shard else 0 + if index < 0 or index >= count: + raise MetricError("Invalid shard index") + dp_key = (index, replica) + if dp_key in dp_copies and dp_copies[dp_key] != (shape, values): + raise MetricError( + f"Data-parallel replicas differ: {key}, shard {index}, TP/PP {replica}" + ) + dp_copies[dp_key] = (shape, values) + if index in parts and parts[index] != (shape, values): + previous_replica = representatives[index] + if ( + previous_replica[1] != replica[1] + or previous_replica[0] == replica[0] + ): + raise MetricError( + f"Replicated pipeline tensor differs: {key}, shard {index}" + ) + old_shape, old_values = parts[index] + atol, rtol = replica_tolerance + integral = tensor_dtypes[key] == "bool" or tensor_dtypes[ + key + ].startswith(("int", "uint")) + if old_shape != shape or any( + abs(a - b) > (0 if integral else atol + rtol * abs(a)) + for a, b in zip(old_values, values) + ): + raise MetricError( + f"Tensor-parallel replicas exceed configured tolerance: {key}, shard {index}" + ) + if index not in parts: + parts[index] = (shape, values) + representatives[index] = replica + if count > 1: + candidate = (shard["axis"], shard.get("part_sizes")) + if descriptor is not None and descriptor != candidate: + raise MetricError("Inconsistent shard layout") + descriptor = candidate + if set(parts) != set(range(count)): + raise MetricError(f"Missing tensor shards: {key}") + tensors[key] = ( + parts[0] + if count == 1 + else _join_shards([parts[i] for i in range(count)], *descriptor) + ) + if count == 1: + # Preserve each TP representative for direct oracle comparison. + # DP copies of a given TP/PP representative were checked bit-exact. + replica_variants[key] = list(dp_copies.values()) + optimizer_contracts = { + (document.get("optimizer_type"), document.get("optimizer_stateful")) + for document in documents + } + if len(optimizer_contracts) != 1: + raise MetricError("Snapshot ranks disagree on optimizer state contract") + optimizer_type, optimizer_stateful = next(iter(optimizer_contracts)) + if optimizer_stateful is False and any(key[0] == "optimizer" for key in tensors): + raise MetricError( + "Stateless optimizer manifest contradicts stored optimizer tensors" + ) + return { + "step": next(iter(steps)), + "ranks": sorted(ranks), + "optimizer_type": optimizer_type, + "optimizer_stateful": optimizer_stateful, + "tensor_dtypes": tensor_dtypes, + "replica_variants": replica_variants, + }, tensors + + +def compare_snapshots( + reference_path, + candidate_path, + *, + dtype, + atol=None, + rtol=None, + required_kinds=("parameter", "gradient"), + expected_world_sizes=(None, None), +): + try: + if not required_kinds or not set(required_kinds).issubset( + {"parameter", "gradient", "optimizer"} + ): + raise MetricError("Snapshot comparison needs known nonempty tensor kinds") + atol, rtol = _tolerances(dtype, atol, rtol) + ref_doc, reference = _load_snapshot( + reference_path, expected_world_sizes[0], (atol, rtol), set(required_kinds) + ) + cand_doc, candidate = _load_snapshot( + candidate_path, expected_world_sizes[1], (atol, rtol), set(required_kinds) + ) + if ref_doc.get("step") != cand_doc.get("step"): + raise MetricError("Snapshot global steps differ") + reference = { + key: value for key, value in reference.items() if key[0] in required_kinds + } + candidate = { + key: value for key, value in candidate.items() if key[0] in required_kinds + } + if reference.keys() != candidate.keys(): + raise MetricError("Logical tensor name/kind sets differ") + kinds = {key[0] for key in reference} + missing = set(required_kinds) - kinds + if "optimizer" in required_kinds: + if (ref_doc["optimizer_type"], ref_doc["optimizer_stateful"]) != ( + cand_doc["optimizer_type"], + cand_doc["optimizer_stateful"], + ): + raise MetricError("Optimizer state contracts differ") + if ( + ref_doc["optimizer_type"] == "sgd" + and ref_doc["optimizer_stateful"] is False + ): + missing.discard("optimizer") + if missing: + raise MetricError(f"Missing snapshot kinds: {sorted(missing)}") + reasons, count, worst = [], 0, 0.0 + for key, (shape, expected) in reference.items(): + storage_dtype = ref_doc["tensor_dtypes"][key] + if storage_dtype != cand_doc["tensor_dtypes"][key]: + raise MetricError( + f"Reference/candidate tensor storage dtypes differ: {key}" + ) + exact = storage_dtype == "bool" or storage_dtype.startswith(("int", "uint")) + expected_variants = ref_doc["replica_variants"].get( + key, [(shape, expected)] + ) + actual_variants = cand_doc["replica_variants"].get(key, [candidate[key]]) + for reference_shape, expected_values in expected_variants: + for variant_index, (actual_shape, actual) in enumerate(actual_variants): + if reference_shape != actual_shape: + raise MetricError(f"Tensor shape differs: {key}") + for index, (a, b) in enumerate(zip(expected_values, actual)): + difference = abs(a - b) + worst = max(worst, difference) + count += 1 + limit = 0 if exact else atol + rtol * abs(a) + if difference > limit and len(reasons) < 50: + reasons.append( + f"{key}[{index}] replica {variant_index}: {b!r} vs {a!r}" + ) + return result( + "FAIL" if reasons else "PASS", + reasons, + elements=count, + tensors=len(reference), + max_absolute_error=worst, + ) + except ( + MetricError, + OSError, + ValueError, + KeyError, + TypeError, + SyntaxError, + struct.error, + ) as error: + return result("FAIL", [str(error)]) + + +def performance_summary(runs, *, warmup=10, window=30, min_repeats=3): + if warmup < 0 or window < 1 or min_repeats < 1: + raise MetricError("Invalid performance sampling contract") + if len(runs) < min_repeats: + raise MetricError( + f"Need {min_repeats} independent repetitions, received {len(runs)}" + ) + medians, peaks, throughputs, within_jitter = [], [], [], [] + for run in runs: + if run.metadata.get("oracle") or any( + record.get("oracle") for record in run.records + ): + raise MetricError( + "Independent numerical oracle placeholder timings are not performance evidence" + ) + _validate(run.records, required=("step_ms",)) + selected = run.records[warmup : warmup + window] + if len(selected) != window: + raise MetricError(f"Need {warmup + window} steps per repetition") + times = [record["step_ms"] for record in selected] + median = statistics.median(times) + medians.append(median) + within_jitter.append((max(times) - min(times)) / median) + if all("peak_memory_bytes" in record for record in selected): + peaks.append(max(record["peak_memory_bytes"] for record in selected)) + if all("global_tokens" in record for record in selected): + throughputs.append( + sum(record["global_tokens"] for record in selected) + / (sum(times) / 1000) + ) + center = statistics.median(medians) + jitter = (max(medians) - min(medians)) / center + return { + "step_ms": center, + "repeat_medians": medians, + "jitter": jitter, + "within_run_jitter": max(within_jitter), + "peak_memory_bytes": max(peaks) if len(peaks) == len(runs) else None, + "tokens_per_second": ( + statistics.median(throughputs) if len(throughputs) == len(runs) else None + ), + } + + +def compare_performance( + reference_runs, + candidate_runs, + *, + warmup=10, + window=30, + min_repeats=3, + threshold=0.05, + max_jitter=0.05, + memory_threshold=0.05, + require_memory=True, +): + try: + for value in (threshold, max_jitter, memory_threshold): + _finite(value, "performance threshold") + if value < 0: + raise MetricError("Performance thresholds must be nonnegative") + options = dict(warmup=warmup, window=window, min_repeats=min_repeats) + ref, cand = performance_summary(reference_runs, **options), performance_summary( + candidate_runs, **options + ) + if require_memory and ( + ref["peak_memory_bytes"] is None or cand["peak_memory_bytes"] is None + ): + raise MetricError("Missing peak memory evidence") + if ( + max( + ref["jitter"], + cand["jitter"], + ref["within_run_jitter"], + cand["within_run_jitter"], + ) + > max_jitter + ): + return result( + "INCONCLUSIVE", + ["Performance noise exceeds calibrated jitter limit"], + reference=ref, + candidate=cand, + ) + slowdown = cand["step_ms"] / ref["step_ms"] - 1 + reasons = [] + if slowdown > threshold: + reasons.append( + f"Step time regression {slowdown:.2%} exceeds {threshold:.2%}" + ) + if ( + ref["peak_memory_bytes"] is not None + and cand["peak_memory_bytes"] is not None + ): + if ref["peak_memory_bytes"] <= 0 or cand["peak_memory_bytes"] <= 0: + raise MetricError("Peak memory must be positive") + growth = cand["peak_memory_bytes"] / ref["peak_memory_bytes"] - 1 + if growth > memory_threshold: + reasons.append( + f"Peak memory regression {growth:.2%} exceeds {memory_threshold:.2%}" + ) + return result( + "FAIL" if reasons else "PASS", + reasons, + reference=ref, + candidate=cand, + slowdown=slowdown, + ) + except (MetricError, ValueError, TypeError, KeyError) as error: + return result("FAIL", [str(error)]) + + +def load_distributed_metrics(prefix, *, expected_steps, dtype, world_size): + """Validate every rank and reduce critical-path timings; never average ranks. + + Loss must already be globally DP-reduced by the training program. Select + final pipeline stage and verify all tensor/data parallel replicas agree. + """ + if world_size < 1: + raise MetricError("world_size must be positive") + streams = [] + expected_steps = list(expected_steps) + if not expected_steps: + raise MetricError("Empty expected step list") + for rank in range(world_size): + path = Path(str(prefix) + f".rank_{rank}.jsonl") + run = load_metrics( + path, + expected_steps=expected_steps, + dtype=dtype, + required=("lr", "global_tokens", "step_ms", "consumed_samples"), + ) + complete = run.metadata.get("complete", {}) + if complete.get("rank") != rank or complete.get("step") != expected_steps[-1]: + raise MetricError(f"Rank {rank}: missing/mismatched completion record") + if ( + "consumed_samples" in complete + and complete["consumed_samples"] != run.records[-1]["consumed_samples"] + ): + raise MetricError(f"Rank {rank}: completion sample cursor mismatch") + for record in run.records: + if "world_size" in record and record["world_size"] != world_size: + raise MetricError( + f"Rank {rank}: configured world size disagrees with stream" + ) + if "dtype" not in record: + raise MetricError(f"Rank {rank}: missing explicit dtype") + if isinstance(record.get("rank"), bool) or record.get("rank") != rank: + raise MetricError(f"Rank {rank}: stream rank mismatch") + if record.get("pp_rank", 0) == record.get("pp", 1) - 1: + _finite(record.get("loss"), f"rank {rank}/loss") + streams.append(run) + combined = [] + for index, step in enumerate(expected_steps): + rows = [stream.records[index] for stream in streams] + for key in ( + "lr", + "global_tokens", + "consumed_samples", + "dtype", + "dp", + "tp", + "pp", + "sp", + ): + if len({row.get(key) for row in rows}) > 1: + raise MetricError(f"Step {step}: ranks disagree on {key}") + final = [row for row in rows if row.get("pp_rank", 0) == row.get("pp", 1) - 1] + if not final: + raise MetricError("No final pipeline rank loss") + loss = final[0]["loss"] + atol, rtol = TOLERANCES[normalize_dtype(dtype)] + if any(abs(row["loss"] - loss) > atol + rtol * abs(loss) for row in final): + raise MetricError( + f"Step {step}: globally reduced losses disagree across ranks" + ) + record = dict(final[0]) + record.pop("rank", None) + record["step_ms"] = max(row["step_ms"] for row in rows) + if all( + "peak_memory_bytes" in row and row["peak_memory_bytes"] is not None + for row in rows + ): + record["peak_memory_bytes"] = max(row["peak_memory_bytes"] for row in rows) + record["tokens_per_second"] = record["global_tokens"] / ( + record["step_ms"] / 1000 + ) + combined.append(record) + return Metrics( + combined, + normalize_dtype(dtype), + {"world_size": world_size, "all_ranks_complete": True}, + ) + + +def expected_learning_rate(step, args): + """Independent closed-form oracle for the LR used by optimizer step (1-based).""" + t = step - 1 + peak = float(args.get("learning_rate", 1e-4)) + minimum = float(args.get("min_lr", 0)) + style = args.get("lr_decay_style", "constant") + warmup = int(args.get("lr_warmup_iters", 0)) + decay = int(args.get("lr_decay_iters", 0)) or int(args["num_iteration"]) + initial = float(args.get("lr_warmup_init", 0)) + if style == "none": + return peak + if decay <= warmup or warmup < 0: + raise MetricError("Invalid LR schedule interval") + if t < warmup: + return initial + (peak - initial) * t / warmup + if style == "constant": + return peak + fraction = min((t - warmup) / (decay - warmup), 1) + if style == "linear": + return peak + (minimum - peak) * fraction + if style == "cosine": + return minimum + 0.5 * (peak - minimum) * (1 + math.cos(math.pi * fraction)) + if style == "inverse-square-root": + return ( + minimum + if t > decay + else max(minimum, peak * math.sqrt(max(warmup, 1) / max(t, 1))) + ) + raise MetricError(f"Unsupported LR style {style}") + + +def check_lr_schedule(metrics, args, *, atol=1e-9, rtol=1e-5): + try: + _validate(metrics.records, required=("lr",)) + errors = [] + for record in metrics.records: + expected = expected_learning_rate(record["step"], args) + if abs(record["lr"] - expected) > atol + rtol * abs(expected): + errors.append( + f"Step {record['step']}: LR {record['lr']} != analytical {expected}" + ) + return result("FAIL" if errors else "PASS", errors, steps=len(metrics.records)) + except (MetricError, ValueError, KeyError) as error: + return result("FAIL", [str(error)]) + + +def check_lora_frozen(initial_path, trained_path): + try: + _, initial = _load_snapshot(initial_path) + _, trained = _load_snapshot(trained_path) + initial = { + key: value for key, value in initial.items() if key[0] == "parameter" + } + trained = { + key: value for key, value in trained.items() if key[0] == "parameter" + } + if initial.keys() != trained.keys(): + raise MetricError("LoRA initial/final parameter sets differ") + adapters = [ + key for key in initial if re.search(r"lora[_\.]?[ab](?:$|\.)", key[1], re.I) + ] + frozen = [key for key in initial if key not in adapters] + if not adapters or not frozen: + raise MetricError( + "LoRA check requires both adapter and frozen base tensors" + ) + errors = [ + f"Frozen base parameter changed: {key[1]}" + for key in frozen + if initial[key] != trained[key] + ] + changed = [key[1] for key in adapters if initial[key] != trained[key]] + if not changed: + errors.append("No LoRA adapter parameter changed") + return result( + "FAIL" if errors else "PASS", + errors, + frozen_tensors=len(frozen), + changed_adapters=changed, + ) + except (MetricError, OSError, ValueError, KeyError, TypeError) as error: + return result("FAIL", [str(error)]) + + +def compare_snapshot_trajectories( + reference_root, + candidate_root, + *, + steps, + dtype, + required_kinds=("parameter", "gradient"), + expected_world_sizes=(None, None), + atol=None, + rtol=None, + check_initial=False, +): + """Compare complete logical tensors at every requested optimizer step. + + check_initial is for equivalent runs, not resumed runs (whose initial + snapshot is the loaded checkpoint step). It proves globally equal weights. + """ + steps = list(steps) + if not steps or len(steps) != len(set(steps)): + return result("FAIL", ["Snapshot trajectory needs a nonempty unique step list"]) + outcomes = [] + if check_initial: + outcomes.append( + compare_snapshots( + Path(reference_root) / "step_0", + Path(candidate_root) / "step_0", + dtype=dtype, + required_kinds=("parameter",), + expected_world_sizes=expected_world_sizes, + atol=0, + rtol=0, + ) + ) + for step in steps: + checked = compare_snapshots( + Path(reference_root) / f"step_{step}", + Path(candidate_root) / f"step_{step}", + dtype=dtype, + required_kinds=required_kinds, + expected_world_sizes=expected_world_sizes, + atol=atol, + rtol=rtol, + ) + checked["data"]["step"] = step + outcomes.append(checked) + reasons = [reason for outcome in outcomes for reason in outcome["reasons"]] + return result( + "PASS" if all(outcome["status"] == "PASS" for outcome in outcomes) else "FAIL", + reasons[:50], + steps=steps, + checks=outcomes, + ) + + +def check_memory_growth( + metrics, + *, + warmup=10, + window=20, + metric="peak_memory_bytes", + max_growth_bytes=64 * 1024 * 1024, + max_growth_fraction=0.1, + min_increasing_fraction=0.75, +): + """Detect sustained allocator-memory drift after warmup, not every leak. + + Per-step peaks are reset by training. Compare robust first/last windows, + then require continued final-window slope and mostly nondecreasing samples + before labeling a sustained leak. A one-time allocator pool expansion is + reported separately as INCONCLUSIVE if it exceeds the budget. + """ + try: + if ( + warmup < 0 + or window < 2 + or max_growth_bytes < 0 + or not 0 <= max_growth_fraction + ): + raise MetricError("Invalid memory growth contract") + if not 0 <= min_increasing_fraction <= 1: + raise MetricError("Invalid sustained-growth fraction") + _finite(max_growth_bytes, "memory growth bytes") + _finite(max_growth_fraction, "memory growth fraction") + _validate(metrics.records, required=(metric,)) + values = [row[metric] for row in metrics.records[warmup:]] + if len(values) < window * 2: + raise MetricError( + f"Memory drift needs warmup + two windows ({warmup + window * 2} steps)" + ) + if min(values) <= 0: + raise MetricError("Memory evidence must be positive") + initial, final = statistics.median(values[:window]), statistics.median( + values[-window:] + ) + growth = final - initial + allowance = max(max_growth_bytes, initial * max_growth_fraction) + differences = [right - left for left, right in zip(values, values[1:])] + fraction = sum(delta >= 0 for delta in differences) / len(differences) + slope = statistics.median( + [ + right - left + for left, right in zip(values[-window:], values[-window + 1 :]) + ] + ) + slope_limit = allowance / (len(values) - window) + data = dict( + metric=metric, + initial_window_median=initial, + final_window_median=final, + growth_bytes=growth, + allowance_bytes=allowance, + final_slope_bytes_per_step=slope, + slope_limit_bytes_per_step=slope_limit, + nondecreasing_fraction=fraction, + interpretation="allocator peak drift; does not establish absence of all leaks", + ) + if growth > allowance: + if slope > slope_limit and fraction >= min_increasing_fraction: + return result( + "FAIL", ["Sustained allocator memory growth exceeds budget"], **data + ) + return result( + "INCONCLUSIVE", + [ + "Memory grew beyond budget without a stable leak slope; inspect allocator expansion/noise" + ], + **data, + ) + return result("PASS", **data) + except (MetricError, ValueError, TypeError, KeyError) as error: + return result("FAIL", [str(error)]) diff --git a/tests/regression/environments/cuda.json b/tests/regression/environments/cuda.json new file mode 100644 index 000000000..66a6dbf65 --- /dev/null +++ b/tests/regression/environments/cuda.json @@ -0,0 +1,17 @@ +{ + "schema_version": 1, + "backend": "cuda", + "environment_class": "cuda-reference", + "exclusive_required_for_performance": true, + "minimum_gpus": 1, + "performance_identity_fields": [ + "gpu_name", + "gpu_count", + "driver", + "cuda", + "nccl", + "cpu", + "build_options" + ], + "fixture_root_environment": "INFINITRAIN_FIXTURE_ROOT" +} diff --git a/tests/regression/execution.py b/tests/regression/execution.py new file mode 100644 index 000000000..c2e8c85a6 --- /dev/null +++ b/tests/regression/execution.py @@ -0,0 +1,458 @@ +"""Process, build and resource management for local/CI regression runs. + +This module intentionally depends only on Python's standard library. It never +executes shell configuration and never deletes an existing build or checkpoint. +""" + +import contextlib +import datetime +import fcntl +import hashlib +import json +import os +from pathlib import Path +import platform +import shlex +import signal +import socket +import subprocess +import threading +import time +import uuid + +_active_processes = set() +_process_mutex = threading.Lock() + + +def terminate_all(): + with _process_mutex: + processes = list(_active_processes) + for process in processes: + _terminate_group(process) + + +def atomic_json(path, value): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + "." + uuid.uuid4().hex + ".tmp") + tmp.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n" + ) + os.replace(str(tmp), str(path)) + + +def unique_run_id(): + return ( + datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%dT%H%M%S") + + "-" + + uuid.uuid4().hex[:8] + ) + + +def safe_case_path(root, case_id): + parts = case_id.split("/") + if not parts or any(p in ("", ".", "..") or "\\" in p for p in parts): + raise ValueError("Unsafe case identifier: " + case_id) + candidate = Path(root).joinpath(*parts) + if Path(root).resolve() not in candidate.resolve().parents: + raise ValueError("Case path escapes run directory") + return candidate + + +def command_text(argv): + return " ".join(shlex.quote(str(a)) for a in argv) + + +def _terminate_group(process): + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + pass + try: + process.wait(timeout=1) + except subprocess.TimeoutExpired: + pass + # The leader may already have exited while a descendant ignored SIGTERM. + # Its original process group must still be killed before a GPU lease ends. + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + pass + process.wait() + + +def run_process(argv, *, cwd, log_path, timeout, env=None, headers=None): + """Run argv without a shell; timeout includes the entire process group.""" + argv = [str(a) for a in argv] + if not argv or timeout <= 0: + raise ValueError("A command and a positive timeout are required") + log_path = Path(log_path) + log_path.parent.mkdir(parents=True, exist_ok=True) + started = time.monotonic() + result = { + "argv": argv, + "cwd": str(cwd), + "log": str(log_path), + "status": "ERROR", + "headers": dict(headers or {}), + "started_at": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + with log_path.open("w") as log: + for key, value in result["headers"].items(): + log.write("[{}] {}\n".format(key, value)) + log.write("[COMMAND] " + command_text(argv) + "\n") + log.write("[ARGV_JSON] " + json.dumps(argv, ensure_ascii=False) + "\n") + log.write("[STARTED_AT] " + result["started_at"] + "\n") + log.write("[CWD] " + str(cwd) + "\n") + log.flush() + try: + process = subprocess.Popen( + argv, + cwd=str(cwd), + env=env, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + except OSError as exc: + result.update(reason=str(exc), returncode=None, elapsed_seconds=0) + log.write(str(exc) + "\n") + return result + with _process_mutex: + _active_processes.add(process) + try: + returncode = process.wait(timeout=timeout) + result.update( + returncode=returncode, status="PASS" if returncode == 0 else "FAIL" + ) + if returncode: + result["reason"] = "Process exited with status {}".format(returncode) + except subprocess.TimeoutExpired: + _terminate_group(process) + result.update( + returncode=process.returncode, + status="ERROR", + reason="timeout after {}s".format(timeout), + ) + except BaseException: + _terminate_group(process) + raise + finally: + # A failed launcher can exit before its rank processes. No child + # owned by this invocation may outlive the invocation's resources. + _terminate_group(process) + with _process_mutex: + _active_processes.discard(process) + result["elapsed_seconds"] = time.monotonic() - started + result["completed_at"] = datetime.datetime.now(datetime.timezone.utc).isoformat() + with log_path.open("a") as log: + log.write("[COMPLETED_AT] " + result["completed_at"] + "\n") + log.write("[EXIT_CODE] " + str(result.get("returncode")) + "\n") + return result + + +def probe(argv): + try: + output = subprocess.check_output(argv, stderr=subprocess.DEVNULL, timeout=15) + return output.decode("utf-8", errors="replace").strip() + except (OSError, subprocess.SubprocessError): + return None + + +def discover_gpus(): + visible = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible is not None: + return [ + s.strip() for s in visible.split(",") if s.strip() and s.strip() != "-1" + ] + output = probe(["nvidia-smi", "--query-gpu=index", "--format=csv,noheader"]) + return output.splitlines() if output else [] + + +def environment_fingerprint(repo, build_dir, backend): + cache = Path(build_dir) / "CMakeCache.txt" + options = {} + if cache.exists(): + for line in cache.read_text(errors="replace").splitlines(): + if ":" in line and "=" in line and not line.startswith(("//", "#")): + left, value = line.split("=", 1) + name = left.split(":", 1)[0] + if name in { + "CMAKE_BUILD_TYPE", + "CMAKE_CXX_COMPILER", + "CMAKE_CXX_FLAGS", + "CMAKE_CUDA_FLAGS", + "CMAKE_CUDA_COMPILER", + "USE_CUDA", + "USE_NCCL", + "USE_OMP", + "PROFILE_MODE", + }: + options[name] = value + gpu = probe( + [ + "nvidia-smi", + "--query-gpu=name,memory.total,driver_version", + "--format=csv,noheader", + ] + ) + cpu = "unknown" + cpuinfo = Path("/proc/cpuinfo") + if cpuinfo.exists(): + cpu = next( + ( + s.split(":", 1)[1].strip() + for s in cpuinfo.read_text().splitlines() + if s.startswith("model name") + ), + "unknown", + ) + result = { + "backend": backend, + "system": platform.system(), + "machine": platform.machine(), + "cpu": cpu, + "gpu_inventory": gpu.splitlines() if gpu else [], + "gpu_topology": ( + probe(["nvidia-smi", "topo", "-m"]) if backend == "cuda" else None + ), + "cuda_compiler": probe(["nvcc", "--version"]) if backend == "cuda" else None, + "compiler": probe(["c++", "--version"]), + "build_options": options, + "nccl": ( + probe(["pkg-config", "--modversion", "nccl"]) if backend == "cuda" else None + ), + "container_image": os.environ.get("INFINITRAIN_TEST_IMAGE", "unspecified"), + } + return result + + +def source_metadata(repo): + repo = Path(repo).resolve() + git = [ + "git", + "-c", + "safe.directory=" + str(repo), + "-c", + "safe.directory=" + str(Path(repo) / ".git"), + "-C", + str(repo), + ] + listed = probe( + git + + [ + "ls-files", + "-co", + "--exclude-standard", + "--", + "infini_train", + "example", + "tests", + "cmake", + "tools", + "CMakeLists.txt", + ] + ) + tree_hash = hashlib.sha256() + for relative in sorted(set((listed or "").splitlines())): + path = Path(repo) / relative + if relative.startswith("tests/regression/") or not path.is_file(): + continue + tree_hash.update(relative.encode() + b"\0" + file_sha256(path).encode() + b"\0") + return { + "commit": probe(git + ["rev-parse", "HEAD"]), + "dirty": bool(probe(git + ["status", "--porcelain"])), + "diff_sha256": hashlib.sha256( + (probe(git + ["diff", "HEAD"]) or "").encode() + ).hexdigest(), + "compiled_source_sha256": tree_hash.hexdigest(), + "hostname": socket.gethostname(), + } + + +def file_sha256(path, cache_dir=None): + """Hash bytes, recording an audit cache without trusting filesystem times. + + Network/container filesystems can preserve all stat fields across a content + change. The caller deduplicates fixture paths within one preflight; a new + invocation must verify bytes again to make its baseline identity trustworthy. + """ + path = Path(path).resolve() + stat = path.stat() + identity = [ + str(path), + stat.st_dev, + stat.st_ino, + stat.st_size, + stat.st_mtime_ns, + stat.st_ctime_ns, + ] + key = hashlib.sha256(json.dumps(identity).encode()).hexdigest() + cached = Path(cache_dir) / (key + ".json") if cache_dir else None + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): + digest.update(chunk) + value = digest.hexdigest() + if cached: + try: + atomic_json(cached, {"identity": identity, "sha256": value}) + except OSError: + # The cache is optional; network/container ownership may make it + # read-only even though the source file is readable. + pass + return value + + +class ResourcePool: + """Atomic GPU-set leases shared with other regression runner processes.""" + + def __init__(self, gpu_ids, lock_dir, wait_seconds=60): + if len(gpu_ids) != len(set(gpu_ids)): + raise ValueError("GPU identifiers must be unique") + self.gpu_ids = list(gpu_ids) + inventory = probe( + ["nvidia-smi", "--query-gpu=index,uuid", "--format=csv,noheader"] + ) + aliases = {} + for line in (inventory or "").splitlines(): + if "," in line: + index, gpu_uuid = (part.strip() for part in line.split(",", 1)) + aliases[index] = gpu_uuid + aliases[gpu_uuid] = gpu_uuid + self.identities = {gpu: aliases.get(gpu, gpu) for gpu in self.gpu_ids} + if len(set(self.identities.values())) != len(self.gpu_ids): + raise ValueError("GPU aliases refer to duplicate physical devices") + self.lock_dir = Path(lock_dir) + self.lock_dir.mkdir(parents=True, exist_ok=True) + self.wait_seconds = wait_seconds + self.mutex = threading.Lock() + + @contextlib.contextmanager + def lease(self, count, exclusive=False): + count = len(self.gpu_ids) if exclusive and count else count + if count < 0 or count > len(self.gpu_ids): + raise RuntimeError( + "Need {} GPUs; only {} configured".format(count, len(self.gpu_ids)) + ) + opened, selected = [], [] + deadline = time.monotonic() + self.wait_seconds + while count: + with self.mutex: + for gpu in self.gpu_ids: + name = hashlib.sha256(self.identities[gpu].encode()).hexdigest()[ + :16 + ] + handle = (self.lock_dir / ("gpu-" + name + ".lock")).open("a+") + try: + fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB) + opened.append(handle) + selected.append(gpu) + except BlockingIOError: + handle.close() + if len(opened) == count: + break + if len(opened) == count: + break + for handle in opened: + fcntl.flock(handle, fcntl.LOCK_UN) + handle.close() + opened, selected = [], [] + if time.monotonic() >= deadline: + raise RuntimeError( + "GPU lease timeout; requested devices are used by another regression run" + ) + time.sleep(0.2) + try: + yield selected + finally: + for handle in opened: + fcntl.flock(handle, fcntl.LOCK_UN) + handle.close() + + +@contextlib.contextmanager +def build_lock(build_dir): + path = Path(build_dir) + path.mkdir(parents=True, exist_ok=True) + with (path / ".regression-build.lock").open("a+") as handle: + fcntl.flock(handle, fcntl.LOCK_EX) + yield + + +def build_project( + repo, build_dir, backend, output_dir, jobs=8, timeout=3600, extra_options=None, + *, profile_mode=False, build_tests=True +): + if backend not in ("cpu", "cuda"): + raise ValueError( + "External backends require --no-build and a provider build directory" + ) + options = [ + "-DCMAKE_BUILD_TYPE=Release", + "-DBUILD_TEST=" + ("ON" if build_tests else "OFF"), + "-DPROFILE_MODE=" + ("ON" if profile_mode else "OFF"), + "-DUSE_CUDA=" + ("ON" if backend == "cuda" else "OFF"), + "-DUSE_NCCL=" + ("ON" if backend == "cuda" else "OFF"), + ] + options += list(extra_options or []) + with build_lock(build_dir): + before = source_metadata(repo) + if not before["commit"]: + return { + "status": "ERROR", + "reason": "Cannot verify source commit; check repository ownership and git availability", + "log": str(Path(output_dir) / "configure.log"), + } + result = run_process( + ["cmake", "-S", str(repo), "-B", str(build_dir)] + options, + cwd=repo, + log_path=Path(output_dir) / "configure.log", + timeout=timeout, + ) + if result["status"] != "PASS": + return result + result = run_process( + ["cmake", "--build", str(build_dir), "--parallel", str(jobs)], + cwd=repo, + log_path=Path(output_dir) / "build.log", + timeout=timeout, + ) + if result["status"] == "PASS": + after = source_metadata(repo) + if before["compiled_source_sha256"] != after["compiled_source_sha256"]: + result.update( + status="ERROR", + reason="Sources changed while building; rebuild before using provenance", + ) + return result + binaries = {} + for name in ("gpt2", "llama3", "qwen3", "mixtral", "infini_run"): + path = Path(build_dir) / name + if path.is_file(): + binaries[name] = file_sha256(path) + atomic_json( + Path(build_dir) / "build-manifest.json", + { + "schema_version": 1, + "source": after, + "binaries": binaries, + "environment": environment_fingerprint(repo, build_dir, backend), + }, + ) + return result + + +def verify_build_manifest(build_dir, binary_names, expected_commit=None): + path = Path(build_dir) / "build-manifest.json" + if not path.is_file(): + return {"verified": False, "reason": "No build-manifest.json"} + data = json.loads(path.read_text()) + if data.get("schema_version") != 1 or not data.get("source", {}).get("commit"): + raise ValueError("Invalid build provenance manifest: " + str(path)) + if expected_commit and data["source"]["commit"] != expected_commit: + raise ValueError("Reference commit differs from build provenance") + for name in binary_names: + if data.get("binaries", {}).get(name) != file_sha256(Path(build_dir) / name): + raise ValueError("Build binary changed after provenance recording: " + name) + return {"verified": True, "manifest": data, "path": str(path)} diff --git a/tests/regression/fixtures/README.md b/tests/regression/fixtures/README.md new file mode 100644 index 000000000..aac762f7a --- /dev/null +++ b/tests/regression/fixtures/README.md @@ -0,0 +1,113 @@ +# Tiny model fixtures + +`generate.py --output PATH` generates all four models and token streams offline, +using the Python standard library. `manifest.json` is the committed version and +SHA-256 contract. The runner verifies every generated file against it. Changing +fixture values requires an intentional manifest review; recording a new run does +not update this contract. The tiny files total less than 1.5 MB and are build artifacts, not source. + +All models use their production Transformer implementation and LLMC loader: + +| Model | Layers | Hidden / heads / KV heads | Vocabulary | FFN | Other | +|---|---:|---|---:|---:|---| +| GPT-2 | 4 | 32 / 4 / 4 | 128 | 128 | Learned positions, GELU, LayerNorm, shared initial embedding values | +| Llama3 | 4 | 32 / 4 / 4 | 128 | 64 | Interleaved RoPE, SwiGLU, RMSNorm | +| Qwen3 | 4 | 32 / 4 / 4 | 128 | 64 | Split-half RoPE, Q/K RMSNorm, SwiGLU | +| Mixtral | 4 | 32 / 4 / 4 | 128 | 64 per expert | Four experts, top-2 routing | + +Context length is 64; scenarios normally use 16. Four layers support PP2×VPP2, +and dimensions support TP2. Mixtral's application currently supports single-rank +training only. These fixtures measure correctness; their throughput does not +represent production-model performance. + +Each token file has 16,384 deterministic tokens. Successive batches differ, so +resume tests exercise the consumed-sample cursor. Weights use a specified 32-bit +integer recurrence and exactly representable binary float scales; generation does +not depend on NumPy, platform transcendental math or a Python PRNG version. + +The 32-element bias/norm tensors deliberately exercise the distributed +optimizer's 64-element alignment padding. This exposed a real framework defect: +`ParamAndGradBuffer::BuildBuckets` allocated aligned parameter starts, while +`ParamAndGradBucket` rebuilt its range map by adding unpadded parameter lengths. +ZeRO optimizer views consequently updated padding or the wrong parameter region. +The fix carries the actual offsets into each bucket. The focused +`OptimizerBucketPaddingTest` validates both ZeRO stages' ranges and independent +expected SGD updates across two shard boundaries. Production tensors with sizes +already divisible by 64 follow the same layout as before. + +## Independent numerical reference + +`reference.py` is an optional **PyTorch + NumPy** oracle used by `tiny_oracle` +in `--alignment deep` mode. Default history mode runs the training cases without +the oracle or tensor snapshots. +It independently decodes LLMC tensors, verifies the framework's complete initial +parameter inventory and values, and implements dense attention, normalizations, +rotary positions, dense expert routing and autograd using PyTorch operations. +It uses the production entrypoints’ SGD for GPT-2 and Adam for the other models. +GPT-2 currently retains tied parameter storage on CPU, while `Module::To` +materializes separate CUDA parameters. The oracle preserves the explicit +`alias_of` snapshot metadata; this known backend difference is not a claim that +the two initialization paths have identical weight-tying behavior. +It compares loss, complete gradients and updated parameters through the normal +snapshot comparison path. The suite uses **operation composition**: independently +compute forward/backward from the preceding verified candidate parameters, check +the gradients, then independently replay SGD/Adam using the captured candidate +gradients as optimizer inputs. Optimizer moments persist in the reference. The +preceding independent update must pass before its candidate value can be used as +the next forward input. Expected final parameters are never copied from the +candidate. This avoids amplifying harmless FP32 gradient cancellation near +Adam's epsilon into misleading update failures, without relaxing gradient or +parameter tolerances. Omitting `--candidate-snapshots` runs a fully independent +trajectory instead; this can be numerically stricter around tiny Adam gradients. It deliberately supports this fixed FP32 architecture, +single rank, no LoRA, and a constant learning rate. Other configurations use +equivalence, scheduler, checkpoint and historical-log checks. + +`--micro-batch-size` preserves gradient accumulation and the loader's complete +microbatch boundaries. In the 200-step stability case (microbatch two, GAS four, +sequence length 16), step 128 crosses the token fixture's epoch boundary. The +independent oracle handles that transition and drops the incomplete epoch tail; +`test_reference_batching.py` asserts concrete expected token and label sequences +at the boundary without depending on framework code or PyTorch. + +```bash +python3 tests/regression/fixtures/reference.py \ + --initial-snapshot artifacts/candidate/snapshots/step_0/rank_0/manifest.json \ + --tokens .cache/regression/fixtures/gpt2_tokens.bin \ + --candidate-snapshots artifacts/candidate/snapshots \ + --output artifacts/reference --steps 4 --batch-size 2 --sequence-length 16 \ + --learning-rate 0.0001 +``` + +The matching `_weights.bin` must be beside the token file. Reference +metrics identify `oracle: pytorch_operations` or `pytorch_trajectory`; their time measures the CPU reference +computation and GPU memory is zero. They must never be used as an InfiniTrain +performance reference. The reference snapshots contain parameters and +gradients, not framework optimizer serialization; checkpoint checks cover that +separately. + +## Runtime artifact contract + +All four model executables accept `--metrics_file PREFIX` and +`--regression_snapshot_dir DIRECTORY`; empty defaults disable their work. +Metrics go to `PREFIX.rank_R.jsonl`, one stream per rank. Each optimizer step +records full-precision JSON numbers, data progress and topology; each successful +rank finishes with a `type: complete` record. Loss is DP-averaged and defined only +on the last pipeline stage; other pipeline stages emit `loss: null`. Step time +includes optimizer completion via device synchronization. Distributed timing is +the maximum rank duration, calculated by the checker. Memory uses the existing +allocator API with 1 MiB resolution, explicitly recorded as +`memory_resolution_bytes: 1048576`. A sub-MiB tiny footprint can report zero; +this cannot establish the absence of small leaks. + +Snapshots live under `DIRECTORY/step_N/rank_R/manifest.json`; each rank writes +one `tensors.bin` bundle of NPY records. Manifest byte offsets and lengths avoid +thousands of separate filesystem opens on network storage. Readers retain +support for individual NPY files from earlier artifacts. Step zero contains +initial parameters; subsequent steps contain parameters, available gradients and +optimizer state. PP local layer names map to global layer names. TP metadata +describes local shapes, the split axis and packed Q/K/V or gate/up blocks. ZeRO-1/2 +reduced gradient shards are gathered only when snapshots are requested, without modifying +training's gradient bindings. Optimizer names include rank because they describe +the rank-local serialized state. NPY widens BF16/FP16 tensors to FP32 but preserves integer counters and FP64 +state exactly; training dtype remains in the manifest. Snapshot I/O is outside step timing, and performance cases disable +snapshots entirely. diff --git a/tests/regression/fixtures/classic.json b/tests/regression/fixtures/classic.json new file mode 100644 index 000000000..2ad56ea1e --- /dev/null +++ b/tests/regression/fixtures/classic.json @@ -0,0 +1,39 @@ +{ + "schema_version": 1, + "description": "Relative to --data-root (the llmc directory). Actual file SHA256 participates in baseline identity; no invented published checksums.", + "files": { + "classic.gpt2.tokens": { + "path": "gpt2/tinyshakespeare/tiny_shakespeare_train.bin" + }, + "classic.gpt2.weights": { + "path": "gpt2/gpt2_124M.bin" + }, + "classic.gpt2.lora": { + "path": "gpt2/gpt2_124M_lora_weights_rank8_alpha16.bin" + }, + "classic.llama3.tokens": { + "path": "llama3/tinyshakespeare/tiny_shakespeare_train.bin" + }, + "classic.llama3.weights": { + "path": "llama3/llama3.2_1B_fp32.bin" + }, + "classic.llama3.lora": { + "path": "llama3/llama3.2_1B_lora_weights_rank8_alpha16.bin" + }, + "classic.qwen3.tokens": { + "path": "qwen3/tinyshakespeare/tiny_shakespeare_train.bin" + }, + "classic.qwen3.weights": { + "path": "qwen3/qwen3-8b-fp32.llmc" + }, + "classic.qwen3.lora": { + "path": "qwen3/qwen3_8b_lora_weights_rank8_alpha16.bin" + }, + "classic.mixtral.tokens": { + "path": "llama3/tinyshakespeare/tiny_shakespeare_train.bin" + }, + "classic.mixtral.weights": { + "path": "mixtral/mixtral_megatron_export.bin" + } + } +} diff --git a/tests/regression/fixtures/generate.py b/tests/regression/fixtures/generate.py new file mode 100644 index 000000000..ab73de0cf --- /dev/null +++ b/tests/regression/fixtures/generate.py @@ -0,0 +1,206 @@ +#!/usr/bin/env python3 +"""Reproducible tiny LLMC fixtures using only the Python standard library. + +The integer PRNG and power-of-two float scale are part of fixture version 1. +Files load through the production checkpoint/data readers, with non-repeating +batches to exercise data-cursor resume. No network or model download is needed. +""" +import argparse +import hashlib +import json +import struct +from pathlib import Path + +MODELS = ("gpt2", "llama3", "qwen3", "mixtral") +CONFIG = dict( + block_size=64, + vocab_size=128, + n_layer=4, + n_head=4, + n_kv_head=4, + n_embd=32, + intermediate_size=64, +) + + +class Weights: + def __init__(self, stream): + self.stream, self.state = stream, 20261009 + + def write(self, count, kind="weight"): + data = [] + for _ in range(count): + self.state = (1664525 * self.state + 1013904223) & 0xFFFFFFFF + data.append( + 1.0 + if kind == "norm" + else 0.0 if kind == "bias" else ((self.state >> 16) - 32768) / 1048576.0 + ) + self.stream.write(struct.pack("<%df" % count, *data)) + + +def weights(model, path): + c = CONFIG + b, v, l, h, k, d, f = [ + c[x] + for x in ( + "block_size", + "vocab_size", + "n_layer", + "n_head", + "n_kv_head", + "n_embd", + "intermediate_size", + ) + ] + header = bytearray(1024) + + def fields(fmt, *values): + struct.pack_into("<" + fmt, header, 0, *values) + + if model == "gpt2": + fields("8i", 20240326, 3, b, v, l, h, d, v) + elif model == "llama3": + fields( + "8ififf4i", + 20240803, + 3, + b, + v, + l, + h, + k, + d, + 0.75, + 32, + 1e-5, + 10000.0, + 0, + 4, + 3, + 2, + ) + elif model == "qwen3": + fields( + "10iff4i", 20240804, 4, b, v, l, h, k, d, f, 1, 1e-6, 10000.0, 0, 4, 3, 0 + ) + else: + fields( + "9iffiff3i", + 20260513, + 2, + b, + v, + l, + h, + k, + d, + 4, + 3.0, + 1.0, + 1, + 1e-5, + 10000.0, + 0, + 2, + f, + ) + with path.open("wb") as stream: + stream.write(header) + w = Weights(stream) + w.write(v * d) + if model == "gpt2": + w.write(b * d) + for count, kind in [ + (d, "norm"), + (d, "bias"), + (3 * d * d, "weight"), + (3 * d, "bias"), + (d * d, "weight"), + (d, "bias"), + (d, "norm"), + (d, "bias"), + (4 * d * d, "weight"), + (4 * d, "bias"), + (4 * d * d, "weight"), + (d, "bias"), + ]: + for _ in range(l): + w.write(count, kind) + w.write(d, "norm") + w.write(d, "bias") + elif model in ("llama3", "qwen3"): + for _ in range(l): + w.write(d, "norm") + if model == "qwen3": + for _ in range(l): + w.write(d // h, "norm") + w.write(d // h, "norm") + for count, kind in [ + (3 * d * d, "weight"), + (d * d, "weight"), + (d, "norm"), + (f * d, "weight"), + (f * d, "weight"), + (d * f, "weight"), + ]: + for _ in range(l): + w.write(count, kind) + w.write(d, "norm") + w.write(v * d) + else: + for _ in range(l): + w.write(d, "norm") + w.write(3 * d * d) + w.write(d * d) + w.write(d, "norm") + w.write(4 * d) + for _ in range(4): + w.write(f * d) + w.write(f * d) + w.write(d * f) + w.write(d, "norm") + w.write(v * d) + + +def tokens(model, path): + header = bytearray(1024) + count, state = 16384, 17 + struct.pack_into( + "<3i", header, 0, 20240520 if model == "gpt2" else 20240801, 1, count + ) + values = [] + for _ in range(count): + state = (1664525 * state + 1013904223) & 0xFFFFFFFF + values.append((state >> 16) % CONFIG["vocab_size"]) + path.write_bytes( + header + + struct.pack("<%d%s" % (count, "H" if model == "gpt2" else "I"), *values) + ) + + +def generate(output): + output = Path(output) + output.mkdir(parents=True, exist_ok=True) + files = {} + for model in MODELS: + for kind, writer in (("weights", weights), ("tokens", tokens)): + path = output / (model + "_" + kind + ".bin") + writer(model, path) + files["tiny." + model + "." + kind] = dict( + path=path.name, + size_bytes=path.stat().st_size, + sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + ) + manifest = dict(schema_version=1, fixture_version=1, config=CONFIG, files=files) + (output / "manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n" + ) + return manifest + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + generate(args.output) diff --git a/tests/regression/fixtures/manifest.json b/tests/regression/fixtures/manifest.json new file mode 100644 index 000000000..55c16e3f4 --- /dev/null +++ b/tests/regression/fixtures/manifest.json @@ -0,0 +1,55 @@ +{ + "config": { + "block_size": 64, + "intermediate_size": 64, + "n_embd": 32, + "n_head": 4, + "n_kv_head": 4, + "n_layer": 4, + "vocab_size": 128 + }, + "files": { + "tiny.gpt2.tokens": { + "path": "gpt2_tokens.bin", + "sha256": "31f1a81cfde4838e1de00f9696e4e0350446e6fc96a67a46065f163725455cb4", + "size_bytes": 33792 + }, + "tiny.gpt2.weights": { + "path": "gpt2_weights.bin", + "sha256": "8bee33df162bd083ccaab3966e95ee7e9be0a81a51eebfaae7a392f5b8695da9", + "size_bytes": 229120 + }, + "tiny.llama3.tokens": { + "path": "llama3_tokens.bin", + "sha256": "5851521cb150014b0edcde33dffe07d0ba603ca6ee651d09db19c442508ad1fd", + "size_bytes": 66560 + }, + "tiny.llama3.weights": { + "path": "llama3_weights.bin", + "sha256": "211d69dd6d2865865687c6cf2ff4f8527548708511e2e0e3fd6a4a02fa95de02", + "size_bytes": 198784 + }, + "tiny.mixtral.tokens": { + "path": "mixtral_tokens.bin", + "sha256": "5851521cb150014b0edcde33dffe07d0ba603ca6ee651d09db19c442508ad1fd", + "size_bytes": 66560 + }, + "tiny.mixtral.weights": { + "path": "mixtral_weights.bin", + "sha256": "0a7f4b14fd8d42f101e55ca3037b9a213190165bead7348d2550daf5d4838364", + "size_bytes": 495744 + }, + "tiny.qwen3.tokens": { + "path": "qwen3_tokens.bin", + "sha256": "5851521cb150014b0edcde33dffe07d0ba603ca6ee651d09db19c442508ad1fd", + "size_bytes": 66560 + }, + "tiny.qwen3.weights": { + "path": "qwen3_weights.bin", + "sha256": "71d3013bf0885f801566f2889fe391d6b947a90e2a32ba16732ca464cf6466b0", + "size_bytes": 199040 + } + }, + "fixture_version": 1, + "schema_version": 1 +} diff --git a/tests/regression/fixtures/reference.py b/tests/regression/fixtures/reference.py new file mode 100644 index 000000000..d1d63791e --- /dev/null +++ b/tests/regression/fixtures/reference.py @@ -0,0 +1,461 @@ +#!/usr/bin/env python3 +"""Independent PyTorch FP32 oracle for the versioned tiny fixtures. + +Uses production-exported initial tensors, independently computes attention, +normalization, routing, gradients and Adam updates. Requires optional torch and +numpy; these are not dependencies of the regression runner or fixture generator. +The oracle intentionally supports only the fixed fixture architecture, unmodified +FP32 parameters, constant LR, and a single logical rank. +""" +import argparse +import io +import json +import math +import struct +import time +from pathlib import Path + + +def batch_token_offsets( + num_tokens, sequence_length, micro_batch_size, global_batch_size, step +): + """Complete microbatch offsets, with data wrapping only between batches. + + Samples start every sequence_length tokens and labels need a following + token. The dataset reserves its final sequence; incomplete batch tails + are dropped before the next epoch starts. + """ + values = (num_tokens, sequence_length, micro_batch_size, global_batch_size, step) + if any(type(value) is not int or value <= 0 for value in values): + raise ValueError("batch dimensions and step must be positive integers") + if global_batch_size % micro_batch_size: + raise ValueError("global batch must contain complete microbatches") + complete_batches = (num_tokens // sequence_length - 1) // micro_batch_size + if complete_batches < 1: + raise ValueError("token fixture has no complete microbatch") + accumulation = global_batch_size // micro_batch_size + first_batch = (step - 1) * accumulation + return [ + ((first_batch + micro) % complete_batches) * micro_batch_size * sequence_length + for micro in range(accumulation) + ] + + +def read_snapshot(manifest_path): + """Read legacy NPY files or one cached bundle per rank without repeated I/O.""" + import numpy as np + + manifest_path = Path(manifest_path) + manifest = json.loads(manifest_path.read_text()) + cache, arrays = {}, {} + for item in manifest["tensors"]: + path = manifest_path.parent / item["path"] + if path not in cache: + cache[path] = path.read_bytes() + payload = cache[path] + if "offset" in item or "length" in item: + offset, length = item["offset"], item["length"] + if ( + type(offset) is not int + or type(length) is not int + or offset < 0 + or length <= 0 + or offset + length > len(payload) + ): + raise ValueError("invalid snapshot bundle extent") + payload = payload[offset : offset + length] + arrays[item["kind"], item["name"]] = np.load( + io.BytesIO(payload), allow_pickle=False + ) + return manifest, arrays + + +def run( + initial_snapshot, + tokens, + output, + steps=4, + batch_size=2, + sequence_length=16, + learning_rate=1e-4, + candidate_snapshots=None, + micro_batch_size=None, +): + import numpy as np + import torch + import torch.nn.functional as F + + torch.set_num_threads(1) + torch.use_deterministic_algorithms(True) + micro_batch_size = batch_size if micro_batch_size is None else micro_batch_size + initial_snapshot, output = Path(initial_snapshot), Path(output) + manifest, initial_arrays = read_snapshot(initial_snapshot) + model = manifest["model"] + if ( + manifest["step"] != 0 + or manifest["tp"] != 1 + or manifest["pp"] != 1 + or manifest["dtype"] != "float32" + ): + raise ValueError("oracle requires step-0, single-rank FP32 snapshot") + initial = {} + for item in manifest["tensors"]: + if item["kind"] == "parameter": + if "lora_" in item["name"]: + raise ValueError("oracle does not support LoRA") + initial[item["name"]] = initial_arrays["parameter", item["name"]] + # Decode the versioned LLMC fixture independently, so a loader regression + # cannot pass by feeding the same incorrectly loaded weights to both paths. + weights_path = Path(tokens).with_name(model + "_weights.bin") + data = weights_path.read_bytes() + offset = 1024 + expected = {} + + def read(name, shape): + nonlocal offset + count = math.prod(shape) + expected[name] = ( + np.frombuffer(data, dtype="