enable bnb 4/8bit inference (#2419)

onmt/model_builder.py

@@ -94,6 +94,9 @@ def load_test_model(opt, model_path=None):
 
     model_opt = ArgumentParser.ckpt_model_opts(checkpoint["opt"])
 
+    model_opt.quant_layers = opt.quant_layers
+    model_opt.quant_type = opt.quant_type
+
     ArgumentParser.update_model_opts(model_opt)
     ArgumentParser.validate_model_opts(model_opt)
     vocabs = dict_to_vocabs(checkpoint["vocab"])
@@ -118,6 +121,7 @@ def load_test_model(opt, model_path=None):
     else:
         device = torch.device("cpu")
 
+    logger.info("Loading data into the model")
     if "model" in checkpoint.keys():
         # weights are in the .pt file
         model.load_state_dict(

onmt/opts.py

@@ -1133,24 +1133,6 @@ def _add_train_general_opts(parser):
         help="rule of thumb: same value as in main model",
     )
 
-    group.add(
-        "--quant_layers",
-        "-quant_layers",
-        default=[],
-        nargs="+",
-        type=str,
-        help="list of layers to be compressed in 4/8bit.",
-    )
-
-    group.add(
-        "--quant_type",
-        "-quant_type",
-        default="bnb_8bit",
-        choices=["bnb_8bit", "bnb_FP4", "bnb_NF4"],
-        type=str,
-        help="Type of compression.",
-    )
-
     _add_reproducibility_opts(parser)
 
     # Init options
@@ -1533,6 +1515,27 @@ def _add_train_dynamic_data(parser):
     )
 
 
+def _add_quant_opts(parser):
+    group = parser.add_argument_group("Quant options")
+    group.add(
+        "--quant_layers",
+        "-quant_layers",
+        default=[],
+        nargs="+",
+        type=str,
+        help="list of layers to be compressed in 4/8bit.",
+    )
+
+    group.add(
+        "--quant_type",
+        "-quant_type",
+        default="bnb_8bit",
+        choices=["bnb_8bit", "bnb_FP4", "bnb_NF4"],
+        type=str,
+        help="Type of compression.",
+    )
+
+
 def train_opts(parser):
     """All options used in train."""
     # options relate to data preprare
@@ -1541,6 +1544,7 @@ def train_opts(parser):
     model_opts(parser)
     _add_train_general_opts(parser)
     _add_train_dynamic_data(parser)
+    _add_quant_opts(parser)
 
 
 def _add_decoding_opts(parser):
@@ -1828,6 +1832,8 @@ def translate_opts(parser, dynamic=False):
         # Adding options related to Transforms
         _add_dynamic_transform_opts(parser)
 
+    _add_quant_opts(parser)
+
 
 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be