Move embeddings scale initialization in a single function (OpenNMT#834)

reinfer · Jun 17, 2022 · 2c8db5e · 2c8db5e
1 parent 08eed4a
commit 2c8db5e
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 10 deletions.
diff --git a/include/ctranslate2/layers/common.h b/include/ctranslate2/layers/common.h
@@ -32,7 +32,6 @@ namespace ctranslate2 {
       const StorageView& _embeddings;
       const DataType _output_type;
       const StorageView* _qscale;
-      std::unique_ptr<const StorageView> _scale;
     };
 
     // This enum order should remain fixed.

diff --git a/src/layers/common.cc b/src/layers/common.cc
@@ -50,10 +50,6 @@ namespace ctranslate2 {
       , _output_type(get_default_float_type(model.effective_compute_type()))
       , _qscale(model.get_variable_if_exists(scope + "/weight_scale"))
     {
-      if (model.get_flag_with_default(scope + "/multiply_by_sqrt_depth", true)) {
-        const StorageView scale(std::sqrt(static_cast<float>(_embeddings.dim(1))));
-        _scale = std::make_unique<StorageView>(scale.to(_output_type));
-      }
     }
 
     DataType Embeddings::output_type() const {
@@ -81,9 +77,6 @@ namespace ctranslate2 {
       } else {
         _gather_op(_embeddings, ids, output);
       }
-
-      if (_scale)
-        ops::Mul()(output, *_scale, output);
     }
 
 

diff --git a/src/layers/transformer.cc b/src/layers/transformer.cc
@@ -136,13 +136,15 @@ namespace ctranslate2 {
                            const std::string& scope,
                            const Layer& embeddings) {
       const auto* scale = model.get_variable_if_exists(scope + "/scale_embeddings");
+
+      // Backward compatibility with older models.
       if (!scale)
-        return nullptr;
+        scale = model.get_variable_if_exists(scope + "/embeddings/multiply_by_sqrt_depth");
 
       StorageView value;
 
       // The attribute can either be a boolean flag or the actual scale value.
-      if (scale->dtype() == DataType::INT8 && scale->as_scalar<int8_t>())
+      if (!scale || (scale->dtype() == DataType::INT8 && scale->as_scalar<int8_t>()))
         value = StorageView(std::sqrt(static_cast<float>(embeddings.output_size())));
       else if (scale->dtype() != DataType::INT8 && scale->as_scalar<float>() != 1.f)
         value = *scale;