diff --git a/README.md b/README.md index 25aab8a8..395ec57b 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,74 @@ MobileFineTuner is an open-source C++ framework for practical, privacy-preservin Unlike simulation-based or desktop-bound approaches, MobileFineTuner is built around a lean native C++ implementation that eliminates Python runtime overhead in the training path and supports both Full Fine-Tuning (Full-FT) and Parameter-Efficient Fine-Tuning (PEFT/LoRA) under tight resource constraints. +### Experimental: Persistent BF16 Full Fine-Tuning + +This fork adds an experimental persistent-BF16 parameter path for native C++ Full Fine-Tuning. + +#### What changed + +- Model parameters can be stored persistently as **BF16** instead of FP32. +- Forward activations and backward computations use FP32 where required. +- Trainable gradients are maintained in **FP32**. +- Adam optimizer moments remain **FP32**, while parameters are written back to BF16. +- SafeTensors loading can convert FP32 source weights to BF16 parameter storage. +- BF16-aware backward paths were added for LayerNorm, RMSNorm, MatMul, and gradient accumulation. +- BF16 storage, optimizer, tokenizer-parity, and end-to-end training tests are included. + +#### GPT-2 124M validation + +The BF16 path was validated with a native C++ build on an ARM64 Android device using Termux/PRoot Debian. + +GPT-2 configuration: + +- 124,439,808 parameters +- 12 layers +- 768 hidden size +- 12 attention heads +- vocabulary size 50,257 +- context length 1,024 + +Full-model BF16 storage validation: + +- 148 BF16 parameter tensors +- 0 FP32 parameter tensors +- 248,879,616 parameter bytes +- 237.35 MiB parameter storage +- approximately half the parameter-storage footprint of FP32 + +A full 124M forward pass completed successfully with finite FP32 logits. + +A full-model 10-step BF16 Full-FT safety run completed successfully: + +- 10/10 forward passes +- 10/10 backward passes +- 10/10 FP32 Adam updates +- gradients cleared after every step +- parameters remained BF16 throughout the run + +#### Real-data experiment + +A separate 20-step causal Full-FT experiment was run on a Classical Tamil dataset using the native GPT-2 tokenizer. + +Configuration: + +- batch size: 1 +- maximum sequence length: 32 +- learning rate: 5e-6 +- Adam β1: 0.9 +- Adam β2: 0.999 +- weight decay: 0 +- 20 optimization steps + +Recorded loss: + +- Step 1: 8.57498 +- Step 20: 5.37094 + +This demonstrates that the BF16 parameter path can execute a real native training workload on a mobile-class ARM64 device. The loss change alone is not evidence of improved model quality or convergence. + +> **Scope:** This is an experimental engineering extension of MobileFineTuner, not a new optimization algorithm. It is intended for reproducibility, systems experimentation, and further research into memory-constrained on-device training. + ### Verified Scope - Stable C++ operator/autograd/LoRA core with unit tests and installable CMake package. diff --git a/operator/CMakeLists.txt b/operator/CMakeLists.txt index 8a8a068f..637d7a15 100644 --- a/operator/CMakeLists.txt +++ b/operator/CMakeLists.txt @@ -1,5 +1,5 @@ cmake_minimum_required(VERSION 3.10) -project(Operators VERSION 2.0.0 LANGUAGES CXX) +project(Operators VERSION 2.0.0 LANGUAGES C CXX) include(GNUInstallDirs) include(CMakePackageConfigHelpers) @@ -565,3 +565,137 @@ message(STATUS " - Build tests: ${BUILD_TESTS}") message(STATUS " - Profiling: ${ENABLE_PROFILING}") message(STATUS "============================================") message(STATUS "") + +if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/finetune_ops/optim/test_gpt2_fullft_10step.cpp) + add_executable(test_gpt2_fullft_10step + finetune_ops/optim/test_gpt2_fullft_10step.cpp + ) + target_link_libraries(test_gpt2_fullft_10step operators) + add_test( + NAME GPT2FullFT10Step + COMMAND test_gpt2_fullft_10step + ) +endif() + +add_executable(test_gpt2_fullft_124m + finetune_ops/optim/test_gpt2_fullft_124m.cpp +) + +target_link_libraries(test_gpt2_fullft_124m PRIVATE operators) + +add_test( + NAME GPT2FullFT124M + COMMAND test_gpt2_fullft_124m +) + + + +add_executable(test_gpt2_tokenizer_parity + finetune_ops/core/test_gpt2_tokenizer_parity.cpp +) + +target_link_libraries(test_gpt2_tokenizer_parity PRIVATE operators) + +add_test( + NAME GPT2TokenizerParity + COMMAND test_gpt2_tokenizer_parity +) + + +add_executable(test_gpt2_fullft_real + finetune_ops/optim/test_gpt2_fullft_real.cpp +) + +target_link_libraries(test_gpt2_fullft_real PRIVATE operators) + +add_test( + NAME GPT2FullFTReal + COMMAND test_gpt2_fullft_real +) + +add_executable(test_gpt2_bf16_load_only + finetune_ops/optim/test_gpt2_bf16_load_only.cpp +) +target_link_libraries(test_gpt2_bf16_load_only PRIVATE operators) + + +# ------------------------------------------------------------ +# ARM64 native numerical backend +# ------------------------------------------------------------ + +add_executable(test_arm64_neon_gemm + finetune_ops/core/test_arm64_neon_gemm.cpp + finetune_ops/core/arm64_neon_kernels.c +) + +target_link_libraries(test_arm64_neon_gemm + PRIVATE + pthread +) + +add_test( + NAME ARM64NEONGEMM + COMMAND test_arm64_neon_gemm +) + +add_executable(test_adam_bf16_param + finetune_ops/optim/test_adam_bf16_param.cpp +) +target_link_libraries(test_adam_bf16_param PRIVATE operators) + +add_executable(test_bf16_storage_diag + finetune_ops/optim/test_bf16_storage_diag.cpp +) +target_link_libraries(test_bf16_storage_diag PRIVATE operators) + +add_executable(test_bf16_direct_write + finetune_ops/optim/test_bf16_direct_write.cpp +) +target_link_libraries(test_bf16_direct_write PRIVATE operators) + +add_executable(test_adam_bf16_param_strict + finetune_ops/optim/test_adam_bf16_param_strict.cpp +) +target_link_libraries(test_adam_bf16_param_strict PRIVATE operators) + +add_executable(test_bf16_weight_pipeline + finetune_ops/optim/test_bf16_weight_pipeline.cpp +) +target_link_libraries(test_bf16_weight_pipeline PRIVATE operators) + + +add_executable(test_gpt2_bf16_forward_safe + finetune_ops/optim/test_gpt2_bf16_forward_safe.cpp +) +target_link_libraries(test_gpt2_bf16_forward_safe PRIVATE operators) + + +add_executable(test_gpt2_bf16_final_safe + finetune_ops/optim/test_gpt2_bf16_final_safe.cpp +) +target_link_libraries(test_gpt2_bf16_final_safe PRIVATE operators) + + +add_executable(test_gpt2_bf16_10step_safe + finetune_ops/optim/test_gpt2_bf16_10step_safe.cpp +) + +target_link_libraries(test_gpt2_bf16_10step_safe + PRIVATE operators +) + +add_executable(test_gpt2_bf16_real_causal_ft + finetune_ops/optim/test_gpt2_bf16_real_causal_ft.cpp +) + +target_link_libraries(test_gpt2_bf16_real_causal_ft + PRIVATE operators +) + +add_executable(test_gpt2_bf16_ultimate + finetune_ops/optim/test_gpt2_bf16_ultimate.cpp +) + +target_link_libraries(test_gpt2_bf16_ultimate + PRIVATE operators +) diff --git a/operator/finetune_ops/core/backward_functions.cpp b/operator/finetune_ops/core/backward_functions.cpp index 8db51b69..e07758bd 100644 --- a/operator/finetune_ops/core/backward_functions.cpp +++ b/operator/finetune_ops/core/backward_functions.cpp @@ -141,6 +141,13 @@ std::vector MatmulBackward::apply(const TensorPtr& grad_output) { if (grad_b && grad_b->shape() != b_shape) { grad_b = sum_to_shape(grad_b, b_shape); } + + // Parameter gradients are always accumulated in FP32. + // The parameter itself may remain BF16. + if (grad_b && b_ && b_->requires_grad() && + grad_b->dtype() != DType::kFloat32) { + grad_b = cast(grad_b, DType::kFloat32); + } } return {grad_a, grad_b}; @@ -438,12 +445,18 @@ std::vector LayerNormBackward::apply(const TensorPtr& grad_output) { int64_t batch = input_->numel() / D; // Gradients - auto grad_input = zeros(shape, input_->dtype(), input_->device()); - auto grad_weight = zeros(weight_->shape(), weight_->dtype(), weight_->device()); - auto grad_bias = zeros(weight_->shape(), weight_->dtype(), weight_->device()); + auto grad_input = zeros(shape, DType::kFloat32, input_->device()); + auto grad_weight = zeros(weight_->shape(), kFloat32, weight_->device()); + auto grad_bias = zeros(weight_->shape(), kFloat32, weight_->device()); + + // LayerNorm always computes in FP32. BF16 parameters are decoded + // to FP32 for the backward calculation. + auto weight_fp32 = (weight_->dtype() == DType::kFloat32) + ? weight_ + : cast(weight_, DType::kFloat32); const float* x = input_->data(); - const float* w = weight_->data(); + const float* w = weight_fp32->data(); const float* gy = grad_output->data(); float* gx = grad_input->data(); float* gw = grad_weight->data(); @@ -495,10 +508,13 @@ std::vector RMSNormBackward::apply(const TensorPtr& grad_output) { int64_t D = shape.back(); int64_t batch = input_->numel() / D; auto grad_input = zeros(shape, input_->dtype(), input_->device()); - auto grad_weight = zeros(weight_->shape(), weight_->dtype(), weight_->device()); + auto grad_weight = zeros(weight_->shape(), kFloat32, weight_->device()); const float* x = input_->data(); - const float* w = weight_->data(); + auto weight_fp32 = (weight_->dtype() == DType::kFloat32) + ? weight_ + : cast(weight_, DType::kFloat32); + const float* w = weight_fp32->data(); const float* gy = grad_output->data(); float* gx = grad_input->data(); float* gw = grad_weight->data(); diff --git a/operator/finetune_ops/core/memory_manager.h b/operator/finetune_ops/core/memory_manager.h index 55e2ab9e..2c5c9c46 100644 --- a/operator/finetune_ops/core/memory_manager.h +++ b/operator/finetune_ops/core/memory_manager.h @@ -8,11 +8,13 @@ #pragma once +#include #include #include #include #include #include +#include #include namespace ops { diff --git a/operator/finetune_ops/core/ops.cpp b/operator/finetune_ops/core/ops.cpp index d1e46457..b0599c30 100644 --- a/operator/finetune_ops/core/ops.cpp +++ b/operator/finetune_ops/core/ops.cpp @@ -9,6 +9,7 @@ */ #include "ops.h" +#include #include "backward_functions.h" #include "autograd_engine.h" #include @@ -296,10 +297,18 @@ namespace { return result_shape; } + template + inline float read_element_as_float(const TensorPtr& t, int64_t index) { + if constexpr (std::is_same_v) { + return bf16_bits_to_float32(t->data()[index]); + } else { + return t->data()[index]; + } + } + template TensorPtr elementwise_binary_op(const TensorPtr& a, const TensorPtr& b, Op op) { if (!can_broadcast(a, b)) { - // 🔧 添加详细错误信息 std::string msg = "Tensors cannot be broadcasted: shape_a=["; for (size_t i = 0; i < a->shape().size(); ++i) { msg += std::to_string(a->shape()[i]); @@ -314,63 +323,84 @@ namespace { throw TensorError(msg); } + if (!DTypeUtils::is_floating_point(a->dtype()) || + !DTypeUtils::is_floating_point(b->dtype())) { + throw TensorError("elementwise_binary_op: only floating-point tensors are supported"); + } + auto result_shape = broadcast_shapes(a, b); auto result = zeros(result_shape, a->dtype(), a->device()); - if (shapes_equal(a, b)) { - const float* data_a = a->data(); - const float* data_b = b->data(); - float* result_data = result->data(); + auto read_value = [](const TensorPtr& t, int64_t index) -> float { + switch (t->dtype()) { + case kFloat32: + return t->data()[index]; + case kBFloat16: + return bf16_bits_to_float32(t->data()[index]); + case kFloat16: + return fp16_bits_to_float32(t->data()[index]); + default: + throw TensorError("elementwise_binary_op: unsupported dtype"); + } + }; - for (int64_t i = 0; i < a->numel(); ++i) { - result_data[i] = op(data_a[i], data_b[i]); + auto write_value = [](const TensorPtr& t, int64_t index, float value) { + switch (t->dtype()) { + case kFloat32: + t->data()[index] = value; + break; + case kBFloat16: + t->data()[index] = float32_to_bf16_bits(value); + break; + case kFloat16: + t->data()[index] = float32_to_fp16_bits(value); + break; + default: + throw TensorError("elementwise_binary_op: unsupported output dtype"); } - } else { - // Complete broadcast implementsation - const float* data_a = a->data(); - const float* data_b = b->data(); - float* result_data = result->data(); - - auto shape_a = a->shape(); - auto shape_b = b->shape(); - - for (int64_t i = 0; i < result->numel(); ++i) { - // Calculate multidimensional index of current position in result + }; + + auto shape_a = a->shape(); + auto shape_b = b->shape(); + + for (int64_t i = 0; i < result->numel(); ++i) { + int64_t idx_a = 0; + int64_t idx_b = 0; + + if (shapes_equal(a, b)) { + idx_a = i; + idx_b = i; + } else { std::vector result_idx(result_shape.size()); int64_t temp = i; - for (int j = result_shape.size() - 1; j >= 0; --j) { + + for (int j = static_cast(result_shape.size()) - 1; j >= 0; --j) { result_idx[j] = temp % result_shape[j]; temp /= result_shape[j]; } - - // Calculate corresponding indices for a and b (simplified version) - int64_t idx_a = 0, idx_b = 0; - - // Calculate linear index for a + for (size_t dim = 0; dim < shape_a.size(); ++dim) { - int result_dim = dim + (result_shape.size() - shape_a.size()); - if (result_dim >= 0) { - int64_t coord = (shape_a[dim] == 1) ? 0 : result_idx[result_dim]; - idx_a = idx_a * shape_a[dim] + coord; - } + int result_dim = static_cast(dim) + + static_cast(result_shape.size() - shape_a.size()); + int64_t coord = (shape_a[dim] == 1) ? 0 : result_idx[result_dim]; + idx_a = idx_a * shape_a[dim] + coord; } - - // Calculate linear index for b + for (size_t dim = 0; dim < shape_b.size(); ++dim) { - int result_dim = dim + (result_shape.size() - shape_b.size()); - if (result_dim >= 0) { - int64_t coord = (shape_b[dim] == 1) ? 0 : result_idx[result_dim]; - idx_b = idx_b * shape_b[dim] + coord; - } + int result_dim = static_cast(dim) + + static_cast(result_shape.size() - shape_b.size()); + int64_t coord = (shape_b[dim] == 1) ? 0 : result_idx[result_dim]; + idx_b = idx_b * shape_b[dim] + coord; } - - result_data[i] = op(data_a[idx_a], data_b[idx_b]); } + + const float va = read_value(a, idx_a); + const float vb = read_value(b, idx_b); + write_value(result, i, op(va, vb)); } if (a->requires_grad() || b->requires_grad()) { result->set_requires_grad(true); - } return result; @@ -637,9 +667,8 @@ TensorPtr matmul(const TensorPtr& a, const TensorPtr& b) { if (b->dtype() != kFloat32 && !is_lowp_float(b->dtype())) { throw TensorError("matmul: unsupported right operand dtype " + DTypeUtils::to_string(b->dtype())); } - if (is_lowp_float(b->dtype()) && b->requires_grad()) { - throw TensorError("matmul: low-precision trainable right operand is not supported; keep trainable weights FP32"); - } + // BF16 trainable right operands are supported. + // Compute kernels accumulate into FP32 output; gradients remain FP32. int64_t m = shape_a[shape_a.size() - 2]; int64_t k = shape_a[shape_a.size() - 1]; @@ -844,9 +873,8 @@ TensorPtr matmul_rhs_T(const TensorPtr& a, const TensorPtr& b) { if (b->dtype() != kFloat32 && !is_lowp_float(b->dtype())) { throw TensorError("matmul_rhs_T: unsupported right operand dtype " + DTypeUtils::to_string(b->dtype())); } - if (is_lowp_float(b->dtype()) && b->requires_grad()) { - throw TensorError("matmul_rhs_T: low-precision trainable right operand is not supported; keep trainable weights FP32"); - } + // BF16 trainable right operands are supported here as well. + // The result and computed gradients remain FP32. int64_t n = shape_b[0]; int64_t k_b = shape_b[1]; @@ -1595,10 +1623,31 @@ TensorPtr layer_norm(const TensorPtr& input, const TensorPtr& weight, const Tens auto result = zeros(input_shape, input->dtype(), input->device()); const float* input_data = input->data(); - const float* weight_data = weight->data(); - const float* bias_data = bias->data(); + + const float* weight_data_fp32 = + (weight->dtype() == kFloat32) ? weight->data() : nullptr; + const uint16_t* weight_data_bf16 = + (weight->dtype() == kBFloat16) ? weight->data() : nullptr; + + const float* bias_data_fp32 = + (bias->dtype() == kFloat32) ? bias->data() : nullptr; + const uint16_t* bias_data_bf16 = + (bias->dtype() == kBFloat16) ? bias->data() : nullptr; + float* result_data = result->data(); + auto read_weight = [&](int64_t i) -> float { + if (weight_data_fp32) return weight_data_fp32[i]; + if (weight_data_bf16) return bf16_bits_to_float32(weight_data_bf16[i]); + throw TensorError("layer_norm: unsupported weight dtype"); + }; + + auto read_bias = [&](int64_t i) -> float { + if (bias_data_fp32) return bias_data_fp32[i]; + if (bias_data_bf16) return bf16_bits_to_float32(bias_data_bf16[i]); + throw TensorError("layer_norm: unsupported bias dtype"); + }; + int64_t batch_size = input->numel() / normalized_dim; for (int64_t b = 0; b < batch_size; ++b) { @@ -1621,7 +1670,7 @@ TensorPtr layer_norm(const TensorPtr& input, const TensorPtr& weight, const Tens float inv_std = 1.0f / std::sqrt(variance + eps); for (int64_t i = 0; i < normalized_dim; ++i) { float normalized = (batch_input[i] - mean) * inv_std; - batch_result[i] = normalized * weight_data[i] + bias_data[i]; + batch_result[i] = normalized * read_weight(i) + read_bias(i); } } diff --git a/operator/finetune_ops/core/test_gpt2_tokenizer_parity.cpp b/operator/finetune_ops/core/test_gpt2_tokenizer_parity.cpp new file mode 100644 index 00000000..60110c0c --- /dev/null +++ b/operator/finetune_ops/core/test_gpt2_tokenizer_parity.cpp @@ -0,0 +1,39 @@ +#include +#include +#include "../core/tokenizer.h" + +int main() { + try { + const std::string model_dir = "/root/gpt2-tamil-124m"; + const std::string text = "\xE0\xAE\xAF\xE0\xAE\xBE\xE0\xAE\xA4\xE0\xAF\x81\xE0\xAE\xAE\xE0\xAF\x8D\x20\xE0\xAE\x8A\xE0\xAE\xB0\xE0\xAF\x87\x20\xE0\xAE\xAF\xE0\xAE\xBE\xE0\xAE\xB5\xE0\xAE\xB0\xE0\xAF\x81\xE0\xAE\xAE\xE0\xAF\x8D\x20\xE0\xAE\x95\xE0\xAF\x87\xE0\xAE\xB3\xE0\xAE\xBF\xE0\xAE\xB0\xE0\xAF\x8D"; + + ops::TokenizerLoadOptions opts; + opts.model_type = "gpt2"; + + auto tok = ops::TokenizerFactory::from_pretrained(model_dir, opts); + + if (!tok) { + throw std::runtime_error("Tokenizer creation failed"); + } + + auto ids = tok->encode(text); + + std::cout << "Native tokenizer vocab: " + << tok->get_vocab_size() << "\n"; + + std::cout << "Token count: " << ids.size() << "\n"; + std::cout << "IDs: ["; + + for (size_t i = 0; i < ids.size(); ++i) { + if (i) std::cout << ", "; + std::cout << ids[i]; + } + + std::cout << "]\n"; + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 1; + } +} diff --git a/operator/finetune_ops/graph/safetensors_loader.cpp b/operator/finetune_ops/graph/safetensors_loader.cpp index fa207fa3..7d9ac6d2 100644 --- a/operator/finetune_ops/graph/safetensors_loader.cpp +++ b/operator/finetune_ops/graph/safetensors_loader.cpp @@ -167,12 +167,14 @@ TensorPtr SafeTensorsReader::read_tensor_data(const SafeTensorInfo& info, bool t numel *= dim; } + // element_size describes the bytes stored in the source SafeTensors file. + // target_dtype describes how the tensor will be represented in RAM. size_t element_size = 4; DType target_dtype = kFloat32; if (info.dtype == "F32") { element_size = 4; - target_dtype = kFloat32; + target_dtype = preserve_low_precision ? kBFloat16 : kFloat32; } else if (info.dtype == "F16") { element_size = 2; target_dtype = preserve_low_precision ? kFloat16 : kFloat32; @@ -242,8 +244,21 @@ TensorPtr SafeTensorsReader::read_tensor_data(const SafeTensorInfo& info, bool t }; if (info.dtype == "F32") { - std::memcpy(tensor->data(), raw_data.data(), byte_size); - transpose_buffer_float(tensor->data()); + const float* fp32_data = + reinterpret_cast(raw_data.data()); + + if (target_dtype == kBFloat16) { + uint16_t* bf16_data = tensor->data(); + + for (int64_t i = 0; i < numel; ++i) { + bf16_data[i] = float32_to_bf16_bits(fp32_data[i]); + } + + transpose_buffer_u16(bf16_data); + } else { + std::memcpy(tensor->data(), raw_data.data(), byte_size); + transpose_buffer_float(tensor->data()); + } } else if (info.dtype == "F16") { const uint16_t* fp16_data = reinterpret_cast(raw_data.data()); if (target_dtype == kFloat16) { @@ -304,6 +319,11 @@ SafeTensorsReader::load_tensors_mapped( it->second.shape.size() == 2; bool preserve_low_precision = !options.auto_promote_fp16; + + if (options.convert_f32_to_bf16 && it->second.dtype == "F32") { + preserve_low_precision = true; + } + if (options.auto_promote_fp16) { for (const auto& needle : options.preserve_low_precision_key_substrings) { if ((!needle.empty()) && @@ -471,12 +491,12 @@ GPT2KeyMapper::generate_gpt2_mapping(int num_layers) { std::unordered_map mapping; // Embeddings - mapping["wte.weight"] = "wte.weight"; - mapping["wpe.weight"] = "wpe.weight"; + mapping["wte.weight"] = "transformer.wte.weight"; + mapping["wpe.weight"] = "transformer.wpe.weight"; // Transformer blocks for (int i = 0; i < num_layers; ++i) { - std::string hf_prefix = "h." + std::to_string(i) + "."; + std::string hf_prefix = "transformer.h." + std::to_string(i) + "."; std::string internal_prefix = "blocks." + std::to_string(i) + "."; // LayerNorm 1 @@ -501,8 +521,8 @@ GPT2KeyMapper::generate_gpt2_mapping(int num_layers) { } // Final LayerNorm - mapping["ln_f.weight"] = "ln_f.weight"; - mapping["ln_f.bias"] = "ln_f.bias"; + mapping["ln_f.weight"] = "transformer.ln_f.weight"; + mapping["ln_f.bias"] = "transformer.ln_f.bias"; // lm_head (typically tied with wte; enable below if loading separately) // mapping["lm_head.weight"] = "lm_head.weight"; diff --git a/operator/finetune_ops/graph/safetensors_loader.h b/operator/finetune_ops/graph/safetensors_loader.h index a8c2a616..b5ecf5f6 100644 --- a/operator/finetune_ops/graph/safetensors_loader.h +++ b/operator/finetune_ops/graph/safetensors_loader.h @@ -35,6 +35,7 @@ struct SafeTensorInfo { struct SafeTensorsLoadOptions { bool transpose_linear = true; // auto-transpose Linear weights [out,in]→[in,out] bool auto_promote_fp16 = true; // auto-promote FP16 to FP32 + bool convert_f32_to_bf16 = false; // convert F32 weights to BF16 while loading bool verbose = true; // print load logs bool strict_shape_check = true; // strict shape validation // Preserve original F16/BF16 storage for matching internal or HF keys even diff --git a/operator/finetune_ops/optim/adam.cpp b/operator/finetune_ops/optim/adam.cpp index 72881dfc..aa8eda0d 100644 --- a/operator/finetune_ops/optim/adam.cpp +++ b/operator/finetune_ops/optim/adam.cpp @@ -27,64 +27,102 @@ void Adam::step(const std::vector& parameters, if (parameters.size() != gradients.size()) { throw std::runtime_error("Parameters and gradients size mismatch"); } - + for (size_t i = 0; i < parameters.size(); ++i) { auto& param = parameters[i]; auto& grad = gradients[i]; - - if (!grad) continue; // Skip if no gradient - - // Initialize state if needed + + if (!param) continue; + if (!grad) continue; + + if (param->dtype() != kFloat32 && param->dtype() != kBFloat16) { + throw std::runtime_error( + "Adam: parameter dtype must be float32 or bfloat16"); + } + + if (grad->dtype() != kFloat32) { + throw std::runtime_error( + "Adam: gradients must remain float32"); + } + if (states_.find(param) == states_.end()) { init_state(param); } - + auto& state = states_[param]; state.step++; - - // Get current parameter and gradient data + const float* grad_data = grad->data(); - float* param_data = param->data(); + + float* param_fp32 = nullptr; + uint16_t* param_bf16 = nullptr; + + if (param->dtype() == kFloat32) { + param_fp32 = param->data(); + } else { + param_bf16 = param->data(); + } + float* m_data = state.m[0]->data(); float* v_data = state.v[0]->data(); + float* v_hat_data = nullptr; - if (adam_config_.amsgrad && !state.v_hat.empty()) { v_hat_data = state.v_hat[0]->data(); } - - // Compute bias correction factors - float bias_correction1 = compute_bias_correction1(state.step); - float bias_correction2 = compute_bias_correction2(state.step); - - // Update parameters + + const float bias_correction1 = + compute_bias_correction1(state.step); + const float bias_correction2 = + compute_bias_correction2(state.step); + for (int64_t j = 0; j < param->numel(); ++j) { - float grad_val = grad_data[j]; - - // Apply weight decay if specified + const float param_value = + (param->dtype() == kBFloat16) + ? bf16_bits_to_float32(param_bf16[j]) + : param_fp32[j]; + + float grad_value = grad_data[j]; + if (adam_config_.weight_decay > 0.0f) { - grad_val += adam_config_.weight_decay * param_data[j]; + grad_value += + adam_config_.weight_decay * param_value; } - - // Update biased first moment estimate - m_data[j] = adam_config_.beta1 * m_data[j] + (1.0f - adam_config_.beta1) * grad_val; - - // Update biased second raw moment estimate - v_data[j] = adam_config_.beta2 * v_data[j] + (1.0f - adam_config_.beta2) * grad_val * grad_val; - - float v_corrected = v_data[j] / bias_correction2; - - // AMSGrad variant + + m_data[j] = + adam_config_.beta1 * m_data[j] + + (1.0f - adam_config_.beta1) * grad_value; + + v_data[j] = + adam_config_.beta2 * v_data[j] + + (1.0f - adam_config_.beta2) * + grad_value * grad_value; + + float v_corrected = + v_data[j] / bias_correction2; + if (adam_config_.amsgrad && v_hat_data) { - v_hat_data[j] = std::max(v_hat_data[j], v_corrected); + v_hat_data[j] = + std::max(v_hat_data[j], v_corrected); v_corrected = v_hat_data[j]; } - - // Compute bias-corrected first moment estimate - float m_corrected = m_data[j] / bias_correction1; - - // Update parameters - param_data[j] -= adam_config_.learning_rate * m_corrected / (std::sqrt(v_corrected) + adam_config_.epsilon); + + const float m_corrected = + m_data[j] / bias_correction1; + + const float updated_value = + param_value - + adam_config_.learning_rate * + m_corrected / + (std::sqrt(v_corrected) + + adam_config_.epsilon); + + if (param->dtype() == kBFloat16) { + param_bf16[j] = + float32_to_bf16_bits(updated_value); + } else { + param_fp32[j] = updated_value; + } } } } diff --git a/operator/finetune_ops/optim/test_adam_bf16_param.cpp b/operator/finetune_ops/optim/test_adam_bf16_param.cpp new file mode 100644 index 00000000..058f28d1 --- /dev/null +++ b/operator/finetune_ops/optim/test_adam_bf16_param.cpp @@ -0,0 +1,87 @@ +#include +#include +#include +#include + +#include "adam.h" +#include "../core/tensor.h" +#include "../core/dtype.h" + +using namespace ops; + +int main() { + try { + auto param = std::make_shared( + std::vector{4}, kBFloat16, kCPU); + + auto* p = param->data(); + p[0] = float32_to_bf16_bits(1.0f); + p[1] = float32_to_bf16_bits(2.0f); + p[2] = float32_to_bf16_bits(3.0f); + p[3] = float32_to_bf16_bits(4.0f); + + param->set_requires_grad(true); + + auto grad = std::make_shared( + std::vector{4}, kFloat32, kCPU); + + float* g = grad->data(); + g[0] = 1.0f; + g[1] = 2.0f; + g[2] = 3.0f; + g[3] = 4.0f; + + AdamConfig cfg; + cfg.learning_rate = 0.1f; + cfg.beta1 = 0.9f; + cfg.beta2 = 0.999f; + cfg.epsilon = 1e-8f; + cfg.weight_decay = 0.0f; + cfg.amsgrad = false; + + Adam adam(cfg); + + std::vector parameters{param}; + std::vector gradients{grad}; + + adam.step(parameters, gradients); + + std::cout << "dtype=" << DTypeUtils::to_string(param->dtype()) << "\n"; + std::cout << "updated:"; + for (int i = 0; i < 4; ++i) { + std::cout << " " + << bf16_bits_to_float32(p[i]); + } + std::cout << "\n"; + + // Step 1 of Adam with positive gradients should move every value down. + bool moved = true; + const float expected_before[] = {1.0f, 2.0f, 3.0f, 4.0f}; + + for (int i = 0; i < 4; ++i) { + float v = bf16_bits_to_float32(p[i]); + if (!(v < expected_before[i])) { + moved = false; + } + } + + if (!moved) { + std::cerr << "[FAIL] BF16 parameters were not updated downward.\n"; + return 1; + } + + if (param->dtype() != kBFloat16) { + std::cerr << "[FAIL] Parameter dtype changed unexpectedly.\n"; + return 2; + } + + std::cout << "[PASS] BF16 parameter updated by FP32 Adam math.\n"; + std::cout << "[PASS] Parameter storage remains BF16.\n"; + std::cout << "[PASS] This test uses 4 elements only.\n"; + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 3; + } +} diff --git a/operator/finetune_ops/optim/test_adam_bf16_param_strict.cpp b/operator/finetune_ops/optim/test_adam_bf16_param_strict.cpp new file mode 100644 index 00000000..561faa8c --- /dev/null +++ b/operator/finetune_ops/optim/test_adam_bf16_param_strict.cpp @@ -0,0 +1,93 @@ +#include +#include +#include +#include +#include + +#include "adam.h" +#include "../core/tensor.h" +#include "../core/dtype.h" + +using namespace ops; + +int main() { + try { + const float initial[4] = {1.0f, 2.0f, 3.0f, 4.0f}; + const float grads[4] = {1.0f, 1.0f, 1.0f, 1.0f}; + + auto param = std::make_shared( + std::vector{4}, kBFloat16, kCPU); + + auto* p = param->data(); + + for (int i = 0; i < 4; ++i) { + p[i] = float32_to_bf16_bits(initial[i]); + } + + auto grad = std::make_shared( + std::vector{4}, kFloat32, kCPU); + + float* g = grad->data(); + for (int i = 0; i < 4; ++i) { + g[i] = grads[i]; + } + + AdamConfig cfg; + cfg.learning_rate = 0.1f; + cfg.beta1 = 0.9f; + cfg.beta2 = 0.999f; + cfg.epsilon = 1e-8f; + cfg.weight_decay = 0.0f; + cfg.amsgrad = false; + + Adam adam(cfg); + + param->set_requires_grad(true); + + adam.step({param}, {grad}); + + /* + * For the first Adam step with beta1=0.9, beta2=0.999 and + * identical positive gradients, the bias-corrected update is + * exactly approximately lr = 0.1. + * + * Therefore expected BF16 values are: + * BF16(0.9), BF16(1.9), BF16(2.9), BF16(3.9) + */ + bool ok = true; + + std::cout << "dtype=" << DTypeUtils::to_string(param->dtype()) << "\n"; + + for (int i = 0; i < 4; ++i) { + const float actual = bf16_bits_to_float32(p[i]); + const float expected = + bf16_bits_to_float32( + float32_to_bf16_bits(initial[i] - 0.1f)); + + const float diff = std::fabs(actual - expected); + + std::cout << "[" << i << "] actual=" << std::setprecision(9) + << actual + << " expected=" << expected + << " diff=" << diff + << "\n"; + + if (diff > 1e-6f) { + ok = false; + } + } + + if (!ok) { + std::cerr << "[FAIL] BF16 Adam update does not match reference.\n"; + return 1; + } + + std::cout << "[PASS] BF16 Adam update matches FP32 reference.\n"; + std::cout << "[PASS] Parameter storage remains BF16.\n"; + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 2; + } +} diff --git a/operator/finetune_ops/optim/test_bf16_direct_write.cpp b/operator/finetune_ops/optim/test_bf16_direct_write.cpp new file mode 100644 index 00000000..f989f4e5 --- /dev/null +++ b/operator/finetune_ops/optim/test_bf16_direct_write.cpp @@ -0,0 +1,55 @@ +#include +#include +#include +#include + +#include "../core/tensor.h" +#include "../core/dtype.h" + +using namespace ops; + +int main() { + try { + auto t = std::make_shared( + std::vector{4}, kBFloat16, kCPU); + + uint16_t* p = t->data(); + + const float values[4] = {0.9f, 1.9f, 2.9f, 3.9f}; + + for (int i = 0; i < 4; ++i) { + p[i] = float32_to_bf16_bits(values[i]); + } + + std::cout << "dtype=" << DTypeUtils::to_string(t->dtype()) << "\n"; + + bool ok = true; + + for (int i = 0; i < 4; ++i) { + float decoded = bf16_bits_to_float32(p[i]); + + std::cout << "[" << i << "] bits=0x" + << std::hex << std::setw(4) << std::setfill('0') << p[i] + << std::dec + << " value=" << std::setprecision(9) << decoded + << "\n"; + + if (decoded != bf16_bits_to_float32( + float32_to_bf16_bits(values[i]))) { + ok = false; + } + } + + if (!ok) { + std::cerr << "[FAIL] Direct BF16 storage test failed.\n"; + return 1; + } + + std::cout << "[PASS] Direct BF16 storage/write/read works.\n"; + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 2; + } +} diff --git a/operator/finetune_ops/optim/test_bf16_storage_diag.cpp b/operator/finetune_ops/optim/test_bf16_storage_diag.cpp new file mode 100644 index 00000000..401c2f14 --- /dev/null +++ b/operator/finetune_ops/optim/test_bf16_storage_diag.cpp @@ -0,0 +1,64 @@ +#include +#include +#include +#include + +#include "../core/tensor.h" +#include "../core/dtype.h" +#include "adam.h" + +using namespace ops; + +static void dump(const char* tag, const TensorPtr& t) { + auto* u = t->data(); + std::cout << tag << "\n"; + for (int i = 0; i < 4; ++i) { + float f = bf16_bits_to_float32(u[i]); + std::cout << " [" << i << "] bits=0x" + << std::hex << std::setw(4) << std::setfill('0') << u[i] + << std::dec << " value=" << std::setprecision(9) << f + << "\n"; + } +} + +int main() { + try { + auto t = std::make_shared( + std::vector{4}, kBFloat16, kCPU); + + auto* u = t->data(); + const float initial[4] = {1.0f, 2.0f, 3.0f, 4.0f}; + + for (int i = 0; i < 4; ++i) { + u[i] = float32_to_bf16_bits(initial[i]); + } + + dump("BEFORE ADAM:", t); + + auto g = std::make_shared( + std::vector{4}, kFloat32, kCPU); + + float* gd = g->data(); + for (int i = 0; i < 4; ++i) gd[i] = 1.0f; + + AdamConfig cfg; + cfg.learning_rate = 0.1f; + cfg.beta1 = 0.9f; + cfg.beta2 = 0.999f; + cfg.epsilon = 1e-8f; + cfg.weight_decay = 0.0f; + cfg.amsgrad = false; + + Adam adam(cfg); + + t->set_requires_grad(true); + adam.step({t}, {g}); + + dump("AFTER ADAM:", t); + + return 0; + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 1; + } +} diff --git a/operator/finetune_ops/optim/test_bf16_weight_pipeline.cpp b/operator/finetune_ops/optim/test_bf16_weight_pipeline.cpp new file mode 100644 index 00000000..401690f3 --- /dev/null +++ b/operator/finetune_ops/optim/test_bf16_weight_pipeline.cpp @@ -0,0 +1,144 @@ +#include +#include +#include +#include +#include +#include + +#include "../core/ops.h" +#include "../core/tensor.h" +#include "../core/dtype.h" +#include "adam.h" + +using namespace ops; + +static float read_bf16(const TensorPtr& t, int64_t i) { + return bf16_bits_to_float32(t->data()[i]); +} + +int main() { + try { + // Small FP32 activation. + auto x = std::make_shared( + std::vector{2, 3}, kFloat32, kCPU); + + float* xd = x->data(); + xd[0] = 1.0f; xd[1] = 2.0f; xd[2] = 3.0f; + xd[3] = 4.0f; xd[4] = 5.0f; xd[5] = 6.0f; + + // Trainable BF16 weight. + auto w = std::make_shared( + std::vector{3, 2}, kBFloat16, kCPU); + + uint16_t* wd = w->data(); + const float initial[] = { + 0.10f, 0.20f, + 0.30f, 0.40f, + 0.50f, 0.60f + }; + + for (int i = 0; i < 6; ++i) { + wd[i] = float32_to_bf16_bits(initial[i]); + } + + w->set_requires_grad(true); + + // This is the important operation: + // FP32 activation × BF16 trainable weight -> FP32 output. + auto y = matmul(x, w); + + if (y->dtype() != kFloat32) { + std::cerr << "[FAIL] Matmul output is not FP32.\n"; + return 1; + } + + // Stable target with same shape. + auto target = std::make_shared( + std::vector{2, 2}, kFloat32, kCPU); + + float* td = target->data(); + for (int i = 0; i < 4; ++i) td[i] = 0.0f; + + auto loss = mse_loss(y, target, "mean"); + + if (!loss) { + std::cerr << "[FAIL] Loss creation failed.\n"; + return 2; + } + + loss->backward(); + + auto grad = w->grad(); + + if (!grad) { + std::cerr << "[FAIL] BF16 weight received no gradient.\n"; + return 3; + } + + if (grad->dtype() != kFloat32) { + std::cerr << "[FAIL] BF16 weight gradient is not FP32.\n"; + return 4; + } + + std::cout << "loss=" << std::setprecision(9) + << loss->item() << "\n"; + + std::cout << "weight_before:"; + for (int i = 0; i < 6; ++i) { + std::cout << " " << read_bf16(w, i); + } + std::cout << "\n"; + + AdamConfig cfg; + cfg.learning_rate = 1e-2f; + cfg.beta1 = 0.9f; + cfg.beta2 = 0.999f; + cfg.epsilon = 1e-8f; + cfg.weight_decay = 0.0f; + cfg.amsgrad = false; + + Adam adam(cfg); + adam.step({w}, {grad}); + + std::cout << "weight_after:"; + for (int i = 0; i < 6; ++i) { + std::cout << " " << read_bf16(w, i); + } + std::cout << "\n"; + + bool changed = false; + + for (int i = 0; i < 6; ++i) { + float after = read_bf16(w, i); + float before = + bf16_bits_to_float32( + float32_to_bf16_bits(initial[i])); + + if (after != before) { + changed = true; + break; + } + } + + if (!changed) { + std::cerr << "[FAIL] BF16 weights did not change after Adam.\n"; + return 5; + } + + if (w->dtype() != kBFloat16) { + std::cerr << "[FAIL] Weight dtype changed.\n"; + return 6; + } + + std::cout << "[PASS] FP32 activation × BF16 trainable weight.\n"; + std::cout << "[PASS] BF16 weight received FP32 gradient.\n"; + std::cout << "[PASS] FP32 Adam updated BF16 weight.\n"; + std::cout << "[PASS] Weight storage remains BF16.\n"; + + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 10; + } +} diff --git a/operator/finetune_ops/optim/test_gpt2_bf16_10step_safe.cpp b/operator/finetune_ops/optim/test_gpt2_bf16_10step_safe.cpp new file mode 100644 index 00000000..947cc2be --- /dev/null +++ b/operator/finetune_ops/optim/test_gpt2_bf16_10step_safe.cpp @@ -0,0 +1,510 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../graph/gpt2_model.h" +#include "../graph/safetensors_loader.h" +#include "../core/tensor.h" +#include "../core/ops.h" +#include "../core/dtype.h" +#include "adam.h" + +using namespace ops; + +static long long mem_available_kb() { + std::ifstream f("/proc/meminfo"); + std::string key, unit; + long long value = 0; + + while (f >> key >> value >> unit) { + if (key == "MemAvailable:") + return value; + } + + return -1; +} + +static void print_mem(const char* label) { + const long long kb = mem_available_kb(); + + std::cout << "[MEM] " << label << ": " + << kb / 1024.0 << " MiB available\n"; +} + +static bool finite_f32(const TensorPtr& t) { + if (!t || t->dtype() != kFloat32) + return false; + + const float* p = t->data(); + + for (int64_t i = 0; i < t->numel(); ++i) { + if (!std::isfinite(p[i])) + return false; + } + + return true; +} + +static bool finite_bf16(const TensorPtr& t) { + if (!t || t->dtype() != kBFloat16) + return false; + + const uint16_t* p = t->data(); + + for (int64_t i = 0; i < t->numel(); ++i) { + if (!std::isfinite(bf16_bits_to_float32(p[i]))) + return false; + } + + return true; +} + +int main() { + constexpr int STEPS = 10; + constexpr int64_t SEQ = 8; + constexpr int64_t VOCAB = 50257; + + // Safety thresholds. + constexpr long long HARD_STOP_KB = 600LL * 1024LL; + constexpr long long STEP_MIN_KB = 1000LL * 1024LL; + constexpr long long START_MIN_KB = 1800LL * 1024LL; + + std::atomic stop_watchdog{false}; + + std::thread watchdog([&]() { + while (!stop_watchdog.load()) { + const long long kb = mem_available_kb(); + + if (kb > 0 && kb < HARD_STOP_KB) { + std::cerr + << "\n[SAFE ABORT] MemAvailable < 600 MiB\n"; + std::_Exit(99); + } + + std::this_thread::sleep_for( + std::chrono::milliseconds(250)); + } + }); + + auto finish = [&](int rc) { + stop_watchdog = true; + watchdog.join(); + return rc; + }; + + try { + const std::string model_dir = + "/root/gpt2-tamil-124m"; + + std::cout + << "============================================\n" + << " SAFE 124M BF16 10-STEP FINE-TUNE\n" + << "============================================\n"; + + print_mem("start"); + + if (mem_available_kb() < START_MIN_KB) { + std::cout + << "[SAFE STOP] Need >= 1.8 GiB at startup.\n"; + return finish(0); + } + + // -------------------------------------------------------- + // MODEL + // -------------------------------------------------------- + + GPT2Config config = + GPT2Config::from_pretrained(model_dir); + + GPT2Model model(config); + + if (config.tie_word_embeddings) + model.tie_weights(); + + // -------------------------------------------------------- + // BF16 LOAD + // -------------------------------------------------------- + + SafeTensorsModelReader reader(model_dir); + reader.parse_headers(); + + auto mapping = + GPT2KeyMapper::generate_gpt2_mapping( + config.n_layer); + + SafeTensorsLoadOptions options; + options.transpose_linear = false; + options.auto_promote_fp16 = true; + options.convert_f32_to_bf16 = true; + options.verbose = false; + + auto tensors = + reader.load_tensors_mapped( + mapping, + options); + + size_t total_params = 0; + size_t bf16_tensors = 0; + size_t fp32_tensors = 0; + + for (const auto& kv : tensors) { + if (!kv.second) + continue; + + total_params += + static_cast(kv.second->numel()); + + if (kv.second->dtype() == kBFloat16) + ++bf16_tensors; + else if (kv.second->dtype() == kFloat32) + ++fp32_tensors; + + model.assign_weight(kv.first, kv.second); + } + + std::cout + << "Parameters: " << total_params << "\n" + << "BF16 tensors: " << bf16_tensors << "\n" + << "FP32 tensors: " << fp32_tensors << "\n"; + + if (total_params != 124439808 || + bf16_tensors != 148 || + fp32_tensors != 0) { + + std::cerr + << "[FAIL] Unexpected parameter dtype state.\n"; + return finish(2); + } + + std::cout + << "[PASS] All 124M parameters are BF16.\n"; + + print_mem("after BF16 load"); + + // -------------------------------------------------------- + // FULL PARAMETER TRAINING + // -------------------------------------------------------- + + for (auto& p : model.parameters()) { + if (p) + p->set_requires_grad(true); + } + + // -------------------------------------------------------- + // INPUT + // -------------------------------------------------------- + + const int64_t ids[SEQ] = { + 288, 275, 271, 265, + 281, 262, 742, 302 + }; + + auto input_ids = + std::make_shared( + std::vector{1, SEQ}, + ids, + kInt64, + kCPU); + + // -------------------------------------------------------- + // TARGET + // -------------------------------------------------------- + // + // Keep the existing tiny deterministic MSE objective. + // This is a training-path stress test, not the final + // corpus LM objective. + // + + auto target = + std::make_shared( + std::vector{ + 1, + SEQ, + VOCAB + }, + kFloat32, + kCPU); + + float* target_data = target->data(); + + std::fill( + target_data, + target_data + target->numel(), + 0.0f); + + const int64_t target_ids[SEQ] = { + 275, 271, 265, 281, + 262, 742, 302, 288 + }; + + for (int64_t i = 0; i < SEQ; ++i) { + target_data[ + i * VOCAB + target_ids[i] + ] = 1.0f; + } + + // -------------------------------------------------------- + // ADAM + // -------------------------------------------------------- + + AdamConfig adam_cfg; + adam_cfg.learning_rate = 1e-5f; + adam_cfg.beta1 = 0.9f; + adam_cfg.beta2 = 0.999f; + adam_cfg.epsilon = 1e-8f; + adam_cfg.weight_decay = 0.0f; + adam_cfg.amsgrad = false; + + Adam adam(adam_cfg); + + // -------------------------------------------------------- + // TRAINING LOOP + // -------------------------------------------------------- + + float first_loss = 0.0f; + float last_loss = 0.0f; + + for (int step = 1; step <= STEPS; ++step) { + std::cout + << "\n--------------------------------------------\n" + << "STEP " << step << "/" << STEPS << "\n" + << "--------------------------------------------\n"; + + print_mem("before step"); + + if (mem_available_kb() < STEP_MIN_KB) { + std::cout + << "[SAFE STOP] < 1 GiB before step.\n" + << "Completed " << (step - 1) + << " steps safely.\n"; + return finish(0); + } + + // ---------------------------------------------------- + // FORWARD + // ---------------------------------------------------- + + auto logits = + model.forward(input_ids); + + if (!logits || + logits->dtype() != kFloat32 || + logits->shape().size() != 3 || + logits->shape()[0] != 1 || + logits->shape()[1] != SEQ || + logits->shape()[2] != VOCAB || + !finite_f32(logits)) { + + std::cerr + << "[FAIL] Invalid logits at step " + << step << "\n"; + return finish(4); + } + + // ---------------------------------------------------- + // LOSS + // ---------------------------------------------------- + + auto loss = + mse_loss( + logits, + target, + "mean"); + + if (!loss || + !std::isfinite(loss->item())) { + + std::cerr + << "[FAIL] Invalid loss at step " + << step << "\n"; + return finish(5); + } + + const float loss_value = + loss->item(); + + if (step == 1) + first_loss = loss_value; + + last_loss = loss_value; + + std::cout + << "Loss: " + << loss_value + << "\n"; + + // ---------------------------------------------------- + // BACKWARD + // ---------------------------------------------------- + + if (mem_available_kb() < STEP_MIN_KB) { + std::cout + << "[SAFE STOP] < 1 GiB before backward.\n"; + return finish(0); + } + + loss->backward(); + + std::vector params; + std::vector grads; + + size_t grad_count = 0; + size_t invalid_grads = 0; + + for (const auto& p : model.parameters()) { + if (!p) + continue; + + auto g = p->grad(); + + if (!g) + continue; + + ++grad_count; + + if (g->dtype() != kFloat32 || + !finite_f32(g)) { + + ++invalid_grads; + continue; + } + + params.push_back(p); + grads.push_back(g); + } + + std::cout + << "Gradients: " + << grad_count + << "\n"; + + std::cout + << "Invalid gradients: " + << invalid_grads + << "\n"; + + if (grad_count == 0 || + invalid_grads != 0) { + + std::cerr + << "[FAIL] Gradient validation failed.\n"; + return finish(6); + } + + // ---------------------------------------------------- + // ADAM + // ---------------------------------------------------- + + if (mem_available_kb() < STEP_MIN_KB) { + std::cout + << "[SAFE STOP] < 1 GiB before Adam.\n"; + return finish(0); + } + + adam.step(params, grads); + + // ---------------------------------------------------- + // PARAMETER CHECK + // ---------------------------------------------------- + + size_t bf16_params = 0; + size_t invalid_params = 0; + + for (const auto& p : params) { + if (!p) + continue; + + if (p->dtype() != kBFloat16) { + ++invalid_params; + continue; + } + + ++bf16_params; + + if (!finite_bf16(p)) + ++invalid_params; + } + + std::cout + << "BF16 parameters: " + << bf16_params + << "\n"; + + std::cout + << "Invalid parameters: " + << invalid_params + << "\n"; + + if (bf16_params != params.size() || + invalid_params != 0) { + + std::cerr + << "[FAIL] Parameter validation failed.\n"; + return finish(7); + } + + // ---------------------------------------------------- + // CLEAR GRADIENTS + // ---------------------------------------------------- + + for (auto& p : model.parameters()) { + if (p) + p->zero_grad(); + } + + // Release graph references before next step. + loss.reset(); + logits.reset(); + + print_mem("after step"); + } + + // -------------------------------------------------------- + // FINAL RESULT + // -------------------------------------------------------- + + std::cout + << "\n============================================\n" + << "10-STEP BF16 FULL-FT COMPLETE\n" + << "============================================\n" + << "First loss: " << first_loss << "\n" + << "Last loss: " << last_loss << "\n" + << "Delta: " + << (last_loss - first_loss) + << "\n"; + + if (!std::isfinite(first_loss) || + !std::isfinite(last_loss)) { + + std::cerr + << "[FAIL] Non-finite training loss.\n"; + return finish(8); + } + + std::cout + << "[PASS] 10 forward passes.\n" + << "[PASS] 10 backward passes.\n" + << "[PASS] 10 FP32-Adam updates.\n" + << "[PASS] Gradients cleared every step.\n" + << "[PASS] Parameters remained BF16.\n"; + + print_mem("final"); + + return finish(0); + + } catch (const std::exception& e) { + std::cerr + << "\n[ERROR] " + << e.what() + << "\n"; + + return finish(10); + } +} diff --git a/operator/finetune_ops/optim/test_gpt2_bf16_final_safe.cpp b/operator/finetune_ops/optim/test_gpt2_bf16_final_safe.cpp new file mode 100644 index 00000000..86709828 --- /dev/null +++ b/operator/finetune_ops/optim/test_gpt2_bf16_final_safe.cpp @@ -0,0 +1,370 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../graph/gpt2_model.h" +#include "../graph/safetensors_loader.h" +#include "../core/tensor.h" +#include "../core/ops.h" +#include "../core/dtype.h" +#include "adam.h" + +using namespace ops; + +static long long mem_available_kb() { + std::ifstream f("/proc/meminfo"); + std::string key, unit; + long long value; + + while (f >> key >> value >> unit) { + if (key == "MemAvailable:") + return value; + } + + return -1; +} + +static void print_mem(const char* phase) { + const auto kb = mem_available_kb(); + std::cout << "[MEM] " << phase << ": " + << (kb / 1024.0) << " MiB available\n"; +} + +static bool finite_tensor(const TensorPtr& t) { + if (!t || t->dtype() != kFloat32) + return false; + + const float* p = t->data(); + + for (int64_t i = 0; i < t->numel(); ++i) { + if (!std::isfinite(p[i])) + return false; + } + + return true; +} + +int main() { + constexpr long long HARD_STOP_KB = 600LL * 1024LL; + constexpr long long ADAM_MIN_KB = 1500LL * 1024LL; + + std::atomic stop_watchdog{false}; + + // Hard memory watchdog. + std::thread watchdog([&]() { + while (!stop_watchdog.load()) { + const long long kb = mem_available_kb(); + + if (kb > 0 && kb < HARD_STOP_KB) { + std::cerr + << "\n[SAFE ABORT] MemAvailable dropped below 600 MiB.\n"; + + std::_Exit(99); + } + + std::this_thread::sleep_for( + std::chrono::milliseconds(250)); + } + }); + + try { + const std::string model_dir = "/root/gpt2-tamil-124m"; + + std::cout << "============================================\n"; + std::cout << " FINAL 124M BF16 WEIGHT TEST\n"; + std::cout << "============================================\n"; + + print_mem("start"); + + // Require a reasonable starting margin. + if (mem_available_kb() < 1800LL * 1024LL) { + std::cout << "[SAFE STOP] Less than 1.8 GiB available.\n"; + stop_watchdog = true; + watchdog.join(); + return 0; + } + + // ---------- MODEL ---------- + GPT2Config config = GPT2Config::from_pretrained(model_dir); + GPT2Model model(config); + + if (config.tie_word_embeddings) + model.tie_weights(); + + // Load real checkpoint as BF16. + SafeTensorsModelReader reader(model_dir); + reader.parse_headers(); + + auto mapping = + GPT2KeyMapper::generate_gpt2_mapping(config.n_layer); + + SafeTensorsLoadOptions options; + options.transpose_linear = false; + options.auto_promote_fp16 = true; + options.convert_f32_to_bf16 = true; + options.verbose = false; + + auto tensors = + reader.load_tensors_mapped(mapping, options); + + size_t total_params = 0; + size_t bf16_tensors = 0; + size_t fp32_tensors = 0; + + for (const auto& kv : tensors) { + if (!kv.second) + continue; + + total_params += static_cast(kv.second->numel()); + + if (kv.second->dtype() == kBFloat16) + ++bf16_tensors; + else if (kv.second->dtype() == kFloat32) + ++fp32_tensors; + + model.assign_weight(kv.first, kv.second); + } + + std::cout << "Parameters: " << total_params << "\n"; + std::cout << "BF16 tensors: " << bf16_tensors << "\n"; + std::cout << "FP32 tensors: " << fp32_tensors << "\n"; + + if (total_params != 124439808 || + bf16_tensors != 148 || + fp32_tensors != 0) { + std::cerr << "[FAIL] BF16 weight verification failed.\n"; + stop_watchdog = true; + watchdog.join(); + return 2; + } + + std::cout << "[PASS] All 124M model parameters are BF16.\n"; + print_mem("after BF16 load"); + + // ---------- ENABLE TRAINING ---------- + for (auto& p : model.parameters()) { + if (p) + p->set_requires_grad(true); + } + + // Tiny real input: 8 known Tamil GPT-2 token IDs. + const int64_t ids[8] = { + 288, 275, 271, 265, + 281, 262, 742, 302 + }; + + auto input_ids = std::make_shared( + std::vector{1, 8}, + ids, + kInt64, + kCPU); + + std::cout << "\n===== FORWARD =====\n"; + print_mem("before forward"); + + auto logits = model.forward(input_ids); + + if (!logits || + logits->dtype() != kFloat32 || + logits->shape().size() != 3 || + logits->shape()[0] != 1 || + logits->shape()[1] != 8 || + !finite_tensor(logits)) { + std::cerr << "[FAIL] Forward produced invalid logits.\n"; + stop_watchdog = true; + watchdog.join(); + return 3; + } + + std::cout << "Logits shape: [" + << logits->shape()[0] << "," + << logits->shape()[1] << "," + << logits->shape()[2] << "]\n"; + + std::cout << "[PASS] 124M BF16-weight forward.\n"; + print_mem("after forward"); + + if (mem_available_kb() < 1000LL * 1024LL) { + std::cout + << "[SAFE STOP] Less than 1 GiB remains before backward.\n"; + stop_watchdog = true; + watchdog.join(); + return 0; + } + + // ---------- LOSS ---------- + auto target = std::make_shared( + logits->shape(), + kFloat32, + kCPU); + + float* target_data = target->data(); + + for (int64_t i = 0; i < target->numel(); ++i) + target_data[i] = 0.0f; + + auto loss = mse_loss(logits, target, "mean"); + + if (!loss || !std::isfinite(loss->item())) { + std::cerr << "[FAIL] Loss is invalid.\n"; + stop_watchdog = true; + watchdog.join(); + return 4; + } + + std::cout << "Loss: " << loss->item() << "\n"; + + // ---------- BACKWARD ---------- + std::cout << "\n===== BACKWARD =====\n"; + print_mem("before backward"); + + if (mem_available_kb() < 900LL * 1024LL) { + std::cout + << "[SAFE STOP] Less than 900 MiB remains before backward.\n"; + stop_watchdog = true; + watchdog.join(); + return 0; + } + + loss->backward(); + + size_t grad_count = 0; + size_t bad_grad_count = 0; + + for (const auto& p : model.parameters()) { + if (!p) + continue; + + auto g = p->grad(); + + if (g) { + ++grad_count; + + if (g->dtype() != kFloat32 || + !finite_tensor(g)) { + ++bad_grad_count; + } + } + } + + std::cout << "Gradients present: " << grad_count << "\n"; + std::cout << "Invalid gradients: " << bad_grad_count << "\n"; + + print_mem("after backward"); + + if (grad_count == 0 || bad_grad_count != 0) { + std::cerr << "[FAIL] Backward gradient validation failed.\n"; + stop_watchdog = true; + watchdog.join(); + return 5; + } + + std::cout << "[PASS] 124M backward with BF16 weights.\n"; + + // ---------- OPTIONAL ONE ADAM STEP ---------- + std::cout << "\n===== ONE ADAM STEP =====\n"; + + const long long before_adam = mem_available_kb(); + + std::cout << "MemAvailable before Adam: " + << before_adam / 1024.0 << " MiB\n"; + + if (before_adam < ADAM_MIN_KB) { + std::cout + << "[SAFE STOP] Not enough RAM margin for full FP32 Adam state.\n"; + std::cout + << "[RESULT] Forward + backward passed; Adam step skipped safely.\n"; + + stop_watchdog = true; + watchdog.join(); + return 0; + } + + std::vector params; + std::vector grads; + + for (const auto& p : model.parameters()) { + if (!p) + continue; + + auto g = p->grad(); + + if (g) { + params.push_back(p); + grads.push_back(g); + } + } + + AdamConfig adam_cfg; + adam_cfg.learning_rate = 1e-5f; + adam_cfg.beta1 = 0.9f; + adam_cfg.beta2 = 0.999f; + adam_cfg.epsilon = 1e-8f; + adam_cfg.weight_decay = 0.0f; + adam_cfg.amsgrad = false; + + Adam adam(adam_cfg); + + adam.step(params, grads); + + print_mem("after Adam"); + + // Verify parameters remain BF16 and finite. + size_t still_bf16 = 0; + size_t invalid_params = 0; + + for (const auto& p : params) { + if (p->dtype() == kBFloat16) { + ++still_bf16; + + const uint16_t* d = p->data(); + + for (int64_t j = 0; j < p->numel(); ++j) { + float v = bf16_bits_to_float32(d[j]); + + if (!std::isfinite(v)) { + ++invalid_params; + break; + } + } + } + } + + std::cout << "BF16 parameters after Adam: " + << still_bf16 << "\n"; + std::cout << "Invalid parameter tensors: " + << invalid_params << "\n"; + + if (still_bf16 != params.size() || + invalid_params != 0) { + std::cerr << "[FAIL] BF16 parameters became invalid after Adam.\n"; + stop_watchdog = true; + watchdog.join(); + return 6; + } + + std::cout << "[PASS] One full FP32-Adam update completed.\n"; + std::cout << "[PASS] Parameters remain BF16.\n"; + + std::cout << "\n============================================\n"; + std::cout << " FINAL RESULT: ALL REQUESTED STAGES PASSED\n"; + std::cout << "============================================\n"; + + stop_watchdog = true; + watchdog.join(); + return 0; + + } catch (const std::exception& e) { + std::cerr << "\n[ERROR] " << e.what() << "\n"; + stop_watchdog = true; + watchdog.join(); + return 10; + } +} diff --git a/operator/finetune_ops/optim/test_gpt2_bf16_forward_safe.cpp b/operator/finetune_ops/optim/test_gpt2_bf16_forward_safe.cpp new file mode 100644 index 00000000..3b2d533e --- /dev/null +++ b/operator/finetune_ops/optim/test_gpt2_bf16_forward_safe.cpp @@ -0,0 +1,189 @@ +#include +#include +#include +#include +#include +#include + +#include "../graph/gpt2_model.h" +#include "../graph/safetensors_loader.h" +#include "../core/tensor.h" +#include "../core/dtype.h" + +using namespace ops; + +static long long mem_available_kb() { + std::ifstream f("/proc/meminfo"); + std::string key; + long long value; + std::string unit; + + while (f >> key >> value >> unit) { + if (key == "MemAvailable:") { + return value; + } + } + + return -1; +} + +int main() { + try { + const long long available_kb = mem_available_kb(); + + std::cout << "===== BF16 124M FORWARD SAFETY CHECK =====\n"; + std::cout << "MemAvailable: " << available_kb << " kB\n"; + + // Conservative phone-safety gate. + // We require about 2 GiB available before attempting this forward. + if (available_kb > 0 && available_kb < 2LL * 1024LL * 1024LL) { + std::cout << "[SAFE STOP] Less than 2 GiB available.\n"; + std::cout << "[SAFE STOP] Skipping 124M forward.\n"; + return 0; + } + + const std::string model_dir = "/root/gpt2-tamil-124m"; + + GPT2Config config = GPT2Config::from_pretrained(model_dir); + GPT2Model model(config); + + if (config.tie_word_embeddings) { + model.tie_weights(); + } + + // Forward-only test: no autograd graph for model parameters. + for (auto& p : model.parameters()) { + if (p) { + p->set_requires_grad(false); + } + } + + SafeTensorsModelReader reader(model_dir); + reader.parse_headers(); + + auto mapping = + GPT2KeyMapper::generate_gpt2_mapping(config.n_layer); + + SafeTensorsLoadOptions options; + options.transpose_linear = false; + options.auto_promote_fp16 = true; + options.convert_f32_to_bf16 = true; + options.verbose = false; + + auto tensors = + reader.load_tensors_mapped(mapping, options); + + size_t parameter_bytes = 0; + size_t parameter_count = 0; + size_t bf16_count = 0; + size_t fp32_count = 0; + + for (const auto& kv : tensors) { + const auto& t = kv.second; + if (!t) continue; + + parameter_count += static_cast(t->numel()); + parameter_bytes += + static_cast(t->numel()) * + DTypeUtils::size_of(t->dtype()); + + if (t->dtype() == kBFloat16) { + ++bf16_count; + } else if (t->dtype() == kFloat32) { + ++fp32_count; + } + + model.assign_weight(kv.first, t); + } + + std::cout << "Parameters: " << parameter_count << "\n"; + std::cout << "BF16 tensors: " << bf16_count << "\n"; + std::cout << "FP32 tensors: " << fp32_count << "\n"; + std::cout << "Weight MiB: " + << (static_cast(parameter_bytes) / + (1024.0 * 1024.0)) + << "\n"; + + if (bf16_count == 0 || fp32_count != 0) { + std::cerr << "[FAIL] Weight dtype verification failed.\n"; + return 2; + } + + /* + * Very small sequence deliberately chosen for the first real-model + * forward. This keeps activation memory low. + */ + const int64_t ids[8] = { + 288, 275, 271, 265, + 281, 262, 742, 302 + }; + + auto input_ids = std::make_shared( + std::vector{1, 8}, + ids, + kInt64, + kCPU); + + std::cout << "Running forward: batch=1 seq=8\n"; + + auto logits = model.forward(input_ids); + + if (!logits) { + std::cerr << "[FAIL] Forward returned null.\n"; + return 3; + } + + const auto& shape = logits->shape(); + + std::cout << "Logits shape: ["; + for (size_t i = 0; i < shape.size(); ++i) { + if (i) std::cout << ","; + std::cout << shape[i]; + } + std::cout << "]\n"; + + if (shape.size() != 3 || + shape[0] != 1 || + shape[1] != 8) { + std::cerr << "[FAIL] Unexpected logits shape.\n"; + return 4; + } + + if (logits->dtype() != kFloat32) { + std::cerr << "[FAIL] Logits are not FP32.\n"; + return 5; + } + + const float* data = logits->data(); + bool finite = true; + + for (int i = 0; i < 16; ++i) { + if (!std::isfinite(data[i])) { + finite = false; + break; + } + } + + if (!finite) { + std::cerr << "[FAIL] Non-finite logits detected.\n"; + return 6; + } + + std::cout << "First logits:"; + for (int i = 0; i < 8; ++i) { + std::cout << " " << data[i]; + } + std::cout << "\n"; + + std::cout << "[PASS] 124M BF16 weights loaded.\n"; + std::cout << "[PASS] 124M forward completed.\n"; + std::cout << "[PASS] FP32 logits are finite.\n"; + std::cout << "[PASS] No backward/optimizer/training executed.\n"; + + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 10; + } +} diff --git a/operator/finetune_ops/optim/test_gpt2_bf16_load_only.cpp b/operator/finetune_ops/optim/test_gpt2_bf16_load_only.cpp new file mode 100644 index 00000000..fd7e2467 --- /dev/null +++ b/operator/finetune_ops/optim/test_gpt2_bf16_load_only.cpp @@ -0,0 +1,108 @@ +#include +#include +#include +#include + +#include "../graph/safetensors_loader.h" +#include "../graph/gpt2_model.h" + +using namespace ops; + +int main() { + const std::string model_dir = "/root/gpt2-tamil-124m"; + + try { + GPT2Config config = GPT2Config::from_pretrained(model_dir); + GPT2Model model(config); + + if (config.tie_word_embeddings) { + model.tie_weights(); + } + + SafeTensorsModelReader reader(model_dir); + reader.parse_headers(); + + auto mapping = GPT2KeyMapper::generate_gpt2_mapping(config.n_layer); + + SafeTensorsLoadOptions options; + options.transpose_linear = false; + options.auto_promote_fp16 = true; + options.convert_f32_to_bf16 = true; + options.verbose = false; + + /* + * Critical safety feature: + * request low-precision storage for every mapped GPT-2 weight. + * + * The loader must support F32 -> BF16 conversion for this to pass. + */ + options.preserve_low_precision_key_substrings.clear(); + + auto tensors = reader.load_tensors_mapped(mapping, options); + + size_t total_params = 0; + size_t total_bytes = 0; + size_t bf16_tensors = 0; + size_t fp32_tensors = 0; + + for (const auto& kv : tensors) { + const auto& name = kv.first; + const auto& tensor = kv.second; + + total_params += static_cast(tensor->numel()); + total_bytes += static_cast(tensor->numel()) * + DTypeUtils::size_of(tensor->dtype()); + + if (tensor->dtype() == kBFloat16) { + ++bf16_tensors; + } else if (tensor->dtype() == kFloat32) { + ++fp32_tensors; + } + + std::cout << "[BF16-LOAD] " << name + << " dtype=" << DTypeUtils::to_string(tensor->dtype()) + << " numel=" << tensor->numel() + << " bytes=" + << (tensor->numel() * DTypeUtils::size_of(tensor->dtype())) + << "\n"; + } + + const double mib = static_cast(total_bytes) / (1024.0 * 1024.0); + const double fp32_mib = + static_cast(total_params * sizeof(float)) / + (1024.0 * 1024.0); + + std::cout << "\n========== BF16 LOAD CHECK ==========\n"; + std::cout << "Parameters: " << total_params << "\n"; + std::cout << "BF16 tensors: " << bf16_tensors << "\n"; + std::cout << "FP32 tensors: " << fp32_tensors << "\n"; + std::cout << "Parameter bytes: " << total_bytes << "\n"; + std::cout << "Parameter MiB: " << mib << "\n"; + std::cout << "FP32 baseline: " << fp32_mib << " MiB\n"; + + if (bf16_tensors == 0) { + std::cerr << "[FAIL] No BF16 tensors were produced.\n"; + return 2; + } + + if (fp32_tensors != 0) { + std::cerr << "[FAIL] Some parameters are still FP32.\n"; + return 3; + } + + if (total_bytes >= total_params * sizeof(float)) { + std::cerr << "[FAIL] Parameter storage was not reduced.\n"; + return 4; + } + + std::cout << "[PASS] Parameters loaded as BF16.\n"; + std::cout << "[PASS] No FP32 parameter tensors remain in this load result.\n"; + std::cout << "[PASS] Test exits before forward/backward/optimizer.\n"; + + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << "\n"; + return 1; + } +} diff --git a/operator/finetune_ops/optim/test_gpt2_bf16_real_causal_ft.cpp b/operator/finetune_ops/optim/test_gpt2_bf16_real_causal_ft.cpp new file mode 100644 index 00000000..e80fa452 --- /dev/null +++ b/operator/finetune_ops/optim/test_gpt2_bf16_real_causal_ft.cpp @@ -0,0 +1,1025 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../graph/gpt2_model.h" +#include "../graph/safetensors_loader.h" +#include "../core/tensor.h" +#include "../core/ops.h" +#include "../core/dtype.h" +#include "../core/tokenizer.h" +#include "../core/lm_loss.h" +#include "adam.h" + +using namespace ops; + +static long long mem_available_kb() { + std::ifstream f("/proc/meminfo"); + + std::string key; + std::string unit; + long long value = 0; + + while (f >> key >> value >> unit) { + if (key == "MemAvailable:") + return value; + } + + return -1; +} + +static void print_mem(const char* name) { + const long long kb = mem_available_kb(); + + std::cout + << "[MEM] " + << name + << ": " + << kb / 1024.0 + << " MiB available\n"; +} + +static bool finite_f32(const TensorPtr& t) { + if (!t || t->dtype() != kFloat32) + return false; + + const float* p = t->data(); + + for (int64_t i = 0; i < t->numel(); ++i) { + if (!std::isfinite(p[i])) + return false; + } + + return true; +} + +static bool finite_bf16(const TensorPtr& t) { + if (!t || t->dtype() != kBFloat16) + return false; + + const uint16_t* p = t->data(); + + for (int64_t i = 0; i < t->numel(); ++i) { + if (!std::isfinite(bf16_bits_to_float32(p[i]))) + return false; + } + + return true; +} + +/* + * Dataset format: + * + * {"classic":"Tamil text...","description":"..."} + * + * Only "classic" is used. + * + * This parser handles the JSON escaping used by the current + * local dataset, including \n, \r, \t, \" and \\. + */ +static bool extract_json_string( + const std::string& line, + const std::string& field, + std::string& result) +{ + const std::string key = "\"" + field + "\""; + + const size_t key_pos = + line.find(key); + + if (key_pos == std::string::npos) + return false; + + const size_t colon = + line.find(':', key_pos + key.size()); + + if (colon == std::string::npos) + return false; + + const size_t first_quote = + line.find('"', colon + 1); + + if (first_quote == std::string::npos) + return false; + + result.clear(); + + bool escaped = false; + + for (size_t i = first_quote + 1; + i < line.size(); + ++i) { + + const char c = line[i]; + + if (escaped) { + switch (c) { + case 'n': + result.push_back('\n'); + break; + + case 'r': + result.push_back('\r'); + break; + + case 't': + result.push_back('\t'); + break; + + case '"': + result.push_back('"'); + break; + + case '\\': + result.push_back('\\'); + break; + + case '/': + result.push_back('/'); + break; + + default: + result.push_back(c); + break; + } + + escaped = false; + continue; + } + + if (c == '\\') { + escaped = true; + continue; + } + + if (c == '"') + break; + + result.push_back(c); + } + + return !result.empty(); +} + +static std::vector load_classic_records( + const std::string& path) +{ + std::ifstream f(path); + + if (!f) { + throw std::runtime_error( + "Cannot open training dataset: " + path); + } + + std::vector records; + + std::string line; + + while (std::getline(f, line)) { + + if (line.empty()) + continue; + + std::string classic; + + if (extract_json_string( + line, + "classic", + classic)) { + + if (!classic.empty()) + records.push_back(classic); + } + } + + return records; +} + +int main() { + /* + * Conservative first real-corpus run. + * + * 40 real training records exist locally. + * We perform 20 optimizer updates, cycling through the records. + * + * MAX_SEQ is intentionally 32 for the first real-data run. + * The previous 124M tests were sequence length 8; increasing + * sequence length substantially increases activation memory. + */ + constexpr int MAX_STEPS = 20; + constexpr int64_t MAX_SEQ = 32; + constexpr int64_t VOCAB_SIZE = 50257; + + /* + * Safety limits. + */ + constexpr long long START_MIN_KB = + 1800LL * 1024LL; + + constexpr long long STEP_MIN_KB = + 1000LL * 1024LL; + + constexpr long long HARD_STOP_KB = + 600LL * 1024LL; + + std::atomic stop_watchdog{false}; + + std::thread watchdog([&]() { + while (!stop_watchdog.load()) { + + const long long kb = + mem_available_kb(); + + if (kb > 0 && + kb < HARD_STOP_KB) { + + std::cerr + << "\n[SAFE ABORT] " + << "MemAvailable < 600 MiB\n"; + + std::_Exit(99); + } + + std::this_thread::sleep_for( + std::chrono::milliseconds(250)); + } + }); + + auto finish = [&](int code) { + stop_watchdog = true; + watchdog.join(); + return code; + }; + + try { + const std::string model_dir = + "/root/gpt2-tamil-124m"; + + const std::string train_path = + "/root/classical-tamil-ppe-smoke/train.jsonl"; + + std::cout + << "============================================\n" + << " NATIVE 124M BF16 REAL-DATA FINE-TUNING\n" + << "============================================\n"; + + print_mem("start"); + + if (mem_available_kb() < START_MIN_KB) { + + std::cout + << "[SAFE STOP] " + << "Need >= 1.8 GiB at startup.\n"; + + return finish(0); + } + + // ======================================================== + // REAL DATASET + // ======================================================== + + const auto records = + load_classic_records(train_path); + + std::cout + << "Training records: " + << records.size() + << "\n"; + + if (records.empty()) + throw std::runtime_error( + "Training dataset contains no classic records."); + + // ======================================================== + // NATIVE GPT-2 TOKENIZER + // ======================================================== + + TokenizerLoadOptions tokenizer_options; + + tokenizer_options.model_type = "gpt2"; + + auto tokenizer = + TokenizerFactory::from_pretrained( + model_dir, + tokenizer_options); + + if (!tokenizer) + throw std::runtime_error( + "Native GPT-2 tokenizer failed to load."); + + std::cout + << "[PASS] Native GPT-2 tokenizer loaded.\n"; + + // ======================================================== + // MODEL + // ======================================================== + + GPT2Config config = + GPT2Config::from_pretrained( + model_dir); + + GPT2Model model(config); + + if (config.tie_word_embeddings) + model.tie_weights(); + + // ======================================================== + // BF16 MODEL LOAD + // ======================================================== + + SafeTensorsModelReader reader( + model_dir); + + reader.parse_headers(); + + auto mapping = + GPT2KeyMapper::generate_gpt2_mapping( + config.n_layer); + + SafeTensorsLoadOptions load_options; + + load_options.transpose_linear = false; + load_options.auto_promote_fp16 = true; + load_options.convert_f32_to_bf16 = true; + load_options.verbose = false; + + auto tensors = + reader.load_tensors_mapped( + mapping, + load_options); + + size_t total_params = 0; + size_t bf16_tensors = 0; + size_t fp32_tensors = 0; + + for (const auto& kv : tensors) { + + if (!kv.second) + continue; + + total_params += + static_cast( + kv.second->numel()); + + if (kv.second->dtype() == kBFloat16) + ++bf16_tensors; + + else if (kv.second->dtype() == kFloat32) + ++fp32_tensors; + + model.assign_weight( + kv.first, + kv.second); + } + + std::cout + << "Parameters: " + << total_params + << "\n"; + + std::cout + << "BF16 tensors: " + << bf16_tensors + << "\n"; + + std::cout + << "FP32 tensors: " + << fp32_tensors + << "\n"; + + if (total_params != 124439808 || + bf16_tensors != 148 || + fp32_tensors != 0) { + + std::cerr + << "[FAIL] Expected 124M BF16 model.\n"; + + return finish(2); + } + + std::cout + << "[PASS] 124M parameters are BF16.\n"; + + print_mem("after model load"); + + // ======================================================== + // FULL PARAMETER TRAINING + // ======================================================== + + for (auto& p : + model.parameters()) { + + if (p) + p->set_requires_grad(true); + } + + // ======================================================== + // ADAM + // ======================================================== + + AdamConfig adam_config; + + adam_config.learning_rate = 5e-6f; + adam_config.beta1 = 0.9f; + adam_config.beta2 = 0.999f; + adam_config.epsilon = 1e-8f; + adam_config.weight_decay = 0.0f; + adam_config.amsgrad = false; + + Adam adam(adam_config); + + float first_loss = 0.0f; + float last_loss = 0.0f; + + bool have_loss = false; + + int completed_steps = 0; + + // ======================================================== + // REAL-DATA CAUSAL-LM LOOP + // ======================================================== + + for (int step = 0; + step < MAX_STEPS; + ++step) { + + std::cout + << "\n--------------------------------------------\n" + << "STEP " + << (step + 1) + << "/" + << MAX_STEPS + << "\n" + << "--------------------------------------------\n"; + + print_mem("before step"); + + if (mem_available_kb() < STEP_MIN_KB) { + + std::cout + << "[SAFE STOP] " + << "Memory below 1 GiB.\n"; + + break; + } + + const std::string& text = + records[ + static_cast(step) + % records.size() + ]; + + // ==================================================== + // TOKENIZE REAL TAMIL TEXT + // ==================================================== + + const std::vector ids = + tokenizer->encode(text); + + if (ids.size() < 2) { + + std::cout + << "[SKIP] Record produced <2 tokens.\n"; + + continue; + } + + /* + * We need N input tokens and N target tokens: + * + * input : t0 t1 t2 ... t(N-1) + * target: t1 t2 t3 ... tN + * + * Therefore usable token count is MAX_SEQ. + */ + const size_t usable = + std::min( + ids.size(), + static_cast( + MAX_SEQ + 1)); + + if (usable < 2) { + + std::cout + << "[SKIP] Record too short.\n"; + + continue; + } + + const int64_t seq_len = + static_cast( + usable - 1); + + std::vector + input_values( + static_cast( + seq_len)); + + std::vector + label_values( + static_cast( + seq_len)); + + for (int64_t i = 0; + i < seq_len; + ++i) { + + input_values[ + static_cast(i)] + = static_cast( + ids[ + static_cast(i)]); + + label_values[ + static_cast(i)] + = static_cast( + ids[ + static_cast(i + 1)]); + } + + // ==================================================== + // INPUT / LABEL TENSORS + // ==================================================== + + auto input_ids = + std::make_shared( + std::vector{ + 1, + seq_len + }, + input_values.data(), + kInt64, + kCPU); + + auto labels = + std::make_shared( + std::vector{ + 1, + seq_len + }, + label_values.data(), + kInt32, + kCPU); + + // ==================================================== + // NATIVE GPT-2 FORWARD + // ==================================================== + + auto logits = + model.forward(input_ids); + + if (!logits || + logits->dtype() != kFloat32 || + logits->shape().size() != 3 || + logits->shape()[0] != 1 || + logits->shape()[1] != seq_len || + logits->shape()[2] != VOCAB_SIZE || + !finite_f32(logits)) { + + std::cerr + << "[FAIL] Invalid logits at step " + << (step + 1) + << "\n"; + + return finish(4); + } + + // ==================================================== + // REAL CAUSAL LANGUAGE-MODELING OBJECTIVE + // ==================================================== + + auto loss = + lm_cross_entropy( + logits, + labels, + -100, + "mean"); + + if (!loss || + !std::isfinite( + loss->item())) { + + std::cerr + << "[FAIL] Invalid causal-LM loss at step " + << (step + 1) + << "\n"; + + return finish(5); + } + + const float loss_value = + loss->item(); + + if (!have_loss) { + first_loss = loss_value; + have_loss = true; + } + + last_loss = loss_value; + + std::cout + << "Tokens: " + << seq_len + << "\n"; + + std::cout + << "Causal-LM loss: " + << loss_value + << "\n"; + + // ==================================================== + // BACKWARD + // ==================================================== + + if (mem_available_kb() < STEP_MIN_KB) { + + std::cout + << "[SAFE STOP] " + << "Memory below 1 GiB before backward.\n"; + + break; + } + + loss->backward(); + + std::vector + params; + + std::vector + grads; + + params.reserve( + model.parameters().size()); + + grads.reserve( + model.parameters().size()); + + size_t grad_count = 0; + size_t invalid_grads = 0; + + for (const auto& p : + model.parameters()) { + + if (!p) + continue; + + auto g = p->grad(); + + if (!g) + continue; + + ++grad_count; + + if (g->dtype() != kFloat32 || + !finite_f32(g)) { + + ++invalid_grads; + continue; + } + + params.push_back(p); + grads.push_back(g); + } + + std::cout + << "Gradients: " + << grad_count + << "\n"; + + std::cout + << "Invalid gradients: " + << invalid_grads + << "\n"; + + if (grad_count == 0 || + invalid_grads != 0) { + + std::cerr + << "[FAIL] Invalid gradient set.\n"; + + return finish(6); + } + + // ==================================================== + // FP32 ADAM → BF16 PARAMETERS + // ==================================================== + + if (mem_available_kb() < STEP_MIN_KB) { + + std::cout + << "[SAFE STOP] " + << "Memory below 1 GiB before Adam.\n"; + + break; + } + + adam.step( + params, + grads); + + // ==================================================== + // VERIFY BF16 PARAMETERS + // ==================================================== + + size_t bf16_params = 0; + size_t invalid_params = 0; + + for (const auto& p : + params) { + + if (!p || + p->dtype() != kBFloat16) { + + ++invalid_params; + continue; + } + + ++bf16_params; + + if (!finite_bf16(p)) + ++invalid_params; + } + + std::cout + << "BF16 parameters: " + << bf16_params + << "\n"; + + std::cout + << "Invalid parameters: " + << invalid_params + << "\n"; + + if (bf16_params != params.size() || + invalid_params != 0) { + + std::cerr + << "[FAIL] BF16 parameter verification failed.\n"; + + return finish(7); + } + + // ==================================================== + // CLEAR GRADIENTS + // ==================================================== + + for (auto& p : + model.parameters()) { + + if (p) + p->zero_grad(); + } + + /* + * Release this step's graph before the next dataset + * record. This is essential for stable long-running + * training on the phone. + */ + loss.reset(); + logits.reset(); + input_ids.reset(); + labels.reset(); + + ++completed_steps; + + print_mem("after step"); + } + + // ======================================================== + // RESULT + // ======================================================== + + std::cout + << "\n============================================\n" + << "NATIVE REAL-DATA BF16 TRAINING RESULT\n" + << "============================================\n"; + + std::cout + << "Completed steps: " + << completed_steps + << "\n"; + + if (have_loss) { + + std::cout + << "First causal-LM loss: " + << first_loss + << "\n"; + + std::cout + << "Last causal-LM loss: " + << last_loss + << "\n"; + + std::cout + << "Loss delta: " + << (last_loss - first_loss) + << "\n"; + } + + if (completed_steps == 0) { + + std::cerr + << "[FAIL] No training steps completed.\n"; + + return finish(8); + } + + if (!have_loss || + !std::isfinite(first_loss) || + !std::isfinite(last_loss)) { + + std::cerr + << "[FAIL] Invalid final loss.\n"; + + return finish(9); + } + + std::cout + << "[PASS] Real JSONL Classical-Tamil data.\n" + << "[PASS] Native GPT-2 tokenizer.\n" + << "[PASS] Causal next-token targets.\n" + << "[PASS] Native lm_cross_entropy.\n" + << "[PASS] Native backward.\n" + << "[PASS] FP32 gradients.\n" + << "[PASS] FP32 Adam.\n" + << "[PASS] BF16 parameter write-back.\n" + << "[PASS] Gradient reset every step.\n"; + + // ======================================================== + // SAVE TRAINED BF16 MODEL + // ======================================================== + // + // Native checkpoint format: + // + // magic[8] + // tensor_count(uint64) + // + // repeated: + // key_length(uint64) + // key bytes + // dtype(uint32) + // ndim(uint64) + // shape[ndim](int64) + // raw tensor bytes(uint64 + data) + // + // No FP32 conversion is performed. + // + const std::string output_checkpoint = + "/root/tamil_gpt2_bf16_trained.bin"; + + std::cout + << "\n===== SAVING TRAINED BF16 MODEL =====\n"; + + { + std::ofstream out( + output_checkpoint, + std::ios::binary); + + if (!out) + throw std::runtime_error( + "Cannot create trained BF16 checkpoint."); + + const char magic[8] = { + 'M','F','T','B','F','1','6','1' + }; + + out.write( + magic, + sizeof(magic)); + + uint64_t tensor_count = 0; + + for (const auto& kv : tensors) { + if (kv.second) + ++tensor_count; + } + + out.write( + reinterpret_cast(&tensor_count), + sizeof(tensor_count)); + + for (const auto& kv : tensors) { + const std::string& key = kv.first; + const TensorPtr& t = kv.second; + + if (!t) + continue; + + if (t->dtype() != kBFloat16) { + throw std::runtime_error( + "Attempted to save non-BF16 tensor: " + key); + } + + const uint64_t key_len = + static_cast(key.size()); + + out.write( + reinterpret_cast(&key_len), + sizeof(key_len)); + + out.write( + key.data(), + static_cast(key.size())); + + const uint32_t dtype = + static_cast(t->dtype()); + + out.write( + reinterpret_cast(&dtype), + sizeof(dtype)); + + const uint64_t ndim = + static_cast( + t->shape().size()); + + out.write( + reinterpret_cast(&ndim), + sizeof(ndim)); + + for (const auto dim : t->shape()) { + const int64_t d = + static_cast(dim); + + out.write( + reinterpret_cast(&d), + sizeof(d)); + } + + const uint64_t data_bytes = + static_cast( + t->numel() * sizeof(uint16_t)); + + out.write( + reinterpret_cast(&data_bytes), + sizeof(data_bytes)); + + out.write( + reinterpret_cast( + t->data()), + static_cast( + data_bytes)); + + if (!out) + throw std::runtime_error( + "Write failed for tensor: " + key); + } + + out.flush(); + + if (!out) + throw std::runtime_error( + "Failed flushing trained checkpoint."); + } + + // Verify the file exists and is non-empty. + { + std::ifstream in( + output_checkpoint, + std::ios::binary | std::ios::ate); + + if (!in) + throw std::runtime_error( + "Cannot reopen saved checkpoint."); + + const auto size = in.tellg(); + + if (size <= 0) + throw std::runtime_error( + "Saved checkpoint is empty."); + + std::cout + << "Saved checkpoint: " + << output_checkpoint + << "\n"; + + std::cout + << "Checkpoint bytes: " + << size + << "\n"; + } + + print_mem("final"); + + return finish(0); + + } catch (const std::exception& e) { + + std::cerr + << "\n[ERROR] " + << e.what() + << "\n"; + + return finish(10); + } +} diff --git a/operator/finetune_ops/optim/test_gpt2_bf16_ultimate.cpp b/operator/finetune_ops/optim/test_gpt2_bf16_ultimate.cpp new file mode 100644 index 00000000..6689b8c6 --- /dev/null +++ b/operator/finetune_ops/optim/test_gpt2_bf16_ultimate.cpp @@ -0,0 +1,1124 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "../graph/gpt2_model.h" +#include "../graph/safetensors_loader.h" +#include "../core/tensor.h" +#include "../core/ops.h" +#include "../core/dtype.h" +#include "../core/tokenizer.h" +#include "../core/lm_loss.h" +#include "adam.h" + +using namespace ops; + +static long long mem_available_kb() { + std::ifstream f("/proc/meminfo"); + std::string key, unit; + long long value = 0; + + while (f >> key >> value >> unit) { + if (key == "MemAvailable:") + return value; + } + + return -1; +} + +static void print_mem(const char* label) { + const auto kb = mem_available_kb(); + + std::cout + << "[MEM] " + << label + << ": " + << kb / 1024.0 + << " MiB available\n"; +} + +static bool finite_f32(const TensorPtr& t) { + if (!t || t->dtype() != kFloat32) + return false; + + const float* p = t->data(); + + for (int64_t i = 0; i < t->numel(); ++i) { + if (!std::isfinite(p[i])) + return false; + } + + return true; +} + +static bool finite_bf16(const TensorPtr& t) { + if (!t || t->dtype() != kBFloat16) + return false; + + const uint16_t* p = t->data(); + + for (int64_t i = 0; i < t->numel(); ++i) { + if (!std::isfinite(bf16_bits_to_float32(p[i]))) + return false; + } + + return true; +} + +static uint64_t bf16_hash(const TensorPtr& t) { + if (!t || t->dtype() != kBFloat16) + return 0; + + const uint16_t* p = t->data(); + + uint64_t h = 1469598103934665603ULL; + + for (int64_t i = 0; i < t->numel(); ++i) { + h ^= static_cast(p[i]); + h *= 1099511628211ULL; + } + + return h; +} + +static bool extract_json_string( + const std::string& line, + const std::string& field, + std::string& result) +{ + const std::string key = "\"" + field + "\""; + + const size_t key_pos = line.find(key); + + if (key_pos == std::string::npos) + return false; + + const size_t colon = + line.find(':', key_pos + key.size()); + + if (colon == std::string::npos) + return false; + + const size_t first_quote = + line.find('"', colon + 1); + + if (first_quote == std::string::npos) + return false; + + result.clear(); + + bool escaped = false; + + for (size_t i = first_quote + 1; + i < line.size(); + ++i) { + + const char c = line[i]; + + if (escaped) { + switch (c) { + case 'n': + result.push_back('\n'); + break; + case 'r': + result.push_back('\r'); + break; + case 't': + result.push_back('\t'); + break; + case '"': + result.push_back('"'); + break; + case '\\': + result.push_back('\\'); + break; + case '/': + result.push_back('/'); + break; + default: + result.push_back(c); + break; + } + + escaped = false; + continue; + } + + if (c == '\\') { + escaped = true; + continue; + } + + if (c == '"') + break; + + result.push_back(c); + } + + return !result.empty(); +} + +static std::vector load_classic_records( + const std::string& path) +{ + std::ifstream f(path); + + if (!f) + throw std::runtime_error( + "Cannot open dataset: " + path); + + std::vector rows; + + std::string line; + + while (std::getline(f, line)) { + if (line.empty()) + continue; + + std::string text; + + if (extract_json_string( + line, + "classic", + text)) { + + rows.push_back(text); + } + } + + return rows; +} + +static double evaluate_dataset( + GPT2Model& model, + Tokenizer& tokenizer, + const std::vector& records, + int64_t max_seq) +{ + double total_loss = 0.0; + int used = 0; + + // Evaluation should not retain parameter gradients. + for (auto& p : model.parameters()) { + if (p) + p->set_requires_grad(false); + } + + for (const auto& text : records) { + const auto ids = tokenizer.encode(text); + + if (ids.size() < 2) + continue; + + const size_t usable = + std::min( + ids.size(), + static_cast(max_seq + 1)); + + if (usable < 2) + continue; + + const int64_t S = + static_cast(usable - 1); + + std::vector input_values( + static_cast(S)); + + std::vector label_values( + static_cast(S)); + + for (int64_t i = 0; i < S; ++i) { + input_values[ + static_cast(i)] + = static_cast( + ids[static_cast(i)]); + + label_values[ + static_cast(i)] + = static_cast( + ids[static_cast(i + 1)]); + } + + auto input_ids = + std::make_shared( + std::vector{1, S}, + input_values.data(), + kInt64, + kCPU); + + auto labels = + std::make_shared( + std::vector{1, S}, + label_values.data(), + kInt32, + kCPU); + + auto logits = + model.forward(input_ids); + + if (!logits || + logits->dtype() != kFloat32 || + !finite_f32(logits)) { + + throw std::runtime_error( + "Validation produced invalid logits."); + } + + auto loss = + lm_cross_entropy( + logits, + labels, + -100, + "mean"); + + if (!loss || + !std::isfinite(loss->item())) { + + throw std::runtime_error( + "Validation produced invalid loss."); + } + + total_loss += + static_cast(loss->item()); + + ++used; + + loss.reset(); + logits.reset(); + } + + // Clear anything left and re-enable training. + for (auto& p : model.parameters()) { + if (p) { + p->zero_grad(); + p->set_requires_grad(true); + } + } + + if (used == 0) + throw std::runtime_error( + "Validation dataset produced zero usable samples."); + + return total_loss / + static_cast(used); +} + +int main() { + constexpr int MAX_STEPS = 100; + constexpr int64_t MAX_SEQ = 32; + constexpr int64_t VOCAB = 50257; + + constexpr long long START_MIN_KB = + 1800LL * 1024LL; + + constexpr long long STEP_MIN_KB = + 1000LL * 1024LL; + + constexpr long long HARD_STOP_KB = + 600LL * 1024LL; + + const std::string model_dir = + "/root/gpt2-tamil-124m"; + + const std::string train_path = + "/root/classical-tamil-ppe-smoke/train.jsonl"; + + const std::string valid_path = + "/root/classical-tamil-ppe-smoke/valid.jsonl"; + + const std::string checkpoint_path = + "/root/tamil_gpt2_bf16_ultimate.ckpt"; + + std::atomic stop_watchdog{false}; + + std::thread watchdog([&]() { + while (!stop_watchdog.load()) { + const long long kb = + mem_available_kb(); + + if (kb > 0 && + kb < HARD_STOP_KB) { + + std::cerr + << "\n[SAFE ABORT] " + << "MemAvailable < 600 MiB\n"; + + std::_Exit(99); + } + + std::this_thread::sleep_for( + std::chrono::milliseconds(250)); + } + }); + + auto finish = [&](int rc) { + stop_watchdog = true; + watchdog.join(); + return rc; + }; + + try { + std::cout + << "====================================================\n" + << " ULTIMATE NATIVE 124M BF16 FINE-TUNE TEST\n" + << "====================================================\n"; + + print_mem("startup"); + + if (mem_available_kb() < START_MIN_KB) { + std::cout + << "[SAFE STOP] " + << "Startup memory below 1.8 GiB.\n"; + + return finish(0); + } + + // ======================================================== + // DATA + // ======================================================== + + auto train_records = + load_classic_records(train_path); + + auto valid_records = + load_classic_records(valid_path); + + std::cout + << "Training records: " + << train_records.size() + << "\n"; + + std::cout + << "Validation records: " + << valid_records.size() + << "\n"; + + if (train_records.empty() || + valid_records.empty()) { + + throw std::runtime_error( + "Training or validation dataset is empty."); + } + + // ======================================================== + // TOKENIZER + // ======================================================== + + TokenizerLoadOptions tok_options; + tok_options.model_type = "gpt2"; + + auto tokenizer = + TokenizerFactory::from_pretrained( + model_dir, + tok_options); + + if (!tokenizer) + throw std::runtime_error( + "Native GPT-2 tokenizer failed."); + + std::cout + << "[PASS] Native GPT-2 tokenizer loaded.\n"; + + // ======================================================== + // MODEL + // ======================================================== + + GPT2Config config = + GPT2Config::from_pretrained( + model_dir); + + GPT2Model model(config); + + if (config.tie_word_embeddings) + model.tie_weights(); + + // ======================================================== + // BF16 LOAD + // ======================================================== + + SafeTensorsModelReader reader(model_dir); + + reader.parse_headers(); + + auto mapping = + GPT2KeyMapper::generate_gpt2_mapping( + config.n_layer); + + SafeTensorsLoadOptions options; + + options.transpose_linear = false; + options.auto_promote_fp16 = true; + options.convert_f32_to_bf16 = true; + options.verbose = false; + + auto tensors = + reader.load_tensors_mapped( + mapping, + options); + + size_t parameter_count = 0; + size_t bf16_count = 0; + size_t fp32_count = 0; + + for (const auto& kv : tensors) { + if (!kv.second) + continue; + + parameter_count += + static_cast( + kv.second->numel()); + + if (kv.second->dtype() == kBFloat16) + ++bf16_count; + + else if (kv.second->dtype() == kFloat32) + ++fp32_count; + + model.assign_weight( + kv.first, + kv.second); + } + + if (parameter_count != 124439808 || + bf16_count != 148 || + fp32_count != 0) { + + std::cerr + << "[FAIL] Invalid 124M BF16 model.\n"; + + return finish(2); + } + + std::cout + << "[PASS] " + << parameter_count + << " parameters / " + << bf16_count + << " BF16 tensors / " + << fp32_count + << " FP32 parameter tensors.\n"; + + print_mem("after load"); + + // ======================================================== + // FULL FINE-TUNING + // ======================================================== + + for (auto& p : model.parameters()) { + if (p) + p->set_requires_grad(true); + } + + // ======================================================== + // BASELINE VALIDATION + // ======================================================== + + std::cout + << "\n===== BASELINE VALIDATION =====\n"; + + const double baseline_loss = + evaluate_dataset( + model, + *tokenizer, + valid_records, + MAX_SEQ); + + std::cout + << "Baseline validation loss: " + << baseline_loss + << "\n"; + + // ======================================================== + // ADAM + // ======================================================== + + AdamConfig adam_cfg; + + adam_cfg.learning_rate = 5e-6f; + adam_cfg.beta1 = 0.9f; + adam_cfg.beta2 = 0.999f; + adam_cfg.epsilon = 1e-8f; + adam_cfg.weight_decay = 0.0f; + adam_cfg.amsgrad = false; + + Adam adam(adam_cfg); + + // ======================================================== + // DETERMINISTIC SHUFFLING + // ======================================================== + + std::vector order( + train_records.size()); + + for (size_t i = 0; i < order.size(); ++i) + order[i] = i; + + std::mt19937 rng(123456789); + + // ======================================================== + // TRAIN + // ======================================================== + + int completed = 0; + double first_train_loss = 0.0; + double last_train_loss = 0.0; + + for (int step = 0; + step < MAX_STEPS; + ++step) { + + const int display_step = + step + 1; + + if (step % static_cast( + train_records.size()) == 0) { + + std::shuffle( + order.begin(), + order.end(), + rng); + } + + print_mem("before step"); + + if (mem_available_kb() < STEP_MIN_KB) { + + std::cout + << "[SAFE STOP] " + << "Memory below 1 GiB.\n"; + + break; + } + + const size_t record_index = + order[ + static_cast(step) + % order.size() + ]; + + const std::string& text = + train_records[record_index]; + + const auto ids = + tokenizer->encode(text); + + if (ids.size() < 2) { + std::cout + << "[SKIP] Too few tokens.\n"; + continue; + } + + const size_t usable = + std::min( + ids.size(), + static_cast( + MAX_SEQ + 1)); + + if (usable < 2) { + std::cout + << "[SKIP] Too few usable tokens.\n"; + continue; + } + + const int64_t S = + static_cast( + usable - 1); + + std::vector input_values( + static_cast(S)); + + std::vector label_values( + static_cast(S)); + + for (int64_t i = 0; + i < S; + ++i) { + + input_values[ + static_cast(i)] + = static_cast( + ids[static_cast(i)]); + + label_values[ + static_cast(i)] + = static_cast( + ids[static_cast(i + 1)]); + } + + auto input_ids = + std::make_shared( + std::vector{1, S}, + input_values.data(), + kInt64, + kCPU); + + auto labels = + std::make_shared( + std::vector{1, S}, + label_values.data(), + kInt32, + kCPU); + + // ==================================================== + // FORWARD + // ==================================================== + + auto logits = + model.forward(input_ids); + + if (!logits || + logits->dtype() != kFloat32 || + logits->shape().size() != 3 || + logits->shape()[0] != 1 || + logits->shape()[1] != S || + logits->shape()[2] != VOCAB || + !finite_f32(logits)) { + + std::cerr + << "[FAIL] Invalid logits at step " + << display_step + << "\n"; + + return finish(4); + } + + // ==================================================== + // REAL CAUSAL LM OBJECTIVE + // ==================================================== + + auto loss = + lm_cross_entropy( + logits, + labels, + -100, + "mean"); + + if (!loss || + !std::isfinite( + loss->item())) { + + std::cerr + << "[FAIL] Invalid causal-LM loss at step " + << display_step + << "\n"; + + return finish(5); + } + + const double loss_value = + static_cast( + loss->item()); + + if (completed == 0) + first_train_loss = loss_value; + + last_train_loss = loss_value; + + // ==================================================== + // BACKWARD + // ==================================================== + + loss->backward(); + + std::vector params; + std::vector grads; + + size_t invalid_grads = 0; + + for (const auto& p : + model.parameters()) { + + if (!p) + continue; + + auto g = p->grad(); + + if (!g) + continue; + + if (g->dtype() != kFloat32 || + !finite_f32(g)) { + + ++invalid_grads; + continue; + } + + params.push_back(p); + grads.push_back(g); + } + + if (params.empty() || + invalid_grads != 0) { + + std::cerr + << "[FAIL] Invalid gradient set at step " + << display_step + << "\n"; + + return finish(6); + } + + // ==================================================== + // ADAM + // ==================================================== + + if (mem_available_kb() < STEP_MIN_KB) { + + std::cout + << "[SAFE STOP] " + << "Memory below 1 GiB before Adam.\n"; + + break; + } + + adam.step( + params, + grads); + + // ==================================================== + // BF16 PARAMETER INTEGRITY + // ==================================================== + + size_t bf16_params = 0; + size_t invalid_params = 0; + + for (const auto& p : params) { + + if (!p || + p->dtype() != kBFloat16) { + + ++invalid_params; + continue; + } + + ++bf16_params; + + if (!finite_bf16(p)) + ++invalid_params; + } + + if (bf16_params != params.size() || + invalid_params != 0) { + + std::cerr + << "[FAIL] BF16 parameter corruption at step " + << display_step + << "\n"; + + return finish(7); + } + + // ==================================================== + // CLEAR GRAPH + // ==================================================== + + for (auto& p : + model.parameters()) { + + if (p) + p->zero_grad(); + } + + loss.reset(); + logits.reset(); + input_ids.reset(); + labels.reset(); + + ++completed; + + std::cout + << "Step " + << display_step + << "/" + << MAX_STEPS + << " | loss=" + << loss_value + << " | tokens=" + << S + << " | BF16 params=" + << bf16_params + << "\n"; + + print_mem("after step"); + + // Validation checkpoints. + if (completed == 50) { + + std::cout + << "\n===== MID-RUN VALIDATION =====\n"; + + const double val_loss = + evaluate_dataset( + model, + *tokenizer, + valid_records, + MAX_SEQ); + + std::cout + << "Validation loss @ step 50: " + << val_loss + << "\n"; + } + } + + if (completed == 0) { + + std::cerr + << "[FAIL] Zero training steps completed.\n"; + + return finish(8); + } + + // ======================================================== + // FINAL VALIDATION + // ======================================================== + + std::cout + << "\n===== FINAL VALIDATION =====\n"; + + const double final_val_loss = + evaluate_dataset( + model, + *tokenizer, + valid_records, + MAX_SEQ); + + std::cout + << "Baseline validation loss: " + << baseline_loss + << "\n"; + + std::cout + << "Final validation loss: " + << final_val_loss + << "\n"; + + // ======================================================== + // FINAL MODEL INTEGRITY + // ======================================================== + + size_t final_bf16 = 0; + size_t invalid_final = 0; + uint64_t aggregate_hash = 0; + + for (const auto& p : + model.parameters()) { + + if (!p) + continue; + + if (p->dtype() != kBFloat16) { + ++invalid_final; + continue; + } + + ++final_bf16; + + if (!finite_bf16(p)) { + ++invalid_final; + continue; + } + + aggregate_hash ^= + bf16_hash(p) + + 0x9e3779b97f4a7c15ULL + + (aggregate_hash << 6) + + (aggregate_hash >> 2); + } + + std::cout + << "\nFinal BF16 parameter tensors: " + << final_bf16 + << "\n"; + + std::cout + << "Invalid final tensors: " + << invalid_final + << "\n"; + + std::cout + << "Parameter hash: 0x" + << std::hex + << aggregate_hash + << std::dec + << "\n"; + + if (final_bf16 != 148 || + invalid_final != 0) { + + std::cerr + << "[FAIL] Final BF16 integrity check failed.\n"; + + return finish(9); + } + + // ======================================================== + // NATIVE BF16 CHECKPOINT + // ======================================================== + // + // Custom compact checkpoint: + // + // magic + // number of tensors + // for every model parameter: + // numel + // raw BF16 values + // + // It is intentionally a runtime checkpoint rather than a + // SafeTensors file. It proves that the trained BF16 state + // can be serialized without converting back to FP32. + // + // ======================================================== + + std::cout + << "\n===== BF16 CHECKPOINT WRITE =====\n"; + + { + std::ofstream out( + checkpoint_path, + std::ios::binary); + + if (!out) + throw std::runtime_error( + "Cannot create BF16 checkpoint."); + + const uint64_t magic = + 0x42463136544E3031ULL; + + const uint64_t count = + static_cast( + model.parameters().size()); + + out.write( + reinterpret_cast(&magic), + sizeof(magic)); + + out.write( + reinterpret_cast(&count), + sizeof(count)); + + for (const auto& p : + model.parameters()) { + + if (!p || + p->dtype() != kBFloat16) { + + throw std::runtime_error( + "Non-BF16 parameter during checkpoint."); + } + + const uint64_t n = + static_cast( + p->numel()); + + out.write( + reinterpret_cast(&n), + sizeof(n)); + + out.write( + reinterpret_cast( + p->data()), + static_cast( + n * sizeof(uint16_t))); + } + + out.flush(); + } + + std::ifstream checkpoint_in( + checkpoint_path, + std::ios::binary | std::ios::ate); + + if (!checkpoint_in) + throw std::runtime_error( + "Cannot reopen checkpoint."); + + const auto checkpoint_size = + checkpoint_in.tellg(); + + std::cout + << "Checkpoint size: " + << checkpoint_size + << " bytes\n"; + + if (checkpoint_size <= 0) + throw std::runtime_error( + "Checkpoint is empty."); + + std::cout + << "[PASS] BF16 checkpoint written.\n"; + + // ======================================================== + // END + // ======================================================== + + print_mem("final"); + + std::cout + << "\n====================================================\n" + << " ULTIMATE TEST RESULT\n" + << "====================================================\n"; + + std::cout + << "[PASS] Real Classical-Tamil training data.\n" + << "[PASS] Native GPT-2 tokenizer.\n" + << "[PASS] Native causal next-token objective.\n" + << "[PASS] Native lm_cross_entropy.\n" + << "[PASS] Native forward/backward.\n" + << "[PASS] FP32 gradients.\n" + << "[PASS] FP32 Adam state/update.\n" + << "[PASS] BF16 parameter storage throughout training.\n" + << "[PASS] Gradient reset after every step.\n" + << "[PASS] Memory watchdog remained active.\n" + << "[PASS] Final BF16 parameter integrity.\n" + << "[PASS] Native BF16 checkpoint serialization.\n"; + + std::cout + << "\nCompleted steps: " + << completed + << "\n"; + + std::cout + << "First training loss: " + << first_train_loss + << "\n"; + + std::cout + << "Last training loss: " + << last_train_loss + << "\n"; + + std::cout + << "Validation baseline: " + << baseline_loss + << "\n"; + + std::cout + << "Validation final: " + << final_val_loss + << "\n"; + + std::cout + << "====================================================\n"; + + return finish(0); + + } catch (const std::exception& e) { + + std::cerr + << "\n[FATAL] " + << e.what() + << "\n"; + + return finish(10); + } +}