fix stack logging for python bindings

Author: ngc92

scripts/train.py

@@ -177,7 +177,7 @@ def main():
 
         # Log allocator stats
         for idx in range(config.gpus):
-            logger.log_allocator(trainer.get_allocator_info(idx))
+            logger.log_allocator(trainer, idx)
 
         # calculate the expected time at peak flops for speed-of-light estimation
         logger.set_expected_time_per_token(trainer)

src/binding/binding.cpp

@@ -338,6 +338,13 @@ NB_MODULE(_pyllmq, m) {
                 res["pageable"] = size.PageableHost;
                 ret[nb::cast(name)] = res;
             }
+
+            auto stack = trainer->get_stack(gpu_id);
+            for (const auto& [name, size] : stack) {
+                nb::dict res;
+                res["stack"] = size;
+                ret[nb::cast(name)] = res;
+            }
             return ret;
             }, nb::arg("gpu_id") = 0, "Get the current memory allocations for the given GPU")
         ;
@@ -441,17 +448,22 @@ NB_MODULE(_pyllmq, m) {
              "Log GPU utilization state")
         .def("log_allocator", [](TrainingRunLogger* logger, const nb::dict& stats) {
             std::vector<std::pair<std::string, sSegmentMemory>> cpp_stats;
+            std::vector<std::pair<std::string, long>> cpp_stack;
             cpp_stats.reserve(stats.size());
             for (auto item : stats) {
                 std::string key = nb::cast<std::string>(item.first);
                 nb::dict value = nb::cast<nb::dict>(item.second);
-                long device = nb::cast<long>(value["device"]);
-                long managed = nb::cast<long>(value["managed"]);
-                long pinned = nb::cast<long>(value["pinned"]);
-                long pageable = nb::cast<long>(value["pageable"]);
-                cpp_stats.emplace_back(key, sSegmentMemory{device, managed, pinned, pageable});
+                if (value.contains("stack")) {
+                    cpp_stack.emplace_back(key, nb::cast<long>(value["stack"]));
+                } else {
+                    long device = nb::cast<long>(value["device"]);
+                    long managed = nb::cast<long>(value["managed"]);
+                    long pinned = nb::cast<long>(value["pinned"]);
+                    long pageable = nb::cast<long>(value["pageable"]);
+                    cpp_stats.emplace_back(key, sSegmentMemory{device, managed, pinned, pageable});
+                }
             }
-            logger->log_allocator(cpp_stats);
+            logger->log_allocator(cpp_stats, cpp_stack);
         }, nb::arg("stats"), "Log memory allocator statistics")
          .def("set_expected_time_per_token", [](TrainingRunLogger* logger, const MultiGPUPyTrainer* trainer){
              auto& config = trainer->config();

src/binding/py_train.cpp

@@ -14,6 +14,7 @@
 #include "utilities/comm.h"
 #include "kernels/kernels.h"
 #include "models/llama_gradients.h"
+#include "models/llama_run_state.h"
 
 MultiGPUPyTrainer::MultiGPUPyTrainer(int ngpus, LLamaConfig config, LLamaOptions options, int batch_size, int seq_len, int grad_accum, bool memcpy_all_gather, bool memcpy_send_recv) :
     mConfig(config), mOptions(options), B(batch_size), T(seq_len), mGradAccumulation(grad_accum)
@@ -249,6 +250,14 @@ std::vector<std::pair<std::string, sSegmentMemory>> MultiGPUPyTrainer::get_alloc
     return result;
 }
 
+std::vector<std::pair<std::string, long>> MultiGPUPyTrainer::get_stack_info(int gpu_id) {
+    std::vector<std::pair<std::string, long>> result;
+    run_work([&result](sThreadContext& ctx) {
+        result = ctx.Model->run_state().Stack.get_allocation_stats();
+    }, gpu_id);
+    return result;
+}
+
 std::vector<std::pair<std::string, Tensor>> MultiGPUPyTrainer::get_gradients(int gpu_id) {
     std::vector<std::pair<std::string, Tensor>> result;
     run_work([&result](sThreadContext& ctx) {

src/binding/py_train.h

@@ -63,6 +63,7 @@ class MultiGPUPyTrainer
     const LLamaOptions& options() const { return mOptions; }
 
     std::vector<std::pair<std::string, sSegmentMemory>> get_allocations(int gpu_id);
+    std::vector<std::pair<std::string, long>> get_stack_info(int gpu_id);
     std::vector<std::pair<std::string, Tensor>> get_gradients(int gpu_id);
 
 private:

src/training/logging.cpp

@@ -321,7 +321,10 @@ void TrainingRunLogger::log_line(std::string_view line) {
     mFirst = false;
 }
 
-void TrainingRunLogger::log_allocator(const std::vector<std::pair<std::string, sSegmentMemory>>& stats, const DeviceMemoryStack& stack) {
+void TrainingRunLogger::log_allocator(
+        const std::vector<std::pair<std::string, sSegmentMemory>>& stats,
+        const std::vector<std::pair<std::string, long>>& stack_info)
+{
     if (mRank != 0) return;
     std::string stat_str = "[";
     bool first = true;
@@ -342,7 +345,7 @@ void TrainingRunLogger::log_allocator(const std::vector<std::pair<std::string, s
             printf("  %16s: %6zu | %7zu | %6zu \n", name.c_str(), amount.OnDevice / 1024 / 1024, amount.Managed / 1024 / 1024, amount.PinnedHost / 1024 / 1024);
         }
         printf("\n");
-        for (auto& [ptr, amount, name]: stack.get_high_mark()) {
+        for (const auto& [name, amount]: stack_info) {
             std::string stack_name = fmt::format("stack.{}", name);
             int mib = static_cast<int>(amount / 1024 / 1024);
             if(mib > 0) {

src/training/logging.h

@@ -43,7 +43,10 @@ class TrainingRunLogger
     void log_step(int step, float epoch, int step_tokens, int duration_ms, float norm, float loss, float lr);
     void log_eval(int step, float epoch, int eval_tokens, int duration_ms, float loss);
     void log_gpu_state(int step, int gpu_id, const GPUUtilInfo& gpu_util);
-    void log_allocator(const std::vector<std::pair<std::string, sSegmentMemory>>& stats, const DeviceMemoryStack& stack);
+    void log_allocator(
+        const std::vector<std::pair<std::string, sSegmentMemory>>& stats,
+        const std::vector<std::pair<std::string, long>>& stack_info
+        );
 
     // call at the beginning and end of a section of processing.
     // will record the time between the two calls

src/utilities/stack.cpp

@@ -40,6 +40,14 @@ void DeviceMemoryStack::free(std::byte* ptr) {
     mAlloc.pop_back();
 }
 
+std::vector<std::pair<std::string, long>> DeviceMemoryStack::get_allocation_stats() const {
+    std::vector<std::pair<std::string, long>> result;
+    for (auto& [ptr, amount, name]: get_high_mark()) {
+        result.emplace_back(name, amount);
+    }
+    return result;
+}
+
 void DeviceMemoryStack::_track_max() {
     if(bytes_used() > mMaxUtilization) {
         mMaxUtilization = bytes_used();

src/utilities/stack.h

@@ -35,6 +35,8 @@ class DeviceMemoryStack {
     const AllocationList& get_high_mark() const { return mHighMark; }
     void set_high_mark(const AllocationList& list) { mHighMark = list; }
 
+    std::vector<std::pair<std::string, long>> get_allocation_stats() const;
+
 private:
     int mDeviceID;
     std::byte* mBackingMemory;

train.cpp

@@ -429,7 +429,7 @@ void TrainingRunner::run_training(int argc, const char** argv, NCCLCommunicator&
 
     logger.log_dataset(train_loader, test_loader);
 
-    logger.log_allocator(model.get_allocator().get_allocation_segments(), model.run_state().Stack);
+    logger.log_allocator(model.get_allocator().get_allocation_segments(), model.run_state().Stack.get_allocation_stats());
 
     Tensor inputs = model.get_input_buffer();
     Tensor targets = model.get_target_buffer();