[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: ba87e63
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,14 +26,14 @@
 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
 from xtuner.v1.train.trainer import XTunerMeta, ExpInfo, ExpHistory, GitInfo
 from xtuner.v1.utils.device import get_device
 from xtuner.v1.datasets.dataloader import Dataloader
 from torch.optim.lr_scheduler import SequentialLR
+from xtuner.v1.model.utils.misc import ModelForwardExtraLogInfo
 
 
 DEVICE = get_device()
@@ -55,10 +55,8 @@ def grad_accumulation_steps(self, *args, **kwargs):
 
     def train_step(self, *args, **kwargs):
         self.train_step_calls += 1
-        return (
-            {"local_loss": 1.0, "reduced_llm_loss": 0.8},
-            {"step_consumed_tokens": 100, "grad_norm": torch.tensor(1.0), "efficient_attn_ratio": 0.5}
-        )
+        return {"total_loss": 1.8, "step_consumed_tokens": 100, "grad_norm": torch.tensor(1.0), "efficient_attn_ratio": 0.5, "logs_info": {"local_loss": 1.0, "reduced_llm_loss": 0.8}, "extra_info": ModelForwardExtraLogInfo()}
+        
 
     def save_hf(self, hf_path):
         self.save_hf_calls.append(hf_path)
@@ -647,14 +645,12 @@ 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
             checkpoint_function_call_times += 1
 
-        def train_step_hook(loss_log, other_log, step, epoch, total_step, total_epoch):
+        def train_step_hook(train_step_info, step, epoch, total_step, total_epoch):
             nonlocal train_step_function_call_times
             train_step_function_call_times += 1
 
@@ -673,10 +669,7 @@ def connect_trainer(self, trainer: Trainer):
             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)
-
+            def __call__(self, train_step_info, step, epoch, total_step, total_epoch):
                 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)