Updated implementation with updating statistics tensors

Author: szaman19

include/lbann/layers/regularizers/distconv/distconv_layer_norm.hpp

@@ -27,7 +27,7 @@
 #ifndef LBANN_LAYERSE_REGULARIZERS_DISTCONV_LAYER_NORM
 #define LBANN_LAYERSE_REGULARIZERS_DISTCONV_LAYER_NORM
 
-#if LBANN_HAS_DISTCONV
+#ifdef LBANN_HAS_DISTCONV
 
 namespace distconv {
 template <typename Backend, typename DataType>
@@ -39,21 +39,17 @@ class LayerNormalization
   using DCTensor = tensor::Tensor<DataType, LocaleMPI, Allocator>;
 
 public:
-  LayerNormalization(Backend& backend,
-                     Datatype epsilon,
-                     size_t max_mini_batch_size)
-    : m_backend(backend),
-      m_epsilon(epsilon),
-      m_max_mini_batch_size(max_mini_batch_size)
+  LayerNormalization(Backend& backend, DataType epsilon)
+    : m_backend(backend), m_epsilon(epsilon)
   {}
 
   template <typename Allocator>
   void calculate_forward_stats(const DCTensor<Allocator>& input,
-                               DC<Allocator>& statistics);
+                               DCTensor<Allocator>& statistics);
 
   template <typename Allocator>
   void apply_normalization(const DCTensor<Allocator>& input,
-                           const DCTensor<Allocator>& statistics,
+                           DCTensor<Allocator>& statistics,
                            DCTensor<Allocator>& output);
 
   template <typename Allocator>
@@ -74,10 +70,9 @@ class LayerNormalization
 
 private:
   DataType m_epsilon;
-  size_t m_max_mini_batch_size;
 
 }; // class definition LayerNorm
 } // namespace distconv
 
-#endif // LBANN_HAS_DISTONV
-#endif // LBANN_LAYERSE_REGULARIZERS_DISTCONV_LAYER_NORM
+#endif // LBANN_HAS_DISTCONV
+#endif // LBANN_LAYERS_REGULARIZERS_DISTCONV_LAYER_NORM

include/lbann/layers/regularizers/layer_norm.hpp

@@ -35,19 +35,20 @@
 #include <memory>
 
 #ifdef LBANN_HAS_DISTCONV
-#include "lbann/utils/distconv.hpp"
 #include "lbann/layers/data_type_distconv_adapter.hpp"
 #include "lbann/layers/regularizers/distconv/distconv_layer_norm.hpp"
+#include "lbann/utils/distconv.hpp"
 #endif // LBANN_HAS_DISTCONV
 
 namespace lbann {
 
 #ifdef LBANN_HAS_DISTCONV
 namespace dc {
-using Shape =  ::distconv::tensor::Shape;
-using Backend= ::distconv::BackendDNNLib; 
+using Shape = ::distconv::tensor::Shape;
+using Backend = ::distconv::BackendDNNLib;
 template <typename TensorDataType>
-using LayerNormalization = ::distconv::LayerNormalization<Backend, TensorDataType>;
+using LayerNormalization =
+  ::distconv::LayerNormalization<Backend, TensorDataType>;
 } // namespace dc
 
 template <typename TensorDataType, data_layout Layout, El::Device Device>
@@ -66,12 +67,10 @@ class layer_norm_distconv_adapter
 
   void setup_distributions(tensor_overlap_constraints& constraints) override;
   void setup_layer(size_t workspace_capacity) override;
-  void setup_fp_tensors() override;
-  void setup_bp_tensors() override;
 
   void fp_compute();
   void bp_compute();
-  
+
   TensorDevType m_statistics;
   TensorDevType m_statistics_grad;
   std::unique_ptr<dc::LayerNormalization<TensorDataType>> m_layer_norm_operator;
@@ -172,14 +171,11 @@ class layer_norm_layer : public data_type_layer<TensorDataType>
   std::unique_ptr<AbsDistMatType> m_statistics_gradient;
 };
 
+LBANN_DEFINE_LAYER_BUILDER(layer_norm);
 
-#endif // LBANN_HAS_DISTCONV
-
-  LBANN_DEFINE_LAYER_BUILDER(layer_norm);
-
-  // =========================================================
-  // Explicit template instantiation
-  // =========================================================
+// =========================================================
+// Explicit template instantiation
+// =========================================================
 
 #ifndef LBANN_LAYER_NORM_LAYER_INSTANTIATE
 #define PROTO_DEVICE(T, Device)                                                \

include/lbann/layers/regularizers/layer_norm_impl.hpp

@@ -33,8 +33,7 @@
 #include "lbann/layers/data_type_distconv_adapter.hpp"
 #endif
 
-namespace lbann{
-
+namespace lbann {
 
 // =========================================================
 // Implementation
@@ -135,7 +134,6 @@ void layer_norm_layer<TensorDataType, Layout, Device>::setup_data(
   m_statistics_gradient.reset(AbsDistMatrixType::Instantiate(dist));
 }
 
-
 #ifdef LBANN_HAS_DISTCONV
 
 // =============================================================
@@ -174,57 +172,53 @@ layer_norm_layer<TensorDataType, Layout, Device>::get_distconv_adapter()
     layer_norm_distconv_adapter<TensorDataType, Layout, Device>&>(
     static_cast<const layer_norm_layer<TensorDataType, Layout, Device>&>(*this)
       .get_distconv_adapter());
+}
 
 // =============================================================
 // LayerNorm DistConv Adapter implementation
 // =============================================================
 
-  template <typename TensorDataType, data_layout Layout, El::Device Device>
-  void layer_norm_distconv_adapter<TensorDataType, Layout, Device>::
-    setup_distributions(tensor_overlap_constraints & constraints)
-  {
-    data_type_distconv_adapter<TensorDataType>::setup_distributions(
-      constraints);
-    // no overlap needed
-    for (auto& d : this->m_prev_activations_dists) {
-      d.clear_overlap();
-      constraints.mark_updated(d);
-      constraints.mark_invariant(d);
-    }
-    for (auto& d : this->m_activations_dists) {
-      d.clear_overlap();
-      constraints.mark_updated(d);
-      constraints.mark_invariant(d);
-    }
-    for (auto& d : this->m_prev_error_signals_dists) {
-      d.clear_overlap();
-      constraints.mark_updated(d);
-      constraints.mark_invariant(d);
-    }
-    for (auto& d : this->m_error_signals_dists) {
-      d.clear_overlap();
-      constraints.mark_updated(d);
-      constraints.mark_invariant(d);
-    }
+template <typename TensorDataType, data_layout Layout, El::Device Device>
+void layer_norm_distconv_adapter<TensorDataType, Layout, Device>::
+  setup_distributions(tensor_overlap_constraints& constraints)
+{
+  data_type_distconv_adapter<TensorDataType>::setup_distributions(constraints);
+  // no overlap needed
+  for (auto& d : this->m_prev_activations_dists) {
+    d.clear_overlap();
+    constraints.mark_updated(d);
+    constraints.mark_invariant(d);
   }
-
-  template <typename TensorDataType, data_layout Layout, El::Device Device>
-  void layer_norm_distconv_adapter<TensorDataType, Layout, Device>::setup_layer(
-    size_t workspace_capacity)
-  {
-    data_type_distconv_adapter<TensorDataType>::setup_layer(workspace_capacity);
-    auto& layer = dynamic_cast<
-      channelwise_fully_connected_layer<TensorDataType, Layout, Device>&>(
-      this->layer());
-    const auto max_mini_batch_size =
-      layer.get_model()->m_max_mini_batch_size_distconv;
-
-    m_layer_norm_operator =
-      make_unique<dc::LayerNormalization<TensorDataType>>(dc::get_backend(),
-                                                          layer.m_epsilon,
-                                                          max_mini_batch_size);
+  for (auto& d : this->m_activations_dists) {
+    d.clear_overlap();
+    constraints.mark_updated(d);
+    constraints.mark_invariant(d);
+  }
+  for (auto& d : this->m_prev_error_signals_dists) {
+    d.clear_overlap();
+    constraints.mark_updated(d);
+    constraints.mark_invariant(d);
   }
+  for (auto& d : this->m_error_signals_dists) {
+    d.clear_overlap();
+    constraints.mark_updated(d);
+    constraints.mark_invariant(d);
+  }
+}
+
+template <typename TensorDataType, data_layout Layout, El::Device Device>
+void layer_norm_distconv_adapter<TensorDataType, Layout, Device>::setup_layer(
+  size_t workspace_capacity)
+{
+  data_type_distconv_adapter<TensorDataType>::setup_layer(workspace_capacity);
+  auto& layer = dynamic_cast<layer_norm_layer<TensorDataType, Layout, Device>&>(
+    this->layer());
+
+  m_layer_norm_operator =
+    make_unique<dc::LayerNormalization<TensorDataType>>(dc::get_backend(),
+                                                        layer.m_epsilon);
+}
 
-#endif LBANN_HAS_DISTCONV
+#endif // LBANN_HAS_DISTCONV
 } // namespace lbann
 #endif // LBANN_LAYER_REGULARIZER_LAYER_NORM_IMPL_HPP_INCLUDED