Fix inference issues FP8-trained models

Signed-off-by: Jan Lasek <[email protected]>
NVIDIA · Oct 1, 2024 · 56d6e6f · 56d6e6f
1 parent 86408cc
commit 56d6e6f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemo/deploy/nlp/megatronllm_deployable.py b/nemo/deploy/nlp/megatronllm_deployable.py
@@ -157,6 +157,8 @@ def _load_from_nemo_checkpoint(self, nemo_checkpoint_filepath: str, num_devices:
             custom_config.activations_checkpoint_granularity = None
             custom_config.activations_checkpoint_method = None
             custom_config.dist_ckpt_load_strictness = StrictHandling.LOG_ALL.value
+            if custom_config.get("fp8", False):
+                custom_config.fp8 = False
 
             self.model = MegatronGPTModel.restore_from(
                 nemo_checkpoint_filepath, trainer=trainer, override_config_path=custom_config