Training of multi-talent model

[plbenveniste]

In this issue, I experiment training finetuning a pretrained Multi-Talent model using this repo: https://github.com/MIC-DKFZ/MultiTalent

[plbenveniste]

Code used:

conda create -n venv_multitalent python=3.10
conda activate venv_multitalent
conda install pytorch torchvision pytorch-cuda=12.1 -c pytorch -c nvidia

git clone https://github.com/MIC-DKFZ/MultiTalent.git
cd MultiTalent/
pip install -e .

I copied dataset 901 (from nnUnet experiments to the following location).

export nnUNet_raw="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_raw"
export nnUNet_results="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_results"
export nnUNet_preprocessed="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_preprocessed"

Then, preprocess the dataset to the MultiTalent format:

prepare_MT_training MT_data901 200 -d 901 --verify_dataset_integrity

[plbenveniste]

The problem with the repo is that it didn't give the updated weights of the model which would work with nnUNetv2. (issue opened here: MIC-DKFZ/MultiTalent#15)

Therefore, I tried the following:

• checkout main -> checked the main branch which uses nnunet (v1)
• in setup.py changed “sklearn” to “scikit-learn”
• then run pip install -e .
• then changed things in the dataset.json and folder organisation such as here: Training of a MEDNEXT model #31 (comment)
  • Error was File "/home/plbenveniste/miniconda3/envs/venv_multitalent/lib/python3.9/site-packages/nnunet/preprocessing/sanity_checks.py", line 110, in verify_dataset_integrity expected_test_identifiers = [i.split("/")[-1][:-7] for i in test_cases]
    changed it to -> expected_test_identifiers = [i[“image”].split("/")[-1][:-7] for i in test_cases]

Then ran:

export nnUNet_raw_data_base="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_raw"
export RESULTS_FOLDER="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_results"
export nnUNet_preprocessed="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_preprocessed"

nnUNet_plan_and_preprocess -t 901  -pl3d ExperimentPlanner3D_v21_Pretrained --planner2d None -overwrite_plans ~/net/ms-lesion-agnostic/multiTalent/pretrained_MT100/Task100_MultiTalent/MultiTalent_tainer_resenc_ddp__MultiTalent_resenc_bs4/plans.pkl -overwrite_plans_identifier PlansPretrainedMultiTalentResEncBs4  --verify_dataset_integrity

To train:

CUDA_VISIBLE_DEVICES=3 nnUNet_train 3d_fullres nnUNetTrainerV2_warmupsegheads 901 0 -p  nnUNetPlans_pretrained_PlansPretrainedMultiTalentResEncBs4_plans_3D  -pretrained_weights  ~/net/ms-lesion-agnostic/multiTalent/pretrained_MT100/Task100_MultiTalent/MultiTalent_tainer_resenc_ddp__MultiTalent_resenc_bs4/all/model_final_checkpoint.model

I am still having issues because I think that the plans used is not the one corresponding to the model saved. Need to look into this,
Error message is the following:

Traceback (most recent call last):
  File "/home/plbenveniste/miniconda3/envs/venv_multitalent/bin/nnUNet_train", line 33, in <module>
    sys.exit(load_entry_point('nnunet', 'console_scripts', 'nnUNet_train')())
  File "/data/users/pierrelouis/ms-lesion-agnostic/multiTalent/MultiTalent/nnunet/run/run_training.py", line 162, in main
    trainer.initialize(not validation_only)
  File "/data/users/pierrelouis/ms-lesion-agnostic/multiTalent/MultiTalent/nnunet/training/network_training/nnUNet_variants/pretraining/nnUNetTrainerV2_warmup.py", line 82, in initialize
    super().initialize(training, force_load_plans)
  File "/data/users/pierrelouis/ms-lesion-agnostic/multiTalent/MultiTalent/nnunet/training/network_training/nnUNetTrainerV2.py", line 123, in initialize
    self.initialize_network()
  File "/data/users/pierrelouis/ms-lesion-agnostic/multiTalent/MultiTalent/nnunet/training/network_training/nnUNetTrainerV2.py", line 156, in initialize_network
    self.network = Generic_UNet(self.num_input_channels, self.base_num_features, self.num_classes,
  File "/data/users/pierrelouis/ms-lesion-agnostic/multiTalent/MultiTalent/nnunet/network_architecture/generic_UNet.py", line 301, in __init__
    self.conv_kwargs['kernel_size'] = self.conv_kernel_sizes[num_pool]
IndexError: list index out of range

[plbenveniste]

They answered the issue that I opened MIC-DKFZ/MultiTalent#15 (comment) and shared the weights of the model trained with nnUNetv2 (https://zenodo.org/records/13753413). 😃

[plbenveniste]

I finally succeeded!! Here is how I did it.
In am fine-tuning the model Dataset617_nativect/MultiTalent_trainer_4000ep__nnUNetResEncUNetL1x1x1_Plans_bs24__3d_fullres on my dataset901.

conda activate venv_nnunet_new

export nnUNet_raw="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_raw"
export nnUNet_results="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_results"
export nnUNet_preprocessed="/home/plbenveniste/net/ms-lesion-agnostic/multiTalent/nnUNet_preprocessed"

nnUNetv2_plan_and_preprocess -d 901 -c 3d_fullres --verify_dataset_integrity

nnUNetv2_preprocess -d 901 -plans_name nnUNetResEncUNetL1x1x1_Plans_bs24

CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 901 3d_fullres 0 -pretrained_weights /home/plbenveniste/net/ms-lesion-agnostic/multiTalent/pretrained_model/Dataset617_nativect/MultiTalent_trainer_4000ep__nnUNetResEncUNetL1x1x1_Plans_bs24__3d_fullres/fold_all/checkpoint_final.pth -tr nnUNetTrainerDiceCELoss_noSmooth -p nnUNetResEncUNetL1x1x1_Plans_bs24

I had the following issue:

 File "/home/plbenveniste/miniconda3/envs/venv_nnunet_new/lib/python3.9/site-packages/torch/_dynamo/output_graph.py", line 1465, in _call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e) from e
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
AssertionError: increase TRITON_MAX_BLOCK['X'] to 4096

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

It didn't work with the following code block in nnUNetTrainerDiceLoss.py:

import torch._dynamo
torch._dynamo.config.suppress_errors = True

To solve it did the following (the cloning is only because I wanted to preserve my old environment for other scripts)

conda create --name venv_nnunet_multitalent_torch_modified --clone venv_nnunet_new
vim /home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/_inductor/runtime/hints.py

Then I changed X in line … to (source: pytorch/pytorch#145074 (comment)):

TRITON_MAX_BLOCK = {
    "X": 4096,
    "Y": 1024,
    "Z": 1024,
    "R": 4096 * 16,  # * 16 is multi-kernel only
}

Then I had a memory problem: So I told it to use 2 GPUS:

CUDA_VISIBLE_DEVICES=2,1 nnUNetv2_train 901 3d_fullres 0 -pretrained_weights /home/plbenveniste/net/ms-lesion-agnostic/multiTalent/pretrained_model/Dataset617_nativect/MultiTalent_trainer_4000ep__nnUNetResEncUNetL1x1x1_Plans_bs24__3d_fullres/fold_all/checkpoint_final.pth -tr nnUNetTrainerDiceCELoss_noSmooth -p nnUNetResEncUNetL1x1x1_Plans_bs24 -num_gpus 2

But I still memory problem, so I reduced batch_size to 6 in nnUNetResEncUNetL1x1x1_Plans_bs24.json

[plbenveniste]

Process crashed with the following error message:

Epoch 453
Current learning rate: 0.00581
2025-04-12 03:48:04.017921:
2025-04-12 03:48:04.077450: Epoch 453
2025-04-12 03:48:04.078914: Current learning rate: 0.00581
[rank1]:[E412 04:21:24.254321149 ProcessGroupNCCL.cpp:616] [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=4128661, OpType=ALLGATHER, NumelIn=3, NumelOut=6, Timeout(ms)=600000) ran for 600060 milliseconds before timing out.
[rank1]:[E412 04:21:25.573289685 ProcessGroupNCCL.cpp:1785] [PG ID 0 PG GUID 0(default_pg) Rank 1] Exception (either an error or timeout) detected by watchdog at work: 4128661, last enqueued NCCL work: 4128690, last completed NCCL work: 4128660.
[rank1]:[E412 04:21:25.573431017 ProcessGroupNCCL.cpp:1834] [PG ID 0 PG GUID 0(default_pg) Rank 1] Timeout at NCCL work: 4128661, last enqueued NCCL work: 4128690, last completed NCCL work: 4128660.
[rank1]:[E412 04:21:25.573459227 ProcessGroupNCCL.cpp:630] [Rank 1] Some NCCL operations have failed or timed out. Due to the asynchronous nature of CUDA kernels, subsequent GPU operations might run on corrupted/incomplete data.
[rank1]:[E412 04:21:25.573475777 ProcessGroupNCCL.cpp:636] [Rank 1] To avoid data inconsistency, we are taking the entire process down.
[rank1]:[E412 04:21:25.687311148 ProcessGroupNCCL.cpp:1595] [PG ID 0 PG GUID 0(default_pg) Rank 1] Process group watchdog thread terminated with exception: [Rank 1] Watchdog caught collective operation timeout: WorkNCCL(SeqNum=4128661, OpType=ALLGATHER, NumelIn=3, NumelOut=6, Timeout(ms)=600000) ran for 600060 milliseconds before timing out.
Exception raised from checkTimeout at /opt/conda/conda-bld/pytorch_1729647350383/work/torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:618 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x96 (0x7f5a1af3a446 in /home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/lib/libc10.so)
frame #1: c10d::ProcessGroupNCCL::WorkNCCL::checkTimeout(std::optional<std::chrono::duration<long, std::ratio<1l, 1000l> > >) + 0x282 (0x7f5a1c2a07f2 in /home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/lib/libtorch_cuda.so)
frame #2: c10d::ProcessGroupNCCL::watchdogHandler() + 0x233 (0x7f5a1c2a7c33 in /home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/lib/libtorch_cuda.so)
frame #3: c10d::ProcessGroupNCCL::ncclCommWatchdog() + 0x14d (0x7f5a1c2a969d in /home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/lib/libtorch_cuda.so)
frame #4: <unknown function> + 0x145c0 (0x7f5a6c1dd5c0 in /home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/lib/libtorch.so)
frame #5: <unknown function> + 0x8609 (0x7f5a75a6f609 in /lib/x86_64-linux-gnu/libpthread.so.0)
frame #6: clone + 0x43 (0x7f5a7583a133 in /lib/x86_64-linux-gnu/libc.so.6)

W0412 04:21:32.117290 1930420 site-packages/torch/multiprocessing/spawn.py:160] Terminating process 1930487 via signal SIGTERM
Traceback (most recent call last):
  File "/home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/bin/nnUNetv2_train", line 8, in <module>
    sys.exit(run_training_entry())
  File "/data/users/pierrelouis/ms-lesion-agnostic/nnunet_experiments/nnUNet/nnunetv2/run/run_training.py", line 275, in run_training_entry
    run_training(args.dataset_name_or_id, args.configuration, args.fold, args.tr, args.p, args.pretrained_weights,
  File "/data/users/pierrelouis/ms-lesion-agnostic/nnunet_experiments/nnUNet/nnunetv2/run/run_training.py", line 178, in run_training
    mp.spawn(run_ddp,
  File "/home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 328, in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method="spawn")
  File "/home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 284, in start_processes
    while not context.join():
  File "/home/plbenveniste/miniconda3/envs/venv_nnunet_multitalent_torch_modified/lib/python3.9/site-packages/torch/multiprocessing/spawn.py", line 184, in join
    raise ProcessExitedException(
torch.multiprocessing.spawn.ProcessExitedException: process 1 terminated with signal SIGABRT

[plbenveniste]

I retrained it on a single GPU and reduced the batch further to 2 so that the experiment is comparable to my other trainings.

CUDA_VISIBLE_DEVICES=2 nnUNetv2_train 901 3d_fullres 0 -pretrained_weights /home/plbenveniste/net/ms-lesion-agnostic/multiTalent/pretrained_model/Dataset617_nativect/MultiTalent_trainer_4000ep__nnUNetResEncUNetL1x1x1_Plans_bs24__3d_fullres/fold_all/checkpoint_final.pth -tr nnUNetTrainerDiceCELoss_noSmooth -p nnUNetResEncUNetL1x1x1_Plans_bs24

Training completed! Evaluation of the model in progress!

CUDA_VISIBLE_DEVICES=3 nnUNetv2_predict -i ~/net/ms-lesion-agnostic/multiTalent/nnUNet_raw/Dataset901_msLesionAgnostic/imagesTs/  -o ~/net/ms-lesion-agnostic/multiTalent/nnUNet_results/Dataset901_msLesionAgnostic/nnUNetTrainerDiceCELoss_noSmooth__nnUNetResEncUNetL1x1x1_Plans_znorm_bs24__3d_fullres/predictionImagesTsFold0ChkFinal   -d 901 -c 3d_fullres -f 0 -chk checkpoint_final.pth -tr nnUNetTrainerDiceCELoss_noSmooth -p nnUNetResEncUNetL1x1x1_Plans_znorm_bs24

[plbenveniste]

The model was evaluated, and the results are much better!