llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Gemma training script

Browse Source
This commit is contained in:
Niels Geens
2025-05-20 20:46:16 +02:00
parent d4a4902b94
commit f2c70cd05d
14 changed files with 394 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.gitmodules vendored Normal file
View File

@@ -0,0 +1,3 @@
[submodule "modules/llama.cpp"]
path = modules/llama.cpp
url = https://github.com/ggml-org/llama.cpp.git

70
Dockerfile
View File

@@ -31,6 +31,11 @@ RUN rm -rf /var/lib/apt/lists/*
RUN curl https://sh.rustup.rs -sSf | bash -s -- -y
ENV PATH="/root/.cargo/bin:${PATH}"
# Install Node.js
RUN curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash
ENV NVM_DIR=/root/.nvm
RUN . "$NVM_DIR/nvm.sh" && nvm install node
# Install llama.cpp
RUN curl -O https://git.nielsgeens.be/api/packages/llm/generic/llama.cpp/b5269/llama.cpp.tar
RUN tar -xf /llama.cpp.tar -C /usr/local/bin --wildcards --no-anchored "llama-*"
@@ -39,75 +44,10 @@ RUN rm llama.cpp.tar
# Create directory structure
RUN mkdir -p \
/root/sia \
/root/sia/scripts \
/root/data/iterations \
/root/data/user \
/root/data/tasks \
/root/data/environment \
/root/models \
/root/desktop \
/root/venvs
# Tool ITB setup
FROM base AS itb-env
RUN python3 -m venv /root/venvs/itb --system-site-packages
COPY ./tools/itb/setup.py /root/sia/tools/itb/setup.py
RUN /root/venvs/itb/bin/python /root/sia/tools/itb/setup.py egg_info
RUN --mount=type=cache,target=/root/.cache/pip \
/root/venvs/itb/bin/pip install -r *.egg-info/requires.txt
RUN rm -rf *.egg-info/
# SIA core setup
FROM base AS sia-env
RUN python3 -m venv /root/venvs/sia --system-site-packages
COPY ./setup.py /root/sia/setup.py
RUN /root/venvs/sia/bin/python /root/sia/setup.py egg_info
RUN sed -i '/\/root\/sia\/lib/d' *.egg-info/requires.txt
RUN --mount=type=cache,target=/root/.cache/pip \
/root/venvs/sia/bin/pip install -r *.egg-info/requires.txt
RUN rm -rf *.egg-info/
# Notebook setup
FROM base AS notebook-env
RUN python3 -m venv /root/venvs/notebook --system-site-packages
RUN --mount=type=cache,target=/root/.cache/pip \
/root/venvs/notebook/bin/pip install \
jupyter \
ipykernel \
ipywidgets
RUN /root/venvs/notebook/bin/ipython kernel install --name=notebook
# Web frontend build
FROM base
# AS web-build
ENV NODE_VERSION=16.13.0
RUN curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.0/install.sh | bash
ENV NVM_DIR=/root/.nvm
RUN . "$NVM_DIR/nvm.sh" && nvm install ${NODE_VERSION}
RUN . "$NVM_DIR/nvm.sh" && nvm use v${NODE_VERSION}
RUN . "$NVM_DIR/nvm.sh" && nvm alias default v${NODE_VERSION}
ENV PATH="/root/.nvm/versions/node/v${NODE_VERSION}/bin/:${PATH}"
FROM node:20-alpine AS web-build
WORKDIR /app
COPY web/package.json ./
RUN npm install
RUN rm -rf /root/.npm/_cacache
COPY web .
RUN npm run build
# Final image
FROM base
COPY --from=itb-env /root/venvs/itb /root/venvs/itb
COPY --from=sia-env /root/venvs/sia /root/venvs/sia
COPY --from=notebook-env /root/venvs/notebook /root/venvs/notebook
COPY --from=web-build /app/dist /root/static/
RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /root/.bashrc
RUN echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /etc/profile
WORKDIR /root/desktop
ENTRYPOINT ["/bin/bash", "-c"]

6
lib/llm_engine_utils/pyproject.toml
View File

@@ -7,7 +7,7 @@ name = "llm_engine_utils"
version = "0.1.0"
requires-python = ">=3.8"
[project.optional-dependencies]
dataset = [
"torch>=4.0.0",
dependencies = [
"datasets>=3.0.0",
"torch>=2.0.0",
]

5
lib/llm_engine_utils/src/llm_engine_utils/__init__.py
View File

@@ -1,7 +1,4 @@
try:
from . import dataset
except ImportError:
pass
from . import dataset
from . import iterators
from . import protocol
from .llm_engine import LlmEngine

92
llm_config.toml
View File

@@ -1,91 +1,3 @@
gemma = """/root/venvs/gemma_infer/bin/gemma_infer --model /root/models/gemma3_1b.gguf --tokenizer /root/models/gemma3_1b_tokenizer"""
gemma = """gemma_infer --model /root/models/current/model.gguf --tokenizer /root/models/current/tokenizer"""
#mistral = """/root/venvs/mistral_infer/bin/mistral_infer"""
default = """#!/bin/bash
# Mock LLM engine for testing
# This script simulates an LLM engine subprocess that responds to the three commands:
# <token_limit/>, <token_count>...</token_count>, and <infer_xml>...</infer_xml>
# Function to read XML input until a complete closing tag is found
read_xml_input() {
local input=""
local line
local char_count=0
# Read until we get a complete XML command
while IFS= read -r line; do
input="${input}${line}"
char_count=$((char_count + ${#line}))
# Debug the actual content (first 30 chars)
if [[ "$input" == *"<token_limit/>"* ]]; then
printf "1024"
printf "\\004" # EOT character (hex 04)
return
fi
if [[ "$input" == *"<token_count>"*"</token_count>"* ]]; then
printf "405"
printf "\\004" # EOT character (hex 04)
return
fi
if [[ "$input" == *"<infer_xml>"*"</infer_xml>"* ]]; then
generate_response
return
fi
done
}
# Function to generate a response token by token
generate_response() {
printf "<"
sleep 0.3
printf "reason"
sleep 0.3
printf "ing"
sleep 0.3
printf ">"
sleep 0.3
printf "This"
sleep 0.3
printf " is"
sleep 0.3
printf " a"
sleep 0.3
printf " test"
sleep 0.3
printf " response."
sleep 0.3
printf "</"
sleep 0.3
printf "reason"
sleep 0.3
printf "ing"
sleep 0.3
printf ">"
sleep 0.3
printf "\\004"
}
# Main loop - keep reading input and responding
iteration=0
while true; do
iteration=$((iteration + 1))
read_xml_input
done"""
default = """/root/sia/scripts/mock_llm.sh"""

1
modules/llama.cpp Submodule

Submodule modules/llama.cpp added at be0239693c

47
scripts/bootstrap.sh
View File

@@ -1,9 +1,12 @@
#!/bin/bash
#SIA_INSTALL_NO_ITB=1
#SIA_INSTALL_NO_TRAIN=1
#SIA_INSTALL_NO_NOTEBOOK=1
#SIA_INSTALL_NO_WEB=1
SIA_INSTALL_NO_NOTEBOOK=1
#SIA_INSTALL_NO_CORE=1
SIA_INSTALL_NO_ITB=1
SIA_INSTALL_NO_MISTRAL_INFER=1
#SIA_INSTALL_NO_GEMMA_INFER=1
#SIA_INSTALL_NO_GEMMA_TRAIN=1
trap 'return 1' ERR
@@ -29,11 +32,14 @@ apt-get install -y \
build-essential \
cmake \
cuda-toolkit \
curl \
git \
gnupg \
jq \
libcurl4-nss-dev \
python3-dev \
python3-venv \
tmux \
vim \
wget \
;
@@ -48,18 +54,15 @@ if [ -z "${SIA_INSTALL_NO_ITB}" ]; then
fi
# Install Rust
if [ -z "${SIA_INSTALL_NO_TRAIN}" ]; then
curl https://sh.rustup.rs -sSf | bash -s -- -y
export PATH="/root/.cargo/bin:${PATH}"
fi
curl https://sh.rustup.rs -sSf | bash -s -- -y
export PATH="/root/.cargo/bin:${PATH}"
# Install Node.js
if [ -z "${SIA_INSTALL_NO_CORE}" ]; then
echo "Installing Node.js..."
curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash
export NVM_DIR="/root/.nvm"
[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"
nvm install node
export NVM_DIR=/root/.nvm
. "$NVM_DIR/nvm.sh" && nvm install node
fi
# Install llama.cpp
@@ -87,36 +90,14 @@ if [ ! -d "/root/sia" ]; then
git config --global core.editor vim
fi
# Build web interface
if [ -z "${SIA_INSTALL_NO_CORE}" ]; then
echo "Building web interface"
cd "/root/sia/web"
npm install
npm run build
ln -s "/root/sia/web/dist" "/root/static"
fi
# Install SIA dependencies
cd /root/desktop
source "/root/sia/scripts/install.sh"
# Add venvs to path in .bashrc
if ! grep -q "source /root/sia/scripts/add_venvs_to_path.sh" /root/.bashrc; then
echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /root/.bashrc
fi
# Add venvs to path in profile
if ! grep -q "source /root/sia/scripts/add_venvs_to_path.sh" /etc/profile; then
echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /etc/profile
fi
# Add venvs to path in current shell
source /root/sia/scripts/add_venvs_to_path.sh
# Finetune model
if [ -z "${SIA_INSTALL_NO_TRAIN}" ]; then
echo "Finetuning model..."
train /root/models/bootstrap
/root/sia/scripts/train.sh /root/models/bootstrap
ln -s /root/models/bootstrap /root/models/current
fi

44
scripts/install.sh
View File

@@ -1,46 +1,68 @@
#!/bin/bash
#SIA_INSTALL_NO_WEB=1
SIA_INSTALL_NO_NOTEBOOK=1
#SIA_INSTALL_NO_CORE=1
SIA_INSTALL_NO_ITB=1
SIA_INSTALL_NO_MISTRAL_INFER=1
#SIA_INSTALL_NO_GEMMA_INFER=1
SIA_INSTALL_NO_TRAIN=1
#SIA_INSTALL_NO_GEMMA_TRAIN=1
cd "/root/desktop"
if [ -z "${SIA_INSTALL_NO_WEB}" ]; then
echo "Building web interface"
(
cd "/root/sia/web"
npm install
npm run build
ln -s "/root/sia/web/dist" "/root/static"
)
fi
if [ -z "${SIA_INSTALL_NO_NOTEBOOK}" ]; then
echo "Installing venv for running notebooks..."
echo "Installing venv for running notebooks"
python3 -m venv /root/venvs/notebook
/root/venvs/notebook/bin/pip install jupyter ipykernel ipywidgets
/root/venvs/notebook/bin/ipython kernel install --name=notebook
fi
if [ -z "${SIA_INSTALL_NO_CORE}" ]; then
echo "Installing SIA core..."
echo "Installing SIA core"
python3 -m venv /root/venvs/sia
/root/venvs/sia/bin/pip install -e /root/sia
fi
if [ -z "${SIA_INSTALL_NO_ITB}" ]; then
echo "Installing ITB tool..."
echo "Installing ITB tool"
python3 -m venv /root/venvs/itb
/root/venvs/itb/bin/pip install -e /root/sia/tools/itb
fi
if [ -z "${SIA_INSTALL_NO_MISTRAL_INFER}" ]; then
echo "Installing venv for mistral inference..."
echo "Installing venv for mistral inference"
python3 -m venv /root/venvs/mistral_infer
/root/venvs/mistral_infer/bin/pip install -e /root/sia/tools/mistral_infer
fi
if [ -z "${SIA_INSTALL_NO_GEMMA_INFER}" ]; then
echo "Installing venv for gemma inference..."
echo "Installing venv for gemma inference"
python3 -m venv /root/venvs/gemma_infer
/root/venvs/gemma_infer/bin/pip install -e /root/sia/tools/gemma_infer
fi
if [ -z "${SIA_INSTALL_NO_TRAIN}" ]; then
echo "Installing Train tool..."
python3 -m venv /root/venvs/train
/root/venvs/train/bin/pip install -e /root/sia/tools/train
/root/venvs/train/bin/ipython kernel install --name=train
if [ -z "${SIA_INSTALL_NO_GEMMA_TRAIN}" ]; then
echo "Installing venv for gemma training"
python3 -m venv /root/venvs/gemma_train
/root/venvs/gemma_train/bin/pip install -e /root/sia/tools/gemma_train
fi
if ! grep -q "source /root/sia/scripts/add_venvs_to_path.sh" /root/.bashrc; then
echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /root/.bashrc
fi
if ! grep -q "source /root/sia/scripts/add_venvs_to_path.sh" /etc/profile; then
echo 'source /root/sia/scripts/add_venvs_to_path.sh' >> /etc/profile
fi
source /root/sia/scripts/add_venvs_to_path.sh

89
scripts/mock_llm.sh Normal file
View File

@@ -0,0 +1,89 @@
#!/bin/bash
# Mock LLM engine for testing
# This script simulates an LLM engine subprocess that responds to the three commands:
# <token_limit/>, <token_count>...</token_count>, and <infer_xml>...</infer_xml>
# Function to read XML input until a complete closing tag is found
read_xml_input() {
local input=""
local line
local char_count=0
# Read until we get a complete XML command
while IFS= read -r line; do
input="${input}${line}"
char_count=$((char_count + ${#line}))
# Debug the actual content (first 30 chars)
if [[ "$input" == *"<token_limit/>"* ]]; then
printf "1024"
printf "\\004" # EOT character (hex 04)
return
fi
if [[ "$input" == *"<token_count>"*"</token_count>"* ]]; then
printf "405"
printf "\\004" # EOT character (hex 04)
return
fi
if [[ "$input" == *"<infer_xml>"*"</infer_xml>"* ]]; then
generate_response
return
fi
done
}
# Function to generate a response token by token
generate_response() {
printf "<"
sleep 0.3
printf "reason"
sleep 0.3
printf "ing"
sleep 0.3
printf ">"
sleep 0.3
printf "This"
sleep 0.3
printf " is"
sleep 0.3
printf " a"
sleep 0.3
printf " test"
sleep 0.3
printf " response."
sleep 0.3
printf "</"
sleep 0.3
printf "reason"
sleep 0.3
printf "ing"
sleep 0.3
printf ">"
sleep 0.3
printf "\\004"
}
# Main loop - keep reading input and responding
iteration=0
while true; do
iteration=$((iteration + 1))
read_xml_input
done

17
scripts/train.sh Normal file
View File

@@ -0,0 +1,17 @@
#!/bin/bash
set -e
SIA_DIR="/root/sia"
OUTPUT_DIR="${1:-/root/models/$(cd "$SIA_DIR" && git rev-parse HEAD)}"
echo "Output dir: $OUTPUT_DIR"
if [ -n "$(cd "$SIA_DIR" && git status --porcelain)" ]; then
echo "Uncommitted changes in SIA directory"
#exit 1
fi
mkdir -p "$OUTPUT_DIR"
train_gemma --output-dir "$OUTPUT_DIR"

21
tools/gemma_train/pyproject.toml Normal file
View File

@@ -0,0 +1,21 @@
[build-system]
requires = ["setuptools>=42", "wheel"]
build-backend = "setuptools.build_meta"
[project]
name = "gemma_train"
version = "0.1.0"
requires-python = ">=3.8"
dependencies = [
"accelerate>=0.26.0",
"bitsandbytes>=0.45.0",
"llama-cpp-scripts @ file:///root/sia/modules/llama.cpp",
"llm_engine_utils @ file:///root/sia/lib/llm_engine_utils",
"trl>=0.17.0",
"peft>=0.15.0",
"python-dotenv>=1.0.0",
]
[project.scripts]
gemma_train = "gemma_train.__main__:main"

153
tools/gemma_train/src/gemma_train/__main__.py Normal file
View File

@@ -0,0 +1,153 @@
from llm_engine_utils.dataset import Dataset
from pathlib import Path
from peft import LoraConfig, AutoPeftModelForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
from trl import SFTTrainer
import os
import sys
import torch
from .config import Config
def main():
config = Config()
train(config)
merge(config)
os.system(f"cp -r {config.output_dir}/tokenizer/* {config.output_dir}/merged")
convert_to_gguf(config)
def train(config: Config):
tokenizer = AutoTokenizer.from_pretrained(config.model, token=config.api_key)
tokenizer.save_pretrained(config.output_dir/"tokenizer")
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16
)
model = AutoModelForCausalLM.from_pretrained(
config.model,
quantization_config=bnb_config,
device_map="auto",
token=config.api_key,
attn_implementation='eager',
)
dataset = Dataset(config.config_path)
dataset.validate()
dataset = dataset.to_transformers_dataset(tokenizer)
lora_config = LoraConfig(
r=4,
lora_alpha=4,
target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
training_args = TrainingArguments(
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=1,
max_steps=1,
learning_rate=1e-3,
fp16=True,
logging_steps=1,
save_strategy="steps",
save_steps=1,
output_dir=config.output_dir/"lora",
optim="paged_adamw_8bit",
seed=42,
group_by_length=True,
)
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
args=training_args,
peft_config=lora_config,
formatting_func=format_sia_example,
)
trainer.train()
trainer.model.save_pretrained(config.output_dir/"lora_adapter")
def merge(config: Config):
adapted_model = AutoPeftModelForCausalLM.from_pretrained(
config.output_dir/"lora_adapter",
torch_dtype=torch.float16, # Use float16 for better compatibility
device_map="auto",
offload_folder="offload",
token=config.api_key,
)
merged_model = adapted_model.merge_and_unload()
merged_model.save_pretrained(
config.output_dir/"merged",
safe_serialization=True
)
def convert_to_gguf(config: Config):
"""Convert the merged model to GGUF format using llama.cpp's convert_hf_to_gguf script."""
print("Converting merged model to GGUF format...")
# Add path to llama.cpp directory
sys.path.append("./llama.cpp")
try:
# Import the necessary components from the conversion script
from convert_hf_to_gguf import ModelBase, gguf, ModelType
# Set up paths
dir_model = config.output_dir / "merged"
fname_out = config.output_dir / "model.gguf"
output_type = gguf.LlamaFileType.MOSTLY_Q8_0 # Using Q8_0 quantization
# Run the conversion with torch inference mode
with torch.inference_mode():
# Load hyperparameters
hparams = ModelBase.load_hparams(dir_model)
model_architecture = hparams["architectures"][0]
model_type = ModelType.TEXT
print(f"Model architecture: {model_architecture}")
try:
# Get the appropriate model class
model_class = ModelBase.from_model_architecture(model_architecture, model_type=model_type)
# Create model instance
model_instance = model_class(
dir_model,
output_type,
fname_out,
is_big_endian=False,
use_temp_file=False,
eager=False
)
# Export the model
print(f"Exporting model to GGUF format...")
model_instance.write()
print(f"Model successfully exported to {model_instance.fname_out}")
except NotImplementedError:
print(f"Error: Model architecture {model_architecture} is not supported for GGUF conversion")
print("Skipping GGUF conversion")
except ImportError as e:
print(f"Error importing conversion script: {e}")
print("Make sure llama.cpp is properly installed and accessible")
print("Skipping GGUF conversion")
except Exception as e:
print(f"Error during GGUF conversion: {e}")
print("Skipping GGUF conversion")
def format_sia_example(example):
return example['messages'].removeprefix("<bos>")
if __name__ == "__main__":
exit(main())

48
tools/gemma_train/src/gemma_train/config.py Normal file
View File

@@ -0,0 +1,48 @@
from pathlib import Path
import argparse
import os
class Config:
def __init__(self):
parser = argparse.ArgumentParser(description='Train Gemma model and convert to gguf')
parser.add_argument(
'--config',
type=Path,
default=Path('/root/sia/training/config.yaml'),
help='Path to config file (default: /root/sia/training/config.yaml)'
)
parser.add_argument(
'--model',
type=str,
default='google/gemma-3-1b-it',
help='Base model for fine-tuning (default: google/gemma-3-1b-it)'
)
parser.add_argument(
'--api-key',
type=str,
default=os.environ.get('SIA_HF_API_KEY'),
help='Huggingface API key (optional, env: SIA_HF_API_KEY)'
)
parser.add_argument(
'--output-dir',
type=Path,
default=Path('/root/models/current'),
help='Output directory for fine-tuned model and converted gguf (default: /root/models/current)'
)
self.args = parser.parse_args()
@property
def config_path(self) -> Path:
return self.args.config
@property
def model(self) -> str:
return self.args.model
@property
def api_key(self) -> str:
return self.args.api_key
@property
def output_dir(self) -> Path:
return self.args.output_dir

3
tools/notebook/requirements.txt Normal file
View File

@@ -0,0 +1,3 @@
jupyter
ipykernel
ipywidgets