Separated out layer norm implementation from class definition

Author: szaman19

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

@@ -24,3 +24,60 @@
 // permissions and limitations under the license.
 ////////////////////////////////////////////////////////////////////////////////
 
+#ifndef LBANN_LAYERSE_REGULARIZERS_DISTCONV_LAYER_NORM
+#define LBANN_LAYERSE_REGULARIZERS_DISTCONV_LAYER_NORM
+
+#if LBANN_HAS_DISTCONV
+
+namespace distconv {
+template <typename Backend, typename DataType>
+class LayerNorm
+{
+  using LocaleMPI = tensor::LocaleMPI;
+
+  template <typename Allocator>
+  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)
+  {}
+
+  template <typename Allocator>
+  void calculate_forward_stats(const DCTensor<Allocator>& input,
+                               DC<Allocator>& statistics);
+
+  template <typename Allocator>
+  void apply_normalization(const DCTensor<Allocator>& input,
+                           const DCTensor<Allocator>& statistics,
+                           DCTensor<Allocator>& output);
+
+  template <typename Allocator>
+  void calculate_backward_stats(const DCTensor<Allocator>& input,
+                                const DCTensor<Allocator>& output_grad,
+                                const DCTensor<Allocator>& statistics,
+                                DCTensor<Allocator>& statistics_grad);
+
+  template <typename Allocator>
+  void apply_grad(const DCTensor<Allocator>& input,
+                  const DCTensor<Allocator>& output_grad,
+                  const DCTensor<Allocator>& statistics,
+                  const DCTensor<Allocator>& statistics_grad,
+                  DCTensor<Allocator>& input_grad);
+
+protected:
+  Backend& m_backend;
+
+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

include/lbann/layers/regularizers/layer_norm.hpp

@@ -35,16 +35,19 @@
 #include <memory>
 
 #ifdef LBANN_HAS_DISTCONV
+#include "lbann/utils/distconv.hpp"
 #include "lbann/layers/data_type_distconv_adapter.hpp"
-#include "lbann/layers/regularizeres/distconv/distconv_layer_norm.hpp"
+#include "lbann/layers/regularizers/distconv/distconv_layer_norm.hpp"
 #endif // LBANN_HAS_DISTCONV
 
 namespace lbann {
 
 #ifdef LBANN_HAS_DISTCONV
 namespace dc {
+using Shape =  ::distconv::tensor::Shape;
+using Backend= ::distconv::BackendDNNLib; 
 template <typename TensorDataType>
-using LayerNorm = ::distconv::LayerNorm<Backend, TensorDataType>;
+using LayerNormalization = ::distconv::LayerNormalization<Backend, TensorDataType>;
 } // namespace dc
 
 template <typename TensorDataType, data_layout Layout, El::Device Device>
@@ -63,11 +66,15 @@ 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();
-
-  std::unique_ptr<dc::LayerNorm<TensorDataType>> m_layer_norm_operator;
+  
+  TensorDevType m_statistics;
+  TensorDevType m_statistics_grad;
+  std::unique_ptr<dc::LayerNormalization<TensorDataType>> m_layer_norm_operator;
 }; // class definition channelwise_fully_connected_distconv_adapter
 
 #endif // LBANN_HAS_DISTCONV
@@ -140,7 +147,7 @@ class layer_norm_layer : public data_type_layer<TensorDataType>
   get_distconv_adapter() override;
   const layer_norm_distconv_adapter<TensorDataType, Layout, Device>&
   get_distconv_adapter() const override;
-#endif
+#endif // LBANN_HAS_DISTCONV
 
 private:
   using AbsDistMatType = El::AbstractDistMatrix<TensorDataType>;
@@ -160,216 +167,14 @@ class layer_norm_layer : public data_type_layer<TensorDataType>
   std::unique_ptr<AbsDistMatType> m_statistics_gradient;
 };
 
-// =========================================================
-// Implementation
-// =========================================================
-
-template <typename T, data_layout L, El::Device D>
-void layer_norm_layer<T, L, D>::write_specific_proto(
-  lbann_data::Layer& proto) const
-{
-  proto.set_datatype(proto::ProtoDataType<T>);
-  auto* msg = proto.mutable_layer_norm();
-  msg->mutable_epsilon()->set_value(m_epsilon);
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-layer_norm_layer<TensorDataType, Layout, Device>::layer_norm_layer(
-  TensorDataType epsilon)
-  : data_type_layer<TensorDataType>(nullptr), m_epsilon(epsilon)
-{}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-layer_norm_layer<TensorDataType, Layout, Device>::layer_norm_layer(
-  const layer_norm_layer<TensorDataType, Layout, Device>& other)
-  : data_type_layer<TensorDataType>(other),
-    m_epsilon(other.m_epsilon),
-    m_statistics(other.m_statistics ? other.m_statistics->Copy() : nullptr),
-    m_statistics_gradient(other.m_statistics_gradient
-                            ? other.m_statistics_gradient->Copy()
-                            : nullptr)
-{}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-layer_norm_layer<TensorDataType, Layout, Device>&
-layer_norm_layer<TensorDataType, Layout, Device>::operator=(
-  const layer_norm_layer<TensorDataType, Layout, Device>& other)
-{
-  data_type_layer<TensorDataType>::operator=(other);
-  m_epsilon = other.m_epsilon;
-  m_statistics.reset(other.m_statistics ? other.m_statistics->Copy() : nullptr);
-  m_statistics_gradient.reset(other.m_statistics_gradient
-                                ? other.m_statistics_gradient->Copy()
-                                : nullptr);
-  return *this;
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-layer_norm_layer<TensorDataType, Layout, Device>*
-layer_norm_layer<TensorDataType, Layout, Device>::copy() const
-{
-  return new layer_norm_layer(*this);
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-std::string layer_norm_layer<TensorDataType, Layout, Device>::get_type() const
-{
-  return "layer norm";
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-data_layout
-layer_norm_layer<TensorDataType, Layout, Device>::get_data_layout() const
-{
-  return Layout;
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-El::Device
-layer_norm_layer<TensorDataType, Layout, Device>::get_device_allocation() const
-{
-  return Device;
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-description
-layer_norm_layer<TensorDataType, Layout, Device>::get_description() const
-{
-  auto desc = data_type_layer<TensorDataType>::get_description();
-  desc.add("Epsilon", m_epsilon);
-  return desc;
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-void layer_norm_layer<TensorDataType, Layout, Device>::setup_dims(
-  DataReaderMetaData& dr_metadata)
-{
-  data_type_layer<TensorDataType>::setup_dims(dr_metadata);
-  this->set_output_dims(this->get_input_dims());
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-void layer_norm_layer<TensorDataType, Layout, Device>::setup_data(
-  size_t max_mini_batch_size)
-{
-  data_type_layer<TensorDataType>::setup_data(max_mini_batch_size);
-  auto dist = this->get_prev_activations().DistData();
-  dist.colDist = El::STAR;
-  m_statistics.reset(AbsDistMatrixType::Instantiate(dist));
-  m_statistics_gradient.reset(AbsDistMatrixType::Instantiate(dist));
-}
-
-#ifdef LBANN_HAS_DISTCONV
-
-// =============================================================
-// DistConv-enabled Scatter member functions
-// =============================================================
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-bool
-layer_norm_layer<TensorDataType, Layout, Device>
-::is_distconv_supported() const {
-  return Device==El::Device::GPU && Layout == data_layout::DATA_PARALLEL;
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-void
-layer_norm_layer<TensorDataType,Layout,Device>
-::setup_distconv_adapter(const DataReaderMetaData& dr_metadata){
-  this->get_distconv_adapter_ptr() = std::make_unique<layer_norm_distconv_adapter<
-    TensorDataType, Layout, Device>>(*this);
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-const layer_norm_distconv_adapter <TensorDataType, Layout, Device>&
-layer_norm_layer<TensorDataType, Layout, Device>
-::get_distconv_adapter() const{
-  return dynamic_cast<const layer_norm_distconv_adapter<
-  TensorDataType, Layout, Device>&>(data_type_layer<TensorDataType>::get_distconv_adapter());
-}
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-layer_norm_distconv_adapter <TensorDataType, Layout, Device>&
-layer_norm_layer<TensorDataType, Layout, Device>
-::get_distconv_adapter(){
-  return const_cast<layer_norm_distconv_adapter<TensorDataType, Layout, Device>&>(
-    static_cast<const layer_norm_layer<TensorDataType, Layout, Device>&>(*this).get_distconv_adapter());
-
-
-// =============================================================
-// Scatter 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_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());
-  m_layer_norm_operator = make_unique<dc::Scatter<TensorDataType>>(dc::get_backend(),
-                                                                   layer.m_epsilon);
-}
 
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-void
-layer_norm_distconv_adapter<TensorDataType, Layout, Device>
-::fp_compute(){
-  // Compute the forward pass
-  m_layer_norm_operator->forward(this->get_prev_activations(0),
-                              this-m_epsilon);
-}
-
-
-template <typename TensorDataType, data_layout Layout, El::Device Device>
-void
-layer_norm_distconv_adapter<TensorDataType, Layout, Device>
-::bp_compute(){
-  // Compute the backward pass
-  m_layer_norm_operator->backward(this->get_prev_error_signals(0));  // Indices gradient. Will be 0'ed out
-}
-
-#define PROTO_DEVICE(T, Device)                       \
-  template class layer_norm_distconv_adapter<         \
-    T,data_layout::DATA_PARALLEL, Device>
-#include "lbann/macros/instantiate_device.hpp"
-#undef PROTO_DEVICE
 #endif // LBANN_HAS_DISTCONV
 
-LBANN_DEFINE_LAYER_BUILDER(layer_norm);
+  LBANN_DEFINE_LAYER_BUILDER(layer_norm);
 
-// =========================================================
-// Explicit template instantiation
-// =========================================================
+  // =========================================================
+  // Explicit template instantiation
+  // =========================================================
 
 #ifndef LBANN_LAYER_NORM_LAYER_INSTANTIATE
 #define PROTO_DEVICE(T, Device)                                                \