Add multi-objective global memory search algorithm

.gitignore

@@ -170,3 +170,6 @@ docker/config
 node_modules/
 package.json
 yarn.lock
+
+# Logs
+logs/

CMakeLists.txt

@@ -14,6 +14,9 @@ set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} ${CMAKE_CURRENT_LIST_DIR}/cmake)
 set(FLEXFLOW_ROOT ${CMAKE_CURRENT_LIST_DIR})
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -UNDEBUG")
 
+# Export Clang Compilation Database
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+
 # Set a default build type if none was specified
 set(default_build_type "Debug")
 if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)

cmake/cuda.cmake

@@ -3,6 +3,7 @@ set(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
 set(CUDA_ROOT ${CUDA_PATH})
 set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_PATH})
 list(APPEND CMAKE_PREFIX_PATH ${CUDA_ROOT})
+message( STATUS "CMAKE_PREFIX_PATH" : "${CMAKE_PREFIX_PATH}" )
 find_package(CUDA REQUIRED)
 
 if(CUDA_FOUND)

include/flexflow/config.h

@@ -117,6 +117,7 @@ class FFConfig {
   int epochs, batchSize, printFreq;
   // int inputHeight, inputWidth;
   int numNodes, cpusPerNode, workersPerNode;
+  float device_mem; // The device (GPU) memory threshold; given by -ll:fsize
   float learningRate, weightDecay;
   size_t workSpaceSize;
   Legion::Context lg_ctx;
@@ -155,6 +156,7 @@ class FFConfig {
   int base_optimize_threshold;
   bool enable_control_replication;
   int python_data_loader_type;
+  bool perform_memory_search{false};
 };
 
 class FFIterationConfig {

include/flexflow/graph.h

@@ -17,6 +17,7 @@
 #define _FLEXFLOW_GRAPH_H_
 #include "flexflow/basic_graph.h"
 #include "flexflow/graph_structures.h"
+#include "flexflow/memory_optimization.h"
 #include "flexflow/model.h"
 #include "flexflow/utils/dot/dot_file.h"
 #include "flexflow/utils/recursive_logger.h"
@@ -109,6 +110,32 @@ struct GraphCostResult {
   friend std::ostream &operator<<(std::ostream &, GraphCostResult const &);
 };
 
+/**
+ * @brief Experimental. Hold the cost information of a PCG. To be merged with
+ * GraphCostResult.
+ */
+struct GraphCostResultWithMemory {
+  float cost;           ///< Run time cost
+  MemoryUsage mem_cost; ///< Memory usage
+  ///< Corresponding machine views (device placement views)
+  std::unordered_map<Node, MachineView> views;
+
+  /**
+   * @brief Get the multi-objective cost that combines the run time and memory
+   * cost.
+   *
+   * @return float Numerical value to represent the overall cost
+   */
+  float get_multi_obj_cost() const;
+
+  static GraphCostResultWithMemory invalid();
+
+  bool operator<(GraphCostResultWithMemory const &other) const;
+
+  friend std::ostream &operator<<(std::ostream &,
+                                  GraphCostResultWithMemory const &);
+};
+
 template <typename T>
 T sequence_cost(T const &first, T const &second);
 
@@ -195,6 +222,20 @@ class SearchHelper {
   template <typename T>
   void add_operator_cost(NodeAssignment const &, float, T *) const;
 
+  template <typename T>
+  void add_sink_node_costs(NodeAssignment const &sink,
+                           CostMetrics metrics,
+                           T *result) const;
+
+  /**
+   * @brief Add run time cost and memory cost of the operator to the graph cost.
+   * This is a temp workaround and should be refactored eventually.
+   */
+  void add_operator_cost_with_memory(NodeAssignment const &node,
+                                     float node_run_time_cost,
+                                     MemoryUsage node_mem_cost,
+                                     GraphCostResultWithMemory *cost) const;
+
   template <typename T>
   float get_cost(T const &) const;
 
@@ -204,6 +245,8 @@ class SearchHelper {
 public:
   mutable std::unique_ptr<RecursiveLogger> logger;
 
+  void clear_cache();
+
 private:
   template <typename T>
   T execute_nonsequence_split(std::unique_ptr<Graph> const &first_graph,
@@ -255,6 +298,8 @@ class Graph {
   Graph subgraph(std::unordered_set<Node> const &nodes) const;
   void contract_out_node(Node const &);
   float optimal_cost() const;
+  // Experimental. To be merged with optimal_cost().
+  float optimal_cost_with_memory(float const run_time_cost_factor) const;
   std::unordered_map<Node, MachineView> optimal_views() const;
   void remove_input_nodes();
   void duplicate_input_node(Node const &);
@@ -330,6 +375,21 @@ struct GraphOptimizeResult {
   friend std::ostream &operator<<(std::ostream &, GraphOptimizeResult const &);
 };
 
+/**
+ * @brief Experimental. Hold the optimization results with memory information.
+ * To be merged with GraphOptimizeResult.
+ */
+struct GraphOptimizeResultWithMemory {
+  tl::optional<Graph> graph; ///< Optimized PCG
+  float cost;                ///< Run time cost
+  MemoryUsage mem_cost;      ///< Memory usage
+  ///< Corresponding machine views (device placement views)
+  std::unordered_map<Node, MachineView> views;
+
+  friend std::ostream &operator<<(std::ostream &,
+                                  GraphOptimizeResultWithMemory const &);
+};
+
 namespace Utils {
 template <>
 struct GraphStructure<FlexFlow::PCG::Graph> {

include/flexflow/memory_optimization.h

@@ -0,0 +1,166 @@
+/**
+ * @file memory_optimization.h
+ * @brief Memory optimization related stuff
+ *
+ * @copyright Copyright 2022 CMU, Facebook, LANL, MIT, and Stanford
+ * (alphabetical)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cassert>
+#include <string>
+
+namespace FlexFlow {
+
+enum class MemoryUsageType {
+  // Use global memory of a PCG as the measure of memory usage. No device
+  // mapping consideration.
+  GLOBAL,
+
+  // Use the max of peak per-device memory usage among devices as the measure.
+  // Need associated device mapping views.
+  PER_DEVICE_MAX,
+
+  // Use detailed per-device memory usage as the measure. Need associated device
+  // mapping views.
+  PER_DEVICE_ALL,
+};
+
+enum class MemorySearchAlgo {
+  // Multiple objective DP search. Combine memory cost and run time cost into
+  // one single cost function and add a factor to balance them.
+  MULTI_OBJECTIVE,
+};
+
+/**
+ * @brief Config class to control memory optimizations. This should be put into
+ * config.h and be stored in FFConfig. But for easy turnaround, put this here
+ * for now.
+ */
+class MemoryOptimConfig {
+public:
+  MemoryUsageType mem_usage_type;   ///< How to represent memory cost
+  MemorySearchAlgo mem_search_algo; ///< How to search for the optimal schedule
+  float run_time_cost_factor; ///< The weight factor of run time cost in the
+                              ///< overall cost function; used in
+                              ///< MULTI_OBJECTIVE algorithm
+
+  MemoryOptimConfig()
+      : mem_usage_type{MemoryUsageType::GLOBAL},
+        mem_search_algo{MemorySearchAlgo::MULTI_OBJECTIVE},
+        run_time_cost_factor{0.5} {}
+  MemoryOptimConfig(float factor)
+      : mem_usage_type{MemoryUsageType::GLOBAL},
+        mem_search_algo{MemorySearchAlgo::MULTI_OBJECTIVE},
+        run_time_cost_factor{factor} {}
+};
+
+/**
+ * @brief Hold the result (including memory information) of a graph_optimize on
+ * a PCG.
+ */
+class MemorySearchResult {
+public:
+  float run_time_cost{};
+  float memory_cost{};
+  float search_time{};
+  ///< The max of per-device memory usage among all devices
+  float max_per_device_mem_all_deivces{0.0};
+};
+
+namespace PCG {
+
+/**
+ * @brief Class to hold memory usage information of a (sub-)PCG.
+ */
+class MemoryUsage {
+public:
+  MemoryUsageType usage_type; ///< What "num" means
+  float num;                  ///< The numerical number of memory usage
+
+  // May need this in the future, but not for now.
+  // std::vector<float> nums;     ///< Detailed number of usage for all devices
+
+  ///
+  /// Public APIs
+  ///
+  MemoryUsage() : usage_type{MemoryUsageType::GLOBAL}, num{0.0} {}
+  MemoryUsage(MemoryUsageType _usage_type, float _num)
+      : usage_type{_usage_type}, num{_num} {}
+
+  std::string to_string() const {
+    std::string type_name;
+    switch (usage_type) {
+      case MemoryUsageType::GLOBAL:
+        type_name = "GLOBAL";
+        break;
+      case MemoryUsageType::PER_DEVICE_MAX:
+        type_name = "PER_DEVICE_MAX";
+        break;
+      case MemoryUsageType::PER_DEVICE_ALL:
+        // Not supporting detailed per-device memory usage now.
+        assert(false);
+        break;
+    }
+    return "(MemoryUsageType:" + type_name + ", Usage:" + std::to_string(num) +
+           ")";
+  }
+
+  MemoryUsage &operator+=(MemoryUsage const &rhs) {
+    assert(usage_type == rhs.usage_type);
+
+    // Handle the merge of memory usage differently here.
+    switch (usage_type) {
+      case MemoryUsageType::GLOBAL:
+        num += rhs.num;
+        break;
+      case MemoryUsageType::PER_DEVICE_MAX:
+        num = std::max(num, rhs.num);
+        break;
+      case MemoryUsageType::PER_DEVICE_ALL:
+        // Not supporting detailed per-device memory usage now.
+        assert(false);
+        break;
+    }
+
+    return *this;
+  }
+
+  /**
+   * @brief Combine the memory usage of two PCGs flexibly based on
+   * MemoryUsageType.
+   */
+  friend MemoryUsage operator+(MemoryUsage lhs, MemoryUsage const &rhs) {
+    lhs += rhs;
+    return lhs;
+  }
+
+  friend std::ostream &operator<<(std::ostream &s, MemoryUsage const &usage) {
+    s << usage.to_string();
+    return s;
+  }
+};
+
+/**
+ * @brief The choice of memory optimizations applied to a Graph.
+ */
+class MemOptDecision {
+public:
+private:
+};
+
+} // namespace PCG
+} // namespace FlexFlow

include/flexflow/model.h

@@ -760,6 +760,13 @@ class FFModel {
                       bool only_data_parallel,
                       std::unique_ptr<PCG::Graph> &best_graph,
                       std::unordered_map<PCG::Node, MachineView> &optimal_view);
+  void graph_optimize(size_t budget,
+                      bool only_data_parallel,
+                      std::unique_ptr<PCG::Graph> &best_graph,
+                      std::unordered_map<PCG::Node, MachineView> &optimal_view,
+                      bool perform_memory_search,
+                      MemoryOptimConfig new_config,
+                      MemorySearchResult &search_result);
   void mcmc_optimize(std::map<Op const *, ParallelConfig> &best,
                      size_t budget,
                      float alpha,
@@ -797,6 +804,11 @@ class FFModel {
 public:
   void set_iteration_config_sequence_length(int seq_length);
 
+  /**
+   * @brief Clear the cache of the GraphSearchHelper and SearchHelper.
+   */
+  void clear_graph_search_cache();
+
 public:
   size_t op_global_guid, layer_global_guid;
   size_t tensor_global_guid, parallel_tensor_global_guid, node_global_guid;

include/flexflow/parallel_tensor.h

@@ -56,9 +56,10 @@ struct ParallelDim {
     return false;
   }
 
-  int size = 0;
-  int degree = UNKNOWN_DEGREE;
-  int parallel_idx = UNKNOWN_INDEX;
+  int size = 0;                     // Actual size of tensor
+  int degree = UNKNOWN_DEGREE;      // Degree of sharding
+  int parallel_idx = UNKNOWN_INDEX; // Runtime information, unique id of each
+                                    // degree of sharding
   bool is_replica_dim = false;
 };
 

include/flexflow/simulator.h

@@ -17,6 +17,7 @@
 
 #include "config.h"
 #include "ffconst.h"
+#include "flexflow/memory_optimization.h"
 #include "flexflow/operator_params.h"
 #include "flexflow/utils/hash_utils.h"
 #include "mpark/variant.hpp"
@@ -53,10 +54,17 @@ class FFModel;
  */
 struct CostMetrics {
   /**
-   * @brief Return the sum of the memory usage recorded in this CostMetrics.
+   * @brief Return the sum of inputs_memory, outputs_memory, and weights_memory
+   * recorded in this CostMetrics.
    */
   size_t total_memory() const;
 
+  /**
+   * @brief Return the sum of memory recorded in this CostMetrics, but in MB,
+   * instead of Bytes.
+   */
+  float total_memory_as_mb() const;
+
   /**
    * @brief Get the incremental difference between the total memory in
    * CostMetrics and sim->offset.
@@ -76,6 +84,8 @@ struct CostMetrics {
   // 2. we call Simulator::free_all before measuring an operator
   // Therefore, the current memory usage of an operator is (size_t)sim->offset
   size_t inputs_memory = 0, outputs_memory = 0, weights_memory = 0;
+  ///< Real memory usage of Op* considering parallelization over devices
+  size_t op_total_mem;
 };
 
 class Device {