Skip to content
All Projects
Medical LLM architecture and training pipeline
EPFL··Academic·coursework

LLMs from Scratch & Medical Fine-Tuning

Implemented a GPT-2 architecture from scratch in PyTorch, then fine-tuned Mistral models on medical conversational and MCQ datasets using LoRA for clinical question answering.

mlllmnlp
  • ▸Implemented full GPT-2 decoder-only transformer from scratch including multi-head attention with causal masking, learned positional embeddings, and layer normalization
  • ▸Fine-tuned Mistral-7B and TinyMistral-248M on 105K medical conversations and 182K multiple-choice examples (MedMCQA) using LoRA with 8-bit quantization
  • ▸Built containerized training pipeline deployed on EPFL's Run:AI GPU cluster with multi-GPU support via DeepSpeed
  • ▸Developed instruction-following fine-tuning with loss masking on prompt tokens for targeted response learning
Stack
PythonPyTorchHugging FaceDocker
RoleTeam member
Team2 people
Source Code

Overview

Grade: 5.75 / 6 | EPFL semester project (12 ECTS), supervised by Yousra El-Bachir and Oleg Bakhteev under Prof. Mathieu Salzmann

This was an EPFL semester project completed with Ali Ridha Mrad and split into two phases: first implement a GPT-2 style language model entirely from scratch to understand transformer internals, then apply efficient fine-tuning techniques to adapt larger models for the medical domain.

The code is split across two repos: the base LLM implementation and the medical fine-tuning pipeline.

What I Built

1. GPT-2 Architecture from Scratch

GPT-2 decoder-only transformer architecture
The transformer architecture I implemented from scratch, showing multi-head attention, feed-forward layers with GELU, and layer normalization.

Following Sebastian Raschka's "Build a Large Language Model from Scratch," I implemented the full decoder-only transformer in PyTorch:

  • Multi-head self-attention with causal masking
  • Learnable token and positional embeddings
  • Feed-forward networks with GELU activation
  • Residual connections and layer normalization
  • Support for GPT-2 Small (124M params) and Medium (355M params)

Tokenization uses OpenAI's tiktoken BPE tokenizer. The model supports loading official GPT-2 pretrained weights and autoregressive text generation.

I then fine-tuned this for two downstream tasks: spam classification (with a linear classification head) and instruction following (using structured prompt formatting with loss computed only on response tokens).

2. Medical Fine-tuning with LoRA

For the medical domain, I fine-tuned Mistral-7B and TinyMistral-248M using parameter-efficient LoRA adapters rather than updating all weights:

  • LoRA config: rank 8, alpha 16, dropout 0.05, targeting Q/K/V/O projection layers
  • Quantization: 8-bit loading via BitsAndBytes for memory efficiency
  • Data: 105,488 medical conversations combined with 182,822 multiple-choice examples (MedMCQA)
  • Prompt format: <s>[INST] instruction [/INST] response </s> with loss computed only on the response tokens (instruction tokens masked out)
  • Infrastructure: Dockerized on NVIDIA PyTorch base image (CUDA 12.6), deployed to EPFL's Run:AI cluster with multi-GPU support via DeepSpeed

Technical Details

Base model (from-scratch implementation):

  • Architecture: Decoder-only transformer (GPT-2)
  • Parameters: 124M (Small) / 355M (Medium)
  • Attention heads: 12 (Small) / 16 (Medium)
  • Hidden dimension: 768 (Small) / 1024 (Medium)
  • Layers: 12 (Small) / 24 (Medium)

Medical fine-tuning:

  • Base models: Mistral-7B, TinyMistral-248M
  • Method: LoRA with 8-bit quantization
  • Training data: 105,488 medical conversations (instruction-response pairs) + 182,822 multiple-choice examples (MedMCQA)
  • Compute: EPFL's Run:AI cluster (NVIDIA GPUs), with gradient accumulation and configurable hyperparameters

Challenges & Tradeoffs

  • Full fine-tuning vs. LoRA: With limited cluster compute, LoRA made it feasible to fine-tune 7B-parameter models by only training a small number of adapter parameters
  • Quantization trade-offs: 8-bit quantization reduced memory requirements significantly but required careful handling (Flash Attention 2 caused instability on A100s, so standard attention was used)
  • Loss masking: Computing loss only on response tokens (not instruction/prompt tokens) was important for teaching the model to generate useful responses rather than memorizing prompts

Results

The fine-tuned models produce coherent medical responses to patient queries, handling topics like symptom assessment, medication information, and general health guidance. The focus of this project was the from-scratch implementation and the training infrastructure, so I validated the fine-tuned models qualitatively rather than against a held-out benchmark. Below are two examples of the model in action:

A patient asks about diabetes, the model explains Type 1 vs Type 2 and provides dietary advice

A patient describes back pain, the model asks follow-up questions and suggests treatment steps

What I Learned

  • Implementing attention from scratch gave me a mechanical understanding of transformers that using a library does not
  • LoRA is effective: adapting a 7B model's behavior only requires training less than 1% of its parameters
  • Infrastructure matters: containerizing the training pipeline and configuring cluster jobs (Run:AI, NAS storage, GPU scheduling) was a significant part of the real work
  • Loss masking on prompt tokens is a small implementation detail that makes a big difference in fine-tuning quality