optimizer.hpp

Go to the documentation of this file.

// Copyright (c) 2014-2023, Lawrence Livermore National Security, LLC.
// Produced at the Lawrence Livermore National Laboratory.
// Written by the LBANN Research Team (B. Van Essen, et al.) listed in
// the CONTRIBUTORS file. <lbann-dev@llnl.gov>
//
// LLNL-CODE-697807.
// All rights reserved.
//
// This file is part of LBANN: Livermore Big Artificial Neural Network
// Toolkit. For details, see http://software.llnl.gov/LBANN or
// https://github.com/LLNL/LBANN.
//
// Licensed under the Apache License, Version 2.0 (the "Licensee"); 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.

#ifndef LBANN_OPTIMIZERS_OPTIMIZER_HPP_INCLUDED
#define LBANN_OPTIMIZERS_OPTIMIZER_HPP_INCLUDED

#include "lbann/base.hpp"
#include "lbann/comm_nb_request.hpp"
#include "lbann/utils/cloneable.hpp"
#include "lbann/utils/compiler_control.hpp"
#ifdef LBANN_HAS_GPU
#include "lbann/utils/gpu/helpers.hpp"
#endif // LBANN_HAS_GPU
#include "lbann/utils/description.hpp"
#include "lbann/utils/memory.hpp"

#include <memory>
#include <string>
#include <typeindex>
#include <unordered_set>

namespace lbann_data {
class Optimizer;
}

namespace lbann {

enum class optimizer_gradient_status
{
  ready,
  cleared,
  allreduce_needed,
  allreduce_started,
};

std::string to_string(optimizer_gradient_status status);

// Forward declarations
class lbann_comm;
class persist;
class weights;

class optimizer : public Cloneable<HasAbstractFunction<optimizer>>
{
public:

  optimizer();
  virtual ~optimizer() = default;


  virtual std::string get_type() const = 0;
  virtual description get_description() const;

  virtual double get_learning_rate() const = 0;
  virtual void set_learning_rate(double) = 0;


  virtual void setup(weights* w) = 0;

  template <typename TensorDataType>
  void add_to_gradient(El::AbstractDistMatrix<TensorDataType> const& contrib,
                       TensorDataType scale = 1.f,
                       bool allreduce_needed = false);

  void clear_gradient();

  El::Int get_num_gradient_sources() const;
  void add_gradient_source(const void* source);

  void remove_gradient_source(const void* source);

  virtual void step() = 0;

  template <typename TensorDataType>
  El::AbstractDistMatrix<TensorDataType>&
  get_gradient_buffer(TensorDataType& buf_scale,
                      TensorDataType& in_scale,
                      bool allreduce_needed = false);



  lbann_comm& get_comm() { return *m_comm; }

  const lbann_comm& get_comm() const { return *m_comm; }



  EvalType get_step_time() const { return m_step_time; }

  virtual void reset_counters() { m_step_time = 0; }



  template <class Archive>
  void serialize(Archive& ar);


  virtual void write_proto(lbann_data::Optimizer& proto) const = 0;

  class GradientHelper
  {
  public:
    virtual ~GradientHelper() = default;
    optimizer_gradient_status get_status() const noexcept { return status_; }
    void set_status(optimizer_gradient_status s) noexcept { status_ = s; }
    virtual El::BaseDistMatrix& gradient() noexcept = 0;
    virtual El::BaseDistMatrix const& gradient() const noexcept = 0;
    virtual void start_allreduce(lbann_comm&) = 0;
    virtual void complete_allreduce(lbann_comm&) = 0;
    virtual void clear() = 0;

  private:
    optimizer_gradient_status status_ = optimizer_gradient_status::cleared;
  }; // class GradientHelper

  template <typename TensorDataType>
  class GradientHelperImpl : public GradientHelper
  {
  public:
    using AbsDistMatType = El::AbstractDistMatrix<TensorDataType>;

  public:
    GradientHelperImpl(El::Int height, El::Int width, El::DistData dist_data)
      : gradient_{AbsDistMatType::Instantiate(dist_data)}
    {
      El::Zeros(*gradient_, height, width);
    }
    AbsDistMatType& gradient() noexcept override { return *gradient_; }
    AbsDistMatType const& gradient() const noexcept override
    {
      return *gradient_;
    }
    void start_allreduce(lbann_comm& comm) override;
    void complete_allreduce(lbann_comm& comm) override;
    void clear() override;

  private:
    std::unique_ptr<AbsDistMatType> gradient_;
    Al::request allreduce_req_;
  }; // class GradientHelperImpl

  optimizer(const optimizer& other);
  optimizer& operator=(const optimizer& other);

  optimizer_gradient_status get_gradient_status() const
  {
    return m_gradient_status;
  }
  void set_gradient_status(const optimizer_gradient_status status)
  {
    m_gradient_status = status;
  }
  std::unordered_set<const void*>& get_gradient_sources()
  {
    return m_gradient_sources;
  }
  void set_comm(lbann_comm& comm) { m_comm = &comm; }

  void set_step_time(EvalType time) { m_step_time = time; }

  void inc_step_time(EvalType time) { m_step_time += time; }

  virtual std::tuple<El::Int, El::Int, El::DistData>
  get_matrix_info() const = 0;

  template <typename TensorDataType>
  void accumulate_all_gradient_contributions(
    El::AbstractDistMatrix<TensorDataType>& gradient);

  void start_gradient_allreduce();

  void finish_gradient_allreduce();

private:
  lbann_comm* m_comm;

  std::unordered_set<const void*> m_gradient_sources;

  optimizer_gradient_status m_gradient_status =
    optimizer_gradient_status::cleared;

  EvalType m_step_time = 0;

  using gradient_manager_type = GradientHelper;
  using gradient_manager_ptr = std::unique_ptr<gradient_manager_type>;
  std::unordered_map<std::type_index, gradient_manager_ptr> gradients_;
};

} // namespace lbann

#endif // LBANN_OPTIMIZERS_OPTIMIZER_HPP_INCLUDED