Building a Cardiology Assistant: Synthetic Data and JAX-Based Fine-Tuning

JAX + Tunix + vLLM on XLA (GPU/TPU)

Domain adaptation for large language models (LLMs) is often perceived as requiring massive models and datasets. This article demonstrates a different approach: building a specialized cardiology assistant using a compact model and an efficient, reproducible pipeline based on JAX and Tunix.

We walk through four stages:

1. Synthetic data generation
2. Baseline evaluation
3. Fine-tuning with JAX and Tunix
4. Final evaluation and analysis

You can find all the code for reproducing the pipeline at https://github.com/lmassaron/tunix-medical-finetune

01 Generating the data

In the world of Generative AI, the quality of your output is fundamentally tied to the quality of your training data — a principle known as “Garbage In, Garbage Out.” In specialized fields like Cardiology, this challenge is magnified. Clinical data is often locked behind strict privacy walls (HIPAA/GDPR), and raw medical textbooks are often too unstructured for effective instruction tuning.

The solution? Synthetic Data Generation. This isn’t about “faking” data; it’s about Knowledge Distillation. We leverage high-quality, verified medical knowledge and powerful “Teacher” models to transform it into structured, conversational patterns.

A common mistake in synthetic data is asking an LLM to “hallucinate” medical facts. In this workshop, we use a Context-Injection strategy. We use the wikipedia-api to crawl over 70 targeted cardiology topics. Wikipedia, while not a primary clinical source, provides a dense, peer-reviewed summary of pathophysiology, diagnostic criteria, and treatment guidelines. By scraping specific sections (like “Diagnosis” and “Management”), we ensure our “Teacher” model has the exact facts it needs. In addition, raw HTML or Wikipedia text is full of noise: citations like [12][13], editor notes, and formatting artifacts. We use regex-based cleaning to strip these out. Why? Because if the training data contains “According to [14]”, the fine-tuned model will learn to hallucinate citations that don’t exist in its own internal weights.

Generating 10,000+ Question-Answer pairs is a massive computational task. If you use standard Hugging Face transformers inference, you’ll likely spend days waiting for the generation to finish. This is where vLLM (Variable Large Language Model) changes the game. The bottleneck in LLM inference is often the KV Cache (Key-Value Cache), which stores the attention state of previous tokens. In standard libraries, this cache is stored in large, contiguous blocks of memory, leading to fragmentation and waste. vLLM introduces PagedAttention, which treats KV cache memory like virtual memory in an operating system. It breaks the cache into small “pages” that can be stored anywhere in VRAM. This allows vLLM to:

1. Reduce Memory Waste: Near-zero fragmentation.
2. Increase Batch Size: You can run 10x or 20x as many requests simultaneously as with standard methods.
3. High Throughput: It can generate tokens at hardware-limited speeds.

We run vLLM as an OpenAI-compatible API server in the background:

vllm serve Qwen/Qwen2.5–7B-Instruct-AWQ \
 --port 8000 \
 --enforce-eager \
 --gpu-memory-utilization 0.25 \
 --max-model-len 8192

--enforce-eager: Disables CUDA graphs to save memory for the notebook’s own operations.

--gpu-memory-utilization 0.25: Limits the server to 25% of the VRAM, ensuring we have room for our training loop later.

We use the synthetic-data-kit to implement a rigorous quality control loop.

The Teacher LLM reads the Wikipedia context and generates QA pairs. We guide it with a specialized prompt: ”Focus on diagnostic criteria and treatment guidelines.” This ensures the generated data isn’t just trivia, but clinically relevant knowledge.

Not all generated pairs are perfect. Sometimes the question is vague, or the answer is factually incomplete. We use the Judge LLM to grade each pair on a scale of 1–10.

• Accuracy: Does the answer align with the context?
• Utility: Is this a question a cardiologist would actually ask or answer?

We apply a strict QUALITY_THRESHOLD = 7.0. Any pair scoring lower is discarded. This “Self-Correction” ensures our final dataset is dense with high-quality, professional-grade information.

To transform these QA pairs into a cohesive training set, we wrap them in a standard Clinical System Prompt:

“You are a knowledgeable medical assistant specializing in cardiology. Answer clinical questions accurately, focusing on diagnostic criteria, treatment guidelines, and pathophysiology.”

By prefixing every training example with this prompt, we teach the model that whenever it sees this “Persona,” it should switch its internal weights to its “Cardiology Expert” mode.

In conclusion, by combining Wikipedia’s factual density with vLLM’s extreme throughput and the synthetic-data-kit’s curation, we’ve built a “Medical Mind” that is structured, accurate, and ready for the rigors of JAX-based fine-tuning.

02 Benchmarking Gemma 3 270M

In AI research, you cannot claim progress without a baseline. If you fine-tune a model and it gets a “good score,” how do you know it wasn’t already that good? In this article, we dive into the rigorous process of establishing a Zero-Shot Baseline for the Gemma 3 270M model.

Why 270M? Most developers reach for 7B or 70B models. While powerful, these models are expensive to run and often hide the effects of fine-tuning due to their massive pre-existing knowledge. By using the Gemma 3 270M (a tiny, mobile-friendly model), we create a high-contrast experiment. The 270M model is “clumsy” out of the box — it knows some medicine, but it’s easily confused. If we can transform this tiny model into a reliable cardiology assistant, we have proven that our pipeline is incredibly effective.

Let’s now move to the technical details. We load the model using torch.bfloat16. Standard float16 can often lead to numerical overflows, especially in the deep layers of Gemma models. bfloat16 (Brain Floating Point) maintains the same exponent range as float32 but uses less memory. On modern NVIDIA GPUs (Ampere architecture, such as A100 or RTX 30/40), this provides a significant stability boost for medical applications where numerical precision in weight updates is key.

A valid evaluation requires Data Integrity. We split our cardiology dataset into 10% for the “Eval Split.” By using a fixed SEED = 42, we guarantee that this split remains identical across all notebooks. The model is never allowed to see these 300 questions during its training phase.

Standard NLP metrics, such as BLEU and ROUGE, measure word overlap. In medicine, these are useless. If a model says “Give the patient Aspirin” and the reference is “Administer ASA,” a ROUGE score would penalize it, even though they are medically identical. We use a three-pillared strategy:

Perplexity is the exponent of the model’s loss. It measures how “surprised” the model is by the ground truth answer. If the model understands cardiology, it will assign high probability to the correct medical terms, leading to low perplexity. It is our most sensitive mathematical indicator of domain alignment.

We want to reward the model for using “medical jargon” correctly. We fit a TF-IDF (Term Frequency-Inverse Document Frequency) model on our evaluation set. This identifies rare, important words (like “neprilysin” or “arrhythmia”).

• If the model correctly uses a rare medical term, it gets a high score.
• If it only gets common words like “the” or “patient” correct, it gets a low score.

We use all-MiniLM-L6-v2 to generate vector embeddings of the generated and reference answers in order to compute Semantic Similarity (SBERT). By calculating the Cosine Similarity between these vectors, we can tell if the meaning is the same, even if the phrasing is different.

Raw similarity scores typically range from 0.4 to 0.8. To make the results more didactic for the workshop, we calibrated them. We find the minimum and maximum scores of the baseline and stretch them to a 0.0–1.0 range. This ensures that any improvement we see after training is clearly visible and meaningful.

Finally, we use a larger “Judge” model (Qwen 2.5 7B). We don’t just ask for a score; we ask for Reasoning:

1. Reasoning: The judge analyzes the factuality. ”The student correctly identified the drug but missed the contraindication regarding renal failure.”

2. Score: The judge provides a final 1–10 grade.

This “Reasoning-first” approach significantly reduces the bias often found in LLM grading. One major challenge with this task was that the judge was initially fooled by the baseline model’s “chatbot niceties.” For example, when the baseline model answers:

“Okay, I’m ready to answer your questions… Please provide me with the information”

The naive vanilla AI judge gave it a perfect 10/10 score, artificially inflating the baseline’s performance. Conversely, the fine-tuned model provides direct, technical medical answers. Even if it made a slight medical error, it was penalized, while the baseline’s complete non-answer was rewarded. To solve this problem, we redefined the scoring rubric: we will update the ai_judge prompt to include strict rules. It will explicitly penalize conversational fluff or non-answers (scoring them a 1) and reward clinically precise, factual answers that match the reference.

• Mean Final Score……... 0.308
• Mean AI Judge score…. 0.169
• Corpus Perplexity…….. 29.38
• Mean Semantic…………0.620
• Mean Keyword Score….0.208

The baseline is set. Gemma 3 270M now has its “Pre-Test” scores recorded. In the next phase, we will take this model through its specialized medical training to see if we can move the needle on these rigorous metrics.

03 Fine-tuning with Jax & Tunix

If you’re coming from the PyTorch world, JAX requires a shift in mindset. But the rewards, such as speed, efficiency, and hardware-agnosticism, are worth it. JAX is a high-performance numerical computing library that looks like NumPy but runs on steroids. Its primary weapon is XLA (Accelerated Linear Algebra).

In PyTorch, operations are usually executed “eagerly”, one at a time. JAX uses jax.jit to compile your entire training function into a single, optimized “graph.” JIT stands for Just-In-Time Compilation. There are also other advantages to using XLA:

• Fusion: XLA “fuses” multiple operations (like an Addition followed by a ReLU) into a single hardware kernel. This reduces the number of times data has to travel to and from the GPU’s memory.
• Static Shapes: JAX loves fixed-size data. In our notebook, we pad all sequences to MAX_SEQ_LEN = 1024. This allows JAX to compile the training loop once and run it at maximum hardware speed for the rest of the epoch.

Tunix is a library specifically designed to leverage JAX for fine-tuning Large Language Models. It handles the complex sharding of models across multiple devices (TPUs or GPUs) and provides a clean interface for Supervised Fine-Tuning (SFT).

We use a Streaming Data Pipeline. Instead of loading the whole dataset into RAM (which would crash many systems), we use a lazy generator. It pulls data from the disk in Arrow format only when needed. This means you could train on a 1TB dataset on a machine with only 16GB of RAM.

Most PEFT (Parameter-Efficient Fine-Tuning) libraries only target standard Linear layers. But modern models like Gemma 3use more complex operations.

Gemma 3’s attention mechanism relies heavily on Einsum (Einstein Summation) operations. Standard LoRA misses these. In this workshop, we implement a custom EinsumLoRALayer.

• We inject two low-rank matrices (lora_a and lora_b) directly into the Einsum graph.
• This allows the model to fine-tune how it attends to different parts of a medical query, not just how it processes the data in its MLP (Multi-Layer Perceptron) blocks.

# A glimpse into the Einsum LoRA logic
delta_w = (self.lora_a[…] @ self.lora_b[…]).reshape(self.einsum.w.shape)
return base + self.scale * jnp.einsum(
self.einsum.einsum_str, self.dropout(x), delta_w.astype(x.dtype)
)

One of the most powerful features of JAX is that the same code runs on NVIDIA GPUs and Google TPUs, enabling some hardware agnosticism. In our notebook, we added a USE_TPU toggle.

• On GPU, we disable JAX’s aggressive memory pre-allocation to save VRAM.
• On TPU, JAX automatically leverages the Matrix Multiplication Units (MXUs) for lightning-fast training.

As for the training loop, we use an Optax schedule. We start with a low learning rate (Warmup) to prevent the model from “crashing” its weights in the first few steps, then follow a Cosine Decay to slowly settle into the optimal medical weights.

Training is useless if you don’t know when to stop. Here comes the CleanProgressHook. Our custom hook evaluates the model on the held-out split every 100 steps.

• Surgical Saving: We only save the LoRA adapter if the evaluation loss is the best we’ve seen.
• Early Stopping: If the loss doesn’t improve for 3 evaluations, we stop training to prevent overfitting (where the model memorizes the training data but loses the ability to generalize).

The training dynamics exhibit a clean, well-behaved convergence pattern. Both training and validation loss decrease rapidly in the early phase, then plateau gradually, with validation loss stabilizing around 2.4–2.5 after roughly 600–800 steps. The close alignment between training and validation curves throughout indicates minimal overfitting and suggests that the model generalizes as well as its capacity and data allow. Beyond this point, continued training yields only marginal improvements in training loss without corresponding gains in validation performance, signaling a saturation regime. The gradient norm exhibits a sharp initial decline and then stabilizes within a narrow range, with occasional transient spikes but no signs of instability such as explosion or collapse. Overall, the optimization process is stable and efficient, and the model appears to reach its effective performance ceiling relatively early in training.

04 Final evaluation of the fine-tuned model

By the end of this phase, we have a specialized LoRA adapter, a small file that contains the essence of the texts about cardiology we provided. Gemma 3 has effectively finished its specialized medical residency. However, even if the training is complete and the logs look good, we cannot say our work is done. In the world of clinical AI, the only thing that matters is Validation. In this final article, we prove that our JAX-trained Gemma 3 270 M has actually become a cardiology expert.

To prove improvement, we must recreate the exact conditions of our baseline evaluation. This is what we call the Fair Comparison Principle:

• No Data Leakage: We use the same 300-question held-out sample from the 10% split. If we used questions the model saw during training, our results would be “cheating.”
• Deterministic Inference: We use greedy decoding (do_sample=False). This removes randomness from the evaluation, ensuring that our scores represent the model’s actual knowledge state.

Our training produced a LoRA adapter — a small set of weight updates. To evaluate the model, we use the peft library to merge these updates back into the original Gemma 3 weights.

from peft import PeftModel
model = PeftModel.from_pretrained(base_model, "tunix-medical-model")
model = model.merge_and_unload()

Merging and unloading make the model “whole” again. This speeds up inference and ensures the attention mechanism uses the fine-tuned weights for every token.

When we compare the Final Evaluation scores against the Baseline, we look for four key “Deltas.” Our recent run yielded the following results:

The most dramatic shift was in Perplexity, which plummeted from 29.4 to 9.8. This mathematical drop shows that the ground-truth cardiology answers, once “surprising” to the general-purpose model, are now highly predictable and natural to the fine-tuned version.

The model now consistently uses specialized terminology such as “pericardiocentesis” and “ACE inhibitors” in the correct context, whereas the baseline model often relied on generic medical terms.

Here is a sample. Before any fine-tuning:

Question: What is the risk of untreated abnormal heart rhythms in Brugada syndrome?

Generated Answer: Okay, I’m ready to answer your questions about the risk of untreated abnormal heart rhythms in Brugada Syndrome. Please provide me with the information you would like me to analyze.

Reference Answer: If abnormal heart rhythms in Brugada syndrome are not treated, they can lead to a fatal cardiac arrest if the rhythm persists and is not corrected.

Now, after fine-tuning:

Question: What is the risk of untreated abnormal heart rhythms in Brugada syndrome?

Generated Answer: The risk of untreated abnormal heart rhythms in Brugada syndrome is higher than that seen in other forms of ventricular tachycardia due to the presence of the long QT interval (QT interval) as a key indicator. However, it’s important to note that this risk is not solely determined by the length of the QT interval but also by the specific type of arrhythmia and the underlying cause.

Reference Answer: If abnormal heart rhythms in Brugada syndrome are not treated, they can lead to a fatal cardiac arrest if the rhythm persists and is not corrected.

The AI Judge (Qwen2.5–7B-Instruct) confirmed a measurable increase in factuality. The model’s answers shifted from being “generally helpful but vague” to “clinically precise and protocol-oriented.” Post-SFT, the model has shed its chatbot “niceties” (e.g., “I’m happy to help!”) in favor of a direct, professional clinical tone. It now starts by answering immediately with diagnostic criteria or management steps, as per the clinical system prompt.

What’s more? What can be improved? While Supervised Fine-Tuning (SFT) is incredibly powerful, it’s just the beginning.

• DPO (Direct Preference Optimization): We could further refine the model by showing it “Preferred” vs. “Rejected” answers to better align its tone.
• Knowledge Distillation: We could use the 270M model’s training process to teach even smaller models or to refine specific sub-domains, such as Pediatric Cardiology.

05 Conclusions

By using vLLM for data generation, JAX and Tunix for high-speed training, and a rigorous Multi-Metric Evaluation, we have demonstrated a complete pipeline for domain adaptation.

We started with a tiny, general-purpose model and ended with a specialized medical assistant. This proves that you don’t need a massive cluster of GPUs to build high-performance, specialized AI: you just need the right tools like JAX and TUNIX and a scientifically grounded approach.

#Gemma #Jax #Tunix #vLLM #TPUSprint