Fix model training instance abstraction

lib/local-execution/include/local-execution/local_training_backing.h

@@ -5,6 +5,7 @@
 #include "local-execution/model_training_instance.dtg.h"
 #include "local-execution/task_registry.h"
 #include "pcg/computation_graph.dtg.h"
+#include "pcg/optimizer_attrs.dtg.h"
 
 namespace FlexFlow {
 
@@ -16,7 +17,8 @@ struct LocalTrainingBacking {
                        ComputationGraph const &,
                        TensorBackingMap const &,
                        RuntimeArgConfig const &,
-                       std::optional<ModelTrainingInstance> &);
+                       std::optional<ModelTrainingInstance> const &,
+                       std::optional<OptimizerAttrs> const &);
 
   void execute_init();
   PerLayerElapsedTime execute_forward();
@@ -38,6 +40,7 @@ struct LocalTrainingBacking {
   TaskRegistry task_registry;
   LocalSlotsBacking local_slots_backing;
   std::optional<ModelTrainingInstance> training_instance;
+  std::optional<OptimizerAttrs> optimizer_attrs;
 };
 
 } // namespace FlexFlow

lib/local-execution/include/local-execution/model_training_instance.struct.toml

@@ -9,8 +9,7 @@ features = [
 
 includes = [
   "op-attrs/ops/loss_functions/loss_attrs.dtg.h",
-  "pcg/tensor_guid_t.dtg.h",
-  "pcg/optimizer_attrs.dtg.h",
+  "pcg/tensor_guid_t.dtg.h"
 ]
 
 [[fields]]
@@ -24,7 +23,3 @@ type = "::FlexFlow::tensor_guid_t"
 [[fields]]
 name = "logit_tensor"
 type = "::FlexFlow::tensor_guid_t"
-
-[[fields]]
-name = "optimizer_attrs"
-type = "::FlexFlow::OptimizerAttrs"

lib/local-execution/include/local-execution/non_graph_tensor_guid_t.struct.toml

@@ -8,10 +8,6 @@ features = [
   "json",
 ]
 
-includes = [
-  "<cstddef>",
-]
-
 [[fields]]
 name = "raw_uid"
-type = "size_t"
+type = "int"

lib/local-execution/src/local_cost_estimator.cc

@@ -76,11 +76,13 @@ CostDetails LocalCostEstimator::estimate_cost(
                            get_vector_piece_attrs(outputs));
 
   std::optional<ModelTrainingInstance> model_training_instance = std::nullopt;
+  std::optional<OptimizerAttrs> optimizer_attrs = std::nullopt;
   LocalTrainingBacking local_backing(allocator,
                                      cg_builder.computation_graph,
                                      tensor_backing_map,
                                      this->runtime_arg_config,
-                                     model_training_instance);
+                                     model_training_instance,
+                                     optimizer_attrs);
 
   local_backing.execute_init();
   PerLayerElapsedTime fwd = local_backing.execute_forward();

lib/local-execution/src/local_training_backing.cc

@@ -1,10 +1,10 @@
 #include "local-execution/local_training_backing.h"
 #include "local-execution/loss_functions.h"
-#include "local-execution/model_training_instance.h"
 #include "local-execution/optimizer.h"
 #include "local-execution/task_invocation.h"
 #include "local-execution/task_signature_impl.h"
 #include "pcg/computation_graph.h"
+#include "pcg/optimizer_attrs.h"
 #include "utils/containers/contains.h"
 #include "utils/containers/contains_key.h"
 #include "utils/containers/get_only.h"
@@ -18,11 +18,12 @@ LocalTrainingBacking::LocalTrainingBacking(
     ComputationGraph const &computation_graph,
     TensorBackingMap const &tensor_backing_mapping,
     RuntimeArgConfig const &runtime_arg_config,
-    std::optional<ModelTrainingInstance> &training_instance)
+    std::optional<ModelTrainingInstance> const &training_instance,
+    std::optional<OptimizerAttrs> const &optimizer_attrs)
     : allocator(allocator), computation_graph(computation_graph),
       local_slots_backing(tensor_backing_mapping, runtime_arg_config),
       task_registry(empty_task_registry()),
-      training_instance(training_instance) {
+      training_instance(training_instance), optimizer_attrs(optimizer_attrs) {
 
   for (layer_guid_t const &node :
        topological_ordering(this->computation_graph)) {
@@ -38,8 +39,8 @@ LocalTrainingBacking::LocalTrainingBacking(
 
     // allocate optimizer buffers
     if (attrs.has<WeightAttrs>() && this->training_instance.has_value()) {
-      OptimizerAttrs attrs = this->training_instance.value().optimizer_attrs;
-      TaskSignature sig = get_update_signature(attrs);
+      assert(this->optimizer_attrs.has_value());
+      TaskSignature sig = get_update_signature(this->optimizer_attrs.value());
       tensor_guid_t weight_tensor =
           get_only(get_outgoing_tensors(this->computation_graph, node));
       this->local_slots_backing.allocate_optimizer_tensors(
@@ -153,7 +154,7 @@ PerLayerElapsedTime LocalTrainingBacking::execute_backward() {
 
 void LocalTrainingBacking::execute_update() {
   assert(this->training_instance.has_value());
-  OptimizerAttrs attrs = this->training_instance.value().optimizer_attrs;
+  assert(this->optimizer_attrs.has_value());
 
   for (layer_guid_t const &node :
        topological_ordering(this->computation_graph)) {
@@ -166,18 +167,19 @@ void LocalTrainingBacking::execute_update() {
           this->local_slots_backing.weight_optimizer_tensor_guids.at(node);
 
       // get invocation
-      TaskInvocation invocation =
-          get_update_invocation(attrs, weight_tensor, grad_buffer_tensors);
+      TaskInvocation invocation = get_update_invocation(
+          this->optimizer_attrs.value(), weight_tensor, grad_buffer_tensors);
       // assert(is_invocation_valid(get_update_signature(attrs), invocation));
 
       // execute update
       TaskArgumentAccessor accessor = this->get_task_arg_accessor(invocation);
-      TaskImplFunction update_impl_fn = get_update_task_impl(attrs);
+      TaskImplFunction update_impl_fn =
+          get_update_task_impl(this->optimizer_attrs.value());
       update_impl_fn.get<GenericTaskImplFunction>().function_ptr(accessor);
     }
   }
 
-  this->training_instance = next(this->training_instance.value());
+  this->optimizer_attrs = next(this->optimizer_attrs.value());
 }
 
 TaskArgumentAccessor LocalTrainingBacking::get_task_arg_accessor(

lib/local-execution/test/src/test_loss_e2e.cc

@@ -59,13 +59,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
               LossAttrs{SparseCategoricalCrossEntropyLossAttrs{
                   /*replace_labels=*/false}},
               label_tensor,
-              logit_tensor,
-              optimizer_attrs};
+              logit_tensor};
       LocalTrainingBacking local_backing(allocator,
                                          cg_builder.computation_graph,
                                          tensor_backing_map,
                                          runtime_arg_config,
-                                         model_training_instance);
+                                         model_training_instance,
+                                         optimizer_attrs);
       local_backing.execute_init();
       local_backing.execute_forward();
       local_backing.execute_backward();
@@ -83,13 +83,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
             ModelTrainingInstance{LossAttrs{NonconfigurableLossAttrs{
                                       LossFunction::CATEGORICAL_CROSSENTROPY}},
                                   label_tensor,
-                                  logit_tensor,
-                                  optimizer_attrs};
+                                  logit_tensor};
         LocalTrainingBacking local_backing(allocator,
                                            cg_builder.computation_graph,
                                            tensor_backing_map,
                                            runtime_arg_config,
-                                           model_training_instance);
+                                           model_training_instance,
+                                           optimizer_attrs);
         local_backing.execute_init();
         local_backing.execute_forward();
         local_backing.execute_backward();
@@ -101,13 +101,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
                 LossAttrs{NonconfigurableLossAttrs{
                     LossFunction::MEAN_SQUARED_ERROR_AVG_REDUCE}},
                 label_tensor,
-                logit_tensor,
-                optimizer_attrs};
+                logit_tensor};
         LocalTrainingBacking local_backing(allocator,
                                            cg_builder.computation_graph,
                                            tensor_backing_map,
                                            runtime_arg_config,
-                                           model_training_instance);
+                                           model_training_instance,
+                                           optimizer_attrs);
         local_backing.execute_init();
         local_backing.execute_forward();
         local_backing.execute_backward();
@@ -118,13 +118,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
             ModelTrainingInstance{
                 LossAttrs{NonconfigurableLossAttrs{LossFunction::IDENTITY}},
                 label_tensor,
-                logit_tensor,
-                optimizer_attrs};
+                logit_tensor};
         LocalTrainingBacking local_backing(allocator,
                                            cg_builder.computation_graph,
                                            tensor_backing_map,
                                            runtime_arg_config,
-                                           model_training_instance);
+                                           model_training_instance,
+                                           optimizer_attrs);
         local_backing.execute_init();
         local_backing.execute_forward();
         local_backing.execute_backward();

lib/local-execution/test/src/test_update_e2e.cc

@@ -58,13 +58,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
                 LossAttrs{NonconfigurableLossAttrs{
                     LossFunction::MEAN_SQUARED_ERROR_AVG_REDUCE}},
                 label_tensor,
-                logit_tensor,
-                optimizer_attrs};
+                logit_tensor};
         LocalTrainingBacking local_backing(allocator,
                                            cg_builder.computation_graph,
                                            tensor_backing_map,
                                            runtime_arg_config,
-                                           model_training_instance);
+                                           model_training_instance,
+                                           optimizer_attrs);
         local_backing.execute_init();
         local_backing.execute_forward();
         local_backing.execute_backward();
@@ -81,13 +81,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
                 LossAttrs{NonconfigurableLossAttrs{
                     LossFunction::MEAN_SQUARED_ERROR_AVG_REDUCE}},
                 label_tensor,
-                logit_tensor,
-                optimizer_attrs};
+                logit_tensor};
         LocalTrainingBacking local_backing(allocator,
                                            cg_builder.computation_graph,
                                            tensor_backing_map,
                                            runtime_arg_config,
-                                           model_training_instance);
+                                           model_training_instance,
+                                           optimizer_attrs);
         local_backing.execute_init();
         local_backing.execute_forward();
         local_backing.execute_backward();
@@ -109,13 +109,13 @@ TEST_SUITE(FF_CUDA_TEST_SUITE) {
               LossAttrs{NonconfigurableLossAttrs{
                   LossFunction::MEAN_SQUARED_ERROR_AVG_REDUCE}},
               label_tensor,
-              logit_tensor,
-              optimizer_attrs};
+              logit_tensor};
       LocalTrainingBacking local_backing(allocator,
                                          cg_builder.computation_graph,
                                          tensor_backing_map,
                                          runtime_arg_config,
-                                         model_training_instance);
+                                         model_training_instance,
+                                         optimizer_attrs);
       local_backing.execute_init();
       local_backing.execute_forward();
       local_backing.execute_backward();

lib/pcg/include/pcg/optimizer_attrs.h

@@ -0,0 +1,13 @@
+
+#ifndef _FLEXFLOW_PCG_OPTIMIZER_ATTRS_H
+#define _FLEXFLOW_PCG_OPTIMIZER_ATTRS_H
+
+#include "pcg/optimizer_attrs.dtg.h"
+
+namespace FlexFlow {
+
+OptimizerAttrs next(OptimizerAttrs const &old);
+
+} // namespace FlexFlow
+
+#endif

lib/pcg/src/pcg/optimizer_attrs.cc

@@ -0,0 +1,24 @@
+#include "pcg/optimizer_attrs.h"
+
+namespace FlexFlow {
+
+OptimizerAttrs next(OptimizerAttrs const &old_attrs) {
+  if (old_attrs.has<AdamOptimizerAttrs>()) {
+    AdamOptimizerAttrs old = old_attrs.get<AdamOptimizerAttrs>();
+    double new_beta1_t = old.beta_t * old.beta1;
+    double new_beta2_t = old.beta2_t * old.beta2;
+    double new_alpha_t = old.alpha * sqrt(1 - new_beta2_t) / (1 - new_beta1_t);
+    return OptimizerAttrs{AdamOptimizerAttrs{old.alpha,
+                                             old.beta1,
+                                             old.beta2,
+                                             old.weight_decay,
+                                             new_alpha_t,
+                                             new_beta1_t,
+                                             new_beta2_t,
+                                             old.epsilon}};
+  } else {
+    return old_attrs;
+  }
+}
+
+} // namespace FlexFlow