[Refactor] Unify TrainEngine by moving model-specific logic to model layer

Previously, we had two separate train engines:
- `TrainEngine` for regular models
- `VisionComposeTrainEngine` for vision-language models

This duplication led to:
- Code maintenance overhead (242 lines of duplicated logic)
- Tight coupling between engine and model-specific details
- Difficulty in extending to new model types

- **Remove** `VisionComposeTrainEngine` entirely (242 lines deleted)
- **Add** `pre_micro_batch_forward()` and `post_micro_batch_forward()` hooks to `BaseModel`
  - `pre_micro_batch_forward()`: Compute data batch statistics before forward pass
  - `post_micro_batch_forward()`: Aggregate micro-batch results and compute metrics
- **Unify** `TrainEngine` to handle all model types through the new hook system

- **BaseModel**:
  - Add `DataBatchInfo` and `BatchForwardInfo` TypedDicts for return types
  - Implement default `pre_micro_batch_forward()` to compute token statistics
  - Implement default `post_micro_batch_forward()` to aggregate losses and extra info
  - Add overload type hints for `__call__` to improve type inference

- **MoE Model**:
  - Override `post_micro_batch_forward()` to handle MoE-specific logic:
    - Compute maxvio for router load balancing
    - Update router bias based on expert load
  - Add `need_update_bias` property for cleaner code
  - Properly scale balancing_loss and z_loss by batch_size

- **ComposeModel**:
  - Override `pre_micro_batch_forward()` to compute image token statistics
  - Add `ComposeDataBatchInfo` with `step_consumed_img_tokens` field

- **TrainEngine**:
  - Simplify `train_step()` to delegate statistics to model hooks
  - Replace `LossLog` and `OtherLog` with unified `TrainStepInfo`
  - Add `_get_total_loss()` to aggregate all losses (with TODO for future refactor)
  - Remove all model-specific branching logic

- **EngineConfig**:
  - Remove conditional logic for VisionComposeTrainEngine
  - Use single TrainEngine.build() path

- Update to use `TrainStepInfo` instead of separate `LossLog` and `OtherLog`
- Simplify hook signatures (from 2 params to 1)
- Remove conditional engine instantiation logic

- Replace `VisionComposeTrainEngine` imports with `TrainEngine`
- Update test assertions to use new `TrainStepInfo` structure
- Remove TypeAdapter validation for deprecated types

Currently, `TrainEngine._get_total_loss()` aggregates losses by iterating
through ModelOutputs fields. This is pragmatic but not ideal:
- **Pros**: Avoids large-scale changes to model forward() logic
- **Cons**: Engine knows about loss field names (coupling)
- **Future**: Model should return total_loss directly (see TODO comment)

`loss_ctx.batch_size` represents the full gradient accumulation batch size,
not intra_layer_micro_batch. This is correctly set in `CELossContext.build_batches()`
and used for scaling balancing_loss and z_loss.

The pre/post hooks provide clean extension points:
- Subclasses can override to add model-specific logic
- Default implementations in BaseModel handle common cases
- No conditional logic needed in engine layer

1. **Code Reduction**: -242 lines (VisionComposeTrainEngine removed)
2. **Better Separation of Concerns**: Engine focuses on training orchestration, models handle their own statistics
3. **Extensibility**: New model types can override hooks without changing engine
4. **Type Safety**: Unified TrainStepInfo with clear field definitions
5. **Maintainability**: Single engine implementation to maintain

- Loss reduce logic still needs clarification (minor issue, doesn't affect training)
- TODO added for future refactor: move total_loss aggregation to model layer
- All format issues (extra blank lines, class formatting) fixed

ghstack-source-id: 9ce752d
Pull-Request: #1518

Author: HAOCHENYE

tests/engine/test_dense_train_engine.py

@@ -90,7 +90,7 @@ def warmup_fn(x):
             seq_ctx = seq_ctx_list[0]
             loss_ctx = loss_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            loss_log = engine.train_step(engine_input)["logs_info"]
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()

tests/engine/test_moe_train_engine.py

@@ -99,7 +99,7 @@ def warmup_fn(x):
             loss_ctx = loss_ctx_list[0]
             seq_ctx = seq_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            loss_log = engine.train_step(engine_input)["logs_info"]
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()
@@ -190,7 +190,7 @@ def warmup_fn(x):
             loss_ctx = loss_ctx_list[0]
             seq_ctx = seq_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            loss_log = engine.train_step(engine_input)["logs_info"]
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()

tests/engine/test_moe_train_engine_float8.py

@@ -93,7 +93,7 @@ def warmup_fn(x):
             loss_ctx = loss_ctx_list[0]
             seq_ctx = seq_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            loss_log = engine.train_step(engine_input)["logs_info"]
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()
@@ -171,7 +171,7 @@ def warmup_fn(x):
             loss_ctx = loss_ctx_list[0]
             seq_ctx = seq_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            loss_log = engine.train_step(engine_input)["logs_info"]
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()
@@ -270,11 +270,11 @@ def warmup_fn(x):
             loss_ctx = loss_ctx_list[0]
             seq_ctx = seq_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            logs_info = engine.train_step(engine_input)["logs_info"]
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()
-            losses.append(loss_log["reduced_llm_loss"])
+            losses.append(logs_info["reduced_llm_loss"])
         losses_ref = torch.tensor([2.41, 2.41, 2.47, 2.42, 2.44, 2.44, 2.42, 2.38, 2.31, 2.30])
         losses = torch.tensor(losses)
         self._check_loss_curve(losses, losses_ref)

tests/model/test_qwen3_tile_embedding.py

@@ -19,7 +19,6 @@
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.config import FSDPConfig, LRConfig, AdamWConfig
 from xtuner.v1.engine.train_engine import TrainEngine
-from xtuner.v1.engine.vision_compose_train_engine import VisionComposeTrainEngine
 from torch.optim.lr_scheduler import LambdaLR
 from xtuner.v1.utils import pad_to_max_length
 from xtuner.v1.utils.device import get_device
@@ -85,7 +84,7 @@ def warmup_fn(x):
             loss_ctx = loss_ctx_list[0]
             seq_ctx = seq_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            engine.train_step(engine_input)
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()
@@ -116,7 +115,7 @@ def test_qwen3vl_tie_embedding(self, device, tp_size):
             cpu_offload=False,
             tp_size=tp_size
         )
-        engine = VisionComposeTrainEngine(
+        engine = TrainEngine(
             model_cfg=dense_cfg, optim_cfg=optim_cfg, fsdp_cfg=fsdp_cfg
         )
         engine.from_hf(hf_path=QWEN3_VL_DENSE_PATH)
@@ -160,7 +159,7 @@ def warmup_fn(x):
             loss_ctx = loss_ctx_list[0]
             seq_ctx = seq_ctx_list[0]
             engine_input = [ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx)]
-            loss_log, _ = engine.train_step(engine_input)
+            engine.train_step(engine_input)
             grad_norm = engine.clip_grad_norm()
             engine.step_optimizer(grad_norm)
             lr_scheduler.step()

tests/train/test_trainer.py

@@ -26,7 +26,6 @@
 from xtuner.v1.datasets import FTDPTokenizeFnConfig
 from xtuner.v1.datasets.sft_tokenize_fn import OpenaiTokenizeFunctionConfig
 from xtuner.v1.train.trainer import TrainerConfig
-from xtuner.v1.engine.train_engine import LossLog, OtherLog
 from xtuner.v1.loss import CELossConfig
 from xtuner._testing import DeterministicDDPTestCase
 from unittest import TestCase
@@ -647,8 +646,6 @@ def test_hooks_config(self):
         self.create_pg(DEVICE)
         checkpoint_function_call_times = 0
         train_step_function_call_times = 0
-        losslog_adapater = TypeAdapter(LossLog)
-        otherlog_adapter = TypeAdapter(OtherLog)
 
         def checkpoint_hook(checkpoint, step, epoch, total_step, total_epoch):
             nonlocal checkpoint_function_call_times
@@ -674,9 +671,6 @@ def __init__(self) -> None:
                 self.count = 0
 
             def __call__(self, loss_log, other_log, step, epoch, total_step, total_epoch):
-                losslog_adapater.validate_python(loss_log)
-                otherlog_adapter.validate_python(other_log)
-
                 assert self.trainer().cur_step == step
                 assert self.trainer().cur_epoch == epoch
                 assert self.trainer().total_step == total_step

xtuner/v1/engine/__init__.py

@@ -1,13 +1,9 @@
 from xtuner.v1.engine.config import EngineConfig
 
-from .train_engine import LossLog, OtherLog, TrainEngine
-from .vision_compose_train_engine import VisionComposeTrainEngine
+from .train_engine import TrainEngine
 
 
 __all__ = [
     "TrainEngine",
     "EngineConfig",
-    "VisionComposeTrainEngine",
-    "LossLog",
-    "OtherLog",
 ]

xtuner/v1/engine/config.py

@@ -4,9 +4,7 @@
 
 from xtuner.v1.config import FSDPConfig, OptimConfig
 from xtuner.v1.engine.train_engine import TrainEngine
-from xtuner.v1.engine.vision_compose_train_engine import VisionComposeTrainEngine
 from xtuner.v1.model.base import BaseModel, ConfigDict
-from xtuner.v1.model.compose.base import BaseComposeConfig
 
 
 @runtime_checkable
@@ -27,7 +25,4 @@ class EngineConfig(PydanticBaseModel):
     model_cfg: ModelConfigProto
 
     def build(self):
-        if isinstance(self.model_cfg, BaseComposeConfig):
-            return VisionComposeTrainEngine(model_cfg=self.model_cfg, optim_cfg=self.optim_cfg, fsdp_cfg=self.fsdp_cfg)
-        else:
-            return TrainEngine(model_cfg=self.model_cfg, optim_cfg=self.optim_cfg, fsdp_cfg=self.fsdp_cfg)
+        return TrainEngine(model_cfg=self.model_cfg, optim_cfg=self.optim_cfg, fsdp_cfg=self.fsdp_cfg)

xtuner/v1/engine/train_engine.py

@@ -8,7 +8,6 @@
 import torch
 import torch.distributed as dist
 import torch.distributed.checkpoint as dcp
-from pydantic import ConfigDict
 from safetensors import safe_open
 from torch.distributed.checkpoint.state_dict import (
     StateDictOptions,
@@ -21,42 +20,33 @@
 from torch.utils._foreach_utils import (
     _device_has_foreach_support,
 )
-from typing_extensions import NotRequired, TypedDict
 
 from xtuner.v1.config import FSDPConfig, OptimConfig
 from xtuner.v1.data_proto.sequence_context import SequenceContext
-from xtuner.v1.model.base import BaseModel, ModelItem, XTunerBaseModelConfig
-from xtuner.v1.model.utils import ModelForwardExtraLogInfo
-from xtuner.v1.module.router import NoAuxRouterConfig
+from xtuner.v1.model.base import (
+    BaseModel,
+    BatchForwardInfo,
+    DataBatchInfo,
+    ModelItem,
+    ModelOutputs,
+    XTunerBaseModelConfig,
+)
 from xtuner.v1.profiler.prober import ProberList
 from xtuner.v1.utils import get_device, get_logger, get_torch_device_module, profile_time_and_memory
 from xtuner.v1.utils.grad_norm import cal_grad_norm
 
 
+class TrainStepInfo(DataBatchInfo, BatchForwardInfo):
+    total_loss: float
+
+
 logger = get_logger()
 DEVICE = get_device()
 DEVICE_MODULE = get_torch_device_module()
 
 threading_lock = threading.Lock()
 
 
-class LossLog(TypedDict):
-    __pydantic_config__ = ConfigDict(arbitrary_types_allowed=True)  # type: ignore[misc]
-    local_loss: float
-    reduced_llm_loss: float
-    reduced_balancing_loss: NotRequired[float]
-    reduced_z_loss: NotRequired[float]
-
-
-class OtherLog(TypedDict):
-    __pydantic_config__ = ConfigDict(arbitrary_types_allowed=True)  # type: ignore[misc]
-    maxvio: NotRequired[float]
-    step_consumed_tokens: int
-    step_consumed_img_tokens: NotRequired[int]
-    extra_info: ModelForwardExtraLogInfo
-    efficient_attn_ratio: float
-
-
 class CPUThreadTaskCoordinator:
     def __init__(self, futures, callback):
         self.futures = futures
@@ -206,66 +196,36 @@ def grad_accumulation_steps(self, data_batches_len: int):
         intra_layer_micro_batch = self.intra_layer_micro_batch
         return data_batches_len // intra_layer_micro_batch
 
-    def train_step(self, data_batches: list[ModelItem]) -> tuple[LossLog, OtherLog]:
+    def train_step(self, data_batches: list[ModelItem]) -> TrainStepInfo:
         """Perform a training step with the given data batches and mesh.
 
         Args:
             data_batches (List[Dict]): The input data batches for the training step.
         """
         self._maybe_precompute_float8_dynamic_scale_for_fsdp()
 
-        loss_log: LossLog = {}  # type: ignore[typeddict-item]
-        other_log: OtherLog = {}  # type: ignore[typeddict-item]
         intra_layer_micro_batch = self.intra_layer_micro_batch
         assert len(data_batches) % intra_layer_micro_batch == 0, (
             f"data_batches length {len(data_batches)} is not divisible by intra_layer_micro_batch {intra_layer_micro_batch}"
         )
         iters_per_step = self.grad_accumulation_steps(len(data_batches))
 
-        moe_need_update_bias = (
-            isinstance(getattr(self.model_cfg, "router", None), NoAuxRouterConfig)
-            and self.model_cfg.router.router_bias_update_speed > 0
-        )
-        moe_need_log_maxvio = getattr(self.model_cfg, "router", None) is not None
-
-        if moe_need_log_maxvio:
-            tokens_per_expert_global_for_bias = torch.zeros(
-                self.model_cfg.num_hidden_layers - self.model_cfg.first_k_dense_replace,
-                self.model_cfg.n_routed_experts,
-                dtype=torch.int64,
-                device=DEVICE,
-            )
-
-        step_loss = torch.tensor(0.0, device=DEVICE)
-        step_llm_loss = torch.tensor(0.0, device=DEVICE)
-        step_balancing_loss: torch.Tensor | None = None
-        step_z_loss: torch.Tensor | None = None
-        step_consumed_tokens = torch.tensor(0, device=DEVICE)
-
         if self._count == 0:
             logger.info(f"grad_accumulation_steps: {iters_per_step}")
             self._count += 1
 
-        train_engine_extra_info = ModelForwardExtraLogInfo()
         micro_batch_iter = 0
-        efficient_forward_tokens = torch.tensor(0, device=DEVICE, dtype=torch.long)
-        total_forward_tokens = torch.tensor(0, device=DEVICE, dtype=torch.long)
+        micro_batch_results = []
+
+        data_batch_info = self.model.pre_micro_batch_forward(data_batches)
+        total_loss = torch.tensor(0.0, device=DEVICE)
+
         for i in range(0, len(data_batches), intra_layer_micro_batch):
             ProberList.set_micro_batch_iter(micro_batch_iter)
             micro_batch_iter += 1
             data_batch = data_batches[i : i + intra_layer_micro_batch]
-            seq_ctx_list = []
-            loss_ctx_list = []
-            for data in data_batch:
-                seq_ctx = data["seq_ctx"]
-                loss_ctx = data["loss_ctx"]
-                seq_ctx_list.append(seq_ctx)
-                loss_ctx_list.append(loss_ctx)
-                step_consumed_tokens += seq_ctx.mask.sum()
-
-                num_tokens = seq_ctx.cu_seq_lens_k[1:] - seq_ctx.cu_seq_lens_k[:-1]
-                efficient_forward_tokens += (num_tokens.long() ** 2).sum()
-                total_forward_tokens += (num_tokens.long().sum()) ** 2
+            seq_ctx_list = [i["seq_ctx"] for i in data_batch]
+            loss_ctx_list = [i["loss_ctx"] for i in data_batch]
 
             if self.intra_layer_micro_batch == 1:
                 output = self.model(seq_ctx=seq_ctx_list[0], loss_ctx=loss_ctx_list[0])
@@ -278,78 +238,16 @@ def train_step(self, data_batches: list[ModelItem]) -> tuple[LossLog, OtherLog]:
                 )
             output.free_nongrad_feature()
 
-            # llm loss has been global averaged
-            llm_loss = output["loss"]
-            step_llm_loss += llm_loss.detach().clone()
-
-            loss = llm_loss
-            if "extra_info" in output:
-                train_engine_extra_info.append(output["extra_info"])
-
-            if "balancing_loss" in output:
-                balancing_loss = output["balancing_loss"] / iters_per_step
-                loss = loss + balancing_loss
-                if step_balancing_loss is None:
-                    step_balancing_loss = balancing_loss
-                else:
-                    step_balancing_loss += balancing_loss
-
-            if "z_loss" in output:
-                z_loss = output["z_loss"] / iters_per_step
-                loss = loss + z_loss
+            micro_batch_results.append(output)
 
-                if step_z_loss is None:
-                    step_z_loss = z_loss
-                else:
-                    step_z_loss += z_loss
-
-            if moe_need_log_maxvio:
-                assert "tokens_per_expert_global" in output, "tokens_per_expert_global is required for bias update."
-                tokens_per_expert_global_for_bias += output["tokens_per_expert_global"]
-
-            del output
+            loss = self._get_total_loss(output)
             loss.backward()
+            total_loss += loss.detach()
             # call dump_forward_records after backward to record the recomputed activations
             ProberList.after_micro_iter_forward()
-            step_loss += loss.detach().clone()
-
-        if moe_need_log_maxvio:
-            avg_count_load = tokens_per_expert_global_for_bias.float().mean(1)
-            max_load_i, _ = torch.max(tokens_per_expert_global_for_bias, dim=1)
-            maxvio_all_layers = (max_load_i - avg_count_load) / avg_count_load
-            maxvio = maxvio_all_layers.mean()
-            if moe_need_update_bias:
-                self.model.update_bias(tokens_per_expert_global_for_bias, avg_count_load)  # type: ignore
-            other_log["maxvio"] = maxvio.item()
-
-        reduced_llm_loss = step_llm_loss
-        dist.all_reduce(reduced_llm_loss.div_(dist.get_world_size()))
-
-        loss_log["local_loss"] = step_loss.item()
-        loss_log["reduced_llm_loss"] = reduced_llm_loss.item()
-        if step_balancing_loss is not None:
-            reduced_balancing_loss = step_balancing_loss
-            dist.all_reduce(reduced_balancing_loss.div_(dist.get_world_size()))
-            loss_log["reduced_balancing_loss"] = reduced_balancing_loss.item()
-        if step_z_loss is not None:
-            reduced_z_loss = step_z_loss
-            dist.all_reduce(reduced_z_loss.div_(dist.get_world_size()))
-            loss_log["reduced_z_loss"] = reduced_z_loss.item()
-        other_log["step_consumed_tokens"] = int(step_consumed_tokens.item())
-        other_log["extra_info"] = train_engine_extra_info
-        other_log["efficient_attn_ratio"] = (efficient_forward_tokens / total_forward_tokens).item()
-
-        extra_info = other_log.get("extra_info", {})  # type: ignore
-
-        # TODO: @duanyanhui `extra_info` should be redesigned.
-        if not isinstance(extra_info, ModelForwardExtraLogInfo):
-            extra_info = ModelForwardExtraLogInfo(extra_info)
-        loss_log.update(extra_info.get())
-
-        if "maxvio" in other_log:
-            loss_log["maxvio"] = other_log["maxvio"]  # type: ignore
-        loss_log["efficient_attn_ratio"] = other_log["efficient_attn_ratio"]  # type: ignore
-        return loss_log, other_log
+
+        batch_forward_info = self.model.post_micro_batch_forward(micro_batch_results)
+        return TrainStepInfo(total_loss=total_loss.item(), **data_batch_info, **batch_forward_info)
 
     def from_hf(self, hf_path: str | Path, strict: bool = False):
         self.model.from_hf(hf_path=hf_path, strict=strict)
@@ -529,3 +427,17 @@ def _maybe_precompute_float8_dynamic_scale_for_fsdp(self):
         for model in self.model.modules():
             if isinstance(model, BaseModel) and model.float8_handler is not None:
                 model.float8_handler.precompute_float8_dynamic_scale_for_fsdp(model)
+
+    def _get_total_loss(self, model_outputs: ModelOutputs) -> torch.Tensor:
+        # TODO: This logic should be moved into the model layer. The model should be responsible
+        # for aggregating all losses (CE loss, balancing loss, z loss, etc.) and returning a
+        # single total_loss. The engine should only call model.forward() and use the returned
+        # total_loss directly, rather than iterating through fields to sum losses here.
+        # This would provide better separation of concerns and make the loss computation logic
+        # more explicit and maintainable.
+        loss = torch.tensor(0.0, device=DEVICE)
+        for key in model_outputs.model_fields:
+            value = getattr(model_outputs, key)
+            if "loss" in key and isinstance(value, torch.Tensor):
+                loss += value
+        return loss