fix for latest mcore

maanug-nv · maanug-nv · commit 75c5fe3e60cb · 2025-03-18T15:52:37.000-07:00
Signed-off-by: Maanu Grover &lt;maanug@nvidia.com&gt;
diff --git a/nemo/tron/checkpointing.py b/nemo/tron/checkpointing.py
@@ -40,9 +40,9 @@
     FullyParallelLoadStrategyWrapper,
     FullyParallelSaveStrategyWrapper,
 )
+from megatron.core.fp8_utils import is_float8tensor
 from megatron.core.num_microbatches_calculator import update_num_microbatches
 from megatron.core.rerun_state_machine import get_rerun_state_machine
-from megatron.core.utils import is_float8tensor
 
 from nemo.tron import fault_tolerance
 from nemo.tron.config import ConfigContainer
diff --git a/nemo/tron/model.py b/nemo/tron/model.py
@@ -20,8 +20,8 @@
     TorchFullyShardedDataParallel,
 )
 from megatron.core.enums import ModelType
+from megatron.core.fp8_utils import is_float8tensor
 from megatron.core.transformer.module import Float16Module
-from megatron.core.utils import is_float8tensor
 
 from nemo.collections.llm.gpt.model.base import GPTConfig
 from nemo.collections.llm.t5.model.t5 import T5Config

Original file line number	Diff line number	Diff line change
`@@ -20,8 +20,8 @@`
`20`	`20`	`TorchFullyShardedDataParallel,`
`21`	`21`	`)`
`22`	`22`	`from megatron.core.enums import ModelType`
	`23`	`+from megatron.core.fp8_utils import is_float8tensor`
`23`	`24`	`from megatron.core.transformer.module import Float16Module`
`24`		`-from megatron.core.utils import is_float8tensor`
`25`	`25`
`26`	`26`	`from nemo.collections.llm.gpt.model.base import GPTConfig`
`27`	`27`	`from nemo.collections.llm.t5.model.t5 import T5Config`