[model] Support deepseek-v4

README.md

@@ -127,7 +127,7 @@ The following is the list of models supported by MCore-Bridge:
 | Series     | model_type                                            |
 | -------- | ------------------------------------------------------------ |
 | Qwen     | qwen2, qwen2_moe<br />qwen3, qwen3_moe, qwen3_next |
-| DeepSeek | deepseek_v3, deepseek_v32                                    |
+| DeepSeek | deepseek_v3, deepseek_v32, deepseek_v4                        |
 | GLM      | glm4, glm4_moe, glm4_moe_lite<br />glm_moe_dsa |
 | MiniMax  | minimax_m2                                                   |
 | Kimi     | kimi_k2, kimi_k25                                   |

README_zh.md

@@ -123,7 +123,7 @@ uv pip install -e . --torch-backend=auto
 | 系列     | model_type                                                   |
 | -------- | ------------------------------------------------------------ |
 | Qwen     | qwen2, qwen2_moe<br />qwen3, qwen3_moe, qwen3_next |
-| DeepSeek | deepseek_v3, deepseek_v32                                    |
+| DeepSeek | deepseek_v3, deepseek_v32, deepseek_v4                     |
 | GLM      | glm4, glm4_moe, glm4_moe_lite<br />glm_moe_dsa |
 | MiniMax  | minimax_m2                                                   |
 | Kimi     | kimi_k2, kimi_k25                                   |

src/mcore_bridge/model/modules/transformer_layer.py

@@ -191,13 +191,19 @@ def can_recompute_pre_mlp_layernorm_for_cudagraph():
             if 'mlp' in self.config.recompute_modules:
                 if not self.is_moe_layer:
                     self.recompute_mlp = True
-        if hasattr(self.config, 'fine_grained_activation_offloading'):
-            self.offload_attn_norm = (
-                self.config.fine_grained_activation_offloading and 'attn_norm' in self.config.offload_modules
-                and not isinstance(self.input_layernorm, IdentityOp))
-            self.offload_mlp_norm = (
-                self.config.fine_grained_activation_offloading and 'mlp_norm' in self.config.offload_modules
-                and not isinstance(self.pre_mlp_layernorm, IdentityOp))
+        if hasattr(self, '_set_offload_modules'):
+            from megatron.core.transformer.transformer_layer import _get_offloading_interface
+            self._set_offload_modules()
+            self.off_interface = _get_offloading_interface()
+            self.mlp_norm_manager = None
+        else:
+            if hasattr(self.config, 'fine_grained_activation_offloading'):
+                self.offload_attn_norm = (
+                    self.config.fine_grained_activation_offloading and 'attn_norm' in self.config.offload_modules
+                    and not isinstance(self.input_layernorm, IdentityOp))
+                self.offload_mlp_norm = (
+                    self.config.fine_grained_activation_offloading and 'mlp_norm' in self.config.offload_modules
+                    and not isinstance(self.pre_mlp_layernorm, IdentityOp))
 
         # @jcasper how should we handle nvfuser?
         # Set bias+dropout+add fusion grad_enable execution handler.