Fix get device type error (#60)

Author: lxd-cumt

megatron/core/distributed/finalize_model_grads.py

@@ -24,7 +24,16 @@
 from megatron.plugin.platform import get_platform
 
 cur_platform = get_platform()
-# FlagScale End
+
+
+try:
+    from megatron.plugin.utils import get_device_type_for_comm
+except ImportError:
+
+    def get_device_type_for_comm(group):
+        """Fallback: return current platform device name for communication."""
+        return cur_platform.device_name()
+
 
 from .. import parallel_state
 from ..transformer.moe.moe_utils import get_updated_expert_bias