Transformer updates and memory optimizations (#2370)

*  Save memory on backpropagation by releasing gradients in GradientManager as soon as possible
* Save more memory in FSDP by synchronizing previous outstanding async communication calls and freeing up local gradient contributions
* FSDP: release full weight views after backprop
* Minor tweaks to transformer training script

Author: tbennun

applications/nlp/transformer/parallelism.py

@@ -56,9 +56,11 @@ def apply_fsdp_allweights(model: lbann.Model, args: argparse.Namespace):
 
     # Loop over all weights
     for layer in model.layers:
+        # As a heuristic, only shard the first set of weights (i.e., no
+        # biases) and skip layer normalization
+        if 'LayerNorm' in str(type(layer)):
+            continue
         if layer.weights:
-            # As a heuristic, only shard the first set of weights (i.e., no
-            # biases)
             if len(layer.weights) > 0:
                 layer.weights[0].sharded = True
 

applications/nlp/transformer/trainer.py

@@ -65,12 +65,13 @@ def construct_training_task(model: lbann.Model,
 # Data reader
 # ----------------------------------------------
 def make_data_reader(dataset_name: str, fraction: float, validate: bool,
-                     val_fraction: float):
+                     val_fraction: float, always_shuffle: bool):
     reader = lbann.reader_pb2.DataReader()
     _reader = reader.reader.add()
     _reader.name = 'python'
     _reader.role = 'train'
-    _reader.shuffle = False if 'pretokenized' in dataset_name else True
+    _reader.shuffle = (True if always_shuffle
+                       or 'pretokenized' not in dataset_name else False)
     _reader.fraction_of_data_to_use = fraction
     _reader.python.module = dataset_name
     _reader.python.module_dir = os.path.join(
@@ -124,7 +125,8 @@ def make_batch_script(model: lbann.Model,
                             training_algo=algo)
     reader = make_data_reader(dataset_name, args.dataset_fraction,
                               not args.skip_validation,
-                              args.validation_set_fraction)
+                              args.validation_set_fraction,
+                              args.always_shuffle)
 
     # Optimizer with learning rate schedule
     if args.optimizer.lower() == 'adamw':
@@ -188,6 +190,10 @@ def make_batch_script(model: lbann.Model,
                 epoch_interval=1,
             ))
 
+    if args.validate_every > 0:
+        model.callbacks.append(
+            lbann.CallbackEvaluateProgress(batch_interval=args.validate_every))
+
     # Print a progress bar
     if args.progress:
         model.callbacks.append(
@@ -237,6 +243,13 @@ def add_training_arguments(parser: argparse.ArgumentParser):
                         action="store_true",
                         default=False,
                         help="Do not run validation (default: false)")
+    parser.add_argument(
+        "--always-shuffle",
+        action="store_true",
+        default=False,
+        help=
+        "Always shuffle training dataset, even if pretokenized (default: false)"
+    )
     parser.add_argument(
         "--validation-set-fraction",
         type=float,
@@ -248,6 +261,10 @@ def add_training_arguments(parser: argparse.ArgumentParser):
         default=False,
         help="Save prototext experiment file instead of protobin (slower but "
         "debuggable) (default: false)")
+    parser.add_argument("--validate-every",
+                        type=int,
+                        default=100,
+                        help="Run validation every N steps (default: 100)")
 
 
 # ----------------------------------------------

include/lbann/optimizers/optimizer_impl.hpp

@@ -46,29 +46,33 @@ class GradientHelperImpl : public optimizer::GradientHelper
                      El::DistData grad_dist_data,
                      bool sharded_weights)
     : local_gradient_contrib_{AbsDistMatType::Instantiate(dist_data)},
-      local_contrib_dist_{dist_data},
       global_gradient_{AbsDistMatType::Instantiate(grad_dist_data)},
-      global_dist_{grad_dist_data},
       sharded_weights_{sharded_weights}
   {
     ensure_gradient_memory(height, width);
     El::Zeros(*local_gradient_contrib_, height, width);
-    if (grad_dist_data != dist_data) {
+    if (sharded_weights) {
       El::Zeros(*global_gradient_, height, width);
     }
   }
 
   void ensure_gradient_memory(El::Int height, El::Int width) override
   {
 #if defined(LBANN_HAS_GPU)
-    local_gradient_contrib_->Matrix().SetMemoryMode(1);
+    static const char* e = std::getenv("LBANN_USE_DIRECT_FOR_CONTRIB");
+    if (e != nullptr && e[0] == '1') {
+      local_gradient_contrib_->Matrix().SetMemoryMode(0);
+    }
+    else {
+      local_gradient_contrib_->Matrix().SetMemoryMode(1);
+    }
 #endif // LBANN_HAS_GPU
 
     if (local_gradient_contrib_->Width() == 0) {
       local_gradient_contrib_->Resize(height, width);
       // If distribution is the same, have global gradient matrix view the
       // local contributions.
-      if (local_contrib_dist_ == global_dist_) {
+      if (!sharded_weights_) {
         El::View(*global_gradient_, *local_gradient_contrib_);
       }
     }
@@ -96,6 +100,13 @@ class GradientHelperImpl : public optimizer::GradientHelper
 
   void start_sync(lbann_comm& comm) override
   {
+    // Complete outstanding synchronization of the same data type
+    static GradientHelperImpl<TensorDataType>* lastsync = nullptr;
+    if (lastsync != nullptr) {
+      lastsync->complete_sync(comm);
+      lastsync = nullptr;
+    }
+
     switch (this->get_status()) {
     case optimizer_gradient_status::sync_needed:
       // Sharded gradients are produced from a reduce-scatter on the local
@@ -122,6 +133,7 @@ class GradientHelperImpl : public optimizer::GradientHelper
         */
       }
       this->set_status(optimizer_gradient_status::sync_started);
+      lastsync = this;
       break;
     case optimizer_gradient_status::ready:
     case optimizer_gradient_status::cleared:
@@ -166,19 +178,19 @@ class GradientHelperImpl : public optimizer::GradientHelper
   void clear() override
   {
     this->set_status(optimizer_gradient_status::cleared);
+    local_gradient_contrib_->Empty();
+    global_gradient_->Empty();
   }
 
 private:
   /** Matches the distribution of gathered (unsharded) weights in backprop. */
   std::unique_ptr<AbsDistMatType> local_gradient_contrib_;
-  El::DistData local_contrib_dist_;
 
   /** Matches the distribution of data_type_optimizer<T>::m_gradient (i.e.,
    *  post synchronization). Will view said matrix if only one data type
    *  exists.
    */
   std::unique_ptr<AbsDistMatType> global_gradient_;
-  El::DistData global_dist_;
 
   Al::request sync_req_;
   bool sharded_weights_;
@@ -218,6 +230,8 @@ optimizer::get_gradient_buffer(TensorDataType& buf_scale,
   // If the manager hasn't been created, let's make it.
   auto mat_info = this->get_matrix_info();
   if (!grad_mgr_ptr) {
+    // If our optimizer contains a gradient of the same data type, reuse (view)
+    // it in the gradient manager
     grad_mgr_ptr = std::make_unique<GradMgrType>(std::get<HEIGHT>(mat_info),
                                                  std::get<WIDTH>(mat_info),
                                                  std::get<DISTDATA_L>(mat_info),
@@ -319,13 +333,13 @@ void optimizer::accumulate_all_gradient_contributions(
   // Handle the case that only 1 update of a different type is needed.
   if (num_updates == 1UL &&
       this->m_local_gradient_contributions.size() == 1UL) {
-    auto const& grad_mgr =
-      *(this->m_local_gradient_contributions.begin()->second);
+    auto& grad_mgr = *(this->m_local_gradient_contributions.begin()->second);
     if (grad_mgr.get_status() != optimizer_gradient_status::ready) {
       LBANN_ERROR("Expected ready status. Got: ",
                   to_string(grad_mgr.get_status()));
     }
     El::Copy(grad_mgr.global_gradient(), gradient);
+    grad_mgr.clear();
   }
   else if (this->m_local_gradient_contributions.size() > 1UL) {
     // Need a temporary matrix for the type-casted copy.
@@ -335,14 +349,15 @@ void optimizer::accumulate_all_gradient_contributions(
     for (auto const& grad_mgr_v : this->m_local_gradient_contributions) {
       if (grad_mgr_v.first == this_type_idx)
         continue;
-      auto const& grad_mgr = *(grad_mgr_v.second);
+      auto& grad_mgr = *(grad_mgr_v.second);
       if (grad_mgr.get_status() != optimizer_gradient_status::ready) {
         LBANN_ERROR("Expected ready status. Got: ",
                     to_string(grad_mgr.get_status()));
       }
       auto const& grad_base = grad_mgr.global_gradient();
       El::Copy(grad_base, *tmp);
       El::Axpy(one, *tmp, gradient);
+      grad_mgr.clear();
     }
   }
 }

python/lbann/models/transformer.py

@@ -29,20 +29,23 @@ def __init__(self, normalized_shape, name=None, builtin=True):
         self.name = (name if name else f'layernorm{LayerNorm.global_count}')
         self.builtin = builtin
 
-        if not self.builtin:
-            # Initialize weights
-            self.weight = lbann.Weights(
-                initializer=lbann.ConstantInitializer(value=1),
-                name=f'{self.name}_weight',
-            )
-            self.bias = lbann.Weights(
-                initializer=lbann.ConstantInitializer(value=0),
-                name=f'{self.name}_bias',
-            )
+        # Initialize weights
+        self.weight = lbann.Weights(
+            initializer=lbann.ConstantInitializer(value=1),
+            name=f'{self.name}_weight',
+        )
+        self.bias = lbann.Weights(
+            initializer=lbann.ConstantInitializer(value=0),
+            name=f'{self.name}_bias',
+        )
 
     def forward(self, x):
         if self.builtin:
-            return lbann.LayerNorm(x, scale=True, bias=True, name=self.name)
+            return lbann.LayerNorm(x,
+                                   scale=True,
+                                   bias=True,
+                                   name=self.name,
+                                   weights=[self.weight, self.bias])
 
         # Normalization
         x = lbann.InstanceNorm(x)

src/layers/data_type_layer.cpp

@@ -329,6 +329,11 @@ void data_type_layer<InputTensorDataType,
   }
 #endif // defined(LBANN_HAS_GPU) && defined(LBANN_DEBUG)
 
+  // Release the now-unnecessary full weight views
+  for (size_t i = 0; i < this->num_weights(); ++i) {
+    this->get_weights(i).release_full_weights();
+  }
+
   // Release activation memory as necessary
   model* m = this->get_model();
   if (m != nullptr) {

src/weights/data_type_weights.cpp

@@ -239,13 +239,17 @@ void data_type_weights<TensorDataType>::do_setup_()
   }
 
   // Construct matrix for weight values
-  // If sharded, use STAR_VC distribution (column distributed)
+  // If sharded, use STAR_VC distribution (column distributed) or VC_STAR (row
+  // distributed) if width=1.
   auto matrix_dist = this->get_matrix_distribution();
+  bool must_use_vc_star = (this->get_matrix_width() == 1);
   m_values.reset(AbsDistMatrixType::Instantiate(
     *matrix_dist.grid,
     matrix_dist.root,
-    this->is_sharded() ? El::STAR : matrix_dist.colDist,
-    this->is_sharded() ? El::VC : matrix_dist.rowDist,
+    this->is_sharded() ? (must_use_vc_star ? El::VC : El::STAR)
+                       : matrix_dist.colDist,
+    this->is_sharded() ? (must_use_vc_star ? El::STAR : El::VC)
+                       : matrix_dist.rowDist,
     (matrix_dist.blockHeight == 1 && matrix_dist.blockWidth == 1 ? El::ELEMENT
                                                                  : El::BLOCK),
     matrix_dist.device));