DiNTS forward() breaks the torch.compile graph on every topology branch
维护者通常 1 天内回复
评估
- 难度
- 4/5
- 预计耗时
- 3-5 天
- 新手友好度
- 68/100
- Issue 类型
- 缺陷
- 描述清晰度
- 描述清楚
- 活跃度
- 活跃
调研方向
Start in monai/networks/nets/dints.py, focusing on DiNTS.forward and TopologyInstance.forward and their node_a/arch_code_a control flow. Run the provided TorchDynamo reproducer first, then verify that the topology behavior is preserved and compilation reaches the stated 1 graph with 0 breaks for the sample configuration.
由索引模型根据 Issue 内容生成。
描述
Describe the bug
DiNTS.forward and TopologyInstance.forward pick which cells to run by indexing the node_a / arch_code_a tensors inside Python control flow:
# monai/networks/nets/dints.py
if self.node_a[0][d]: ...
elif self.node_a[blk_idx + 1][res_idx]: ...
for res_idx, activation in enumerate(self.arch_code_a[blk_idx].data): ...
Each is a data-dependent tensor read, so TorchDynamo cannot constant-fold the branch and splits the graph instead: a 6-block / 3-depth DiNTS compiles into 14 graphs with 13 breaks. These flags are fixed once a model is deployed, so the breaks are pure overhead — they block fusion across the network and scale with num_blocks, so deeper searched architectures fragment further.
To Reproduce
-
Install MONAI from
dev, with a torch 2.x providing TorchDynamo. -
Run:
import torch, torch._dynamo as dyn from monai.networks.nets.dints import DiNTS, TopologyInstance grid = dict(channel_mul=0.2, num_blocks=6, num_depths=3, use_downsample=True, spatial_dims=3, device="cpu") m = DiNTS(dints_space=TopologyInstance(**grid), in_channels=1, num_classes=2, spatial_dims=3, use_downsample=True).eval() expl = dyn.explain(m)(torch.randn(1, 1, 32, 32, 32)) print(expl.graph_count, expl.graph_break_count) -
Observed:
14 13. (expl.break_reasonsis not populated on this torch version, so only the counts are available.)
Expected behavior
The topology flags are constant for the lifetime of a deployed model, so the branches should be constant-foldable and the network should compile to 1 0 — one graph, no breaks.
Screenshots
N/A — the reproducer output above is the full evidence.
Environment
MONAI version: 1.6.1rc0+8.g860506514
MONAI rev id : 860506514fb1fba41db1578e0eeacfa952583a13 (dev)
Pytorch version: 2.12.1+rocm7.1
Numpy version: 2.5.2
Reproduced on CPU. The graph breaks are introduced at Dynamo trace time, so they are independent of
the accelerator and of the torch build — this is not platform-specific.
python -c "import monai; monai.config.print_debug_info()"
================================
Printing MONAI config...
================================
MONAI version: 1.6.1rc0+8.g860506514
Numpy version: 2.5.2
Pytorch version: 2.12.1+rocm7.1
MONAI flags: HAS_EXT=False, USE_COMPILED=False, USE_META_DICT=False
MONAI rev id: 860506514fb1fba41db1578e0eeacfa952583a13
MONAI __file__: /home/user/upstream_monai/upstream-fork-monai/monai/__init__.py
Optional dependencies:
Pytorch Ignite version: 0.5.5
ITK version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: 5.4.2
scikit-image version: 0.26.0
scipy version: 1.18.1
Pillow version: 12.3.0
Tensorboard version: NOT INSTALLED or UNKNOWN VERSION.
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: 0.27.1+rocm7.1
tqdm version: 4.70.1
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: NOT INSTALLED or UNKNOWN VERSION.
pandas version: NOT INSTALLED or UNKNOWN VERSION.
einops version: 0.8.2
transformers version: NOT INSTALLED or UNKNOWN VERSION.
mlflow version: NOT INSTALLED or UNKNOWN VERSION.
pynrrd version: NOT INSTALLED or UNKNOWN VERSION.
clearml version: NOT INSTALLED or UNKNOWN VERSION.
================================
Printing system config...
================================
`psutil` required for `print_system_info`
================================
Printing GPU config...
================================
Num GPUs: 1
Has CUDA: True
CUDA version: None
cuDNN enabled: True
NVIDIA_TF32_OVERRIDE: None
TORCH_ALLOW_TF32_CUBLAS_OVERRIDE: None
cuDNN version: 3005001
Current device: 0
Library compiled for CUDA architectures: ['gfx900', 'gfx906', 'gfx908', 'gfx90a', 'gfx942', 'gfx1030', 'gfx1100', 'gfx1101', 'gfx1102', 'gfx1200', 'gfx1201', 'gfx950', 'gfx1150', 'gfx1151']
GPU 0 Name: AMD Instinct MI300X
GPU 0 Is integrated: False
GPU 0 Is multi GPU board: False
GPU 0 Multi processor count: 304
GPU 0 Total memory (GB): 192.0
GPU 0 CUDA capability (maj.min): 9.4
- 主要语言
- Python
- 星标
- 8.8k
- 派生
- 1.6k
- 平均合并
- 22 小时 56 分钟
- 30 天内合并 PR
- 7
环境准备
- 提供 Dockerfile 或 Docker Compose 文件
- 有 Pull Request 模板
- 阅读贡献指南
从这里开始
- 先读完整个 Issue，再读项目的贡献指南。
- 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
- Fork 仓库，在一个分支上完成修改。
- 提交 Pull Request，并在描述里引用这个 Issue 编号。
Project-MONAI/MONAI 的其他 Issue
-
DiNTS: missing graph-break regression test (follow-up to #9145)可能已有人在做 @nilapate 于 3 天前认领。 未关闭
难度 2/5 1-3 小时 新手友好度 66/100
Project-MONAI/MONAI#9159 ·
维护者通常 1 天内回复
-
ProbMapProducer raises IndexError on numpy 2.x: SlidingPatchWSIDataset emits float locations可能已有人在做 @nilapate 于 14 天前认领。 未关闭
难度 2/5 1-3 小时 新手友好度 82/100
Project-MONAI/MONAI#9135 ·
维护者通常 1 天内回复
-
RandGridPatch crashes for documented 0 or None patch_size entries可能已有人在做 @Rajioba1 于 70 天前认领。 未关闭
难度 2/5 1-3 小时 新手友好度 78/100
Project-MONAI/MONAI#9046 ·
维护者通常 1 天内回复
-
Incorrect type hint in `pad_func`可能已有人在做 @Harshal875 于 76 天前认领。 未关闭
难度 1/5 1 小时以内 新手友好度 88/100
Project-MONAI/MONAI#9026 ·
维护者通常 1 天内回复
-
Remove MMAR Code可能已有人在做 @KumoLiu 于 2 天前认领。 未关闭
Project-MONAI/MONAI#9162 · 已指派 2 人 ·
维护者通常 1 天内回复
查看 Project-MONAI/MONAI 的全部 Issue
相似的 Issue
-
难度 2/5 1-3 小时 新手友好度 70/100
-
good first issue No-AI
难度 2/5 1-3 小时 新手友好度 69/100
-
documentation
难度 2/5 1-3 小时 新手友好度 78/100
-
🐛 Bug: Calendar work items cannot be opened on mobile可能已有人在做 关联的 PR 仍在进行中或已合并。 未关闭
难度 1/5 1 小时以内 新手友好度 75/100
维护者通常 1 天内回复
-
bug
难度 2/5 1-3 小时 新手友好度 72/100
维护者通常 2 天内回复