Fine-Tuning Llama 3.2 11B Vision on an Astronomy Dataset with Unsloth

Fine-tuning Llama 3.2 11B Vision on an astronomy dataset using Unsloth for high-speed cosmic image analysis

From raw image-caption pairs to a deployed multimodal model in under 15 minutes of training

There is a certain category of problem in applied machine learning that initially sounds unreasonable: take an 11-billion-parameter multimodal language model, adapt it to a highly specific scientific domain, do it on a single-GPU free-tier environment, and make the whole thing production-ready by the end of your session. That is exactly what this project does.

The target domain is astronomy. The task is image captioning, or more precisely, visual question answering with a scientific framing. The model should look at an image of the Hubble Space Telescope, a Martian crater, or the Milky Way, and produce an accurate, domain-aware description. The end result is sitting on Hugging Face at AIOmarRehan/Llama-3.2-11B-Vision-LoRA-on-Astronomy-with-Unsloth, served through a Gradio interface called AstroVision Chat.

This post walks through the full pipeline, including dataset analysis, data cleaning, model setup, training, evaluation, and deployment. All of it runs on a Tesla T4. All of it uses Unsloth.

The Dataset

The starting point is a custom astronomy dataset stored as a ZIP file on Google Drive. It contains 250 image-caption pairs, structured as a JSON file alongside an images/ directory. The JSON schema is straightforward: each record has an image_id, a text caption, and a relative image path.

After loading and running an initial check, the dataset is clean out of the box:

Missing values in each column:
image_id    0
text        0
image       0
dtype: int64

250 entries, 250 unique captions, no missing values, no empty strings. This is a well-curated dataset, which makes the EDA phase useful for understanding the data distribution rather than firefighting obvious quality issues.

Exploratory Data Analysis

Caption Statistics

df['text_length'] = df['text'].apply(len)
df['word_count'] = df['text'].apply(lambda x: len(x.split()))

print("Average caption length (characters):", df['text_length'].mean())
print("Average caption length (words):", df['word_count'].mean())

The average caption is 93.9 characters long and 15.2 words. That is a reasonably dense description, enough context for a model to learn domain-specific vocabulary.

This histogram shows the distribution of word counts across all 250 captions. The distribution is roughly bell-shaped and centered around 15 words, with very few captions below 10 or above 25 words. The absence of extreme outliers confirms the dataset is consistent.

This bar chart displays the distribution of caption word counts in the dataset, showing that most captions contain between 14 and 16 words

Word Frequency

def clean_text(text):
    text = text.lower()
    text = re.sub(r'[^a-z\s]', '', text)
    return text

all_words = " ".join(df['text'].apply(clean_text)).split()
word_freq = Counter(all_words).most_common(20)

The top words after cleaning are what you would expect: “the”, “a”, “of”, “showing”. Digging past the stopwords, domain-specific terms appear quickly: “mars” (96 occurrences), “earth” (61), “space” (53), “milky” (50), “hubble” (43), “telescope” (39).

This bar chart shows the 20 most common words across all cleaned captions. After generic connector words, the next tier is purely astronomical, confirming the domain specificity of the dataset. The balance between “mars” and “earth” as the most common proper nouns is notable.

This bar chart ranks the 20 most frequent words found in the dataset captions, highlighting a focus on astronomical terms like “Mars,” “Earth,” and “Hubble”

Image Resolution

Astronomy datasets often contain a mix of archival images, telescope photography, and rover photographs. These come from heterogeneous sources with wildly different resolutions.

for img_name in df['image']:
    img_path = os.path.join('/content/my_data/astronomy_dataset', img_name)
    with Image.open(img_path) as img:
        w, h = img.size
        widths.append(w)
        heights.append(h)

The most common resolution is 300x168 pixels (15 images), followed by 225x225 (13 images) and 720x1280 (12 images). The most striking observation is the presence of both portrait and landscape orientations, and resolutions ranging from small thumbnails to 1920x1080 and 2000x2000 full-resolution images.

This scatter plot shows each image plotted by its width on the x-axis and height on the y-axis. Most images cluster in the lower-left, confirming that the majority are small to medium resolution. A handful of outliers appear in the top-right, representing the high-resolution images (1920x1080, 2000x2000).

Scatter plot showing dataset image resolutions

These two histograms show the marginal distributions of width and height separately. Both are right-skewed, with most images below 750px in each dimension. The tails extend toward high resolutions for both axes.

Histograms showing image width and height

Sample Images

6-panel random sample grid

Randomly sampling from the full dataset gives a representative feel: some images are high-contrast telescope photographs, others are NASA press-release style shots of planetary surfaces, and some are orbital perspectives of Earth.

A more structured view groups one image per semantic category:

5-panel one-per-category grid (Earth, Mars, Hubble, Milky Way, Mars Rover)

Category Distribution

An automatic labeling step assigns each caption to one of five categories based on keyword matching in the text:

def detect_label(text):
    text_lower = text.lower()
    if "earth" in text_lower:
        return "Earth"
    elif "mars" in text_lower and "rover" not in text_lower:
        return "Mars"
    elif "hubble" in text_lower:
        return "Hubble"
    elif "milky" in text_lower:
        return "Milky Way"
    elif "rover" in text_lower or "laboratory" in text_lower:
        return "Mars Rover"
    else:
        return "Unknown"

The resulting distribution:

Earth         77
Mars Rover    46
Milky Way     45
Mars          42
Hubble        28
Unknown       12

The bar chart makes the imbalance visually obvious: Earth accounts for nearly a third of the dataset, while Hubble and Unknown together make up only 16%.

Label / Class Distribution bar chart

Earth dominates with 77 images, roughly 31% of the dataset. Hubble is underrepresented at 28 images. The 12 “Unknown” images are cases where none of the five keywords appear in the caption, often because the caption describes a feature (like a continent or a crater field) without naming the subject directly.

Data Cleaning

The dataset is clean in the statistical sense, but the pipeline still applies normalization steps that matter for fine-tuning quality.

Text Normalization

def clean_caption(text):
    text = text.lower()
    text = re.sub(r'[^a-z\s]', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text

This converts all text to lowercase and strips punctuation and digits. For a captioning task, this is a reasonable tradeoff. The model will generate descriptions that rely on vocabulary rather than punctuation precision.

Before and after:

Before: 'A natural color image of Mars displaying the Tharsis volcanic plateau...'
After:  'a natural color image of mars displaying the tharsis volcanic plateau...'

The normalization also caught a minor artifact: the hyphenated term “wide-angle” becomes “wideangle” after punctuation removal. This is acceptable since the word still carries its semantic meaning.

Duplicate Removal and Image Validation

After normalization, zero duplicate captions remain. All 250 images pass validation:

def is_valid_image(image_path):
    try:
        with Image.open(image_path) as img:
            img.verify()
        return True
    except Exception as e:
        return False

Invalid or missing images found: 0
Remaining valid images: 250

There is also a commented-out resolution filter in the notebook that would have removed images below 224x224 pixels. It was intentionally left disabled. Given the small dataset size, removing low-resolution images would reduce training data with no guaranteed improvement in output quality.

Model Setup

Loading Llama 3.2 11B Vision with Unsloth

from unsloth import FastVisionModel
import torch

model, tokenizer = FastVisionModel.from_pretrained(
    "unsloth/Llama-3.2-11B-Vision-Instruct",
    load_in_4bit=True,
    use_gradient_checkpointing="unsloth",
)

Unsloth patches the model at load time for faster forward and backward passes. The 4-bit quantization (NF4) brings the 11B parameter model well within the 14.7 GB VRAM of a Tesla T4. Without this, you would need an A100 or a multi-GPU setup just to load the model weights.

The startup log confirms the environment:

Tesla T4. Num GPUs = 1. Max memory: 14.741 GB. Platform: Linux.
Torch: 2.8.0+cu126. CUDA: 7.5. Transformers: 4.56.2.

Two safetensor shards are downloaded: model-00001-of-00002.safetensors (4.97 GB) and model-00002-of-00002.safetensors (2.94 GB).

LoRA Configuration

model = FastVisionModel.get_peft_model(
    model,
    finetune_vision_layers     = True,
    finetune_language_layers   = True,
    finetune_attention_modules = True,
    finetune_mlp_modules       = True,
    r = 16,
    lora_alpha = 16,
    lora_dropout = 0,
    bias = "none",
    random_state = 3407,
    use_rslora = False,
)

All four parameter groups are included: vision layers, language layers, attention modules, and MLP modules. The rank is 16, which is a standard starting point that provides enough capacity without excessive memory overhead. With lora_alpha = r, the effective scaling factor is 1.0.

The result: 67,174,400 trainable parameters out of 10,737,395,235 total. That is 0.63% of the model being updated during training. The remaining 99.37% stays frozen.

Dataset Formatting

The SFT training pipeline requires the dataset to be structured as multi-turn conversations. Each sample becomes a user/assistant exchange:

def convert_to_conversation(sample):
    conversation = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "You are an expert astronomer. Describe accurately what you see in this image."},
                {"type": "image", "image": os.path.join(base_path, sample["image"])},
            ],
        },
        {
            "role": "assistant",
            "content": [{"type": "text", "text": sample["text"]}],
        },
    ]
    return {"messages": conversation}

The system instruction is embedded in the user turn: “You are an expert astronomer. Describe accurately what you see in this image.” This framing anchors the model’s output style toward scientific accuracy rather than general image description.

The dataset is split 80/10/10:

Train: 200
Valid: 25
Test:  25

Baseline: What the Model Says Before Fine-Tuning

Before starting the training loop, it is worth running a single inference pass on the base model. This establishes a baseline and often reveals interesting things about what the model already knows.

Prompt: a satellite image of Earth (images/001.jpg). Instruction: "You are an expert astronomer. Describe accurately what you see in this image."

The base model responds:

“From the perspective of the moon’s surface, the Earth appears as a majestic blue and white globe, illuminated by the radiant sun, set against the inky blackness of space, speckled with stars. The moon’s surface, dominated by dark grey terrain, is characterized by numerous craters and rugged ridges… awe-inspiring beauty and a reminder of our world’s place in the universe.”

The hallucination here is instructive. The model correctly identifies Earth as a blue and white globe, but it invents a perspective (from the moon’s surface) that is not supported by the input image. It produces fluent, even eloquent prose, but the spatial and observational details are fabricated. This is the classic base model problem: high fluency, unreliable grounding.

This is exactly the gap that fine-tuning is designed to close.

Training

from unsloth.trainer import UnslothVisionDataCollator
from trl import SFTTrainer, SFTConfig

FastVisionModel.for_training(model)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    data_collator=UnslothVisionDataCollator(model, tokenizer),
    train_dataset=hf_dataset,
    args=SFTConfig(
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        max_steps=30,
        learning_rate=2e-4,
        logging_steps=1,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir="outputs",
        remove_unused_columns=False,
        dataset_text_field="",
        dataset_kwargs={"skip_prepare_dataset": True},
        max_length=2048,
    ),
)

The effective batch size is 8 (2 per device times 4 gradient accumulation steps). With 250 samples at batch size 8, you get roughly 31 gradient updates per epoch, which matches the max_steps = 30 setting.

adamw_8bit is the optimizer, with Unsloth's 8-bit quantization of optimizer states. This meaningfully reduces the memory footprint of the optimizer itself, not just the model weights.

GPU Memory

Before training:

GPU = Tesla T4. Max memory = 14.741 GB.
8.582 GB of memory reserved.

After training:

Peak reserved memory = 10.049 GB.
Peak reserved memory for training = 1.467 GB.
Peak reserved memory % of max memory = 68.17 %.

The training itself consumed only 1.467 GB on top of the base reservation. Loading an 11B model, running LoRA fine-tuning, and staying under 70% of a 14 GB GPU is a genuinely remarkable outcome.

Total training time:

686.0076 seconds used for training.
11.43 minutes used for training.

Under twelve minutes on a free-tier T4.

Evaluation

After training, the model is evaluated on the 25-sample test set using BLEU and ROUGE metrics.

from evaluate import load

bleu = load("bleu")
rouge = load("rouge")

bleu.add_batch(predictions=predictions, references=[[r] for r in references])
bleu_score = bleu.compute()

rouge.add_batch(predictions=predictions, references=references)
rouge_score = rouge.compute()

Results:

BLEU:   0.0537
ROUGE-1: 0.2658
ROUGE-2: 0.0979
ROUGE-L: 0.2383

A BLEU score of 0.054 is low in absolute terms. For context, state-of-the-art image captioning models on COCO achieve BLEU-4 scores in the 35–40 range. However, those models are trained on 100,000+ caption pairs, with multiple reference captions per image for evaluation.

With 200 training samples and a single reference per test image, a BLEU-4 of 0.054 is a more meaningful number than it might first appear.

The ROUGE-1 score of 0.266 tells a more useful story. It measures unigram overlap between predicted and reference captions, regardless of exact phrase matching. This suggests the fine-tuned model is producing domain-relevant vocabulary at a meaningful rate.

The per-image comparisons are where the quality becomes tangible:

Image 1
Predicted: "a photograph of the hubble space telescope in orbit above earth
            illuminated by sunlight"
Reference:  "an image showing hubbles metallic body glinting in sunlight
            with the darkness of space behind it"

Image 2
Predicted: "a photograph of earth taken from space showing the continent of
            africa surrounded by ocean"
Reference:  "a distant perspective of earth capturing the blue oceans and
            swirling atmosphere against the blackness of space"

Image 3
Predicted: "an orbital image of the mars spacecraft orbiting above the
            martian surface with the martian terrain visible beneath"
Reference:  "a detailed orbital view of mars highlighting dark dune fields
            and dustfilled craters"

The model is no longer hallucinating. It correctly identifies Hubble in Image 1, correctly identifies Africa in Image 2, and correctly identifies Mars and orbital context in Image 3. The specific details differ from the reference captions, but the subject matter, perspective, and domain language are accurate. That is a qualitative improvement that the BLEU score does not fully capture.

Visual Inference Results

The final inference test uses a Hubble image from the test set and compares the ground truth caption to the model’s output side by side.

plt.imshow(image)
plt.axis("off")
plt.title("Test Image Example", fontsize=14)
plt.show()

print("Ground Truth Caption:")
print(true_caption)

print("\nModel Prediction:")
print(generated_text)

Image used to test the model of the “Ground Truth Caption” and “Model Prediction”

Ground Truth Caption:
an image showing hubbles metallic body glinting in sunlight with the darkness of space behind it

Model Prediction:
user

You are an expert astronomer. Describe accurately what you see in this image.assistant

a close up view of the hubble space telescope in orbit above the earth

Ground Truth: “an image showing hubbles metallic body glinting in sunlight with the darkness of space behind it”

Model Prediction: “a close up view of the hubble space telescope in orbit above the earth”

The model’s prediction is less poetic than the reference. It does not mention the glinting sunlight or the metallic body. But it is factually correct: it correctly names the Hubble Space Telescope, correctly identifies the setting as Earth orbit, and the perspective as a close-up view. Before fine-tuning, the model was inventing vantage points that did not exist in the image. Now it is grounded.

Saving and Pushing to Hugging Face

After training, the LoRA adapter is merged back into the base model weights at full 16-bit precision:

model.save_pretrained_merged("unsloth_finetune", tokenizer)

Unsloth handles the dequantization and merge automatically. The resulting folder, unsloth_finetune, contains five safetensor shards totaling roughly 21 GB:

model-00001-of-00005.safetensors — 4.99 GB
model-00002-of-00005.safetensors — 4.97 GB
model-00003-of-00005.safetensors — 4.92 GB
model-00004-of-00005.safetensors — 5.00 GB
model-00005-of-00005.safetensors — 1.47 GB

The merge process takes roughly 7 minutes. The final step pushes the adapter to Hugging Face:

model.push_to_hub("AIOmarRehan/Llama-3.2-11B-Vision-LoRA-on-Astronomy-with-Unsloth", tokenizer=tokenizer)

The adapter file itself is 269 MB. Anyone with access to the base unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit model can load the fine-tuned version by merging this adapter at inference time.

Deployment

The deployment notebook separates cleanly from the training notebook. It loads the base model and merges the LoRA adapter from Hugging Face in a single step:

from unsloth import FastVisionModel

base_model_id = "unsloth/Llama-3.2-11B-Vision-Instruct-bnb-4bit"
model, tokenizer = FastVisionModel.from_pretrained(
    base_model_id,
    load_in_4bit=True,
    use_gradient_checkpointing="unsloth",
)

lora_repo_id = "AIOmarRehan/Llama-3.2-11B-Vision-LoRA-on-Astronomy-with-Unsloth"
model = FastVisionModel.get_peft_model(model, lora_adapter=lora_repo_id)
model.eval()
FastVisionModel.for_inference(model)

This is the correct pattern for production deployment. Rather than serving the full 21 GB merged model, you serve the 4-bit base model (roughly 8 GB) plus the 269 MB adapter. The total memory footprint stays manageable.

Gradio Interface

import gradio as gr

def chat_with_image(prompt, img):
    if img is None:
        return "Please upload an image!"

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": prompt}
            ]
        }
    ]

    input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
    inputs = tokenizer(img, input_text, add_special_tokens=False, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=512, temperature=1.5)

    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    if "assistant" in generated_text:
        generated_text = generated_text.split("assistant")[-1].strip()

    return generated_text

iface = gr.Interface(
    fn=chat_with_image,
    inputs=[
        gr.Textbox(label="Ask something"),
        gr.Image(type="pil", label="Upload an image")
    ],
    outputs=gr.Textbox(label="Model Response", lines=15),
    title="AstroVision Chat",
    description="Ask the model about your image. Llama 3.2 11B Vision LoRA in action!",
    allow_flagging="never"
)

iface.launch()

The interface accepts any image and a free-form text prompt. The extraction logic strips the system and user portions of the decoded output and returns only the assistant response. The temperature=1.5 setting at inference time produces varied, natural-sounding descriptions rather than deterministic repetitions.

What This Project Actually Demonstrates

There are a few things worth being explicit about here.

First, 250 samples is a small dataset. The fact that this fine-tuning run produces coherent, domain-grounded captions from 200 training examples says more about the quality of the base model and LoRA as an adaptation method than it does about the volume of training data. The base Llama 3.2 11B Vision model already understands what a telescope looks like and what orbit means. The fine-tuning teaches it to apply that knowledge through the specific lens of this dataset’s vocabulary and style.

Second, the memory efficiency here is real and not a marketing claim. 8.58 GB reserved before training, 10.05 GB peak during training, on a 14.74 GB GPU, with an 11B parameter model. That window is narrow but workable, and it means this entire pipeline runs for free on Google Colab.

Third, BLEU and ROUGE are inadequate evaluation metrics for this type of task, and the notebook is honest about that by printing the raw prediction-reference pairs alongside the aggregate scores. The per-image comparisons show a model that is correctly grounded in the visual content, even when its exact phrasing diverges from the reference. A more rigorous evaluation would involve human ratings or a VQA benchmark, but for a domain-specific proof of concept, the qualitative results are compelling.

The broader lesson is that fine-tuning a large vision-language model on a small, well-curated domain-specific dataset is viable and practical. The barrier is no longer compute. It is dataset quality and the willingness to do the EDA and cleaning work before you ever touch the model.

Conclusion

The complete pipeline covered in this project is:

1. Domain-specific dataset loading and EDA.
2. Text normalization, language filtering, duplicate removal, and image integrity validation.
3. 4-bit quantized model loading with Unsloth.
4. LoRA adapter configuration targeting all parameter groups.
5. Conversation-format dataset preparation for SFT.
6. 11.43 minutes of training on a free-tier Tesla T4.
7. BLEU/ROUGE evaluation with qualitative per-sample inspection.
8. Model merging at 16-bit precision and push to Hugging Face Hub.
9. Deployment via a Gradio interface with the base model plus adapter pattern.

The fine-tuned model is available at https://huggingface.co/AIOmarRehan/Llama-3.2-11B-Vision-LoRA-on-Astronomy-with-Unsloth. Anyone with a GPU and the base Llama 3.2 11B Vision model can reproduce the deployment notebook in a few minutes.

The code is straightforward enough that swapping in a different domain dataset, changing the system prompt, and re-running the training loop should take less than an hour. That is the real promise of this stack: rapid domain adaptation of production-grade multimodal models without requiring production-grade infrastructure.

Connect With Me If you’d like to follow my work or get in touch, you can find me here:

You can find more details about the project through these links, including the source code and additional explanations. The code is also available for testing on Hugging Face.

• LinkedIn: [My LinkedIn Profile]
• GitHub: [My GitHub Profile]
• Kaggle: [My Kaggle Profile]
• Hugging Face: [My Hugging Face Profile]
• My Website: *[My Official Website]*