diff --git a/pkg/aikit2llb/finetune/convert.go b/pkg/aikit2llb/finetune/convert.go
index 2b9564b9..a070b148 100644
--- a/pkg/aikit2llb/finetune/convert.go
+++ b/pkg/aikit2llb/finetune/convert.go
@@ -12,7 +12,7 @@ import (
 )
 
 const (
-	unslothCommitOrTag = "27fa021a7bb959a53667dd4e7cdb9598c207aa0d" // June-2024
+	unslothCommitOrTag = "fb77505f8429566f5d21d6ea5318c342e8a67991" // September-2024
 	nvidiaMknod        = "mknod --mode 666 /dev/nvidiactl c 195 255 && mknod --mode 666 /dev/nvidia-uvm c 235 0 && mknod --mode 666 /dev/nvidia-uvm-tools c 235 1 && mknod --mode 666 /dev/nvidia0 c 195 0 && nvidia-smi"
 	sourceVenv         = ". .venv/bin/activate"
 )
@@ -46,7 +46,7 @@ func Aikit2LLB(c *config.FineTuneConfig) llb.State {
 	if c.Target == utils.TargetUnsloth {
 		// installing unsloth and its dependencies
 		// uv does not support installing xformers via unsloth pyproject
-		state = state.Run(utils.Shf("pip install --upgrade pip uv && uv venv --system-site-packages && %[1]s && uv pip install packaging torch==2.3.0 ipython ninja packaging bitsandbytes setuptools==69.5.1 wheel psutil && uv pip install flash-attn --no-build-isolation && python -m pip install 'unsloth[cu121_ampere_torch230] @ git+https://github.com/unslothai/unsloth.git@%[2]s'", sourceVenv, unslothCommitOrTag)).Root()
+		state = state.Run(utils.Shf("pip install --upgrade pip uv && uv venv --system-site-packages && %[1]s && uv pip install --upgrade --force-reinstall packaging torch==2.4.0 ipython ninja packaging bitsandbytes setuptools==69.5.1 wheel psutil transformers==4.44.2 numpy==2.0.2 && uv pip install flash-attn --no-build-isolation && python -m pip install 'unsloth[cu121_ampere_torch240] @ git+https://github.com/unslothai/unsloth.git@%[2]s'", sourceVenv, unslothCommitOrTag)).Root()
 
 		version := version.Version
 		if version == "" {
@@ -66,7 +66,7 @@ func Aikit2LLB(c *config.FineTuneConfig) llb.State {
 		state = state.Run(utils.Shf("%[1]s && %[2]s && python -m target_unsloth", nvidiaMknod, sourceVenv), llb.Security(llb.SecurityModeInsecure)).Root()
 
 		// copy gguf to scratch which will be the output
-		const inputFile = "*.gguf"
+		const inputFile = "model/*.gguf"
 		copyOpts := []llb.CopyOption{}
 		copyOpts = append(copyOpts, &llb.CopyInfo{AllowWildcard: true})
 		outputFile := fmt.Sprintf("%s-%s.gguf", c.Output.Name, c.Output.Quantize)
diff --git a/pkg/finetune/target_unsloth.py b/pkg/finetune/target_unsloth.py
index e010d773..e466bb58 100644
--- a/pkg/finetune/target_unsloth.py
+++ b/pkg/finetune/target_unsloth.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from unsloth import is_bfloat16_supported
+from transformers import TrainingArguments, DataCollatorForSeq2Seq
 from unsloth import FastLanguageModel
 import torch
 from trl import SFTTrainer
@@ -71,13 +73,14 @@ def formatting_prompts_func(examples):
 else:
     dataset = load_dataset(source, split = "train")
 
-dataset = dataset.map(formatting_prompts_func, batched = True)
+dataset = dataset.map(formatting_prompts_func, batched=True)
 
 trainer = SFTTrainer(
     model=model,
     train_dataset=dataset,
     dataset_text_field="text",
     max_seq_length=max_seq_length,
+    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),
     tokenizer=tokenizer,
     dataset_num_proc = 2,
     packing = cfg.get('packing'), # Can make training 5x faster for short sequences.
@@ -87,8 +90,8 @@ def formatting_prompts_func(examples):
         warmup_steps=cfg.get('warmupSteps'),
         max_steps=cfg.get('maxSteps'),
         learning_rate = cfg.get('learningRate'),
-        fp16=not torch.cuda.is_bf16_supported(),
-        bf16=torch.cuda.is_bf16_supported(),
+        fp16=not is_bfloat16_supported(),
+        bf16=is_bfloat16_supported(),
         logging_steps=cfg.get('loggingSteps'),
         optim=cfg.get('optimizer'),
         weight_decay = cfg.get('weightDecay'),
diff --git a/test/aikitfile-unsloth-custom.yaml b/test/aikitfile-unsloth-custom.yaml
index 2f255c6b..ca23fd7a 100644
--- a/test/aikitfile-unsloth-custom.yaml
+++ b/test/aikitfile-unsloth-custom.yaml
@@ -18,12 +18,8 @@ config: |
   - name: custom
     backend: llama
     parameters:
-      top_k: 80
-      temperature: 0.8
-      top_p: 0.7
       model: model-q4_k_m.gguf
     context_size: 4096
-    gpu_layers: 35
     f16: true
     mmap: true
     template:
diff --git a/test/aikitfile-unsloth.yaml b/test/aikitfile-unsloth.yaml
index b23c3f18..bd75ab64 100644
--- a/test/aikitfile-unsloth.yaml
+++ b/test/aikitfile-unsloth.yaml
@@ -1,6 +1,6 @@
 #syntax=aikit:test
 apiVersion: v1alpha1
-baseModel: unsloth/llama-3-8b-bnb-4bit
+baseModel: unsloth/Meta-Llama-3.1-8B
 datasets:
   - source: "yahma/alpaca-cleaned"
     type: alpaca