fix-loader-fp16

libai/models/utils/model_loader/base_loader.py

@@ -23,7 +23,7 @@
 from termcolor import colored
 
 import libai.utils.distributed as dist
-from libai.config import LazyCall
+from libai.config import LazyCall, try_get_key
 from libai.models.build import build_model
 
 logger = logging.getLogger(__name__)
@@ -389,8 +389,7 @@ def _convert_tensor(self, tensor):
         Returns:
             flow.Tensor: The target tensor.
         """
-        tensor = tensor.float()
-        return flow.Tensor(tensor.detach().cpu().numpy())
+        return flow.tensor(tensor.detach().cpu().numpy())
 
     def _convert_tensors(self, torch_state_dict):
 
@@ -578,6 +577,10 @@ def load(self):
         else:
             self.model = build_model(LazyCall(self.model)(cfg=self.libai_cfg))
 
+        # Convert to fp16
+        if try_get_key(self.libai_cfg, "amp_enabled"):
+            self.model.half()
+
         # State_dict to global
         logger.info("transfering state_dict local to global...")
         flow_state_dict = self._state_dict_to_global(flow_state_dict, mode="pytorch")

projects/GLM/infer_glm.py

@@ -43,7 +43,7 @@
     output_dropout_prob=0,
 )
 model = loader.load()
-model = model.half().cuda()
+model = model.cuda()
 model.eval()
 
 dist.set_device_type("cuda")