Checklist

• 1. I have searched the existing issues (https://gitcode.com/mindspore/hyper-parallel/issues)
• 2. I have read the relevant documentation.
• 3. I have created a minimal reproduction case that clearly demonstrates the issue, including a complete code example and the error message with full traceback and error logs.

🐛 Describe the bug

*真实 QwenImageTransformer2DModel（Diffusers）单卡 NPU 训练正常，但 2 卡 DDP/FSDP 均在启动阶段被系统 SIGKILL (-9)，疑似 OOM。
单卡正常：
运行指令：
export HYPER_PARALLEL_PLATFORM=torch
export RANK=0 WORLD_SIZE=1 MASTER_ADDR=127.0.0.1 MASTER_PORT=29500
python scripts/train_dit.py examples/dit_qwen_image/train.yaml
输出：
100 Step Loss Curve:
Step 10: loss=1.3253, grad_norm=1.1204
Step 20: loss=1.2697, grad_norm=1.0466
Step 30: loss=1.2316, grad_norm=0.9812
Step 40: loss=1.2025, grad_norm=0.9304
Step 50: loss=1.1733, grad_norm=0.8647
Step 60: loss=1.1649, grad_norm=0.8617
Step 70: loss=1.1173, grad_norm=0.7675
Step 80: loss=1.0763, grad_norm=0.6891
Step 90: loss=1.0644, grad_norm=0.6610
Step 100: loss=1.0468, grad_norm=0.6196
Training completed

2卡报错：
运行指令：
torchrun --nproc_per_node=2 scripts/train_dit.py examples/dit_qwen_image/train.yaml

错误日志：
完整的错误日志太长，这里我贴一些我认为比较关键的片段，由数字顺序逐一列出
①
Root Cause (first observed failure):
[0]:
exitcode : -9 (pid: 537)
traceback : Signal 9 (SIGKILL) received by PID 537

②
[INFO] httpx: HTTP Request: HEAD https://hf-mirror.com/Qwen/Qwen-Image/resolve/main/transformer/config.json "HTTP/1.1 307 Temporary Redirect"
The config attributes {'pooled_projection_dim': 768} were passed ...

③
[WARNING] ... Unknown config key 'in_channels' for ModelConfig ignored.
[WARNING] ... Unknown config key 'height' for ModelConfig ignored.
[WARNING] ... Unknown config key 'patch_size' for ModelConfig ignored.

④
Config: data.type=dummy_dit, model.name=qwen_image_dit, model.num_hidden_layers=1

⑤
[rank0]: Model built on npu: QwenImageDiT
[rank0]: Truncated transformer blocks: 60 -> 2
[rank0]: Model params: 720.2M
...
W0622 ... Sending process ... closing signal SIGTERM
E0622 ... failed (exitcode: -9) local_rank: 1 ...
Root Cause: Signal 9 (SIGKILL) received by PID ...

已尝试但均失败的方案：

• 降低分辨率到 128×128 / 64×64
• 关闭 FSDP，只用 DDP
• 砍层到 num_layers=1
• init_device 改为 meta（触发 Cannot copy out of meta tensor）
• parallelize_fn 手动 fully_shard 分片
• 开启 gradient_checkpointing

Expected behavior

2 卡 DDP 应正常训练，loss 与单卡对齐（100 step 平均 diff < 5e-3）。

Additional context

我尝试运行以下脚本：

python -c "
import torch
import torch_npu
from hyper_parallel.models.qwen_image_dit import QwenImageDiT

cfg = {'num_layers': 2}
model = QwenImageDiT(cfg).to('npu:0')

# 打印显存占用
print('Allocated:', torch_npu.npu.memory_allocated() / 1024**3, 'GB')
print('Reserved:', torch_npu.npu.memory_reserved() / 1024**3, 'GB')

# 模拟一个 step
latent = torch.randn(1, 64, 32, 32).to('npu:0')
t = torch.randint(0, 1000, (1,)).to('npu:0')
cond = torch.randn(1, 77, 3584).to('npu:0')
target = torch.randn(1, 64, 32, 32).to('npu:0')

out = model(latent, t, cond, target_noise=target)
out.loss.backward()

print('After forward+backward:')
print('Allocated:', torch_npu.npu.memory_allocated() / 1024**3, 'GB')
print('Reserved:', torch_npu.npu.memory_reserved() / 1024**3, 'GB')
"

得到的输出是：

/usr/local/python3.11.14/lib/python3.11/site-packages/huggingface_hub/utils/_validators.py:205: UserWarning: The `local_dir_use_symlinks` argument is deprecated and ignored in `hf_hub_download`. Downloading to a local directory does not use symlinks anymore.
  warnings.warn(
The config attributes {'pooled_projection_dim': 768} were passed to QwenImageTransformer2DModel, but are not expected and will be ignored. Please verify your config.json configuration file.
Truncated transformer blocks: 60 -> 2
Model params: 720.2M
Allocated: 2.6829380989074707 GB
Reserved: 2.76171875 GB
[W621 14:46:24.632268673 TensorFactories.cpp:340] Warning: Cannot create tensor with interal format while allow_internel_format=False, tensor will be created with base format. (function operator())
After forward+backward:
Allocated: 5.058430194854736 GB
Reserved: 5.681640625 GB

单卡显存诊断：模型加载后 ~2.7GB，1 step forward+backward 后 ~5.7GB。64GB HBM 理论上非常充裕。
2 卡 OOM 发生在模型构建完成后、训练循环开始前，无 Python traceback，直接系统级 SIGKILL。
关联 PR: #876 (feat/dit-trainer-v2)
关联 Issue: #2100

Environment info

复现命令：

git clone https://gitcode.com/luguobin-2026/hyper-parallel.git
cd hyper-parallel && git checkout feat/dit-trainer-v2
pip install -e . && pip install diffusers torchdata
export HYPER_PARALLEL_PLATFORM=torch
torchrun --nproc_per_node=2 scripts/train_dit.py examples/dit_qwen_image/train.yaml

Thanks for contributing 🎉!