benchmark

Author: HAOCHENYE

xtuner/v1/float8/fsdp_utils.py

@@ -22,7 +22,6 @@
 DEVICE_MODULE = get_torch_device_module()
 
 
-@maybe_compile(fullgraph=True)
 def tensor_to_per_block_fp8_devided_64_scales(
     tensor: "WeightWithDynamicTilewiseFloat8CastTensor",
     reduce_mesh_devided_64: Optional[DeviceMesh] = None,
@@ -224,7 +223,6 @@ def cast_to_per_block_fp8_with_scales(
     return tensor_bits_fp8
 
 
-@maybe_compile(fullgraph=True)
 def cast_to_per_block_fp8_devided_64_with_scales(
     tensor: torch.Tensor,
     scales: torch.Tensor,

xtuner/v1/model/moe/deepseek_v3.py

@@ -6,6 +6,7 @@
 from transformers.models.deepseek_v3 import DeepseekV3Config as HFDeepseekV3Config
 from xtuner.v1.model.moe.moe import BalancingLossConfig, MoEConfig, ZLossConfig
 from xtuner.v1.module.attention import MLAConfig
+from xtuner.v1.module.router.greedy import GreedyRouterConfig
 from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig
 from xtuner.v1.utils import get_logger
 
@@ -89,7 +90,7 @@ class DeepSeekV3Config(MoEConfig):
     num_experts_per_tok: int = 8
     hidden_factor: float = 1.0
     moe_intermediate_size: int = 2048
-    router: NoAuxRouterConfig = NoAuxRouterConfig(
+    router: NoAuxRouterConfig | GreedyRouterConfig = NoAuxRouterConfig(
         n_group=8,
         topk_group=4,
         scoring_func="sigmoid",

xtuner/v1/model/moe/moe.py

@@ -305,6 +305,7 @@ def _micro_batch_forward(
         cat_position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None
         cat_hidden_states: torch.Tensor | None = None
 
+        moe_forawrd = False
         for idx, decoder_layer in self.layers.items():
             layer_idx = int(idx)
 
@@ -322,8 +323,9 @@ def _micro_batch_forward(
                     seq_ctx=cat_seq_ctx,
                 )
             else:
-                if cat_hidden_states is not None:
+                if cat_hidden_states is not None and not moe_forawrd:
                     hidden_states_list = list(cat_hidden_states.chunk(len(seq_ctx_list), dim=1))
+                    moe_forawrd = True
 
                 layer_results = decoder_layer(
                     *hidden_states_list,

xtuner/v1/module/dispatcher/deepep.py

@@ -60,6 +60,7 @@ class DeepEPPreCombineResult(PreCombineResult):
 
 class DeepEPCombineResult(CombineResult):
     forward_finished_event: EventOverlap | None
+    backward_previous_event: EventOverlap | None
 
 
 DeepEPPostCombineResult = PostCombineResult
@@ -144,6 +145,7 @@ def forward(
         handle: DeepEPHandle,
         group: dist.ProcessGroup,
         forward_previous_event: EventOverlap | None = None,
+        backward_previous_event: EventOverlap | None = None,
         backward_finished_event: EventOverlap | None = None,
     ) -> tuple[torch.Tensor, EventOverlap]:
         combined_x, event = combine_forward(x, num_experts, handle, group, forward_previous_event)
@@ -152,17 +154,18 @@ def forward(
         ctx.group = group
         ctx.num_experts = num_experts
         ctx.backward_finished_event = backward_finished_event
+        ctx.backward_previous_event = backward_previous_event
         return combined_x, event
 
     @staticmethod
     def backward(  # type: ignore[invalid-override]
         ctx, grad_combined_x: torch.Tensor, *args
-    ) -> tuple[torch.Tensor | tuple[torch.Tensor, torch.Tensor], None, None, None, None, None]:
+    ) -> tuple[torch.Tensor | tuple[torch.Tensor, torch.Tensor], None, None, None, None, None, None]:
         # load saved comm handle
         handle = ctx.saved_tensors
-        grad_x, event = combine_backward(grad_combined_x, ctx.num_experts, handle, ctx.group, buffer_capture())
+        grad_x, event = combine_backward(grad_combined_x, ctx.num_experts, handle, ctx.group, ctx.backward_previous_event)
         ctx.backward_finished_event.event = event.event
-        return grad_x, None, None, None, None, None
+        return grad_x, None, None, None, None, None, None
 
 
 _async_combine = copy_method_signature(DeepEPCombine.forward)(DeepEPCombine.apply)
@@ -388,24 +391,30 @@ def combine(
         decoding: bool = False,
     ) -> CombineResult:
         if async_op:
+            backward_previous_event = EventOverlap(None)
             assert pre_combined["forward_finished_event"] is not None, "Please use `async_op=True` for combine!"
             pre_combined["forward_finished_event"].current_stream_wait()
+        else:
+            backward_previous_event = None
 
         combined_hidden_states, event = _async_combine(
             pre_combined["hidden_states"],
             self._n_routed_experts,
             dispatched["handle"],
             self._process_group,
             pre_combined["forward_finished_event"],
+            backward_previous_event,
             pre_combined["backward_previous_event"],
         )
         if not async_op:
             event.current_stream_wait()
 
+
         if not decoding:
             return DeepEPCombineResult(
                 hidden_states=combined_hidden_states,
                 forward_finished_event=event,
+                backward_previous_event=backward_previous_event,
             )
         else:
             raise NotImplementedError
@@ -424,6 +433,17 @@ def combine_postprocess(
         hidden_states = combined["hidden_states"]
         forward_previous_event = combined["forward_finished_event"]
 
+        hidden_states = hidden_states.view_as(hidden_states)
+
+        if hidden_states.grad_fn is not None:
+            hidden_states.grad_fn.register_hook(
+                get_backward_hook(
+                    backward_finished_event=combined["backward_previous_event"],
+                    name="DeeEPDispatcher.combine_postprocess",
+                    debug=XTUNER_DISPATCHER_DEBUG,
+                )
+            )
+
         if async_op:
             assert forward_previous_event is not None, "Please use `async_op=True` for combine!"
             forward_previous_event.current_stream_wait()

xtuner/v1/module/router/noaux_router.py

@@ -59,9 +59,6 @@ def __init__(
         )
 
     def forward(self, logits) -> RouterResults:
-        if os.getenv("XTUNER_ROUTER_DEBUG") == "true":
-            noise = torch.randn_like(logits) * 50
-            logits = logits + noise
 
         if self.scoring_func == "sigmoid":
             scores = logits.sigmoid()
@@ -71,6 +68,10 @@ def forward(self, logits) -> RouterResults:
 
         scores_for_choice = scores + self.e_score_correction_bias.unsqueeze(0)
 
+        if os.getenv("XTUNER_ROUTER_DEBUG") == "true":
+            noise = torch.randn_like(scores) * 50
+            scores_for_choice= scores + noise
+
         # select top-k experts
         # (only applicable when ep_size >= 64. when ep_size=32 (4 nodes), there is no need to employ this strategy)
         _, topk_idx = torch.topk(scores_for_choice, k=self.top_k, dim=-1)

xtuner/v1/ops/comm/deepep_op.py

@@ -20,7 +20,8 @@
 _low_latency_buffer: Optional[Buffer] = None
 # Set the number of SMs to use
 # NOTES: this is a static variable
-Buffer.set_num_sms(24)
+# Buffer.set_num_sms(24)
+Buffer.set_num_sms(20)
 
 
 # You may call this function at the framework initialization

xtuner/v1/train/trainer.py

@@ -43,6 +43,7 @@
     record_git_info,
 )
 from xtuner.v1.utils.device import get_device, get_torch_device_module
+import gc
 
 from .toy_tokenizer import UTF8ByteTokenizer
 
@@ -142,7 +143,7 @@ class TrainerConfig(BaseModel):
     hf_interval: int | None = None
     hf_max_keep: int | None = None
     exp_tracker: Literal["tensorboard", "jsonl"] = "jsonl"
-    profile_step: int | None = None
+    profile_step: list | int | None = None
     profile_time: bool = True
     profile_memory: bool = False
     intra_layer_micro_batch: int = 1
@@ -237,7 +238,7 @@ def __init__(
         hf_interval: int | None = None,
         hf_max_keep: int | None = None,
         exp_tracker: Literal["tensorboard", "jsonl"] = "jsonl",
-        profile_step: int | None = None,
+        profile_step: list | None = None,
         profile_time: bool = True,
         profile_memory: bool = False,
         intra_layer_micro_batch: int = 1,
@@ -257,6 +258,8 @@ def __init__(
 
         self._micro_batch_size: int | None = None
 
+        if type(profile_step) is int:
+            profile_step = [profile_step]
         self._profile_step = profile_step
         self._profile_time = profile_time
         self._profile_memory = profile_memory
@@ -492,6 +495,9 @@ def fit(self):
 
             time_before_get_data = time.time()
 
+            if self.cur_step % 50 == 0:
+                gc.collect()
+
     @property
     def world_size(self) -> int:
         """Get the total number of processes in the distributed training group.
@@ -944,7 +950,7 @@ def _init_xtuner_meta(self, work_dir: Path, auto_resume: bool) -> XTunerMeta:
     @contextmanager
     def _maybe_profiling(self):
         """Check if profiling is enabled and perform profiling if necessary."""
-        if self._profile_step is not None and self._cur_step == self._profile_step:
+        if self._profile_step is not None and self._cur_step in self._profile_step:
             with contextlib.ExitStack() as stack:
                 if self._profile_time:
                     time_dir = self.exp_dir / self._PROFILE_TIME_PATH / f"step-{self._cur_step}"
@@ -1177,6 +1183,7 @@ def _resume_dataloader(self, dataloader_path: Path):
         self._dataloader.load_state_dict(dataloader_state)
 
     def _setup_env(self):
+        gc.disable()
         os.environ["TOKENIZERS_PARALLELISM"] = "true"
 
         log_str = "\n============XTuner Training Environment============\n"