Hugging Face's logo

ItachiSan
/
LiquidCode-2.5-24B-A2B

Text Generation
Transformers
Safetensors
English
lfm2_moe
text-generation-inference
unsloth
conversational
Model card Files
xet
 Community
LiquidCode-2.5-24B-A2B / notebook.py
ItachiSan's picture
ItachiSan
Upload the training notebook
26ec365 verified
Raw History Blame Contribute Delete
35.6 kB
# /// script
# requires-python = ">=3.13"
# dependencies = [
# "datasets==4.8.5",
# "huggingface-hub==1.33.0",
# "torch==2.14.1",
# "transformers==5.17.0",
# "trl==1.15.0",
# "unsloth==2026.11.2",
# "unsloth-zoo==2026.11.3",
# ]
# ///
import marimo
__generated_with = "0.25.2"
app = marimo.App(width="medium", auto_download=["html"])
@app.cell(hide_code=True)
def _():
# Let's start!
import marimo as mo
mo.md("# Unsloth Finetuning notebook!")
return (mo,)
@app.cell(hide_code=True)
def _(mo):
# Definition of mo_exec and mo_exec stream
from subprocess import (
run as _run,
Popen as _Popen,
PIPE as _PIPE,
STDOUT as _STDOUT,
)
def mo_exec(*process_args, cwd=None):
cmdline = " ".join(process_args).split(" ")
result = _run(
cmdline,
capture_output=True,
text=True,
cwd=cwd,
)
result_text = f"""
**Command**: _{" ".join(process_args)}_
**Exit code**: {result.returncode}
"""
if result.stdout:
result_text += f"""
**stdout**
```
{result.stdout}
```
"""
if result.stderr:
result_text += f"""
**stderr**
```
{result.stderr}
```
"""
return mo.md(result_text)
def mo_stream(*process_args, cwd=None):
print(f"🚀 Running: {' '.join(process_args)} ...")
# Run process with real-time output piping
cmdline = " ".join(process_args).split(" ")
process = _Popen(
cmdline,
stdout=_PIPE,
stderr=_STDOUT,
text=True,
cwd=cwd,
bufsize=1,
errors="replace", # Replace invalid bytes
)
# Print lines as they come in so you see live progress
for line in process.stdout:
print(line, end="")
process.wait()
if process.returncode != 0:
raise RuntimeError(
f"❌ Command failed with exit code {process.returncode}"
)
return mo_exec, mo_stream
@app.cell
def _(mo_stream):
# Important setup!
import os
# Avoid Unsloth custom modules compilation. Make sure we can run properly
os.environ["UNSLOTH_FULLGRAPH"] = "0"
os.environ["UNSLOTH_NO_COMPILE"] = "1"
os.environ["UNSLOTH_DISABLE_MOE_GROUPED_MM"] = "1"
# Extra debug messages in case of errors
os.environ.setdefault("TORCHDYNAMO_VERBOSE", "1")
os.environ.setdefault("TORCH_LOGS", "+dynamo")
# Hide HuggingFace download messages
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
# Set prebuilt llama.cpp location
llama_cpp_dir = "/marimo/unsloth/llama.cpp"
os.environ.setdefault("UNSLOTH_LLAMA_CPP_PATH", llama_cpp_dir)
# Do the ugly. Install Unsloth stuff via CLI
mo_stream("unsloth install-kernels")
from huggingface_hub import HfApi
hf = HfApi() # HF_TOKEN is in .env
print("Connected to HuggingFace!")
return hf, llama_cpp_dir, os
@app.cell
def _(mo_stream):
# Show required wheels
mo_stream("unsloth install-kernels --dry-run")
print(f"""
These are the packages to be installed via the cell below or the *Packages* section.
If needed, install these dependencies (found by hand; better the one from the cell above):
- causal-conv1d: https://github.com/protopia-ai/causal-conv1d/releases/download/v1.6.2.post7/causal_conv1d-1.6.2.post7+cu130torch2.14cxx11abiTRUE-cp313-cp313-linux_x86_64.whl
- xformers: https://download.pytorch.org/whl/cu130/xformers-0.0.35-py39-none-manylinux_2_28_x86_64.whl
- flash-attn: https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.10.0/flash_attn-2.8.3%2Bcu130torch2.14-cp313-cp313-linux_x86_64.whl
""")
return
@app.cell
def _(mo_stream):
# Check Unsloth status
mo_stream("unsloth doctor")
return
@app.cell
def _(mo_exec):
mo_exec("env")
return
@app.cell
def _(mo, os):
# Show system status
import sys
import torch
# Silly dependency, but ensures we run the cell with the os import which has all the important stuff.
_unsloth_patch = True if os.getenv("UNSLOTH_PATCHED") == "1" else False
fa_avail, xf_avail, cc_avail = False, False, False
try:
import flash_attn as _fa
fa_avail = _fa.__version__
except ImportError:
pass
try:
import xformers as _xf
xf_avail = _xf.__version__
except ImportError:
pass
try:
import causal_conv1d as _cc
cc_avail = _cc.__version__
except ImportError:
pass
mo.md(f"""
## System status
|**Metric** |**Value** |
|----------------|-------------------------------|
|Python version |{sys.version} |
|PyTorch version |{torch.__version__} |
|CUDA available |{torch.cuda.is_available()} |
|CUDA version |{torch.version.cuda} |
|GPU |{torch.cuda.get_device_name(0)}|
|Flash Attention |{fa_avail} |
|XFormers |{xf_avail} |
|causal-conv1d |{cc_avail} |
|Unsloth applied |{_unsloth_patch} |
""")
return (torch,)
@app.cell(hide_code=True)
def _():
# Define models and datasets
# Here are all the models and datasets we want to use for our finetuning
MODELS = [
"LiquidAI/LFM2.5-24B-A2B",
"LiquidAI/LFM2.5-8B-A1B",
"LiquidAI/LFM2.5-2.6B",
"LiquidAI/LFM3-1.7B",
]
DATASETS = [
"nvidia/Nemotron-SFT-SWE-v3",
"empero-ai/gpt-5.6-luna-sft-900x",
]
return DATASETS, MODELS
@app.cell(hide_code=True)
def _(DATASETS, MODELS, mo):
# Nice dropdown for selection.
model_selection = mo.ui.dropdown(
options=MODELS, value=MODELS[0], label="Model to finetune"
)
dataset_selection = mo.ui.dropdown(
options=DATASETS,
value=DATASETS[0],
label="Dataset to use for finetuning",
)
mo.vstack([model_selection, dataset_selection])
return dataset_selection, model_selection
@app.cell(hide_code=True)
def _(dataset_selection, mo, model_selection):
# Fill the required variables for all next sections!
base_model_name, dataset_name = (
model_selection.value,
dataset_selection.value,
)
# Definition of the place where the model ends!
model_base_name = "-".join(base_model_name.split("-")[1:])
model_name = f"LiquidCode-2.5-{model_base_name}"
model_repo = f"ItachiSan/{model_name}"
model_gguf_repo = model_repo + "-GGUF"
mo.md(f"""
Selected combination:
- Base model: {base_model_name}
- Dataset: {dataset_name}
- Model name: {model_name}
- Model repository: {model_repo}
""")
return (
base_model_name,
dataset_name,
model_base_name,
model_gguf_repo,
model_name,
model_repo,
)
@app.cell
def _(mo):
mo.md("""
# Setting up crucial parameters!
""")
return
@app.function(hide_code=True)
# Function to calculate the training config.
def calculate_training_config(
vram_gb: float,
model_params_b: float,
num_gpus: int = 1,
target_gbs: int = 128,
max_seq_length: int = 2048,
use_4bit_qlora: bool = True,
):
"""
Calculates safe per-device batch size and gradient accumulation steps.
"""
print(f"--- Training Configuration Calculator ---")
print(f"Hardware: {vram_gb} GB VRAM | GPUs: {num_gpus}")
print(
f"Model Size: {model_params_b}B parameters | Max Seq Length: {max_seq_length}"
)
print(f"Target Global Batch Size (GBS): {target_gbs}\n")
# 1. Estimate Base Model VRAM Footprint
if use_4bit_qlora:
base_model_vram = (
model_params_b * 0.55
) # ~0.55 GB per billion params in 4-bit
else:
base_model_vram = (
model_params_b * 2.0
) # ~2 GB per billion params in 16-bit
overhead_vram = 4.0 # CUDA context, system overhead, LoRA adapters
available_vram_for_activations = (
vram_gb - base_model_vram - overhead_vram
)
if available_vram_for_activations <= 0:
return "❌ Error: Model is too large to fit in this VRAM size even with quantization!"
# 2. Approximate VRAM consumed per sequence at given sequence length
# (Empirical scaling factor for 4k/8k context transformers)
vram_per_sequence = (
(max_seq_length / 4096) * (model_params_b / 24) * 2.5
)
# 3. Calculate safe max per-device batch size (leaving a 15% safety buffer)
safe_vram_pool = available_vram_for_activations * 0.85
estimated_bs = int(safe_vram_pool / vram_per_sequence)
# Bound batch size between 1 and 32, rounding to clean numbers if possible
per_device_bs = max(1, min(estimated_bs, 32))
# 4. Calculate Gradient Accumulation Steps to match Target GBS
total_effective_bs_per_step = per_device_bs * num_gpus
if target_gbs % total_effective_bs_per_step != 0:
# Round up or find closest clean division
grad_accum_steps = max(
1, round(target_gbs / total_effective_bs_per_step)
)
actual_gbs = per_device_bs * grad_accum_steps * num_gpus
else:
grad_accum_steps = target_gbs // total_effective_bs_per_step
actual_gbs = target_gbs
print(f"✅ Recommended Settings:")
print(f" per_device_train_batch_size = {per_device_bs}")
print(f" gradient_accumulation_steps = {grad_accum_steps}")
print(
f" Resulting Global Batch Size = {actual_gbs} (Target was {target_gbs})"
)
print(
f" Estimated VRAM utilization: ~{int(base_model_vram + (per_device_bs * vram_per_sequence) + overhead_vram)} GB / {vram_gb} GB\n"
)
return per_device_bs, grad_accum_steps
@app.cell(hide_code=True)
def _(mo_exec):
# GPU check
mo_exec("nvidia-smi")
return
@app.cell(hide_code=True)
def _(mo, model_base_name):
mo.md(f"""
Planned model: {model_base_name}
GPU available: see above.
Use the data for the next calculation.
""")
return
@app.cell
def _(model_base_name):
# --- Example: Your exact setup (96GB VRAM, 24B Model, 4096 Seq Length) ---
max_seq_length = 2048 # Adjust as needed (e.g., 4096)
_target_gbs = 128 # Not sure how to get this, got it from Gemini
_gpu_size_gb = 96.0 # GPU VRAM GB, get it from nvidia-smi block above
# Detecting model size from the name
_model_params_b = model_base_name.rstrip("B")
if "-" in _model_params_b:
_model_params_b = _model_params_b.split("-")[0].rstrip("B")
_model_params_b = float(_model_params_b)
batch_size, gradient_steps = calculate_training_config(
vram_gb=_gpu_size_gb,
model_params_b=_model_params_b,
num_gpus=1,
target_gbs=_target_gbs,
max_seq_length=max_seq_length,
use_4bit_qlora=True,
)
# Nice info print
print(f"""\
Configured parameters:
- Total model size: {_model_params_b}
- Available VRAM: {_gpu_size_gb}
- Sequence length: {max_seq_length}
- Batch size: {batch_size}
- Gradient steps: {gradient_steps}
""")
return batch_size, gradient_steps, max_seq_length
@app.cell(hide_code=True)
def _(mo):
mo.md("""
# Training the model!
""")
return
@app.cell
def _(base_model_name, max_seq_length):
# Load the model
from unsloth import FastLanguageModel
dtype = None # None for auto-detection
load_in_4bit = True # Use 4-bit quantization to save memory
# 1. Load base model and tokenizer
base_model, tokenizer = FastLanguageModel.from_pretrained(
model_name=base_model_name, # Or your specific base model path
max_seq_length=max_seq_length,
dtype=dtype,
load_in_4bit=load_in_4bit,
)
base_model, tokenizer
return FastLanguageModel, base_model, tokenizer
@app.cell
def _(base_model):
# Try detecting modules automatically
detected_modules = set()
# Simple map of modules to skip. Their bool value is to print the item once.
_modules_to_skip = {
# Modules to skip for LFM2 MoE
"q_layernorm": False,
"k_layernorm": False,
"experts": False,
"gate": False,
}
for _layer in base_model.model.layers:
for _name, _ in _layer.named_modules():
# The modules are in form
# sublayer(?) . module
# If no dot is there, we have nothing to do here.
if "." not in _name:
continue
_, _module_name = _name.split(".", maxsplit=1)
# Handle cleanly re-detection
if "." in _module_name:
_module_name = _module_name.split(".", maxsplit=1)[0]
if _module_name in _modules_to_skip:
if _modules_to_skip[_module_name] is False:
print("Skip:", _module_name)
_modules_to_skip[_module_name] = True
continue
if _module_name not in detected_modules:
print("Found:", _module_name)
detected_modules.add(_module_name)
detected_modules
return (detected_modules,)
@app.cell
def _(FastLanguageModel, base_model, detected_modules, tokenizer):
# 2. Configure the LoRA model and get the adapters
# There are hard-coded values found by inspection and chat with AI
_target_modules = [
## Default values for Unsloth, with extra comments from Gemini
# 1. Standard Attention Projections
"q_proj",
"k_proj",
"v_proj",
"o_proj",
# Unused in Lfm2Moe, but keep for safety
"gate_proj",
"up_proj",
"down_proj",
## Added extra modules that should properly target MoE
# 2. Hybrid Short Convolution Layers (Local temporal mixing)
"in_proj",
"out_proj",
"conv",
# 3. MoE Expert Projections & Router (Layer naming for LFM MoE)
"ffn_gate_exps",
"ffn_up_exps",
"ffn_down_exps",
"ffn_gate_inp",
# 4. Dense FFN layers (Present in non-expert blocks)
# "ffn_gate", "ffn_up", "ffn_down",
"w1",
"w2",
"w3", # Fixed after checking the values below
]
# Stick also out own modules for safety.
_target_modules = sorted(
list(set(_target_modules + list(detected_modules)))
)
print("Training following modules:", _target_modules)
# Now, get the trainable model
model = FastLanguageModel.get_peft_model(
base_model,
target_modules=_target_modules,
r=32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
lora_alpha=64,
lora_dropout=0, # Optimized to 0 by default
bias="none", # Supports "none", "all"
use_gradient_checkpointing="unsloth", # Saves VRAM
random_state=3407,
use_rslora=False, # Rank stabilized LoRA support
loftq_config=None, # LoftQ
)
model, tokenizer
return (model,)
@app.cell
def _(model):
_trainable_params = 0
_total_params = 0
for _name, _param in model.named_parameters():
_total_params += _param.numel()
if _param.requires_grad:
_trainable_params += _param.numel()
print(f"Total parameters: {_total_params:,}")
print(
f"Trainable parameters: {_trainable_params:,} ({100 * _trainable_params / _total_params:.2f}%)"
)
return
@app.cell(hide_code=True)
def _(tokenizer):
print("Tokenizer template:")
print(tokenizer.get_chat_template())
return
@app.cell(hide_code=True)
def _():
# Shared preprocessing functions
import json
def rjl(value):
"""
Hand-made Recursive JSON Loads.
This is because AI sucks at writing this.
JSON allows for
- iterables: dict, list
- values: bool, null, int, float, string
Some datasets have JSON nested in strings and the default json.loads does
not handle those. This function does.
Args:
value: The value to convert. Does nothing unless it is a JSON string.
Returns:
The parsed Python object, or the original value if not a JSON string.
"""
# Handle iterables first. We want to be sure to loop when needed.
if isinstance(value, list):
return [rjl(item) for item in value]
if isinstance(value, dict):
return {key: rjl(val) for key, val in value.items()}
# Other simple case: if not string, no need to json.loads
if not isinstance(value, str):
return value
# Here we have a string, let's see if it is a JSON
try:
decoded = json.loads(value)
except (json.JSONDecodeError, TypeError):
# The string is not a JSON, so return it.
return value
# Now we decoded something, but we don't know what.
# So, let's recurse! This way we handle all cases
return rjl(decoded)
return (rjl,)
@app.cell(hide_code=True)
def _(c, rjl):
# Nemotron specific preprocessing
def _fix_nemotron_chat(chat):
fixed_chat = []
for message in chat:
fixed_message = message.copy()
# Ensure tool calls are data and not JSON strings
fixed_message["tool_calls"] = rjl(message["tool_calls"])
# If we are an assistant, ensure the tool calls is a list otherwise the template will fail
if fixed_message["role"] == "assistant":
if fixed_message["tool_calls"] is None:
fixed_message["tool_calls"] = []
fixed_chat.append(fixed_message)
return fixed_chat
def preprocess_nemotron_sft(examples):
# First, get the proper column.
chats = examples["messages"]
# Check if we are batching or not and return the result accordingly
batched = not isinstance(chats[0], dict)
return (
[_fix_nemotron_chat(c) for c in chats]
if batched
else _fix_nemotron_chat(c)
)
return (preprocess_nemotron_sft,)
@app.cell
def _(tokenizer):
# Process the dataset if needed
def generate_prompts_text(
examples,
dataset_name: str | None = None,
dataset_functions: dict[str, any] | None = None,
):
# If we need to do any extra processing, here we do:
if dataset_name and dataset_functions:
examples = dataset_functions[dataset_name](examples)
return {
"text": tokenizer.apply_chat_template(
examples,
# We tokenize later with Unsloth
tokenize=False,
# Not sure why. Perhaps check?
add_generation_prompt=False,
)
}
return (generate_prompts_text,)
@app.cell
def _(model_name):
dataset_repo = f"ItachiSan/{model_name}-SFT-chats"
print(f"Dataset repo: {dataset_repo}")
return (dataset_repo,)
@app.cell
def _(
dataset_name,
dataset_repo,
generate_prompts_text,
model_name,
os,
preprocess_nemotron_sft,
):
# v2: do evaluation
from datasets import load_dataset
# Fetch the dataset online. Disable only the first time we process the dataset.
_fetch_online = True
# Whether to push the dataset for refetching. Requires the repo to exist.
_push_online = True
# Choose how much to use for evaluation
eval_split_size = 0.01 # 1%
# Good number to detect how many processes we can run
num_proc = int(os.cpu_count() * 0.75) + 1
if not _fetch_online:
# Load the dataset, efficiently!
dataset = load_dataset(dataset_name, split="train", num_proc=num_proc)
# Format conversations using the model's chat template.
# Adjust the column name (e.g., "conversations" or "messages")
# to match your dataset's schema.
# IMPORTANT CHANGES!
# We need to drop the other columns, because... we suck
# The main issue is that, in our patching of the Nemotron dataset, we convert data.
# Then, pyarrow complains the JSON data we parse to build the text has different shape (duh...)
# We bypass the problem by dropping all the other columns, because yes.
chats = dataset.map(
generate_prompts_text,
fn_kwargs=dict(
dataset_name=dataset_name,
dataset_functions={
## Add special handlers connections in here if needed.
# Select proper column and decode JSON strings
"nvidia/Nemotron-SFT-SWE-v3": preprocess_nemotron_sft,
},
),
batched=True,
remove_columns=dataset.column_names,
num_proc=num_proc,
)
chats_split = chats.train_test_split(
test_size=0.01, # 1% for test size can also be an integer for # of rows
shuffle=True, # Should always set to True!
seed=3407,
)
if _push_online:
chats_split.push_to_hub(
f"ItachiSan/{model_name}-SFT-chats",
private=True,
num_proc=num_proc,
)
else:
chats_split = load_dataset(dataset_repo, num_proc=num_proc)
# Split the dataset
train_dataset = chats_split["train"] # Dataset for training
eval_dataset = chats_split["test"] # Dataset for evaluation
train_dataset, eval_dataset
return eval_dataset, num_proc, train_dataset
@app.cell
def _(train_dataset):
# Peek into the chats!
from random import randint as _randint
_chat_n = _randint(0, train_dataset.num_rows)
print(f"Example chat (row {_chat_n}):\n{train_dataset[_chat_n]['text']}")
return
@app.cell
def _(
batch_size,
eval_dataset,
gradient_steps,
max_seq_length,
model,
model_repo,
num_proc,
tokenizer,
torch,
train_dataset,
):
# Configure the SFTTrainer
from trl import SFTTrainer, SFTConfig
# Using the formula in the tool section, we can calculate the following.
training_steps = int(
round((train_dataset.num_rows / (batch_size * gradient_steps)) / 10)
* 10
)
warmup_steps = training_steps * 0.05
print(
f"Training {training_steps} steps with batch: {batch_size}, gradient: {gradient_steps}"
)
_sft_args = SFTConfig(
# Gemini recommended params
# Numbers in comments are earlier ones, now they are calculated via function above.
# Math is easy: 16x8 is based on GPU performance
# Then, n_steps are dataset_size / (16x8)
# Then warmup steps is 5% of the buffer
per_device_train_batch_size=batch_size,
gradient_accumulation_steps=gradient_steps,
max_steps=training_steps,
warmup_steps=warmup_steps,
# Log not too frequently (original: 1)
logging_steps=5,
# Better on long trainings (original: linear)
lr_scheduler_type="cosine",
# Same as before
learning_rate=2e-4,
fp16=not torch.cuda.is_bf16_supported(),
bf16=torch.cuda.is_bf16_supported(),
optim="adamw_8bit",
weight_decay=0.01,
seed=3407,
output_dir="outputs",
report_to="none",
# We do long training. Ensure we can resume.
save_strategy="steps",
save_steps=50,
save_total_limit=5,
# Store checkpoints on Hugging Face
push_to_hub=True,
hub_model_id=model_repo,
hub_strategy="checkpoint",
# v2: do eval
eval_strategy="steps", # Runs eval every few steps or epochs.
eval_steps=0.05, # Evaluate every 5%
per_device_eval_batch_size=batch_size, # Ensure we eval quickly
eval_accumulation_steps=gradient_steps, # You can increase this include of batch_size
# TODO: add early stopping: https://unsloth.ai/docs/basics/finetuning-from-last-checkpoint#how-do-i-do-early-stopping
)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
dataset_text_field="text", # Here we point at our single column
max_seq_length=max_seq_length,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
dataset_num_proc=num_proc,
packing=False, # TODO: check wheter packing caused some issues or not
args=_sft_args,
)
trainer
return (trainer,)
@app.cell(hide_code=True)
def _(mo, model_repo, os):
# Reload checkpoints if any
from huggingface_hub import snapshot_download
repo_id = model_repo
local_output_dir = "outputs"
# Automatically pull existing checkpoints from your HF Hub repo back into the sandbox
try:
print(
f"📥 Checking Hugging Face Hub for existing checkpoints in {repo_id}..."
)
snapshot_download(
repo_id=repo_id,
local_dir=local_output_dir,
allow_patterns=[
"checkpoint-*/**",
"last-checkpoint/**",
"*.json",
"*.safetensors",
],
)
print("✅ Checkpoints synced successfully from Hub.")
except Exception as e:
print(f"⚠️ No remote checkpoints found or repository is empty yet: {e}")
# Detect the checkpoint if any
existing_checkpoint = False
if os.path.isdir(local_output_dir):
existing_checkpoints = [
os.path.join(local_output_dir, d)
for d in os.listdir(local_output_dir)
if d.startswith("checkpoint-") or d == "last-checkpoint"
]
existing_checkpoints.sort(
key=lambda x: (
int(x.split("-")[-1])
if x.startswith("checkpoint") # covers local checkpoints
else 42069
) # for 'last-checkpoint' we use a very huge number
)
if existing_checkpoints: # pythonic len(list) > 0
existing_checkpoint = existing_checkpoints[-1]
mo.md(f"Checkpoint: {existing_checkpoint}")
return (existing_checkpoint,)
@app.cell
def _(mo):
mo.md("""
# DO IT!
""")
return
@app.cell
def _(os):
# Print env vars if needed to check if some things need to change
for _k, _v in sorted(os.environ.items(), key=lambda k: k[0]):
print(f"{_k}={_v}")
return
@app.cell
def _(existing_checkpoint, model, tokenizer, trainer):
# Finetune the model
_force_training = False
_starting_point = False
if not _force_training:
_starting_point = existing_checkpoint
# Do the thing!
trainer_stats = trainer.train(resume_from_checkpoint=_starting_point)
lora_output_dir = "/marimo/model/lora"
# Save the LoRA adapters and tokenizer
model.save_pretrained(lora_output_dir)
tokenizer.save_pretrained(lora_output_dir)
print("Model saved to", lora_output_dir)
trainer_stats
return
@app.cell
def _(model, model_repo, tokenizer):
# Publish all the things!
model.push_to_hub(model_repo)
tokenizer.push_to_hub(model_repo)
model.push_to_hub_merged(
model_repo,
tokenizer=tokenizer,
save_method="merged_16bit",
)
return
@app.cell(hide_code=True)
def _(mo):
mo.md("""
# Testing the model!
""")
return
@app.cell(hide_code=True)
def _(os, torch):
# If you get TorchDynamo errors like:
# File "/tmp/torchinductor_root/r7/cr75u36gyrpc7ppml5xne6rruogr56gjefh5axla5phsc2ix5gjp.py"
#, line 101, in call
# buf2 = torch.ops.unsloth_zoo.routed_nf4_gate_up.default(many_args...)
#NameError: name 's78' is not defined. Did you mean: 's68'?
# Run this cell. It clears the corrupted TorchInductor cache and reset dynamo graphs
import shutil
shutil.rmtree('/tmp/torchinductor_root', ignore_errors=True)
torch._dynamo.reset()
# Also, disable all compilation and Dynamo tracing for inference to avoid shape generation bugs
os.environ["UNSLOTH_NO_COMPILE"] = "1"
os.environ["UNSLOTH_FULLGRAPH"] = "0"
# If truly needed, run also the next line
#os.environ["TORCH_COMPILE"] = "0"
return
@app.cell
def _(model_repo, os):
# If available, load the LoRA directly. Otherwise, we fetch from the repo.
# Used both for chat and export
loadable_model_path = model_repo
_lora_model_path = "/marimo/model/lora"
if os.path.isfile(_lora_model_path + os.sep + "config.json"):
loadable_model_path = _lora_model_path
print("Model to load:", loadable_model_path)
return (loadable_model_path,)
@app.cell
def _(FastLanguageModel, loadable_model_path):
# Load base model + adapters directly
test_model, test_tokenizer = FastLanguageModel.from_pretrained(
model_name=loadable_model_path,
max_seq_length=2048,
dtype=None,
load_in_4bit=True,
)
# Enable fast inference mode
FastLanguageModel.for_inference(test_model)
test_model, test_tokenizer
return test_model, test_tokenizer
@app.cell
def _(mo, test_model, test_tokenizer):
# Hook our model in a chat box!
def ui_chat_with_our_model(messages: list[mo.ai.ChatMessage], config: mo.ai.ChatModelConfig):
# Match the custom messages objects to things we can work with:
chat_history = [
{"role": message.role, "content": message.content}
for message in messages
]
# 1. Get those inputs matched
inputs = test_tokenizer.apply_chat_template(
chat_history,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt",
).to("cuda")
# 3. Generate response
return test_model.generate(
input_ids=inputs,
use_cache=True,
max_new_tokens=config.max_tokens,
temperature=config.temperature,
top_p=config.top_p,
top_k=config.top_k,
#frequency_penalty=config.frequency_penalty,
#presence_penalty=config.presence_penalty,
)
mo.ui.chat(
ui_chat_with_our_model,
show_configuration_controls=True,
)
return
@app.cell
def _(mo):
mo.md("""
# Export time!
""")
return
@app.cell
def _(FastLanguageModel, loadable_model_path):
# Wire the model for export
print("🔄 Loading uploaded model weights for export to GGUF...")
export_model, export_tokenizer = FastLanguageModel.from_pretrained(
model_name=loadable_model_path,
max_seq_length=2048,
dtype=None,
load_in_4bit=True,
)
print(
"✅ Model loaded! Variables are `export_model` and `export_tokenizer`."
)
export_model, export_tokenizer
return export_model, export_tokenizer
@app.cell(hide_code=True)
def _(mo, os):
# Define change_dir to allow Unsloth to run properly.
from contextlib import contextmanager
@contextmanager
def change_dir(path):
prev_dir = os.getcwd()
if not os.path.isdir(path):
os.mkdir()
os.chdir(path)
try:
yield
finally:
os.chdir(prev_dir)
mo.md("""
`change_dir` is now defined.
Remember to use it for Unsloth GGUF operations as such:
```python
with change_dir("/home/marimo/.unsloth/llama.cpp"):
do_the_thing.
```
""")
return
@app.cell(hide_code=True)
def _(manual_model, manual_tokenizer):
print("🔄 Merging adapters to 16-bit safetensors...")
# Save the merged model as standard HF format first (bypasses GGUF conversion script)
manual_model.save_pretrained_merged(
"merged_model",
manual_tokenizer,
save_method="merged_16bit", # Merges LoRA into base weights
)
print("✅ Merged model successfully saved to ./merged_model_16bit!")
return
@app.cell
def _(llama_cpp_dir):
# Retrieve Unsloth llama.cpp
# Original referece here:
# https://github.com/unslothai/unsloth-zoo/blob/main/unsloth_zoo/llama_cpp.py
print("🔄 Fetching Unsloth llama.cpp ...")
from unsloth_zoo.llama_cpp import install_llama_cpp
install_llama_cpp(
llama_cpp_folder=llama_cpp_dir, gpu_support=True, print_output=True
)
print(f"✅ llama.cpp is now available at {llama_cpp_dir}")
return
@app.cell
def _(export_model, export_tokenizer):
# First, export the model as merged bits
model_merged_dir = "/marimo/model/merged"
export_model.save_pretrained_merged(
model_merged_dir, export_tokenizer, save_method="merged_16bit"
)
return (model_merged_dir,)
@app.cell
def _(llama_cpp_dir, mo_stream, model_merged_dir, model_name, os):
# Second, convert to F16 GGUF as first step
model_gguf_dir = "/marimo/model/gguf"
os.makedirs(model_gguf_dir, exist_ok=True)
def model_gguf_file(quant):
return f"{model_gguf_dir}/{model_name}-{quant}.gguf"
os.makedirs(model_gguf_dir, exist_ok=True)
_f16_model_path = model_gguf_file("F16")
if not os.path.isfile(_f16_model_path):
mo_stream(
"python",
f"{llama_cpp_dir}/convert_hf_to_gguf.py",
model_merged_dir,
"--outfile",
model_gguf_file("F16"),
"--outtype",
"f16",
"--split-max-size",
"50G",
cwd=llama_cpp_dir,
)
else:
print("Conversion to F16 GGUF already done!")
return model_gguf_dir, model_gguf_file
@app.cell
def _(llama_cpp_dir, mo_stream, model_gguf_file, os):
# Now quantize all the files!
quantize_quants = [
"Q2_K",
"Q3_K_M",
"Q4_K_M",
"Q5_K_M",
"Q6_K",
"Q8_0",
"TQ1_0",
"TQ2_0",
]
for _quant in quantize_quants:
# Handle special quant file naming
_model_gguf_file = model_gguf_file(_quant)
if _quant == "Q2_K":
_model_gguf_file = _model_gguf_file.replace(_quant, "Q2_K_M")
elif _quant == "Q6_K":
_model_gguf_file = _model_gguf_file.replace(_quant, "Q6_K_M")
elif _quant == "Q8_0":
_model_gguf_file = _model_gguf_file.replace(_quant, "Q8_K_M")
# Avoid doing the work again.
if os.path.isfile(_model_gguf_file):
print(
f"✅ File for quant {_quant} already exists: {_model_gguf_file}"
)
continue
print(f"✨ Working on quant: {_quant}")
mo_stream(
f"{llama_cpp_dir}/llama-quantize",
model_gguf_file("F16"),
_model_gguf_file,
_quant, # quant type
"16", # number of threads
)
print(f"✅ File for quant {_quant} now at: {_model_gguf_file}")
return
@app.cell
def _(hf, model_gguf_dir, model_gguf_repo):
# Finally push the stuff!
print("Starting the upload...")
hf.upload_folder(
folder_path=model_gguf_dir,
repo_id=model_gguf_repo,
repo_type="model",
allow_patterns=["*.gguf", "*.json", "*.md", "*.txt"],
)
print("All GGUF files uploaded successfully!")
return
@app.cell(hide_code=True)
def _(mo):
mo.md("""
# Experiments zone!
""")
return
@app.cell(hide_code=True)
def _(mo):
mo.md(r"""
Nothing here for now...
""")
return
if __name__ == "__main__":
app.run()