diff --git a/src/torch4ms/autograd/forward_extractor.py b/src/torch4ms/autograd/forward_extractor.py
index b76d17b09..0e2ce1c08 100644
--- a/src/torch4ms/autograd/forward_extractor.py
+++ b/src/torch4ms/autograd/forward_extractor.py
@@ -6,7 +6,6 @@
 """
 from typing import Callable, Dict, List, Tuple
 import torch
-from torch.nn.utils import stateless as torch_stateless
 from torch4ms.tensor import Tensor as Torch4msTensor
 from mindspore import Tensor as ms_Tensor
 
@@ -129,58 +128,22 @@ def get_ms_forward_fn(self, env, include_buffers: bool = False) -> Callable:
         """
         获取用于 MindSpore GradOperation 的 forward 函数。
 
-        输入为 MindSpore Tensor，内部会转换为 torch.Tensor 执行前向（因为 functional_call 需要），
-        并将结果转换回 MindSpore Tensor。
+        输入为 MindSpore Tensor，内部先转为 torch4ms.Tensor 视图，再在 env 下执行
+        functional forward。这样模块内部算子仍通过 torch4ms dispatch 路由到 MindSpore，
+        避免退化为纯 PyTorch 计算图。
         """
         base_forward = self.get_forward_fn(include_buffers=include_buffers)
 
         def forward_fn(*ms_args):
-            # 将 MindSpore Tensor 转为 torch.Tensor（functional_call 需要）
-            # 先转为 torch4ms.Tensor，再转为 torch.Tensor
+            # 仅做同构包装：MindSpore Tensor -> torch4ms.Tensor（不复制，不降级为纯 torch.Tensor）
             t4_args = env.ms2t_iso(ms_args)
-            torch_args = []
-            for arg in t4_args:
-                if isinstance(arg, Torch4msTensor):
-                    # 转换为 torch.Tensor（用于 functional_call）
-                    from torch4ms.ops import mappings
-                    torch_args.append(mappings.ms2t(arg._elem))
-                else:
-                    torch_args.append(arg)
-            
-            # 调用 base_forward（它期望 torch.Tensor）
+
+            # 在 env 下调用 base_forward，让模块内部 op 继续走 torch4ms dispatch
             with env:
-                res = base_forward(*torch_args)
-            
-            # 将结果转换回 MindSpore Tensor
-            if isinstance(res, Torch4msTensor):
-                # 如果已经是 torch4ms.Tensor，直接返回其 _elem
-                return res._elem
-            elif isinstance(res, torch.Tensor):
-                # 检查是否是 meta 设备上的 tensor
-                # 注意：torch4ms.Tensor 的 device 属性返回字符串，不是 torch.device
-                # 所以我们需要检查 res 是否是普通的 torch.Tensor 且在 meta 设备上
-                try:
-                    device_type = res.device.type if hasattr(res.device, 'type') else str(res.device)
-                    if device_type == 'meta':
-                        # meta tensor 无法直接转换
-                        # 这种情况可能发生在某些操作没有被正确拦截的情况下
-                        # 尝试检查是否有 _elem 属性（可能是 torch4ms.Tensor 但没有被正确识别）
-                        if hasattr(res, '_elem') and isinstance(res._elem, ms_Tensor):
-                            return res._elem
-                        raise RuntimeError(
-                            "Cannot convert meta tensor to MindSpore tensor. "
-                            "This may indicate that the forward function returned a meta tensor. "
-                            "Please ensure all operations are performed on real tensors. "
-                            "If you're using torch4ms, make sure the environment is enabled."
-                        )
-                except (AttributeError, TypeError):
-                    # 如果 device 属性不是标准的 torch.device，尝试直接转换
-                    pass
-                
-                from torch4ms.ops import mappings
-                return mappings.t2ms(res)
-            else:
-                return res
+                res = base_forward(*t4_args)
+
+            # 同构回收：torch4ms / torch 结果统一映射回 MindSpore
+            return env.t2ms_iso(res)
 
         return forward_fn
     
@@ -232,9 +195,11 @@ def forward_fn(*args):
             # 合并参数和 buffers
             params_and_buffers = {**param_dict, **buffer_dict}
             
-            # 使用 functional_call 调用模块
-            with torch_stateless._reparametrize_module(module, params_and_buffers):
-                return module.forward(*input_args)
+            # 使用 torch.func.functional_call（与 torchax 路径一致）
+            # tie_weights=False 避免共享参数在 functional_call 中被强制绑定。
+            return torch.func.functional_call(
+                module, params_and_buffers, input_args, {}, tie_weights=False
+            )
         
         return forward_fn
     
diff --git a/src/torch4ms/autograd/ms_autograd_function.py b/src/torch4ms/autograd/ms_autograd_function.py
index 4ae2235d2..36fa7557c 100644
--- a/src/torch4ms/autograd/ms_autograd_function.py
+++ b/src/torch4ms/autograd/ms_autograd_function.py
@@ -363,11 +363,9 @@ def extract_and_wrap_loss_fn(module: torch.nn.Module, loss_fn: Callable, *module
     # 获取 MindSpore 版本的输入和标签
     input_tensors_ms = tuple(inp._elem if isinstance(inp, Tensor) else inp for inp in input_tensors)
     label_tensors_ms = tuple(label._elem if isinstance(label, Tensor) else label for label in label_tensors) if label_tensors else None
-    
+
     # 对于 Linear 模块，直接使用 MindSpore 算子实现（避免转换问题）
     if isinstance(module, torch.nn.Linear):
-        from mindspore import ops
-        
         # 获取参数
         params = extractor.get_trainable_params()
         param_tensors = []
@@ -394,7 +392,7 @@ def model_forward_fn(*ms_params):
             return output_ms
     else:
         # 其他模块使用 extractor 的方法
-        model_forward_fn = extractor.get_ms_forward_fn(env, include_buffers=False)
+        model_forward_base_fn = extractor.get_ms_forward_fn(env, include_buffers=False)
         # 获取参数
         params = extractor.get_trainable_params()
         
@@ -409,48 +407,131 @@ def model_forward_fn(*ms_params):
                 param_tensors.append(Tensor(ms_param, env, requires_grad=True))
             else:
                 raise TypeError(f"Unsupported parameter type: {type(param)}")
-    
-    # 创建完整的 loss 函数（MindSpore 语义）
-    # 这个函数包含：模型前向 + loss 计算
-    # 注意：我们需要将 loss_fn 的操作也转换为 MindSpore 语义
+
+        def model_forward_fn(*ms_params):
+            return model_forward_base_fn(*ms_params, *input_tensors_ms)
+    
+    def _to_ms_scalar_or_tensor(value):
+        if isinstance(value, ms_Tensor):
+            return value
+        if isinstance(value, Tensor):
+            return value._elem
+        if isinstance(value, torch.Tensor):
+            from torch4ms.ops import mappings
+            return mappings.t2ms(value)
+        try:
+            return ms_Tensor(value)
+        except Exception:
+            return value
+
+    def _apply_reduction(ms_loss, reduction):
+        if reduction == "mean":
+            return ops.reduce_mean(ms_loss)
+        if reduction == "sum":
+            return ops.reduce_sum(ms_loss)
+        return ms_loss
+
+    def _convert_known_torch_loss_to_ms(ms_output, ms_target):
+        """
+        针对常用回归 / 二分类 / 多分类损失做显式转换：
+        - nn.MSELoss
+        - nn.L1Loss
+        - nn.SmoothL1Loss
+        - nn.HuberLoss
+        - nn.BCEWithLogitsLoss（基础无权重场景）
+        """
+        if ms_target is None:
+            return None
+
+        # 仅针对 torch.nn.Module 形式的 loss 做稳定转换
+        if not isinstance(loss_fn, torch.nn.Module):
+            return None
+
+        reduction = getattr(loss_fn, "reduction", "mean")
+        diff = ms_output - ms_target
+
+        # -------- 回归类 --------
+        if isinstance(loss_fn, torch.nn.MSELoss):
+            per_elem = ops.square(diff)
+            return _apply_reduction(per_elem, reduction)
+
+        if isinstance(loss_fn, torch.nn.L1Loss):
+            per_elem = ops.abs(diff)
+            return _apply_reduction(per_elem, reduction)
+
+        if isinstance(loss_fn, torch.nn.SmoothL1Loss):
+            beta = float(getattr(loss_fn, "beta", 1.0))
+            abs_diff = ops.abs(diff)
+            if beta == 0.0:
+                per_elem = abs_diff
+            else:
+                per_elem = ops.where(
+                    abs_diff < beta,
+                    0.5 * ops.square(diff) / beta,
+                    abs_diff - 0.5 * beta,
+                )
+            return _apply_reduction(per_elem, reduction)
+
+        if isinstance(loss_fn, torch.nn.HuberLoss):
+            delta = float(getattr(loss_fn, "delta", 1.0))
+            abs_diff = ops.abs(diff)
+            per_elem = ops.where(
+                abs_diff <= delta,
+                0.5 * ops.square(diff),
+                delta * (abs_diff - 0.5 * delta),
+            )
+            return _apply_reduction(per_elem, reduction)
+
+        # -------- 二分类：BCEWithLogitsLoss（基础配置） --------
+        if isinstance(loss_fn, torch.nn.BCEWithLogitsLoss):
+            weight = getattr(loss_fn, "weight", None)
+            pos_weight = getattr(loss_fn, "pos_weight", None)
+            # 复杂配置暂时交给 fallback
+            if weight is not None or pos_weight is not None:
+                return None
+            x = ms_output
+            y = ms_target
+            zero = ops.zeros_like(x)
+            max_x0 = ops.maximum(x, zero)
+            neg_abs_x = -ops.abs(x)
+            # log(1 + exp(-|x|)) 数值安全版本
+            log_term = ops.log(ops.exp(neg_abs_x) + 1.0)
+            per_elem = max_x0 - x * y + log_term
+            return _apply_reduction(per_elem, reduction)
+
+        return None
+
+    # 创建完整的 loss 函数（模型前向 + loss）
+    # 优先按 PyTorch 风格调用 loss_fn，让 loss 内部算子走 torch4ms 拦截转换。
+    # 如果用户传入的是 MindSpore 风格 loss_fn（期望 ms.Tensor），则回退到 ms 调用。
     def full_loss_fn(*ms_params):
-        """完整的 loss 函数：模型前向 + loss 计算"""
+        """完整的 loss 函数：模型前向 + loss 计算。"""
         # 模型前向
-        output_ms = model_forward_fn(*ms_params, *input_tensors_ms)
-        
-        # 调用 loss 函数（直接使用 MindSpore Tensor）
-        # 假设 loss_fn 是 MindSpore 语义的函数
-        if label_tensors_ms:
-            loss_ms = loss_fn(output_ms, label_tensors_ms[0])
-        else:
-            loss_ms = loss_fn(output_ms)
-        
-        
-        # 确保返回 MindSpore Tensor
-        if isinstance(loss_ms, ms_Tensor):
-            return loss_ms
-        else:
-            # 如果不是 MindSpore Tensor，尝试转换
-            try:
-                return ms_Tensor(loss_ms)
-            except:
-                return loss_ms
-        
-        # 转换回 MindSpore Tensor
-        # 注意：loss_t4 可能已经是 MindSpore Tensor（如果 loss_fn 被正确拦截）
-        if isinstance(loss_t4, Tensor):
-            return loss_t4._elem
-        elif isinstance(loss_t4, torch.Tensor):
-            from torch4ms.ops import mappings
-            return mappings.t2ms(loss_t4)
-        elif isinstance(loss_t4, ms_Tensor):
-            return loss_t4
-        else:
-            # 尝试转换为 MindSpore Tensor
-            try:
-                return ms_Tensor(loss_t4)
-            except:
-                return loss_t4
+        output_ms = model_forward_fn(*ms_params)
+        target_ms = label_tensors_ms[0] if label_tensors_ms else None
+
+        # 0) 常用损失显式转换（先满足 MSE 等核心需求）
+        known_loss_ms = _convert_known_torch_loss_to_ms(output_ms, target_ms)
+        if known_loss_ms is not None:
+            return known_loss_ms
+
+        # 1) 优先 PyTorch 风格 loss（通过 env.dispatch 拦截转换）
+        try:
+            output_t4 = Tensor(output_ms, env, requires_grad=False)
+            label_t4 = label_tensors[0] if label_tensors else None
+            with env:
+                if label_t4 is not None:
+                    loss_val = loss_fn(output_t4, label_t4)
+                else:
+                    loss_val = loss_fn(output_t4)
+            return _to_ms_scalar_or_tensor(loss_val)
+        except Exception:
+            # 2) 回退 MindSpore 风格 loss（兼容现有调用）
+            if label_tensors_ms:
+                loss_val = loss_fn(output_ms, label_tensors_ms[0])
+            else:
+                loss_val = loss_fn(output_ms)
+            return _to_ms_scalar_or_tensor(loss_val)
     
     # 包装为 autograd 函数（只传入 params）
     loss_output = ms2t_autograd(full_loss_fn, *param_tensors)
diff --git a/src/torch4ms/ops/maten.py b/src/torch4ms/ops/maten.py
index debd494bb..a5c7f8162 100644
--- a/src/torch4ms/ops/maten.py
+++ b/src/torch4ms/ops/maten.py
@@ -86,6 +86,12 @@ def _aten_div(x, y, rounding_mode=None):
 @op(torch.ops.aten.pow)
 @op(torch.ops.aten.pow.Scalar)
 def _aten_pow(x, y):
+    if not isinstance(y, ms.Tensor):
+        y = ms.Tensor(float(y), dtype=x.dtype)
+    else:
+        y_dtype = getattr(y, "dtype", None)
+        if y_dtype is not None and "Int" in str(y_dtype):
+            y = ops.cast(y, x.dtype)
     return ops.pow(x, y)
 
 
@@ -294,7 +300,29 @@ def _aten_log_softmax(x, dim=-1, dtype=None):
 
 @op(torch.ops.aten.layer_norm)
 def _aten_layer_norm(x, normalized_shape, weight=None, bias=None, eps=1e-5):
-    return ops.layer_norm(x, normalized_shape, weight, bias, eps)
+    # 手工实现 layer_norm，避免部分 MindSpore 版本 layer_norm 内核不可用
+    if isinstance(normalized_shape, int):
+        normalized_shape = (normalized_shape,)
+    else:
+        normalized_shape = tuple(normalized_shape)
+
+    ndim = int(x.ndim)
+    k = len(normalized_shape)
+    axes = tuple(range(ndim - k, ndim))
+
+    mean = ops.mean(x, axis=axes, keep_dims=True)
+    centered = x - mean
+    var = ops.mean(centered * centered, axis=axes, keep_dims=True)
+    y = centered / ops.sqrt(var + eps)
+
+    if weight is not None:
+        w_shape = (1,) * (ndim - k) + normalized_shape
+        y = y * ops.reshape(weight, w_shape)
+    if bias is not None:
+        b_shape = (1,) * (ndim - k) + normalized_shape
+        y = y + ops.reshape(bias, b_shape)
+
+    return y
 
 
 @op(torch.ops.aten.batch_norm)
@@ -303,11 +331,11 @@ def _aten_batch_norm(x, running_mean, running_var, weight=None, bias=None,
     # MindSpore的BatchNorm接口略有不同，这里使用MindSpore的reduce操作实现
     if training:
         # 训练模式下的BatchNorm - 计算当前batch的均值和方差
-        # 使用MindSpore的reduce_mean计算均值
-        mean = ops.reduce_mean(x, axis=(0, 2, 3), keep_dims=True)
+        # 使用 MindSpore 的 mean 计算均值（该版本参数名为 keep_dims）
+        mean = ops.mean(x, axis=(0, 2, 3), keep_dims=True)
         # 计算方差：E[(x - E[x])^2]
         centered = x - mean
-        var = ops.reduce_mean(centered * centered, axis=(0, 2, 3), keep_dims=True)
+        var = ops.mean(centered * centered, axis=(0, 2, 3), keep_dims=True)
         # 更新running_mean和running_var（简化实现）
         if running_mean is not None:
             running_mean = running_mean * (1 - momentum) + mean.squeeze() * momentum
@@ -345,7 +373,11 @@ def _aten_silu(x):
 @op(torch.ops.aten.dropout)
 def _aten_dropout(x, p=0.5, training=True):
     if training:
-        return ops.dropout(x, p)
+        out = ops.dropout(x, p)
+        # 部分 MindSpore 版本返回 (output, mask)
+        if isinstance(out, tuple):
+            return out[0]
+        return out
     return x
 
 
@@ -684,49 +716,20 @@ def _aten_convolution(input, weight, bias, stride, padding, dilation, transposed
         output_padding = (0,) * num_shape_dim
     
     # 根据空间维度数量选择对应的卷积操作
-    # 使用MindSpore的nn层API（参考 torchax 的实现方式，统一使用nn层以确保兼容性）
-    import mindspore.nn as nn
+    # 直接使用 ops.convXd，让 weight/bias 作为显式输入参与图计算，避免 set_data 造成梯度链断裂。
     
     if num_shape_dim == 1:
-        # 1D卷积 - 通过转换为2D卷积实现（参考 MindSpore 内部实现）
-        # 输入: [N, C_in, L] -> [N, C_in, 1, L]
-        # 权重: [C_out, C_in, K] -> [C_out, C_in, 1, K]
-        # 使用 Conv2D 进行卷积
-        # 输出: [N, C_out, 1, L_out] -> [N, C_out, L_out]
-        
-        # 扩展输入维度：在位置2插入维度1
-        input_2d = ops.expand_dims(input, 2)  # [N, C_in, L] -> [N, C_in, 1, L]
-        
-        # 扩展权重维度：在位置2插入维度1
-        weight_2d = ops.expand_dims(weight, 2)  # [C_out, C_in, K] -> [C_out, C_in, 1, K]
-        
-        # 准备padding：对于2D，padding格式为(上,下,左,右)，1D只需要左右padding
-        pad_val = int(padding[0])
-        pad_2d = (0, 0, pad_val, pad_val)  # (上=0, 下=0, 左=pad, 右=pad)
-        
-        # 使用Conv2D进行卷积
-        conv2d_op = nn.Conv2d(
-            in_channels=int(weight.shape[1] * groups),
-            out_channels=int(weight.shape[0]),
-            kernel_size=(1, int(weight.shape[2])),  # (height=1, width=kernel_size)
-            stride=(1, int(stride[0])),  # (height_stride=1, width_stride=stride)
+        pad_val = int(padding[0]) if len(padding) > 0 else 0
+        res = ops.conv1d(
+            input,
+            weight,
+            bias=bias,
+            stride=int(stride[0]),
             pad_mode='pad',
-            padding=pad_2d,
-            dilation=(1, int(dilation[0])),  # (height_dilation=1, width_dilation=dilation)
-            group=groups,
-            has_bias=(bias is not None)
+            padding=pad_val,
+            dilation=int(dilation[0]),
+            groups=groups,
         )
-        
-        # 设置权重和偏置
-        conv2d_op.weight.set_data(weight_2d)
-        if bias is not None:
-            conv2d_op.bias.set_data(bias)
-        
-        # 执行2D卷积
-        res_2d = conv2d_op(input_2d)  # [N, C_out, 1, L_out]
-        
-        # 压缩输出维度：移除位置2的维度
-        res = ops.squeeze(res_2d, 2)  # [N, C_out, 1, L_out] -> [N, C_out, L_out]
     
     elif num_shape_dim == 2:
         # 2D卷积 - 使用 nn.Conv2d
@@ -741,25 +744,16 @@ def _aten_convolution(input, weight, bias, stride, padding, dilation, transposed
         else:
             pad_val = tuple(padding[:2])
         
-        # 创建临时Conv2d层
-        conv2d_layer = nn.Conv2d(
-            in_channels=int(weight.shape[1] * groups),
-            out_channels=int(weight.shape[0]),
-            kernel_size=tuple(weight.shape[2:]),
+        res = ops.conv2d(
+            input,
+            weight,
+            bias=bias,
             stride=tuple(stride),
             pad_mode='pad',
             padding=pad_val,
             dilation=tuple(dilation),
-            group=groups,
-            has_bias=(bias is not None)
+            groups=groups,
         )
-        
-        # 设置权重和偏置
-        conv2d_layer.weight.set_data(weight)
-        if bias is not None:
-            conv2d_layer.bias.set_data(bias)
-        
-        res = conv2d_layer(input)
     
     elif num_shape_dim == 3:
         # 3D卷积 - 使用 nn.Conv3d
@@ -777,25 +771,16 @@ def _aten_convolution(input, weight, bias, stride, padding, dilation, transposed
         else:
             pad_val = tuple(padding[:3])
         
-        # 创建临时Conv3d层
-        conv3d_layer = nn.Conv3d(
-            in_channels=int(weight.shape[1] * groups),
-            out_channels=int(weight.shape[0]),
-            kernel_size=tuple(weight.shape[2:]),
+        res = ops.conv3d(
+            input,
+            weight,
+            bias=bias,
             stride=tuple(stride),
             pad_mode='pad',
             padding=pad_val,
             dilation=tuple(dilation),
-            group=groups,
-            has_bias=(bias is not None)
+            groups=groups,
         )
-        
-        # 设置权重和偏置
-        conv3d_layer.weight.set_data(weight)
-        if bias is not None:
-            conv3d_layer.bias.set_data(bias)
-        
-        res = conv3d_layer(input)
     
     else:
         raise NotImplementedError(f"Convolution not supported for {num_shape_dim}D spatial dimensions (weight.ndim={weight.ndim})")
@@ -1114,9 +1099,11 @@ def _aten_gather(x, dim, index):
 @op(torch.ops.aten.embedding)
 def _aten_embedding(weight, indices, padding_idx=-1, scale_grad_by_freq=False, sparse=False):
     """嵌入层"""
-    # 使用 MindSpore 的 embedding lookup
-    # 注意：MindSpore 的 embedding 可能需要不同的实现方式
-    return ops.gather(weight, indices, axis=0)
+    vocab_size = weight.shape[0]
+    on_value = ms.Tensor(1.0, dtype=weight.dtype)
+    off_value = ms.Tensor(0.0, dtype=weight.dtype)
+    one_hot = ops.one_hot(indices, vocab_size, on_value, off_value)
+    return ops.matmul(one_hot, weight)
 
 
 @op(torch.ops.aten.embedding_renorm_)
diff --git a/src/torch4ms/ops/mten.py b/src/torch4ms/ops/mten.py
index 335e25b15..2871a7c7f 100644
--- a/src/torch4ms/ops/mten.py
+++ b/src/torch4ms/ops/mten.py
@@ -87,6 +87,12 @@ def _aten_div(x, y, rounding_mode=None):
 @op(torch.ops.aten.pow)
 @op(torch.ops.aten.pow.Scalar)
 def _aten_pow(x, y):
+    if not isinstance(y, ms.Tensor):
+        y = ms.Tensor(float(y), dtype=x.dtype)
+    else:
+        y_dtype = getattr(y, "dtype", None)
+        if y_dtype is not None and "Int" in str(y_dtype):
+            y = ops.cast(y, x.dtype)
     return ops.pow(x, y)
 
 
@@ -295,7 +301,28 @@ def _aten_log_softmax(x, dim=-1, dtype=None):
 
 @op(torch.ops.aten.layer_norm)
 def _aten_layer_norm(x, normalized_shape, weight=None, bias=None, eps=1e-5):
-    return ops.layer_norm(x, normalized_shape, weight, bias, eps)
+    if isinstance(normalized_shape, int):
+        normalized_shape = (normalized_shape,)
+    else:
+        normalized_shape = tuple(normalized_shape)
+
+    ndim = int(x.ndim)
+    k = len(normalized_shape)
+    axes = tuple(range(ndim - k, ndim))
+
+    mean = ops.mean(x, axis=axes, keep_dims=True)
+    centered = x - mean
+    var = ops.mean(centered * centered, axis=axes, keep_dims=True)
+    y = centered / ops.sqrt(var + eps)
+
+    if weight is not None:
+        w_shape = (1,) * (ndim - k) + normalized_shape
+        y = y * ops.reshape(weight, w_shape)
+    if bias is not None:
+        b_shape = (1,) * (ndim - k) + normalized_shape
+        y = y + ops.reshape(bias, b_shape)
+
+    return y
 
 
 @op(torch.ops.aten.batch_norm)
@@ -305,7 +332,7 @@ def _aten_batch_norm(x, running_mean, running_var, weight=None, bias=None,
     if training:
         # 训练模式下的BatchNorm
         mean = ops.mean(x, (0, 2, 3), keep_dims=True)
-        var = ops.var(x, (0, 2, 3), keep_dims=True)
+        var = ops.var(x, (0, 2, 3), keepdims=True)
         # 更新running_mean和running_var（简化实现）
         if running_mean is not None:
             running_mean = running_mean * (1 - momentum) + mean.squeeze() * momentum
@@ -343,7 +370,10 @@ def _aten_silu(x):
 @op(torch.ops.aten.dropout)
 def _aten_dropout(x, p=0.5, training=True):
     if training:
-        return ops.dropout(x, p)
+        out = ops.dropout(x, p)
+        if isinstance(out, tuple):
+            return out[0]
+        return out
     return x
 
 
diff --git a/src/torch4ms/ops/mtorch.py b/src/torch4ms/ops/mtorch.py
index 9de8b6771..c907afedc 100644
--- a/src/torch4ms/ops/mtorch.py
+++ b/src/torch4ms/ops/mtorch.py
@@ -16,7 +16,7 @@
 import torch4ms.tensor
 from torch4ms.view import View, NarrowInfo
 import torch.utils._pytree as pytree
-from torch4ms.ops.maten import _aten_convolution
+from torch4ms.ops.maten import _aten_convolution, _aten_batch_norm, _aten_dropout, _aten_layer_norm
 
 
 def register_function(torch_func, **kwargs):
@@ -173,6 +173,183 @@ def one_hot(tensor, num_classes=-1):
     return mops.one_hot(tensor, num_classes, on_value=1, off_value=0).astype(mnp.int64)
 
 
+@register_function(torch.transpose)
+def functional_transpose(input, dim0, dim1):
+    rank = int(input.ndim)
+    d0 = int(dim0)
+    d1 = int(dim1)
+    if d0 < 0:
+        d0 += rank
+    if d1 < 0:
+        d1 += rank
+    perm = list(range(rank))
+    perm[d0], perm[d1] = perm[d1], perm[d0]
+    return mops.transpose(input, tuple(perm))
+
+
+@register_function(torch.permute)
+def functional_permute(input, dims):
+    return mops.transpose(input, tuple(dims))
+
+
+@register_function(torch.reshape)
+def functional_reshape(input, *shape):
+    if len(shape) == 1 and isinstance(shape[0], (list, tuple)):
+        target_shape = tuple(shape[0])
+    else:
+        target_shape = tuple(shape)
+    return mops.reshape(input, target_shape)
+
+
+@register_function(torch.Tensor.view)
+def tensor_view(input, *shape):
+    if len(shape) == 1 and isinstance(shape[0], (list, tuple)):
+        target_shape = tuple(shape[0])
+    else:
+        target_shape = tuple(shape)
+    return mops.reshape(input, target_shape)
+
+
+@register_function(torch.Tensor.contiguous)
+def tensor_contiguous(input, memory_format=None):
+    # MindSpore tensors are contiguous in this path; keep semantic no-op.
+    return input
+
+
+@register_function(torch.Tensor.float)
+def tensor_float(input):
+    return mops.cast(input, ms.float32)
+
+
+@register_function(torch.flatten)
+def functional_flatten(input, start_dim=0, end_dim=-1):
+    return mops.flatten(input, start_dim, end_dim)
+
+
+@register_function(torch.chunk)
+def functional_chunk(input, chunks, dim=0):
+    return mops.chunk(input, chunks, axis=dim)
+
+
+@register_function(torch.split)
+def functional_split(input, split_size_or_sections, dim=0):
+    return mops.split(input, split_size_or_sections, axis=dim)
+
+
+@register_function(torch.stack)
+def functional_stack(tensors, dim=0):
+    return mops.stack(tuple(tensors), axis=dim)
+
+
+@register_function(torch.cat)
+def functional_cat(tensors, dim=0):
+    return mops.concat(tuple(tensors), axis=dim)
+
+
+@register_function(torch.matmul)
+def functional_matmul(input, other):
+    return mops.matmul(input, other)
+
+
+@register_function(torch.softmax)
+def functional_softmax(input, dim=None, dtype=None):
+    axis = -1 if dim is None else dim
+    x = input
+    if dtype is not None:
+        x = x.astype(mappings.t2ms_dtype(dtype))
+    return mops.softmax(x, axis=axis)
+
+
+@register_function(torch.nn.functional.softmax)
+def functional_softmax_f(input, dim=None, _stacklevel=3, dtype=None):
+    axis = -1 if dim is None else dim
+    x = input
+    if dtype is not None:
+        x = x.astype(mappings.t2ms_dtype(dtype))
+    return mops.softmax(x, axis=axis)
+
+
+@register_function(torch.mul)
+def functional_mul(input, other):
+    return mops.mul(input, other)
+
+
+@register_function(torch.pow)
+def functional_pow(input, exponent):
+    exp = exponent
+    if not isinstance(exp, ms.Tensor):
+        exp = ms.Tensor(float(exp), dtype=input.dtype)
+    else:
+        exp_dtype = getattr(exp, "dtype", None)
+        if exp_dtype is not None and "Int" in str(exp_dtype):
+            exp = mops.cast(exp, input.dtype)
+    return mops.pow(input, exp)
+
+
+@register_function(torch.rsqrt)
+def functional_rsqrt(input):
+    return mops.rsqrt(input)
+
+
+@register_function(torch.sqrt)
+def functional_sqrt(input):
+    return mops.sqrt(input)
+
+
+@register_function(torch.nn.functional.gelu)
+def functional_gelu(input, approximate="none"):
+    approx = "tanh" if approximate == "tanh" else "none"
+    return mops.gelu(input, approximate=approx)
+
+
+@register_function(torch.nn.functional.embedding)
+def functional_embedding(input, weight, padding_idx=-1, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False):
+    vocab_size = weight.shape[0]
+    on_value = ms.Tensor(1.0, dtype=weight.dtype)
+    off_value = ms.Tensor(0.0, dtype=weight.dtype)
+    one_hot = mops.one_hot(input, vocab_size, on_value, off_value)
+    return mops.matmul(one_hot, weight)
+
+
+@register_function(torch.mean)
+def functional_mean(input, dim=None, keepdim=False, dtype=None):
+    x = input
+    if dtype is not None:
+        x = x.astype(mappings.t2ms_dtype(dtype))
+    if dim is None:
+        return mops.mean(x)
+    axis = tuple(dim) if isinstance(dim, (list, tuple)) else dim
+    return mops.mean(x, axis=axis, keep_dims=keepdim)
+
+
+# 补充 Tensor 方法入口，避免 x.transpose()/x.reshape() 走回退路径
+if hasattr(torch.Tensor, "transpose"):
+    register_function(torch.Tensor.transpose)(functional_transpose)
+if hasattr(torch.Tensor, "permute"):
+    register_function(torch.Tensor.permute)(functional_permute)
+if hasattr(torch.Tensor, "reshape"):
+    register_function(torch.Tensor.reshape)(functional_reshape)
+if hasattr(torch.Tensor, "flatten"):
+    register_function(torch.Tensor.flatten)(functional_flatten)
+if hasattr(torch.Tensor, "chunk"):
+    register_function(torch.Tensor.chunk)(functional_chunk)
+if hasattr(torch.Tensor, "split"):
+    register_function(torch.Tensor.split)(functional_split)
+if hasattr(torch.Tensor, "pow"):
+    register_function(torch.Tensor.pow)(functional_pow)
+if hasattr(torch.Tensor, "mean"):
+    register_function(torch.Tensor.mean)(functional_mean)
+if hasattr(torch.Tensor, "sqrt"):
+    register_function(torch.Tensor.sqrt)(functional_sqrt)
+
+
+@register_function(torch.relu)
+@register_function(torch.nn.functional.relu)
+def functional_relu(input, inplace=False):
+    # MindSpore relu 没有 inplace 语义，忽略该参数。
+    return mops.relu(input)
+
+
 @register_function(torch.nn.functional.pad)
 def pad(tensor, pad, mode="constant", value=None):
     # MindSpore的pad接口需要不同的参数格式
@@ -234,6 +411,107 @@ def scaled_dot_product_attention(
     )
 
 
+@register_function(torch.nn.functional.multi_head_attention_forward, is_mindspore_function=False)
+def functional_multi_head_attention_forward(
+    query,
+    key,
+    value,
+    embed_dim_to_check,
+    num_heads,
+    in_proj_weight,
+    in_proj_bias,
+    bias_k,
+    bias_v,
+    add_zero_attn,
+    dropout_p,
+    out_proj_weight,
+    out_proj_bias,
+    training=True,
+    key_padding_mask=None,
+    need_weights=True,
+    attn_mask=None,
+    use_separate_proj_weight=False,
+    q_proj_weight=None,
+    k_proj_weight=None,
+    v_proj_weight=None,
+    static_k=None,
+    static_v=None,
+    average_attn_weights=True,
+    is_causal=False,
+):
+    """
+    轻量 MHA 前向实现，优先覆盖 TransformerEncoderLayer 的常见路径：
+    - use_separate_proj_weight=False
+    - bias_k/bias_v/static_k/static_v/add_zero_attn 不启用
+    - need_weights=False
+    """
+    if use_separate_proj_weight:
+        raise torch4ms.tensor.OperatorNotFound(
+            "multi_head_attention_forward with separate proj weights is not supported yet."
+        )
+    if bias_k is not None or bias_v is not None or static_k is not None or static_v is not None:
+        raise torch4ms.tensor.OperatorNotFound(
+            "multi_head_attention_forward with bias_k/bias_v/static_k/static_v is not supported yet."
+        )
+    if add_zero_attn:
+        raise torch4ms.tensor.OperatorNotFound(
+            "multi_head_attention_forward with add_zero_attn=True is not supported yet."
+        )
+    if key_padding_mask is not None:
+        raise torch4ms.tensor.OperatorNotFound(
+            "multi_head_attention_forward with key_padding_mask is not supported yet."
+        )
+
+    # query/key/value shape: [L, N, E]
+    # in-projection: [L, N, 3E] -> chunk to q/k/v
+    qkv = torch.nn.functional.linear(query, in_proj_weight, in_proj_bias)
+    q, k, v = torch.chunk(qkv, 3, dim=-1)
+
+    L, N, E = q.shape
+    head_dim = E // num_heads
+    if head_dim * num_heads != E:
+        raise ValueError(f"embed_dim {E} not divisible by num_heads {num_heads}")
+
+    def _to_heads(x):
+        # [L, N, E] -> [N, H, L, D]
+        x = torch.transpose(x, 0, 1)
+        x = torch.reshape(x, (N, L, num_heads, head_dim))
+        x = torch.transpose(x, 1, 2)
+        return x
+
+    qh = _to_heads(q)
+    kh = _to_heads(k)
+    vh = _to_heads(v)
+
+    if attn_mask is not None:
+        raise torch4ms.tensor.OperatorNotFound(
+            "multi_head_attention_forward with attn_mask is not supported yet."
+        )
+    if is_causal:
+        raise torch4ms.tensor.OperatorNotFound(
+            "multi_head_attention_forward with is_causal=True is not supported yet."
+        )
+
+    scale = 1.0 / math.sqrt(head_dim)
+    attn_scores = torch.matmul(qh, torch.transpose(kh, -2, -1)) * scale
+    attn_probs = torch.softmax(attn_scores, dim=-1)
+    if training and dropout_p > 0:
+        attn_probs = torch.nn.functional.dropout(attn_probs, p=dropout_p, training=True)
+    attn_out = torch.matmul(attn_probs, vh)
+
+    # [N, H, L, D] -> [L, N, E]
+    attn_out = torch.transpose(attn_out, 1, 2)
+    attn_out = torch.reshape(attn_out, (N, L, E))
+    attn_out = torch.transpose(attn_out, 0, 1)
+
+    out = torch.nn.functional.linear(attn_out, out_proj_weight, out_proj_bias)
+
+    if need_weights:
+        # 暂不返回 attention 权重（测试路径中 need_weights=False）
+        return out, None
+    return out, None
+
+
 @register_function(
     torch.Tensor.__getitem__,
     is_mindspore_function=False,
@@ -845,6 +1123,7 @@ def _functional_max_pool2d(
     return output
 
 
+@register_function(torch.conv1d)
 @register_function(torch.nn.functional.conv1d)
 def functional_conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
     """
@@ -884,6 +1163,7 @@ def functional_conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1,
     return result
 
 
+@register_function(torch.conv2d)
 @register_function(torch.nn.functional.conv2d)
 def functional_conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
     """
@@ -923,6 +1203,7 @@ def functional_conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1,
     return result
 
 
+@register_function(torch.conv3d)
 @register_function(torch.nn.functional.conv3d)
 def functional_conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
     """
@@ -960,3 +1241,41 @@ def functional_conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1,
     if isinstance(input, torch4ms.tensor.Tensor):
         return input._env.ms2t_iso(result)
     return result
+
+
+@register_function(torch.nn.functional.batch_norm)
+def functional_batch_norm(input, running_mean, running_var, weight=None, bias=None,
+                          training=False, momentum=0.1, eps=1e-5):
+    result = _aten_batch_norm(
+        input, running_mean, running_var, weight, bias, training, momentum, eps
+    )
+    if isinstance(input, torch4ms.tensor.Tensor):
+        return input._env.ms2t_iso(result)
+    return result
+
+
+@register_function(torch.nn.functional.dropout)
+def functional_dropout(input, p=0.5, training=True, inplace=False):
+    # MindSpore 路径无 inplace 语义，这里忽略 inplace 参数。
+    result = _aten_dropout(input, p, training)
+    if isinstance(input, torch4ms.tensor.Tensor):
+        return input._env.ms2t_iso(result)
+    return result
+
+
+if hasattr(torch, "batch_norm"):
+    register_function(torch.batch_norm)(functional_batch_norm)
+if hasattr(torch, "dropout"):
+    register_function(torch.dropout)(functional_dropout)
+
+
+@register_function(torch.nn.functional.layer_norm)
+def functional_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-5):
+    result = _aten_layer_norm(input, normalized_shape, weight, bias, eps)
+    if isinstance(input, torch4ms.tensor.Tensor):
+        return input._env.ms2t_iso(result)
+    return result
+
+
+if hasattr(torch, "layer_norm"):
+    register_function(torch.layer_norm)(functional_layer_norm)
diff --git a/src/torch4ms/tensor.py b/src/torch4ms/tensor.py
index 45e04c941..d0a81d999 100644
--- a/src/torch4ms/tensor.py
+++ b/src/torch4ms/tensor.py
@@ -160,8 +160,42 @@ def size(self, dim=None):
         return self._elem.shape[dim]
 
     def flatten(self, start_dim=0, end_dim=-1):
-        # 使用MindSpore的reshape操作
-        return self._env.ms2t_iso(mnp.reshape(self._elem, (-1,)))
+        # 对齐 PyTorch 语义：仅展平 [start_dim, end_dim] 区间
+        shape = tuple(int(s) for s in self._elem.shape)
+        ndim = len(shape)
+        if ndim == 0:
+            # 标量 flatten 后为一维长度 1
+            res = mnp.reshape(self._elem, (1,))
+            out = self._env.ms2t_iso(res)
+            if self._requires_grad:
+                out._requires_grad = True
+            return out
+
+        start = int(start_dim)
+        end = int(end_dim)
+        if start < 0:
+            start += ndim
+        if end < 0:
+            end += ndim
+
+        start = max(0, min(start, ndim - 1))
+        end = max(0, min(end, ndim - 1))
+        if end < start:
+            # 与 PyTorch 行为一致：空区间不改变形状
+            out = self._env.ms2t_iso(self._elem)
+            if self._requires_grad:
+                out._requires_grad = True
+            return out
+
+        flat_dim = 1
+        for s in shape[start : end + 1]:
+            flat_dim *= s
+        new_shape = shape[:start] + (flat_dim,) + shape[end + 1 :]
+        res = mnp.reshape(self._elem, new_shape)
+        out = self._env.ms2t_iso(res)
+        if self._requires_grad:
+            out._requires_grad = True
+        return out
 
     # ========= 基本算术运算支持 =========
     def _binary_op(self, other, ms_op, op_name=None):
@@ -274,12 +308,22 @@ def __rfloordiv__(self, other):
 
     def __pow__(self, other):
         """支持 x ** y。"""
+        if not isinstance(other, (Tensor, ms_Tensor, torch.Tensor, np.ndarray)):
+            try:
+                other = ms_Tensor(float(other), dtype=self._elem.dtype)
+            except Exception:
+                pass
         return self._binary_op(other, ops.pow)
 
     def __rpow__(self, other):
         """支持 y ** x。"""
         if isinstance(other, Tensor):
             return other._binary_op(self, ops.pow)
+        if not isinstance(other, (ms_Tensor, torch.Tensor, np.ndarray)):
+            try:
+                other = ms_Tensor(float(other), dtype=self._elem.dtype)
+            except Exception:
+                pass
         return self._binary_op(other, lambda a, b: ops.pow(b, a))
 
     # 比较运算，返回布尔 Tensor
@@ -311,9 +355,9 @@ def type_as(self, other):
         # 确保other是Tensor类型
         if not isinstance(other, Tensor):
             raise TypeError(f"Expected Tensor, got {type(other).__name__}")
-        # 获取目标数据类型（MindSpore dtype，供 mnp.astype 使用）
+        # 获取目标数据类型并使用 MindSpore cast
         target_ms_dtype = other._elem.dtype
-        return self._env.ms2t_iso(mnp.astype(self._elem, target_ms_dtype))
+        return self._env.ms2t_iso(ops.cast(self._elem, target_ms_dtype))
 
     __torch_function__ = torch._C._disabled_torch_function_impl
 
diff --git a/tests/torch4ms/test_loss_interception.py b/tests/torch4ms/test_loss_interception.py
new file mode 100644
index 000000000..c8fcbbe51
--- /dev/null
+++ b/tests/torch4ms/test_loss_interception.py
@@ -0,0 +1,113 @@
+import torch
+import torch.nn as nn
+
+import torch4ms
+from torch4ms.autograd.ms_autograd_function import extract_and_wrap_loss_fn
+
+
+class SimpleRegNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc = nn.Linear(3, 1)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+class SimpleClsNet(nn.Module):
+    def __init__(self, num_classes: int = 4):
+        super().__init__()
+        self.fc = nn.Linear(3, num_classes)
+
+    def forward(self, x):
+        return self.fc(x)
+
+
+def _init_linear(module: nn.Module):
+    with torch.no_grad():
+        for p in module.parameters():
+            p.fill_(0.5)
+
+
+def _compare_loss(loss_name: str, model_ctor, loss_ctor, inputs, targets, atol=1e-5):
+    print("=" * 60)
+    print(f"Testing loss interception for {loss_name}")
+
+    # Torch reference
+    model_ref = model_ctor()
+    _init_linear(model_ref)
+    loss_ref_fn = loss_ctor()
+
+    out_ref = model_ref(inputs)
+    loss_ref = loss_ref_fn(out_ref, targets)
+    loss_ref_val = float(loss_ref.item())
+    print(f"[torch] loss = {loss_ref_val}")
+
+    # torch4ms path
+    model = model_ctor()
+    _init_linear(model)
+    loss_torch_fn = loss_ctor()
+
+    env = torch4ms.default_env()
+    with env:
+        loss_wrapper = extract_and_wrap_loss_fn(model, loss_torch_fn, inputs, targets)
+        loss_t4 = loss_wrapper.output
+
+    # 转回 torch 标量比较数值
+    loss_t4_val = float(loss_t4.torch().item())
+    print(f"[torch4ms] loss = {loss_t4_val}")
+
+    diff = abs(loss_t4_val - loss_ref_val)
+    print(f"[compare] |diff| = {diff}")
+    if diff < atol:
+        print(f"[OK] {loss_name} forward loss matches within atol={atol}")
+    else:
+        print(f"[WARN] {loss_name} forward loss mismatch (>|{atol}|)")
+
+
+def main():
+    # 公共输入
+    x_reg = torch.tensor([[1.0, 2.0, 3.0],
+                          [4.0, 5.0, 6.0]], dtype=torch.float32)
+    y_reg = torch.tensor([[10.0],
+                          [20.0]], dtype=torch.float32)
+
+    y_bce = torch.tensor([[0.0],
+                          [1.0]], dtype=torch.float32)
+
+    # 回归类损失
+    _compare_loss("MSELoss",
+                  SimpleRegNet,
+                  lambda: nn.MSELoss(reduction="mean"),
+                  x_reg,
+                  y_reg)
+
+    _compare_loss("L1Loss",
+                  SimpleRegNet,
+                  lambda: nn.L1Loss(reduction="mean"),
+                  x_reg,
+                  y_reg)
+
+    _compare_loss("SmoothL1Loss",
+                  SimpleRegNet,
+                  lambda: nn.SmoothL1Loss(beta=1.0, reduction="mean"),
+                  x_reg,
+                  y_reg)
+
+    _compare_loss("HuberLoss",
+                  SimpleRegNet,
+                  lambda: nn.HuberLoss(delta=1.0, reduction="mean"),
+                  x_reg,
+                  y_reg)
+
+    # 二分类 BCEWithLogits
+    _compare_loss("BCEWithLogitsLoss",
+                  SimpleRegNet,
+                  lambda: nn.BCEWithLogitsLoss(reduction="mean"),
+                  x_reg,
+                  y_bce)
+
+
+if __name__ == "__main__":
+    main()
+
diff --git a/tests/torch4ms/test_train_cnn.py b/tests/torch4ms/test_train_cnn.py
new file mode 100644
index 000000000..9ea457c93
--- /dev/null
+++ b/tests/torch4ms/test_train_cnn.py
@@ -0,0 +1,162 @@
+import os
+os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import math
+
+import torch4ms
+from torch4ms.autograd.ms_autograd_function import extract_and_wrap_loss_fn
+from torch4ms.optim import Torch4msOptimizer
+
+NUM_EPOCHS = 2
+LR = 0.01
+
+
+def get_dataloader():
+    torch.manual_seed(0)
+    x = torch.randn(4, 1, 8, 8, dtype=torch.float32)
+    y = torch.randn(4, 1, dtype=torch.float32)
+    return [(x, y)]
+
+
+class SimpleCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2d(1, 4, kernel_size=3, padding=1)
+        self.bn = nn.BatchNorm2d(4)
+        self.relu = nn.ReLU()
+        self.drop = nn.Dropout(p=0.2)
+        self.fc = nn.Linear(4 * 8 * 8, 1)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        x = x.flatten(1)
+        x = self.drop(x)
+        x = self.fc(x)
+        return x
+
+
+def init_model(model: nn.Module):
+    torch.manual_seed(42)
+    with torch.no_grad():
+        for p in model.parameters():
+            p.uniform_(-0.1, 0.1)
+
+
+def _print_grad_stats(prefix: str, model: nn.Module):
+    for name, param in model.named_parameters():
+        if param.grad is None:
+            print(f"{prefix} {name}.grad = None")
+        else:
+            print(f"{prefix} {name}.grad norm = {param.grad.norm().item():.8f}")
+
+
+def _collect_grad_norms(model: nn.Module):
+    result = {}
+    for name, param in model.named_parameters():
+        if param.grad is None:
+            result[name] = None
+        else:
+            result[name] = float(param.grad.norm().item())
+    return result
+
+
+def _is_valid_norm(v):
+    return (v is not None) and (not math.isnan(v)) and (not math.isinf(v))
+
+
+def _status_by_pair(torch4ms_v, torch_v):
+    if not _is_valid_norm(torch4ms_v):
+        return "BAD_INVALID"
+    if torch_v is None:
+        return "UNKNOWN"
+    # torch 明显有梯度而 torch4ms 几乎为 0，判为不一致
+    if abs(torch_v) > 1e-6 and abs(torch4ms_v) < 1e-8:
+        return "MISMATCH_ZERO_GRAD"
+    # 二者都非零时给一个粗粒度相对差异判断
+    denom = max(abs(torch_v), 1e-12)
+    rel_diff = abs(torch4ms_v - torch_v) / denom
+    if rel_diff < 0.15:
+        return "OK"
+    return "WARN_DIFF"
+
+
+def train_with_torch4ms_cnn():
+    print("=== torch4ms CNN path ===")
+    dataloader = get_dataloader()
+    criterion = nn.MSELoss()
+
+    model = SimpleCNN()
+    init_model(model)
+    optimizer = Torch4msOptimizer(optim.SGD(model.parameters(), lr=LR), model)
+
+    env = torch4ms.default_env()
+    last_grad_norms = {}
+    with env:
+        for epoch in range(NUM_EPOCHS):
+            for step, (inputs, labels) in enumerate(dataloader):
+                print(f"[torch4ms] epoch={epoch}, step={step}")
+                loss_wrapper = extract_and_wrap_loss_fn(model, criterion, inputs, labels)
+                loss_t4 = loss_wrapper.output
+                print(f"[torch4ms] loss = {loss_t4}")
+
+                loss_t4.backward(module=model)
+                _print_grad_stats("[torch4ms]", model)
+                last_grad_norms = _collect_grad_norms(model)
+
+                optimizer.step()
+                optimizer.zero_grad()
+
+    return last_grad_norms
+
+
+def train_with_torch_cnn():
+    print("=== torch CNN ===")
+    dataloader = get_dataloader()
+    criterion = nn.MSELoss()
+
+    model = SimpleCNN()
+    init_model(model)
+    optimizer = optim.SGD(model.parameters(), lr=LR)
+
+    last_grad_norms = {}
+    for epoch in range(NUM_EPOCHS):
+        for step, (inputs, labels) in enumerate(dataloader):
+            optimizer.zero_grad()
+            out = model(inputs)
+            loss = criterion(out, labels)
+            print(f"[torch] epoch={epoch}, step={step}, loss = {loss.item()}")
+            loss.backward()
+            _print_grad_stats("[torch]", model)
+            last_grad_norms = _collect_grad_norms(model)
+            optimizer.step()
+
+    return last_grad_norms
+
+
+def main():
+    t4_norms = train_with_torch4ms_cnn()
+    print()
+    torch_norms = train_with_torch_cnn()
+
+    print()
+    print("=== stability regression summary ===")
+    for name in sorted(set(t4_norms.keys()) | set(torch_norms.keys())):
+        t4_v = t4_norms.get(name)
+        th_v = torch_norms.get(name)
+        status = _status_by_pair(t4_v, th_v)
+        print(f"[summary] {name}: torch4ms={t4_v}, torch={th_v}, torch4ms_status={status}")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as exc:
+        import traceback
+
+        print(f"[ERROR] {exc}")
+        traceback.print_exc()
diff --git a/tests/torch4ms/test_train_transformer.py b/tests/torch4ms/test_train_transformer.py
new file mode 100644
index 000000000..a28d1f167
--- /dev/null
+++ b/tests/torch4ms/test_train_transformer.py
@@ -0,0 +1,161 @@
+import os
+os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
+
+import math
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+import torch4ms
+from torch4ms.autograd.ms_autograd_function import extract_and_wrap_loss_fn
+from torch4ms.optim import Torch4msOptimizer
+
+NUM_EPOCHS = 2
+LR = 0.01
+ATOL_REL = 0.2
+
+
+def get_dataloader():
+    torch.manual_seed(0)
+    # [batch, seq, in_dim]
+    x = torch.randn(4, 8, 16, dtype=torch.float32)
+    y = torch.randn(4, 1, dtype=torch.float32)
+    return [(x, y)]
+
+
+class TinyTransformerRegressor(nn.Module):
+    def __init__(self, in_dim=16, d_model=32, nhead=4, ff_dim=64, num_layers=2):
+        super().__init__()
+        self.in_proj = nn.Linear(in_dim, d_model)
+        layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=ff_dim,
+            dropout=0.0,
+            batch_first=True,
+            activation="relu",
+        )
+        self.encoder = nn.TransformerEncoder(layer, num_layers=num_layers)
+        self.head = nn.Linear(d_model, 1)
+
+    def forward(self, x):
+        x = self.in_proj(x)
+        x = self.encoder(x)
+        x = torch.mean(x, dim=1)
+        x = self.head(x)
+        return x
+
+
+def init_model(model: nn.Module):
+    torch.manual_seed(42)
+    with torch.no_grad():
+        for p in model.parameters():
+            p.uniform_(-0.1, 0.1)
+
+
+def _grad_norms(model: nn.Module):
+    norms = {}
+    for name, p in model.named_parameters():
+        if p.grad is None:
+            norms[name] = None
+        else:
+            norms[name] = float(p.grad.norm().item())
+    return norms
+
+
+def _print_grad_norms(prefix: str, norms):
+    for name, v in norms.items():
+        if v is None:
+            print(f"{prefix} {name}.grad = None")
+        else:
+            print(f"{prefix} {name}.grad norm = {v:.8f}")
+
+
+def _is_valid(v):
+    return (v is not None) and (not math.isnan(v)) and (not math.isinf(v))
+
+
+def _status_pair(t4, th):
+    if not _is_valid(t4):
+        return "BAD_INVALID"
+    if th is None:
+        return "UNKNOWN"
+    if abs(th) > 1e-6 and abs(t4) < 1e-8:
+        return "MISMATCH_ZERO_GRAD"
+    rel = abs(t4 - th) / max(abs(th), 1e-12)
+    return "OK" if rel < ATOL_REL else "WARN_DIFF"
+
+
+def train_with_torch4ms():
+    print("=== torch4ms TinyTransformer ===")
+    dataloader = get_dataloader()
+    criterion = nn.MSELoss()
+
+    model = TinyTransformerRegressor()
+    init_model(model)
+    optimizer = Torch4msOptimizer(optim.SGD(model.parameters(), lr=LR), model)
+
+    last_norms = {}
+    env = torch4ms.default_env()
+    with env:
+        for epoch in range(NUM_EPOCHS):
+            for step, (inputs, labels) in enumerate(dataloader):
+                print(f"[torch4ms] epoch={epoch}, step={step}")
+                loss_wrapper = extract_and_wrap_loss_fn(model, criterion, inputs, labels)
+                loss_t4 = loss_wrapper.output
+                print(f"[torch4ms] loss = {loss_t4}")
+
+                loss_t4.backward(module=model)
+                last_norms = _grad_norms(model)
+                _print_grad_norms("[torch4ms]", last_norms)
+
+                optimizer.step()
+                optimizer.zero_grad()
+    return last_norms
+
+
+def train_with_torch():
+    print("=== torch TinyTransformer ===")
+    dataloader = get_dataloader()
+    criterion = nn.MSELoss()
+
+    model = TinyTransformerRegressor()
+    init_model(model)
+    optimizer = optim.SGD(model.parameters(), lr=LR)
+
+    last_norms = {}
+    for epoch in range(NUM_EPOCHS):
+        for step, (inputs, labels) in enumerate(dataloader):
+            optimizer.zero_grad()
+            out = model(inputs)
+            loss = criterion(out, labels)
+            print(f"[torch] epoch={epoch}, step={step}, loss = {loss.item()}")
+            loss.backward()
+            last_norms = _grad_norms(model)
+            _print_grad_norms("[torch]", last_norms)
+            optimizer.step()
+    return last_norms
+
+
+def main():
+    t4_norms = train_with_torch4ms()
+    print()
+    th_norms = train_with_torch()
+
+    print()
+    print("=== transformer stability summary ===")
+    all_names = sorted(set(t4_norms.keys()) | set(th_norms.keys()))
+    for name in all_names:
+        t4 = t4_norms.get(name)
+        th = th_norms.get(name)
+        status = _status_pair(t4, th)
+        print(f"[summary] {name}: torch4ms={t4}, torch={th}, status={status}")
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as exc:
+        import traceback
+        print(f"[ERROR] {exc}")
+        traceback.print_exc()