Update

Author: Cydral

dlib/cuda/cpu_dlib.h

@@ -786,7 +786,8 @@ namespace dlib
             const tensor& input_tensor,
             const tensor& output_tensor,
             tensor& grad,
-            double& loss
+            double& loss,
+            long ignore_index
         ) const
         {
             DLIB_CASSERT(output_tensor.k() == 1);
@@ -796,16 +797,46 @@ namespace dlib
             const long batch_size = output_tensor.num_samples();
             const long seq_len = output_tensor.nr();
             const long vocab_size = output_tensor.nc();
-        
-            // Normalization over all positions
-            const double scale = 1.0 / (batch_size * seq_len);
-        
-            loss = 0.0;
+                
             const float* out_data = output_tensor.host();
             const float* in_data = input_tensor.host();
-            float* g = grad.host();
+            float* g = grad.host();            
 
             std::fill(g, g + grad.size(), 0.0f);
+
+            long valid_tokens = 0;
+
+            if (ignore_index < 0)
+            {
+                valid_tokens = batch_size * seq_len;
+            }
+            else {
+                for (long i = 0; i < batch_size; ++i)
+                {
+                    for (long t = 0; t < seq_len; ++t)
+                    {
+                        unsigned long target_class;
+                        if (t < seq_len - 1) {
+                            target_class = static_cast<unsigned long>(
+                                in_data[tensor_index(input_tensor, i, 0, t + 1, 0)]
+                                );
+                        }
+                        else
+                            target_class = *(truth + i);
+
+                        if (static_cast<long>(target_class) != ignore_index)
+                            valid_tokens++;
+                    }
+                }
+            }
+            if (valid_tokens == 0)
+            {
+                loss = 0.0;
+                return;
+            }
+
+            const double scale = 1.0 / valid_tokens;
+            loss = 0.0;
 
             for (long i = 0; i < batch_size; ++i)
             {
@@ -825,6 +856,9 @@ namespace dlib
                         target_class = *(truth + i);
                     }
 
+                    if (ignore_index >= 0 && static_cast<long>(target_class) == ignore_index)
+                        continue;
+
                     DLIB_CASSERT(target_class < static_cast<unsigned long>(vocab_size));
 
                     // Find max logit for numerical stability

dlib/cuda/cuda_dlib.cu

@@ -3139,6 +3139,39 @@ namespace dlib
 
     // ----------------------------------------------------------------------------------------
 
+        __global__ void _cuda_count_valid_tokens(
+            float* valid_count,
+            const unsigned long* truth,
+            const float* input_data,
+            size_t batch_size,
+            size_t seq_len,
+            long ignore_index
+        )
+        {
+            float count = 0.0f;
+
+            for (auto sample_idx : grid_stride_range(0, batch_size))
+            {
+                for (size_t t = 0; t < seq_len; ++t)
+                {
+                    unsigned long target_class;
+                    if (t < seq_len - 1) {
+                        const size_t input_idx = sample_idx * seq_len + (t + 1);
+                        target_class = static_cast<unsigned long>(input_data[input_idx]);
+                    }
+                    else {
+                        target_class = truth[sample_idx];
+                    }
+
+                    if (ignore_index < 0 || static_cast<long>(target_class) != ignore_index) {
+                        count += 1.0f;
+                    }
+                }
+            }
+
+            warp_reduce_atomic_add(*valid_count, count);
+        }
+
         __global__ void _cuda_compute_loss_cross_entropy_per_logit(
             float* loss_out,
             float* g,
@@ -3148,7 +3181,8 @@ namespace dlib
             size_t batch_size,
             size_t seq_len,
             size_t vocab_size,
-            const float scale
+            float scale,
+            long ignore_index
         )
         {
             float total_loss = 0;
@@ -3158,7 +3192,6 @@ namespace dlib
                 for (size_t t = 0; t < seq_len; ++t)
                 {
                     unsigned long target_class;
-
                     if (t < seq_len - 1) {
                         const size_t input_idx = sample_idx * seq_len + (t + 1);
                         target_class = static_cast<unsigned long>(input_data[input_idx]);
@@ -3168,7 +3201,15 @@ namespace dlib
                     }
 
                     const size_t base_idx = sample_idx * seq_len * vocab_size + t * vocab_size;
-                    float max_val = out_data[base_idx + 0];
+
+                    if (ignore_index >= 0 && static_cast<long>(target_class) == ignore_index) {
+                        for (size_t c = 0; c < vocab_size; ++c) {
+                            g[base_idx + c] = 0.0f;
+                        }
+                        continue;
+                    }
+
+                    float max_val = out_data[base_idx];
                     for (size_t c = 1; c < vocab_size; ++c)
                     {
                         max_val = ::max(max_val, out_data[base_idx + c]);
@@ -3210,7 +3251,8 @@ namespace dlib
             const tensor& input_tensor,
             const tensor& subnetwork_output,
             tensor& gradient,
-            double& loss
+            double& loss,
+            long ignore_index
         )
         {
             CHECK_CUDA(cudaMemset(gradient.device(), 0, gradient.size() * sizeof(float)));
@@ -3220,7 +3262,35 @@ namespace dlib
             const long seq_len = subnetwork_output.nr();
             const long vocab_size = subnetwork_output.nc();
 
-            const double scale = 1.0 / (batch_size * seq_len);
+            double scale;
+            if (ignore_index < 0)
+            {
+                scale = 1.0 / (batch_size * seq_len);
+            }
+            else {
+                cuda_data_void_ptr count_buf = device_global_buffer(sizeof(float));
+                auto valid_count_ptr = static_pointer_cast<float>(count_buf, 1);
+                CHECK_CUDA(cudaMemset(valid_count_ptr, 0, sizeof(float)));
+
+                launch_kernel(_cuda_count_valid_tokens, max_jobs(batch_size),
+                    valid_count_ptr.data(),
+                    truth_buffer.data(),
+                    input_tensor.device(),
+                    batch_size,
+                    seq_len,
+                    ignore_index
+                );
+
+                float valid_count;
+                dlib::cuda::memcpy(&valid_count, valid_count_ptr);
+
+                if (valid_count == 0) {
+                    loss = 0.0;
+                    return;
+                }
+
+                scale = 1.0 / valid_count;
+            }
 
             launch_kernel(_cuda_compute_loss_cross_entropy_per_logit, max_jobs(batch_size),
                 loss_work_buffer.data(),
@@ -3231,7 +3301,8 @@ namespace dlib
                 batch_size,
                 seq_len,
                 vocab_size,
-                scale
+                static_cast<float>(scale),
+                ignore_index
             );
 
             float floss;

dlib/cuda/cuda_dlib.h

@@ -682,7 +682,8 @@ namespace dlib
                 const tensor& input_tensor,        // Source tokens
                 const tensor& subnetwork_output,   // Logits
                 tensor& gradient,
-                double& loss
+                double& loss,
+                long ignore_index
                 ) const
             {
                 const size_t bytes_per_sample = sizeof(unsigned long);
@@ -697,7 +698,7 @@ namespace dlib
                 }
 
                 auto truth_buf = static_pointer_cast<const unsigned long>(buf, subnetwork_output.num_samples());
-                do_work(loss_buf, truth_buf, input_tensor, subnetwork_output, gradient, loss);
+                do_work(loss_buf, truth_buf, input_tensor, subnetwork_output, gradient, loss, ignore_index);
             }
 
         private:
@@ -707,7 +708,8 @@ namespace dlib
                 const tensor& input_tensor,
                 const tensor& subnetwork_output,
                 tensor& gradient,
-                double& loss
+                double& loss,
+                long ignore_index
             );
 
             mutable cuda_data_void_ptr buf;

dlib/dnn/loss.h

@@ -918,6 +918,11 @@ namespace dlib
         typedef unsigned long training_label_type;
         typedef unsigned long output_label_type;
 
+        loss_cross_entropy_per_logit_() : ignore_index_(-1) {}
+
+        void set_ignore_index(long idx) { ignore_index_ = idx; }
+        long get_ignore_index() const { return ignore_index_; }
+
         template <typename SUB_TYPE, typename label_iterator>
         void to_label(
             const tensor& input_tensor,
@@ -977,38 +982,43 @@ namespace dlib
 
             double loss = 0.0;
 #ifdef DLIB_USE_CUDA
-            cuda_compute(truth, input_tensor, output_tensor, grad, loss);
+            cuda_compute(truth, input_tensor, output_tensor, grad, loss, ignore_index_);
 #else
-            cpu_compute(truth, input_tensor, output_tensor, grad, loss);
+            cpu_compute(truth, input_tensor, output_tensor, grad, loss, ignore_index_);
 #endif
             return loss;
         }
 
-        friend void serialize(const loss_cross_entropy_per_logit_&, std::ostream& out)
+        friend void serialize(const loss_cross_entropy_per_logit_& item, std::ostream& out)
         {
             serialize("loss_cross_entropy_per_logit_", out);
+            serialize(item.ignore_index_, out);
         }
 
-        friend void deserialize(loss_cross_entropy_per_logit_&, std::istream& in)
+        friend void deserialize(loss_cross_entropy_per_logit_& item, std::istream& in)
         {
             std::string version;
             deserialize(version, in);
             if (version != "loss_cross_entropy_per_logit_")
                 throw serialization_error("Unexpected version found while deserializing dlib::loss_cross_entropy_per_logit_.");
+            deserialize(item.ignore_index_, in);
         }
 
-        friend std::ostream& operator<<(std::ostream& out, const loss_cross_entropy_per_logit_&)
+        friend std::ostream& operator<<(std::ostream& out, const loss_cross_entropy_per_logit_& item)
         {
             out << "loss_cross_entropy_per_logit";
+            out << " (ignore_index=" << item.ignore_index_ << ")";
             return out;
         }
 
-        friend void to_xml(const loss_cross_entropy_per_logit_& /*item*/, std::ostream& out)
+        friend void to_xml(const loss_cross_entropy_per_logit_& item, std::ostream& out)
         {
-            out << "<loss_cross_entropy_per_logit/>\n";
+            out << "<loss_cross_entropy_per_logit ignore_index='" << item.ignore_index_ << "'/>\n";
         }
 
         private:
+            long ignore_index_;
+
 #ifdef DLIB_USE_CUDA
             cuda::compute_loss_cross_entropy_per_logit cuda_compute;
 #else

examples/slm_advanced_train_ex.cpp

@@ -449,6 +449,7 @@ int main(int argc, char** argv)
             // Build and train the network
             using net_type = my_transformer::network_type<true>;
             net_type net;
+            layer<0>(net).loss_details().set_ignore_index(tokenizer.get_special_token_id("<pad>"));
             cout << my_transformer::model_info::describe() << endl;
 
             // Tokenizer stored with model for simplified inference

examples/slm_chatbot_ex.cpp

@@ -267,7 +267,7 @@ int main(int argc, char** argv)
             trainer.set_mini_batch_size(batch_size);
             trainer.set_max_num_epochs(max_epochs);
             trainer.set_iterations_without_progress_threshold(patience);
-            trainer.set_synchronization_file("chkpt-" + finetuned_model, std::chrono::minutes(5));
+            trainer.set_synchronization_file("chkpt-" + finetuned_model, std::chrono::minutes(25));
             trainer.be_quiet();
 
             // Load tokenizer & model
@@ -282,6 +282,7 @@ int main(int argc, char** argv)
                 cout << "Pre-trained tokenizer not found at: " << tokenizer_file << endl;
                 return 1;
             }
+            layer<0>(net).loss_details().set_ignore_index(tokenizer.get_special_token_id("<pad>"));
 
             // Load Q&A datasets for fine-tuning
             cout << "Loading Q&A training datasets...\n";
@@ -348,10 +349,10 @@ int main(int argc, char** argv)
             qa_tokens.clear();
 
             cout << "Applying freezing strategy...\n";
-            set_all_learning_rate_multipliers(net, 0.1);
-            layer<1>(net).layer_details().set_learning_rate_multiplier(1.0);    // linear
-            layer<2>(net).layer_details().set_learning_rate_multiplier(0.5);    // rms_norm
-            layer<111>(net).layer_details().set_learning_rate_multiplier(0.5);  // embeddings
+            //set_all_learning_rate_multipliers(net, 0.1);
+            //layer<1>(net).layer_details().set_learning_rate_multiplier(1.0);    // linear
+            //layer<2>(net).layer_details().set_learning_rate_multiplier(0.5);    // rms_norm
+            //layer<111>(net).layer_details().set_learning_rate_multiplier(0.5);  // embeddings
             cout << "Fine-tuning learning rate strategy applied.\n\n";
             cout << net << endl;
 

examples/slm_mixture_of_experts_ex.cpp

@@ -792,6 +792,7 @@ int main(int argc, char** argv)
             // Build and train the network
             using net_type = my_transformer::network_type<true>;
             net_type net;
+            layer<0>(net).loss_details().set_ignore_index(tokenizer.get_special_token_id("<pad>"));
             cout << my_transformer::model_info::describe() << endl;
 
             // Tokenizer stored with model for simplified inference
@@ -805,7 +806,7 @@ int main(int argc, char** argv)
             trainer.set_learning_rate_shrink_factor(0.1);
             trainer.set_mini_batch_size(batch_size);
             trainer.set_iterations_without_progress_threshold(patience);
-            trainer.set_synchronization_file("chkpt-" + model_file, std::chrono::minutes(5));
+            trainer.set_synchronization_file("chkpt-" + model_file, std::chrono::minutes(25));
             trainer.be_quiet();
             cout << net << endl << endl; // Show the model architecture
             cout << "Starting training...\n";