Dsv4 fix

flagscale/models/megatron/deepseek_v4/deepseek_model.py

@@ -206,16 +206,17 @@ def build_schedule_plan(
         """
         Adaptation of overlap_moe_expert_parallel_comm.
         """
-        # Precompute the engram_hash_iput_ids, it will be used to create a TransformerChunkSchedulePlan.
-        engram_hash_input_ids = LazyHashInputIds(
-            hash_mapping=self.engram_hash,
-            input_ids=input_ids,
-            hash_stream=self._hash_stream,
-        )
-        if extra_block_kwargs is None:
-            extra_block_kwargs = {
-                "engram_hash_input_ids": engram_hash_input_ids,
-            }
+        if self.config.use_engram:
+            # Precompute the engram_hash_iput_ids, it will be used to create a TransformerChunkSchedulePlan.
+            engram_hash_input_ids = LazyHashInputIds(
+                hash_mapping=self.engram_hash,
+                input_ids=input_ids,
+                hash_stream=self._hash_stream,
+            )
+            if extra_block_kwargs is None:
+                extra_block_kwargs = {
+                    "engram_hash_input_ids": engram_hash_input_ids,
+                }
         return super().build_schedule_plan(
             input_ids,
             position_ids,

flagscale/models/megatron/deepseek_v4/deepseek_transformer_layer.py

@@ -61,7 +61,7 @@ def _get_submodules_under_cudagraphs(self):
         """Include the DeepSeek-specific submodules in cudagraph pre-forward hooks."""
         submodules = super()._get_submodules_under_cudagraphs()
 
-        if isinstance(self.engram, IdentityOp):
+        if self.engram is None:
             return submodules
 
         try:
@@ -78,13 +78,11 @@ def forward(self, *args, **kwargs):
         self._deepseek_engram_hash_input_ids = kwargs.pop(
             "engram_hash_input_ids", getattr(self, "_deepseek_engram_hash_input_ids", None)
         )
-        self._mhc_recompute_manager = kwargs.pop("mhc_recompute_manager", None)
 
         try:
             return super().forward(*args, **kwargs)
         finally:
             self._deepseek_engram_hash_input_ids = None
-            self._mhc_recompute_manager = None
 
     def _forward_attention(
         self,
@@ -127,6 +125,7 @@ def _forward_attention(
             sequence_len_offset=sequence_len_offset,
             padding_mask=padding_mask,
             input_ids=input_ids,
+            mhc_recompute_manager=mhc_recompute_manager,
             inference_params=inference_params,
         )
 

flagscale/train/megatron/training/arguments.py

@@ -2073,6 +2073,7 @@ def _add_network_size_args(parser):
         "persist_layer_norm",
         "bias_dropout_fusion",
         "apply_rope_fusion",
+        "apply_dsa_kernel_fusion",
         # FlagScale-specific: manually added in arguments_fs.py
         "enable_hetero",
         "hetero_pipeline_layer_split",
@@ -3217,6 +3218,12 @@ def _add_mla_args(parser):
                        help="Dimension of the head in the V projection.")
     group.add_argument('--rotary-scaling-factor', type=float, default=1.0,
                        help="Rotary scaling factor for the rotary embeddings.")
+    group.add_argument(
+        '--original-max-position-embeddings',
+        type=int,
+        default=4096,
+        help="Original maximum position embeddings for the original model, used by yarn.",
+    )
     group.add_argument('--mscale', type=float, default=1.0,
                        help="Mscale for YaRN RoPE in multi-latent attention.")
     group.add_argument('--mscale-all-dim', type=float, default=0.0,

flagscale/train/megatron/training/arguments_fs.py

@@ -244,8 +244,6 @@ def pre_validate_args(self):
                 accumulated_world_size += temp_world_size
                 current_process_mesh_idx += 1
         # DeepSeek-V4 Temporary
-        if self.args.enable_hyper_connections:
-            assert not self.args.overlap_moe_expert_parallel_comm, "Hyper-connection is not supported with overlap_moe_expert_parallel_comm yet!"
         if self.args.experimental_attention_variant == "dsv4_hybrid":
             assert self.args.context_parallel_size == 1, "Context parallelism is not supported with dsv4_hybrid attention variant yet!"
 
@@ -461,10 +459,6 @@ def _parse_recompute_refined_config(recom_config, recom_config_name):
             assert (
                 self.args.use_distributed_optimizer
             ), "When use engram, distributed_optimizer must be enabled, because there is a bug caused by allreduce grad norm in model parallel group when do not use distributed_optimizer. We have not found a pretty solution yet, so disable it temporarily."
-        assert not (
-            args.pipeline_model_parallel_size == 1
-            and args.overlap_moe_expert_parallel_comm
-        ), "When no pipeline and enable overlap_moe_expert_parallel_comm, a bug will occur, it will be fixed in a later version."
 
 
 def _add_hetero_args(parser):
@@ -813,11 +807,6 @@ def _add_distributed_args(parser):
         action='store_true',
         help='Indicate whether not running on a shared file system.',
     )
-    group.add_argument(
-        '--use-padded-layerwise-optimizer',
-        action='store_true',
-        help='Enable pad when use layer-wise optimizer.'
-    )
     return parser