<commit_before>#pragma once

#include <future>
#include <thread>

#include <boost/filesystem.hpp>
#include <boost/thread/locks.hpp>
#include <boost/thread/shared_mutex.hpp>

#include "3rd_party/threadpool.h"
#include "common/definitions.h"
#include "data/batch_generator.h"
#include "models/model_base.h"
#include "optimizers/optimizers.h"
#include "training/dropper.h"
#include "training/scheduler.h"
#include "training/sparse_tensor.h"
#include "training/training.h"
#include "training/validator.h"
#include "training/graph_group.h"

namespace marian {

template <class Builder>
class SyncGraphGroup : public GraphGroup {
public:
  typedef Builder builder_type;
  typedef typename Builder::dataset_type dataset_type;

  virtual void setScheduler(Ptr<Scheduler<dataset_type>> scheduler) {
    scheduler_ = scheduler;
    // optimizer has to be registered last to see a change of learning rate
    scheduler_->registerTrainingObserver(scheduler_);
    scheduler_->registerTrainingObserver(opt_);
  }

private:
  std::vector<Ptr<models::ModelBase>> builders_;
  std::vector<Ptr<ExpressionGraph>> graphs_;
  std::vector<size_t> devices_;

  std::vector<Tensor> params_;
  std::vector<Ptr<TensorAllocator>> paramsAllocs_;

  std::vector<Tensor> grads_;
  std::vector<Tensor> tmpTensors_;
  std::vector<Ptr<TensorAllocator>> gradsAllocs_;

  std::vector<Ptr<OptimizerBase>> shardOpt_;

  int shardSize_;
  bool first_{true};

  Ptr<Scheduler<dataset_type>> scheduler_;

  void execute(Ptr<data::Batch> batch) {

    if(first_) {

      for(size_t i = 0; i < graphs_.size(); ++i) {
        // takes care of thead_local stuff
        THREAD_GUARD(builders_[i]->build(graphs_[i], batch);
                     graphs_[i]->forward(););
      }

      if(params_.size() == 0) {
        int totalSize = graphs_[0]->params()->vals()->size();
        shardSize_ = ceil(totalSize / devices_.size());

        int pos = 0;
        for(auto device : devices_) {
          int __size__ = min(shardSize_, totalSize);
          totalSize -= __size__;


          auto paramsAlloc = New<TensorAllocator>(device);
          paramsAllocs_.push_back(paramsAlloc);

          paramsAlloc->reserveExact(__size__ * sizeof(float));

          Tensor param;
          paramsAlloc->allocate(param, {1, __size__});
          param->copyFrom(graphs_[0]->params()->vals()->subtensor(pos, __size__));

          params_.push_back(param);

          pos += __size__;
        }
      }

      if(grads_.size() == 0) {
        int totalSize = graphs_[0]->params()->vals()->size();

        for(auto device : devices_) {
          int __size__ = min(shardSize_, totalSize);
          totalSize -= __size__;

          auto gradsAlloc = New<TensorAllocator>(device);
          gradsAllocs_.push_back(gradsAlloc);

          gradsAlloc->reserveExact(2 * __size__ * sizeof(float));

          Tensor grad, tmp;
          gradsAlloc->allocate(grad, {1, __size__});
          gradsAlloc->allocate(tmp, {1, __size__});
          grads_.push_back(grad);
          tmpTensors_.push_back(tmp);
        }
      }

      first_ = false;
    }

    std::vector<Ptr<data::Batch>> batches = batch->split(devices_.size());
    std::vector<float> costs(devices_.size());

    {
      auto task = [this, &costs, batches](size_t idx) {
        auto graph = graphs_[idx];
        auto costNode = builders_[idx]->build(graph, batches[idx]);

        graph->forward();
        costs[idx] = costNode->scalar();
        graph->backward();<commit_msg>init graphs in sync sgd<commit_after>#pragma once

#include <future>
#include <thread>

#include <boost/filesystem.hpp>
#include <boost/thread/locks.hpp>
#include <boost/thread/shared_mutex.hpp>

#include "3rd_party/threadpool.h"
#include "common/definitions.h"
#include "data/batch_generator.h"
#include "models/model_base.h"
#include "optimizers/optimizers.h"
#include "training/dropper.h"
#include "training/scheduler.h"
#include "training/sparse_tensor.h"
#include "training/training.h"
#include "training/validator.h"
#include "training/graph_group.h"

namespace marian {

template <class Builder>
class SyncGraphGroup : public GraphGroup {
public:
  typedef Builder builder_type;
  typedef typename Builder::dataset_type dataset_type;

  virtual void setScheduler(Ptr<Scheduler<dataset_type>> scheduler) {
    scheduler_ = scheduler;
    // optimizer has to be registered last to see a change of learning rate
    scheduler_->registerTrainingObserver(scheduler_);
    scheduler_->registerTrainingObserver(opt_);
  }

private:
  std::vector<Ptr<models::ModelBase>> builders_;
  std::vector<Ptr<ExpressionGraph>> graphs_;
  std::vector<size_t> devices_;

  std::vector<Tensor> params_;
  std::vector<Ptr<TensorAllocator>> paramsAllocs_;

  std::vector<Tensor> grads_;
  std::vector<Tensor> tmpTensors_;
  std::vector<Ptr<TensorAllocator>> gradsAllocs_;

  std::vector<Ptr<OptimizerBase>> shardOpt_;

  int shardSize_;
  bool first_{true};

  Ptr<Scheduler<dataset_type>> scheduler_;

  void execute(Ptr<data::Batch> batch) {
    std::vector<Ptr<data::Batch>> batches = batch->split(devices_.size());

    if(first_) {

      for(size_t i = 0; i < graphs_.size(); ++i) {
        // takes care of thead_local stuff
        THREAD_GUARD(builders_[i]->build(graphs_[i], batches[i]);
                     graphs_[i]->forward(););
        if(i > 0)
          graphs_[i]->params()->vals()->copyFrom(graphs_[0]->params()->vals());
      }

      if(params_.size() == 0) {
        int totalSize = graphs_[0]->params()->vals()->size();
        shardSize_ = ceil(totalSize / devices_.size());

        int pos = 0;
        for(auto device : devices_) {
          int __size__ = min(shardSize_, totalSize);
          totalSize -= __size__;


          auto paramsAlloc = New<TensorAllocator>(device);
          paramsAllocs_.push_back(paramsAlloc);

          paramsAlloc->reserveExact(__size__ * sizeof(float));

          Tensor param;
          paramsAlloc->allocate(param, {1, __size__});
          param->copyFrom(graphs_[0]->params()->vals()->subtensor(pos, __size__));

          params_.push_back(param);

          pos += __size__;
        }
      }

      if(grads_.size() == 0) {
        int totalSize = graphs_[0]->params()->vals()->size();

        for(auto device : devices_) {
          int __size__ = min(shardSize_, totalSize);
          totalSize -= __size__;

          auto gradsAlloc = New<TensorAllocator>(device);
          gradsAllocs_.push_back(gradsAlloc);

          gradsAlloc->reserveExact(2 * __size__ * sizeof(float));

          Tensor grad, tmp;
          gradsAlloc->allocate(grad, {1, __size__});
          gradsAlloc->allocate(tmp, {1, __size__});
          grads_.push_back(grad);
          tmpTensors_.push_back(tmp);
        }
      }

      first_ = false;
    }

    std::vector<float> costs(devices_.size());

    {
      auto task = [this, &costs, batches](size_t idx) {
        auto graph = graphs_[idx];
        auto costNode = builders_[idx]->build(graph, batches[idx]);

        graph->forward();
        costs[idx] = costNode->scalar();
        graph->backward();
{'fraction_non_alphanumeric': 0.14154800507542648, 'fraction_numerical': 0.003665585788805865, 'mean_word_length': 3.2479041916167666, 'pattern_counts': {'":': 0, '<': 65, '<?xml version=': 0, '>': 117, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3199533', 'n_tokens_mistral': 2250, 'n_tokens_neox': 2154, 'n_words': 420}