Help on preparing dataset for finetuning MRI DICOM 3D data

[Ribin-Baby]

Hi, I was trying to finetune 'google/medgemma-1.5-4b-it' model on 3D brain MRI DICOM data. In the repository i can find examples on finetuning on single image slices but to finetune on 3D data, here is how i have curated the data,

{'images': [<PIL.Image.Image image mode=RGB size=224x224>,
  <PIL.Image.Image image mode=RGB size=224x224>,
  <PIL.Image.Image image mode=RGB size=224x224>,
  <PIL.Image.Image image mode=RGB size=224x224>,
  <PIL.Image.Image image mode=RGB size=224x224>,
  <PIL.Image.Image image mode=RGB size=224x224>,
  <PIL.Image.Image image mode=RGB size=224x224>,
  <PIL.Image.Image image mode=RGB size=224x224>],
 'messages': [{'role': 'system',
   'content': [{'type': 'text', 'text': 'You are an expert radiologist.'}]},
  {'role': 'user',
   'content': [{'type': 'image'},
    {'type': 'text', 'text': 'Slice 1'},
    {'type': 'image'},
    {'type': 'text', 'text': 'Slice 2'},
    {'type': 'image'},
    {'type': 'text', 'text': 'Slice 3'},
    {'type': 'image'},
    {'type': 'text', 'text': 'Slice 4'},
    {'type': 'image'},
    {'type': 'text', 'text': 'Slice 5'},
    {'type': 'image'},
    {'type': 'text', 'text': 'Slice 6'},
    {'type': 'image'},
    {'type': 'text', 'text': 'Slice 7'},
...
    {'type': 'text',
     'text': 'Analyze this brain MRI and provide the most likely diagnosis in one short phrase.'}]},
  {'role': 'assistant',
   'content': [{'type': 'text',
     'text': 'acute ischemic stroke left precentral gyrus'}]}]}    

The issue i was facing is was not able to include complete slices for finetuning, if i take all slices i was getting Out of memory error even with 2 L40 GPU with 48GB VRAM.

Can you provide an example training script for finetuning for DICOM dataset, is the way i prepared data was correct or not ?

[srikantak-221]

Hi @Ribin-Baby ,

Thanks for sharing the details of your setup and the dataset structure.

The way you prepared the dataset is generally aligned with how MedGemma expects multimodal inputs. MedGemma supports prompts that interleave images and text, and providing the MRI slices as a list under images while referencing them through "type": "image" entries inside the messages structure is a valid approach. Adding slice labels such as “Slice 1”, “Slice 2”, etc., also helps preserve ordering information so the model can interpret the images as part of a sequence rather than unrelated images.

The out-of-memory error you encountered is expected when attempting to pass an entire 3D MRI volume to the model during fine-tuning. Each image slice is processed by the vision encoder and converted into visual tokens before being fed into the language model. When many slices are included, the number of tokens in the sequence increases quickly. During training, the transformer must store activations and gradients for all tokens across multiple layers, which causes GPU memory usage to grow significantly. As a result, including all slices from a full MRI series can exceed available memory even on GPUs with large VRAM.

In practice, when adapting vision-language models to volumetric medical imaging, it is common to reduce the number of slices per example rather than using the full scan. Typical approaches include uniformly sampling a subset of slices from the volume (for example 16–64 slices), removing empty slices at the top and bottom of the scan, or selecting slices that contain the most anatomical information. These preprocessing steps significantly reduce the token length while still preserving the relevant clinical context for the model.

On the training side, parameter-efficient fine-tuning methods are usually recommended for models of this size. Instead of updating all model weights, techniques such as LoRA combined with low-precision loading (for example 4-bit quantisation) allow you to train lightweight adapter layers while keeping the base model frozen. This greatly reduces GPU memory usage and makes fine-tuning feasible even when working with multimodal inputs. Additional techniques such as gradient checkpointing and small per-device batch sizes are also commonly used to keep memory usage manageable.

For a complete reference on running and fine-tuning MedGemma models, you may find the official documentation helpful. It includes supported workflows, example usage, and recommended training setups.

Please refer here for Generic Fine tuning Guide.

Also refer here for a Quick start guide on DICOM.

Overall, your dataset format is conceptually correct, and the main adjustment needed is controlling the number of slices included per training example and using memory-efficient fine-tuning techniques. This should help avoid the out-of-memory issue while still allowing the model to learn from the MRI data.

[Ribin-Baby]

In GCP model garden we have this model medgemma-1.5-4b-it-dicom available to deploy which can handle 3D DICOM datas.
https://console.cloud.google.com/vertex-ai/publishers/google/model-garden/medgemma
Which can handle upto 125 slices on 2 x L4 NVIDIA GPUs (vLLM 36k context 125 image limit (2 NVIDIA_L4; g2-standard-24).

How it was able to handle 3D data ? can we get the inferencing script for it ?

[srikantak-221]

Hi again, thanks for the follow-up question.
This is a good observation, and the difference you noticed mainly comes from how training workloads differ from inference workloads, as well as how the serving stack is optimised in the Model Garden deployment.

When you fine-tune a model, the GPU must store several additional tensors beyond the model weights. In addition to the forward activations, the system must keep gradients and optimiser states for every trainable parameter. For transformer models, the memory required for activations grows with the sequence length and number of layers, and multimodal inputs increase that sequence length further because every image is converted into visual tokens before being processed by the language model. If a full MRI or CT series contains many slices, the resulting sequence can become very long, and the memory required to store activations during back-propagation grows quickly. This is why training with large numbers of slices often leads to out-of-memory errors even on GPUs with large VRAM.

Inference, on the other hand, only requires the forward pass. The system does not need to store gradients or optimiser states, which significantly reduces memory usage. The primary additional memory structure used during inference is the key–value (KV) cache, which stores attention states so the model can efficiently process long sequences during generation. Because the KV cache grows linearly with the sequence length and does not require gradient storage, inference can handle much longer contexts than training in many cases.

The Model Garden deployment you mentioned uses an optimised serving stack based on vLLM, which is designed specifically for efficient large-context inference. vLLM implements advanced memory management techniques for the KV cache so that long sequences can be served efficiently across one or more GPUs. This allows deployments to support relatively large context windows and multiple images within a single request while keeping GPU memory usage manageable. As a result, configurations such as two L4 GPUs can support inference workloads with larger numbers of image slices than what would typically be practical during fine-tuning.

Regarding how the model handles 3D medical data, the model itself still operates on individual 2D images rather than directly consuming a volumetric tensor. Each slice of the MRI or CT study is passed through the vision encoder and converted into a set of visual tokens. These visual tokens are then inserted into the model’s input sequence alongside the text prompt. By presenting the slices in a consistent order within the prompt, the model can reason across the sequence of images and effectively interpret them as a structured representation of the volume. In practice, this means the system is performing reasoning over an ordered set of slices rather than performing native 3D convolution or volumetric processing.

For running inference with the deployed model, the recommended approach is to use the example provided in the official MedGemma resources for Model Garden deployments. The notebook demonstrates how to invoke the deployed endpoint using the Vertex AI SDK and send requests containing multimodal inputs.

Reference notebook: https://github.com/google-health/medgemma/blob/main/notebooks/quick_start_with_model_garden.ipynb

This notebook shows the typical workflow for interacting with the deployed model and should help illustrate how inference requests are structured when working with the Model Garden endpoint.

Overall, the key difference you are seeing is that the Model Garden deployment is optimised for serving long multimodal contexts during inference, whereas fine-tuning requires significantly more memory because of the additional tensors needed for back-propagation.