[tt-train] Add kahan summation in AdamW (#15518)

### Problem description
bfloat16 has wide range, but is not as precise as float32. Some of the
parameters are not updated due to small magnitude of the gradients
(after multiplication by learning rate), as an example: gamma in
LayerNorm.

### What's changed
Add kahan summation flag to enable kahan summation when update weights. 

### Checklist
- [x] Post commit CI passes
https://github.com/tenstorrent/tt-metal/actions/runs/12128966379
- [x] New/Existing tests provide coverage for changes

<img width="1620" alt="image"
src="https://github.com/user-attachments/assets/b19ecc48-0a74-49b5-b55a-60198eb8ce9d">

tt-train/sources/examples/nano_gpt/main.cpp

@@ -141,6 +141,8 @@ struct TrainingConfig {
     uint32_t max_steps = 5000;
     float learning_rate = 3e-4F;
     float weight_decay = 1e-2F;
+    // works only for AdamW
+    bool use_kahan_summation = false;
     std::string model_path;
     std::string data_path;
     ttml::models::gpt2::TransformerConfig transformer_config;
@@ -157,6 +159,7 @@ TrainingConfig parse_config(const YAML::Node &yaml_config) {
     config.max_steps = training_config["max_steps"].as<uint32_t>();
     config.learning_rate = training_config["learning_rate"].as<float>();
     config.weight_decay = training_config["weight_decay"].as<float>();
+    config.use_kahan_summation = training_config["use_kahan_summation"].as<bool>(config.use_kahan_summation);
     config.model_path = training_config["model_path"].as<std::string>("");
     config.data_path = training_config["data_path"].as<std::string>(std::string(DATA_FOLDER) + "/shakespeare.txt");
     config.transformer_config = ttml::models::gpt2::read_config(training_config["transformer_config"]);
@@ -295,9 +298,11 @@ int main(int argc, char **argv) {
     auto adamw_params = ttml::optimizers::AdamWConfig();
     adamw_params.lr = config.learning_rate;
     adamw_params.weight_decay = config.weight_decay;
+    adamw_params.use_kahan_summation = config.use_kahan_summation;
     fmt::print("AdamW configuration:\n");
     fmt::print("    Learning rate: {}\n", adamw_params.lr);
     fmt::print("    Weight decay: {}\n", adamw_params.weight_decay);
+    fmt::print("    Use Kahan summation: {}\n", adamw_params.use_kahan_summation);
     auto optimizer = ttml::optimizers::AdamW(model->parameters(), adamw_params);
 
     if (!config.model_path.empty() && std::filesystem::exists(config.model_path)) {

tt-train/sources/ttml/optimizers/adamw.cpp

@@ -16,13 +16,18 @@ namespace {
 
 const std::string kFirstMoment = "first_moment/";
 const std::string kSecondMoment = "second_moment/";
+const std::string kKahanCompensation = "kahan_compensation/";
 
 }  // namespace
 
 namespace ttml::optimizers {
 
 MorehAdamW::MorehAdamW(autograd::NamedParameters parameters, const AdamWConfig& config) :
     OptimizerBase(std::move(parameters)), m_config(config) {
+    if (m_config.use_kahan_summation) {
+        throw std::runtime_error("MorehAdamW: Kahan summation is not supported. Use default AdamW instead.");
+    }
+
     for (const auto& [key, tensor_ptr] : m_parameters) {
         if (tensor_ptr->get_requires_grad()) {
             m_first_moment.emplace(
@@ -137,6 +142,13 @@ AdamW::AdamW(autograd::NamedParameters parameters, const AdamWConfig& config) :
                 autograd::create_tensor(
                     core::zeros_like(tensor_ptr->get_value(autograd::PreferredPrecision::FULL)),
                     /* requires_grad */ false));
+            if (m_config.use_kahan_summation) {
+                m_kahan_compensation.emplace(
+                    key,
+                    autograd::create_tensor(
+                        core::zeros_like(tensor_ptr->get_value(autograd::PreferredPrecision::FULL)),
+                        /* requires_grad */ false));
+            }
         }
     }
 }
@@ -188,11 +200,29 @@ void AdamW::step() {
         // weights -= lr * first_moment_hat / (sqrt(second_moment_hat) + epsilon)
         first_moment_ptr->set_value(first_moment);
         second_moment_ptr->set_value(second_moment);
-        tensor_ptr->set_value(ttnn::subtract(
-            tensor_ptr->get_value(autograd::PreferredPrecision::FULL),
-            ttnn_fixed::divide(
-                ttnn::multiply(first_moment_hat, m_config.lr),
-                ttnn::add(ttnn::sqrt(second_moment_hat), m_config.epsilon))));
+
+        auto update_tensor = ttnn_fixed::divide(
+            ttnn::multiply(first_moment_hat, -m_config.lr), ttnn::add(ttnn::sqrt(second_moment_hat), m_config.epsilon));
+
+        if (!m_config.use_kahan_summation) {
+            tensor_ptr->set_value(ttnn::add(tensor_ptr->get_value(autograd::PreferredPrecision::FULL), update_tensor));
+        } else {
+            auto value_tensor = tensor_ptr->get_value(autograd::PreferredPrecision::FULL);
+
+            const auto& kahan_compensation_ptr = m_kahan_compensation.at(key);
+            // A running compensation for lost low-order bits
+            auto compensation_tensor = kahan_compensation_ptr->get_value(autograd::PreferredPrecision::FULL);
+            // Adjust the update with the compensation
+            auto adjusted_update = ttnn::subtract(update_tensor, compensation_tensor);
+            // Update the value with the adjusted update
+            auto result = ttnn::add(value_tensor, adjusted_update);
+            // (result - value_tensor) cancels the high-order part of adjusted_update;
+            // subtracting adjusted_update recovers negative (low part of adjusted_update)
+            compensation_tensor = ttnn::subtract(ttnn::subtract(result, value_tensor), adjusted_update);
+
+            tensor_ptr->set_value(result);
+            kahan_compensation_ptr->set_value(compensation_tensor);
+        }
     }
 }
 
@@ -206,6 +236,10 @@ void AdamW::step() {
         state_dict.emplace(kSecondMoment + key, second_moment);
     }
 
+    for (const auto& [key, kahan_compensation] : m_kahan_compensation) {
+        state_dict.emplace(kKahanCompensation + key, kahan_compensation);
+    }
+
     return state_dict;
 }
 
@@ -215,6 +249,8 @@ void AdamW::set_state_dict(const autograd::NamedParameters& dict) {
             m_first_moment[key.substr(kFirstMoment.size())] = tensor;
         } else if (key.starts_with(kSecondMoment)) {
             m_second_moment[key.substr(kSecondMoment.size())] = tensor;
+        } else if (key.starts_with(kKahanCompensation)) {
+            m_kahan_compensation[key.substr(kKahanCompensation.size())] = tensor;
         } else {
             throw std::runtime_error(fmt::format("AdamW: Invalid key in state dict. Key = {}", key));
         }

tt-train/sources/ttml/optimizers/adamw.hpp

@@ -16,6 +16,9 @@ struct AdamWConfig {
     float epsilon{1e-8F};
     float weight_decay{0.01F};
     // TODO: add amsgrad
+
+    // flag to enable kahan summation to reduce floating point errors
+    bool use_kahan_summation{false};
 };
 
 class MorehAdamW : public OptimizerBase {
@@ -58,6 +61,7 @@ class AdamW : public OptimizerBase {
     AdamWConfig m_config;
     autograd::NamedParameters m_first_moment;
     autograd::NamedParameters m_second_moment;
+    autograd::NamedParameters m_kahan_compensation;
 };
 
 }  // namespace ttml::optimizers