fix(engine): guard megatron-bridge on NPU and restore bridge_type field

Combined megatron-bridge compatibility fixes for NPU:

- Make megatron.bridge imports defensive in megatron_engine.py and
  megatron_utils/megatron_lora.py (the latter runs a module-level monkey
  patch at import time) so modules load on NPU where the CUDA-only
  megatron-bridge package is unavailable.
- Restore the bridge_type field on MegatronEngineConfig (an earlier rebase
  silently landed it inside MindSpeedEngineConfig after as_dict, which
  made MegatronEngineConfig(bridge_type=...) raise TypeError).

Author: Adiactive

areal/api/cli_args.py

@@ -921,6 +921,22 @@ class MegatronEngineConfig:
     account_for_embedding_in_pipeline_split: bool = field(default=False)
     account_for_loss_in_pipeline_split: bool = field(default=False)
 
+    # Bridge backend used for HF<->Megatron conversion/model creation.
+    bridge_type: str = field(
+        default="mbridge",
+        metadata={
+            "help": "Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'.",
+            "choices": ["mbridge", "megatron-bridge"],
+        },
+    )
+
+    use_mbridge_save: bool = field(
+        default=False,
+        metadata={
+            "help": "Use mbridge's save method to save gpu memory when saving weights."
+        },
+    )
+
 
 @dataclass
 class MindSpeedEngineConfig:
@@ -976,22 +992,6 @@ class MindSpeedEngineConfig:
     def as_dict(self) -> dict[str, Any]:
         return asdict(self)
 
-    # Bridge backend used for HF<->Megatron conversion/model creation.
-    bridge_type: str = field(
-        default="mbridge",
-        metadata={
-            "help": "Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'.",
-            "choices": ["mbridge", "megatron-bridge"],
-        },
-    )
-
-    use_mbridge_save: bool = field(
-        default=False,
-        metadata={
-            "help": "Use mbridge's save method to save gpu memory when saving weights."
-        },
-    )
-
 
 class SchedulingStrategyType(str, Enum):
     separation = "separation"

areal/engine/megatron_engine.py

@@ -44,8 +44,6 @@ def _ms_sanitized_get_full_args():
 import mbridge
 import torch
 import torch.distributed as dist
-from megatron.bridge import AutoBridge as MegatronBridgeAutoBridge
-from megatron.bridge.peft.lora import LoRA as MegatronBridgeLoRA
 from megatron.core import parallel_state as mpu
 from megatron.core import tensor_parallel
 from megatron.core.distributed import DistributedDataParallel as DDP
@@ -109,7 +107,7 @@ def _ms_sanitized_get_full_args():
     configure_pipeline_layer_splits,
 )
 from areal.infra.dist_rollout import DistRolloutCoordinator
-from areal.infra.platforms import current_platform
+from areal.infra.platforms import current_platform, is_npu_available
 from areal.models.mcore.hf_load import load_weights_from_hf_with_mbridge_fast
 from areal.models.mcore.hf_save import (
     save_critic_value_head,
@@ -268,6 +266,12 @@ def __init__(self, config: TrainEngineConfig):
         )
         self.quantization_config: dict[str, int | str | list[str]] | None = None
         self.bridge_cls: str = getattr(self.mcore_config, "bridge_type", "mbridge")
+        if self.bridge_cls == "megatron-bridge" and is_npu_available:
+            raise ValueError(
+                "bridge_type='megatron-bridge' is not supported on NPU "
+                "(megatron-bridge has CUDA-only dependencies). "
+                "Use bridge_type='mbridge' instead."
+            )
         self.bridge_lora: MegatronBridgeLoRA | None = None
         self.is_vision_model: bool = False
         self.processor = None
@@ -335,6 +339,8 @@ def _apply_megatron_bridge_lora(self) -> None:
                 "linear_fc1",
                 "linear_fc2",
             ]
+        from megatron.bridge.peft.lora import LoRA as MegatronBridgeLoRA
+
         self.bridge_lora = MegatronBridgeLoRA(
             target_modules=target_modules,
             dim=self.config.lora_rank,
@@ -619,6 +625,8 @@ def _build_hf_mcore_bridge(self):
                 raise NotImplementedError(
                     "Tree training is not supported with bridge_type='megatron-bridge'."
                 )
+            from megatron.bridge import AutoBridge as MegatronBridgeAutoBridge
+
             self.bridge = MegatronBridgeAutoBridge.from_hf_pretrained(
                 self.config.path,
                 trust_remote_code=True,

areal/engine/megatron_utils/megatron_lora.py

@@ -277,7 +277,11 @@ def _build_adapter_config_dict(
 
 def _monkey_patch_save_hf_adapter():
     """Add save_hf_adapter to AutoBridge when megatron-bridge does not provide it."""
-    from megatron.bridge import AutoBridge
+    try:
+        from megatron.bridge import AutoBridge
+    except ImportError:
+        # megatron-bridge is not installed (e.g. NPU environment); nothing to patch.
+        return
 
     if hasattr(AutoBridge, "save_hf_adapter"):
         # Already exists, no need to patch

docs/en/cli_reference.md

@@ -953,6 +953,7 @@ Refer to Megatron-LM documentation for implementation details.
 | `num_layers_in_last_pipeline_stage`        | integer \| None                                                      | `None`       | Number of layers in the last pipeline stage                                                                                                                                                       |
 | `account_for_embedding_in_pipeline_split`  | boolean                                                              | `False`      | -                                                                                                                                                                                                 |
 | `account_for_loss_in_pipeline_split`       | boolean                                                              | `False`      | -                                                                                                                                                                                                 |
+| `bridge_type`                              | string                                                               | `"mbridge"`  | Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'. **Choices:** `mbridge`, `megatron-bridge`                                                                             |
 
 (section-mind-speed-engine)=
 

docs/zh/cli_reference.md

@@ -951,6 +951,7 @@ Refer to Megatron-LM documentation for implementation details.
 | `num_layers_in_last_pipeline_stage`        | integer \| None                                                      | `None`       | Number of layers in the last pipeline stage                                                                                                                                                       |
 | `account_for_embedding_in_pipeline_split`  | boolean                                                              | `False`      | -                                                                                                                                                                                                 |
 | `account_for_loss_in_pipeline_split`       | boolean                                                              | `False`      | -                                                                                                                                                                                                 |
+| `bridge_type`                              | string                                                               | `"mbridge"`  | Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'. **Choices:** `mbridge`, `megatron-bridge`                                                                             |
 
 (section-mind-speed-engine)=