MedNeXt (versions 1 and 2) experiments, VRAM and runtime reporting

[Shrajan]

Hi nnUNet team,

I have some questions regarding the experimental setup and reported results in the MedNeXt-v1 [1], MedNeXt-v2 [2] and nnU-Net Revisited[3] papers. I have been trying to carefully reproduce the experiments using the official nnU-Net framework, MedNeXt repository, and settings.

1. I would like to ask about GPU VRAM usage reported in Table 1 of the nnU-Net Revisited paper [3]. I reproduced the experiments using the official code, same presets and patch size, checkpointing and AMP enabled, and on an A100 40 GB GPU. The runtime scaling between ResEnc L and MedNeXt L k3 is very similar to what is reported in the paper, so I believe my setup is aligned with the benchmark. However, the allocated memory for MedNeXt L k3 is noticeably higher in my experiments, and in my case MedNeXt uses more memory than ResEnc L, while in Table 1 [3] it is reported to use less. MedNeXt L k5 also utilises similar GPU VRAM as MedNeXt L k3, but of course, takes longer to run.

For ResEnc UNet L on A100 40 GB, my results match closely:

Training time ≈ 28 hours (paper: 35 hours)
Allocated CUDA memory ≈ 21.5 GB (paper: 22.7 GB)
Reserved CUDA memory ≈ 22.6 GB
Nvidia-smi ≈ 23.5 GB

However, for MedNeXt L k3 on A100 40 GB, I observe:

Training time ≈ 60 hours (paper: 68 hours)
Allocated CUDA memory ≈ 24.9 GB (paper: 17.3 GB)
Reserved CUDA memory ≈ 33.3 GB
Nvidia-smi ≈ 34.6 GB

I even trianed the MedNeXt models with a RTX 5090 (32 GB) GPU, and even then the VRAM usage did not go below 20 GB.

2. Could you please clarify:
   a) Whether any additional settings were used for MedNeXt [3] in the benchmark?
   b) How exactly VRAM was measured for Table 1 (allocated vs reserved memory or Nvidia-smi)[3]?
   c) Whether a slightly different MedNeXt configuration or commit version was used internally?
   d) I also wanted to confirm whether the reported runtime in Table 1 [3] is measured per fold (average), right? Not aggregated over all 5 cross-validation folds.

3. Finally, I noticed that the DSC values of nnUNet, ResEnc UNet L and MedNeXt-v1 trained from scratch differ slightly between the MedNeXt-v1 [1] and MedNeXt-v2 [2] papers (for example on AMOS, KiTS and ACDC). I understand that different splits or protocols may have been used, but I would appreciate clarification on this as well.

Thank you very much for your time and help.

References:
[1] MedNeXt v1 - https://conferences.miccai.org/2023/papers/410-Paper1656.html
[2] MedNeXt v2 - https://arxiv.org/pdf/2512.17774
[3] nnU-Net Revisited - https://papers.miccai.org/miccai-2024/paper/2847_paper.pdf

[FabianIsensee]

@saikat-roy

[FabianIsensee]

One note on VRAM consumption of traninigs. cudnn.benchmark is always active in nnU-Net. This will try different kernels for each operation at the start of the training and pick the one that is fastest in the current setup. What kernel is picked can change the amount of VRAM used in training. This very much depends on the GPU architecture, cuDNN version and probably also driver version. All we could do for the paper is report how the numbers were at the time of running these experiments.
Note that sometimes certain kernels that use a lot of VRAM are not picked on smaller GPUs but will be picked on GPUs with more VRAM, which is why a training that requires > 24GB on GPU A might still be able to run on GPU B even if B only has < 24GB of VRAM.

[FabianIsensee]

MedNext uses gradient checkpointing. This could also play a role. Again, let me refer to Saikat because he is the expert in this

[Shrajan]

Hi @FabianIsensee,

Thank you very much for the detailed explanation. Indeed, I have also observed that the system settings (GPU architecture, CPU cores, etc.) and package versions impact the overall VRAM consumption and training time, and this is apparent for MedNeXt models.

Please note that my original post was an attempt to do a fair comparision under the same GPU settings mentioned in your publised papers [1,2,3].
i) The experimental results listed in my original post for MedNeXt models were obtained with gradient checkpointing.
ii) The experimental results listed in the original post for ResEnc L, MedNeXt L k3 and MedNeXt L k5 were the best ones (lowest possible VRAM consumption and training times) across all the different versions of cuDNN, NVIDIA drivers, and PyTorch packages. As a final note, the best results were obtained with the latest versions.

The following are a few things I left out in my original post that were part of my experiments:

1. I compared different versions of cuDNN, NVIDIA drivers, and PyTorch (2.2, 2.3, 2.7, 2.10) packages.
2. I have also tried to train ResEnc L, MedNeXt L k3 and MedNeXt L k5 with RTX 3090 (24 GB) and RTX 4090 (24 GB), but was able to fit only ResEnc L in the two GPUs. Training with MedNeXt L k3 and MedNeXt L k5 would throw out of memory errors even with gradient checkpointing enabled.
3. Without gradient checkpointing, the MedNeXt models would not fit in a A100 40 GB GPU - I tried with a A100 80 GB GPU that consumed approximately 52 GB VRAM with a shorter training duration. However, I did not try with an NVIDIA RTX 6000 Ada (48 GB) GPU.

Therefore, some more information from your end is crucial to complete this comparsion and get a better picture. Also, it would be helpful if I could get answers for questions 2(b) and 2(d) from the original post.

Thank you very much.

[saikat-roy]

Hi @Shrajan, sorry for the late reply. Can I ask you what your patch size was for the VRAM calculations in your work? You mentioned it was the same - was it specifically 128x128x128?

Also, one thing I have seen in MedNeXt that might not be an issue in ResEncL is fragmentation. Specifically, I had started using PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to overcome some of the issues with PyTorch reserved memory. Could you do this and see if there's any effect to a) your VRAM usage and b) You inability to run on 40 GB A100s.

a) Whether any additional settings were used for MedNeXt [3] in the benchmark?
b) How exactly VRAM was measured for Table 1 (allocated vs reserved memory or Nvidia-smi)[3]?
c) Whether a slightly different MedNeXt configuration or commit version was used internally?
d) I also wanted to confirm whether the reported runtime in Table 1 [3] is measured per fold (average), right? Not aggregated over all 5 cross-validation folds.

I think regarding the benchmarking in nnUNet Revisited itself, @constantinulrich is the best person to answer the settings when computing memory usage. I know that project used the original MedNeXt repository (based on nnUNet-v1) and not for example, porting the MedNeXt architecture to the newer nnUNet-v2. But that in principle should not affect the computation of VRAM usage during training, particularly via nvidia-smi.

Additionally, the results in the MedNeXt-v2 paper for your mentioned datasets are fold 0 results for a smaller ablation experiment, while those in the nnUNet Revisited are 5 fold CV results. That is why they differ. I also used a older version of AMOS with 200 samples, so that yet again is another point of difference.

[Shrajan]

Hi @saikat-roy, thanks for getting back to me.

Yes, I used the same patch-size of 128x128x128, as mentioned in your papers. I will try a few experiments with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True, and get back to you on that one.

Regarding the experiments pertaining to MedNeXt models, I trained with three implementations: (i) original MedNeXt repository (based on nnUNet-v1), (ii) TriALS (based on nnUNet-v2), and (iii) my own version of MedNeXt (have not implemented gradient checkpointing script) using dynamic-network-architectures based on nnUNet-v2. All three implementations have similar results and VRAM consumption when trained without gradient checkpointing. Furthermore, models from original MedNeXt repository and TriALS have near identical results and VRAM consumption when trained with gradient checkpointing enabled.

Finally, I was aware that you used the older version of AMOS with 200 samples (CT modality: only train, no validation) in both MedNeXt-v1 and v2 papers. That said, thank you clarifying the usage of fold 0 results listed in Tables 2 and 3 of the MedNeXt-v2 paper, as I was unable to find this information in the current version.

[saikat-roy]

Ah yes you're right, we do mention the 200 samples in the dataset size of AMOS in the MedNeXt-v2 paper, but we indeed don't specifically highlight why it is lower. The paper is also a preprint currently and will probably go through a few more iterations in the future.